Add background job to clear unreferenced state groups

changelog.d/18154.feature

@@ -0,0 +1 @@
+Add background job to clear unreferenced state groups.

synapse/storage/controllers/purge_events.py

@@ -36,6 +36,8 @@
 class PurgeEventsStorageController:
     """High level interface for purging rooms and event history."""
 
+    CLEAR_UNREFERENCED_STATE_GROUPS_PERIOD_MS = 60 * 1000
+
     def __init__(self, hs: "HomeServer", stores: Databases):
         self.stores = stores
 
@@ -44,6 +46,12 @@ def __init__(self, hs: "HomeServer", stores: Databases):
                 self._delete_state_groups_loop, 60 * 1000
             )
 
+        self._last_checked_state_group = 0
+        if hs.config.worker.run_background_tasks:
+            self._clear_unreferenced_state_loop_call = hs.get_clock().looping_call(
+                self._clear_unreferenced_state_groups_loop, self.CLEAR_UNREFERENCED_STATE_GROUPS_PERIOD_MS
+            )
+
     async def purge_room(self, room_id: str) -> None:
         """Deletes all record of a room"""
 
@@ -203,3 +211,58 @@ async def _delete_state_groups(
             room_id,
             groups_to_sequences,
         )
+
+    @wrap_as_background_process("_clear_unreferenced_state_groups_loop")
+    async def _clear_unreferenced_state_groups_loop(self) -> None:
+        """Background task that deletes any state groups that may be pending
+        deletion."""
+
+        # Look for state groups that can be cleaned up.
+        search_limit = 100
+        next_set = await self.stores.main.get_state_groups(
+            self._last_checked_state_group + 1, search_limit
+        )
+        if len(next_set) < search_limit:
+            self._last_checked_state_group = 0
+        else:
+            self._last_checked_state_group = list(next_set)[-1]
+
+        if len(next_set) == 0:
+            return
+
+        referenced = await self.stores.main.get_referenced_state_groups(next_set)
+        next_set -= referenced
+
+        if len(next_set) == 0:
+            return
+
+        referenced = await self.stores.state.get_next_state_groups(next_set)
+        next_set -= set(referenced.values())
+
+        if len(next_set) == 0:
+            return
+
+        # Find all state groups that can be deleted if the original set is deleted.
+        # This set includes the original set, as well as any state groups that would
+        # become unreferenced upon deleting the original set.
+        to_delete = await self._find_unreferenced_groups(next_set)
+
+        if len(to_delete) == 0:
+            return
+
+        pending_deletions = await self.stores.state_deletion.get_pending_deletions(
+            to_delete
+        )
+        to_delete -= pending_deletions.keys()
+
+        if len(to_delete) == 0:
+            return
+
+        # TODO: state_groups_pending_deletion table is never cleaned up after deletion!
+        # TODO: should we be cleaning up any state_group_edges that are dangling after
+        # the deletion as well?
+
+        # Mark the state groups for deletion by the deletion background task.
+        await self.stores.state_deletion.mark_state_groups_as_pending_deletion(
+            to_delete
+        )

synapse/storage/databases/main/state.py

@@ -675,6 +675,45 @@ async def get_referenced_state_groups(
 
         return {row[0] for row in rows}
 
+    async def get_state_groups(
+        self,
+        initial_state_group: int,
+        limit: int,
+    ) -> Set[int]:
+        """Get a list of stored state groups
+
+        Args:
+            initial_state_group: get state groups starting with this one.
+            limit: the maximum number of state groups to return.
+
+        Returns:
+            The set of stored state groups following the initial_state_group.
+        """
+        return await self.db_pool.runInteraction(
+            "get_state_groups",
+            self._get_state_groups_txn,
+            initial_state_group,
+            limit,
+        )
+
+    def _get_state_groups_txn(
+        self,
+        txn: LoggingTransaction,
+        initial_state_group: int,
+        limit: int,
+    ) -> Set[int]:
+        sql = f"""
+        SELECT id from state_groups
+        WHERE id > {initial_state_group}
+        LIMIT {limit}
+        """
+
+        txn.execute(sql)
+
+        rows = txn.fetchall()
+
+        return {row[0] for row in rows}
+
     async def update_state_for_partial_state_event(
         self,
         event: EventBase,

tests/storage/test_purge.py

@@ -270,3 +270,88 @@ def test_purge_unreferenced_state_group(self) -> None:
             )
         )
         self.assertEqual(len(state_groups), 1)
+
+    def test_clear_unreferenced_state_groups(self) -> None:
+        """Test that any unreferenced state groups are automatically cleaned up.
+        """
+
+        self.helper.send(self.room_id, body="test1")
+        state1 = self.helper.send_state(
+            self.room_id, "org.matrix.test", body={"number": 2}
+        )
+        state2 = self.helper.send_state(
+            self.room_id, "org.matrix.test", body={"number": 3}
+        )
+        self.helper.send(self.room_id, body="test4")
+        last = self.helper.send(self.room_id, body="test5")
+
+        # Create an unreferenced state group that has a prev group of one of the
+        # to-be-purged events.
+        prev_group = self.get_success(
+            self.store._get_state_group_for_event(state1["event_id"])
+        )
+        unreferenced_state_group = self.get_success(
+            self.state_store.store_state_group(
+                event_id=last["event_id"],
+                room_id=self.room_id,
+                prev_group=prev_group,
+                delta_ids={("org.matrix.test", ""): state2["event_id"]},
+                current_state_ids=None,
+            )
+        )
+
+        another_unreferenced_state_group = self.get_success(
+            self.state_store.store_state_group(
+                event_id=last["event_id"],
+                room_id=self.room_id,
+                prev_group=unreferenced_state_group,
+                delta_ids={("org.matrix.test", ""): state2["event_id"]},
+                current_state_ids=None,
+            )
+        )
+
+        # Advance so that the background job to clear unreferenced state groups runs
+        self.reactor.advance(
+            1 + self._storage_controllers.purge_events.CLEAR_UNREFERENCED_STATE_GROUPS_PERIOD_MS / 1000
+        )
+
+        # Advance so that the background jobs to delete the state groups runs
+        self.reactor.advance(
+            1 + self.state_deletion_store.DELAY_BEFORE_DELETION_MS / 1000
+        )
+
+        # We expect that the unreferenced state group has been deleted.
+        row = self.get_success(
+            self.state_store.db_pool.simple_select_one_onecol(
+                table="state_groups",
+                keyvalues={"id": unreferenced_state_group},
+                retcol="id",
+                allow_none=True,
+                desc="test_purge_unreferenced_state_group",
+            )
+        )
+        self.assertIsNone(row)
+
+        # We expect that the other unreferenced state group has also been deleted.
+        row = self.get_success(
+            self.state_store.db_pool.simple_select_one_onecol(
+                table="state_groups",
+                keyvalues={"id": another_unreferenced_state_group},
+                retcol="id",
+                allow_none=True,
+                desc="test_purge_unreferenced_state_group",
+            )
+        )
+        self.assertIsNone(row)
+
+        # We expect there to now only be one state group for the room, which is
+        # the state group of the last event (as the only outlier).
+        state_groups = self.get_success(
+            self.state_store.db_pool.simple_select_onecol(
+                table="state_groups",
+                keyvalues={"room_id": self.room_id},
+                retcol="id",
+                desc="test_purge_unreferenced_state_group",
+            )
+        )
+        self.assertEqual(len(state_groups), 8)