Refactor persisting profile updates when persisting events

Instead of passing in a previous membership or a list of target users down in the event internal metadata, use the state delta (or rather, previously calculated sliding sync table changes based on state delta) to figure out if a membership really changed. Then add the member lookups to get the right users to update for.

Somewhat WIP, a few tests still fail and this commit misses a bunch of docstrings and/or cleanups of old functionality.

synapse/storage/databases/main/events.py

@@ -42,6 +42,7 @@ from synapse.api.constants import (
     EventContentFields,
     EventTypes,
     Membership,
+    ProfileUpdateAction,
     RelationTypes,
 )
 from synapse.api.errors import PartialStateConflictError
@@ -76,6 +77,7 @@ from synapse.types import (
     MutableStateMap,
     StateMap,
     StrCollection,
+    UserID,
 )
 from synapse.types.handlers import SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES
 from synapse.types.state import StateFilter
@@ -1209,14 +1211,6 @@ class PersistEventsStore:
                 txn, [ev for ev, _ in events_and_contexts]
             )
 
-        if self._msc4429_enabled:
-            self.store.record_profile_updates_for_membership_changes_from_events(
-                txn=txn,
-                events=[
-                    ev for ev, _ in events_and_contexts if ev.type == EventTypes.Member
-                ],
-            )
-
         # We only update the sliding sync tables for non-backfilled events.
         self._update_sliding_sync_tables_with_new_persisted_events_txn(
             txn, room_id, events_and_contexts
@@ -2127,6 +2121,102 @@ class PersistEventsStore:
             txn, {m for m in members_to_cache_bust if not self.hs.is_mine_id(m)}
         )
 
