Speed up calculating push actions in large rooms (#13973)

We move the expensive check of visibility to after calculating push actions, avoiding the expensive check for users who won't get pushed anyway.

I think this should have a big impact on rooms with large numbers of local users that have pushed disabled.

3 files changed, 96 insertions, 12 deletions

diff --git a/changelog.d/13973.misc b/changelog.d/13973.misc
new file mode 100644
index 0000000000..58150a2b35
--- /dev/null
+++ b/changelog.d/13973.misc
@@ -0,0 +1 @@
+Speed up calculating push actions in large rooms.
diff --git a/synapse/push/bulk_push_rule_evaluator.py b/synapse/push/bulk_push_rule_evaluator.py
index 60f3129005..7bfe380543 100644
--- a/synapse/push/bulk_push_rule_evaluator.py
+++ b/synapse/push/bulk_push_rule_evaluator.py
@@ -303,20 +303,10 @@ class BulkPushRuleEvaluator:
             event.room_id, users
         )
 
-        # This is a check for the case where user joins a room without being
-        # allowed to see history, and then the server receives a delayed event
-        # from before the user joined, which they should not be pushed for
-        uids_with_visibility = await filter_event_for_clients_with_state(
-            self.store, users, event, context
-        )
-
         for uid, rules in rules_by_user.items():
             if event.sender == uid:
                 continue
 
-            if uid not in uids_with_visibility:
-                continue
-
             display_name = None
             profile = profiles.get(uid)
             if profile:
@@ -342,6 +332,21 @@ class BulkPushRuleEvaluator:
                 # Push rules say we should notify the user of this event
                 actions_by_user[uid] = actions
 
+        # This is a check for the case where user joins a room without being
+        # allowed to see history, and then the server receives a delayed event
+        # from before the user joined, which they should not be pushed for
+        #
+        # We do this *after* calculating the push actions as a) its unlikely
+        # that we'll filter anyone out and b) for large rooms its likely that
+        # most users will have push disabled and so the set of users to check is
+        # much smaller.
+        uids_with_visibility = await filter_event_for_clients_with_state(
+            self.store, actions_by_user.keys(), event, context
+        )
+
+        for user_id in set(actions_by_user).difference(uids_with_visibility):
+            actions_by_user.pop(user_id, None)
+
         # Mark in the DB staging area the push actions for users who should be
         # notified for this event. (This will then get handled when we persist
         # the event)
diff --git a/tests/push/test_push_rule_evaluator.py b/tests/push/test_push_rule_evaluator.py
index b8308cbc05..8804f0e0d3 100644
--- a/tests/push/test_push_rule_evaluator.py
+++ b/tests/push/test_push_rule_evaluator.py
@@ -19,17 +19,18 @@ import frozendict
 from twisted.test.proto_helpers import MemoryReactor
 
 import synapse.rest.admin
-from synapse.api.constants import EventTypes, Membership
+from synapse.api.constants import EventTypes, HistoryVisibility, Membership
 from synapse.api.room_versions import RoomVersions
 from synapse.appservice import ApplicationService
 from synapse.events import FrozenEvent
 from synapse.push.bulk_push_rule_evaluator import _flatten_dict
 from synapse.push.httppusher import tweaks_for_actions
+from synapse.rest import admin
 from synapse.rest.client import login, register, room
 from synapse.server import HomeServer
 from synapse.storage.databases.main.appservice import _make_exclusive_regex
 from synapse.synapse_rust.push import PushRuleEvaluator
-from synapse.types import JsonDict
+from synapse.types import JsonDict, UserID
 from synapse.util import Clock
 
 from tests import unittest
@@ -437,3 +438,80 @@ class TestBulkPushRuleEvaluator(unittest.HomeserverTestCase):
         )
 
         self.assertEqual(len(users_with_push_actions), 0)
+
+
+class BulkPushRuleEvaluatorTestCase(unittest.HomeserverTestCase):
+    servlets = [
+        admin.register_servlets,
+        login.register_servlets,
+        room.register_servlets,
+    ]
+
+    def prepare(
+        self, reactor: MemoryReactor, clock: Clock, homeserver: HomeServer
+    ) -> None:
+        self.main_store = homeserver.get_datastores().main
+
+        self.user_id1 = self.register_user("user1", "password")
+        self.tok1 = self.login(self.user_id1, "password")
+        self.user_id2 = self.register_user("user2", "password")
+        self.tok2 = self.login(self.user_id2, "password")
+
+        self.room_id = self.helper.create_room_as(tok=self.tok1)
+
+        # We want to test history visibility works correctly.
+        self.helper.send_state(
+            self.room_id,
+            EventTypes.RoomHistoryVisibility,
+            {"history_visibility": HistoryVisibility.JOINED},
+            tok=self.tok1,
+        )
+
+    def get_notif_count(self, user_id: str) -> int:
+        return self.get_success(
+            self.main_store.db_pool.simple_select_one_onecol(
+                table="event_push_actions",
+                keyvalues={"user_id": user_id},
+                retcol="COALESCE(SUM(notif), 0)",
+                desc="get_staging_notif_count",
+            )
+        )
+
+    def test_plain_message(self) -> None:
+        """Test that sending a normal message in a room will trigger a
+        notification
+        """
+
+        # Have user2 join the room and cle
+        self.helper.join(self.room_id, self.user_id2, tok=self.tok2)
+
+        # They start off with no notifications, but get them when messages are
+        # sent.
+        self.assertEqual(self.get_notif_count(self.user_id2), 0)
+
+        user1 = UserID.from_string(self.user_id1)
+        self.create_and_send_event(self.room_id, user1)
+
+        self.assertEqual(self.get_notif_count(self.user_id2), 1)
+
+    def test_delayed_message(self) -> None:
+        """Test that a delayed message that was from before a user joined
+        doesn't cause a notification for the joined user.
+        """
+        user1 = UserID.from_string(self.user_id1)
+
+        # Send a message before user2 joins
+        event_id1 = self.create_and_send_event(self.room_id, user1)
+
+        # Have user2 join the room
+        self.helper.join(self.room_id, self.user_id2, tok=self.tok2)
+
+        # They start off with no notifications
+        self.assertEqual(self.get_notif_count(self.user_id2), 0)
+
+        # Send another message that references the event before the join to
+        # simulate a "delayed" event
+        self.create_and_send_event(self.room_id, user1, prev_event_ids=[event_id1])
+
+        # user2 should not be notified about it, because they can't see it.
+        self.assertEqual(self.get_notif_count(self.user_id2), 0)