Sliding Sync: Return room tags in account data extension (#17707)

The account data extension was also updated to avoid copies when we pull
the data out of the cache.

Fix #17694

Author: MadLittleMods

changelog.d/17707.feature

@@ -0,0 +1 @@
+Return room tags in Sliding Sync account data extension.

synapse/handlers/sliding_sync/extensions.py

@@ -14,7 +14,19 @@
 
 import itertools
 import logging
-from typing import TYPE_CHECKING, AbstractSet, Dict, Mapping, Optional, Sequence, Set
+from typing import (
+    TYPE_CHECKING,
+    AbstractSet,
+    ChainMap,
+    Dict,
+    List,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Sequence,
+    Set,
+    cast,
+)
 
 from typing_extensions import assert_never
 
@@ -381,29 +393,47 @@ async def get_account_data_extension_response(
                 )
             )
 
+            # TODO: This should take into account the `from_token` and `to_token`
             have_push_rules_changed = await self.store.have_push_rules_changed_for_user(
                 user_id, from_token.stream_token.push_rules_key
             )
             if have_push_rules_changed:
-                global_account_data_map = dict(global_account_data_map)
                 # TODO: This should take into account the `from_token` and `to_token`
                 global_account_data_map[
                     AccountDataTypes.PUSH_RULES
                 ] = await self.push_rules_handler.push_rules_for_user(sync_config.user)
         else:
             # TODO: This should take into account the `to_token`
