From 63fb78682b05525a25b73c65a19b44167f1ee5c2 Mon Sep 17 00:00:00 2001 From: Ben Papillon Date: Wed, 16 Sep 2026 17:22:55 -0700 Subject: [PATCH] give a lease-extend joiner its own shortfall --- conformance/SPEC.md | 13 ++- src/schematic/leases/lease_manager.py | 89 +++++++++++++--- tests/lease_support.py | 23 ++++- tests/leases/test_check_and_track.py | 51 ++++++++++ tests/leases/test_lease_manager.py | 140 ++++++++++++++++++++++++++ 5 files changed, 298 insertions(+), 18 deletions(-) diff --git a/conformance/SPEC.md b/conformance/SPEC.md index ed23a10..8f7e51d 100644 --- a/conformance/SPEC.md +++ b/conformance/SPEC.md @@ -277,6 +277,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). @@ -443,7 +450,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). diff --git a/src/schematic/leases/lease_manager.py b/src/schematic/leases/lease_manager.py index caeb02e..a563730 100644 --- a/src/schematic/leases/lease_manager.py +++ b/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: diff --git a/tests/lease_support.py b/tests/lease_support.py index 26a19cb..5c3e24e 100644 --- a/tests/lease_support.py +++ b/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 @@ -150,10 +151,28 @@ 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 the extend, for emulating the slot's lease being - # replaced while a check waits on the extend wire call. + # The same seam on the extend: for emulating the slot's lease being + # replaced while a check waits on the extend wire call, or 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, company_id: str, diff --git a/tests/leases/test_check_and_track.py b/tests/leases/test_check_and_track.py index 0fe8bbd..9114081 100644 --- a/tests/leases/test_check_and_track.py +++ b/tests/leases/test_check_and_track.py @@ -436,6 +436,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() diff --git a/tests/leases/test_lease_manager.py b/tests/leases/test_lease_manager.py index a507ad9..1f848c8 100644 --- a/tests/leases/test_lease_manager.py +++ b/tests/leases/test_lease_manager.py @@ -21,7 +21,9 @@ LeaseState, LeaseStore, RedisLeaseStore, + lease_key, ) +from schematic.leases.lease_manager import _Flight CONFIG = LeaseConfig(lease_duration=300, reservation_ttl=60, lease_size=1000, low_water_mark=0.25) @@ -44,6 +46,24 @@ def _make_manager(clock: VirtualClock) -> tuple[LeaseManager, InMemoryLeaseStore return manager, store, wire +async def _settle() -> None: + """Let just-started tasks run on to their next suspension point.""" + for _ in range(10): + await asyncio.sleep(0) + + +async def _drawn_down_lease(store: LeaseStore, clock: VirtualClock) -> None: + """A live 1000-credit lease with 200 left: under the 25% water mark.""" + await store.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=clock() + 300, + ) + await store.try_reserve("co_1", "ct_1", 800) + + async def test_acquire_installs_the_lease(clock: VirtualClock) -> None: manager, store, wire = _make_manager(clock) wire.acquire_responses.append(_lease(clock)) @@ -147,6 +167,126 @@ async def test_extend_is_sized_to_the_shortfall(clock: VirtualClock) -> None: assert entry is not None and entry.local_remaining_credits == 5000 +async def test_a_joiner_whose_shortfall_outran_the_flight_tops_up(clock: VirtualClock) -> None: + # A water-mark extend, asking for one tranche, is in flight when a check + # needing 5000 arrives. Taking the tranche would leave that check's + # post-extend retry failing with the credits sitting on the server, so the + # joiner waits the flight out and tops up the difference. + manager, store, wire = _make_manager(clock) + await _drawn_down_lease(store, clock) + arrived, release = wire.hold_extend() + wire.extend_responses.append({"lease": {"granted_total": 2000, "expires_at": clock() + 600}}) + wire.extend_responses.append({"lease": {"granted_total": 5800, "expires_at": clock() + 600}}) + + watermark = asyncio.ensure_future(manager.maybe_extend("co_1", "ct_1")) + await arrived.wait() + assert wire.extend_calls[0]["additional_amount"] == 1000 + + joiner = asyncio.ensure_future(manager.maybe_extend("co_1", "ct_1", 5000)) + await _settle() + # Still one wire call: the joiner waits the flight out rather than racing a + # second extend onto the same lease. + assert len(wire.extend_calls) == 1 + + release.set() + await watermark + joined = await joiner + + # Exactly one follow-up, sized against the slot the flight just moved: + # 5000 required less the 200 left plus the 1000 granted. + assert len(wire.extend_calls) == 2 + assert wire.extend_calls[1]["additional_amount"] == 3800 + assert joined is not None and joined.local_remaining_credits == 5000 + + +async def test_a_joiner_the_flight_already_covers_shares_the_one_wire_call(clock: VirtualClock) -> None: + # The common case, and the fan-out the follow-up must not introduce: the + # joiner's shortfall of 700 fits inside the tranche the flight asked for. + manager, store, wire = _make_manager(clock) + await _drawn_down_lease(store, clock) + arrived, release = wire.hold_extend() + wire.extend_responses.append({"lease": {"granted_total": 2000, "expires_at": clock() + 600}}) + + watermark = asyncio.ensure_future(manager.maybe_extend("co_1", "ct_1")) + await arrived.wait() + joiner = asyncio.ensure_future(manager.maybe_extend("co_1", "ct_1", 900)) + await _settle() + + release.set() + first = await watermark + joined = await joiner + + assert len(wire.extend_calls) == 1 + assert joined == first + assert joined is not None and joined.local_remaining_credits == 1200 + + +async def test_two_watermark_joiners_share_the_one_wire_call(clock: VirtualClock) -> None: + # Neither carries a required figure, so both ask for the same tranche and + # one wire call serves them, which is the point of single-flight. + manager, store, wire = _make_manager(clock) + await _drawn_down_lease(store, clock) + arrived, release = wire.hold_extend() + wire.extend_responses.append({"lease": {"granted_total": 2000, "expires_at": clock() + 600}}) + + first = asyncio.ensure_future(manager.maybe_extend("co_1", "ct_1")) + await arrived.wait() + joiners = [asyncio.ensure_future(manager.maybe_extend("co_1", "ct_1")) for _ in range(2)] + await _settle() + + release.set() + results = await asyncio.gather(first, *joiners) + + assert len(wire.extend_calls) == 1 + assert [entry.local_remaining_credits for entry in results if entry] == [1200] * 3 + + +async def test_the_follow_up_never_chains(clock: VirtualClock) -> None: + # A company whose balance cannot reach the request would otherwise spin: + # the follow-up resolves short and the caller's retry reports insufficient + # balance, as it should. + manager, store, wire = _make_manager(clock) + await _drawn_down_lease(store, clock) + arrived, release = wire.hold_extend() + wire.extend_responses.append({"lease": {"granted_total": 2000, "expires_at": clock() + 600}}) + # The server grants what it has, still far short of the ask. + wire.extend_responses.append({"lease": {"granted_total": 3000, "expires_at": clock() + 600}}) + + watermark = asyncio.ensure_future(manager.maybe_extend("co_1", "ct_1")) + await arrived.wait() + joiner = asyncio.ensure_future(manager.maybe_extend("co_1", "ct_1", 50_000)) + await _settle() + + release.set() + await watermark + joined = await joiner + + assert len(wire.extend_calls) == 2 + assert joined is not None and joined.local_remaining_credits == 2200 + + +async def test_the_flight_cleanup_leaves_a_follow_up_registered(clock: VirtualClock) -> None: + # A follow-up registers under the key of the flight it waited out, so that + # flight's cleanup has to check identity before dropping the entry. + manager, store, wire = _make_manager(clock) + await _drawn_down_lease(store, clock) + arrived, release = wire.hold_extend() + wire.extend_responses.append({"lease": {"granted_total": 2000, "expires_at": clock() + 600}}) + + extending = asyncio.ensure_future(manager.maybe_extend("co_1", "ct_1")) + await arrived.wait() + key = lease_key("co_1", "ct_1") + landed: "asyncio.Future[Optional[LeaseState]]" = asyncio.get_running_loop().create_future() + landed.set_result(None) + follow_up = _Flight(task=landed, requested_additional=3800) + manager._inflight_extend[key] = follow_up + + release.set() + await extending + + assert manager._inflight_extend.get(key) is follow_up + + async def test_never_extends_an_expired_lease(clock: VirtualClock) -> None: manager, store, wire = _make_manager(clock) await store.replace(