+        if self._msc4429_enabled:
+            # Handle changes to the profile updates stream.
+            # We've already done a bunch of work calculating the changes needed
+            # for the sliding sync tables, so we may as well re-use that information
+            # here to avoid parsing the state delta again, and handling various
+            # edge cases.
+            profile_update_joins = {
+                c.user_id
+                for c in sliding_sync_table_changes.to_insert_membership_snapshots
+                if self.hs.is_mine_id(c.user_id) and c.membership == Membership.JOIN
+            }
+            profile_update_leaves = {
+                c.user_id
+                for c in sliding_sync_table_changes.to_insert_membership_snapshots
+                if self.hs.is_mine_id(c.user_id) and c.membership == Membership.LEAVE
+            }
+            if profile_update_joins:
+                # Write the profile updates for JOIN events
+                self.store.record_profile_updates_from_membership_changes(
+                    txn=txn,
+                    joined_users=profile_update_joins,
+                )
+            if profile_update_leaves:
+                # Write the profile updates for LEAVE events
+                for user_id in profile_update_leaves:
+                    self._record_profile_updates_for_user_left_room_txn(
+                        txn=txn,
+                        user_id=UserID.from_string(user_id),
+                        room_id=room_id,
+                    )
+
+    def _record_profile_updates_for_user_left_room_txn(
+        self,
+        txn: LoggingTransaction,
+        user_id: UserID,
+        room_id: str,
+    ) -> None:
+        """
+        Record updates into the profile updates stream for when a user leaves
+        the last room with a set of users.
+
+        Doing this clears all old rows from the `profile_updates_per_user` table,
+        to avoid exposing any profile field changes past the point of not being
+        in any common rooms with the user.
+
+        Note, this method lives here in the events store file due to the profile
+        store not having access to the membership store, which we need to re-use
+        the `do_users_share_a_room_txn` method there.
+
+        Args:
+            user_id: The user who left the last common room with a set of users.
+            users_to_update: The set of users who no longer share rooms with the user.
+        """
+        rows = self.db_pool.simple_select_onecol_txn(
+            txn=txn,
+            table="local_current_membership",
+            retcol="user_id",
+            keyvalues={
+                "membership": Membership.JOIN,
+                "room_id": room_id,
+            },
+        )
+        # For each user check if we still share rooms
+        users_sharing_rooms = self.store.do_users_share_a_room_txn(
+            txn=txn,
+            user_id=user_id.to_string(),
+            other_user_ids=set(rows),
+        )
+        users_no_longer_sharing_rooms = set(rows) - set(users_sharing_rooms.keys())
+
+        # First clear the previous rows from the table
+        user_clause, user_args = make_in_list_sql_clause(
+            txn.database_engine,
+            "user_id",
+            users_no_longer_sharing_rooms,
+        )
+        txn.execute(
+            f"""
+                DELETE FROM profile_updates_per_user
+                    WHERE {user_clause}
+                    AND stream_id IN (
+                        SELECT stream_id FROM profile_updates WHERE user_id = ?
+                    )
+            """,
+            (*user_args, user_id.to_string()),
+        )
+
+        # Now record the "left room" action in the stream
+        self.store.record_profile_updates_txn(
+            txn=txn,
+            user_id=user_id,
+            action=ProfileUpdateAction.LEFT_ROOM,
+            field_names=[],
+            target_users=users_no_longer_sharing_rooms,
+        )
+
     @classmethod
     def _get_relevant_sliding_sync_current_state_event_ids_txn(
         cls, txn: LoggingTransaction, room_id: str

synapse/storage/databases/main/profile.py

@@ -24,9 +24,13 @@ from typing import TYPE_CHECKING, Collection, Iterable, cast
 import attr
 from canonicaljson import encode_canonical_json
 
-from synapse.api.constants import Membership, ProfileFields, ProfileUpdateAction
+from synapse.api.constants import (
+    EventTypes,
+    Membership,
+    ProfileFields,
+    ProfileUpdateAction,
+)
 from synapse.api.errors import Codes, StoreError
-from synapse.events import EventBase
 from synapse.replication.tcp.streams._base import ProfileUpdatesStream
 from synapse.storage._base import SQLBaseStore, make_in_list_sql_clause
 from synapse.storage.database import (
@@ -751,73 +755,73 @@ class ProfileWorkerStore(SQLBaseStore):
             return None
 
         # Record updates in the profile updates stream
-        stream_id = self._record_profile_updates_txn(
+        stream_id = self.record_profile_updates_txn(
             txn=txn,
             user_id=user_id,
             action=ProfileUpdateAction.UPDATE,
             field_names=[field_name],
-            target_users=target_users,
         )
 
         return stream_id
 
-    def record_profile_updates_for_membership_changes_from_events(
-        self, *, txn: LoggingTransaction, events: list[EventBase]
+    def record_profile_updates_from_membership_changes(
+        self, *, txn: LoggingTransaction, joined_users: set[str]
     ) -> None:
         """
-        Record profile updates from membership events.
+        Record profile updates from membership changes.
+
+        Note, currently only local user changes are reported.
 
         Adds "joined_room" and "left_room" profile update actions into the
         profile updates stream tables.
 
         Args:
             txn: The transaction to use.
-            events: A list of membership events.
+            changes: A list of changes to room membership state, a tuple of
+                "user ID" and "membership".
         """
         if not self._msc4429_enabled:
             return
 
         assert self._is_events_writer
 
-        for event in events:
-            if (
-                event.membership not in (Membership.JOIN, Membership.LEAVE)
-                or not event.internal_metadata.membership_update_users_in_shared_rooms
-            ):
-                continue
-            if (
-                event.membership == Membership.JOIN
-                and event.internal_metadata.prev_membership != Membership.JOIN
-            ):
-                self._record_profile_updates_txn(
-                    txn=txn,
-                    user_id=UserID.from_string(event.sender),
-                    action=ProfileUpdateAction.JOINED_ROOM,
-                    target_users=set(
-                        event.internal_metadata.membership_update_users_in_shared_rooms
-                    ),
-                    field_names=None,
-                )
-            elif (
-                event.membership == Membership.LEAVE
-                and event.internal_metadata.prev_membership != Membership.LEAVE
-            ):
-                self._record_profile_updates_for_user_left_room_txn(
-                    txn=txn,
-                    user_id=UserID.from_string(event.sender),
-                    users_to_update=set(
-                        event.internal_metadata.membership_update_users_in_shared_rooms
-                    ),
-                )
+        # Get rooms for all the users in the changes
+        rows = self.db_pool.simple_select_many_txn(
+            txn=txn,
+            table="current_state_events",
+            column="state_key",
+            iterable=joined_users,
+            retcols=(
+                "state_key",
+                "room_id",
+            ),
+            keyvalues={
+                "type": EventTypes.Member,
+                "membership": Membership.JOIN,
+            },
+        )
+        user_rooms: dict[str, set[str]] = {user_id: set() for user_id in joined_users}
+        for state_key, room_id in rows:
+            user_rooms[state_key].add(room_id)
 
-    def _record_profile_updates_txn(
+        for user_id in joined_users:
+            self.record_profile_updates_txn(
+                txn=txn,
+                user_id=UserID.from_string(user_id),
+                action=ProfileUpdateAction.JOINED_ROOM,
+                field_names=None,
+                user_rooms=user_rooms[user_id],
+            )
+
+    def record_profile_updates_txn(
         self,
         *,
         txn: LoggingTransaction,
         user_id: UserID,
         action: ProfileUpdateAction,
         field_names: list[str] | None,
-        target_users: set[str],
+        target_users: set[str] | None = None,
+        user_rooms: set[str] | None = None,
     ) -> int | None:
         """
         Record updates into the profile updates stream tables.
@@ -841,6 +845,38 @@ class ProfileWorkerStore(SQLBaseStore):
         else:
             assert not field_names
 
+        if not target_users:
+            if not user_rooms:
+                rows = self.db_pool.simple_select_onecol_txn(
+                    txn=txn,
+                    table="current_state_events",
+                    keyvalues={
+                        "type": EventTypes.Member,
+                        "membership": Membership.JOIN,
+                        "state_key": user_id.to_string(),
+                    },
+                    retcol="room_id",
+                )
+                user_rooms = set(rows)
+
+            rows = self.db_pool.simple_select_many_txn(
+                txn=txn,
+                table="local_current_membership",
+                column="room_id",
+                iterable=user_rooms,
+                retcols=("user_id",),
+                keyvalues={
+                    "membership": Membership.JOIN,
+                },
+            )
+            target_users = {row[0] for row in rows}
+        if action in (ProfileUpdateAction.JOINED_ROOM, ProfileUpdateAction.LEFT_ROOM):
+            target_users.discard(user_id.to_string())
+            if not target_users:
+                # No point writing an update for ourselves, if a membership change and no
+                # other users interested
+                return None
+
         # Record the profile update
         inserted_ts = self.clock.time_msec()
         if field_names:
@@ -966,12 +1002,11 @@ class ProfileWorkerStore(SQLBaseStore):
             if not self._msc4429_enabled:
                 return None
 
-            stream_id = self._record_profile_updates_txn(
+            stream_id = self.record_profile_updates_txn(
                 txn=txn,
                 user_id=user_id,
                 action=ProfileUpdateAction.UPDATE,
                 field_names=[field_name],
-                target_users=target_users,
             )
             return stream_id
 
@@ -1009,12 +1044,11 @@ class ProfileWorkerStore(SQLBaseStore):
                 (user_id.to_string(),),
             )
             # Update the profile updates stream
-            stream_id = self._record_profile_updates_txn(
+            stream_id = self.record_profile_updates_txn(
                 txn=txn,
                 user_id=user_id,
                 action=ProfileUpdateAction.UPDATE,
                 field_names=field_names,
-                target_users=target_users,
             )
             return stream_id
 
@@ -1023,50 +1057,6 @@ class ProfileWorkerStore(SQLBaseStore):
             _delete_profile_txn,
         )
 
-    def _record_profile_updates_for_user_left_room_txn(
-        self,
-        txn: LoggingTransaction,
-        user_id: UserID,
-        users_to_update: set[str],
-    ) -> None:
-        """
-        Record updates into the profile updates stream for when a user leaves
-        the last room with a set of users.
-
-        Doing this clears all old rows from the `profile_updates_per_user` table,
-        to avoid exposing any profile field changes past the point of not being
-        in any common rooms with the user.
-
-        Args:
-            user_id: The user who left the last common room with a set of users.
-            users_to_update: The set of users who no longer share rooms with the user.
-        """
-        # First clear the previous rows from the table
-        user_clause, user_args = make_in_list_sql_clause(
-            txn.database_engine,
-            "user_id",
-            users_to_update,
-        )
-        txn.execute(
-            f"""
-                DELETE FROM profile_updates_per_user
-                    WHERE {user_clause}
-                    AND stream_id IN (
-                        SELECT stream_id FROM profile_updates WHERE user_id = ?
-                    )
-            """,
-            (*user_args, user_id.to_string()),
-        )
-
-        # Now record the "left room" action in the stream
-        self._record_profile_updates_txn(
-            txn=txn,
-            user_id=user_id,
-            action=ProfileUpdateAction.LEFT_ROOM,
-            field_names=[],
-            target_users=users_to_update,
-        )
-
 
 class ProfileStore(ProfileWorkerStore):
     pass

synapse/storage/databases/main/roommember.py

@@ -851,6 +851,49 @@ class RoomMemberWorkerStore(EventsWorkerStore, CacheInvalidationWorkerStore):
     ) -> bool:
         raise NotImplementedError()
 
+    def do_users_share_a_room_txn(
+        self,
+        txn: LoggingTransaction,
+        user_id: str,
+        other_user_ids: Collection[str],
+        exclude_room_id: str | None = None,
+    ) -> dict[str, bool]:
+        clause, args = make_in_list_sql_clause(
+            self.database_engine, "state_key", other_user_ids
+        )
+        # Build SQL args based on whether we are excluding a room ID or not
+        exclude_room_id_clause = ""
+        sql_args = [
+            user_id,
+        ]
+        if exclude_room_id:
+            exclude_room_id_clause = "AND room_id != ?"
+            sql_args.extend([exclude_room_id, *args, exclude_room_id])
+        else:
+            sql_args.extend([*args])
+
+        # This query works by fetching both the list of rooms for the target
+        # user and the set of other users, and then checking if there is any
+        # overlap.
+        sql = f"""
+            SELECT DISTINCT b.state_key
+            FROM (
+                SELECT room_id FROM current_state_events
+                WHERE type = 'm.room.member' AND membership = 'join' AND state_key = ?
+                    {exclude_room_id_clause}
+            ) AS a
+            INNER JOIN (
+                SELECT room_id, state_key FROM current_state_events
+                WHERE type = 'm.room.member' AND membership = 'join' AND {clause}
+                    {exclude_room_id_clause}
+            ) AS b using (room_id)
+        """
+        txn.execute(
+            sql,
+            sql_args,
+        )
+        return {u: True for (u,) in txn}
+
     @cachedList(
         cached_method_name="does_pair_of_users_share_a_room",
         list_name="other_user_ids",
@@ -870,49 +913,14 @@ class RoomMemberWorkerStore(EventsWorkerStore, CacheInvalidationWorkerStore):
         room.
         """
 
-        def do_users_share_a_room_txn(
-            txn: LoggingTransaction, user_ids: Collection[str]
-        ) -> dict[str, bool]:
-            clause, args = make_in_list_sql_clause(
-                self.database_engine, "state_key", user_ids
-            )
-            # Build SQL args based on whether we are excluding a room ID or not
-            exclude_room_id_clause = ""
-            sql_args = [
-                user_id,
-            ]
-            if exclude_room_id:
-                exclude_room_id_clause = "AND room_id != ?"
-                sql_args.extend([exclude_room_id, *args, exclude_room_id])
-            else:
-                sql_args.extend([*args])
-
-            # This query works by fetching both the list of rooms for the target
-            # user and the set of other users, and then checking if there is any
-            # overlap.
-            sql = f"""
-                SELECT DISTINCT b.state_key
-                FROM (
-                    SELECT room_id FROM current_state_events
-                    WHERE type = 'm.room.member' AND membership = 'join' AND state_key = ?
-                        {exclude_room_id_clause}
-                ) AS a
-                INNER JOIN (
-                    SELECT room_id, state_key FROM current_state_events
-                    WHERE type = 'm.room.member' AND membership = 'join' AND {clause}
-                        {exclude_room_id_clause}
-                ) AS b using (room_id)
-            """
-            txn.execute(
-                sql,
-                sql_args,
-            )
-            return {u: True for (u,) in txn}
-
         to_return = {}
         for batch_user_ids in batch_iter(other_user_ids, 1000):
             res = await self.db_pool.runInteraction(
-                "do_users_share_a_room", do_users_share_a_room_txn, batch_user_ids
+                "do_users_share_a_room",
+                self.do_users_share_a_room_txn,
+                user_id,
+                batch_user_ids,
+                exclude_room_id,
             )
             to_return.update(res)