From 5bd998c197eebabf667aca233c367e19d53a332a Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 10 Jul 2026 07:30:49 +0000 Subject: [PATCH] Exactly-once for workflow transactions recovered before commit A workflow's idempotent transaction call executed (and committed) twice when a failure hit the window between the transaction's response returning to the workflow and its commit becoming durable: a transaction's response is returned right after two phase commit's durable prepare, with the commit finishing asynchronously. On restart the transaction was recovered as prepared and committed during recovery, but the recovered workflow re-executed concurrently and its duplicate call raced that commit and missed deduplication, so the user's transaction method ran a second time on top of the first commit. This made `test_workflow_idempotent_writer_and_transaction_survive_recovery` flaky on slow runners (`AssertionError: 2 != 1`); delaying every participant commit by one second reproduced the failure deterministically. Two gaps combined to cause the missed deduplication: 1. `DatabaseService::RecoverTransactionIdempotentMutations` scanned a recovered transaction's write batch only under the `idempotent-mutation-v4` key prefix, so the workflow-scoped and expiring uncommitted idempotent mutations introduced in 6a7e71429 recovered as an empty `uncommitted_idempotent_mutations`. Now the scan covers every idempotent mutation key prefix. 2. The in-flight duplicate-call guard in `check_for_idempotent_mutation()` (#5361) matched only `transaction.idempotency_key`, which is not persisted and is therefore `None` for a recovered transaction. Now the guard also matches the transaction's recovered uncommitted idempotent mutation keys, so a duplicate call awaits the recovered transaction and then serves its cached response. With (1), a recovered transaction carrying idempotent mutations can commit before any call has populated the per-state idempotent mutations cache, so `transaction_participant_commit()` now updates the cache entry only when present instead of assuming it exists. Added a deterministic regression test that pins the interleaving by delaying participant commits, in the style of the existing coordinator/participant failure tests. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01UPXVU5ytyWJxSL4YChKhm6 --- reboot/aio/state_managers.py | 32 ++++-- reboot/server/database.cc | 59 ++++++----- tests/reboot/tasks_tests.py | 187 +++++++++++++++++++++++++++++++++++ 3 files changed, 246 insertions(+), 32 deletions(-) diff --git a/reboot/aio/state_managers.py b/reboot/aio/state_managers.py index 793de8cd..fe0ec285 100644 --- a/reboot/aio/state_managers.py +++ b/reboot/aio/state_managers.py @@ -4754,7 +4754,17 @@ async def check_for_idempotent_mutation( for transaction in self._lookup_participant_transactions( state_type_name, state_ref ).values(): - if transaction.idempotency_key == context.idempotency_key: + if ( + transaction.idempotency_key == context.idempotency_key or + # A prepared transaction recovered after a restart + # has `idempotency_key == None`, but its recovered + # uncommitted idempotent mutations are keyed by the + # idempotency keys of the calls that produced them, + # including the transaction's own, so a key match + # means this call is a duplicate of that + # transaction. + context.idempotency_key in transaction.idempotent_mutations + ): await transaction break @@ -6015,13 +6025,19 @@ async def transaction_participant_commit( # users can tradeoff performance for better error # detection of their own code). if len(transaction.idempotent_mutations) > 0: - # Invariant here is that if this transaction - # includes idempotent mutations then we should - # have already recovered them from the - # database (we'll get a `KeyError` here if - # this invariant is broken). - self._idempotent_mutations[state_type][ - state_ref].update( + # A transaction that executed on this server + # has always populated this cache entry via + # `check_for_idempotent_mutation()`, but a + # transaction recovered after a restart + # commits without any prior lookup, so the + # entry may be absent. In that case skipping + # the update is safe: any later lookup + # recovers from the database, which as of the + # commit above includes these mutations. + idempotent_mutations = self._idempotent_mutations[ + state_type].get(state_ref) + if idempotent_mutations is not None: + idempotent_mutations.update( transaction.idempotent_mutations ) diff --git a/reboot/server/database.cc b/reboot/server/database.cc index e6811276..1449cde8 100644 --- a/reboot/server/database.cc +++ b/reboot/server/database.cc @@ -3614,8 +3614,6 @@ grpc::Status DatabaseService::RecoverTransactionIdempotentMutations( std::unique_ptr iterator( CHECK_NOTNULL(batch->NewIterator(*column_family_handle))); - iterator->Seek(rocksdb::Slice(IDEMPOTENT_MUTATION_KEY_PREFIX)); - // When iterating via the transaction's write batch directly we get // `rocksdb::WriteEntry`s which may exist for the same key multiple // times. In the case of an idempotent mutation, we never expect it @@ -3623,36 +3621,49 @@ grpc::Status DatabaseService::RecoverTransactionIdempotentMutations( // each idempotency key we find in a set. std::set idempotency_keys; - while (iterator->Valid()) { - rocksdb::WriteEntry entry = iterator->Entry(); + // Idempotent mutations are stored under a scope-specific key + // prefix (see `MakeIdempotentMutationKey`), so we scan the write + // batch once per prefix. + for (const std::string_view prefix : { + std::string_view(IDEMPOTENT_MUTATION_KEY_PREFIX), + std::string_view(EXPIRING_IDEMPOTENT_MUTATION_KEY_PREFIX), + std::string_view(WORKFLOW_IDEMPOTENT_MUTATION_KEY_PREFIX), + std::string_view(WORKFLOW_EXPIRING_IDEMPOTENT_MUTATION_KEY_PREFIX), + std::string_view(WORKFLOW_ITERATION_IDEMPOTENT_MUTATION_KEY_PREFIX), + }) { + iterator->Seek(rocksdb::Slice(prefix.data(), prefix.size())); + + while (iterator->Valid()) { + rocksdb::WriteEntry entry = iterator->Entry(); + + // Make sure this is still a key we care about. + if (entry.key.ToStringView().find(prefix) != 0) { + break; + } - // Make sure this is still a key we care about. - if (entry.key.ToStringView().find(IDEMPOTENT_MUTATION_KEY_PREFIX) != 0) { - break; - } + // We are only expecting idempotent mutations to be inserted into + // the database during a transaction. + CHECK_EQ(entry.type, rocksdb::kPutRecord); - // We are only expecting idempotent mutations to be inserted into - // the database during a transaction. - CHECK_EQ(entry.type, rocksdb::kPutRecord); - - IdempotentMutation idempotent_mutation; + IdempotentMutation idempotent_mutation; - CHECK(idempotent_mutation.ParseFromArray( - entry.value.data(), - entry.value.size())); + CHECK(idempotent_mutation.ParseFromArray( + entry.value.data(), + entry.value.size())); - // Idempotent mutation should be for this transaction's state ref! - CHECK_EQ(idempotent_mutation.state_type(), transaction.state_type()); - CHECK_EQ(idempotent_mutation.state_ref(), transaction.state_ref()); + // Idempotent mutation should be for this transaction's state ref! + CHECK_EQ(idempotent_mutation.state_type(), transaction.state_type()); + CHECK_EQ(idempotent_mutation.state_ref(), transaction.state_ref()); - CHECK(idempotency_keys.count(idempotent_mutation.key()) == 0); + CHECK(idempotency_keys.count(idempotent_mutation.key()) == 0); - idempotency_keys.insert(idempotent_mutation.key()); + idempotency_keys.insert(idempotent_mutation.key()); - *transaction.add_uncommitted_idempotent_mutations() = - std::move(idempotent_mutation); + *transaction.add_uncommitted_idempotent_mutations() = + std::move(idempotent_mutation); - iterator->Next(); + iterator->Next(); + } } return grpc::Status::OK; diff --git a/tests/reboot/tasks_tests.py b/tests/reboot/tasks_tests.py index d3c1f720..3c663779 100644 --- a/tests/reboot/tasks_tests.py +++ b/tests/reboot/tasks_tests.py @@ -17,6 +17,7 @@ TASKS_RESPONSES_CACHE_TARGET_CAPACITY, ) from reboot.aio.internals.tasks_dispatcher import TasksDispatcher +from reboot.aio.state_managers import SidecarStateManager from reboot.aio.tests import Reboot from reboot.aio.types import StateRef from reboot.aio.workflows import until_changes @@ -1331,6 +1332,192 @@ async def workflow( self.assertEqual(writer_call_count, 1) self.assertEqual(transaction_call_count, 1) + async def test_workflow_transaction_uncommitted_at_failure_survives_recovery( + self, + ) -> None: + """Test that a workflow's transaction whose two phase commit + was durably prepared but whose commit had not yet become + durable at the time of a failure is executed exactly once: on + recovery the transaction is recovered as prepared and + committed, and the workflow's re-executed call must reuse its + response rather than execute the transaction again. + + A transaction's response is returned to its caller once all + participants have durably prepared, while the commit becomes + durable asynchronously. To make this interleaving + deterministic we block every participant commit on an event + that we only set once `rbt.down()`, `rbt.up()`, and the + re-executed workflow's duplicate transaction call reaching + its idempotency check have all happened. The application and + its state managers run in the test's own process and event + loop, so `asyncio.Event`s and `mock.patch` reach them + directly. + """ + # Every `transaction_participant_commit()` blocks on this + # event, so the transaction is uncommitted when `rbt.down()` + # happens and the recovered transaction is still pending + # while the duplicate call checks idempotency. + proceed_with_commit = asyncio.Event() + + # Set when the re-executed workflow's duplicate transaction + # call after recovery reaches its idempotency check. + transaction_reattempted = asyncio.Event() + + # Set by the test once `rbt.down()` has happened, so that + # only post-recovery idempotency checks signal + # `transaction_reattempted`. + recovered = False + + transaction_participant_commit = ( + SidecarStateManager.transaction_participant_commit + ) + + async def blocked_transaction_participant_commit( + state_manager: SidecarStateManager, + transaction, + ): + await proceed_with_commit.wait() + return await transaction_participant_commit( + state_manager, transaction + ) + + check_for_idempotent_mutation = ( + SidecarStateManager.check_for_idempotent_mutation + ) + + async def observed_check_for_idempotent_mutation( + state_manager: SidecarStateManager, + context, + ): + if recovered and isinstance(context, TransactionContext): + transaction_reattempted.set() + return await check_for_idempotent_mutation(state_manager, context) + + writer_and_transaction_called = asyncio.Event() + proceed_after_down = asyncio.Event() + + writer_call_count = 0 + transaction_call_count = 0 + + class TestServicer(GeneralServicer): + + def authorizer(self): + return allow() + + async def constructor_writer( + self, + context: WriterContext, + state: General.State, + request: GeneralRequest, + ) -> GeneralResponse: + return GeneralResponse() + + async def writer( + self, + context: WriterContext, + state: General.State, + request: GeneralRequest, + ) -> GeneralResponse: + nonlocal writer_call_count + writer_call_count += 1 + return GeneralResponse() + + async def transaction( + self, + context: TransactionContext, + state: General.State, + request: GeneralRequest, + ) -> GeneralResponse: + nonlocal transaction_call_count + transaction_call_count += 1 + return GeneralResponse() + + @classmethod + async def workflow( + cls, + context: WorkflowContext, + request: GeneralRequest, + ) -> GeneralResponse: + g = General.ref() + + # Make writer and transaction calls which are + # idempotent by default. + await g.writer(context) + await g.transaction(context) + + # Signal that both have been called. + writer_and_transaction_called.set() + + # Wait until the test tells us to proceed (after + # `rbt.down()`/`rbt.up()`). + await proceed_after_down.wait() + + return GeneralResponse() + + with mock.patch( + 'reboot.aio.state_managers.SidecarStateManager' + '.transaction_participant_commit', + blocked_transaction_participant_commit, + ), mock.patch( + 'reboot.aio.state_managers.SidecarStateManager' + '.check_for_idempotent_mutation', + observed_check_for_idempotent_mutation, + ): + revision = await self.rbt.up( + Application(servicers=[TestServicer]), + # Need to disable effect validation because this test + # relies on setting some nonlocal variables. + effect_validation=EffectValidation.DISABLED, + ) + + context = self.rbt.create_external_context(name=self.id()) + + # Construct. + g, _ = await General.constructor_writer(context) + + # Spawn the workflow. + task = await g.spawn().workflow(context) + + # Wait for both mutations to be called. + await writer_and_transaction_called.wait() + + self.assertEqual(writer_call_count, 1) + self.assertEqual(transaction_call_count, 1) + + # Simulate failure. The transaction's commit is blocked + # on `proceed_with_commit` so the transaction is durably + # prepared but uncommitted. + await self.rbt.down() + + # Reset events for the re-run after recovery. + writer_and_transaction_called.clear() + + recovered = True + + await self.rbt.up(revision=revision) + + # Wait for the re-executed workflow's duplicate + # transaction call to reach its idempotency check while + # the recovered transaction is still pending, then let + # the recovered transaction commit. + await transaction_reattempted.wait() + proceed_with_commit.set() + + # Let the workflow proceed to completion. + proceed_after_down.set() + + await task + + # We should have tried to call both mutations again. + await writer_and_transaction_called.wait() + + # But they still should only have been called once even + # after recovery because the workflow's re-executed calls + # must reuse the recovered mutations, including the + # transaction's, which only commits during recovery. + self.assertEqual(writer_call_count, 1) + self.assertEqual(transaction_call_count, 1) + async def test_control_loop_per_iteration_idempotency_survives_recovery( self, ) -> None: