diff --git a/reboot/aio/state_managers.py b/reboot/aio/state_managers.py index 74fa62e1..e8b9c18c 100644 --- a/reboot/aio/state_managers.py +++ b/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! diff --git a/tests/reboot/state_manager_tests.py b/tests/reboot/state_manager_tests.py index b8db8840..0db920eb 100644 --- a/tests/reboot/state_manager_tests.py +++ b/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