Sliding sync: Add connection tracking to the `account_data` extension (#17695)

This is basically exactly the same logic as for receipts. Essentially we
just need to track which room account data we have and haven't sent down
to clients, and use that when we pull stuff out.

I think this just needs a couple of extra tests written

---------

Co-authored-by: Eric Eastwood <eric.eastwood@beta.gouv.fr>

1 files changed, 37 insertions, 0 deletions

diff --git a/synapse/storage/databases/main/sliding_sync.py b/synapse/storage/databases/main/sliding_sync.py

index 83939d10b0..f2df37fec1 100644
--- a/synapse/storage/databases/main/sliding_sync.py
+++ b/synapse/storage/databases/main/sliding_sync.py
@@ -267,6 +267,15 @@ class SlidingSyncStore(SQLBaseStore):
                 (have_sent_room.status.value, have_sent_room.last_token)
             )
 
+        for (
+            room_id,
+            have_sent_room,
+        ) in per_connection_state.account_data._statuses.items():
+            key_values.append((connection_position, "account_data", room_id))
+            value_values.append(
+                (have_sent_room.status.value, have_sent_room.last_token)
+            )
+
         self.db_pool.simple_upsert_many_txn(
             txn,
             table="sliding_sync_connection_streams",
@@ -407,6 +416,7 @@ class SlidingSyncStore(SQLBaseStore):
         # Now look up the per-room stream data.
         rooms: Dict[str, HaveSentRoom[str]] = {}
         receipts: Dict[str, HaveSentRoom[str]] = {}
+        account_data: Dict[str, HaveSentRoom[str]] = {}
 
         receipt_rows = self.db_pool.simple_select_list_txn(
             txn,
@@ -427,6 +437,8 @@ class SlidingSyncStore(SQLBaseStore):
                 rooms[room_id] = have_sent_room
             elif stream == "receipts":
                 receipts[room_id] = have_sent_room
+            elif stream == "account_data":
+                account_data[room_id] = have_sent_room
             else:
                 # For forwards compatibility we ignore unknown streams, as in
                 # future we want to be able to easily add more stream types.
@@ -435,6 +447,7 @@ class SlidingSyncStore(SQLBaseStore):
         return PerConnectionStateDB(
             rooms=RoomStatusMap(rooms),
             receipts=RoomStatusMap(receipts),
+            account_data=RoomStatusMap(account_data),
             room_configs=room_configs,
         )
 
@@ -452,6 +465,7 @@ class PerConnectionStateDB:
 
     rooms: "RoomStatusMap[str]"
     receipts: "RoomStatusMap[str]"
+    account_data: "RoomStatusMap[str]"
 
     room_configs: Mapping[str, "RoomSyncConfig"]
 
@@ -484,10 +498,21 @@ class PerConnectionStateDB:
             for room_id, status in per_connection_state.receipts.get_updates().items()
         }
 
+        account_data = {
+            room_id: HaveSentRoom(
+                status=status.status,
+                last_token=(
+                    str(status.last_token) if status.last_token is not None else None
+                ),
+            )
+            for room_id, status in per_connection_state.account_data.get_updates().items()
+        }
+
         log_kv(
             {
                 "rooms": rooms,
                 "receipts": receipts,
+                "account_data": account_data,
                 "room_configs": per_connection_state.room_configs.maps[0],
             }
         )
@@ -495,6 +520,7 @@ class PerConnectionStateDB:
         return PerConnectionStateDB(
             rooms=RoomStatusMap(rooms),
             receipts=RoomStatusMap(receipts),
+            account_data=RoomStatusMap(account_data),
             room_configs=per_connection_state.room_configs.maps[0],
         )
 
@@ -524,8 +550,19 @@ class PerConnectionStateDB:
             for room_id, status in self.receipts._statuses.items()
         }
 
+        account_data = {
+            room_id: HaveSentRoom(
+                status=status.status,
+                last_token=(
+                    int(status.last_token) if status.last_token is not None else None
+                ),
+            )
+            for room_id, status in self.account_data._statuses.items()
+        }
+
         return PerConnectionState(
             rooms=RoomStatusMap(rooms),
             receipts=RoomStatusMap(receipts),
+            account_data=RoomStatusMap(account_data),
             room_configs=self.room_configs,
         )