Skip to content

Commit c4bd449

Browse files
showelltimabbott
authored andcommitted
peformance: Cache user mentions for multiple PMs.
It's slightly annoying to plumb Optional[MentionBackend] down the stack, but it's a one-time change. I tried to make the cache code relatively unobtrusive for the single-message use case. We should be able to eliminate redundant stream queries using similar techniques. I considered caching at the level of rendering the message itself, but this involves nearly as much plumbing, and you have to account for the fact that several users on your realm may have distinct default languages (French, Spanish, Russian, etc.), so you would not eliminate as many query hops. Also, if multiple streams were involved, users would get slightly different messages based on their prior subscriptions.
1 parent c644826 commit c4bd449

File tree

4 files changed

+72
-22
lines changed

4 files changed

+72
-22
lines changed

zerver/lib/actions.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1855,6 +1855,7 @@ def build_message_send_dict(
18551855
realm: Optional[Realm] = None,
18561856
widget_content_dict: Optional[Dict[str, Any]] = None,
18571857
email_gateway: bool = False,
1858+
mention_backend: Optional[MentionBackend] = None,
18581859
) -> SendMessageRequest:
18591860
"""Returns a dictionary that can be passed into do_send_messages. In
18601861
production, this is always called by check_message, but some
@@ -1863,7 +1864,9 @@ def build_message_send_dict(
18631864
if realm is None:
18641865
realm = message.sender.realm
18651866

1866-
mention_backend = MentionBackend(realm.id)
1867+
if mention_backend is None:
1868+
mention_backend = MentionBackend(realm.id)
1869+
18671870
mention_data = MentionData(
18681871
mention_backend=mention_backend,
18691872
content=message.content,
@@ -3303,6 +3306,7 @@ def check_message(
33033306
email_gateway: bool = False,
33043307
*,
33053308
skip_stream_access_check: bool = False,
3309+
mention_backend: Optional[MentionBackend] = None,
33063310
) -> SendMessageRequest:
33073311
"""See
33083312
https://zulip.readthedocs.io/en/latest/subsystems/sending-messages.html
@@ -3428,6 +3432,7 @@ def check_message(
34283432
realm=realm,
34293433
widget_content_dict=widget_content_dict,
34303434
email_gateway=email_gateway,
3435+
mention_backend=mention_backend,
34313436
)
34323437

34333438
if stream is not None and message_send_dict.rendering_result.mentions_wildcard:
@@ -3444,6 +3449,7 @@ def _internal_prep_message(
34443449
addressee: Addressee,
34453450
content: str,
34463451
email_gateway: bool = False,
3452+
mention_backend: Optional[MentionBackend] = None,
34473453
) -> Optional[SendMessageRequest]:
34483454
"""
34493455
Create a message object and checks it, but doesn't send it or save it to the database.
@@ -3473,6 +3479,7 @@ def _internal_prep_message(
34733479
content,
34743480
realm=realm,
34753481
email_gateway=email_gateway,
3482+
mention_backend=mention_backend,
34763483
)
34773484
except JsonableError as e:
34783485
logging.exception(
@@ -3528,7 +3535,11 @@ def internal_prep_stream_message_by_name(
35283535

35293536

35303537
def internal_prep_private_message(
3531-
realm: Realm, sender: UserProfile, recipient_user: UserProfile, content: str
3538+
realm: Realm,
3539+
sender: UserProfile,
3540+
recipient_user: UserProfile,
3541+
content: str,
3542+
mention_backend: Optional[MentionBackend] = None,
35323543
) -> Optional[SendMessageRequest]:
35333544
"""
35343545
See _internal_prep_message for details of how this works.
@@ -3540,6 +3551,7 @@ def internal_prep_private_message(
35403551
sender=sender,
35413552
addressee=addressee,
35423553
content=content,
3554+
mention_backend=mention_backend,
35433555
)
35443556

35453557

zerver/lib/mention.py

Lines changed: 51 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -37,27 +37,61 @@ def Q(self) -> Q:
3737
raise AssertionError("totally empty filter makes no sense")
3838

3939

40-
@dataclass
4140
class MentionBackend:
42-
realm_id: int
41+
def __init__(self, realm_id: int) -> None:
42+
self.realm_id = realm_id
43+
self.user_cache: Dict[Tuple[int, str], FullNameInfo] = {}
4344

