Prune old rows in device_lists_changes_in_room table. (#19473)

Fixes #13043

The usages of the table mostly already correctly handled if we don't
have old entries, as that was needed when we first added the table.

I arbitrarily set the prune time to 30 days. The only use for old
entries is for sync streams that haven't synced since then, and we
should very rarely see sync streams that haven't been used in 30 days.

Reviewable commit-by-commit.

---------

Co-authored-by: Olivier 'reivilibre' <oliverw@element.io>
Co-authored-by: Olivier 'reivilibre' <olivier@librepush.net>

changelog.d/19473.misc

@@ -0,0 +1 @@
+Reduce database disk space usage by pruning old rows from `device_lists_changes_in_room`.

synapse/handlers/device.py

@@ -58,6 +58,7 @@ from synapse.types import (
     DeviceListUpdates,
     JsonDict,
     JsonMapping,
+    MultiWriterStreamToken,
     ScheduledTask,
     StrCollection,
     StreamKeyType,
@@ -1193,7 +1194,16 @@ class DeviceWriterHandler(DeviceHandler):
         changes = await self.store.get_device_list_changes_in_room(
             room_id, device_lists_stream_id
         )
-        local_changes = {(u, d) for u, d in changes if self.hs.is_mine_id(u)}
+        if changes is not None:
+            local_changes = {(u, d) for u, d in changes if self.hs.is_mine_id(u)}
+        else:
+            # The `device_lists_stream_id` is too old, so we need to fall back
+            # to looking for changes for all local users.
+            local_users = await self.store.get_local_users_in_room(room_id)
+            local_changes = await self.store.get_device_changes_for_users(
+                MultiWriterStreamToken(stream=device_lists_stream_id), local_users
+            )
+
         if not local_changes:
             return
 

synapse/storage/databases/main/devices.py

@@ -79,6 +79,19 @@ DROP_DEVICE_LIST_STREAMS_NON_UNIQUE_INDEXES = (
 
 BG_UPDATE_REMOVE_DUP_OUTBOUND_POKES = "remove_dup_outbound_pokes"
 
+# Background update name for adding an index on
+# `device_lists_changes_in_room.inserted_ts`.
+BG_UPDATE_ADD_INSERTED_TS_INDEX = "device_lists_changes_in_room_inserted_ts_idx"
+
+
+# Prunes entries out of the `device_lists_changes_in_room` table that are more
+# than this old.
+PRUNE_DEVICE_LISTS_CHANGES_IN_ROOM_AGE = Duration(days=30)
+
+# The number of rows to delete at once when pruning old entries out of the
+# `device_lists_changes_in_room` table.
+PRUNE_DEVICE_LISTS_BATCH_SIZE = 1000
+
 
 class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
     _device_list_id_gen: MultiWriterIdGenerator
@@ -194,6 +207,10 @@ class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
             self.clock.looping_call(
                 self._prune_old_outbound_device_pokes, Duration(hours=1)
             )
+            self.clock.looping_call(
+                self._prune_device_lists_changes_in_room,
+                Duration(hours=1),
+            )
 
     def process_replication_rows(
         self, stream_name: str, instance_name: str, token: int, rows: Iterable[Any]
@@ -1143,6 +1160,35 @@ class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
             The set of user_ids whose devices have changed since `from_key` (exclusive)
                 until `to_key` (inclusive).
         """
+        return {
+            user_id
+            for user_id, _ in await self.get_device_changes_for_users(
+                from_key, user_ids, to_key
+            )
+        }
+
+    @cancellable
+    async def get_device_changes_for_users(
+        self,
+        from_key: MultiWriterStreamToken,
+        user_ids: Collection[str],
+        to_key: MultiWriterStreamToken | None = None,
+    ) -> set[tuple[str, str]]:
+        """Get set of user/device ID tuple whose devices have changed since `from_key` that
+        are in the given list of user_ids.
+
+        Args:
+            from_key: The minimum device lists stream token to query device list changes for,
+                exclusive.
+            user_ids: If provided, only check if these users have changed their device lists.
+                Otherwise changes from all users are returned.
+            to_key: The maximum device lists stream token to query device list changes for,
+                inclusive. If None then no upper limit is applied.
+
+        Returns:
+            The set of user/device ID tuples whose devices have changed since `from_key`
+            (exclusive) until `to_key` (inclusive).
+        """
         # Get set of users who *may* have changed. Users not in the returned
         # list have definitely not changed.
         user_ids_to_check = self._device_list_stream_cache.get_entities_changed(
@@ -1156,18 +1202,18 @@ class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
         if to_key is None:
             to_key = self.get_device_stream_token()
 
-        def _get_users_whose_devices_changed_txn(
+        def get_device_changes_for_users_txn(
             txn: LoggingTransaction,
             from_key: MultiWriterStreamToken,
             to_key: MultiWriterStreamToken,
-        ) -> set[str]:
+        ) -> set[tuple[str, str]]:
             sql = """
-                SELECT user_id, stream_id, instance_name
+                SELECT user_id, device_id, stream_id, instance_name
                 FROM device_lists_stream
                 WHERE  ? < stream_id AND stream_id <= ? AND %s
             """
 
-            changes: set[str] = set()
+            changes: set[tuple[str, str]] = set()
 
             # Query device changes with a batch of users at a time
             for chunk in batch_iter(user_ids_to_check, 100):
@@ -1179,8 +1225,8 @@ class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
                     [from_key.stream, to_key.get_max_stream_pos()] + args,
                 )
                 changes.update(
-                    user_id
-                    for (user_id, stream_id, instance_name) in txn
+                    (user_id, device_id)
+                    for (user_id, device_id, stream_id, instance_name) in txn
                     if MultiWriterStreamToken.is_stream_position_in_range(
                         low=from_key,
                         high=to_key,
@@ -1192,8 +1238,8 @@ class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
             return changes
 
         return await self.db_pool.runInteraction(
-            "get_users_whose_devices_changed",
-            _get_users_whose_devices_changed_txn,
+            "get_device_changes_for_users",
+            get_device_changes_for_users_txn,
             from_key,
             to_key,
         )
@@ -1699,17 +1745,22 @@ class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
 
         return devices
 
-    @cached()
-    async def _get_min_device_lists_changes_in_room(self) -> int:
-        """Returns the minimum stream ID that we have entries for
-        `device_lists_changes_in_room`
+    def _get_max_pruned_device_lists_changes_in_room_txn(
+        self, txn: LoggingTransaction
+    ) -> int:
+        """Returns the maximum stream ID that has been pruned from
+        `device_lists_changes_in_room`.
+
+        Any queries for stream IDs less than this value cannot be answered
+        completely, as the data has been deleted.
         """
 
-        return await self.db_pool.simple_select_one_onecol(
-            table="device_lists_changes_in_room",
+        return self.db_pool.simple_select_one_onecol_txn(
+            txn,
+            table="device_lists_changes_in_room_max_pruned_stream_id",
             keyvalues={},
-            retcol="COALESCE(MIN(stream_id), 0)",
-            desc="get_min_device_lists_changes_in_room",
+            retcol="stream_id",
+            allow_none=False,
         )
 
     @cancellable
@@ -1728,55 +1779,54 @@ class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
         if not room_ids:
             return set()
 
-        min_stream_id = await self._get_min_device_lists_changes_in_room()
-
-        # Return early if there are no rows to process in device_lists_changes_in_room
-        if min_stream_id > from_token.stream:
-            return None
-
         changed_room_ids = self._device_list_room_stream_cache.get_entities_changed(
             room_ids, from_token.stream
         )
         if not changed_room_ids:
             return set()
 
-        sql = """
-            SELECT user_id, stream_id, instance_name
-            FROM device_lists_changes_in_room
-            WHERE {clause} AND stream_id > ? AND stream_id <= ?
-        """
-
         def _get_device_list_changes_in_rooms_txn(
             txn: LoggingTransaction,
-            chunk: list[str],
-        ) -> set[str]:
-            clause, args = make_in_list_sql_clause(
-                self.database_engine, "room_id", chunk
+        ) -> set[str] | None:
+            # Check if the from_token is too old (i.e. data has been pruned).
+            max_pruned_stream_id = (
+                self._get_max_pruned_device_lists_changes_in_room_txn(txn)
             )
-            args.append(from_token.stream)
-            args.append(to_token.get_max_stream_pos())
+            if max_pruned_stream_id > from_token.stream:
+                return None
 
-            txn.execute(sql.format(clause=clause), args)
-            return {
-                user_id
-                for (user_id, stream_id, instance_name) in txn
-                if MultiWriterStreamToken.is_stream_position_in_range(
-                    low=from_token,
-                    high=to_token,
-                    instance_name=instance_name,
-                    pos=stream_id,
+            changes: set[str] = set()
+
+            for chunk in batch_iter(changed_room_ids, 1000):
+                clause, args = make_in_list_sql_clause(
+                    self.database_engine, "room_id", chunk
                 )
-            }
+                args.append(from_token.stream)
+                args.append(to_token.get_max_stream_pos())
 
-        changes = set()
-        for chunk in batch_iter(changed_room_ids, 1000):
-            changes |= await self.db_pool.runInteraction(
-                "get_device_list_changes_in_rooms",
-                _get_device_list_changes_in_rooms_txn,
-                chunk,
-            )
+                sql = f"""
+                    SELECT user_id, stream_id, instance_name
+                    FROM device_lists_changes_in_room
+                    WHERE {clause} AND stream_id > ? AND stream_id <= ?
+                """
+                txn.execute(sql, args)
+                changes.update(
+                    user_id
+                    for (user_id, stream_id, instance_name) in txn
+                    if MultiWriterStreamToken.is_stream_position_in_range(
+                        low=from_token,
+                        high=to_token,
+                        instance_name=instance_name,
+                        pos=stream_id,
+                    )
+                )
 
-        return changes
+            return changes
+
+        return await self.db_pool.runInteraction(
+            "get_device_list_changes_in_rooms",
+            _get_device_list_changes_in_rooms_txn,
+        )
 
     async def get_all_device_list_changes(self, from_id: int, to_id: int) -> set[str]:
         """Return the set of rooms where devices have changed since the given
@@ -1785,46 +1835,66 @@ class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
         Will raise an exception if the given stream ID is too old.
         """
 
-        min_stream_id = await self._get_min_device_lists_changes_in_room()
-
-        if min_stream_id > from_id:
-            raise Exception("stream ID is too old")
-
-        sql = """
-            SELECT DISTINCT room_id FROM device_lists_changes_in_room
-            WHERE stream_id > ? AND stream_id <= ?
-        """
-
         def _get_all_device_list_changes_txn(
             txn: LoggingTransaction,
-        ) -> set[str]:
+        ) -> set[str] | None:
+            # Check if the from_token is too old (i.e. data has been pruned).
+            max_pruned_stream_id = (
+                self._get_max_pruned_device_lists_changes_in_room_txn(txn)
+            )
+            if max_pruned_stream_id > from_id:
+                logger.warning(
+                    "Given stream ID is too old %d < %d",
+                    from_id,
+                    max_pruned_stream_id,
+                )
+                return None
+
+            sql = """
+                SELECT DISTINCT room_id FROM device_lists_changes_in_room
+                WHERE stream_id > ? AND stream_id <= ?
+            """
+
             txn.execute(sql, (from_id, to_id))
             return {room_id for (room_id,) in txn}
 
-        return await self.db_pool.runInteraction(
+        room_ids = await self.db_pool.runInteraction(
             "get_all_device_list_changes",
             _get_all_device_list_changes_txn,
         )
 
+        if room_ids is None:
+            raise Exception(f"Given stream ID is too old {from_id}")
+
+        return room_ids
+
     async def get_device_list_changes_in_room(
         self, room_id: str, min_stream_id: int
-    ) -> Collection[tuple[str, str]]:
+    ) -> Collection[tuple[str, str]] | None:
         """Get all device list changes that happened in the room since the given
         stream ID.
 
         Returns:
             Collection of user ID/device ID tuples of all devices that have
-            changed
-        """
-
-        sql = """
-            SELECT DISTINCT user_id, device_id FROM device_lists_changes_in_room
-            WHERE room_id = ? AND stream_id > ?
+            changed, or None if the given stream ID is too old and so a complete
+            list cannot be calculated.
         """
 
         def get_device_list_changes_in_room_txn(
             txn: LoggingTransaction,
-        ) -> Collection[tuple[str, str]]:
+        ) -> Collection[tuple[str, str]] | None:
+            # Check if the from_token is too old (i.e. data has been pruned).
+            max_pruned_stream_id = (
+                self._get_max_pruned_device_lists_changes_in_room_txn(txn)
+            )
+            if max_pruned_stream_id > min_stream_id:
+                return None
+
+            sql = """
+                SELECT DISTINCT user_id, device_id FROM device_lists_changes_in_room
+                WHERE room_id = ? AND stream_id > ?
+            """
+
             txn.execute(sql, (room_id, min_stream_id))
             return cast(Collection[tuple[str, str]], txn.fetchall())
 
@@ -2160,6 +2230,8 @@ class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
 
         encoded_context = json_encoder.encode(context)
 
+        now = self.clock.time_msec()
+
         # The `device_lists_changes_in_room.stream_id` column matches the
         # corresponding `stream_id` of the update in the `device_lists_stream`
         # table, i.e. all rows persisted for the same device update will have
@@ -2175,6 +2247,7 @@ class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
                 "instance_name",
                 "converted_to_destinations",
                 "opentracing_context",
+                "inserted_ts",
             ),
             values=[
                 (
@@ -2186,6 +2259,7 @@ class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
                     # We only need to calculate outbound pokes for local users
                     not self.hs.is_mine_id(user_id),
                     encoded_context,
+                    now,
                 )
                 for room_id in room_ids
                 for device_id, stream_id in zip(device_ids, stream_ids)
@@ -2401,6 +2475,144 @@ class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
             desc="set_device_change_last_converted_pos",
         )
 
+    @wrap_as_background_process("prune_device_lists_changes_in_room")
+    async def _prune_device_lists_changes_in_room(self) -> None:
+        """Delete old entries out of the `device_lists_changes_in_room`, so that
+        the table doesn't grow indefinitely.
+        """
+
+        # Let's only do this pruning if the index on inserted_ts has been
+        # created, otherwise this query will be very inefficient.
+        has_index_been_created = (
+            await self.db_pool.updates.has_completed_background_update(
+                BG_UPDATE_ADD_INSERTED_TS_INDEX
+            )
+        )
+        if not has_index_been_created:
+            return
+
+        prune_before_ts = (
+            self.clock.time_msec() - PRUNE_DEVICE_LISTS_CHANGES_IN_ROOM_AGE.as_millis()
+        )
+
+        # Get stream ID corresponding to the prune_before_ts timestamp. We can
+        # delete all rows with a stream ID less than or equal to this, as they
+        # will be older than the cutoff.
+        #
+        # Some rows will have a NULL inserted_ts (due to being inserted before
+        # the column was added), but we can assume that the timestamp will
+        # monotonically increase with stream ID, so we can safely ignore those
+        # rows when calculating the cutoff stream ID. This means that we may end
+        # up keeping some rows with a non-NULL inserted_ts that are older than
+        # the cutoff, but that's better than accidentally deleting rows that are
+        # newer than the cutoff.
+        cutoff_sql = """
+            SELECT stream_id FROM device_lists_changes_in_room
+            WHERE inserted_ts <= ? AND inserted_ts IS NOT NULL
+            ORDER BY inserted_ts DESC
+            LIMIT 1
+        """
+
+        def get_prune_before_stream_id_txn(txn: LoggingTransaction) -> int | None:
+            txn.execute(cutoff_sql, (prune_before_ts,))
+            row = txn.fetchone()
+            return row[0] if row else None
+
+        prune_before_stream_id = await self.db_pool.runInteraction(
+            "prune_device_lists_changes_in_room_get_stream_id",
+            get_prune_before_stream_id_txn,
+        )
+
+        if prune_before_stream_id is None:
+            return
+
+        # Get the max stream ID in the table so we avoid deleting it. We need
+        # to keep the latest row so that we can calculate the maximum stream ID
+        # used.
+        max_stream_id = await self.db_pool.simple_select_one_onecol(
+            table="device_lists_changes_in_room",
+            keyvalues={},
+            retcol="MAX(stream_id)",
+            desc="prune_device_lists_changes_in_room_get_max_stream_id",
+        )
+        if prune_before_stream_id >= max_stream_id:
+            prune_before_stream_id = max_stream_id - 1
+
+        logger.debug(
+            "Pruning device_lists_changes_in_room before stream ID %d (timestamp %d)",
+            prune_before_stream_id,
+            prune_before_ts,
+        )
+
+        # Now delete all rows with stream_id less than the
+        # prune_before_stream_id.
+        #
+        # We also delete in batches to avoid massive churn when initially
+        # clearing out all the old entries.
+        #
+        # We set a minimum stream ID so that when we delete in batches the
+        # database doesn't have to scan through all the (dead) tuples that were just
+        # deleted to find the next batch to delete.
+
+        # The minimum stream ID to delete in the next batch, c.f. comment above.
+        # We default to 0 here as that is less than all possible stream IDs.
+        min_stream_id = 0
+
+        def prune_device_lists_changes_in_room_txn(txn: LoggingTransaction) -> int:
+            nonlocal min_stream_id
+
+            delete_sql = """
+                DELETE FROM device_lists_changes_in_room
+                WHERE stream_id IN (
+                    SELECT stream_id FROM device_lists_changes_in_room
+                    WHERE ? < stream_id AND stream_id <= ?
+                    ORDER BY stream_id ASC
+                    LIMIT ?
+                )
+                RETURNING stream_id
+            """
+            txn.execute(
+                delete_sql,
+                (min_stream_id, prune_before_stream_id, PRUNE_DEVICE_LISTS_BATCH_SIZE),
+            )
+
+            # We can't use rowcount as that is incorrect on SQLite when using
+            # RETURNING.
+            num_deleted = 0
+            for row in txn:
+                num_deleted += 1
+                min_stream_id = max(min_stream_id, row[0])
+
+            return num_deleted
+
+        num_rows_deleted = 0
+        while True:
+            batch_deleted = await self.db_pool.runInteraction(
+                "prune_device_lists_changes_in_room",
+                prune_device_lists_changes_in_room_txn,
+            )
+            num_rows_deleted += batch_deleted
+            if batch_deleted < PRUNE_DEVICE_LISTS_BATCH_SIZE:
+                break
+
+            # Sleep for a short time to avoid hammering the database too much if
+            # there are a lot of rows to delete.
+            await self.clock.sleep(Duration(milliseconds=100))
+
+        if num_rows_deleted:
+            # Update the max pruned stream ID tracking table so that the
+            # safety check knows data up to this point has been deleted.
+            await self.db_pool.simple_update_one(
+                table="device_lists_changes_in_room_max_pruned_stream_id",
+                keyvalues={},
+                updatevalues={"stream_id": prune_before_stream_id},
+                desc="prune_device_lists_changes_in_room_update_max_pruned",
+            )
+
+            logger.info(
+                "Pruned %d rows from device_lists_changes_in_room", num_rows_deleted
+            )
+
 
 class DeviceBackgroundUpdateStore(SQLBaseStore):
     _instance_name: str
@@ -2459,6 +2671,15 @@ class DeviceBackgroundUpdateStore(SQLBaseStore):
             columns=["room_id", "stream_id"],
         )
 
+        # Add indexes to speed up pruning of device_lists_changes_in_room
+        self.db_pool.updates.register_background_index_update(
+            BG_UPDATE_ADD_INSERTED_TS_INDEX,
+            index_name="device_lists_changes_in_room_inserted_ts_idx",
+            table="device_lists_changes_in_room",
+            columns=["inserted_ts"],
+            where_clause="inserted_ts IS NOT NULL",
+        )
+
     async def _drop_device_list_streams_non_unique_indexes(
         self, progress: JsonDict, batch_size: int
     ) -> int:

synapse/storage/schema/main/delta/94/03_device_lists_room_timestamp.sql

@@ -0,0 +1,18 @@
+--
+-- This file is licensed under the Affero General Public License (AGPL) version 3.
+--
+-- Copyright (C) 2025 Element Creations, Ltd
+--
+-- This program is free software: you can redistribute it and/or modify
+-- it under the terms of the GNU Affero General Public License as
+-- published by the Free Software Foundation, either version 3 of the
+-- License, or (at your option) any later version.
+--
+-- See the GNU Affero General Public License for more details:
+-- <https://www.gnu.org/licenses/agpl-3.0.html>.
+
+ALTER TABLE device_lists_changes_in_room ADD COLUMN inserted_ts BIGINT;
+
+-- Add a background update to add index
+INSERT INTO background_updates (ordering, update_name, progress_json) VALUES
+  (9403, 'device_lists_changes_in_room_inserted_ts_idx', '{}');

synapse/storage/schema/main/delta/94/04_device_lists_changes_max_pruned.sql

@@ -0,0 +1,34 @@
+--
+-- This file is licensed under the Affero General Public License (AGPL) version 3.
+--
+-- Copyright (C) 2026 Element Creations Ltd
+--
+-- This program is free software: you can redistribute it and/or modify
+-- it under the terms of the GNU Affero General Public License as
+-- published by the Free Software Foundation, either version 3 of the
+-- License, or (at your option) any later version.
+--
+-- See the GNU Affero General Public License for more details:
+-- <https://www.gnu.org/licenses/agpl-3.0.html>.
+
+
+-- Tracks the maximum stream_id that has been deleted (pruned) from the
+-- device_lists_changes_in_room table. This is used to determine whether it's
+-- safe to read from that table for a given stream_id — if the requested
+-- stream_id is < the value here, the data has been pruned and the table cannot
+-- provide a complete answer.
+--
+-- We need a separate table, rather than looking at the minimum stream_id in the
+-- device_lists_changes_in_room table, because not all valid stream IDs will
+-- have entries in the table. This could lead to situations where the minimum
+-- stream ID was potentially much more recent than when we actually pruned. This
+-- would cause us to incorrectly think that the table was not safe to read from,
+-- when in fact it was.
+CREATE TABLE IF NOT EXISTS device_lists_changes_in_room_max_pruned_stream_id (
+    Lock CHAR(1) NOT NULL DEFAULT 'X' UNIQUE,
+    stream_id BIGINT NOT NULL
+);
+
+-- We assume that nothing has been deleted from the device_lists_changes_in_room
+-- table, so we can set the initial value to 0.
+INSERT INTO device_lists_changes_in_room_max_pruned_stream_id (stream_id) VALUES (0);

tests/handlers/test_device.py

@@ -21,20 +21,42 @@
 #
 
 from unittest import mock
+from unittest.mock import AsyncMock, Mock, patch
 
+import signedjson.key
+from parameterized import parameterized
+from signedjson.types import SigningKey
+
+from twisted.internet import defer
 from twisted.internet.defer import ensureDeferred
 from twisted.internet.testing import MemoryReactor
 
-from synapse.api.constants import RoomEncryptionAlgorithms
+from synapse.api.constants import EventTypes, JoinRules, RoomEncryptionAlgorithms
 from synapse.api.errors import NotFoundError, SynapseError
+from synapse.api.room_versions import RoomVersions
 from synapse.appservice import ApplicationService
+from synapse.crypto.event_signing import add_hashes_and_signatures
+from synapse.events import EventBase, FrozenEventV3
+from synapse.federation.federation_client import SendJoinResult
+from synapse.federation.transport.client import (
+    StateRequestResponse,
+    TransportLayerClient,
+)
+from synapse.federation.units import Transaction
 from synapse.handlers.device import MAX_DEVICE_DISPLAY_NAME_LEN, DeviceWriterHandler
 from synapse.rest import admin
 from synapse.rest.client import devices, login, register
 from synapse.server import HomeServer
 from synapse.storage.databases.main.appservice import _make_exclusive_regex
-from synapse.types import JsonDict, UserID, create_requester
+from synapse.types import (
+    JsonDict,
+    StateMap,
+    UserID,
+    create_requester,
+    get_domain_from_id,
+)
 from synapse.util.clock import Clock
+from synapse.util.duration import Duration
 from synapse.util.task_scheduler import TaskScheduler
 
 from tests import unittest
@@ -581,3 +603,334 @@ class DehydrationTestCase(unittest.HomeserverTestCase):
         self.assertTrue(len(res["next_batch"]) > 1)
         self.assertEqual(len(res["events"]), 1)
         self.assertEqual(res["events"][0]["content"]["body"], "foo")
+
+
+@patch("synapse.crypto.keyring.Keyring.process_request", AsyncMock(return_value=None))
+class DeviceUnPartialStateTestCase(unittest.HomeserverTestCase):
+    """Tests that local device list changes during partial state are sent to
+    remote servers when the room un-partials."""
+
+    servlets = [
+        admin.register_servlets,
+        login.register_servlets,
+    ]
+
+    # The two remote servers to fake
+    REMOTE1_SERVER_NAME = "remote1"
+    REMOTE1_SERVER_SIGNATURE_KEY = signedjson.key.generate_signing_key("test")
+    REMOTE1_USER = f"@user:{REMOTE1_SERVER_NAME}"
+
+    REMOTE2_SERVER_NAME = "remote2"
+    REMOTE2_SERVER_SIGNATURE_KEY = signedjson.key.generate_signing_key("test")
+    REMOTE2_USER = f"@user:{REMOTE2_SERVER_NAME}"
+
+    def default_config(self) -> JsonDict:
+        config = super().default_config()
+        # Enable federation so that get_device_updates_by_remote works.
+        config["federation_sender_instances"] = ["master"]
+        return config
+
+    def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer:
+        # Mock the federation transport client to prevent actual network calls.
+        self.federation_transport_client = AsyncMock(TransportLayerClient)
+
+        self.federation_transport_client.send_transaction.return_value = {}
+
+        hs = self.setup_test_homeserver(
+            federation_transport_client=self.federation_transport_client,
+        )
+        handler = hs.get_device_handler()
+        assert isinstance(handler, DeviceWriterHandler)
+        self.device_handler = handler
+        self.store = hs.get_datastores().main
+
+        return hs
+
+    def _build_public_room(self) -> StateMap[EventBase]:
+        """Build a public room DAG that has REMOTE1 in it"""
+
+        room_id = f"!room:{self.REMOTE1_SERVER_NAME}"
+        room_version = RoomVersions.V10
+
+        events: list[EventBase] = []
+
+        # First we make the create event
+        create_event_dict: JsonDict = {
+            "auth_events": [],
+            "content": {
+                "creator": self.REMOTE1_USER,
+                "room_version": room_version.identifier,
+            },
+            "depth": 0,
+            "origin_server_ts": 0,
+            "prev_events": [],
+            "room_id": room_id,
+            "sender": self.REMOTE1_USER,
+            "state_key": "",
+            "type": EventTypes.Create,
+        }
+
+        add_hashes_and_signatures(
+            room_version,
+            create_event_dict,
+            self.REMOTE1_SERVER_NAME,
+            self.REMOTE1_SERVER_SIGNATURE_KEY,
+        )
+
+        create_event = FrozenEventV3(create_event_dict, room_version, {}, None)
+        events.append(create_event)
+
+        room_version = self.hs.config.server.default_room_version
+        join_event_dict: JsonDict = {
+            "auth_events": [
+                create_event.event_id,
+            ],
+            "content": {"membership": "join"},
+            "depth": 1,
+            "origin_server_ts": 100,
+            "prev_events": [create_event.event_id],
+            "sender": self.REMOTE1_USER,
+            "state_key": self.REMOTE1_USER,
+            "room_id": room_id,
+            "type": EventTypes.Member,
+        }
+        add_hashes_and_signatures(
+            room_version,
+            join_event_dict,
+            self.hs.hostname,
+            self.hs.signing_key,
+        )
+        join_event = FrozenEventV3(join_event_dict, room_version, {}, None)
+        events.append(join_event)
+
+        # Then set the join rules to public
+        join_rules_event_dict: JsonDict = {
+            "auth_events": [create_event.event_id, join_event.event_id],
+            "content": {"join_rule": JoinRules.PUBLIC},
+            "depth": 2,
+            "origin_server_ts": 200,
+            "prev_events": [join_event.event_id],
+            "room_id": room_id,
+            "sender": self.REMOTE1_USER,
+            "state_key": "",
+            "type": EventTypes.JoinRules,
+        }
+
+        add_hashes_and_signatures(
+            room_version,
+            join_rules_event_dict,
+            self.REMOTE1_SERVER_NAME,
+            self.REMOTE1_SERVER_SIGNATURE_KEY,
+        )
+        join_rules_event = FrozenEventV3(join_rules_event_dict, room_version, {}, None)
+        events.append(join_rules_event)
+
+        return {(event.type, event.state_key): event for event in events}
+
+    def _build_signed_join_event(
+        self,
+        room_id: str,
+        user: str,
+        signing_key: SigningKey,
+        state: StateMap[EventBase],
+    ) -> FrozenEventV3:
+        """Build a join event for the local user, signed by the local server."""
+
+        latest_event = max(state.values(), key=lambda e: e.depth)
+
+        room_version = self.hs.config.server.default_room_version
+        join_event_dict: JsonDict = {
+            "auth_events": [
+                state[(EventTypes.Create, "")].event_id,
+                state[(EventTypes.JoinRules, "")].event_id,
+            ],
+            "content": {"membership": "join"},
+            "depth": latest_event.depth + 1,
+            "origin_server_ts": latest_event.origin_server_ts + 100,
+            "prev_events": [latest_event.event_id],
+            "sender": user,
+            "state_key": user,
+            "room_id": room_id,
+            "type": EventTypes.Member,
+        }
+        add_hashes_and_signatures(
+            room_version,
+            join_event_dict,
+            get_domain_from_id(user),
+            signing_key,
+        )
+        return FrozenEventV3(join_event_dict, room_version, {}, None)
+
+    @parameterized.expand([("not_pruned", False), ("pruned", True)])
+    @patch(
+        "synapse.storage.databases.main.devices.PRUNE_DEVICE_LISTS_CHANGES_IN_ROOM_AGE",
+        Duration(minutes=1),
+    )
+    def test_local_device_changes_sent_to_new_servers_on_un_partial_state(
+        self, _test_suffix: str, prune_device_lists_change_in_room: bool
+    ) -> None:
+        """When a room un-partials, local device list changes made during the
+        partial state period should be sent to remote servers that were NOT
+        known at the time of the partial join.
+
+        We do this by creating a room with one remote server, partialling
+        joining it, then receiving a join event from a second remote server. The
+        second remote server should receive a device list update EDU for any
+        local device changes that happened during the partial state period.
+
+        We parameterize this test over whether during the unpartial process we
+        prune the `device_list_changes_in_room` table, to check that the
+        unpartial process correctly handles the case.
+        """
+
+        local_user = self.register_user("alice", "password")
+        self.login("alice", "password")
+
+        # Build the remote room's state events.
+        room_state = self._build_public_room()
+
+        # Before joining, we mock out the federation endpoints that are used
+        # during the unpartial process, so that we can control when the
+        # unpartial process completes.
+        get_room_state_ids_deferred: defer.Deferred[JsonDict] = defer.Deferred()
+        get_room_state_deferred: defer.Deferred[StateRequestResponse] = defer.Deferred()
+        self.federation_transport_client.get_room_state_ids = Mock(
+            side_effect=[get_room_state_ids_deferred]
+        )
+        self.federation_transport_client.get_room_state = Mock(
+            side_effect=[get_room_state_deferred]
+        )
+
+        # Now make the local server partially join the room.
+        room_id = room_state[(EventTypes.Create, "")].room_id
+        room_version = room_state[(EventTypes.Create, "")].room_version
+
+        local_join_event = self._build_signed_join_event(
+            room_id, local_user, self.hs.signing_key, room_state
+        )
+
+        # Mock the federation client endpoints for the partial join.
+        mock_make_membership_event = AsyncMock(
+            return_value=(self.REMOTE1_SERVER_NAME, local_join_event, room_version)
+        )
+        mock_send_join = AsyncMock(
+            return_value=SendJoinResult(
+                local_join_event,
+                self.REMOTE1_SERVER_NAME,
+                state=list(room_state.values()),
+                auth_chain=list(room_state.values()),
+                partial_state=True,
+                # Only REMOTE1_SERVER_NAME is known at join time.
+                servers_in_room={self.REMOTE1_SERVER_NAME},
+            )
+        )
+
+        fed_handler = self.hs.get_federation_handler()
+        fed_client = self.hs.get_federation_client()
+        with (
+            patch.object(
+                fed_client, "make_membership_event", mock_make_membership_event
+            ),
+            patch.object(fed_client, "send_join", mock_send_join),
+        ):
+            self.get_success(
+                fed_handler.do_invite_join(
+                    [self.REMOTE1_SERVER_NAME], room_id, local_user, {}
+                )
+            )
+
+        # The room should now be in partial state.
+        self.assertTrue(self.get_success(self.store.is_partial_state_room(room_id)))
+
+        # A local device change happens while the room is in partial state.
+        self.get_success(
+            self.store.add_device_change_to_streams(
+                local_user, ["NEW_DEVICE"], [room_id]
+            )
+        )
+
+        if prune_device_lists_change_in_room:
+            # Add a device change for another room, as we won't prune the most
+            # recent change.
+            self.get_success(
+                self.store.add_device_change_to_streams(
+                    "@other:user", ["device1"], ["!some:room"]
+                )
+            )
+
+            # Now prune the device list changes for the room. This simulates the
+            # case where the unpartial process prunes the
+            # `device_list_changes_in_room` table before processing the device
+            # list changes.
+            self.reactor.advance(120)  # Advance past the pruning threshold
+            self.get_success(self.store._prune_device_lists_changes_in_room())
+
+            # Assert we actually pruned the device list changes for the room.
+            room_ids = self.get_success(
+                self.store.db_pool.simple_select_onecol(
+                    table="device_lists_changes_in_room",
+                    keyvalues={},
+                    retcol="room_id",
+                )
+            )
+            self.assertCountEqual(room_ids, ["!some:room"])
+
+        # Join the second server
+        new_state = dict(room_state)
+        new_state[(EventTypes.Member, local_user)] = local_join_event
+        join_event_2 = self._build_signed_join_event(
+            room_id,
+            self.REMOTE2_USER,
+            self.REMOTE2_SERVER_SIGNATURE_KEY,
+            new_state,
+        )
+
+        self.get_success(
+            self.hs.get_federation_event_handler().on_receive_pdu(
+                self.REMOTE2_SERVER_NAME, join_event_2
+            )
+        )
+
+        # Some EDUs may get sent out immediately, such as presence updates.
+        # However, we only care about the device list update EDU sent by the
+        # unpartialling process. Let's wait a few seconds and reset the mock.
+        self.reactor.advance(5)
+        self.federation_transport_client.send_transaction.reset_mock()
+
+        # We now unblock the unpartial processs by returning the room state and
+        # state ids. This should trigger the device list update to be sent to
+        # REMOTE2_SERVER_NAME.
+        self.federation_transport_client.get_room_state_ids.assert_called_once_with(
+            self.REMOTE1_SERVER_NAME,
+            room_id,
+            event_id=local_join_event.prev_event_ids()[0],
+        )
+
+        get_room_state_ids_deferred.callback(
+            {
+                "pdu_ids": [event.event_id for event in room_state.values()],
+                "auth_event_ids": [],
+            }
+        )
+        get_room_state_deferred.callback(
+            StateRequestResponse(
+                state=list(room_state.values()),
+                auth_events=[],
+            )
+        )
+
+        # The device list EDU isn't necessarily sent out immediately
+        self.reactor.advance(30)
+
+        # Check that only one transaction was sent, and that it contains the
+        # device list update EDU for the new device to REMOTE2_SERVER_NAME.
+        self.federation_transport_client.send_transaction.assert_called_once()
+        args, _ = self.federation_transport_client.send_transaction.call_args
+        transaction: Transaction = args[0]
+
+        self.assertEqual(transaction.destination, self.REMOTE2_SERVER_NAME)
+        self.assertEqual(len(transaction.edus), 1)
+
+        edu = transaction.edus[0]
+        self.assertEqual(edu["edu_type"], "m.device_list_update")
+        self.assertEqual(edu["content"]["device_id"], "NEW_DEVICE")

tests/storage/test_devices.py

@@ -19,15 +19,21 @@
 #
 #
 
+import itertools
 from typing import Collection
+from unittest.mock import patch
 
 from twisted.internet.testing import MemoryReactor
 
 import synapse.api.errors
 from synapse.api.constants import EduTypes
 from synapse.server import HomeServer
+from synapse.storage.databases.main.devices import (
+    PRUNE_DEVICE_LISTS_CHANGES_IN_ROOM_AGE,
+)
 from synapse.types import JsonDict
 from synapse.util.clock import Clock
+from synapse.util.duration import Duration
 
 from tests.unittest import HomeserverTestCase
 
@@ -351,3 +357,101 @@ class DeviceStoreTestCase(HomeserverTestCase):
             synapse.api.errors.StoreError,
         )
         self.assertEqual(404, exc.value.code)
+
+    @patch("synapse.storage.databases.main.devices.PRUNE_DEVICE_LISTS_BATCH_SIZE", 5)
+    def test_prune_old_device_lists_changes_in_room(self) -> None:
+        """Test that old entries in the `device_lists_changes_in_room` table are pruned properly."""
+
+        # Pretend the user is in a few rooms.
+        room_ids = [f"!room{i}:test" for i in range(20)]
+
+        # Create a generator for device IDs so we can easily create many unique
+        # device IDs without having to keep track of the count ourselves.
+        device_id_gen = (f"device_id{i}" for i in itertools.count())
+
+        def get_devices_in_room_status() -> tuple[int, str]:
+            """Helper function to get the count of entries in
+            `device_lists_changes_in_room` and the minimum device_id."""
+            return self.get_success(
+                self.store.db_pool.simple_select_one(
+                    table="device_lists_changes_in_room",
+                    keyvalues={},
+                    retcols=("COUNT(*)", "MIN(device_id)"),
+                )
+            )
+
+        # First we add some initial entries to the `device_lists_changes_in_room`.
+        self.get_success(
+            self.store.add_device_change_to_streams(
+                user_id="@user_id:test",
+                device_ids=[next(device_id_gen) for _ in range(10)],
+                room_ids=room_ids,
+            )
+        )
+
+        # Advance the reactor a while, but not long enough to trigger pruning.
+        self.reactor.advance(Duration(hours=1).as_secs())
+
+        # The `device_lists_changes_in_room` table should now have 10 *
+        # len(room_ids) entries, and the minimum device_id should be
+        # `device_id0`.
+        count, min_device_id = get_devices_in_room_status()
+        self.assertEqual(count, 10 * len(room_ids))
+        self.assertEqual(min_device_id, "device_id0")
+
+        # Record the max pruned stream ID before pruning, so we can check
+        # that this correctly updates after pruning.
+        starting_max_pruned_id = self.get_success(
+            self.store.db_pool.runInteraction(
+                "get_max_pruned_device_lists_changes_in_room",
+                self.store._get_max_pruned_device_lists_changes_in_room_txn,
+            )
+        )
+
+        # Now we add some more entries.
+        self.get_success(
+            self.store.add_device_change_to_streams(
+                user_id="@user_id:test",
+                device_ids=[next(device_id_gen) for _ in range(10)],
+                room_ids=room_ids,
+            )
+        )
+
+        # Advance the reactor a while more, so that the first batch of entries is
+        # now old enough to be pruned.
+        self.reactor.advance(
+            (PRUNE_DEVICE_LISTS_CHANGES_IN_ROOM_AGE - Duration(minutes=30)).as_secs()
+        )
+
+        # Advance repeatedly a bit so that the pruning process can run to completion.
+        for _ in range(10):
+            self.reactor.advance(Duration(milliseconds=110).as_secs())
+
+        # Check that the old entries have been pruned, and the new entries are still there.
+        count, min_device_id = get_devices_in_room_status()
+        self.assertEqual(count, 10 * len(room_ids))
+        self.assertEqual(min_device_id, "device_id10")
+
+        # We should always keep the most recent entries, even if they are old enough to be pruned.
+        self.reactor.advance(
+            (PRUNE_DEVICE_LISTS_CHANGES_IN_ROOM_AGE + Duration(minutes=30)).as_secs()
+        )
+
+        # Advance repeatedly a bit so that the pruning process can run to completion.
+        for _ in range(10):
+            self.reactor.advance(Duration(milliseconds=110).as_secs())
+
+        count, min_device_id = get_devices_in_room_status()
+        # We should always keep the most recent entries so that we can
+        # calculate the maximum stream ID used.
+        self.assertEqual(count, len(room_ids))
+        self.assertEqual(min_device_id, "device_id19")
+
+        # Check that the max pruned stream ID has been advanced after pruning.
+        max_pruned_id = self.get_success(
+            self.store.db_pool.runInteraction(
+                "get_max_pruned_device_lists_changes_in_room",
+                self.store._get_max_pruned_device_lists_changes_in_room_txn,
+            )
+        )
+        self.assertGreater(max_pruned_id, starting_max_pruned_id)