diff --git a/tests/storage/test_state_deletion.py b/tests/storage/test_state_deletion.py
new file mode 100644
index 0000000000..a4d318ae20
--- /dev/null
+++ b/tests/storage/test_state_deletion.py
@@ -0,0 +1,475 @@
+#
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright (C) 2025 New Vector, 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>.
+#
+
+
+import logging
+
+from twisted.test.proto_helpers import MemoryReactor
+
+from synapse.rest import admin
+from synapse.rest.client import login, room
+from synapse.server import HomeServer
+from synapse.util import Clock
+
+from tests.test_utils.event_injection import create_event
+from tests.unittest import HomeserverTestCase
+
+logger = logging.getLogger(__name__)
+
+
+class StateDeletionStoreTestCase(HomeserverTestCase):
+ """Tests for the StateDeletionStore."""
+
+ servlets = [
+ admin.register_servlets,
+ room.register_servlets,
+ login.register_servlets,
+ ]
+
+ def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
+ self.store = hs.get_datastores().main
+ self.state_store = hs.get_datastores().state
+ self.state_deletion_store = hs.get_datastores().state_deletion
+ self.purge_events = hs.get_storage_controllers().purge_events
+
+ # We want to disable the automatic deletion of state groups in the
+ # background, so we can do controlled tests.
+ self.purge_events._delete_state_loop_call.stop()
+
+ self.user_id = self.register_user("test", "password")
+ tok = self.login("test", "password")
+ self.room_id = self.helper.create_room_as(self.user_id, tok=tok)
+
+ def check_if_can_be_deleted(self, state_group: int) -> bool:
+ """Check if the state group is pending deletion."""
+
+ state_group_to_sequence_number = self.get_success(
+ self.state_deletion_store.get_pending_deletions([state_group])
+ )
+
+ can_be_deleted = self.get_success(
+ self.state_deletion_store.db_pool.runInteraction(
+ "test_existing_pending_deletion_is_cleared",
+ self.state_deletion_store.get_state_groups_ready_for_potential_deletion_txn,
+ state_group_to_sequence_number,
+ )
+ )
+
+ return state_group in can_be_deleted
+
+ def test_no_deletion(self) -> None:
+ """Test that calling persisting_state_group_references is fine if
+ nothing is pending deletion"""
+ event, context = self.get_success(
+ create_event(
+ self.hs,
+ room_id=self.room_id,
+ type="m.test",
+ sender=self.user_id,
+ )
+ )
+
+ ctx_mgr = self.state_deletion_store.persisting_state_group_references(
+ [(event, context)]
+ )
+
+ self.get_success(ctx_mgr.__aenter__())
+ self.get_success(ctx_mgr.__aexit__(None, None, None))
+
+ def test_no_deletion_error(self) -> None:
+ """Test that calling persisting_state_group_references is fine if
+ nothing is pending deletion, but an error occurs."""
+
+ event, context = self.get_success(
+ create_event(
+ self.hs,
+ room_id=self.room_id,
+ type="m.test",
+ sender=self.user_id,
+ )
+ )
+
+ ctx_mgr = self.state_deletion_store.persisting_state_group_references(
+ [(event, context)]
+ )
+
+ self.get_success(ctx_mgr.__aenter__())
+ self.get_success(ctx_mgr.__aexit__(Exception, Exception("test"), None))
+
+ def test_existing_pending_deletion_is_cleared(self) -> None:
+ """Test that the pending deletion flag gets cleared when the state group
+ gets persisted."""
+
+ event, context = self.get_success(
+ create_event(
+ self.hs,
+ room_id=self.room_id,
+ type="m.test",
+ state_key="",
+ sender=self.user_id,
+ )
+ )
+ assert context.state_group is not None
+
+ # Mark a state group that we're referencing as pending deletion.
+ self.get_success(
+ self.state_deletion_store.mark_state_groups_as_pending_deletion(
+ [context.state_group]
+ )
+ )
+
+ ctx_mgr = self.state_deletion_store.persisting_state_group_references(
+ [(event, context)]
+ )
+
+ self.get_success(ctx_mgr.__aenter__())
+ self.get_success(ctx_mgr.__aexit__(None, None, None))
+
+ # The pending deletion flag should be cleared
+ pending_deletion = self.get_success(
+ self.state_deletion_store.db_pool.simple_select_one_onecol(
+ table="state_groups_pending_deletion",
+ keyvalues={"state_group": context.state_group},
+ retcol="1",
+ allow_none=True,
+ desc="test_existing_pending_deletion_is_cleared",
+ )
+ )
+ self.assertIsNone(pending_deletion)
+
+ def test_pending_deletion_is_cleared_during_persist(self) -> None:
+ """Test that the pending deletion flag is cleared when a state group
+ gets marked for deletion during persistence"""
+
+ event, context = self.get_success(
+ create_event(
+ self.hs,
+ room_id=self.room_id,
+ type="m.test",
+ state_key="",
+ sender=self.user_id,
+ )
+ )
+ assert context.state_group is not None
+
+ ctx_mgr = self.state_deletion_store.persisting_state_group_references(
+ [(event, context)]
+ )
+ self.get_success(ctx_mgr.__aenter__())
+
+ # Mark the state group that we're referencing as pending deletion,
+ # *after* we have started persisting.
+ self.get_success(
+ self.state_deletion_store.mark_state_groups_as_pending_deletion(
+ [context.state_group]
+ )
+ )
+
+ self.get_success(ctx_mgr.__aexit__(None, None, None))
+
+ # The pending deletion flag should be cleared
+ pending_deletion = self.get_success(
+ self.state_deletion_store.db_pool.simple_select_one_onecol(
+ table="state_groups_pending_deletion",
+ keyvalues={"state_group": context.state_group},
+ retcol="1",
+ allow_none=True,
+ desc="test_existing_pending_deletion_is_cleared",
+ )
+ )
+ self.assertIsNone(pending_deletion)
+
+ def test_deletion_check(self) -> None:
+ """Test that the `get_state_groups_that_can_be_purged_txn` check is
+ correct during different points of the lifecycle of persisting an
+ event."""
+ event, context = self.get_success(
+ create_event(
+ self.hs,
+ room_id=self.room_id,
+ type="m.test",
+ state_key="",
+ sender=self.user_id,
+ )
+ )
+ assert context.state_group is not None
+
+ self.get_success(
+ self.state_deletion_store.mark_state_groups_as_pending_deletion(
+ [context.state_group]
+ )
+ )
+
+ # We shouldn't be able to delete the state group as not enough time as passed
+ can_be_deleted = self.check_if_can_be_deleted(context.state_group)
+ self.assertFalse(can_be_deleted)
+
+ # After enough time we can delete the state group
+ self.reactor.advance(
+ 1 + self.state_deletion_store.DELAY_BEFORE_DELETION_MS / 1000
+ )
+ can_be_deleted = self.check_if_can_be_deleted(context.state_group)
+ self.assertTrue(can_be_deleted)
+
+ ctx_mgr = self.state_deletion_store.persisting_state_group_references(
+ [(event, context)]
+ )
+ self.get_success(ctx_mgr.__aenter__())
+
+ # But once we start persisting we can't delete the state group
+ can_be_deleted = self.check_if_can_be_deleted(context.state_group)
+ self.assertFalse(can_be_deleted)
+
+ self.get_success(ctx_mgr.__aexit__(None, None, None))
+
+ # The pending deletion flag should remain cleared after persistence has
+ # finished.
+ can_be_deleted = self.check_if_can_be_deleted(context.state_group)
+ self.assertFalse(can_be_deleted)
+
+ def test_deletion_error_during_persistence(self) -> None:
+ """Test that state groups remain marked as pending deletion if persisting
+ the event fails."""
+
+ event, context = self.get_success(
+ create_event(
+ self.hs,
+ room_id=self.room_id,
+ type="m.test",
+ state_key="",
+ sender=self.user_id,
+ )
+ )
+ assert context.state_group is not None
+
+ # Mark a state group that we're referencing as pending deletion.
+ self.get_success(
+ self.state_deletion_store.mark_state_groups_as_pending_deletion(
+ [context.state_group]
+ )
+ )
+
+ ctx_mgr = self.state_deletion_store.persisting_state_group_references(
+ [(event, context)]
+ )
+
+ self.get_success(ctx_mgr.__aenter__())
+ self.get_success(ctx_mgr.__aexit__(Exception, Exception("test"), None))
+
+ # We should be able to delete the state group after a certain amount of
+ # time
+ self.reactor.advance(
+ 1 + self.state_deletion_store.DELAY_BEFORE_DELETION_MS / 1000
+ )
+ can_be_deleted = self.check_if_can_be_deleted(context.state_group)
+ self.assertTrue(can_be_deleted)
+
+ def test_race_between_check_and_insert(self) -> None:
+ """Check that we correctly handle the race where we go to delete a
+ state group, check that it is unreferenced, and then it becomes
+ referenced just before we delete it."""
+
+ event, context = self.get_success(
+ create_event(
+ self.hs,
+ room_id=self.room_id,
+ type="m.test",
+ state_key="",
+ sender=self.user_id,
+ )
+ )
+ assert context.state_group is not None
+
+ # Mark a state group that we're referencing as pending deletion.
+ self.get_success(
+ self.state_deletion_store.mark_state_groups_as_pending_deletion(
+ [context.state_group]
+ )
+ )
+
+ # Advance time enough so we can delete the state group
+ self.reactor.advance(
+ 1 + self.state_deletion_store.DELAY_BEFORE_DELETION_MS / 1000
+ )
+
+ # Check that we'd be able to delete this state group.
+ state_group_to_sequence_number = self.get_success(
+ self.state_deletion_store.get_pending_deletions([context.state_group])
+ )
+
+ can_be_deleted = self.get_success(
+ self.state_deletion_store.db_pool.runInteraction(
+ "test_existing_pending_deletion_is_cleared",
+ self.state_deletion_store.get_state_groups_ready_for_potential_deletion_txn,
+ state_group_to_sequence_number,
+ )
+ )
+ self.assertCountEqual(can_be_deleted, [context.state_group])
+
+ # ... in the real world we'd check that the state group isn't referenced here ...
+
+ # Now we persist the event to reference the state group, *after* we
+ # check that the state group wasn't referenced
+ ctx_mgr = self.state_deletion_store.persisting_state_group_references(
+ [(event, context)]
+ )
+
+ self.get_success(ctx_mgr.__aenter__())
+ self.get_success(ctx_mgr.__aexit__(Exception, Exception("test"), None))
+
+ # We simulate a pause (required to hit the race)
+ self.reactor.advance(
+ 1 + self.state_deletion_store.DELAY_BEFORE_DELETION_MS / 1000
+ )
+
+ # We should no longer be able to delete the state group, without having
+ # to recheck if its referenced.
+ can_be_deleted = self.get_success(
+ self.state_deletion_store.db_pool.runInteraction(
+ "test_existing_pending_deletion_is_cleared",
+ self.state_deletion_store.get_state_groups_ready_for_potential_deletion_txn,
+ state_group_to_sequence_number,
+ )
+ )
+ self.assertCountEqual(can_be_deleted, [])
+
+ def test_remove_ancestors_from_can_delete(self) -> None:
+ """Test that if a state group is not ready to be deleted, we also don't
+ delete anything that is referenced by it"""
+
+ event, context = self.get_success(
+ create_event(
+ self.hs,
+ room_id=self.room_id,
+ type="m.test",
+ state_key="",
+ sender=self.user_id,
+ )
+ )
+ assert context.state_group is not None
+
+ # Create a new state group that references the one from the event
+ new_state_group = self.get_success(
+ self.state_store.store_state_group(
+ event.event_id,
+ event.room_id,
+ prev_group=context.state_group,
+ delta_ids={},
+ current_state_ids=None,
+ )
+ )
+
+ # Mark them both as pending deletion
+ self.get_success(
+ self.state_deletion_store.mark_state_groups_as_pending_deletion(
+ [context.state_group, new_state_group]
+ )
+ )
+
+ # Advance time enough so we can delete the state group so they're both
+ # ready for deletion.
+ self.reactor.advance(
+ 1 + self.state_deletion_store.DELAY_BEFORE_DELETION_MS / 1000
+ )
+
+ # We can now delete both state groups
+ self.assertTrue(self.check_if_can_be_deleted(context.state_group))
+ self.assertTrue(self.check_if_can_be_deleted(new_state_group))
+
+ # Use the new_state_group to bump its deletion time
+ self.get_success(
+ self.state_store.store_state_group(
+ event.event_id,
+ event.room_id,
+ prev_group=new_state_group,
+ delta_ids={},
+ current_state_ids=None,
+ )
+ )
+
+ # We should now not be able to delete either of the state groups.
+ state_group_to_sequence_number = self.get_success(
+ self.state_deletion_store.get_pending_deletions(
+ [context.state_group, new_state_group]
+ )
+ )
+
+ # We shouldn't be able to delete the state group as not enough time has passed
+ can_be_deleted = self.get_success(
+ self.state_deletion_store.db_pool.runInteraction(
+ "test_existing_pending_deletion_is_cleared",
+ self.state_deletion_store.get_state_groups_ready_for_potential_deletion_txn,
+ state_group_to_sequence_number,
+ )
+ )
+ self.assertCountEqual(can_be_deleted, [])
+
+ def test_newly_referenced_state_group_gets_removed_from_pending(self) -> None:
+ """Check that if a state group marked for deletion becomes referenced
+ (without being removed from pending deletion table), it gets removed
+ from pending deletion table."""
+
+ event, context = self.get_success(
+ create_event(
+ self.hs,
+ room_id=self.room_id,
+ type="m.test",
+ state_key="",
+ sender=self.user_id,
+ )
+ )
+ assert context.state_group is not None
+
+ # Mark a state group that we're referencing as pending deletion.
+ self.get_success(
+ self.state_deletion_store.mark_state_groups_as_pending_deletion(
+ [context.state_group]
+ )
+ )
+
+ # Advance time enough so we can delete the state group so they're both
+ # ready for deletion.
+ self.reactor.advance(
+ 1 + self.state_deletion_store.DELAY_BEFORE_DELETION_MS / 1000
+ )
+
+ # Manually insert into the table to mimic the state group getting used.
+ self.get_success(
+ self.store.db_pool.simple_insert(
+ table="event_to_state_groups",
+ values={"state_group": context.state_group, "event_id": event.event_id},
+ desc="test_newly_referenced_state_group_gets_removed_from_pending",
+ )
+ )
+
+ # Manually run the background task to delete pending state groups.
+ self.get_success(self.purge_events._delete_state_groups_loop())
+
+ # The pending deletion flag should be cleared...
+ pending_deletion = self.get_success(
+ self.state_deletion_store.db_pool.simple_select_one_onecol(
+ table="state_groups_pending_deletion",
+ keyvalues={"state_group": context.state_group},
+ retcol="1",
+ allow_none=True,
+ desc="test_newly_referenced_state_group_gets_removed_from_pending",
+ )
+ )
+ self.assertIsNone(pending_deletion)
+
+ # .. but the state should not have been deleted.
+ state = self.get_success(
+ self.state_store._get_state_for_groups([context.state_group])
+ )
+ self.assertGreater(len(state[context.state_group]), 0)
|