-            all_global_account_data = await self.store.get_global_account_data_for_user(
-                user_id
+            immutable_global_account_data_map = (
+                await self.store.get_global_account_data_for_user(user_id)
             )
 
-            global_account_data_map = dict(all_global_account_data)
-            # TODO: This should take into account the  `to_token`
-            global_account_data_map[
-                AccountDataTypes.PUSH_RULES
-            ] = await self.push_rules_handler.push_rules_for_user(sync_config.user)
+            # Use a `ChainMap` to avoid copying the immutable data from the cache
+            global_account_data_map = ChainMap(
+                {
+                    # TODO: This should take into account the `to_token`
+                    AccountDataTypes.PUSH_RULES: await self.push_rules_handler.push_rules_for_user(
+                        sync_config.user
+                    )
+                },
+                # Cast is safe because `ChainMap` only mutates the top-most map,
+                # see https://github.com/python/typeshed/issues/8430
+                cast(
+                    MutableMapping[str, JsonMapping], immutable_global_account_data_map
+                ),
+            )
 
         # Fetch room account data
-        account_data_by_room_map: Mapping[str, Mapping[str, JsonMapping]] = {}
+        #
+        # List of -> Mapping from room_id to mapping of `type` to `content` of room
+        # account data events.
+        #
+        # This is is a list so we can avoid making copies of immutable data and instead
+        # just provide multiple maps that need to be combined. Normally, we could
+        # reach for `ChainMap` in this scenario, but this is a nested map and accessing
+        # the ChainMap by room_id won't combine the two maps for that room (we would
+        # need a new `NestedChainMap` type class).
+        account_data_by_room_maps: List[Mapping[str, Mapping[str, JsonMapping]]] = []
         relevant_room_ids = self.find_relevant_room_ids_for_extension(
             requested_lists=account_data_request.lists,
             requested_room_ids=account_data_request.rooms,
@@ -418,22 +448,66 @@ async def get_account_data_extension_response(
                         user_id, from_token.stream_token.account_data_key
                     )
                 )
+
+                # Add room tags
+                #
+                # TODO: This should take into account the `from_token` and `to_token`
+                tags_by_room = await self.store.get_updated_tags(
+                    user_id, from_token.stream_token.account_data_key
+                )
+                for room_id, tags in tags_by_room.items():
+                    account_data_by_room_map.setdefault(room_id, {})[
+                        AccountDataTypes.TAG
+                    ] = {"tags": tags}
+
+                account_data_by_room_maps.append(account_data_by_room_map)
             else:
                 # TODO: This should take into account the `to_token`
-                account_data_by_room_map = (
+                immutable_account_data_by_room_map = (
                     await self.store.get_room_account_data_for_user(user_id)
                 )
+                account_data_by_room_maps.append(immutable_account_data_by_room_map)
 
-        # Filter down to the relevant rooms
-        account_data_by_room_map = {
-            room_id: account_data_map
-            for room_id, account_data_map in account_data_by_room_map.items()
-            if room_id in relevant_room_ids
-        }
+                # Add room tags
+                #
+                # TODO: This should take into account the `to_token`
+                tags_by_room = await self.store.get_tags_for_user(user_id)
+                account_data_by_room_maps.append(
+                    {
+                        room_id: {AccountDataTypes.TAG: {"tags": tags}}
+                        for room_id, tags in tags_by_room.items()
+                    }
+                )
+
+        # Filter down to the relevant rooms ... and combine the maps
+        relevant_account_data_by_room_map: MutableMapping[
+            str, Mapping[str, JsonMapping]
+        ] = {}
+        for room_id in relevant_room_ids:
+            # We want to avoid adding empty maps for relevant rooms that have no room
+            # account data so do a quick check to see if it's in any of the maps.
+            is_room_in_maps = False
+            for room_map in account_data_by_room_maps:
+                if room_id in room_map:
+                    is_room_in_maps = True
+                    break
+
+            # If we found the room in any of the maps, combine the maps for that room
+            if is_room_in_maps:
+                relevant_account_data_by_room_map[room_id] = ChainMap(
+                    {},
+                    *(
+                        # Cast is safe because `ChainMap` only mutates the top-most map,
+                        # see https://github.com/python/typeshed/issues/8430
+                        cast(MutableMapping[str, JsonMapping], room_map[room_id])
+                        for room_map in account_data_by_room_maps
+                        if room_map.get(room_id)
+                    ),
+                )
 
         return SlidingSyncResult.Extensions.AccountDataExtension(
             global_account_data_map=global_account_data_map,
-            account_data_by_room_map=account_data_by_room_map,
+            account_data_by_room_map=relevant_account_data_by_room_map,
         )
 
     @trace

synapse/storage/databases/main/account_data.py

@@ -177,7 +177,7 @@ async def get_room_account_data_for_user(
 
         def get_room_account_data_for_user_txn(
             txn: LoggingTransaction,
-        ) -> Dict[str, Dict[str, JsonDict]]:
+        ) -> Dict[str, Dict[str, JsonMapping]]:
             # The 'content != '{}' condition below prevents us from using
             # `simple_select_list_txn` here, as it doesn't support conditions
             # other than 'equals'.
@@ -194,7 +194,7 @@ def get_room_account_data_for_user_txn(
 
             txn.execute(sql, (user_id,))
 
-            by_room: Dict[str, Dict[str, JsonDict]] = {}
+            by_room: Dict[str, Dict[str, JsonMapping]] = {}
             for room_id, account_data_type, content in txn:
                 room_data = by_room.setdefault(room_id, {})
 
@@ -394,7 +394,7 @@ def get_updated_room_account_data_txn(
 
     async def get_updated_global_account_data_for_user(
         self, user_id: str, stream_id: int
-    ) -> Mapping[str, JsonMapping]:
+    ) -> Dict[str, JsonMapping]:
         """Get all the global account_data that's changed for a user.
 
         Args:
@@ -407,7 +407,7 @@ async def get_updated_global_account_data_for_user(
 
         def get_updated_global_account_data_for_user(
             txn: LoggingTransaction,
-        ) -> Dict[str, JsonDict]:
+        ) -> Dict[str, JsonMapping]:
             sql = """
                 SELECT account_data_type, content FROM account_data
                 WHERE user_id = ? AND stream_id > ?
@@ -429,7 +429,7 @@ def get_updated_global_account_data_for_user(
 
     async def get_updated_room_account_data_for_user(
         self, user_id: str, stream_id: int
-    ) -> Dict[str, Dict[str, JsonDict]]:
+    ) -> Dict[str, Dict[str, JsonMapping]]:
         """Get all the room account_data that's changed for a user.
 
         Args:
@@ -442,14 +442,14 @@ async def get_updated_room_account_data_for_user(
 
         def get_updated_room_account_data_for_user_txn(
             txn: LoggingTransaction,
-        ) -> Dict[str, Dict[str, JsonDict]]:
+        ) -> Dict[str, Dict[str, JsonMapping]]:
             sql = """
                 SELECT room_id, account_data_type, content FROM room_account_data
                 WHERE user_id = ? AND stream_id > ?
             """
             txn.execute(sql, (user_id, stream_id))
 
-            account_data_by_room: Dict[str, Dict[str, JsonDict]] = {}
+            account_data_by_room: Dict[str, Dict[str, JsonMapping]] = {}
             for row in txn:
                 room_account_data = account_data_by_room.setdefault(row[0], {})
                 room_account_data[row[1]] = db_to_json(row[2])

synapse/types/handlers/sliding_sync.py

@@ -314,8 +314,8 @@ class AccountDataExtension:
             """The Account Data extension (MSC3959)
 
             Attributes:
-                global_account_data_map: Mapping from `type` to `content` of global account
-                    data events.
+                global_account_data_map: Mapping from `type` to `content` of global
+                    account data events.
                 account_data_by_room_map: Mapping from room_id to mapping of `type` to
                     `content` of room account data events.
             """