From 944661a949996b55e398d28b6f836d73773a5377 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Fri, 6 Sep 2024 00:24:26 -0500
Subject: [PATCH 01/23] Pull from `room_stats_state`

---
 .../databases/main/events_bg_updates.py       | 110 +++++++++++++++---
 1 file changed, 93 insertions(+), 17 deletions(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index e20fc4471ef..1abdc577ff8 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -24,7 +24,12 @@
 
 import attr
 
-from synapse.api.constants import EventContentFields, Membership, RelationTypes
+from synapse.api.constants import (
+    EventContentFields,
+    EventTypes,
+    Membership,
+    RelationTypes,
+)
 from synapse.api.room_versions import KNOWN_ROOM_VERSIONS
 from synapse.events import EventBase, make_event_from_dict
 from synapse.storage._base import SQLBaseStore, db_to_json, make_in_list_sql_clause
@@ -1676,6 +1681,15 @@ def _get_rooms_to_update_txn(txn: LoggingTransaction) -> List[Tuple[str]]:
         most_recent_current_state_delta_stream_id = (
             await self.get_max_stream_id_in_current_state_deltas()
         )
+        # Same as `await self.get_stats_positions()` but get around the method
+        # resolution issues that crop up when we try to mixin `StatsStore` into this
+        # store.
+        stats_stream_position = await self.db_pool.simple_select_one_onecol(
+            table="stats_incremental_position",
+            keyvalues={},
+            retcol="stream_id",
+            desc="stats_incremental_position",
+        )
         for (room_id,) in rooms_to_update:
             current_state_ids_map = await self.db_pool.runInteraction(
                 "_sliding_sync_joined_rooms_bg_update._get_relevant_sliding_sync_current_state_event_ids_txn",
@@ -1689,26 +1703,88 @@ def _get_rooms_to_update_txn(txn: LoggingTransaction) -> List[Tuple[str]]:
             if not current_state_ids_map:
                 continue
 
-            fetched_events = await self.get_events(current_state_ids_map.values())
+            current_state_ids_to_fetch: List[str] = []
 
-            current_state_map: StateMap[EventBase] = {
-                state_key: fetched_events[event_id]
-                for state_key, event_id in current_state_ids_map.items()
-                # `get_events(...)` will filter out events for unknown room versions
-                if event_id in fetched_events
-            }
+            state_insert_values: SlidingSyncStateInsertValues = {}
 
-            # Even if we are joined to the room, this can happen for unknown room
-            # versions (old room versions that aren't known anymore) since
-            # `get_events(...)` will filter out events for unknown room versions
-            if not current_state_map:
-                continue
+            # Check if the current state has been updated since the stats were
+            # generated.
+            state_deltas_since_stats = await self.get_current_state_deltas_for_room(
+                room_id,
+                from_token=RoomStreamToken(stream=stats_stream_position),
+                to_token=None,
+            )
+            can_use_room_stats_state_table = True
+            for state_delta in state_deltas_since_stats:
+                # We only need to check for the state is relevant to the
+                # `sliding_sync_joined_rooms` table.
+                if (
+                    state_delta.event_type,
+                    state_delta.state_key,
+                ) in SLIDING_SYNC_RELEVANT_STATE_SET:
+                    can_use_room_stats_state_table = False
+                    # Once we find one relevant state event that changed, no need to
+                    # look any further
+                    break
+
+            # Avoid fetching full events out as much as possible. We can look-up all of
+            # the relevant state values except for the `tombstone_successor_room_id`
+            # from the `room_stats_state` table.
+            if can_use_room_stats_state_table:
+                (name, room_type, encryption) = cast(
+                    Tuple[Optional[str], Optional[str], Optional[str]],
+                    await self.db_pool.simple_select_one(
+                        table="room_stats_state",
+                        keyvalues={"room_id": room_id},
+                        retcols=["name", "room_type", "encryption"],
+                    ),
+                )
+
+                state_insert_values["room_name"] = name
+                state_insert_values["room_type"] = room_type
+                state_insert_values["is_encrypted"] = encryption is not None
 
-            state_insert_values = (
-                PersistEventsStore._get_sliding_sync_insert_values_from_state_map(
-                    current_state_map
+                # If the room has a tombstone, we still need to fetch it out manually
+                # because this information isn't available in the `room_stats_state`
+                # table.
+                tombstone_event_id = current_state_ids_map.get(
+                    (EventTypes.Tombstone, "")
                 )
-            )
+                if tombstone_event_id is not None:
+                    current_state_ids_to_fetch.append(tombstone_event_id)
+
+            else:
+                current_state_ids_to_fetch.extend(current_state_ids_map.values())
+
+            # Fetch any more events that we need to complete the `state_insert_values`.
+            # For the happy-path, this will just be the tombstone event if the room has
+            # one since we got the rest of the values already from the
+            # `room_stats_state` table. In the case, where we can't use
+            # `room_stats_state`, we just fetch all of the events manually.
+            if current_state_ids_to_fetch:
+                fetched_events = await self.get_events(current_state_ids_map.values())
+
+                # Even if we are joined to the room, `get_events(...)` will filter out
+                # events for unknown room versions that we no longer support. We don't need
+                # to do anything for these unsupported rooms.
+                if not fetched_events:
+                    continue
+
+                current_state_map: StateMap[EventBase] = {
+                    state_key: fetched_events[event_id]
+                    for state_key, event_id in current_state_ids_map.items()
+                    # `get_events(...)` will filter out events for unknown room versions
+                    if event_id in fetched_events
+                }
+
+                more_state_insert_values = (
+                    PersistEventsStore._get_sliding_sync_insert_values_from_state_map(
+                        current_state_map
+                    )
+                )
+                # Update our main insert map
+                state_insert_values.update(more_state_insert_values)
+
             # We should have some insert values for each room, even if they are `None`
             assert state_insert_values
             joined_room_updates[room_id] = state_insert_values

From 5894fc200c639c2c4e0832c8928cd33916b9950d Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Fri, 6 Sep 2024 00:33:55 -0500
Subject: [PATCH 02/23] `concurrently_execute` data fetching

---
 .../databases/main/events_bg_updates.py       | 33 +++++++++++--------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index 1abdc577ff8..bc2cb6b4305 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -57,6 +57,7 @@
 from synapse.types.handlers import SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES
 from synapse.types.state import StateFilter
 from synapse.util import json_encoder
+from synapse.util.async_helpers import concurrently_execute
 from synapse.util.iterutils import batch_iter
 
 if TYPE_CHECKING:
@@ -1659,23 +1660,17 @@ def _get_rooms_to_update_txn(txn: LoggingTransaction) -> List[Tuple[str]]:
 
             return rooms_to_update_rows
 
-        rooms_to_update = await self.db_pool.runInteraction(
+        rooms_to_update_rows = await self.db_pool.runInteraction(
             "_sliding_sync_joined_rooms_bg_update._get_rooms_to_update_txn",
             _get_rooms_to_update_txn,
         )
 
-        if not rooms_to_update:
+        if not rooms_to_update_rows:
             await self.db_pool.updates._end_background_update(
                 _BackgroundUpdates.SLIDING_SYNC_JOINED_ROOMS_BG_UPDATE
             )
             return 0
 
-        # Map from room_id to insert/update state values in the `sliding_sync_joined_rooms` table.
-        joined_room_updates: Dict[str, SlidingSyncStateInsertValues] = {}
-        # Map from room_id to stream_ordering/bump_stamp, etc values
-        joined_room_stream_ordering_updates: Dict[
-            str, _JoinedRoomStreamOrderingUpdate
-        ] = {}
         # As long as we get this value before we fetch the current state, we can use it
         # to check if something has changed since that point.
         most_recent_current_state_delta_stream_id = (
@@ -1690,7 +1685,15 @@ def _get_rooms_to_update_txn(txn: LoggingTransaction) -> List[Tuple[str]]:
             retcol="stream_id",
             desc="stats_incremental_position",
         )
-        for (room_id,) in rooms_to_update:
+
+        # Map from room_id to insert/update state values in the `sliding_sync_joined_rooms` table.
+        joined_room_updates: Dict[str, SlidingSyncStateInsertValues] = {}
+        # Map from room_id to stream_ordering/bump_stamp, etc values
+        joined_room_stream_ordering_updates: Dict[
+            str, _JoinedRoomStreamOrderingUpdate
+        ] = {}
+
+        async def handle_room(room_id: str) -> None:
             current_state_ids_map = await self.db_pool.runInteraction(
                 "_sliding_sync_joined_rooms_bg_update._get_relevant_sliding_sync_current_state_event_ids_txn",
                 PersistEventsStore._get_relevant_sliding_sync_current_state_event_ids_txn,
@@ -1701,7 +1704,7 @@ def _get_rooms_to_update_txn(txn: LoggingTransaction) -> List[Tuple[str]]:
             # `sliding_sync_joined_rooms` table so we should skip and b) we won't have
             # any `current_state_events` for the room.
             if not current_state_ids_map:
-                continue
+                return
 
             current_state_ids_to_fetch: List[str] = []
 
@@ -1768,7 +1771,7 @@ def _get_rooms_to_update_txn(txn: LoggingTransaction) -> List[Tuple[str]]:
                 # events for unknown room versions that we no longer support. We don't need
                 # to do anything for these unsupported rooms.
                 if not fetched_events:
-                    continue
+                    return
 
                 current_state_map: StateMap[EventBase] = {
                     state_key: fetched_events[event_id]
@@ -1824,6 +1827,10 @@ def _get_rooms_to_update_txn(txn: LoggingTransaction) -> List[Tuple[str]]:
                 )
             )
 
+        await concurrently_execute(
+            handle_room, [row[0] for row in rooms_to_update_rows], 10
+        )
+
         def _fill_table_txn(txn: LoggingTransaction) -> None:
             # Handle updating the `sliding_sync_joined_rooms` table
             #
@@ -1897,14 +1904,14 @@ def _fill_table_txn(txn: LoggingTransaction) -> None:
                 txn,
                 table="sliding_sync_joined_rooms_to_recalculate",
                 keys=("room_id",),
-                values=rooms_to_update,
+                values=rooms_to_update_rows,
             )
 
         await self.db_pool.runInteraction(
             "sliding_sync_joined_rooms_bg_update", _fill_table_txn
         )
 
-        return len(rooms_to_update)
+        return len(rooms_to_update_rows)
 
     async def _sliding_sync_membership_snapshots_bg_update(
         self, progress: JsonDict, batch_size: int

From b4fcbc968622a1fc0c7651e143fbd0bf03f6dfe3 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Fri, 6 Sep 2024 00:48:48 -0500
Subject: [PATCH 03/23] `concurrently_execute` inserting

---
 .../databases/main/events_bg_updates.py       | 150 +++++++++---------
 1 file changed, 76 insertions(+), 74 deletions(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index bc2cb6b4305..940d5570ad2 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -1827,88 +1827,90 @@ async def handle_room(room_id: str) -> None:
                 )
             )
 
-        await concurrently_execute(
-            handle_room, [row[0] for row in rooms_to_update_rows], 10
-        )
-
-        def _fill_table_txn(txn: LoggingTransaction) -> None:
-            # Handle updating the `sliding_sync_joined_rooms` table
-            #
-            for (
-                room_id,
-                update_map,
-            ) in joined_room_updates.items():
-                joined_room_stream_ordering_update = (
-                    joined_room_stream_ordering_updates[room_id]
-                )
-                event_stream_ordering = (
-                    joined_room_stream_ordering_update.most_recent_event_stream_ordering
-                )
-                bump_stamp = joined_room_stream_ordering_update.most_recent_bump_stamp
-
-                # Check if the current state has been updated since we gathered it.
-                # We're being careful not to insert/overwrite with stale data.
-                state_deltas_since_we_gathered_current_state = (
-                    self.get_current_state_deltas_for_room_txn(
-                        txn,
-                        room_id,
-                        from_token=RoomStreamToken(
-                            stream=most_recent_current_state_delta_stream_id
-                        ),
-                        to_token=None,
-                    )
-                )
-                for state_delta in state_deltas_since_we_gathered_current_state:
-                    # We only need to check for the state is relevant to the
-                    # `sliding_sync_joined_rooms` table.
-                    if (
-                        state_delta.event_type,
-                        state_delta.state_key,
-                    ) in SLIDING_SYNC_RELEVANT_STATE_SET:
-                        # Raising exception so we can just exit and try again. It would
-                        # be hard to resolve this within the transaction because we need
-                        # to get full events out that take redactions into account. We
-                        # could add some retry logic here, but it's easier to just let
-                        # the background update try again.
-                        raise Exception(
-                            "Current state was updated after we gathered it to update "
-                            + "`sliding_sync_joined_rooms` in the background update. "
-                            + "Raising exception so we can just try again."
-                        )
+        rooms_to_update = [row[0] for row in rooms_to_update_rows]
+        await concurrently_execute(handle_room, rooms_to_update, 10)
+
+        def _fill_table_for_room_txn(txn: LoggingTransaction, room_id: str) -> None:
+            update_map = joined_room_updates[room_id]
+            joined_room_stream_ordering_update = joined_room_stream_ordering_updates[
+                room_id
+            ]
+            event_stream_ordering = (
+                joined_room_stream_ordering_update.most_recent_event_stream_ordering
+            )
+            bump_stamp = joined_room_stream_ordering_update.most_recent_bump_stamp
 
-                # Since we fully insert rows into `sliding_sync_joined_rooms`, we can
-                # just do everything on insert and `ON CONFLICT DO NOTHING`.
-                #
-                self.db_pool.simple_upsert_txn(
+            # Check if the current state has been updated since we gathered it.
+            # We're being careful not to insert/overwrite with stale data.
+            state_deltas_since_we_gathered_current_state = (
+                self.get_current_state_deltas_for_room_txn(
                     txn,
-                    table="sliding_sync_joined_rooms",
-                    keyvalues={"room_id": room_id},
-                    values={},
-                    insertion_values={
-                        **update_map,
-                        # The reason we're only *inserting* (not *updating*) `event_stream_ordering`
-                        # and `bump_stamp` is because if they are present, that means they are already
-                        # up-to-date.
-                        "event_stream_ordering": event_stream_ordering,
-                        "bump_stamp": bump_stamp,
-                    },
+                    room_id,
+                    from_token=RoomStreamToken(
+                        stream=most_recent_current_state_delta_stream_id
+                    ),
+                    to_token=None,
                 )
+            )
+            for state_delta in state_deltas_since_we_gathered_current_state:
+                # We only need to check for the state is relevant to the
+                # `sliding_sync_joined_rooms` table.
+                if (
+                    state_delta.event_type,
+                    state_delta.state_key,
+                ) in SLIDING_SYNC_RELEVANT_STATE_SET:
+                    # Raising exception so we can just exit and try again. It would
+                    # be hard to resolve this within the transaction because we need
+                    # to get full events out that take redactions into account. We
+                    # could add some retry logic here, but it's easier to just let
+                    # the background update try again.
+                    raise Exception(
+                        "Current state was updated after we gathered it to update "
+                        + "`sliding_sync_joined_rooms` in the background update. "
+                        + "Raising exception so we can just try again."
+                    )
 
-            # Now that we've processed all the room, we can remove them from the
-            # queue.
+            # Handle updating the `sliding_sync_joined_rooms` table
             #
-            # Note: we need to remove all the rooms from the queue we pulled out
-            # from the DB, not just the ones we've processed above. Otherwise
-            # we'll simply keep pulling out the same rooms over and over again.
-            self.db_pool.simple_delete_many_batch_txn(
+            # Since we fully insert rows into `sliding_sync_joined_rooms`, we can
+            # just do everything on insert and `ON CONFLICT DO NOTHING`.
+            #
+            self.db_pool.simple_upsert_txn(
                 txn,
-                table="sliding_sync_joined_rooms_to_recalculate",
-                keys=("room_id",),
-                values=rooms_to_update_rows,
+                table="sliding_sync_joined_rooms",
+                keyvalues={"room_id": room_id},
+                values={},
+                insertion_values={
+                    **update_map,
+                    # The reason we're only *inserting* (not *updating*) `event_stream_ordering`
+                    # and `bump_stamp` is because if they are present, that means they are already
+                    # up-to-date.
+                    "event_stream_ordering": event_stream_ordering,
+                    "bump_stamp": bump_stamp,
+                },
             )
 
-        await self.db_pool.runInteraction(
-            "sliding_sync_joined_rooms_bg_update", _fill_table_txn
+        async def fill_table_for_room(room_id: str) -> None:
+            await self.db_pool.runInteraction(
+                "sliding_sync_joined_rooms_bg_update", _fill_table_for_room_txn, room_id
+            )
+
+        # Since most of the time spent here is probably just latency to/from the
+        # database, let's just run this concurrently.
+        await concurrently_execute(fill_table_for_room, joined_room_updates.keys(), 10)
+
+        # Now that we've processed all of the rooms, we can remove them from the
+        # queue.
+        #
+        # Note: we need to remove all the rooms from the queue we pulled out
+        # from the DB, not just the ones we've processed above. Otherwise
+        # we'll simply keep pulling out the same rooms over and over again.
+        await self.db_pool.simple_delete_many(
+            table="sliding_sync_joined_rooms_to_recalculate",
+            column="room_id",
+            iterable=rooms_to_update,
+            keyvalues={},
+            desc="_sliding_sync_joined_rooms_bg_update: removing rooms that we processed",
         )
 
         return len(rooms_to_update_rows)

From ce7e1b4e67eaf605c442deb6d960e49b79d4d160 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Fri, 6 Sep 2024 01:48:39 -0500
Subject: [PATCH 04/23] Look for previous snapshot

---
 .../databases/main/events_bg_updates.py       | 75 +++++++++++++++++--
 1 file changed, 68 insertions(+), 7 deletions(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index 940d5570ad2..f6c03a06508 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -2051,7 +2051,7 @@ def _find_memberships_to_update_txn(
                 )
                 return 0
 
-        def _find_previous_membership_txn(
+        def _find_previous_invite_knock_membership_txn(
             txn: LoggingTransaction, room_id: str, user_id: str, event_id: str
         ) -> Tuple[str, str]:
             # Find the previous invite/knock event before the leave event
@@ -2109,9 +2109,19 @@ def _find_previous_membership_txn(
         to_insert_membership_snapshots: Dict[
             Tuple[str, str], SlidingSyncMembershipSnapshotSharedInsertValues
         ] = {}
+        # Map from (room_id, user_id) to ...
         to_insert_membership_infos: Dict[
             Tuple[str, str], SlidingSyncMembershipInfoWithEventPos
         ] = {}
+
+        # Map from room_id to ...
+        # Just some convenience maps for easier lookup by `room_id`.
+        last_to_insert_membership_snapshots_by_room_id: Dict[
+            str, SlidingSyncMembershipSnapshotSharedInsertValues
+        ] = {}
+        last_to_insert_membership_infos_by_room_id: Dict[
+            str, SlidingSyncMembershipInfoWithEventPos
+        ] = {}
         for (
             room_id,
             room_id_from_rooms_table,
@@ -2153,9 +2163,56 @@ def _find_previous_membership_txn(
                     },
                 )
 
+            # Check if the current state has been updated since the last snapshot we inserted.
+            can_use_last_snapshot = False
+            # TODO: We could also look in the database to see if a snapshot is available
+            last_to_insert_membership_snapshot = (
+                last_to_insert_membership_snapshots_by_room_id.get(room_id)
+            )
+            last_to_insert_membership_info = (
+                last_to_insert_membership_infos_by_room_id.get(room_id)
+            )
+            if last_to_insert_membership_snapshot is not None:
+                can_use_last_snapshot = True
+
+                # These should go hand-in-hand
+                assert last_to_insert_membership_info is not None
+
+                state_deltas_since_last_snapshot = await self.get_current_state_deltas_for_room(
+                    room_id,
+                    # From the last snapshot we inserted
+                    from_token=RoomStreamToken(
+                        stream=last_to_insert_membership_info.membership_event_stream_ordering,
+                    ),
+                    # To our current membership position
+                    to_token=RoomStreamToken(
+                        stream=membership_event_stream_ordering,
+                    ),
+                )
+                for state_delta in state_deltas_since_last_snapshot:
+                    # We only need to check for the state is relevant to the
+                    # `sliding_sync_joined_rooms` table.
+                    if (
+                        state_delta.event_type,
+                        state_delta.state_key,
+                    ) in SLIDING_SYNC_RELEVANT_STATE_SET:
+                        can_use_last_snapshot = False
+                        # Once we find one relevant state event that changed, no need to
+                        # look any further
+                        break
+
             # Map of values to insert/update in the `sliding_sync_membership_snapshots` table
             sliding_sync_membership_snapshots_insert_map: SlidingSyncMembershipSnapshotSharedInsertValues = {}
-            if membership == Membership.JOIN:
+            # If none of the relevant state has changed since the last snapshot, we can
+            # just re-use the last snapshot we inserted.
+            if can_use_last_snapshot:
+                # Based on the logic above, we should have a last snapshot to work from
+                assert last_to_insert_membership_snapshot is not None
+
+                sliding_sync_membership_snapshots_insert_map = (
+                    last_to_insert_membership_snapshot
+                )
+            elif membership == Membership.JOIN:
                 # If we're still joined, we can pull from current state.
                 current_state_ids_map: StateMap[
                     str
@@ -2240,8 +2297,8 @@ def _find_previous_membership_txn(
                         invite_or_knock_event_id,
                         invite_or_knock_membership,
                     ) = await self.db_pool.runInteraction(
-                        "sliding_sync_membership_snapshots_bg_update._find_previous_membership",
-                        _find_previous_membership_txn,
+                        "sliding_sync_membership_snapshots_bg_update._find_previous_invite_knock_membership_txn",
+                        _find_previous_invite_knock_membership_txn,
                         room_id,
                         user_id,
                         membership_event_id,
@@ -2327,10 +2384,14 @@ def _find_previous_membership_txn(
                     f"Unexpected membership {membership} ({membership_event_id}) that we don't know how to handle yet"
                 )
 
-            to_insert_membership_snapshots[(room_id, user_id)] = (
+            to_insert_membership_snapshots[
+                (room_id, user_id)
+            ] = last_to_insert_membership_snapshots_by_room_id[room_id] = (
                 sliding_sync_membership_snapshots_insert_map
             )
-            to_insert_membership_infos[(room_id, user_id)] = (
+            to_insert_membership_infos[
+                (room_id, user_id)
+            ] = last_to_insert_membership_infos_by_room_id[room_id] = (
                 SlidingSyncMembershipInfoWithEventPos(
                     user_id=user_id,
                     sender=sender,
@@ -2348,7 +2409,7 @@ def _fill_table_txn(txn: LoggingTransaction) -> None:
             #
             for key, insert_map in to_insert_membership_snapshots.items():
                 room_id, user_id = key
-                membership_info = to_insert_membership_infos[key]
+                membership_info = to_insert_membership_infos[(room_id, user_id)]
                 sender = membership_info.sender
                 membership_event_id = membership_info.membership_event_id
                 membership = membership_info.membership

From 27caa47c50871f05402975a6f673055e3aa2d1ed Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Fri, 6 Sep 2024 02:34:07 -0500
Subject: [PATCH 05/23] Work on snapshots background update

---
 synapse/storage/databases/main/events.py      |   4 +-
 .../databases/main/events_bg_updates.py       | 151 ++++++++++++------
 2 files changed, 101 insertions(+), 54 deletions(-)

diff --git a/synapse/storage/databases/main/events.py b/synapse/storage/databases/main/events.py
index d423d80efa7..8eff4d1791e 100644
--- a/synapse/storage/databases/main/events.py
+++ b/synapse/storage/databases/main/events.py
@@ -137,7 +137,7 @@ class SlidingSyncStateInsertValues(TypedDict, total=False):
     """
 
     room_type: Optional[str]
-    is_encrypted: Optional[bool]
+    is_encrypted: bool
     room_name: Optional[str]
     tombstone_successor_room_id: Optional[str]
 
@@ -150,7 +150,7 @@ class SlidingSyncMembershipSnapshotSharedInsertValues(
     multiple memberships
     """
 
-    has_known_state: Optional[bool]
+    has_known_state: bool
 
 
 @attr.s(slots=True, auto_attribs=True)
diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index f6c03a06508..ff4f064f162 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -1881,10 +1881,9 @@ def _fill_table_for_room_txn(txn: LoggingTransaction, room_id: str) -> None:
                 keyvalues={"room_id": room_id},
                 values={},
                 insertion_values={
+                    # The reason we're only *inserting* (not *updating*) is because if
+                    # they are present, that means they are already up-to-date.
                     **update_map,
-                    # The reason we're only *inserting* (not *updating*) `event_stream_ordering`
-                    # and `bump_stamp` is because if they are present, that means they are already
-                    # up-to-date.
                     "event_stream_ordering": event_stream_ordering,
                     "bump_stamp": bump_stamp,
                 },
@@ -2122,6 +2121,7 @@ def _find_previous_invite_knock_membership_txn(
         last_to_insert_membership_infos_by_room_id: Dict[
             str, SlidingSyncMembershipInfoWithEventPos
         ] = {}
+        # TODO: `concurrently_execute` based on buckets of room_ids
         for (
             room_id,
             room_id_from_rooms_table,
@@ -2226,6 +2226,9 @@ def _find_previous_invite_knock_membership_txn(
                     # we only want some non-membership state
                     await_full_state=False,
                 )
+
+                # TODO: Read the state from the `room_stats_state` table if we can
+
                 # We're iterating over rooms that we are joined to so they should
                 # have `current_state_events` and we should have some current state
                 # for each room
@@ -2404,59 +2407,103 @@ def _find_previous_invite_knock_membership_txn(
                 )
             )
 
+        # Assemble data so it's ready for the batch queries in the transaction
+        key_names = ("room_id", "user_id")
+        key_values: List[Tuple[str, str]] = []
+        insertion_value_names = (
+            "has_known_state",
+            "room_type",
+            "is_encrypted",
+            "room_name",
+            "tombstone_successor_room_id",
+            "sender",
+            "membership_event_id",
+            "membership",
+            "event_stream_ordering",
+            "event_instance_name",
+        )
+        insertion_value_values: List[
+            Tuple[
+                bool,
+                Optional[str],
+                bool,
+                Optional[str],
+                Optional[str],
+                str,
+                str,
+                str,
+                int,
+                str,
+            ]
+        ] = []
+        forgotten_update_query_args: List[Tuple[str, str, str]] = []
+        for key, insert_map in to_insert_membership_snapshots.items():
+            room_id, user_id = key
+            membership_info = to_insert_membership_infos[(room_id, user_id)]
+            sender = membership_info.sender
+            membership_event_id = membership_info.membership_event_id
+            membership = membership_info.membership
+            membership_event_stream_ordering = (
+                membership_info.membership_event_stream_ordering
+            )
+            membership_event_instance_name = (
+                membership_info.membership_event_instance_name
+            )
+
+            key_values.append((room_id, user_id))
+            insertion_value_values.append(
+                (
+                    # `has_known_state` should be set to *some* True/False value
+                    insert_map["has_known_state"],
+                    insert_map.get("room_type"),
+                    insert_map.get("is_encrypted", False),
+                    insert_map.get("room_name"),
+                    insert_map.get("tombstone_successor_room_id"),
+                    sender,
+                    membership_event_id,
+                    membership,
+                    membership_event_stream_ordering,
+                    membership_event_instance_name,
+                )
+            )
+
+            forgotten_update_query_args.append(
+                (
+                    membership_event_id,
+                    room_id,
+                    user_id,
+                )
+            )
+
         def _fill_table_txn(txn: LoggingTransaction) -> None:
             # Handle updating the `sliding_sync_membership_snapshots` table
             #
-            for key, insert_map in to_insert_membership_snapshots.items():
-                room_id, user_id = key
-                membership_info = to_insert_membership_infos[(room_id, user_id)]
-                sender = membership_info.sender
-                membership_event_id = membership_info.membership_event_id
-                membership = membership_info.membership
-                membership_event_stream_ordering = (
-                    membership_info.membership_event_stream_ordering
-                )
-                membership_event_instance_name = (
-                    membership_info.membership_event_instance_name
-                )
+            # We don't need to update in the upsert the state because we never partially
+            # insert/update the snapshots and anything already there is up-to-date
+            # EXCEPT for the `forgotten` field since that is updated out-of-band from
+            # the membership changes. We're using an upsert to avoid unique
+            # violation errors that would happen from directly inserting.
+            self.db_pool.simple_upsert_many_txn(
+                txn,
+                table="sliding_sync_membership_snapshots",
+                key_names=key_names,
+                key_values=key_values,
+                # TODO: Implement these
+                insertion_value_names=insertion_value_names,
+                insertion_value_values=insertion_value_values,
+            )
 
-                # We don't need to upsert the state because we never partially
-                # insert/update the snapshots and anything already there is up-to-date
-                # EXCEPT for the `forgotten` field since that is updated out-of-band
-                # from the membership changes.
-                #
-                # Even though we're only doing insertions, we're using
-                # `simple_upsert_txn()` here to avoid unique violation errors that would
-                # happen from `simple_insert_txn()`
-                self.db_pool.simple_upsert_txn(
-                    txn,
-                    table="sliding_sync_membership_snapshots",
-                    keyvalues={"room_id": room_id, "user_id": user_id},
-                    values={},
-                    insertion_values={
-                        **insert_map,
-                        "sender": sender,
-                        "membership_event_id": membership_event_id,
-                        "membership": membership,
-                        "event_stream_ordering": membership_event_stream_ordering,
-                        "event_instance_name": membership_event_instance_name,
-                    },
-                )
-                # We need to find the `forgotten` value during the transaction because
-                # we can't risk inserting stale data.
-                txn.execute(
-                    """
-                    UPDATE sliding_sync_membership_snapshots
-                    SET
-                        forgotten = (SELECT forgotten FROM room_memberships WHERE event_id = ?)
-                    WHERE room_id = ? and user_id = ?
-                    """,
-                    (
-                        membership_event_id,
-                        room_id,
-                        user_id,
-                    ),
-                )
+            # We need to find the `forgotten` value during the transaction because
+            # we can't risk inserting stale data.
+            txn.execute_batch(
+                """
+                UPDATE sliding_sync_membership_snapshots
+                SET
+                    forgotten = (SELECT forgotten FROM room_memberships WHERE event_id = ?)
+                WHERE room_id = ? and user_id = ?
+                """,
+                forgotten_update_query_args,
+            )
 
         await self.db_pool.runInteraction(
             "sliding_sync_membership_snapshots_bg_update", _fill_table_txn

From 333b4721bef661fdee31fd419d8d8e5f0d4ecafa Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Fri, 6 Sep 2024 15:26:13 -0500
Subject: [PATCH 06/23] TODO and asserts

---
 synapse/storage/databases/main/events_bg_updates.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index ff4f064f162..b4acfcf7853 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -2178,6 +2178,14 @@ def _find_previous_invite_knock_membership_txn(
                 # These should go hand-in-hand
                 assert last_to_insert_membership_info is not None
 
+                # Sanity check that our tokens are in order. We should be iterating in
+                # ascending order.
+                assert (
+                    last_to_insert_membership_info.membership_event_stream_ordering
+                    <= membership_event_stream_ordering
+                )
+
+                # TODO: This only works if someone is from the server is participating in the room.
                 state_deltas_since_last_snapshot = await self.get_current_state_deltas_for_room(
                     room_id,
                     # From the last snapshot we inserted

From f74106bc71ea90604bbbe7b180e670e4d0370274 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Mon, 9 Sep 2024 18:44:15 -0500
Subject: [PATCH 07/23] Implement
 `insertion_value_names`/`insertion_value_values` for
 `simple_upsert_many_txn(...)`

---
 synapse/storage/database.py                   | 68 +++++++++++++++--
 .../databases/main/events_bg_updates.py       |  3 +-
 tests/storage/test__base.py                   | 64 ++++++++++++++++
 tests/storage/test_base.py                    | 75 ++++++++++++++++++-
 4 files changed, 202 insertions(+), 8 deletions(-)

diff --git a/synapse/storage/database.py b/synapse/storage/database.py
index cb4a5857bed..17aacf69806 100644
--- a/synapse/storage/database.py
+++ b/synapse/storage/database.py
@@ -1472,6 +1472,11 @@ async def simple_upsert_many(
         key_values: Collection[Collection[Any]],
         value_names: Collection[str],
         value_values: Collection[Collection[Any]],
+        # Given these are the same type as the normal values, force keyword-only so that
+        # they can't be confused.
+        *,
+        insertion_value_names: Collection[str] = [],
+        insertion_value_values: Collection[Iterable[Any]] = [],
         desc: str,
     ) -> None:
         """
@@ -1497,6 +1502,8 @@ async def simple_upsert_many(
             key_values,
             value_names,
             value_values,
+            insertion_value_names=insertion_value_names,
+            insertion_value_values=insertion_value_values,
             db_autocommit=autocommit,
         )
 
@@ -1508,6 +1515,11 @@ def simple_upsert_many_txn(
         key_values: Collection[Iterable[Any]],
         value_names: Collection[str],
         value_values: Collection[Iterable[Any]],
+        # Given these are the same type as the normal values, force keyword-only so that
+        # they can't be confused.
+        *,
+        insertion_value_names: Collection[str] = [],
+        insertion_value_values: Collection[Iterable[Any]] = [],
     ) -> None:
         """
         Upsert, many times.
@@ -1519,6 +1531,9 @@ def simple_upsert_many_txn(
             value_names: The value column names
             value_values: A list of each row's value column values.
                 Ignored if value_names is empty.
+            insertion_value_names: The value column names to use only when inserting
+            insertion_value_values: A list of each row's value column values to use only
+                when inserting. Ignored if `insertion_value_names` is empty.
         """
         # If there's nothing to upsert, then skip executing the query.
         if not key_values:
@@ -1528,14 +1543,30 @@ def simple_upsert_many_txn(
         # zip() works correctly.
         if not value_names:
             value_values = [() for x in range(len(key_values))]
-        elif len(value_values) != len(key_values):
+        elif len(key_values) != len(value_values):
             raise ValueError(
                 f"{len(key_values)} key rows and {len(value_values)} value rows: should be the same number."
             )
 
+        # No value columns, therefore make a blank list so that the following
+        # zip() works correctly.
+        if not insertion_value_names:
+            insertion_value_values = [() for x in range(len(key_values))]
+        elif len(key_values) != len(insertion_value_values):
+            raise ValueError(
+                f"{len(key_values)} key rows and {len(insertion_value_values)} insertion value rows: should be the same number."
+            )
+
         if table not in self._unsafe_to_upsert_tables:
             return self.simple_upsert_many_txn_native_upsert(
-                txn, table, key_names, key_values, value_names, value_values
+                txn,
+                table,
+                key_names,
+                key_values,
+                value_names,
+                value_values,
+                insertion_value_names=insertion_value_names,
+                insertion_value_values=insertion_value_values,
             )
         else:
             return self.simple_upsert_many_txn_emulated(
@@ -1545,6 +1576,9 @@ def simple_upsert_many_txn(
                 key_values,
                 value_names,
                 value_values,
+                # TODO
+                # insertion_value_names=insertion_value_names,
+                # insertion_value_values=insertion_value_values,
             )
 
     def simple_upsert_many_txn_emulated(
@@ -1555,6 +1589,11 @@ def simple_upsert_many_txn_emulated(
         key_values: Collection[Iterable[Any]],
         value_names: Collection[str],
         value_values: Iterable[Iterable[Any]],
+        # Given these are the same type as the normal values, force keyword-only so that
+        # they can't be confused.
+        *,
+        insertion_value_names: Collection[str] = [],
+        insertion_value_values: Collection[Iterable[Any]] = [],
     ) -> None:
         """
         Upsert, many times, but without native UPSERT support or batching.
@@ -1566,6 +1605,9 @@ def simple_upsert_many_txn_emulated(
             value_names: The value column names
             value_values: A list of each row's value column values.
                 Ignored if value_names is empty.
+            insertion_value_names: The value column names to use only when inserting
+            insertion_value_values: A list of each row's value column values to use only
+                when inserting. Ignored if `insertion_value_names` is empty.
         """
 
         # Lock the table just once, to prevent it being done once per row.
@@ -1573,11 +1615,16 @@ def simple_upsert_many_txn_emulated(
         # the lock is held for the remainder of the current transaction.
         self.engine.lock_table(txn, table)
 
-        for keyv, valv in zip(key_values, value_values):
+        for keyv, valv, insertionv in zip(
+            key_values, value_values, insertion_value_values
+        ):
             _keys = dict(zip(key_names, keyv))
             _vals = dict(zip(value_names, valv))
+            _insertion_vals = dict(zip(insertion_value_names, insertionv))
 
-            self.simple_upsert_txn_emulated(txn, table, _keys, _vals, lock=False)
+            self.simple_upsert_txn_emulated(
+                txn, table, _keys, _vals, insertion_values=_insertion_vals, lock=False
+            )
 
     @staticmethod
     def simple_upsert_many_txn_native_upsert(
@@ -1587,6 +1634,11 @@ def simple_upsert_many_txn_native_upsert(
         key_values: Collection[Iterable[Any]],
         value_names: Collection[str],
         value_values: Iterable[Iterable[Any]],
+        # Given these are the same type as the normal values, force keyword-only so that
+        # they can't be confused.
+        *,
+        insertion_value_names: Collection[str] = [],
+        insertion_value_values: Collection[Iterable[Any]] = [],
     ) -> None:
         """
         Upsert, many times, using batching where possible.
@@ -1598,10 +1650,14 @@ def simple_upsert_many_txn_native_upsert(
             value_names: The value column names
             value_values: A list of each row's value column values.
                 Ignored if value_names is empty.
+            insertion_value_names: The value column names to use only when inserting
+            insertion_value_values: A list of each row's value column values to use only
+                when inserting. Ignored if `insertion_value_names` is empty.
         """
         allnames: List[str] = []
         allnames.extend(key_names)
         allnames.extend(value_names)
+        allnames.extend(insertion_value_names)
 
         if not value_names:
             latter = "NOTHING"
@@ -1612,8 +1668,8 @@ def simple_upsert_many_txn_native_upsert(
 
         args = []
 
-        for x, y in zip(key_values, value_values):
-            args.append(tuple(x) + tuple(y))
+        for x, y, z in zip(key_values, value_values, insertion_value_values):
+            args.append(tuple(x) + tuple(y) + tuple(z))
 
         if isinstance(txn.database_engine, PostgresEngine):
             # We use `execute_values` as it can be a lot faster than `execute_batch`,
diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index 5ed99516cfc..06e0b53ca80 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -2500,7 +2500,8 @@ def _fill_table_txn(txn: LoggingTransaction) -> None:
                 table="sliding_sync_membership_snapshots",
                 key_names=key_names,
                 key_values=key_values,
-                # TODO: Implement these
+                value_names=(),
+                value_values=[],
                 insertion_value_names=insertion_value_names,
                 insertion_value_values=insertion_value_values,
             )
diff --git a/tests/storage/test__base.py b/tests/storage/test__base.py
index 49dc973a369..9f87747652e 100644
--- a/tests/storage/test__base.py
+++ b/tests/storage/test__base.py
@@ -31,6 +31,11 @@
 
 
 class UpdateUpsertManyTests(unittest.HomeserverTestCase):
+    """
+    Integration tests for the "simple" SQL generating methods in SQLBaseStore that
+    actually run against a database.
+    """
+
     def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
         self.storage = hs.get_datastores().main
 
@@ -130,6 +135,65 @@ def test_upsert_many(self) -> None:
             {(1, "user1", "hello"), (2, "user2", "bleb")},
         )
 
+    def test_upsert_many_with_insertion_values(self) -> None:
+        """
+        Upsert_many will only insert the
+        `insertion_value_names`/`insertion_value_values` (not on update/conflict)
+        """
+        # Add some data to an empty table
+        key_names = ["id", "username"]
+        key_values = [[1, "user1"], [2, "user2"]]
+        value_names: List[str] = []
+        value_values: List[List[str]] = []
+        insertion_value_names = ["value"]
+        insertion_value_values = [["hello"], ["there"]]
+
+        self.get_success(
+            self.storage.db_pool.runInteraction(
+                "test",
+                self.storage.db_pool.simple_upsert_many_txn,
+                self.table_name,
+                key_names,
+                key_values,
+                value_names,
+                value_values,
+                insertion_value_names=insertion_value_names,
+                insertion_value_values=insertion_value_values,
+            )
+        )
+
+        # Check results are what we expect
+        self.assertEqual(
+            set(self._dump_table_to_tuple()),
+            {(1, "user1", "hello"), (2, "user2", "there")},
+        )
+
+        # Update only user2
+        key_values = [[2, "user2"]]
+        # Since this row already exists, when we try to insert it again, it should not
+        # insert the value again.
+        insertion_value_values = [["again"]]
+
+        self.get_success(
+            self.storage.db_pool.runInteraction(
+                "test",
+                self.storage.db_pool.simple_upsert_many_txn,
+                self.table_name,
+                key_names,
+                key_values,
+                value_names,
+                value_values,
+                insertion_value_names=insertion_value_names,
+                insertion_value_values=insertion_value_values,
+            )
+        )
+
+        # Check results are what we expect
+        self.assertEqual(
+            set(self._dump_table_to_tuple()),
+            {(1, "user1", "hello"), (2, "user2", "there")},
+        )
+
     def test_simple_update_many(self) -> None:
         """
         simple_update_many performs many updates at once.
diff --git a/tests/storage/test_base.py b/tests/storage/test_base.py
index 9420d03841c..33ba2b2456b 100644
--- a/tests/storage/test_base.py
+++ b/tests/storage/test_base.py
@@ -35,7 +35,12 @@
 
 
 class SQLBaseStoreTestCase(unittest.TestCase):
-    """Test the "simple" SQL generating methods in SQLBaseStore."""
+    """
+    Test the "simple" SQL generating methods in SQLBaseStore.
+
+    Tests that the SQL is generated correctly and that the correct arguments are passed
+    (does not actually run the queries or test the end-result in the database).
+    """
 
     def setUp(self) -> None:
         # This is the Twisted connection pool.
@@ -620,6 +625,74 @@ def test_upsert_many_no_values(
                 [("oldvalue",)],
             )
 
+    @defer.inlineCallbacks
+    def test_upsert_many_no_values_with_insertion_values(
+        self,
+    ) -> Generator["defer.Deferred[object]", object, None]:
+        yield defer.ensureDeferred(
+            self.datastore.db_pool.simple_upsert_many(
+                table="tablename",
+                key_names=["keycol1"],
+                key_values=[["keyval1.1"], ["keyval1.2"]],
+                value_names=[],
+                value_values=[],
+                insertion_value_names=["insertioncol1"],
+                insertion_value_values=[["insertionvalue1"], ["insertionvalue2"]],
+                desc="",
+            )
+        )
+
+        if USE_POSTGRES_FOR_TESTS:
+            self.mock_execute_values.assert_called_once_with(
+                self.mock_txn,
+                "INSERT INTO tablename (keycol1, insertioncol1) VALUES ? ON CONFLICT (keycol1) DO NOTHING",
+                [("keyval1.1", "insertionvalue1"), ("keyval1.2", "insertionvalue2")],
+                template=None,
+                fetch=False,
+            )
+        else:
+            self.mock_txn.executemany.assert_called_once_with(
+                "INSERT INTO tablename (keycol1, insertioncol1) VALUES (?, ?) ON CONFLICT (keycol1) DO NOTHING",
+                [("keyval1.1", "insertionvalue1"), ("keyval1.2", "insertionvalue2")],
+            )
+
+    @defer.inlineCallbacks
+    def test_upsert_many_values_with_insertion_values(
+        self,
+    ) -> Generator["defer.Deferred[object]", object, None]:
+        yield defer.ensureDeferred(
+            self.datastore.db_pool.simple_upsert_many(
+                table="tablename",
+                key_names=["keycol1"],
+                key_values=[["keyval1.1"], ["keyval1.2"]],
+                value_names=["valuecol1"],
+                value_values=[["value1"], ["value2"]],
+                insertion_value_names=["insertioncol1"],
+                insertion_value_values=[["insertionvalue1"], ["insertionvalue2"]],
+                desc="",
+            )
+        )
+
+        if USE_POSTGRES_FOR_TESTS:
+            self.mock_execute_values.assert_called_once_with(
+                self.mock_txn,
+                "INSERT INTO tablename (keycol1, valuecol1, insertioncol1) VALUES ? ON CONFLICT (keycol1) DO UPDATE SET valuecol1=EXCLUDED.valuecol1",
+                [
+                    ("keyval1.1", "value1", "insertionvalue1"),
+                    ("keyval1.2", "value2", "insertionvalue2"),
+                ],
+                template=None,
+                fetch=False,
+            )
+        else:
+            self.mock_txn.executemany.assert_called_once_with(
+                "INSERT INTO tablename (keycol1, valuecol1, insertioncol1) VALUES (?, ?, ?) ON CONFLICT (keycol1) DO UPDATE SET valuecol1=EXCLUDED.valuecol1",
+                [
+                    ("keyval1.1", "value1", "insertionvalue1"),
+                    ("keyval1.2", "value2", "insertionvalue2"),
+                ],
+            )
+
     @defer.inlineCallbacks
     def test_upsert_emulated_no_values_exists(
         self,

From 836bcaefe0210f6868cd5d808b44e1270583ac0b Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Mon, 9 Sep 2024 19:13:28 -0500
Subject: [PATCH 08/23] Fix nothing being upserted

---
 synapse/storage/databases/main/events_bg_updates.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index 06e0b53ca80..f17e0fc3ad5 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -2655,6 +2655,10 @@ def _resolve_stale_data_in_sliding_sync_joined_rooms_table(
                 # No value columns, therefore make a blank list so that the following
                 # zip() works correctly.
                 value_values=[() for x in range(len(chunk))],
+                insertion_value_names=(),
+                # No value columns, therefore make a blank list so that the following
+                # zip() works correctly.
+                insertion_value_values=[() for x in range(len(chunk))],
             )
     else:
         # Avoid adding the background updates when there is no data to run them on (if

From 90371d48af29479a71e2393c1552ee3c098d1e30 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Mon, 9 Sep 2024 19:31:47 -0500
Subject: [PATCH 09/23] Sort by `event_stream_ordering` so we can re-use
 snapshots along the way

---
 synapse/storage/databases/main/events_bg_updates.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index f17e0fc3ad5..22662a389f7 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -1955,7 +1955,11 @@ def _find_memberships_to_update_txn(
             Tuple[str, Optional[str], str, str, str, str, int, Optional[str], bool]
         ]:
             # Fetch the set of event IDs that we want to update
-
+            #
+            # It's important to sort by `event_stream_ordering` *ascending* (oldest to
+            # newest) so that we can re-use the snapshots along the way if nothing has
+            # changed.
+            #
             if initial_phase:
                 # There are some old out-of-band memberships (before
                 # https://github.com/matrix-org/synapse/issues/6983) where we don't have
@@ -1978,7 +1982,7 @@ def _find_memberships_to_update_txn(
                     INNER JOIN events AS e USING (event_id)
                     LEFT JOIN rooms AS r ON (c.room_id = r.room_id)
                     WHERE (c.room_id, c.user_id) > (?, ?)
-                    ORDER BY c.room_id ASC, c.user_id ASC
+                    ORDER BY c.event_stream_ordering ASC, c.room_id ASC, c.user_id ASC
                     LIMIT ?
                     """,
                     (last_room_id, last_user_id, batch_size),

From 2c2146d16748c6ff7b3f418d17431c8204e05c56 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Mon, 9 Sep 2024 20:00:42 -0500
Subject: [PATCH 10/23] Only re-use for join memberships

---
 .../databases/main/events_bg_updates.py       | 61 +++++++++++++------
 1 file changed, 41 insertions(+), 20 deletions(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index 22662a389f7..b05d863bec6 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -2123,10 +2123,10 @@ def _find_previous_invite_or_knock_membership_txn(
 
         # Map from room_id to ...
         # Just some convenience maps for easier lookup by `room_id`.
-        last_to_insert_membership_snapshots_by_room_id: Dict[
+        last_to_insert_join_membership_snapshots_by_room_id: Dict[
             str, SlidingSyncMembershipSnapshotSharedInsertValues
         ] = {}
-        last_to_insert_membership_infos_by_room_id: Dict[
+        last_to_insert_join_membership_infos_by_room_id: Dict[
             str, SlidingSyncMembershipInfoWithEventPos
         ] = {}
         # TODO: `concurrently_execute` based on buckets of room_ids
@@ -2171,34 +2171,50 @@ def _find_previous_invite_or_knock_membership_txn(
                     },
                 )
 
-            # Check if the current state has been updated since the last snapshot we inserted.
+            # Check if we can potentially re-use the last snapshot we inserted for this
+            # room.
+            #
             can_use_last_snapshot = False
             # TODO: We could also look in the database to see if a snapshot is available
-            last_to_insert_membership_snapshot = (
-                last_to_insert_membership_snapshots_by_room_id.get(room_id)
+            last_to_insert_join_membership_snapshot = (
+                last_to_insert_join_membership_snapshots_by_room_id.get(room_id)
             )
-            last_to_insert_membership_info = (
-                last_to_insert_membership_infos_by_room_id.get(room_id)
+            last_to_insert_join_membership_info = (
+                last_to_insert_join_membership_infos_by_room_id.get(room_id)
             )
-            if last_to_insert_membership_snapshot is not None:
+            # We can only re-use the last snapshot if their previous membership was in
+            # the room. Since that's not easy to determine without more extra lookups,
+            # we only apply this to join memberships for now.
+            #
+            # - JOIN: Re-use if no state changes
+            # - LEAVE: Could re-use if they were previously joined otherwise this could
+            #   be a rescinded/rejected invite/knock and we don't want to leak anything.
+            # - INVITE/KNOCK: ❌ We can only rely on the stripped state
+            # - BAN: Could re-use if they were previously joined otherwise they might
+            #   have just been banned after invite/knock and we don't want to leak
+            #   anything.
+            if (
+                last_to_insert_join_membership_snapshot is not None
+                and membership == Membership.JOIN
+            ):
                 can_use_last_snapshot = True
 
                 # These should go hand-in-hand
-                assert last_to_insert_membership_info is not None
+                assert last_to_insert_join_membership_info is not None
 
                 # Sanity check that our tokens are in order. We should be iterating in
                 # ascending order.
                 assert (
-                    last_to_insert_membership_info.membership_event_stream_ordering
+                    last_to_insert_join_membership_info.membership_event_stream_ordering
                     <= membership_event_stream_ordering
                 )
 
-                # TODO: This only works if someone is from the server is participating in the room.
+                # Check if the current state has been updated since the last snapshot we inserted.
                 state_deltas_since_last_snapshot = await self.get_current_state_deltas_for_room(
                     room_id,
                     # From the last snapshot we inserted
                     from_token=RoomStreamToken(
-                        stream=last_to_insert_membership_info.membership_event_stream_ordering,
+                        stream=last_to_insert_join_membership_info.membership_event_stream_ordering,
                     ),
                     # To our current membership position
                     to_token=RoomStreamToken(
@@ -2223,10 +2239,10 @@ def _find_previous_invite_or_knock_membership_txn(
             # just re-use the last snapshot we inserted.
             if can_use_last_snapshot:
                 # Based on the logic above, we should have a last snapshot to work from
-                assert last_to_insert_membership_snapshot is not None
+                assert last_to_insert_join_membership_snapshot is not None
 
                 sliding_sync_membership_snapshots_insert_map = (
-                    last_to_insert_membership_snapshot
+                    last_to_insert_join_membership_snapshot
                 )
             elif membership == Membership.JOIN:
                 # If we're still joined, we can pull from current state.
@@ -2403,14 +2419,10 @@ def _find_previous_invite_or_knock_membership_txn(
                     f"Unexpected membership {membership} ({membership_event_id}) that we don't know how to handle yet"
                 )
 
-            to_insert_membership_snapshots[
-                (room_id, user_id)
-            ] = last_to_insert_membership_snapshots_by_room_id[room_id] = (
+            to_insert_membership_snapshots[(room_id, user_id)] = (
                 sliding_sync_membership_snapshots_insert_map
             )
-            to_insert_membership_infos[
-                (room_id, user_id)
-            ] = last_to_insert_membership_infos_by_room_id[room_id] = (
+            to_insert_membership_infos[(room_id, user_id)] = (
                 SlidingSyncMembershipInfoWithEventPos(
                     user_id=user_id,
                     sender=sender,
@@ -2422,6 +2434,15 @@ def _find_previous_invite_or_knock_membership_txn(
                     or "master",
                 )
             )
+            # Keep track of the last join membership snapshot as we may be able to
+            # re-use all the work we did.
+            if membership == Membership.JOIN:
+                last_to_insert_join_membership_snapshots_by_room_id[room_id] = (
+                    to_insert_membership_snapshots[(room_id, user_id)]
+                )
+                last_to_insert_join_membership_infos_by_room_id[room_id] = (
+                    to_insert_membership_infos[(room_id, user_id)]
+                )
 
         # Assemble data so it's ready for the batch queries in the transaction
         key_names = ("room_id", "user_id")

From 107816373112901f70d9b1f5c5c4da877ce72529 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Mon, 9 Sep 2024 21:31:44 -0500
Subject: [PATCH 11/23] Make sure re-use doesn't blow up

---
 .../databases/main/events_bg_updates.py       |   3 +-
 tests/storage/test_sliding_sync_tables.py     | 216 ++++++++++++++++++
 2 files changed, 218 insertions(+), 1 deletion(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index b05d863bec6..e3c40f97448 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -2175,7 +2175,8 @@ def _find_previous_invite_or_knock_membership_txn(
             # room.
             #
             can_use_last_snapshot = False
-            # TODO: We could also look in the database to see if a snapshot is available
+            # TODO: We could also look in the database to see if a snapshot is
+            # available. Is it worth the lookup time?
             last_to_insert_join_membership_snapshot = (
                 last_to_insert_join_membership_snapshots_by_room_id.get(room_id)
             )
diff --git a/tests/storage/test_sliding_sync_tables.py b/tests/storage/test_sliding_sync_tables.py
index de80ad53cd8..a34334b518b 100644
--- a/tests/storage/test_sliding_sync_tables.py
+++ b/tests/storage/test_sliding_sync_tables.py
@@ -4546,6 +4546,222 @@ def test_membership_snapshots_background_update_forgotten_partial(self) -> None:
             user2_snapshot,
         )
 
+    def test_membership_snapshots_background_update_multiple_memberships(self) -> None:
+        """
+        Test that the background update for `sliding_sync_membership_snapshots` works
+        correctly when we populate multiple memberships in a room. Make sure the
+        snapshot re-use isn't flawed.
+        """
+        user1_id = self.register_user("user1", "pass")
+        _user1_tok = self.login(user1_id, "pass")
+        user2_id = self.register_user("user2", "pass")
+        user2_tok = self.login(user2_id, "pass")
+        user3_id = self.register_user("user3", "pass")
+        user3_tok = self.login(user3_id, "pass")
+        user4_id = self.register_user("user4", "pass")
+        user4_tok = self.login(user4_id, "pass")
+        user5_id = self.register_user("user5", "pass")
+        _user5_tok = self.login(user5_id, "pass")
+        user6_id = self.register_user("user6", "pass")
+        user6_tok = self.login(user6_id, "pass")
+
+        # Create a room
+        room_id = self.helper.create_room_as(user2_id, tok=user2_tok)
+        # Add a room name
+        self.helper.send_state(
+            room_id,
+            EventTypes.Name,
+            {"name": "my super room"},
+            tok=user2_tok,
+        )
+
+        # User1 is invited to the room
+        self.helper.invite(room_id, src=user2_id, targ=user1_id, tok=user2_tok)
+        # User3 joins the room
+        self.helper.join(room_id, user3_id, tok=user3_tok)
+        # User4 joins the room
+        self.helper.join(room_id, user4_id, tok=user4_tok)
+        # User5 is invited to the room
+        self.helper.invite(room_id, src=user2_id, targ=user5_id, tok=user2_tok)
+        # User6 joins the room
+        self.helper.join(room_id, user6_id, tok=user6_tok)
+
+        # Update the room name after all of the memberships have been added
+        self.helper.send_state(
+            room_id,
+            EventTypes.Name,
+            {"name": "my super duper room"},
+            tok=user2_tok,
+        )
+
+        # Clean-up the `sliding_sync_membership_snapshots` table as if the inserts did not
+        # happen during event creation.
+        self.get_success(
+            self.store.db_pool.simple_delete_many(
+                table="sliding_sync_membership_snapshots",
+                column="room_id",
+                iterable=(room_id,),
+                keyvalues={},
+                desc="sliding_sync_membership_snapshots.test_membership_snapshots_background_update_forgotten_missing",
+            )
+        )
+
+        # We shouldn't find anything in the table because we just deleted them in
+        # preparation for the test.
+        sliding_sync_membership_snapshots_results = (
+            self._get_sliding_sync_membership_snapshots()
+        )
+        self.assertIncludes(
+            set(sliding_sync_membership_snapshots_results.keys()),
+            set(),
+            exact=True,
+        )
+
+        # Insert and run the background update.
+        self.get_success(
+            self.store.db_pool.simple_insert(
+                "background_updates",
+                {
+                    "update_name": _BackgroundUpdates.SLIDING_SYNC_MEMBERSHIP_SNAPSHOTS_BG_UPDATE,
+                    "progress_json": "{}",
+                },
+            )
+        )
+        self.store.db_pool.updates._all_done = False
+        self.wait_for_background_updates()
+
+        # Make sure the table is populated
+        sliding_sync_membership_snapshots_results = (
+            self._get_sliding_sync_membership_snapshots()
+        )
+        self.assertIncludes(
+            set(sliding_sync_membership_snapshots_results.keys()),
+            {
+                (room_id, user1_id),
+                (room_id, user2_id),
+                (room_id, user3_id),
+                (room_id, user4_id),
+                (room_id, user5_id),
+                (room_id, user6_id),
+            },
+            exact=True,
+        )
+        state_map = self.get_success(
+            self.storage_controllers.state.get_current_state(room_id)
+        )
+        self.assertEqual(
+            sliding_sync_membership_snapshots_results.get((room_id, user1_id)),
+            _SlidingSyncMembershipSnapshotResult(
+                room_id=room_id,
+                user_id=user1_id,
+                sender=user2_id,
+                membership_event_id=state_map[(EventTypes.Member, user1_id)].event_id,
+                membership=Membership.INVITE,
+                event_stream_ordering=state_map[
+                    (EventTypes.Member, user1_id)
+                ].internal_metadata.stream_ordering,
+                has_known_state=True,
+                room_type=None,
+                # Has the value included in the stripped state at the time of the invite
+                room_name="my super room",
+                is_encrypted=False,
+                tombstone_successor_room_id=None,
+            ),
+        )
+        self.assertEqual(
+            sliding_sync_membership_snapshots_results.get((room_id, user2_id)),
+            _SlidingSyncMembershipSnapshotResult(
+                room_id=room_id,
+                user_id=user2_id,
+                sender=user2_id,
+                membership_event_id=state_map[(EventTypes.Member, user2_id)].event_id,
+                membership=Membership.JOIN,
+                event_stream_ordering=state_map[
+                    (EventTypes.Member, user2_id)
+                ].internal_metadata.stream_ordering,
+                has_known_state=True,
+                room_type=None,
+                # Even though the room name after the room creator joined, the
+                # background update uses current state for the room for joined users.
+                room_name="my super duper room",
+                is_encrypted=False,
+                tombstone_successor_room_id=None,
+            ),
+        )
+        self.assertEqual(
+            sliding_sync_membership_snapshots_results.get((room_id, user3_id)),
+            _SlidingSyncMembershipSnapshotResult(
+                room_id=room_id,
+                user_id=user3_id,
+                sender=user3_id,
+                membership_event_id=state_map[(EventTypes.Member, user3_id)].event_id,
+                membership=Membership.JOIN,
+                event_stream_ordering=state_map[
+                    (EventTypes.Member, user3_id)
+                ].internal_metadata.stream_ordering,
+                has_known_state=True,
+                room_type=None,
+                room_name="my super duper room",
+                is_encrypted=False,
+                tombstone_successor_room_id=None,
+            ),
+        )
+        self.assertEqual(
+            sliding_sync_membership_snapshots_results.get((room_id, user4_id)),
+            _SlidingSyncMembershipSnapshotResult(
+                room_id=room_id,
+                user_id=user4_id,
+                sender=user4_id,
+                membership_event_id=state_map[(EventTypes.Member, user4_id)].event_id,
+                membership=Membership.JOIN,
+                event_stream_ordering=state_map[
+                    (EventTypes.Member, user4_id)
+                ].internal_metadata.stream_ordering,
+                has_known_state=True,
+                room_type=None,
+                room_name="my super duper room",
+                is_encrypted=False,
+                tombstone_successor_room_id=None,
+            ),
+        )
+        self.assertEqual(
+            sliding_sync_membership_snapshots_results.get((room_id, user5_id)),
+            _SlidingSyncMembershipSnapshotResult(
+                room_id=room_id,
+                user_id=user5_id,
+                sender=user2_id,
+                membership_event_id=state_map[(EventTypes.Member, user5_id)].event_id,
+                membership=Membership.INVITE,
+                event_stream_ordering=state_map[
+                    (EventTypes.Member, user5_id)
+                ].internal_metadata.stream_ordering,
+                has_known_state=True,
+                room_type=None,
+                # Has the value included in the stripped state at the time of the invite
+                room_name="my super room",
+                is_encrypted=False,
+                tombstone_successor_room_id=None,
+            ),
+        )
+        self.assertEqual(
+            sliding_sync_membership_snapshots_results.get((room_id, user6_id)),
+            _SlidingSyncMembershipSnapshotResult(
+                room_id=room_id,
+                user_id=user6_id,
+                sender=user6_id,
+                membership_event_id=state_map[(EventTypes.Member, user6_id)].event_id,
+                membership=Membership.JOIN,
+                event_stream_ordering=state_map[
+                    (EventTypes.Member, user6_id)
+                ].internal_metadata.stream_ordering,
+                has_known_state=True,
+                room_type=None,
+                room_name="my super duper room",
+                is_encrypted=False,
+                tombstone_successor_room_id=None,
+            ),
+        )
+
 
 class SlidingSyncTablesCatchUpBackgroundUpdatesTestCase(SlidingSyncTablesTestCaseBase):
     """

From a4cc5eb02c29b7e2aa6d832570f706c7cc091b9e Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Mon, 9 Sep 2024 22:53:52 -0500
Subject: [PATCH 12/23] Get snapshots info concurrently

---
 .../databases/main/events_bg_updates.py       | 576 ++++++++++--------
 1 file changed, 308 insertions(+), 268 deletions(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index e3c40f97448..ca0a627cc60 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -2129,147 +2129,272 @@ def _find_previous_invite_or_knock_membership_txn(
         last_to_insert_join_membership_infos_by_room_id: Dict[
             str, SlidingSyncMembershipInfoWithEventPos
         ] = {}
-        # TODO: `concurrently_execute` based on buckets of room_ids
-        for (
-            room_id,
-            room_id_from_rooms_table,
-            user_id,
-            sender,
-            membership_event_id,
-            membership,
-            membership_event_stream_ordering,
-            membership_event_instance_name,
-            is_outlier,
-        ) in memberships_to_update_rows:
-            # We don't know how to handle `membership` values other than these. The
-            # code below would need to be updated.
-            assert membership in (
-                Membership.JOIN,
-                Membership.INVITE,
-                Membership.KNOCK,
-                Membership.LEAVE,
-                Membership.BAN,
-            )
-
-            # There are some old out-of-band memberships (before
-            # https://github.com/matrix-org/synapse/issues/6983) where we don't have the
-            # corresponding room stored in the `rooms` table`. We have a `FOREIGN KEY`
-            # constraint on the `sliding_sync_membership_snapshots` table so we have to
-            # fix-up these memberships by adding the room to the `rooms` table.
-            if room_id_from_rooms_table is None:
-                await self.db_pool.simple_insert(
-                    table="rooms",
-                    values={
-                        "room_id": room_id,
-                        # Only out-of-band memberships are missing from the `rooms`
-                        # table so that is the only type of membership we're dealing
-                        # with here. Since we don't calculate the "chain cover" for
-                        # out-of-band memberships, we can just set this to `True` as if
-                        # the user ever joins the room, we will end up calculating the
-                        # "chain cover" anyway.
-                        "has_auth_chain_index": True,
-                    },
-                )
 
-            # Check if we can potentially re-use the last snapshot we inserted for this
-            # room.
-            #
-            can_use_last_snapshot = False
-            # TODO: We could also look in the database to see if a snapshot is
-            # available. Is it worth the lookup time?
-            last_to_insert_join_membership_snapshot = (
-                last_to_insert_join_membership_snapshots_by_room_id.get(room_id)
-            )
-            last_to_insert_join_membership_info = (
-                last_to_insert_join_membership_infos_by_room_id.get(room_id)
-            )
-            # We can only re-use the last snapshot if their previous membership was in
-            # the room. Since that's not easy to determine without more extra lookups,
-            # we only apply this to join memberships for now.
-            #
-            # - JOIN: Re-use if no state changes
-            # - LEAVE: Could re-use if they were previously joined otherwise this could
-            #   be a rescinded/rejected invite/knock and we don't want to leak anything.
-            # - INVITE/KNOCK: ❌ We can only rely on the stripped state
-            # - BAN: Could re-use if they were previously joined otherwise they might
-            #   have just been banned after invite/knock and we don't want to leak
-            #   anything.
-            if (
-                last_to_insert_join_membership_snapshot is not None
-                and membership == Membership.JOIN
-            ):
-                can_use_last_snapshot = True
+        async def _handle_memberships_to_update_rows_for_room(
+            memberships_to_update_rows: List[
+                Tuple[str, Optional[str], str, str, str, str, int, Optional[str], bool]
+            ],
+        ) -> None:
+            if not memberships_to_update_rows:
+                return
 
-                # These should go hand-in-hand
-                assert last_to_insert_join_membership_info is not None
+            room_id_for_all_rows: str = memberships_to_update_rows[0][0]
 
-                # Sanity check that our tokens are in order. We should be iterating in
-                # ascending order.
-                assert (
-                    last_to_insert_join_membership_info.membership_event_stream_ordering
-                    <= membership_event_stream_ordering
+            for (
+                room_id,
+                room_id_from_rooms_table,
+                user_id,
+                sender,
+                membership_event_id,
+                membership,
+                membership_event_stream_ordering,
+                membership_event_instance_name,
+                is_outlier,
+            ) in memberships_to_update_rows:
+                # Sanity check that we're working on the same room so the snapshot
+                # re-use logic can work properly (we need to process memberships in a
+                # room sequentially chronologically)
+                assert room_id == room_id_for_all_rows
+
+                # We don't know how to handle `membership` values other than these. The
+                # code below would need to be updated.
+                assert membership in (
+                    Membership.JOIN,
+                    Membership.INVITE,
+                    Membership.KNOCK,
+                    Membership.LEAVE,
+                    Membership.BAN,
                 )
 
-                # Check if the current state has been updated since the last snapshot we inserted.
-                state_deltas_since_last_snapshot = await self.get_current_state_deltas_for_room(
-                    room_id,
-                    # From the last snapshot we inserted
-                    from_token=RoomStreamToken(
-                        stream=last_to_insert_join_membership_info.membership_event_stream_ordering,
-                    ),
-                    # To our current membership position
-                    to_token=RoomStreamToken(
-                        stream=membership_event_stream_ordering,
-                    ),
-                )
-                for state_delta in state_deltas_since_last_snapshot:
-                    # We only need to check for the state is relevant to the
-                    # `sliding_sync_joined_rooms` table.
-                    if (
-                        state_delta.event_type,
-                        state_delta.state_key,
-                    ) in SLIDING_SYNC_RELEVANT_STATE_SET:
-                        can_use_last_snapshot = False
-                        # Once we find one relevant state event that changed, no need to
-                        # look any further
-                        break
-
-            # Map of values to insert/update in the `sliding_sync_membership_snapshots` table
-            sliding_sync_membership_snapshots_insert_map: SlidingSyncMembershipSnapshotSharedInsertValues = {}
-            # If none of the relevant state has changed since the last snapshot, we can
-            # just re-use the last snapshot we inserted.
-            if can_use_last_snapshot:
-                # Based on the logic above, we should have a last snapshot to work from
-                assert last_to_insert_join_membership_snapshot is not None
-
-                sliding_sync_membership_snapshots_insert_map = (
-                    last_to_insert_join_membership_snapshot
+                # There are some old out-of-band memberships (before
+                # https://github.com/matrix-org/synapse/issues/6983) where we don't have the
+                # corresponding room stored in the `rooms` table`. We have a `FOREIGN KEY`
+                # constraint on the `sliding_sync_membership_snapshots` table so we have to
+                # fix-up these memberships by adding the room to the `rooms` table.
+                if room_id_from_rooms_table is None:
+                    await self.db_pool.simple_insert(
+                        table="rooms",
+                        values={
+                            "room_id": room_id,
+                            # Only out-of-band memberships are missing from the `rooms`
+                            # table so that is the only type of membership we're dealing
+                            # with here. Since we don't calculate the "chain cover" for
+                            # out-of-band memberships, we can just set this to `True` as if
+                            # the user ever joins the room, we will end up calculating the
+                            # "chain cover" anyway.
+                            "has_auth_chain_index": True,
+                        },
+                    )
+
+                # Check if we can potentially re-use the last snapshot we inserted for this
+                # room.
+                #
+                can_use_last_snapshot = False
+                # TODO: We could also look in the database to see if a snapshot is
+                # available. Is it worth the lookup time?
+                last_to_insert_join_membership_snapshot = (
+                    last_to_insert_join_membership_snapshots_by_room_id.get(room_id)
                 )
-            elif membership == Membership.JOIN:
-                # If we're still joined, we can pull from current state.
-                current_state_ids_map: StateMap[
-                    str
-                ] = await self.hs.get_storage_controllers().state.get_current_state_ids(
-                    room_id,
-                    state_filter=StateFilter.from_types(
-                        SLIDING_SYNC_RELEVANT_STATE_SET
-                    ),
-                    # Partially-stated rooms should have all state events except for
-                    # remote membership events so we don't need to wait at all because
-                    # we only want some non-membership state
-                    await_full_state=False,
+                last_to_insert_join_membership_info = (
+                    last_to_insert_join_membership_infos_by_room_id.get(room_id)
                 )
+                # We can only re-use the last snapshot if their previous membership was in
+                # the room. Since that's not easy to determine without more extra lookups,
+                # we only apply this to join memberships for now.
+                #
+                # - JOIN: Re-use if no state changes
+                # - LEAVE: Could re-use if they were previously joined otherwise this could
+                #   be a rescinded/rejected invite/knock and we don't want to leak anything.
+                # - INVITE/KNOCK: ❌ We can only rely on the stripped state
+                # - BAN: Could re-use if they were previously joined otherwise they might
+                #   have just been banned after invite/knock and we don't want to leak
+                #   anything.
+                if (
+                    last_to_insert_join_membership_snapshot is not None
+                    and membership == Membership.JOIN
+                ):
+                    can_use_last_snapshot = True
 
-                # TODO: Read the state from the `room_stats_state` table if we can
+                    # These should go hand-in-hand
+                    assert last_to_insert_join_membership_info is not None
 
-                # We're iterating over rooms that we are joined to so they should
-                # have `current_state_events` and we should have some current state
-                # for each room
-                if current_state_ids_map:
-                    try:
-                        fetched_events = await self.get_events(
-                            current_state_ids_map.values()
+                    # Sanity check that our tokens are in order. We should be iterating in
+                    # ascending order.
+                    assert (
+                        last_to_insert_join_membership_info.membership_event_stream_ordering
+                        <= membership_event_stream_ordering
+                    )
+
+                    # Check if the current state has been updated since the last snapshot we inserted.
+                    state_deltas_since_last_snapshot = await self.get_current_state_deltas_for_room(
+                        room_id,
+                        # From the last snapshot we inserted
+                        from_token=RoomStreamToken(
+                            stream=last_to_insert_join_membership_info.membership_event_stream_ordering,
+                        ),
+                        # To our current membership position
+                        to_token=RoomStreamToken(
+                            stream=membership_event_stream_ordering,
+                        ),
+                    )
+                    for state_delta in state_deltas_since_last_snapshot:
+                        # We only need to check for the state is relevant to the
+                        # `sliding_sync_joined_rooms` table.
+                        if (
+                            state_delta.event_type,
+                            state_delta.state_key,
+                        ) in SLIDING_SYNC_RELEVANT_STATE_SET:
+                            can_use_last_snapshot = False
+                            # Once we find one relevant state event that changed, no need to
+                            # look any further
+                            break
+
+                # Map of values to insert/update in the `sliding_sync_membership_snapshots` table
+                sliding_sync_membership_snapshots_insert_map: SlidingSyncMembershipSnapshotSharedInsertValues = {}
+                # If none of the relevant state has changed since the last snapshot, we can
+                # just re-use the last snapshot we inserted.
+                if can_use_last_snapshot:
+                    # Based on the logic above, we should have a last snapshot to work from
+                    assert last_to_insert_join_membership_snapshot is not None
+
+                    sliding_sync_membership_snapshots_insert_map = (
+                        last_to_insert_join_membership_snapshot
+                    )
+                elif membership == Membership.JOIN:
+                    # If we're still joined, we can pull from current state.
+                    current_state_ids_map: StateMap[
+                        str
+                    ] = await self.hs.get_storage_controllers().state.get_current_state_ids(
+                        room_id,
+                        state_filter=StateFilter.from_types(
+                            SLIDING_SYNC_RELEVANT_STATE_SET
+                        ),
+                        # Partially-stated rooms should have all state events except for
+                        # remote membership events so we don't need to wait at all because
+                        # we only want some non-membership state
+                        await_full_state=False,
+                    )
+
+                    # TODO: Read the state from the `room_stats_state` table if we can
+
+                    # We're iterating over rooms that we are joined to so they should
+                    # have `current_state_events` and we should have some current state
+                    # for each room
+                    if current_state_ids_map:
+                        try:
+                            fetched_events = await self.get_events(
+                                current_state_ids_map.values()
+                            )
+                        except (DatabaseCorruptionError, InvalidEventError) as e:
+                            logger.warning(
+                                "Failed to fetch state for room '%s' due to corrupted events. Ignoring. Error: %s",
+                                room_id,
+                                e,
+                            )
+                            continue
+
+                        current_state_map: StateMap[EventBase] = {
+                            state_key: fetched_events[event_id]
+                            for state_key, event_id in current_state_ids_map.items()
+                            # `get_events(...)` will filter out events for unknown room versions
+                            if event_id in fetched_events
+                        }
+
+                        # Can happen for unknown room versions (old room versions that aren't known
+                        # anymore) since `get_events(...)` will filter out events for unknown room
+                        # versions
+                        if not current_state_map:
+                            continue
+
+                        state_insert_values = PersistEventsStore._get_sliding_sync_insert_values_from_state_map(
+                            current_state_map
+                        )
+                        sliding_sync_membership_snapshots_insert_map.update(
+                            state_insert_values
+                        )
+                        # We should have some insert values for each room, even if they are `None`
+                        assert sliding_sync_membership_snapshots_insert_map
+
+                        # We have current state to work from
+                        sliding_sync_membership_snapshots_insert_map[
+                            "has_known_state"
+                        ] = True
+                    else:
+                        # Although we expect every room to have a create event (even
+                        # past unknown room versions since we haven't supported one
+                        # without it), there seem to be some corrupted rooms in
+                        # practice that don't have the create event in the
+                        # `current_state_events` table. The create event does exist
+                        # in the events table though. We'll just say that we don't
+                        # know the state for these rooms and continue on with our
+                        # day.
+                        sliding_sync_membership_snapshots_insert_map[
+                            "has_known_state"
+                        ] = False
+                elif membership in (Membership.INVITE, Membership.KNOCK) or (
+                    membership in (Membership.LEAVE, Membership.BAN) and is_outlier
+                ):
+                    invite_or_knock_event_id = membership_event_id
+                    invite_or_knock_membership = membership
+
+                    # If the event is an `out_of_band_membership` (special case of
+                    # `outlier`), we never had historical state so we have to pull from
+                    # the stripped state on the previous invite/knock event. This gives
+                    # us a consistent view of the room state regardless of your
+                    # membership (i.e. the room shouldn't disappear if your using the
+                    # `is_encrypted` filter and you leave).
+                    if membership in (Membership.LEAVE, Membership.BAN) and is_outlier:
+                        (
+                            invite_or_knock_event_id,
+                            invite_or_knock_membership,
+                        ) = await self.db_pool.runInteraction(
+                            "sliding_sync_membership_snapshots_bg_update._find_previous_invite_or_knock_membership_txn",
+                            _find_previous_invite_or_knock_membership_txn,
+                            room_id,
+                            user_id,
+                            membership_event_id,
+                        )
+
+                    # Pull from the stripped state on the invite/knock event
+                    invite_or_knock_event = await self.get_event(
+                        invite_or_knock_event_id
+                    )
+
+                    raw_stripped_state_events = None
+                    if invite_or_knock_membership == Membership.INVITE:
+                        invite_room_state = invite_or_knock_event.unsigned.get(
+                            "invite_room_state"
+                        )
+                        raw_stripped_state_events = invite_room_state
+                    elif invite_or_knock_membership == Membership.KNOCK:
+                        knock_room_state = invite_or_knock_event.unsigned.get(
+                            "knock_room_state"
                         )
+                        raw_stripped_state_events = knock_room_state
+
+                    sliding_sync_membership_snapshots_insert_map = PersistEventsStore._get_sliding_sync_insert_values_from_stripped_state(
+                        raw_stripped_state_events
+                    )
+
+                    # We should have some insert values for each room, even if no
+                    # stripped state is on the event because we still want to record
+                    # that we have no known state
+                    assert sliding_sync_membership_snapshots_insert_map
+                elif membership in (Membership.LEAVE, Membership.BAN):
+                    # Pull from historical state
+                    state_ids_map = await self.hs.get_storage_controllers().state.get_state_ids_for_event(
+                        membership_event_id,
+                        state_filter=StateFilter.from_types(
+                            SLIDING_SYNC_RELEVANT_STATE_SET
+                        ),
+                        # Partially-stated rooms should have all state events except for
+                        # remote membership events so we don't need to wait at all because
+                        # we only want some non-membership state
+                        await_full_state=False,
+                    )
+
+                    try:
+                        fetched_events = await self.get_events(state_ids_map.values())
                     except (DatabaseCorruptionError, InvalidEventError) as e:
                         logger.warning(
                             "Failed to fetch state for room '%s' due to corrupted events. Ignoring. Error: %s",
@@ -2278,9 +2403,9 @@ def _find_previous_invite_or_knock_membership_txn(
                         )
                         continue
 
-                    current_state_map: StateMap[EventBase] = {
+                    state_map: StateMap[EventBase] = {
                         state_key: fetched_events[event_id]
-                        for state_key, event_id in current_state_ids_map.items()
+                        for state_key, event_id in state_ids_map.items()
                         # `get_events(...)` will filter out events for unknown room versions
                         if event_id in fetched_events
                     }
@@ -2288,11 +2413,11 @@ def _find_previous_invite_or_knock_membership_txn(
                     # Can happen for unknown room versions (old room versions that aren't known
                     # anymore) since `get_events(...)` will filter out events for unknown room
                     # versions
-                    if not current_state_map:
+                    if not state_map:
                         continue
 
                     state_insert_values = PersistEventsStore._get_sliding_sync_insert_values_from_state_map(
-                        current_state_map
+                        state_map
                     )
                     sliding_sync_membership_snapshots_insert_map.update(
                         state_insert_values
@@ -2300,152 +2425,67 @@ def _find_previous_invite_or_knock_membership_txn(
                     # We should have some insert values for each room, even if they are `None`
                     assert sliding_sync_membership_snapshots_insert_map
 
-                    # We have current state to work from
+                    # We have historical state to work from
                     sliding_sync_membership_snapshots_insert_map["has_known_state"] = (
                         True
                     )
                 else:
-                    # Although we expect every room to have a create event (even
-                    # past unknown room versions since we haven't supported one
-                    # without it), there seem to be some corrupted rooms in
-                    # practice that don't have the create event in the
-                    # `current_state_events` table. The create event does exist
-                    # in the events table though. We'll just say that we don't
-                    # know the state for these rooms and continue on with our
-                    # day.
-                    sliding_sync_membership_snapshots_insert_map["has_known_state"] = (
-                        False
-                    )
-            elif membership in (Membership.INVITE, Membership.KNOCK) or (
-                membership in (Membership.LEAVE, Membership.BAN) and is_outlier
-            ):
-                invite_or_knock_event_id = membership_event_id
-                invite_or_knock_membership = membership
-
-                # If the event is an `out_of_band_membership` (special case of
-                # `outlier`), we never had historical state so we have to pull from
-                # the stripped state on the previous invite/knock event. This gives
-                # us a consistent view of the room state regardless of your
-                # membership (i.e. the room shouldn't disappear if your using the
-                # `is_encrypted` filter and you leave).
-                if membership in (Membership.LEAVE, Membership.BAN) and is_outlier:
-                    (
-                        invite_or_knock_event_id,
-                        invite_or_knock_membership,
-                    ) = await self.db_pool.runInteraction(
-                        "sliding_sync_membership_snapshots_bg_update._find_previous_invite_or_knock_membership_txn",
-                        _find_previous_invite_or_knock_membership_txn,
-                        room_id,
-                        user_id,
-                        membership_event_id,
+                    # We don't know how to handle this type of membership yet
+                    #
+                    # FIXME: We should use `assert_never` here but for some reason
+                    # the exhaustive matching doesn't recognize the `Never` here.
+                    # assert_never(membership)
+                    raise AssertionError(
+                        f"Unexpected membership {membership} ({membership_event_id}) that we don't know how to handle yet"
                     )
 
-                # Pull from the stripped state on the invite/knock event
-                invite_or_knock_event = await self.get_event(invite_or_knock_event_id)
-
-                raw_stripped_state_events = None
-                if invite_or_knock_membership == Membership.INVITE:
-                    invite_room_state = invite_or_knock_event.unsigned.get(
-                        "invite_room_state"
-                    )
-                    raw_stripped_state_events = invite_room_state
-                elif invite_or_knock_membership == Membership.KNOCK:
-                    knock_room_state = invite_or_knock_event.unsigned.get(
-                        "knock_room_state"
-                    )
-                    raw_stripped_state_events = knock_room_state
-
-                sliding_sync_membership_snapshots_insert_map = PersistEventsStore._get_sliding_sync_insert_values_from_stripped_state(
-                    raw_stripped_state_events
+                to_insert_membership_snapshots[(room_id, user_id)] = (
+                    sliding_sync_membership_snapshots_insert_map
                 )
-
-                # We should have some insert values for each room, even if no
-                # stripped state is on the event because we still want to record
-                # that we have no known state
-                assert sliding_sync_membership_snapshots_insert_map
-            elif membership in (Membership.LEAVE, Membership.BAN):
-                # Pull from historical state
-                state_ids_map = await self.hs.get_storage_controllers().state.get_state_ids_for_event(
-                    membership_event_id,
-                    state_filter=StateFilter.from_types(
-                        SLIDING_SYNC_RELEVANT_STATE_SET
-                    ),
-                    # Partially-stated rooms should have all state events except for
-                    # remote membership events so we don't need to wait at all because
-                    # we only want some non-membership state
-                    await_full_state=False,
+                to_insert_membership_infos[(room_id, user_id)] = (
+                    SlidingSyncMembershipInfoWithEventPos(
+                        user_id=user_id,
+                        sender=sender,
+                        membership_event_id=membership_event_id,
+                        membership=membership,
+                        membership_event_stream_ordering=membership_event_stream_ordering,
+                        # If instance_name is null we default to "master"
+                        membership_event_instance_name=membership_event_instance_name
+                        or "master",
+                    )
                 )
-
-                try:
-                    fetched_events = await self.get_events(state_ids_map.values())
-                except (DatabaseCorruptionError, InvalidEventError) as e:
-                    logger.warning(
-                        "Failed to fetch state for room '%s' due to corrupted events. Ignoring. Error: %s",
-                        room_id,
-                        e,
+                # Keep track of the last join membership snapshot as we may be able to
+                # re-use all the work we did.
+                if membership == Membership.JOIN:
+                    last_to_insert_join_membership_snapshots_by_room_id[room_id] = (
+                        to_insert_membership_snapshots[(room_id, user_id)]
                     )
-                    continue
-
-                state_map: StateMap[EventBase] = {
-                    state_key: fetched_events[event_id]
-                    for state_key, event_id in state_ids_map.items()
-                    # `get_events(...)` will filter out events for unknown room versions
-                    if event_id in fetched_events
-                }
-
-                # Can happen for unknown room versions (old room versions that aren't known
-                # anymore) since `get_events(...)` will filter out events for unknown room
-                # versions
-                if not state_map:
-                    continue
-
-                state_insert_values = (
-                    PersistEventsStore._get_sliding_sync_insert_values_from_state_map(
-                        state_map
+                    last_to_insert_join_membership_infos_by_room_id[room_id] = (
+                        to_insert_membership_infos[(room_id, user_id)]
                     )
-                )
-                sliding_sync_membership_snapshots_insert_map.update(state_insert_values)
-                # We should have some insert values for each room, even if they are `None`
-                assert sliding_sync_membership_snapshots_insert_map
 
-                # We have historical state to work from
-                sliding_sync_membership_snapshots_insert_map["has_known_state"] = True
-            else:
-                # We don't know how to handle this type of membership yet
-                #
-                # FIXME: We should use `assert_never` here but for some reason
-                # the exhaustive matching doesn't recognize the `Never` here.
-                # assert_never(membership)
-                raise AssertionError(
-                    f"Unexpected membership {membership} ({membership_event_id}) that we don't know how to handle yet"
-                )
+        # Bucket the memberships by room_id so we can process them concurrently.
+        room_id_to_memberships_to_update_rows: Dict[
+            str,
+            List[
+                Tuple[str, Optional[str], str, str, str, str, int, Optional[str], bool]
+            ],
+        ] = {}
+        for row in memberships_to_update_rows:
+            room_id = row[0]
+            room_id_to_memberships_to_update_rows.setdefault(room_id, []).append(row)
 
-            to_insert_membership_snapshots[(room_id, user_id)] = (
-                sliding_sync_membership_snapshots_insert_map
-            )
-            to_insert_membership_infos[(room_id, user_id)] = (
-                SlidingSyncMembershipInfoWithEventPos(
-                    user_id=user_id,
-                    sender=sender,
-                    membership_event_id=membership_event_id,
-                    membership=membership,
-                    membership_event_stream_ordering=membership_event_stream_ordering,
-                    # If instance_name is null we default to "master"
-                    membership_event_instance_name=membership_event_instance_name
-                    or "master",
-                )
-            )
-            # Keep track of the last join membership snapshot as we may be able to
-            # re-use all the work we did.
-            if membership == Membership.JOIN:
-                last_to_insert_join_membership_snapshots_by_room_id[room_id] = (
-                    to_insert_membership_snapshots[(room_id, user_id)]
-                )
-                last_to_insert_join_membership_infos_by_room_id[room_id] = (
-                    to_insert_membership_infos[(room_id, user_id)]
-                )
+        # We can batch process each room concurrently.
+        #
+        # In order for the snapshot re-use logic to work correctly, we need to process
+        # memberships in a room sequentially chronologically.
+        await concurrently_execute(
+            _handle_memberships_to_update_rows_for_room,
+            room_id_to_memberships_to_update_rows.values(),
+            10,
+        )
 
-        # Assemble data so it's ready for the batch queries in the transaction
+        # Re-assemble data so it's ready for the batch queries in the transaction
         key_names = ("room_id", "user_id")
         key_values: List[Tuple[str, str]] = []
         insertion_value_names = (

From 815ef2fb9958068121c07a03a8c0f480b575e793 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Mon, 9 Sep 2024 23:10:25 -0500
Subject: [PATCH 13/23] Read from the `room_stats_state` table

---
 .../databases/main/events_bg_updates.py       | 160 ++++++++++++++----
 1 file changed, 123 insertions(+), 37 deletions(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index ca0a627cc60..6a05014756b 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -1765,7 +1765,7 @@ async def handle_room(room_id: str) -> None:
             # `room_stats_state` table. In the case, where we can't use
             # `room_stats_state`, we just fetch all of the events manually.
             if current_state_ids_to_fetch:
-                fetched_events = await self.get_events(current_state_ids_map.values())
+                fetched_events = await self.get_events(current_state_ids_to_fetch)
 
                 # Even if we are joined to the room, `get_events(...)` will filter out
                 # events for unknown room versions that we no longer support. We don't need
@@ -2112,6 +2112,19 @@ def _find_previous_invite_or_knock_membership_txn(
 
             return event_id, membership
 
+        # As long as we get this value before we fetch the current state, we can use it
+        # to check if something has changed since that point.
+        #
+        # Same as `await self.get_stats_positions()` but get around the method
+        # resolution issues that crop up when we try to mixin `StatsStore` into this
+        # store.
+        stats_stream_position = await self.db_pool.simple_select_one_onecol(
+            table="stats_incremental_position",
+            keyvalues={},
+            retcol="stream_id",
+            desc="stats_incremental_position",
+        )
+
         # Map from (room_id, user_id) to ...
         to_insert_membership_snapshots: Dict[
             Tuple[str, str], SlidingSyncMembershipSnapshotSharedInsertValues
@@ -2275,50 +2288,123 @@ async def _handle_memberships_to_update_rows_for_room(
                         await_full_state=False,
                     )
 
-                    # TODO: Read the state from the `room_stats_state` table if we can
-
                     # We're iterating over rooms that we are joined to so they should
                     # have `current_state_events` and we should have some current state
                     # for each room
                     if current_state_ids_map:
-                        try:
-                            fetched_events = await self.get_events(
-                                current_state_ids_map.values()
-                            )
-                        except (DatabaseCorruptionError, InvalidEventError) as e:
-                            logger.warning(
-                                "Failed to fetch state for room '%s' due to corrupted events. Ignoring. Error: %s",
+                        current_state_ids_to_fetch: List[str] = []
+
+                        # Check if the current state has been updated since the stats were
+                        # generated.
+                        state_deltas_since_stats = (
+                            await self.get_current_state_deltas_for_room(
                                 room_id,
-                                e,
+                                from_token=RoomStreamToken(
+                                    stream=stats_stream_position
+                                ),
+                                to_token=None,
                             )
-                            continue
-
-                        current_state_map: StateMap[EventBase] = {
-                            state_key: fetched_events[event_id]
-                            for state_key, event_id in current_state_ids_map.items()
-                            # `get_events(...)` will filter out events for unknown room versions
-                            if event_id in fetched_events
-                        }
-
-                        # Can happen for unknown room versions (old room versions that aren't known
-                        # anymore) since `get_events(...)` will filter out events for unknown room
-                        # versions
-                        if not current_state_map:
-                            continue
-
-                        state_insert_values = PersistEventsStore._get_sliding_sync_insert_values_from_state_map(
-                            current_state_map
                         )
-                        sliding_sync_membership_snapshots_insert_map.update(
-                            state_insert_values
-                        )
-                        # We should have some insert values for each room, even if they are `None`
-                        assert sliding_sync_membership_snapshots_insert_map
+                        can_use_room_stats_state_table = True
+                        for state_delta in state_deltas_since_stats:
+                            # We only need to check for the state is relevant to the
+                            # `sliding_sync_joined_rooms` table.
+                            if (
+                                state_delta.event_type,
+                                state_delta.state_key,
+                            ) in SLIDING_SYNC_RELEVANT_STATE_SET:
+                                can_use_room_stats_state_table = False
+                                # Once we find one relevant state event that changed, no need to
+                                # look any further
+                                break
+
+                        # Avoid fetching full events out as much as possible. We can look-up all of
+                        # the relevant state values except for the `tombstone_successor_room_id`
+                        # from the `room_stats_state` table.
+                        if can_use_room_stats_state_table:
+                            (name, room_type, encryption) = cast(
+                                Tuple[Optional[str], Optional[str], Optional[str]],
+                                await self.db_pool.simple_select_one(
+                                    table="room_stats_state",
+                                    keyvalues={"room_id": room_id},
+                                    retcols=["name", "room_type", "encryption"],
+                                ),
+                            )
 
-                        # We have current state to work from
-                        sliding_sync_membership_snapshots_insert_map[
-                            "has_known_state"
-                        ] = True
+                            # We have current state to work from
+                            sliding_sync_membership_snapshots_insert_map[
+                                "has_known_state"
+                            ] = True
+                            # And can copy some of the values from the `room_stats_state` table
+                            sliding_sync_membership_snapshots_insert_map[
+                                "room_name"
+                            ] = name
+                            sliding_sync_membership_snapshots_insert_map[
+                                "room_type"
+                            ] = room_type
+                            sliding_sync_membership_snapshots_insert_map[
+                                "is_encrypted"
+                            ] = encryption is not None
+
+                            # If the room has a tombstone, we still need to fetch it out manually
+                            # because this information isn't available in the `room_stats_state`
+                            # table.
+                            tombstone_event_id = current_state_ids_map.get(
+                                (EventTypes.Tombstone, "")
+                            )
+                            if tombstone_event_id is not None:
+                                current_state_ids_to_fetch.append(tombstone_event_id)
+
+                        else:
+                            current_state_ids_to_fetch.extend(
+                                current_state_ids_map.values()
+                            )
+
+                        # Fetch any more events that we need to complete the `state_insert_values`.
+                        # For the happy-path, this will just be the tombstone event if the room has
+                        # one since we got the rest of the values already from the
+                        # `room_stats_state` table. In the case, where we can't use
+                        # `room_stats_state`, we just fetch all of the events manually.
+                        if current_state_ids_to_fetch:
+                            try:
+                                fetched_events = await self.get_events(
+                                    current_state_ids_to_fetch
+                                )
+                            except (DatabaseCorruptionError, InvalidEventError) as e:
+                                logger.warning(
+                                    "Failed to fetch state for room '%s' due to corrupted events. Ignoring. Error: %s",
+                                    room_id,
+                                    e,
+                                )
+                                continue
+
+                            current_state_map: StateMap[EventBase] = {
+                                state_key: fetched_events[event_id]
+                                for state_key, event_id in current_state_ids_map.items()
+                                # `get_events(...)` will filter out events for unknown room versions
+                                if event_id in fetched_events
+                            }
+
+                            # Can happen for unknown room versions (old room versions that aren't known
+                            # anymore) since `get_events(...)` will filter out events for unknown room
+                            # versions
+                            if not current_state_map:
+                                continue
+
+                            more_state_insert_values = PersistEventsStore._get_sliding_sync_insert_values_from_state_map(
+                                current_state_map
+                            )
+                            # Update our main insert map
+                            sliding_sync_membership_snapshots_insert_map.update(
+                                more_state_insert_values
+                            )
+                            # We should have some insert values for each room, even if they are `None`
+                            assert sliding_sync_membership_snapshots_insert_map
+
+                            # We have current state to work from
+                            sliding_sync_membership_snapshots_insert_map[
+                                "has_known_state"
+                            ] = True
                     else:
                         # Although we expect every room to have a create event (even
                         # past unknown room versions since we haven't supported one

From 1688212268517678cc90c61add68d65dd178e905 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Mon, 9 Sep 2024 23:14:38 -0500
Subject: [PATCH 14/23] Uncomment and remove todo

---
 synapse/storage/database.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/synapse/storage/database.py b/synapse/storage/database.py
index 17aacf69806..b1656e93da4 100644
--- a/synapse/storage/database.py
+++ b/synapse/storage/database.py
@@ -1576,9 +1576,8 @@ def simple_upsert_many_txn(
                 key_values,
                 value_names,
                 value_values,
-                # TODO
-                # insertion_value_names=insertion_value_names,
-                # insertion_value_values=insertion_value_values,
+                insertion_value_names=insertion_value_names,
+                insertion_value_values=insertion_value_values,
             )
 
     def simple_upsert_many_txn_emulated(

From f045574426f3e93ae132ce2238a48d3deedd5cb9 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Mon, 9 Sep 2024 23:25:08 -0500
Subject: [PATCH 15/23] Better comments

---
 synapse/storage/databases/main/events_bg_updates.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index 6a05014756b..cd64d702adb 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -2134,7 +2134,7 @@ def _find_previous_invite_or_knock_membership_txn(
             Tuple[str, str], SlidingSyncMembershipInfoWithEventPos
         ] = {}
 
-        # Map from room_id to ...
+        # Map from room_id to the last join membership snapshot/info we inserted.
         # Just some convenience maps for easier lookup by `room_id`.
         last_to_insert_join_membership_snapshots_by_room_id: Dict[
             str, SlidingSyncMembershipSnapshotSharedInsertValues
@@ -2148,8 +2148,8 @@ async def _handle_memberships_to_update_rows_for_room(
                 Tuple[str, Optional[str], str, str, str, str, int, Optional[str], bool]
             ],
         ) -> None:
-            if not memberships_to_update_rows:
-                return
+            # We should have exited earlier if there were no rows to process
+            assert len(memberships_to_update_rows) > 0
 
             room_id_for_all_rows: str = memberships_to_update_rows[0][0]
 
@@ -2164,9 +2164,9 @@ async def _handle_memberships_to_update_rows_for_room(
                 membership_event_instance_name,
                 is_outlier,
             ) in memberships_to_update_rows:
-                # Sanity check that we're working on the same room so the snapshot
-                # re-use logic can work properly (we need to process memberships in a
-                # room sequentially chronologically)
+                # Sanity check that we're working on the same room for all of the rows
+                # so the snapshot re-use logic can work properly (we need to process
+                # memberships in a room sequentially chronologically)
                 assert room_id == room_id_for_all_rows
 
                 # We don't know how to handle `membership` values other than these. The

From a2da5376a39c5ae4861907d0fdc9cd54ebd698b3 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Mon, 9 Sep 2024 23:38:36 -0500
Subject: [PATCH 16/23] Add changelog

---
 changelog.d/17676.misc | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 changelog.d/17676.misc

diff --git a/changelog.d/17676.misc b/changelog.d/17676.misc
new file mode 100644
index 00000000000..756918e2b21
--- /dev/null
+++ b/changelog.d/17676.misc
@@ -0,0 +1 @@
+Pre-populate room data used in experimental [MSC3575](https://github.com/matrix-org/matrix-spec-proposals/pull/3575) Sliding Sync `/sync` endpoint for quick filtering/sorting.

From ca4afad068b2d3b2498e69a65b481b31496646e4 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Mon, 9 Sep 2024 23:42:22 -0500
Subject: [PATCH 17/23] Be more precise

---
 synapse/storage/databases/main/events_bg_updates.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index cd64d702adb..5b3c1b97d11 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -2166,7 +2166,7 @@ async def _handle_memberships_to_update_rows_for_room(
             ) in memberships_to_update_rows:
                 # Sanity check that we're working on the same room for all of the rows
                 # so the snapshot re-use logic can work properly (we need to process
-                # memberships in a room sequentially chronologically)
+                # memberships in a room sequentially by `stream_ordering` ascending)
                 assert room_id == room_id_for_all_rows
 
                 # We don't know how to handle `membership` values other than these. The
@@ -2564,7 +2564,7 @@ async def _handle_memberships_to_update_rows_for_room(
         # We can batch process each room concurrently.
         #
         # In order for the snapshot re-use logic to work correctly, we need to process
-        # memberships in a room sequentially chronologically.
+        # memberships in a room sequentially by `stream_ordering` ascending.
         await concurrently_execute(
             _handle_memberships_to_update_rows_for_room,
             room_id_to_memberships_to_update_rows.values(),

From bb5784b26882252722a02caa434f8b071b33858b Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Tue, 10 Sep 2024 13:08:06 -0500
Subject: [PATCH 18/23] Need to sort by room_id to make progress

---
 synapse/storage/databases/main/events_bg_updates.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index 825fdbc083d..64238e2a0fc 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -1968,8 +1968,8 @@ def _find_memberships_to_update_txn(
             # Fetch the set of event IDs that we want to update
             #
             # It's important to sort by `event_stream_ordering` *ascending* (oldest to
-            # newest) so that we can re-use the snapshots along the way if nothing has
-            # changed.
+            # newest) (at-least per-room) so that we can re-use the snapshots along the
+            # way if nothing has changed.
             #
             if initial_phase:
                 # There are some old out-of-band memberships (before
@@ -1993,7 +1993,7 @@ def _find_memberships_to_update_txn(
                     INNER JOIN events AS e USING (event_id)
                     LEFT JOIN rooms AS r ON (c.room_id = r.room_id)
                     WHERE (c.room_id, c.user_id) > (?, ?)
-                    ORDER BY c.event_stream_ordering ASC, c.room_id ASC, c.user_id ASC
+                    ORDER BY c.room_id ASC, c.event_stream_ordering ASC, c.user_id ASC
                     LIMIT ?
                     """,
                     (last_room_id, last_user_id, batch_size),

From 6f3e8be1dc78233e448cc4cc98f02765a5a8b015 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Tue, 10 Sep 2024 13:41:49 -0500
Subject: [PATCH 19/23] Add assertion message to better debug

---
 synapse/storage/databases/main/events_bg_updates.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index 64238e2a0fc..0e6b2ede274 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -2247,6 +2247,10 @@ async def _handle_memberships_to_update_rows_for_room(
                     assert (
                         last_to_insert_join_membership_info.membership_event_stream_ordering
                         <= membership_event_stream_ordering
+                    ), (
+                        f"Expected to be iterating in ascending order as we process membership snapshots for room {room_id}"
+                        + " but our last join membership was at {last_to_insert_join_membership_info.membership_event_stream_ordering}"
+                        + "and our current membership is at {membership_event_stream_ordering}"
                     )
 
                     # Check if the current state has been updated since the last snapshot we inserted.

From 8477df44b69a50871a0ab1a85114e1056568ca10 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Tue, 10 Sep 2024 13:46:06 -0500
Subject: [PATCH 20/23] Fix f-string

---
 synapse/storage/databases/main/events_bg_updates.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index 0e6b2ede274..ca73bea361a 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -2249,8 +2249,8 @@ async def _handle_memberships_to_update_rows_for_room(
                         <= membership_event_stream_ordering
                     ), (
                         f"Expected to be iterating in ascending order as we process membership snapshots for room {room_id}"
-                        + " but our last join membership was at {last_to_insert_join_membership_info.membership_event_stream_ordering}"
-                        + "and our current membership is at {membership_event_stream_ordering}"
+                        + f" but our last join membership was at {last_to_insert_join_membership_info.membership_event_stream_ordering}"
+                        + f" and our current membership is at {membership_event_stream_ordering}"
                     )
 
                     # Check if the current state has been updated since the last snapshot we inserted.

From 76cc6629deb85ba93ef5c42732da48634f451820 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Tue, 10 Sep 2024 13:55:56 -0500
Subject: [PATCH 21/23] Use `events.stream_ordering` since
 `local_current_membership.event_stream_ordering` is nullable

And seems to actually have `null` values in it
---
 synapse/storage/databases/main/events_bg_updates.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index ca73bea361a..212241962c0 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -1993,7 +1993,7 @@ def _find_memberships_to_update_txn(
                     INNER JOIN events AS e USING (event_id)
                     LEFT JOIN rooms AS r ON (c.room_id = r.room_id)
                     WHERE (c.room_id, c.user_id) > (?, ?)
-                    ORDER BY c.room_id ASC, c.event_stream_ordering ASC, c.user_id ASC
+                    ORDER BY c.room_id ASC, c.stream_ordering ASC, c.user_id ASC
                     LIMIT ?
                     """,
                     (last_room_id, last_user_id, batch_size),

From 2e8f5d66da1be88a054d37b39589803a1bc92b1e Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Tue, 10 Sep 2024 14:00:34 -0500
Subject: [PATCH 22/23] Use correct table

---
 synapse/storage/databases/main/events_bg_updates.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index 212241962c0..b098d018208 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -1993,7 +1993,7 @@ def _find_memberships_to_update_txn(
                     INNER JOIN events AS e USING (event_id)
                     LEFT JOIN rooms AS r ON (c.room_id = r.room_id)
                     WHERE (c.room_id, c.user_id) > (?, ?)
-                    ORDER BY c.room_id ASC, c.stream_ordering ASC, c.user_id ASC
+                    ORDER BY c.room_id ASC, e.stream_ordering ASC, c.user_id ASC
                     LIMIT ?
                     """,
                     (last_room_id, last_user_id, batch_size),

From 6565380a8b2139b3caa0623cb0a135c76187cc81 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Wed, 11 Sep 2024 10:58:07 -0500
Subject: [PATCH 23/23] Rely on `events.stream_ordering`

---
 .../databases/main/events_bg_updates.py        | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index b098d018208..28b01e6b9fc 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -1967,10 +1967,18 @@ def _find_memberships_to_update_txn(
         ]:
             # Fetch the set of event IDs that we want to update
             #
-            # It's important to sort by `event_stream_ordering` *ascending* (oldest to
+            # It's important to sort by `stream_ordering` *ascending* (oldest to
             # newest) (at-least per-room) so that we can re-use the snapshots along the
             # way if nothing has changed.
             #
+            # Also be careful because `local_current_membership.event_stream_ordering`
+            # is nullable and not always filled in. We need to join on `events` to rely
+            # on `events.stream_ordering` instead. Even though the
+            # `events.stream_ordering` also doesn't have a `NOT NULL` constraint, it
+            # doesn't have any rows where this is the case (checked on `matrix.org`).
+            # The fact the `events.stream_ordering` is a nullable column is a holdover
+            # from a rename of the column.
+            #
             if initial_phase:
                 # There are some old out-of-band memberships (before
                 # https://github.com/matrix-org/synapse/issues/6983) where we don't have
@@ -1999,7 +2007,7 @@ def _find_memberships_to_update_txn(
                     (last_room_id, last_user_id, batch_size),
                 )
             elif last_event_stream_ordering is not None:
-                # It's important to sort by `event_stream_ordering` *ascending* (oldest to
+                # It's important to sort by `stream_ordering` *ascending* (oldest to
                 # newest) so that if we see that this background update in progress and want
                 # to start the catch-up process, we can safely assume that it will
                 # eventually get to the rooms we want to catch-up on anyway (see
@@ -2017,13 +2025,13 @@ def _find_memberships_to_update_txn(
                         e.sender,
                         c.event_id,
                         c.membership,
-                        c.event_stream_ordering,
+                        e.stream_ordering,
                         e.instance_name,
                         e.outlier
                     FROM local_current_membership AS c
                     INNER JOIN events AS e USING (event_id)
-                    WHERE event_stream_ordering > ?
-                    ORDER BY event_stream_ordering ASC
+                    WHERE e.stream_ordering > ?
+                    ORDER BY e.stream_ordering ASC
                     LIMIT ?
                     """,
                     (last_event_stream_ordering, batch_size),