Bugfix: workflow transactions recovered before commit were not exactly-once

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
                             )
 

reboot/server/database.cc

@@ -3614,45 +3614,56 @@ grpc::Status DatabaseService::RecoverTransactionIdempotentMutations(
   std::unique_ptr<rocksdb::WBWIIterator> 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
   // to be added more than once, so to assert that invariant we track
   // each idempotency key we find in a set.
   std::set<std::string> 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;

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: