From 7bbcc2ac96010f5f2fa023d6225b111b3968045d Mon Sep 17 00:00:00 2001 From: Steve Howell Date: Tue, 20 Oct 2020 15:46:31 +0000 Subject: [PATCH] refactor: Compute peers for public streams later. This saves us a query for edge cases like when you try to unsubscribe from a public stream that you have already unsubscribed from. But this is mostly to prep for upcoming optimizations. --- zerver/lib/actions.py | 43 ++++++++++-------- zerver/lib/stream_subscription.py | 74 +++++++++++++++---------------- zerver/tests/test_subs.py | 2 +- 3 files changed, 61 insertions(+), 58 deletions(-) diff --git a/zerver/lib/actions.py b/zerver/lib/actions.py index fc8bf41f36..99448bdc91 100644 --- a/zerver/lib/actions.py +++ b/zerver/lib/actions.py @@ -121,7 +121,7 @@ from zerver.lib.sessions import delete_user_sessions from zerver.lib.storage import static_path from zerver.lib.stream_subscription import ( SubInfo, - bulk_get_peers, + bulk_get_private_peers, bulk_get_subscriber_peer_info, get_active_subscriptions_for_stream_id, get_bulk_stream_subscriber_info, @@ -2863,7 +2863,7 @@ def bulk_add_subscriptions( realm=realm, altered_user_dict=altered_user_dict, stream_dict=stream_dict, - peer_id_dict=subscriber_peer_info.peer_ids, + private_peer_dict=subscriber_peer_info.private_peer_dict, ) return ( @@ -2937,7 +2937,7 @@ def send_peer_subscriber_events( realm: Realm, stream_dict: Dict[int, Stream], altered_user_dict: Dict[int, Set[int]], - peer_id_dict: Dict[int, Set[int]], + private_peer_dict: Dict[int, Set[int]], ) -> None: # Send peer_add/peer_remove events to other users who are tracking the # subscribers lists of streams in their browser; everyone for @@ -2952,7 +2952,7 @@ def send_peer_subscriber_events( for stream_id in private_stream_ids: altered_user_ids = altered_user_dict[stream_id] - peer_user_ids = peer_id_dict[stream_id] - altered_user_ids + peer_user_ids = private_peer_dict[stream_id] - altered_user_ids if peer_user_ids: for new_user_id in altered_user_ids: @@ -2970,28 +2970,33 @@ def send_peer_subscriber_events( and not stream_dict[stream_id].is_in_zephyr_realm ] - for stream_id in public_stream_ids: - altered_user_ids = altered_user_dict[stream_id] - peer_user_ids = peer_id_dict[stream_id] - altered_user_ids + if public_stream_ids: + public_peer_ids = set(active_non_guest_user_ids(realm.id)) - if peer_user_ids: - for new_user_id in altered_user_ids: - event = dict( - type="subscription", - op=op, - stream_id=stream_id, - user_id=new_user_id, - ) - send_event(realm, event, peer_user_ids) + for stream_id in public_stream_ids: + altered_user_ids = altered_user_dict[stream_id] + peer_user_ids = public_peer_ids - altered_user_ids + + if peer_user_ids: + for new_user_id in altered_user_ids: + event = dict( + type="subscription", + op=op, + stream_id=stream_id, + user_id=new_user_id, + ) + send_event(realm, event, peer_user_ids) def send_peer_remove_events( realm: Realm, streams: List[Stream], altered_user_dict: Dict[int, Set[int]], ) -> None: - peer_id_dict = bulk_get_peers( + private_streams = [stream for stream in streams if stream.invite_only] + + private_peer_dict = bulk_get_private_peers( realm=realm, - streams=streams, + private_streams=private_streams, ) stream_dict = {stream.id: stream for stream in streams} @@ -3000,7 +3005,7 @@ def send_peer_remove_events( realm=realm, stream_dict=stream_dict, altered_user_dict=altered_user_dict, - peer_id_dict=peer_id_dict, + private_peer_dict=private_peer_dict, ) def get_available_notification_sounds() -> List[str]: diff --git a/zerver/lib/stream_subscription.py b/zerver/lib/stream_subscription.py index 1bf3622b75..a75b41adcd 100644 --- a/zerver/lib/stream_subscription.py +++ b/zerver/lib/stream_subscription.py @@ -6,14 +6,7 @@ from typing import Any, Dict, List, Optional, Set from django.db.models.query import QuerySet -from zerver.models import ( - Realm, - Recipient, - Stream, - Subscription, - UserProfile, - active_non_guest_user_ids, -) +from zerver.models import Realm, Recipient, Stream, Subscription, UserProfile @dataclass @@ -25,7 +18,7 @@ class SubInfo: @dataclass class SubscriberPeerInfo: subscribed_ids: Dict[int, Set[int]] - peer_ids: Dict[int, Set[int]] + private_peer_dict: Dict[int, Set[int]] def get_active_subscriptions_for_stream_id(stream_id: int) -> QuerySet: # TODO: Change return type to QuerySet[Subscription] @@ -134,13 +127,19 @@ def bulk_get_subscriber_peer_info( stream, which we generally send to the person subscribing to the stream. - peer_ids: + private_peer_dict: These are the folks that need to know about a new subscriber. It's usually a superset of the subscribers. + + Note that we only compute this for PRIVATE streams. We + let other code handle peers for public streams, since the + peers for all public streams are actually the same group + of users, and downstream code can use that property of + public streams to avoid extra work. """ subscribed_ids = {} - peer_ids = {} + private_peer_dict = {} private_stream_ids = {stream.id for stream in streams if stream.invite_only} public_stream_ids = {stream.id for stream in streams if not stream.invite_only} @@ -151,48 +150,47 @@ def bulk_get_subscriber_peer_info( realm_admin_ids = {user.id for user in realm.get_admin_users_and_bots()} for stream_id in private_stream_ids: + # This is the same business rule as we use in + # bulk_get_private_peers. Realm admins can see all private stream + # subscribers. subscribed_user_ids = stream_user_ids.get(stream_id, set()) subscribed_ids[stream_id] = subscribed_user_ids - peer_ids[stream_id] = subscribed_user_ids | realm_admin_ids + private_peer_dict[stream_id] = subscribed_user_ids | realm_admin_ids - if public_stream_ids: - non_guests = active_non_guest_user_ids(realm.id) - for stream_id in public_stream_ids: - subscribed_user_ids = stream_user_ids.get(stream_id, set()) - subscribed_ids[stream_id] = subscribed_user_ids - peer_ids[stream_id] = set(non_guests) + for stream_id in public_stream_ids: + subscribed_user_ids = stream_user_ids.get(stream_id, set()) + subscribed_ids[stream_id] = subscribed_user_ids return SubscriberPeerInfo( subscribed_ids=subscribed_ids, - peer_ids=peer_ids, + private_peer_dict=private_peer_dict, ) -def bulk_get_peers( +def bulk_get_private_peers( realm: Realm, - streams: List[Stream], + private_streams: List[Stream], ) -> Dict[int, Set[int]]: - # This is almost a subset of bulk_get_subscriber_peer_info, - # with the nuance that we don't have to query subscribers - # for public streams. (The other functions tries to save - # a query hop.) - peer_ids = {} + if not private_streams: + return {} - private_stream_ids = {stream.id for stream in streams if stream.invite_only} - public_stream_ids = {stream.id for stream in streams if not stream.invite_only} + for stream in private_streams: + # Our caller should only pass us private streams. + assert stream.invite_only - if private_stream_ids: - realm_admin_ids = {user.id for user in realm.get_admin_users_and_bots()} - stream_user_ids = get_user_ids_for_streams(private_stream_ids) + peer_ids: Dict[int, Set[int]] = {} - for stream_id in private_stream_ids: - subscribed_user_ids = stream_user_ids.get(stream_id, set()) - peer_ids[stream_id] = subscribed_user_ids | realm_admin_ids + realm_admin_ids = {user.id for user in realm.get_admin_users_and_bots()} - if public_stream_ids: - non_guests = active_non_guest_user_ids(realm.id) - for stream_id in public_stream_ids: - peer_ids[stream_id] = set(non_guests) + stream_ids = {stream.id for stream in private_streams} + stream_user_ids = get_user_ids_for_streams(stream_ids) + + for stream in private_streams: + # This is the same business rule as we use in + # bulk_get_subscriber_peer_info. Realm admins can see all private + # stream subscribers. + subscribed_user_ids = stream_user_ids.get(stream.id, set()) + peer_ids[stream.id] = subscribed_user_ids | realm_admin_ids return peer_ids diff --git a/zerver/tests/test_subs.py b/zerver/tests/test_subs.py index 87e9c53037..d94e252b2f 100644 --- a/zerver/tests/test_subs.py +++ b/zerver/tests/test_subs.py @@ -1677,7 +1677,7 @@ class StreamAdminTest(ZulipTestCase): fails gracefully. """ result = self.attempt_unsubscribe_of_principal( - query_count=11, + query_count=10, target_users=[self.example_user('cordelia')], is_realm_admin=True, is_subbed=False,