Presume a deadlock only with a holder that held through a grace period

reboot/aio/state_managers.py

@@ -712,8 +712,21 @@ async def claim_ownership(
             transaction_ids = context.transaction_ids
             assert transaction_ids is not None
 
+            # Presume a deadlock only with an owner that has owned this
+            # state for a whole grace period while we waited; an owner
+            # that changed in between is a sign the state is moving
+            # between nested transactions, and we keep waiting.
+            previous_owner_ids = list(self.owner_ids)
+
             def on_grace() -> None:
-                self._abort_if_presumed_sibling_deadlock(transaction_ids)
+                nonlocal previous_owner_ids
+                current_owner_ids = list(self.owner_ids)
+                if current_owner_ids == previous_owner_ids:
+                    self._abort_if_presumed_sibling_deadlock(
+                        transaction_ids,
+                        current_owner_ids,
+                    )
+                previous_owner_ids = current_owner_ids
 
             await self.wait_ownership(
                 lambda: self.is_claimable_by(transaction_ids),
@@ -742,25 +755,28 @@ def on_grace() -> None:
         def _abort_if_presumed_sibling_deadlock(
             self,
             transaction_ids: list[uuid.UUID],
+            owner_ids: list[uuid.UUID],
         ) -> None:
             """Raises `SystemAborted(NestedTransactionShouldRetry(...))`
-            when the nested transaction owning this state is an older
-            sibling of a nested transaction among `transaction_ids`, so
-            that the younger sibling rolls back, ownership of what it
-            claimed returns to the transaction that started it, and the
-            older sibling proceeds: the younger of two siblings waiting
-            on each other's states is always the one to go. A state
-            owned by a younger sibling, or by a descendant of the
-            caller, keeps the caller waiting.
+            when `owner_ids`, the nested transactions that have owned
+            this state for a whole grace period while a call of
+            `transaction_ids` waited, end in an older sibling of a
+            nested transaction among `transaction_ids`, so that the
+            younger sibling rolls back, ownership of what it claimed
+            returns to the transaction that started it, and the older
+            sibling proceeds: the younger of two siblings waiting on
+            each other's states is always the one to go. A state owned
+            by a younger sibling, or by a descendant of the caller,
+            keeps the caller waiting.
             """
             nested_transaction_id = presumed_deadlocked_nested_transaction(
                 transaction_ids,
-                self.owner_ids,
+                owner_ids,
             )
             if nested_transaction_id is None:
                 return
             level = transaction_ids.index(nested_transaction_id)
-            owner_id = self.owner_ids[level]
+            owner_id = owner_ids[level]
             grace_ms = TRANSACTION_DEADLOCK_GRACE // timedelta(milliseconds=1)
             message = (
                 f"Nested transaction {nested_transaction_id} waited longer "
@@ -1695,9 +1711,12 @@ class Lock:
       promoted to `exclusive`. A caller can upgrade and skip other
       `exclusive` waiters, preserving the upgrading transaction's
       shared-consistent view of state. However, at most one upgrade
-      may be pending per lock; a second `upgrade(...)` raises
-      `SystemAborted(Unavailable())` immediately to avoid the deadlock
-      where two shared holders both want to upgrade.
+      may be pending per lock; a second `upgrade(...)` aborts
+      immediately, since two shared holders that both want to upgrade
+      are each waiting for the other's shared hold to go, which is a
+      deadlock by construction: with `SystemAborted(TransactionShouldRetry)`
+      carrying the transaction's age when the upgrader passed its
+      participant, and `SystemAborted(Unavailable())` otherwise.
 
     - `downgrade()` is the inverse: a caller that holds `exclusive` is
       demoted to `shared`, granting any queued `shared` waiters that
@@ -2072,15 +2091,38 @@ async def _wait(
         if upgrade:
             # If we've already got a upgrader that is waiting we fail
             # fast because otherwise each upgrader would sit on their
-            # shared hold and we'd deadlock.
+            # shared hold and we'd deadlock. A transaction is asked to
+            # retry the way a presumed deadlock asks it, carrying its
+            # age, so that its retry skips the backoff and is not the
+            # youngest again; a caller without a participant is asked
+            # to retry as `Unavailable`.
             if self._upgrader is not None:
+                if transaction is None:
+                    raise SystemAborted(
+                        Unavailable(),
+                        message=(
+                            "Cannot upgrade shared lock to exclusive: "
+                            "another transaction is already upgrading "
+                            "the same state; retry the transaction."
+                        ),
+                    )
+                pending = self._upgrader.transaction
+                message = (
+                    f"Transaction {transaction.root_id} cannot upgrade its "
+                    "shared hold to exclusive: " + (
+                        f"transaction {pending.root_id} (age {pending.age})"
+                        if pending is not None else "another transaction"
+                    ) + " is already upgrading the same state, and each "
+                    "would wait for the other's shared hold to go; "
+                    "aborting so that it proceeds. Retry required."
+                )
+                logger.warning(message)
                 raise SystemAborted(
-                    Unavailable(),
-                    message=(
-                        "Cannot upgrade shared lock to exclusive: "
-                        "another transaction is already upgrading "
-                        "the same state; retry the transaction."
+                    TransactionShouldRetry(
+                        reason=TransactionShouldRetry.PRESUMED_DEADLOCK,
+                        retry_age=str(transaction.age),
                     ),
+                    message=message,
                 )
             self._upgrader = waiter
         else:
@@ -2850,23 +2892,67 @@ def __init__(
     def latest_timestamp_ms(self) -> Optional[int]:
         return self._latest_timestamp_ms
 
+    def _presume_deadlock_on_grace(
+        self,
+        state_type: StateTypeName,
+        state_ref: StateRef,
+        transaction: StateManager.Transaction,
+    ) -> Callable[[], None]:
+        """The `on_grace` hook for `transaction` waiting on the lock of
+        `(state_type, state_ref)`: each time a grace period elapses it
+        presumes a deadlock with, and aborts `transaction` for, an
+        older holder that has held the lock since the previous grace
+        period elapsed, or since the wait began. A holder that arrived
+        in between is not a deadlock but the lock changing hands, and
+        `transaction`, queued behind it, keeps waiting.
+
+        TODO: when every state in a cycle lives on this server, the
+        deadlock can be proven rather than presumed, and at once: record
+        on each participant the state it is waiting on (here and in
+        `claim_ownership`), and when a wait begins walk from the lock's
+        holders through what they wait on, and so on, looking for a
+        path back to `transaction`. A cycle found that way aborts its
+        youngest transaction immediately, with no grace period and no
+        false positive; a wait that leaves this server ends the walk
+        and falls back to the grace check. Most cycles in tests and in
+        small deployments are local, so this would resolve them in
+        microseconds instead of a grace period.
+        """
+        lock = self._locks[state_type][state_ref]
+        previous_holders = set(lock.holders)
+
+        def on_grace() -> None:
+            nonlocal previous_holders
+            current_holders = set(lock.holders)
+            self._abort_if_presumed_deadlock(
+                state_type,
+                state_ref,
+                transaction,
+                current_holders & previous_holders,
+            )
+            previous_holders = current_holders
+
+        return on_grace
+
     def _abort_if_presumed_deadlock(
         self,
         state_type: StateTypeName,
         state_ref: StateRef,
         transaction: StateManager.Transaction,
+        holders: set[StateManager.Transaction],
     ) -> None:
         """Raises `SystemAborted(TransactionShouldRetry(...))`, with
-        reason `PRESUMED_DEADLOCK`, when the lock on `(state_type,
-        state_ref)` is held by a transaction older than `transaction`,
-        so that `transaction` aborts and the older one proceeds: the
-        younger of two transactions waiting on each other's states is
-        always the one to go. A lock held only by younger transactions,
-        or by a plain reader or writer outside any transaction (which
-        never waits on anything and so does not identify itself to the
-        lock), keeps `transaction` waiting.
+        reason `PRESUMED_DEADLOCK`, when `holders`, the transactions
+        that have held the lock on `(state_type, state_ref)` for a
+        whole grace period while `transaction` waited, include one
+        older than `transaction`, so that `transaction` aborts and the
+        older one proceeds: the younger of two transactions waiting on
+        each other's states is always the one to go. A lock held only
+        by younger transactions, or by a plain reader or writer outside
+        any transaction (which never waits on anything and so does not
+        identify itself to the lock), keeps `transaction` waiting.
         """
-        for holder in self._locks[state_type][state_ref].holders:
+        for holder in holders:
             # An upgrader is itself among the holders through the
             # shared hold it upgrades from.
             if holder is transaction:
@@ -4272,22 +4358,21 @@ async def _upgrade_lock(
         `(state_type, state_ref)` to exclusive, aborting `transaction`
         with `TransactionShouldRetry` (presumed deadlock) instead if it
         has waited longer than the grace period on a shared holder that
-        is older.
+        is older, or at once if another transaction's upgrade is
+        already pending, since the two would each wait for the other's
+        shared hold to go.
         """
         assert transaction.mode == Lock.Mode.SHARED
 
-        def on_grace() -> None:
-            self._abort_if_presumed_deadlock(
-                state_type,
-                state_ref,
-                transaction,
-            )
-
         await self._locks[state_type][state_ref].upgrade(
             deadline=LOCK_ACQUIRE_DEADLINE_DEFAULT,
             transaction=transaction,
             grace=TRANSACTION_DEADLOCK_GRACE,
-            on_grace=on_grace,
+            on_grace=self._presume_deadlock_on_grace(
+                state_type,
+                state_ref,
+                transaction,
+            ),
         )
         transaction.mode = Lock.Mode.EXCLUSIVE
 
@@ -5441,12 +5526,11 @@ async def _transaction_participant_start(
         # success / failure we resolve `transaction.acquired_lock`
         # so concurrent callers on the same transaction can
         # proceed (or propagate our failure).
-        def on_grace() -> None:
-            self._abort_if_presumed_deadlock(
-                state_type,
-                state_ref,
-                transaction,
-            )
+        on_grace = self._presume_deadlock_on_grace(
+            state_type,
+            state_ref,
+            transaction,
+        )
 
         try:
             if transaction.mode == Lock.Mode.SHARED:

tests/reboot/state_manager_tests.py

@@ -1112,7 +1112,8 @@ async def test_second_upgrade_fast_fails(self) -> None:
         await asyncio.sleep(0)
         self.assertFalse(first_upgrade_task.done())
 
-        # The second upgrade attempt must fail immediately.
+        # The second upgrade attempt must fail immediately; without a
+        # participant it is asked to retry as `Unavailable`.
         with self.assertRaises(SystemAborted) as aborted:
             await lock.upgrade(deadline=None)
         self.assertEqual(type(aborted.exception.error), Unavailable)
@@ -1123,6 +1124,47 @@ async def test_second_upgrade_fast_fails(self) -> None:
         await asyncio.wait_for(first_upgrade_task, timeout=1.0)
         lock.release_exclusive()
 
+    async def test_second_upgrade_by_a_transaction_retries_with_its_age(
+        self,
+    ) -> None:
+        """A transaction whose upgrade meets another pending upgrade is
+        asked to retry the way a presumed deadlock asks it, carrying its
+        age, since the two would each wait for the other's shared hold.
+        """
+        lock = Lock()
+        # Stand-ins for the participants; the lock only reports them.
+        first = unittest.mock.Mock(spec=StateManager.Transaction)
+        first.root_id = uuid7(timestamp_ms=1000)
+        first.age = first.root_id
+        second = unittest.mock.Mock(spec=StateManager.Transaction)
+        second.root_id = uuid7(timestamp_ms=2000)
+        second.age = second.root_id
+        await lock.acquire_shared(deadline=None, transaction=first)
+        await lock.acquire_shared(deadline=None, transaction=second)
+
+        first_upgrade_task = asyncio.create_task(
+            lock.upgrade(deadline=None, transaction=first)
+        )
+        await asyncio.sleep(0)
+        self.assertFalse(first_upgrade_task.done())
+
+        with self.assertRaises(SystemAborted) as aborted:
+            await lock.upgrade(deadline=None, transaction=second)
+        error = aborted.exception.error
+        assert isinstance(error, TransactionShouldRetry)
+        self.assertEqual(
+            error.reason, TransactionShouldRetry.PRESUMED_DEADLOCK
+        )
+        self.assertEqual(error.retry_age, str(second.age))
+        assert aborted.exception.message is not None
+        self.assertIn(str(first.root_id), aborted.exception.message)
+
+        # The second still holds shared; releasing it lets the first
+        # upgrade complete, and then the first releases exclusive.
+        lock.release_shared(transaction=second)
+        await asyncio.wait_for(first_upgrade_task, timeout=1.0)
+        lock.release_exclusive(transaction=first)
+
     async def test_downgrade_when_sole_holder(self) -> None:
         lock = Lock()
         await lock.acquire_exclusive(deadline=None)