From 74fd508b2fc273dda3759e18175dabee3c2de755 Mon Sep 17 00:00:00 2001
From: Tim Abbott <tabbott@humbughq.com>
Date: Fri, 28 Jun 2013 11:16:55 -0400
Subject: [PATCH] Remove subscriptions using bulk queries.

This improves the performance of unsubscribing to N streams by more
than a factor of 10 for large N.

(imported from commit a529e6d3ac4452f49c2294908d275280019bbd05)
---
 zephyr/lib/actions.py | 55 +++++++++++++++++++++++++++++++++++--------
 zephyr/views.py       | 13 +++++-----
 2 files changed, 51 insertions(+), 17 deletions(-)

diff --git a/zephyr/lib/actions.py b/zephyr/lib/actions.py
index 13465cde09..cf0344af17 100644
--- a/zephyr/lib/actions.py
+++ b/zephyr/lib/actions.py
@@ -635,6 +635,50 @@ def do_add_subscription(user_profile, stream, no_log=False):
         notify_new_subscription(user_profile, stream, subscription, no_log)
     return did_subscribe
 
+def notify_subscription_removed(user_profile, stream, no_log=False):
+    if not no_log:
+        log_event({'type': 'subscription_removed',
+                   'user': user_profile.email,
+                   'name': stream.name,
+                   'domain': stream.realm.domain})
+
+    notice = dict(event=dict(type="subscription", op="remove",
+                             subscription=dict(name=stream.name)),
+                  users=[user_profile.id])
+    tornado_callbacks.send_notification(notice)
+
+def bulk_remove_subscriptions(users, streams):
+    recipients_map = bulk_get_recipients(Recipient.STREAM,
+                                         [stream.id for stream in streams])
+    stream_map = {}
+    for stream in streams:
+        stream_map[recipients_map[stream.id].id] = stream
+
+    subs_by_user = dict((user_profile.id, []) for user_profile in users)
+    for sub in Subscription.objects.select_related("user_profile").filter(user_profile__in=users,
+                                                                          recipient__in=recipients_map.values(),
+                                                                          active=True):
+        subs_by_user[sub.user_profile_id].append(sub)
+
+    subs_to_deactivate = []
+    not_subscribed = []
+    for user_profile in users:
+        recipients_to_unsub = set([recipient.id for recipient in recipients_map.values()])
+        for sub in subs_by_user[user_profile.id]:
+            recipients_to_unsub.remove(sub.recipient_id)
+            subs_to_deactivate.append((sub, stream_map[sub.recipient_id]))
+        for recipient_id in recipients_to_unsub:
+            not_subscribed.append((user_profile, stream_map[recipient_id]))
+
+    Subscription.objects.filter(id__in=[sub.id for (sub, stream_name) in
+                                        subs_to_deactivate]).update(active=False)
+
+    for (sub, stream) in subs_to_deactivate:
+        notify_subscription_removed(sub.user_profile, stream)
+
+    return ([(sub.user_profile, stream) for (sub, stream) in subs_to_deactivate],
+            not_subscribed)
+
 def do_remove_subscription(user_profile, stream, no_log=False):
     recipient = get_recipient(Recipient.STREAM, stream.id)
     maybe_sub = Subscription.objects.filter(user_profile=user_profile,
@@ -646,16 +690,7 @@ def do_remove_subscription(user_profile, stream, no_log=False):
     subscription.active = False
     subscription.save(update_fields=["active"])
     if did_remove:
-        if not no_log:
-            log_event({'type': 'subscription_removed',
-                       'user': user_profile.email,
-                       'name': stream.name,
-                       'domain': stream.realm.domain})
-
-        notice = dict(event=dict(type="subscription", op="remove",
-                                 subscription=dict(name=stream.name)),
-                      users=[user_profile.id])
-        tornado_callbacks.send_notification(notice)
+        notify_subscription_removed(user_profile, stream, no_log)
 
     return did_remove
 
diff --git a/zephyr/views.py b/zephyr/views.py
index 2bfaf5b065..8182ff6edb 100644
--- a/zephyr/views.py
+++ b/zephyr/views.py
@@ -21,7 +21,7 @@ from zephyr.models import Message, UserProfile, Stream, Subscription, \
     MAX_SUBJECT_LENGTH, get_stream, bulk_get_streams, UserPresence, \
     get_recipient, valid_stream_name, to_dict_cache_key, to_dict_cache_key_id, \
     extract_message_dict, stringify_message_dict, parse_usermessage_flags
-from zephyr.lib.actions import do_remove_subscription, \
+from zephyr.lib.actions import do_remove_subscription, bulk_remove_subscriptions, \
     do_change_password, create_mit_user_if_needed, do_change_full_name, \
     do_change_enable_desktop_notifications, do_change_enter_sends, do_change_enable_sounds, \
     do_send_confirmation_email, do_activate_user, do_create_user, check_send_message, \
@@ -1240,12 +1240,11 @@ def remove_subscriptions_backend(request, user_profile,
     streams = list_to_streams(streams_raw, user_profile)
 
     result = dict(removed=[], not_subscribed=[])
-    for stream in streams:
-        did_remove = do_remove_subscription(user_profile, stream)
-        if did_remove:
-            result["removed"].append(stream.name)
-        else:
-            result["not_subscribed"].append(stream.name)
+    (removed, not_subscribed) = bulk_remove_subscriptions([user_profile], streams)
+    for (subscriber, stream) in removed:
+        result["removed"].append(stream.name)
+    for (subscriber, stream) in not_subscribed:
+        result["not_subscribed"].append(stream.name)
 
     return json_success(result)