4445
def get_full_name_info_list(self, user_filters: List[UserFilter]) -> List[FullNameInfo]:
45-
q_list = [user_filter.Q() for user_filter in user_filters]
46-
47-
rows = (
48-
UserProfile.objects.filter(
49-
realm_id=self.realm_id,
50-
is_active=True,
51-
)
52-
.filter(
53-
functools.reduce(lambda a, b: a | b, q_list),
46+
result: List[FullNameInfo] = []
47+
unseen_user_filters: List[UserFilter] = []
48+
49+
# Try to get messages from the user_cache first.
50+
# This loop populates two lists:
51+
# - results are the objects we pull from cache
52+
# - unseen_user_filters are filters where need to hit the DB
53+
for user_filter in user_filters:
54+
# We expect callers who take advantage of our user_cache to supply both
55+
# id and full_name in the user mentions in their messages.
56+
if user_filter.id is not None and user_filter.full_name is not None:
57+
user = self.user_cache.get((user_filter.id, user_filter.full_name), None)
58+
if user is not None:
59+
result.append(user)
60+
continue
61+
62+
# BOO! We have to go the database.
63+
unseen_user_filters.append(user_filter)
64+
65+
# Most of the time, we have to go to the database to get user info,
66+
# unless our last loop found everything in the cache.
67+
if unseen_user_filters:
68+
q_list = [user_filter.Q() for user_filter in unseen_user_filters]
69+
70+
rows = (
71+
UserProfile.objects.filter(
72+
realm_id=self.realm_id,
73+
is_active=True,
74+
)
75+
.filter(
76+
functools.reduce(lambda a, b: a | b, q_list),
77+
)
78+
.only(
79+
"id",
80+
"full_name",
81+
)
5482
)
55-
.only(
56-
"id",
57-
"full_name",
58-
)
59-
)
60-
return [FullNameInfo(id=row.id, full_name=row.full_name) for row in rows]
83+
84+
user_list = [FullNameInfo(id=row.id, full_name=row.full_name) for row in rows]
85+
86+
# We expect callers who take advantage of our cache to supply both
87+
# id and full_name in the user mentions in their messages.
88+
for user in user_list:
89+
if user.id is not None and user.full_name is not None:
90+
self.user_cache[(user.id, user.full_name)] = user
91+
92+
result += user_list
93+
94+
return result
6195

6296

6397
def user_mention_matches_wildcard(mention: str) -> bool:

zerver/tests/test_subs.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4971,7 +4971,7 @@ def test_gather_subscriptions(self) -> None:
49714971
streams,
49724972
dict(principals=orjson.dumps(users_to_subscribe).decode()),
49734973
)
4974-
self.assert_length(queries, 50)
4974+
self.assert_length(queries, 48)
49754975

49764976
msg = f"""
49774977
@**King Hamlet|{hamlet.id}** subscribed you to the following streams:

zerver/views/streams.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@
5050
OrganizationOwnerRequired,
5151
ResourceNotFoundError,
5252
)
53-
from zerver.lib.mention import silent_mention_syntax_for_user
53+
from zerver.lib.mention import MentionBackend, silent_mention_syntax_for_user
5454
from zerver.lib.request import REQ, has_request_variables
5555
from zerver.lib.response import json_success
5656
from zerver.lib.retention import parse_message_retention_days
@@ -603,6 +603,9 @@ def send_messages_for_new_subscribers(
603603

604604
newly_created_stream_names = {s.name for s in created_streams}
605605

606+
realm = user_profile.realm
607+
mention_backend = MentionBackend(realm.id)
608+
606609
# Inform the user if someone else subscribed them to stuff,
607610
# or if a new stream was created with the "announce" option.
608611
notifications = []
@@ -633,10 +636,11 @@ def send_messages_for_new_subscribers(
633636

634637
notifications.append(
635638
internal_prep_private_message(
636-
realm=user_profile.realm,
639+
realm=realm,
637640
sender=sender,
638641
recipient_user=recipient_user,
639642
content=msg,
643+
mention_backend=mention_backend,
640644
)
641645
)
642646

0 commit comments

Comments
 (0)