Add a prune job for the profile updates stream tables

Taken and adapted from https://github.com/element-hq/synapse/pull/19473

Didn't add the "safety" table due to afaik that has more relevance with the device updates, and not so much here for the profile updates.

synapse/storage/databases/main/profile.py

@@ -19,6 +19,7 @@
 #
 #
 import json
+import logging
 from typing import TYPE_CHECKING, Collection, Iterable, cast
 
 import attr
@@ -26,6 +27,7 @@ from canonicaljson import encode_canonical_json
 
 from synapse.api.constants import ProfileFields, ProfileUpdateAction
 from synapse.api.errors import Codes, StoreError
+from synapse.metrics.background_process_metrics import wrap_as_background_process
 from synapse.replication.tcp.streams._base import ProfileUpdatesStream
 from synapse.storage._base import SQLBaseStore, make_in_list_sql_clause
 from synapse.storage.database import (
@@ -37,14 +39,24 @@ from synapse.storage.databases.main.roommember import ProfileInfo
 from synapse.storage.engines import PostgresEngine, Sqlite3Engine
 from synapse.storage.util.id_generators import MultiWriterIdGenerator
 from synapse.types import JsonDict, JsonValue, UserID
+from synapse.util.duration import Duration
 
 if TYPE_CHECKING:
     from synapse.server import HomeServer
 
+logger = logging.getLogger(__name__)
 
 # The number of bytes that the serialized profile can have.
 MAX_PROFILE_SIZE = 65536
 
+# Prunes entries out of the `profile_updates` and `profile_updates_per_user` tables
+# that are more than this old.
+PRUNE_PROFILE_UPDATES_AGE = Duration(days=30)
+
+# The number of rows to delete at once when pruning old entries out of the
+# `profile_updates` and `profile_updates_per_user` tables.
+PRUNE_PROFILE_UPDATES_BATCH_SIZE = 1000
+
 
 @attr.s(slots=True, frozen=True, auto_attribs=True)
 class ProfileUpdate:
@@ -95,6 +107,11 @@ class ProfileWorkerStore(SQLBaseStore):
             sequence_name="profile_updates_sequence",
             writers=hs.config.worker.writers.profile_updates,
         )
+        if hs.config.worker.run_background_tasks:
+            self.clock.looping_call(
+                self._prune_profile_updates,
+                Duration(hours=1),
+            )
 
     async def populate_full_user_id_profiles(
         self, progress: JsonDict, batch_size: int
@@ -927,6 +944,129 @@ class ProfileWorkerStore(SQLBaseStore):
             keyvalues={"full_user_id": user_id.to_string()},
         )
 
