give a lease-extend joiner its own shortfall

Author: bpapillon

conformance/SPEC.md

@@ -270,6 +270,13 @@ Rules:
   Sizing to the shortfall matters: a single check needing more than `remaining + lease_size`
   would otherwise fail its post-extend retry forever regardless of server balance.
   `expires_at = now + lease_duration_ms`.
+- A caller that joins an extend already in flight must be sized too. If its own
+  `additional_amount` exceeds the one the in-flight extend asked for, it waits that flight out
+  and then issues **exactly one** further extend, re-sized against the slot the flight just
+  moved; if the flight's ask already covers it, it issues nothing. A joiner that silently
+  inherits a tranche-sized ask fails its post-extend retry with credits sitting on the server.
+  The follow-up never chains — a company whose balance cannot reach the request would otherwise
+  spin.
 - On response, reconcile via the store's `extend` with the server's **total** and new expiry,
   **pinned** to the extended lease's id.
 - Failures resolve to "no lease" without throwing (often fire-and-forget).
@@ -431,7 +438,11 @@ to demonstrate; ports must uphold them and should test them natively.
    stores balances as strings to avoid integer truncation.
 4. **Single-flight.** Per-process, per-slot single-flight for acquire and for extend, tracked
    separately. Best-effort only: duplicate wire calls are safe (idempotent server + keep-first
-   `replace` + reconcile-to-total `extend`).
+   `replace` + reconcile-to-total `extend`). An extend flight carries the `additional_amount` it
+   asked for: a joiner whose required shortfall exceeds that figure waits the flight out and
+   then issues exactly one further extend for the remaining shortfall, while a joiner the flight
+   already covers — every watermark-driven one, the common case — issues nothing and shares the
+   single wire call.
 5. **Concurrent cross-pod extends converge.** Two pods extending from the same stale read must
    not double-count — guaranteed by reconcile-to-total computed inside the store (the sequential
    out-of-order-totals vector pins the arithmetic; the concurrent schedule needs a race).

src/schematic/leases/lease_manager.py

@@ -43,6 +43,19 @@ class LeaseGrant:
     expires_at: float
 
 
+@dataclass
+class _Flight:
+    """An in-flight wire call under single-flight.
+
+    ``requested_additional`` is the additional amount an extend's wire call
+    asked for, the figure a joiner compares its own shortfall against. Acquire
+    flights share the type and leave it unset.
+    """
+
+    task: "asyncio.Future[Optional[LeaseState]]"
+    requested_additional: Optional[float] = None
+
+
 class LeaseWireClient(Protocol):
     """The three lease calls the manager makes.
 
@@ -151,8 +164,8 @@ def __init__(
         self._clock = clock
         # Kept separate so an in-flight extend can never satisfy an acquire,
         # or the other way round.
-        self._inflight_acquire: Dict[str, "asyncio.Future[Optional[LeaseState]]"] = {}
-        self._inflight_extend: Dict[str, "asyncio.Future[Optional[LeaseState]]"] = {}
+        self._inflight_acquire: Dict[str, _Flight] = {}
+        self._inflight_extend: Dict[str, _Flight] = {}
         # Every task shutdown has to wait out, whatever it resolves to: the
         # fire-and-forget work from `_spawn` and the single-flight acquires and
         # extends, which resolve to a LeaseState.
@@ -193,7 +206,7 @@ async def acquire_if_needed(
         key = lease_key(company_id, credit_type_id)
         inflight = self._inflight_acquire.get(key)
         if inflight is not None:
-            return await asyncio.shield(inflight)
+            return await asyncio.shield(inflight.task)
         return await self._single_flight(
             self._inflight_acquire, key, self._acquire(company_id, credit_type_id, timeout)
         )
@@ -254,7 +267,24 @@ async def maybe_extend(
         Triggered by either the low-water-mark ratio (steady-state refresh) or
         a ``required_credits`` hint above the local remaining (a check just
         failed a reserve of that size).
+
+        A caller arriving while an extend is in flight joins it. If its own
+        shortfall is larger than what that extend asked for, it waits the
+        flight out and then issues exactly one follow-up extend for the
+        remaining difference: otherwise it would inherit a tranche-sized ask
+        and fail its post-extend retry with credits still sitting on the
+        server.
         """
