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: