performance: Use SubInfo when removing subscribers.

We get two speedups:

    * The query to get existing subscribers only
      gets the two fields we need.  We no longer
      need all the overhead of user_profile
      and recipient data being returned in the
      query.

    * We avoid Django making extra hops to the
      database to get user info.
This commit is contained in:
Steve Howell
2020-10-16 12:51:01 +00:00
committed by Tim Abbott
parent 73982f6cc9
commit 6d1f9de7d3
3 changed files with 59 additions and 48 deletions

View File

@@ -3021,9 +3021,9 @@ def bulk_remove_subscriptions(users: Iterable[UserProfile],
stream_dict = {stream.id: stream for stream in streams}
existing_subs_by_user = get_bulk_stream_subscriber_info(users, stream_dict)
existing_subs_by_user = get_bulk_stream_subscriber_info(users, streams)
def get_non_subscribed_tups() -> List[Tuple[UserProfile, Stream]]:
def get_non_subscribed_subs() -> List[Tuple[UserProfile, Stream]]:
stream_ids = {stream.id for stream in streams}
not_subscribed: List[Tuple[UserProfile, Stream]] = []
@@ -3032,8 +3032,8 @@ def bulk_remove_subscriptions(users: Iterable[UserProfile],
user_sub_stream_info = existing_subs_by_user[user_profile.id]
subscribed_stream_ids = {
stream.id
for (sub, stream) in user_sub_stream_info
sub_info.stream.id
for sub_info in user_sub_stream_info
}
not_subscribed_stream_ids = stream_ids - subscribed_stream_ids
@@ -3043,17 +3043,17 @@ def bulk_remove_subscriptions(users: Iterable[UserProfile],
return not_subscribed
not_subscribed = get_non_subscribed_tups()
not_subscribed = get_non_subscribed_subs()
subs_to_deactivate: List[Tuple[Subscription, Stream]] = []
subs_to_deactivate: List[SubInfo] = []
sub_ids_to_deactivate: List[int] = []
# This loop just flattens out our data into big lists for
# bulk operations.
for tup_list in existing_subs_by_user.values():
for (sub, stream) in tup_list:
subs_to_deactivate.append((sub, stream))
sub_ids_to_deactivate.append(sub.id)
for sub_infos in existing_subs_by_user.values():
for sub_info in sub_infos:
subs_to_deactivate.append(sub_info)
sub_ids_to_deactivate.append(sub_info.sub.id)
our_realm = users[0].realm
@@ -3069,25 +3069,28 @@ def bulk_remove_subscriptions(users: Iterable[UserProfile],
# Log Subscription Activities in RealmAuditLog
event_time = timezone_now()
event_last_message_id = get_last_message_id()
all_subscription_logs: (List[RealmAuditLog]) = []
for (sub, stream) in subs_to_deactivate:
audit_log_entry = RealmAuditLog(
realm=sub.user_profile.realm,
all_subscription_logs = [
RealmAuditLog(
realm=sub.user.realm,
acting_user=acting_user,
modified_user=sub.user_profile,
modified_stream=stream,
modified_user=sub_info.user,
modified_stream=sub_info.stream,
event_last_message_id=event_last_message_id,
event_type=RealmAuditLog.SUBSCRIPTION_DEACTIVATED,
event_time=event_time)
all_subscription_logs.append(audit_log_entry)
event_time=event_time,
)
for sub in subs_to_deactivate
]
# Now since we have all log objects generated we can do a bulk insert
RealmAuditLog.objects.bulk_create(all_subscription_logs)
altered_user_dict: Dict[int, Set[int]] = defaultdict(set)
streams_by_user: Dict[int, List[Stream]] = defaultdict(list)
for (sub, stream) in subs_to_deactivate:
streams_by_user[sub.user_profile_id].append(stream)
altered_user_dict[stream.id].add(sub.user_profile_id)
for sub_info in subs_to_deactivate:
stream = sub_info.stream
streams_by_user[sub_info.user.id].append(stream)
altered_user_dict[stream.id].add(sub_info.user.id)
for user_profile in users:
if len(streams_by_user[user_profile.id]) == 0:
@@ -3116,7 +3119,7 @@ def bulk_remove_subscriptions(users: Iterable[UserProfile],
do_deactivate_stream(stream, acting_user=acting_user)
return (
[(sub.user_profile, stream) for (sub, stream) in subs_to_deactivate],
[(sub_info.user, sub_info.stream) for sub_info in subs_to_deactivate],
not_subscribed,
)

View File

@@ -2,7 +2,7 @@ import itertools
from collections import defaultdict
from dataclasses import dataclass
from operator import itemgetter
from typing import Any, Dict, List, Optional, Set, Tuple
from typing import Any, Dict, List, Optional, Set
from django.db.models.query import QuerySet
@@ -65,28 +65,36 @@ def get_stream_subscriptions_for_users(user_profiles: List[UserProfile]) -> Quer
)
def get_bulk_stream_subscriber_info(
user_profiles: List[UserProfile],
stream_dict: Dict[int, Stream]) -> Dict[int, List[Tuple[Subscription, Stream]]]:
users: List[UserProfile],
streams: List[Stream],
) -> Dict[int, List[SubInfo]]:
stream_ids = stream_dict.keys()
result: Dict[int, List[Tuple[Subscription, Stream]]] = {
user_profile.id: []
for user_profile in user_profiles
}
stream_ids = {stream.id for stream in streams}
subs = Subscription.objects.filter(
user_profile__in=user_profiles,
user_profile__in=users,
recipient__type=Recipient.STREAM,
recipient__type_id__in=stream_ids,
active=True,
).select_related('user_profile', 'recipient')
).only('user_profile_id', 'recipient_id')
stream_map = {stream.recipient_id: stream for stream in streams}
user_map = {user.id: user for user in users}
result: Dict[int, List[SubInfo]] = {user.id: [] for user in users}
for sub in subs:
user_profile_id = sub.user_profile_id
stream_id = sub.recipient.type_id
stream = stream_dict[stream_id]
result[user_profile_id].append((sub, stream))
user_id = sub.user_profile_id
user = user_map[user_id]
recipient_id = sub.recipient_id
stream = stream_map[recipient_id]
sub_info = SubInfo(
user=user,
sub=sub,
stream=stream,
)
result[user_id].append(sub_info)
return result

View File

@@ -1382,7 +1382,7 @@ class StreamAdminTest(ZulipTestCase):
those you aren't on.
"""
result = self.attempt_unsubscribe_of_principal(
query_count=21,
query_count=20,
target_users=[self.example_user('cordelia')],
is_realm_admin=True,
is_subbed=True,
@@ -1398,14 +1398,14 @@ class StreamAdminTest(ZulipTestCase):
If you're a realm admin, you can remove multiple users from a stream.
TODO: We have too many queries for this situation--each additional
user leads to 9 more queries.
user leads to 8 more queries.
"""
target_users = [
self.example_user(name)
for name in ['cordelia', 'prospero', 'iago', 'hamlet', 'ZOE']
]
result = self.attempt_unsubscribe_of_principal(
query_count=57,
query_count=52,
cache_count=9,
target_users=target_users,
is_realm_admin=True,
@@ -1423,7 +1423,7 @@ class StreamAdminTest(ZulipTestCase):
are on.
"""
result = self.attempt_unsubscribe_of_principal(
query_count=22,
query_count=21,
target_users=[self.example_user('cordelia')],
is_realm_admin=True,
is_subbed=True,
@@ -1440,7 +1440,7 @@ class StreamAdminTest(ZulipTestCase):
streams you aren't on.
"""
result = self.attempt_unsubscribe_of_principal(
query_count=22,
query_count=21,
target_users=[self.example_user('cordelia')],
is_realm_admin=True,
is_subbed=False,
@@ -1457,7 +1457,7 @@ class StreamAdminTest(ZulipTestCase):
You can remove others from public streams you're a stream administrator of.
"""
result = self.attempt_unsubscribe_of_principal(
query_count=21,
query_count=20,
target_users=[self.example_user('cordelia')],
is_realm_admin=False,
is_stream_admin=True,
@@ -1478,7 +1478,7 @@ class StreamAdminTest(ZulipTestCase):
for name in ['cordelia', 'prospero', 'othello', 'hamlet', 'ZOE']
]
result = self.attempt_unsubscribe_of_principal(
query_count=57,
query_count=52,
cache_count=9,
target_users=target_users,
is_realm_admin=False,
@@ -1496,7 +1496,7 @@ class StreamAdminTest(ZulipTestCase):
You can remove others from private streams you're a stream administrator of.
"""
result = self.attempt_unsubscribe_of_principal(
query_count=22,
query_count=21,
target_users=[self.example_user('cordelia')],
is_realm_admin=False,
is_stream_admin=True,
@@ -1518,7 +1518,7 @@ class StreamAdminTest(ZulipTestCase):
def test_admin_remove_others_from_stream_legacy_emails(self) -> None:
result = self.attempt_unsubscribe_of_principal(
query_count=21,
query_count=20,
target_users=[self.example_user('cordelia')],
is_realm_admin=True,
is_subbed=True,
@@ -1532,7 +1532,7 @@ class StreamAdminTest(ZulipTestCase):
def test_admin_remove_multiple_users_from_stream_legacy_emails(self) -> None:
result = self.attempt_unsubscribe_of_principal(
query_count=30,
query_count=28,
target_users=[self.example_user('cordelia'), self.example_user('prospero')],
is_realm_admin=True,
is_subbed=True,
@@ -3258,7 +3258,7 @@ class SubscriptionAPITest(ZulipTestCase):
get_client("website"),
)
self.assert_length(query_count, 60)
self.assert_length(query_count, 54)
self.assert_length(cache_count, 4)
peer_events = [e for e in events