+    @wrap_as_background_process("prune_profile_updates")
+    async def _prune_profile_updates(self) -> None:
+        """Delete old entries out of the `profile_updates` and
+        `profile_updates_per_user` tables, so that the tables don't grow indefinitely.
+        """
+        prune_before_ts = self.clock.time_msec() - PRUNE_PROFILE_UPDATES_AGE.as_millis()
+        cutoff_sql = """
+            SELECT stream_id FROM profile_updates
+            WHERE inserted_ts <= ?
+            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_profile_updates_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="profile_updates",
+            keyvalues={},
+            retcol="MAX(stream_id)",
+            desc="prune_profile_updates_get_max_stream_id",
+        )
+        if prune_before_stream_id >= max_stream_id:
+            prune_before_stream_id = max_stream_id - 1
+
+        logger.debug(
+            "Pruning profile_updates 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_profile_updates_txn(txn: LoggingTransaction) -> int:
+            nonlocal min_stream_id
+
+            assert table in ("profile_updates", "profile_updates_per_user")
+            delete_sql = """
+                    DELETE FROM %s
+                    WHERE stream_id IN (
+                        SELECT stream_id FROM %s
+                        WHERE ? < stream_id AND stream_id <= ?
+                        ORDER BY stream_id ASC
+                        LIMIT ?
+                    )
+                    RETURNING stream_id
+                """ % (table, table)
+            txn.execute(
+                delete_sql,
+                (
+                    min_stream_id,
+                    prune_before_stream_id,
+                    PRUNE_PROFILE_UPDATES_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
+
+        # Do this twice, first for the per_user table, then for the main table
+        for table in ("profile_updates_per_user", "profile_updates"):
+            progress_num_rows_deleted = 0
+            while True:
+                batch_deleted = await self.db_pool.runInteraction(
+                    f"prune_{table}",
+                    prune_profile_updates_txn,
+                )
+
+                finished = batch_deleted < PRUNE_PROFILE_UPDATES_BATCH_SIZE
+
+                progress_num_rows_deleted += batch_deleted
+
+                # Periodically report progress in the logs. We do this either when
+                # we've deleted a significant number of rows or when we've finished
+                # deleting all rows in this round.
+                if finished or progress_num_rows_deleted > 10000:
+                    logger.info(
+                        "Pruned %d rows from %s",
+                        progress_num_rows_deleted,
+                        table,
+                    )
+                    progress_num_rows_deleted = 0
+
+                if finished:
+                    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))
+
+            # Reset the minimum stream id for our next table
+            min_stream_id = 0
+
 
 class ProfileStore(ProfileWorkerStore):
     pass

synapse/storage/schema/main/delta/95/01_profile_updates.sql

@@ -27,12 +27,14 @@ CREATE TABLE profile_updates (
   -- This is only required if "action" is "update"
   field_name TEXT NULL,
 
-  -- Unix timestamp for debugging purposes
+  -- Unix timestamp. Used to determine when to cull rows (to prevent the table
+  -- from growing indefinitely).
   inserted_ts BIGINT NOT NULL
 );
 
 CREATE INDEX profile_updates_by_user ON profile_updates (user_id, stream_id);
 CREATE INDEX profile_updates_by_field ON profile_updates (field_name, stream_id);
+CREATE INDEX profile_updates_inserted_ts ON profile_updates (inserted_ts);
 
 -- Track which local users should receive each profile update.
 CREATE TABLE profile_updates_per_user (
@@ -47,3 +49,4 @@ CREATE TABLE profile_updates_per_user (
 );
 
 CREATE INDEX profile_updates_per_user_by_user_stream ON profile_updates_per_user (user_id, stream_id);
+CREATE INDEX profile_updates_per_user_inserted_ts ON profile_updates_per_user (inserted_ts);

tests/storage/test_profile.py

@@ -18,14 +18,19 @@
 # [This file includes modifications made by New Vector Limited]
 #
 #
+import itertools
+from unittest.mock import patch
 
 from twisted.internet.testing import MemoryReactor
 
+from synapse.api.constants import ProfileUpdateAction
 from synapse.server import HomeServer
 from synapse.storage.database import LoggingTransaction
+from synapse.storage.databases.main.profile import PRUNE_PROFILE_UPDATES_AGE
 from synapse.storage.engines import PostgresEngine
 from synapse.types import UserID
 from synapse.util.clock import Clock
+from synapse.util.duration import Duration
 
 from tests import unittest
 
@@ -132,3 +137,116 @@ class ProfileStoreTestCase(unittest.HomeserverTestCase):
         )
         self.assertEqual(len(res), len(expected_values))
         self.assertEqual(res, expected_values)
+
+    @patch("synapse.storage.databases.main.profile.PRUNE_PROFILE_UPDATES_BATCH_SIZE", 5)
+    def test_prune_profile_updates(self) -> None:
+        """Test that old entries in the `profile_updates` and `profile_updates_per_user`
+        tables are pruned properly."""
+
+        # Create a generator for field names so we can easily create many unique
+        # field names without having to keep track of the count ourselves.
+        field_name_gen = (f"field{i}" for i in itertools.count())
+
+        def get_profile_updates_status() -> tuple[int, str]:
+            """Helper function to get the count of entries in the
+            `profile_updates` table."""
+            return self.get_success(
+                self.store.db_pool.simple_select_one(
+                    table="profile_updates",
+                    keyvalues={},
+                    retcols=("COUNT(*)", "MIN(field_name)"),
+                )
+            )
+
+        def get_profile_updates_per_user_status() -> tuple[int]:
+            """Helper function to get the count of entries in the
+            `profile_updates_per_user` table."""
+            return self.get_success(
+                self.store.db_pool.simple_select_one(
+                    table="profile_updates_per_user",
+                    keyvalues={},
+                    retcols=("COUNT(*)",),
+                )
+            )
+
+        # First add some entries
+        for _ in range(10):
+            stream_id = self.get_success(
+                self.store.add_profile_updates(
+                    user_id=UserID.from_string("@user:test"),
+                    updated_fields=[next(field_name_gen)],
+                    action=ProfileUpdateAction.UPDATE.value,
+                )
+            )
+            self.get_success(
+                self.store.track_profile_updates_per_user(
+                    stream_id=stream_id,
+                    user_ids={"@alice:test", "@bob:test"},
+                )
+            )
+
+        # Advance the reactor a while, but not long enough to trigger pruning.
+        self.reactor.advance(Duration(hours=1).as_secs())
+
+        # The `profile_updates_per_user` table should now have 10 * 2 entries.
+        per_user_count = get_profile_updates_per_user_status()
+        self.assertEqual(per_user_count[0], 20)
+        # The `profile_updates` table should have 10 entries.
+        # and the minimum field name should be `field0`.
+        updates_count, min_field_name = get_profile_updates_status()
+        self.assertEqual(updates_count, 10)
+        self.assertEqual(min_field_name, "field0")
+
+        # Now we add some more entries
+        for _ in range(10):
+            stream_id = self.get_success(
+                self.store.add_profile_updates(
+                    user_id=UserID.from_string("@user:test"),
+                    updated_fields=[next(field_name_gen)],
+                    action=ProfileUpdateAction.UPDATE.value,
+                )
+            )
+            self.get_success(
+                self.store.track_profile_updates_per_user(
+                    stream_id=stream_id,
+                    user_ids={"@alice:test", "@bob:test"},
+                )
+            )
+
+        # Advance the reactor a while more, so that the first batch of entries is
+        # now old enough to be pruned.
+        self.reactor.advance(
+            (PRUNE_PROFILE_UPDATES_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.
+        # The `profile_updates_per_user` table should now have 10 * 2 entries.
+        per_user_count = get_profile_updates_per_user_status()
+        self.assertEqual(per_user_count[0], 20)
+        # The `profile_updates` table should have 10 entries.
+        # and the minimum field name should be `field10`.
+        updates_count, min_field_name = get_profile_updates_status()
+        self.assertEqual(updates_count, 10)
+        self.assertEqual(min_field_name, "field10")
+
+        # We should always keep the most recent entries, even if they are old enough to be pruned.
+        self.reactor.advance(
+            (PRUNE_PROFILE_UPDATES_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())
+
+        # The `profile_updates_per_user` table should now have 2 entries.
+        per_user_count = get_profile_updates_per_user_status()
+        self.assertEqual(per_user_count[0], 2)
+        # The `profile_updates` table should have 1 entry.
+        # and the minimum field name should be `field19`.
+        updates_count, min_field_name = get_profile_updates_status()
+        self.assertEqual(updates_count, 1)
+        self.assertEqual(min_field_name, "field19")