+        return await self._maybe_extend(company_id, credit_type_id, required_credits, timeout, True)
+
+    async def _maybe_extend(
+        self,
+        company_id: str,
+        credit_type_id: str,
+        required_credits: Optional[float],
+        timeout: Optional[float],
+        allow_follow_up: bool,
+    ) -> Optional[LeaseState]:
         try:
             entry = await self._lease_store.get(company_id, credit_type_id)
         except Exception as err:
@@ -274,30 +304,54 @@ async def maybe_extend(
         if not below_watermark and not below_required:
             return entry
 
+        # Size the extend to cover the request that triggered it: a single
+        # check needing more than remaining plus one tranche would otherwise
+        # fail its post-extend retry forever, however much balance the server
+        # has. The steady-state path keeps asking for the configured tranche.
+        # Sized here, one level above the wire call, so the flight registered
+        # below and the request body provably carry the same number for a
+        # joiner to compare against.
+        shortfall = (required_credits - entry.local_remaining_credits) if required_credits is not None else 0.0
+        additional_amount = max(resolved.lease_size, shortfall)
+
         key = lease_key(company_id, credit_type_id)
         inflight = self._inflight_extend.get(key)
         if inflight is not None:
-            return await asyncio.shield(inflight)
+            joined = await asyncio.shield(inflight.task)
+            # The flight already asked for at least what we need: every
+            # watermark-driven joiner, and any check the tranche covers. One
+            # wire call serves all of them, which is the point of single-flight.
+            if additional_amount <= (inflight.requested_additional or 0.0) or not allow_follow_up:
+                return joined
+            # The flight we waited out has settled. Its own cleanup usually
+            # runs first, but leaving it registered would have the follow-up
+            # join a finished flight and issue nothing.
+            if self._inflight_extend.get(key) is inflight:
+                del self._inflight_extend[key]
+            # Our shortfall outran the flight's ask. We waited it out rather
+            # than racing a second extend onto the same lease; now top up the
+            # difference with exactly one more, re-reading the slot the flight
+            # just moved. No follow-up on the follow-up: when the server cannot
+            # cover the request, a chain would spin.
+            return await self._maybe_extend(company_id, credit_type_id, required_credits, timeout, False)
         return await self._single_flight(
-            self._inflight_extend, key, self._extend(entry, resolved, required_credits, timeout)
+            self._inflight_extend,
+            key,
+            self._extend(entry, resolved, additional_amount, timeout),
+            additional_amount,
         )
 
     async def _extend(
         self,
         entry: LeaseState,
         resolved: ResolvedLeaseConfig,
-        required_credits: Optional[float],
+        additional_amount: float,
         timeout: Optional[float] = None,
     ) -> Optional[LeaseState]:
-        # Size the extend to cover the request that triggered it: a single
-        # check needing more than remaining plus one tranche would otherwise
-        # fail its post-extend retry forever, however much balance the server
-        # has. The steady-state path keeps asking for the configured tranche.
-        shortfall = (required_credits - entry.local_remaining_credits) if required_credits is not None else 0.0
         try:
             grant = await self._wire.extend(
                 entry.lease_id,
-                max(resolved.lease_size, shortfall),
+                additional_amount,
                 self._clock() + resolved.lease_duration,
                 timeout,
             )
@@ -395,12 +449,14 @@ def stop(self) -> None:
 
     async def _single_flight(
         self,
-        registry: Dict[str, "asyncio.Future[Optional[LeaseState]]"],
+        registry: Dict[str, _Flight],
         key: str,
         coro: Awaitable[Optional[LeaseState]],
+        requested_additional: Optional[float] = None,
     ) -> Optional[LeaseState]:
         task = asyncio.ensure_future(coro)
-        registry[key] = task
+        flight = _Flight(task=task, requested_additional=requested_additional)
+        registry[key] = flight
         # The registry dedupes concurrent callers and the drain set waits the
         # wire call out; they have different lifetimes. Cancelling a caller
         # cancels its `shield`, not the task, and drops the registry entry the
@@ -411,7 +467,10 @@ async def _single_flight(
         try:
             return await asyncio.shield(task)
         finally:
-            if registry.get(key) is task:
+            # Identity-guarded rather than an unconditional delete: a joiner
+            # whose shortfall outran this flight registers a follow-up under
+            # the same key, and this cleanup must not evict it.
+            if registry.get(key) is flight:
                 del registry[key]
 
     async def _release(self, lease_id: str) -> None:

tests/lease_support.py

@@ -8,6 +8,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import datetime as dt
 from typing import Any, Awaitable, Dict, List, Optional, cast
 
@@ -148,6 +149,26 @@ def __init__(self) -> None:
         # Runs while an acquire is in flight, for emulating a sibling pod
         # winning the race.
         self.during_acquire: Optional[Any] = None
+        # The same seam on extend, for holding one open while another caller
+        # joins it.
+        self.during_extend: Optional[Any] = None
+
+    def hold_extend(self) -> "tuple[asyncio.Event, asyncio.Event]":
+        """Hold the next extend wire call open.
+
+        The first event fires once that call has landed, the second releases
+        it, so a test can place a joining caller against a flight it knows is
+        in flight rather than against a sleep.
+        """
+        arrived = asyncio.Event()
+        release = asyncio.Event()
+
+        async def hold() -> None:
+            arrived.set()
+            await release.wait()
+
+        self.during_extend = hold
+        return arrived, release
 
     async def acquire(
         self,
@@ -195,6 +216,10 @@ async def extend(
                 "timeout": timeout,
             }
         )
+        during = self.during_extend
+        if during is not None:
+            self.during_extend = None
+            await during()
         scripted = self.extend_responses.pop(0) if self.extend_responses else None
         lease = _scripted_lease(scripted, "unscripted extend wire call")
         return LeaseGrant(

tests/leases/test_check_and_track.py

@@ -427,6 +427,57 @@ async def test_extends_once_and_retries_when_the_lease_is_short(self, clock: Vir
         assert len(flow.wire.extend_calls) == 1
         assert await flow.remaining() == 600
 
+    async def test_allows_a_check_needing_more_than_the_extend_in_flight_asked_for(
+        self, clock: VirtualClock
+    ) -> None:
+        # A sub-water-mark check fires a background extend for one tranche, and
+        # a check needing 1500 arrives while it is in flight. Inheriting the
+        # tranche would leave that check at 1200 local and denied for
+        # insufficient balance with the credits sitting on the server.
+        flow = make_flow(clock)
+        arrived, release = flow.wire.hold_extend()
+        for granted_total in (2000, 3000, 4000):
+            flow.wire.extend_responses.append(
+                {"lease": {"granted_total": granted_total, "expires_at": clock() + LEASE_DURATION}}
+            )
+
+        # 80 at a rate of 10 draws 800 of the 1000-credit lease, leaving 200:
+        # below the water mark, so this check's background extend goes out and
+        # is held open.
+        first = await check_with_lease(
+            flow.deps, FLAG_KEY, COMPANY, None, CheckOptions(usage=80, event_subtype=EVENT_SUBTYPE), flow.fallback
+        )
+        assert first.allowed is True
+        await arrived.wait()
+        assert flow.wire.extend_calls[0]["additional_amount"] == LEASE_SIZE
+
+        # 150 at a rate of 10 is 1500 against 200 local: the reserve fails and
+        # the check asks for an extend, joining the tranche-sized flight.
+        second_task = asyncio.ensure_future(
+            check_with_lease(
+                flow.deps,
+                FLAG_KEY,
+                COMPANY,
+                None,
+                CheckOptions(usage=150, event_subtype=EVENT_SUBTYPE),
+                flow.fallback,
+            )
+        )
+        for _ in range(10):
+            await asyncio.sleep(0)
+        assert len(flow.wire.extend_calls) == 1
+
+        release.set()
+        second = await second_task
+
+        assert second.allowed is True
+        assert second.reservation is not None
+        assert second.reservation.credits_reserved == 1500
+        # The follow-up, sized against the slot the first flight moved to 1200.
+        assert len(flow.wire.extend_calls) == 2
+        assert flow.wire.extend_calls[1]["additional_amount"] == LEASE_SIZE
+        await flow.manager._drain_background()
+
     async def test_denies_when_the_retry_after_a_failed_extend_is_still_short(self, clock: VirtualClock) -> None:
         flow = make_flow(clock)
         await flow.check()