From e0d84ab922ab94504af9d5cb1a5656bdfc081d9d Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 23 Jan 2025 14:51:38 +0000 Subject: [PATCH] Add 'state_groups_persisting' table --- synapse/storage/controllers/persist_events.py | 26 +++--- synapse/storage/databases/state/epochs.py | 93 ++++++++++++------- .../state/delta/89/02_state_groups_epochs.sql | 7 ++ 3 files changed, 78 insertions(+), 48 deletions(-) diff --git a/synapse/storage/controllers/persist_events.py b/synapse/storage/controllers/persist_events.py index 4b85e4df76..387a8eab14 100644 --- a/synapse/storage/controllers/persist_events.py +++ b/synapse/storage/controllers/persist_events.py @@ -638,20 +638,18 @@ class EventsPersistenceStorageController: room_id, [e for e, _ in chunk] ) - # TODO: Add a table to track what state groups we're currently - # inserting? There's a race where this transaction takes so long - # that we delete the state groups we're inserting. - await self._state_epoch_store.mark_state_groups_as_used(events_and_contexts) - - await self.persist_events_store._persist_events_and_state_updates( - room_id, - chunk, - state_delta_for_room=state_delta_for_room, - new_forward_extremities=new_forward_extremities, - use_negative_stream_ordering=backfilled, - inhibit_local_membership_updates=backfilled, - new_event_links=new_event_links, - ) + async with self._state_epoch_store.persisting_state_group_references( + events_and_contexts + ): + await self.persist_events_store._persist_events_and_state_updates( + room_id, + chunk, + state_delta_for_room=state_delta_for_room, + new_forward_extremities=new_forward_extremities, + use_negative_stream_ordering=backfilled, + inhibit_local_membership_updates=backfilled, + new_event_links=new_event_links, + ) return replaced_events diff --git a/synapse/storage/databases/state/epochs.py b/synapse/storage/databases/state/epochs.py index e470b24f06..e47e4c31ae 100644 --- a/synapse/storage/databases/state/epochs.py +++ b/synapse/storage/databases/state/epochs.py @@ -13,7 +13,8 @@ # -from typing import TYPE_CHECKING, Collection, Dict, Optional, Set, Tuple +import contextlib +from typing import TYPE_CHECKING, AsyncIterator, Collection, Dict, Optional, Set, Tuple from synapse.events import EventBase from synapse.events.snapshot import EventContext @@ -55,6 +56,10 @@ class StateEpochDataStore: ): self._clock = hs.get_clock() self.db_pool = database + self._instance_name = hs.get_instance_name() + + # TODO: Clear from `state_groups_persisting` any holdovers from previous + # running instance. if hs.config.worker.run_background_tasks: self._clock.looping_call_now(self._advance_state_epoch, 2 * 60 * 1000) @@ -94,39 +99,6 @@ class StateEpochDataStore: desc="get_state_epoch", ) - async def mark_state_groups_as_used( - self, event_and_contexts: Collection[Tuple[EventBase, EventContext]] - ) -> None: - referenced_state_groups: Set[int] = set() - state_epochs = [] - for event, ctx in event_and_contexts: - if ctx.rejected or event.internal_metadata.is_outlier(): - continue - - assert ctx.state_epoch is not None - assert ctx.state_group is not None - - state_epochs.append(ctx.state_epoch) - - referenced_state_groups.add(ctx.state_group) - - if ctx.state_group_before_event: - referenced_state_groups.add(ctx.state_group_before_event) - - if not referenced_state_groups: - # We don't reference any state groups, so nothing to do - return - - assert state_epochs # If we have state groups we have a state epoch - min_state_epoch = min(state_epochs) - - await self.db_pool.runInteraction( - "mark_state_groups_as_used", - self._mark_state_groups_as_used_txn, - min_state_epoch, - referenced_state_groups, - ) - def _mark_state_groups_as_used_txn( self, txn: LoggingTransaction, state_epoch: int, state_groups: Set[int] ) -> None: @@ -176,6 +148,13 @@ class StateEpochDataStore: values=[(state_group,) for state_group in state_groups], ) + self.db_pool.simple_insert_many_txn( + txn, + table="state_groups_persisting", + keys=("state_group", "instance_name"), + values=[(state_group, self._instance_name) for state_group in state_groups], + ) + async def is_state_group_pending_deletion(self, state_group: int) -> bool: """Check if a state group is marked as pending deletion.""" @@ -243,3 +222,49 @@ class StateEpochDataStore: keys=("state_group", "state_epoch"), values=[(state_group, current_epoch) for state_group in new_groups], ) + + @contextlib.asynccontextmanager + async def persisting_state_group_references( + self, event_and_contexts: Collection[Tuple[EventBase, EventContext]] + ) -> AsyncIterator[None]: + referenced_state_groups: Set[int] = set() + state_epochs = [] + for event, ctx in event_and_contexts: + if ctx.rejected or event.internal_metadata.is_outlier(): + continue + + assert ctx.state_epoch is not None + assert ctx.state_group is not None + + state_epochs.append(ctx.state_epoch) + + referenced_state_groups.add(ctx.state_group) + + if ctx.state_group_before_event: + referenced_state_groups.add(ctx.state_group_before_event) + + if not referenced_state_groups: + # We don't reference any state groups, so nothing to do + yield + return + + assert state_epochs # If we have state groups we have a state epoch + min_state_epoch = min(state_epochs) + + await self.db_pool.runInteraction( + "mark_state_groups_as_used", + self._mark_state_groups_as_used_txn, + min_state_epoch, + referenced_state_groups, + ) + + try: + yield None + finally: + await self.db_pool.simple_delete_many( + table="state_groups_persisting", + column="state_group", + iterable=referenced_state_groups, + keyvalues={"instance_name": self._instance_name}, + desc="persisting_state_group_references_delete", + ) diff --git a/synapse/storage/schema/state/delta/89/02_state_groups_epochs.sql b/synapse/storage/schema/state/delta/89/02_state_groups_epochs.sql index ea2a4c2f8a..be00015dcd 100644 --- a/synapse/storage/schema/state/delta/89/02_state_groups_epochs.sql +++ b/synapse/storage/schema/state/delta/89/02_state_groups_epochs.sql @@ -28,3 +28,10 @@ CREATE TABLE IF NOT EXISTS state_groups_pending_deletion ( ); CREATE INDEX state_groups_pending_deletion_epoch ON state_groups_pending_deletion(state_epoch); + + +CREATE TABLE IF NOT EXISTS state_groups_persisting ( + state_group BIGINT NOT NULL, + instance_name TEXT NOT NULL, + PRIMARY KEY (state_group, instance_name) +);