state_managers: shield a transaction's outcome from its waiters

reboot/aio/state_managers.py

@@ -1283,7 +1283,15 @@ def __await__(self):
             """Awaits for transaction to finish, i.e., aborted or committed."""
 
             async def closure():
-                await self._committed
+                # Shield so that cancelling a waiter cancels only that
+                # waiter. `_committed` is shared with every other
+                # waiter and with the participant paths that resolve
+                # it, so a cancellation reaching it would leave
+                # `finished()` reporting a transaction that `commit()`
+                # and `abort()` can no longer resolve. The
+                # coordinator's participants future is shielded for
+                # the same reason.
+                await asyncio.shield(self._committed)
 
             return closure().__await__()
 
@@ -4044,7 +4052,13 @@ async def transactionally(
             # has already started (i.e., calls were made
             # concurrently). We must wait for that first call to
             # acquire the per-state lock before we continue here.
-            await transaction.acquired_lock
+            #
+            # Shield for the same reason as `_committed`:
+            # `acquired_lock` is shared with every other concurrent
+            # call on this state, so a cancellation reaching it would
+            # leave `_transaction_participant_start()` unable to
+            # resolve it.
+            await asyncio.shield(transaction.acquired_lock)
 
             if transaction.finished():
                 # TODO(benh): add a test case for this!

tests/reboot/state_manager_tests.py

@@ -1646,6 +1646,87 @@ async def shared() -> None:
         self.assertFalse(lock.is_locked())
 
 
+class TransactionTest(unittest.IsolatedAsyncioTestCase):
+    """Unit tests for `StateManager.Transaction`, whose outcome every
+    interested party learns by awaiting the transaction itself."""
+
+    def _transaction(self) -> StateManager.Transaction:
+        state_type = StateTypeName('test.v1.Transactional')
+        state_ref = StateRef.from_id(state_type, 'test-1234')
+        root_id = uuid.uuid4()
+        return StateManager.Transaction(
+            transaction_ids=[root_id],
+            coordinator_state_type=state_type,
+            coordinator_state_ref=state_ref,
+            state_type=state_type,
+            state_ref=state_ref,
+            tasks_dispatcher=unittest.mock.MagicMock(spec=TasksDispatcher),
+            mode=Lock.Mode.SHARED,
+            # A first attempt is as old as its root transaction id.
+            age=root_id,
+        )
+
+    async def _waiter(
+        self,
+        transaction: StateManager.Transaction,
+    ) -> asyncio.Task:
+        """Returns a task that is awaiting `transaction`."""
+        task = asyncio.create_task(self._await_transaction(transaction))
+        # Runs the task up to its await on the transaction.
+        await asyncio.sleep(0)
+        return task
+
+    async def _await_transaction(
+        self,
+        transaction: StateManager.Transaction,
+    ) -> None:
+        await transaction
+
+    async def test_cancelled_waiter_leaves_transaction_unfinished(
+        self,
+    ) -> None:
+        """A task cancelled while awaiting a transaction takes the
+        cancellation with it, leaving the transaction free to be
+        aborted (or committed) afterwards.
+
+        A cancellation that reached the transaction's own future
+        instead would make it claim to be finished while the
+        participant's `abort()` raises `InvalidStateError` from inside
+        a section that must not raise, abandoning the participant
+        entry and the state's lock.
+        """
+        transaction = self._transaction()
+        waiter = await self._waiter(transaction)
+
+        waiter.cancel()
+        with self.assertRaises(asyncio.CancelledError):
+            await waiter
+
+        self.assertFalse(transaction.finished())
+
+        transaction.abort()
+        self.assertTrue(transaction.aborted())
+
+    async def test_cancelled_waiter_leaves_other_waiters_waiting(
+        self,
+    ) -> None:
+        """One waiter's cancellation leaves the other waiters to learn
+        the transaction's outcome as usual."""
+        transaction = self._transaction()
+        cancelled_waiter = await self._waiter(transaction)
+        surviving_waiter = await self._waiter(transaction)
+
+        cancelled_waiter.cancel()
+        with self.assertRaises(asyncio.CancelledError):
+            await cancelled_waiter
+
+        transaction.prepare()
+        transaction.commit()
+
+        await surviving_waiter
+        self.assertTrue(transaction.committed())
+
+
 class EffectsRequiresExclusiveTest(unittest.TestCase):
     """Unit tests for `Effects.requires_exclusive()`, which decides
     whether a transaction running in shared mode must upgrade