Track lazy loaded members in SSS separately.

This ensures that the set of required state doesn't keep growing as we
add and remove member state. We then only load them from the DB when
needed, rather than all state for all rooms when we get a request.

Author: erikjohnston

synapse/handlers/sliding_sync/__init__.py

@@ -63,6 +63,7 @@
     HaveSentRoomFlag,
     MutablePerConnectionState,
     PerConnectionState,
+    RoomLazyMembershipChanges,
     RoomSyncConfig,
     SlidingSyncConfig,
     SlidingSyncResult,
@@ -984,14 +985,15 @@ async def get_room_sync_data(
         #
         # Calculate the `StateFilter` based on the `required_state` for the room
         required_state_filter = StateFilter.none()
-        # The requested `required_state_map` with the lazy membership expanded and
-        # `$ME` replaced with the user's ID. This allows us to see what membership we've
-        # sent down to the client in the next request.
-        #
-        # Make a copy so we can modify it. Still need to be careful to make a copy of
-        # the state key sets if we want to add/remove from them. We could make a deep
-        # copy but this saves us some work.
-        expanded_required_state_map = dict(room_sync_config.required_state_map)
+
+        # Keep track of which users' state we may need to fetch. We split this
+        # into explicit users and lazy loaded users.
+        explicit_user_state = set()
+        lazy_load_user_ids = set()
+
+        # Whether lazy-loading of room members is enabled.
+        lazy_load_room_members = False
+
         if room_membership_for_user_at_to_token.membership not in (
             Membership.INVITE,
             Membership.KNOCK,
@@ -1039,7 +1041,6 @@ async def get_room_sync_data(
             else:
                 required_state_types: list[tuple[str, str | None]] = []
                 num_wild_state_keys = 0
-                lazy_load_room_members = False
                 num_others = 0
                 for (
                     state_type,
@@ -1071,6 +1072,10 @@ async def get_room_sync_data(
                                             timeline_event.state_key
                                         )
 
+                            # The client needs to know the membership of everyone in
+                            # the timeline we're returning.
+                            lazy_load_user_ids.update(timeline_membership)
+
                             # Update the required state filter so we pick up the new
                             # membership
                             if limited or initial:
@@ -1095,36 +1100,30 @@ async def get_room_sync_data(
                                 # to state resolution anyway).
                                 required_state_types.append((EventTypes.Member, None))
 
-                            # Add an explicit entry for each user in the timeline
-                            #
-                            # Make a new set or copy of the state key set so we can
-                            # modify it without affecting the original
-                            # `required_state_map`
-                            expanded_required_state_map[EventTypes.Member] = (
-                                expanded_required_state_map.get(
-                                    EventTypes.Member, set()
+                                # Record the extra members we're returning.
+                                lazy_load_user_ids.update(
+                                    state_key
+                                    for event_type, state_key in room_state_delta_id_map
+                                    if event_type == EventTypes.Member
                                 )
-                                | timeline_membership
-                            )
-                        elif state_key == StateValues.ME:
+                        else:
                             num_others += 1
-                            required_state_types.append((state_type, user.to_string()))
+
                             # Replace `$ME` with the user's ID so we can deduplicate
                             # when someone requests the same state with `$ME` or with
                             # their user ID.
-                            #
-                            # Make a new set or copy of the state key set so we can
-                            # modify it without affecting the original
-                            # `required_state_map`
-                            expanded_required_state_map[EventTypes.Member] = (
-                                expanded_required_state_map.get(
-                                    EventTypes.Member, set()
-                                )
-                                | {user.to_string()}
+                            normalized_state_key = state_key
+                            if state_key == StateValues.ME:
+                                normalized_state_key = user.to_string()
+
+                            if state_type == EventTypes.Member:
+                                # Also track explicitly requested member state for
+                                # lazy membership tracking.
+                                explicit_user_state.add(normalized_state_key)
+
+                            required_state_types.append(
+                                (state_type, normalized_state_key)
                             )
-                        else:
-                            num_others += 1
-                            required_state_types.append((state_type, state_key))
 
                 set_tag(
                     SynapseTags.FUNC_ARG_PREFIX
@@ -1142,6 +1141,10 @@ async def get_room_sync_data(
 
                 required_state_filter = StateFilter.from_types(required_state_types)
 
+        # Define `required_user_state` as all user state we want.
+        required_user_state = explicit_user_state | lazy_load_user_ids
+        lazy_load_user_ids -= explicit_user_state
+
         # We need this base set of info for the response so let's just fetch it along
         # with the `required_state` for the room
         hero_room_state = [
@@ -1174,6 +1177,17 @@ async def get_room_sync_data(
         # state as well).
         hero_membership_state: StateMap[EventBase] = {}
 
+        # By default we mark all required user state as being added when lazy
+        # loaded members is enabled.
+        #
+        # We may later update this to account for previously sent members.
+        returned_users = {}
+        if lazy_load_room_members:
+            returned_users = dict.fromkeys(lazy_load_user_ids)
+        new_connection_state.room_lazy_membership[room_id] = RoomLazyMembershipChanges(
+            returned=returned_users
+        )
+
         if initial:
             room_state = await self.get_current_state_at(
                 room_id=room_id,
@@ -1186,19 +1200,38 @@ async def get_room_sync_data(
             # We'll later filter this down so we don't need to do so here.
             hero_membership_state = room_state
         else:
+            assert from_token is not None
             assert from_bound is not None
 
             if prev_room_sync_config is not None:
+                # Fetch which of the needed lazy-loaded members we have already sent.
+                if required_user_state:
+                    previously_returned_user_state = (
+                        await self.store.get_sliding_sync_connection_lazy_members(
+                            connection_position=from_token.connection_position,
+                            room_id=room_id,
+                            user_ids=required_user_state,
+                        )
+                    )
+                else:
+                    previously_returned_user_state = {}
+
                 # Check if there are any changes to the required state config
                 # that we need to handle.
                 changes_return = _required_state_changes(
                     user.to_string(),
                     prev_required_state_map=prev_room_sync_config.required_state_map,
-                    request_required_state_map=expanded_required_state_map,
+                    request_required_state_map=room_sync_config.required_state_map,
+                    previously_returned_user_state=previously_returned_user_state,
+                    required_user_state=lazy_load_user_ids,
                     state_deltas=room_state_delta_id_map,
                 )
                 changed_required_state_map = changes_return.required_state_map_change
 
+                new_connection_state.room_lazy_membership[
+                    room_id
+                ].invalidated = changes_return.lazy_members_invalidated
+
                 if changes_return.added_state_filter:
                     # Some state entries got added, so we pull out the current
                     # state for them. If we don't do this we'd only send down new deltas.
@@ -1309,7 +1342,7 @@ async def get_room_sync_data(
             bump_stamp = 0
 
         room_sync_required_state_map_to_persist: Mapping[str, AbstractSet[str]] = (
-            expanded_required_state_map
+            room_sync_config.required_state_map
         )
         if changed_required_state_map:
             room_sync_required_state_map_to_persist = changed_required_state_map
@@ -1508,17 +1541,27 @@ class _RequiredStateChangesReturn:
             the room config, or None if there is no change.
         added_state_filter: The state filter to use to fetch any additional
             current state that needs to be returned to the client.
+        lazy_members_previously_returned: The set of user IDs we should add to
+            the lazy members cache that we had previously returned.
+        lazy_members_invalidated: The set of user IDs whose membership has
+            changed but we didn't send down, so we need to invalidate them from
+            the cache.
     """
 
     required_state_map_change: Mapping[str, AbstractSet[str]] | None
     added_state_filter: StateFilter
 
+    lazy_members_previously_returned: AbstractSet[str] = frozenset()
+    lazy_members_invalidated: AbstractSet[str] = frozenset()
+
 
 def _required_state_changes(
     user_id: str,
     *,
     prev_required_state_map: Mapping[str, AbstractSet[str]],
     request_required_state_map: Mapping[str, AbstractSet[str]],
+    previously_returned_user_state: Mapping[str, int | None],
+    required_user_state: AbstractSet[str],
     state_deltas: StateMap[str],
 ) -> _RequiredStateChangesReturn:
     """Calculates the changes between the required state room config from the
@@ -1533,13 +1576,60 @@ def _required_state_changes(
     This function tries to ensure to handle the case where a state entry is
     added, removed and then added again to the required state. In that case we
     only want to re-send that entry down sync if it has changed.
+
+    Args:
+        user_id: The user ID of the user making the request.
+        prev_required_state_map: The required state map from the previous
+            request.
+        request_required_state_map: The required state map from the current
+            request.
+        previously_returned_user_state: The set of user IDs whose lazy-loaded
+            membership we have previously returned to the client.
+        required_user_state: The set of user IDs whose lazy-loaded membership
+            is required for this request.
+        state_deltas: The state deltas that have changed in the room since the
+            previous request.
     """
+
+    # First we find any lazy members that have been invalidated due to state
+    # changes that we are not sending down.
+    lazy_members_invalidated = set()
+    for event_type, state_key in state_deltas:
+        if event_type != EventTypes.Member:
+            continue
+
+        if state_key in required_user_state:
+            # We're returning this member change.
+            continue
+
+        if state_key not in previously_returned_user_state:
+            # We've not previously returned this member so nothing to
+            # invalidate.
+            continue
+
+        lazy_members_invalidated.add(state_key)
+
     if prev_required_state_map == request_required_state_map:
-        # There has been no change. Return immediately.
+        # There has been no change in state, just need to check lazy members.
+        newly_returned_lazy_members = (
+            required_user_state - previously_returned_user_state.keys()
+        )
+        if newly_returned_lazy_members:
+            # There are some new lazy members we need to fetch.
+            added_types: list[tuple[str, str | None]] = []
+            for new_user_id in newly_returned_lazy_members:
+                added_types.append((EventTypes.Member, new_user_id))
+
+            added_state_filter = StateFilter.from_types(added_types)
+        else:
+            added_state_filter = StateFilter.none()
+
         return _RequiredStateChangesReturn(
             required_state_map_change=None,
-            added_state_filter=StateFilter.none(),
+            added_state_filter=added_state_filter,
+            lazy_members_invalidated=lazy_members_invalidated,
         )
+
     prev_wildcard = prev_required_state_map.get(StateValues.WILDCARD, set())
     request_wildcard = request_required_state_map.get(StateValues.WILDCARD, set())
 
@@ -1551,6 +1641,7 @@ def _required_state_changes(
         return _RequiredStateChangesReturn(
             required_state_map_change=request_required_state_map,
             added_state_filter=StateFilter.none(),
+            lazy_members_invalidated=lazy_members_invalidated,
         )
 
     # If a event type wildcard has been added or removed we don't try and do
@@ -1561,12 +1652,14 @@ def _required_state_changes(
         return _RequiredStateChangesReturn(
             required_state_map_change=request_required_state_map,
             added_state_filter=StateFilter.all(),
+            lazy_members_invalidated=lazy_members_invalidated,
         )
     if prev_wildcard - request_wildcard:
         # Keys were only removed, so we don't have to fetch everything.
         return _RequiredStateChangesReturn(
             required_state_map_change=request_required_state_map,
             added_state_filter=StateFilter.none(),
+            lazy_members_invalidated=lazy_members_invalidated,
         )
 
     # Contains updates to the required state map compared with the previous room
@@ -1577,6 +1670,11 @@ def _required_state_changes(
     # client. Passed to `StateFilter.from_types(...)`
     added: list[tuple[str, str | None]] = []
 
+    # Record any members that were previously explicitly tracked and should now
+    # be tracked as lazy members. This handles the case of membership changing
+    # from e.g. `{@alice:example.com}` to `{LAZY}`.
+    lazy_members_previously_returned: set[str] = set()
+
     # Convert the list of state deltas to map from type to state_keys that have
     # changed.
     changed_types_to_state_keys: dict[str, set[str]] = {}
@@ -1599,6 +1697,39 @@ def _required_state_changes(
             # Nothing *added*, so we skip. Removals happen below.
             continue
 
+        # Handle the special case of adding LAZY membership, where we want to
+        # remember what explicit members we've previously sent down.
+        if event_type == EventTypes.Member:
+            old_state_key_lazy = StateValues.LAZY in old_state_keys
+            request_state_key_lazy = StateValues.LAZY in request_state_keys
+            if not old_state_key_lazy and request_state_key_lazy:
+                # We're adding a LAZY flag. We therefore add any previously
+                # explicit members we've sent down to lazy cache.
+                for state_key in old_state_keys:
+                    if (
+                        state_key == StateValues.WILDCARD
+                        or state_key == StateValues.LAZY
+                    ):
+                        # Ignore non-user IDs.
+                        continue
+
+                    if state_key == StateValues.ME:
+                        # Normalize to proper user ID
+                        state_key = user_id
+
+                    # We remember the user if either a) they haven't been
+                    # invalidated...
+                    if (EventTypes.Member, state_key) not in state_deltas:
+                        lazy_members_previously_returned.add(state_key)
+
+                    # ...or b) if we are going to send the delta down in this
+                    # sync.
+                    if state_key in required_user_state:
+                        lazy_members_previously_returned.add(state_key)
+
+                changes[event_type] = request_state_keys
+                continue
+
         # We only remove state keys from the effective state if they've been
         # removed from the request *and* the state has changed. This ensures
         # that if a client removes and then re-adds a state key, we only send
@@ -1669,9 +1800,23 @@ def _required_state_changes(
                     # LAZY values should also be ignore for event types that are
                     # not membership.
                     pass
+                elif event_type == EventTypes.Member:
+                    if state_key not in previously_returned_user_state:
+                        # Only add *explicit* members we haven't previously sent
+                        # down.
+                        added.append((event_type, state_key))
                 else:
                     added.append((event_type, state_key))
 
+    # We also need to pull out any lazy members that are now required but
+    # haven't previously been returned.
+    for required_user_id in (
+        required_user_state
+        - previously_returned_user_state.keys()
+        - lazy_members_previously_returned
+    ):
+        added.append((EventTypes.Member, required_user_id))
+
     added_state_filter = StateFilter.from_types(added)
 
     # Figure out what changes we need to apply to the effective required state
@@ -1746,9 +1891,13 @@ def _required_state_changes(
         return _RequiredStateChangesReturn(
             required_state_map_change=new_required_state_map,
             added_state_filter=added_state_filter,
+            lazy_members_invalidated=lazy_members_invalidated,
+            lazy_members_previously_returned=lazy_members_previously_returned,
         )
     else:
         return _RequiredStateChangesReturn(
             required_state_map_change=None,
             added_state_filter=added_state_filter,
+            lazy_members_invalidated=lazy_members_invalidated,
+            lazy_members_previously_returned=lazy_members_previously_returned,
         )