From c993a6de519ef2c98f96be5d015c843976cc0878 Mon Sep 17 00:00:00 2001 From: Ben Papillon Date: Wed, 16 Sep 2026 17:25:04 -0700 Subject: [PATCH] pin the reservation to the lease try_reserve actually charged --- conformance/SPEC.md | 22 ++- conformance/vectors/lease-lifecycle.json | 50 +++++ src/schematic/leases/__init__.py | 3 +- src/schematic/leases/check.py | 32 ++-- src/schematic/leases/lease_store.py | 30 ++- src/schematic/leases/redis_lease_store.py | 35 ++-- tests/conformance/test_vectors.py | 8 +- tests/lease_support.py | 13 +- tests/leases/test_check_and_track.py | 71 ++++++- tests/leases/test_crash_windows.py | 224 ++++++++++++++++++++-- tests/leases/test_lease_store.py | 14 +- tests/leases/test_redis_lease_store.py | 18 +- 12 files changed, 455 insertions(+), 65 deletions(-) diff --git a/conformance/SPEC.md b/conformance/SPEC.md index 6191f10..ed23a10 100644 --- a/conformance/SPEC.md +++ b/conformance/SPEC.md @@ -70,7 +70,7 @@ Store-level (exercise the lease store and reservation store directly): | `advance_clock` | `ms` | — | | `replace_lease` | `lease_id`, `company_id`, `credit_type_id`, `granted_amount`, `expires_at_ms` | `written` (bool) | | `drop_lease` | `company_id`, `credit_type_id` | — | -| `try_reserve` | `company_id`, `credit_type_id`, `credits` | `balance` (post-debit number, or `null`) | +| `try_reserve` | `company_id`, `credit_type_id`, `credits` | `balance` (post-debit number, or `null`), `lease_id`? (the lease actually charged) | | `refund_lease` | `company_id`, `credit_type_id`, `credits`, `pin_lease_id`? | — | | `extend_lease` | `company_id`, `credit_type_id`, `granted_total`, `expires_at_ms`?, `pin_lease_id`? | — | | `get_lease` | `company_id`, `credit_type_id` | `exists`, `lease_id`?, `granted_amount`?, `local_remaining_credits`? | @@ -178,7 +178,14 @@ Returns written/kept so the caller can run the redundant-lease release logic (se never reach the arithmetic — it slips through every comparison and would poison the balance into approving everything). - Otherwise debit and return the **post-debit balance** (so the caller can derive the pre-debit - figure as `returned + credits` without a racy follow-up read). + figure as `returned + credits` without a racy follow-up read) **and the `lease_id` of the lease + that was charged**, read atomically with the debit (in-process: under the same per-slot lock; + Redis: inside the same script). +- The debit is **not keyed by lease id** — it charges whichever lease occupies the slot at that + moment, which need not be the one the caller's acquire returned (the slot's lease can be + replaced in between by expiry + a successor install, by the sweeper, or by a sibling process + sharing the backend). Reporting the charged id is what lets the caller pin its reservation to + the lease it actually drew on; see [check flow](#check-flow) step 9. - Reserving down to exactly 0 is allowed. ### `refund(company, credit, credits, pin_lease_id?)` @@ -319,9 +326,14 @@ Then: Still refused → failure handling, reason `insufficient_lease_balance`. Store error → failure handling, reason `lease_store_error`. 9. **Record the reservation** (TTL = `reservation_ttl_ms` from now) — *after* the debit, *before* - the engine gate. If persisting fails, undo the debit (claim-and-refund; direct refund if - nothing persisted; both pinned to the lease) and go to failure handling - (`lease_store_error`). If even the undo fails, accept the bounded leak. + the engine gate. Pin it to the `lease_id` **`try_reserve` reported**, never to the one the + acquire in step 7 returned: those differ whenever the slot's lease was replaced in the window + between them (which spans the awaited extend in step 8), and a stale pin sends the settle + refund and the sweep refund to a lease that was never charged — `refund`'s pin drops both — + while the Track event bills a released lease. If persisting fails, undo the debit + (claim-and-refund; direct refund if nothing persisted; both pinned to the charged lease) and + go to failure handling (`lease_store_error`). If even the undo fails, accept the bounded + leak. 10. **Engine gate.** Re-run the engine against a company snapshot whose `credit_balances[credit_id]` is substituted with the **pre-reservation** local balance (post-debit balance returned by `try_reserve` + `credit_cost` — exact as of the debit, no diff --git a/conformance/vectors/lease-lifecycle.json b/conformance/vectors/lease-lifecycle.json index 085b5fa..c9f5db1 100644 --- a/conformance/vectors/lease-lifecycle.json +++ b/conformance/vectors/lease-lifecycle.json @@ -221,6 +221,56 @@ } ] }, + { + "name": "try_reserve_reports_the_lease_it_charged", + "description": "try_reserve is not keyed by lease id: it debits whichever lease holds the slot. When the incumbent has expired and a successor was installed, the debit lands on the successor, and try_reserve must report the successor's id \u2014 that is the id the caller pins its reservation to, so the settle and sweep refunds (both pinned) apply and the Track event bills the lease that was actually charged.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "advance_clock", + "ms": 60001 + }, + { + "op": "replace_lease", + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 500, + "expires_at_ms": 180000, + "expect": { "written": true } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 400, "lease_id": "lse_2" } + }, + { + "op": "refund_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "pin_lease_id": "lse_2" + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_2", "local_remaining_credits": 500 } + } + ] + }, { "name": "refund_clamped_at_granted_amount", "description": "A refund can never push the local balance above the granted amount.", diff --git a/src/schematic/leases/__init__.py b/src/schematic/leases/__init__.py index 6dcb083..f7e9845 100644 --- a/src/schematic/leases/__init__.py +++ b/src/schematic/leases/__init__.py @@ -10,7 +10,7 @@ from .check import CheckDataStream, CreditCheckDeps, check_with_lease from .lease_manager import CreditsWireClient, LeaseGrant, LeaseManager, LeaseWireClient -from .lease_store import InMemoryLeaseStore, LeaseStore, lease_key +from .lease_store import InMemoryLeaseStore, LeaseStore, ReserveResult, lease_key from .redis_lease_store import RedisLeaseStore from .redis_reservation_store import RedisReservationStore from .reservation_store import InMemoryReservationStore, ReservationStore @@ -62,6 +62,7 @@ "ReservationConsumeResult", "ReservationRecord", "ReservationStore", + "ReserveResult", "ResolvedLeaseConfig", "build_reservation_track_event", "check_with_lease", diff --git a/src/schematic/leases/check.py b/src/schematic/leases/check.py index 5793cfa..6915ca0 100644 --- a/src/schematic/leases/check.py +++ b/src/schematic/leases/check.py @@ -206,18 +206,22 @@ async def failure(reason: str) -> "CheckResult": return await failure("lease_acquire_failed") # try_reserve is the atomic gate: check and debit in one step, returning - # the post-debit balance so the pre-debit figure needs no second read. + # the post-debit balance so the pre-debit figure needs no second read, plus + # the id of the lease it charged. That id, not lease.lease_id, is what the + # reservation pins to: the debit is not keyed by lease, so the slot's lease + # may have been replaced since the acquire above, over a window that spans + # the awaited extend below. try: - post_reserve_balance = await deps.lease_store.try_reserve(resolved_company.id, credit_id, credit_cost) - if post_reserve_balance is None: + reserve = await deps.lease_store.try_reserve(resolved_company.id, credit_id, credit_cost) + if reserve is None: # Pass the cost as required_credits so a single large request # extends even while the ratio sits above the water mark. await deps.manager.maybe_extend(resolved_company.id, credit_id, credit_cost, options.timeout) - post_reserve_balance = await deps.lease_store.try_reserve(resolved_company.id, credit_id, credit_cost) + reserve = await deps.lease_store.try_reserve(resolved_company.id, credit_id, credit_cost) except Exception as err: log.error(f"Lease check: reserve against {resolved_company.id}/{credit_id} failed: {err}") return await failure("lease_store_error") - if post_reserve_balance is None: + if reserve is None: return await failure("insufficient_lease_balance") # Record the hold after the debit and before the gate. A crash between the @@ -227,7 +231,11 @@ async def failure(reason: str) -> "CheckResult": resolved_config = deps.manager.resolve_config(credit_id) record = ReservationRecord( id=str(uuid.uuid4()), - lease_id=lease.lease_id, + # The lease the debit actually landed on, which may not be the one + # acquire_if_needed handed back. Pinning the acquired id instead would + # send the settle refund, the sweep refund, and the track event's lease + # id to a lease that was never charged. + lease_id=reserve.lease_id, company_id=resolved_company.id, credit_type_id=credit_id, event_subtype=event_subtype, @@ -244,12 +252,14 @@ async def failure(reason: str) -> "CheckResult": log.error(f"Lease check: failed to persist reservation {record.id}: {err}") # Undo the debit rather than strand it until lease expiry. consume # claims whatever slice of the add landed and refunds it; a None says - # nothing landed, so refund the debit directly. Both are pinned to this - # lease. If the undo itself fails, accept the bounded leak: the slice - # comes back at lease expiry, which beats risking a double refund. + # nothing landed, so refund the debit directly. Both are pinned to the + # lease the debit landed on (the record carries that id, so consume + # pins to it too), never to the acquired one. If the undo itself fails, + # accept the bounded leak: the slice comes back at lease expiry, which + # beats risking a double refund. try: if await deps.reservations.consume(record.id, 0) is None: - await deps.lease_store.refund(resolved_company.id, credit_id, credit_cost, lease.lease_id) + await deps.lease_store.refund(resolved_company.id, credit_id, credit_cost, reserve.lease_id) except Exception as undo_err: log.warning( f"Lease check: could not undo the local debit for {record.id} ({undo_err}); " @@ -262,7 +272,7 @@ async def failure(reason: str) -> "CheckResult": # returned plus what it debited, exact as of the debit), and credit_cost # tells the engine what this call costs, so it evaluates the same # arithmetic try_reserve just enforced, plus every non-credit rule. - pre_reservation = post_reserve_balance + credit_cost + pre_reservation = reserve.balance + credit_cost substituted = _substitute_credit_balance(resolved_company, credit_id, pre_reservation) try: result = datastream.evaluate_flag( diff --git a/src/schematic/leases/lease_store.py b/src/schematic/leases/lease_store.py index 0742ff7..587dfc4 100644 --- a/src/schematic/leases/lease_store.py +++ b/src/schematic/leases/lease_store.py @@ -12,6 +12,7 @@ import math import time from contextlib import asynccontextmanager +from dataclasses import dataclass from typing import AsyncIterator, Dict, List, Optional, Tuple from .types import Clock, LeaseState @@ -21,6 +22,15 @@ def lease_key(company_id: str, credit_type_id: str) -> str: return f"{company_id}:{credit_type_id}" +@dataclass(frozen=True) +class ReserveResult: + """What a successful ``try_reserve`` reports: the post-debit balance, and + the lease the credits actually came out of.""" + + balance: float + lease_id: str + + class LeaseStore(abc.ABC): """Backing store for lease slots, shared by the in-memory and Redis backends.""" @@ -52,13 +62,22 @@ async def replace( """ @abc.abstractmethod - async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[float]: - """Atomically check and debit, returning the post-debit balance. + async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[ReserveResult]: + """Atomically check and debit, reporting the post-debit balance and the + lease the debit landed on. ``None`` when there is no lease, it has expired, the balance is short, or ``credits`` is not a finite non-negative number. Returning the balance (rather than a bool) lets the caller derive the pre-debit figure as ``returned + credits`` without a racy follow-up read. + + The debit is NOT keyed by lease id: it charges whichever lease occupies + the slot at that moment, which need not be the one the caller's acquire + handed back, since the slot's lease can be replaced in between by the + sweeper or by a sibling pod on a shared backend. The caller must + therefore pin its reservation to the returned ``lease_id``, never to + the acquired one: the settle refund, the sweep refund, and the track + event's lease id all have to name the lease that was actually charged. """ @abc.abstractmethod @@ -181,7 +200,7 @@ async def replace( ) return True - async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[float]: + async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[ReserveResult]: # NaN passes every comparison below, and a NaN balance would approve # every later reserve, so it never reaches the arithmetic. if not is_finite_non_negative(credits): @@ -196,7 +215,10 @@ async def try_reserve(self, company_id: str, credit_type_id: str, credits: float if entry.local_remaining_credits < credits: return None entry.local_remaining_credits -= credits - return entry.local_remaining_credits + # The lease id is read under the SAME lock as the debit: the caller + # pins its reservation to it, and a read taken after the lock could + # name a lease that replaced this one in between. + return ReserveResult(balance=entry.local_remaining_credits, lease_id=entry.lease_id) async def refund( self, diff --git a/src/schematic/leases/redis_lease_store.py b/src/schematic/leases/redis_lease_store.py index 0accf65..04d77bd 100644 --- a/src/schematic/leases/redis_lease_store.py +++ b/src/schematic/leases/redis_lease_store.py @@ -10,7 +10,7 @@ import time from typing import Any, Dict, List, Optional -from .lease_store import LeaseStore, is_finite_non_negative, lease_key +from .lease_store import LeaseStore, ReserveResult, is_finite_non_negative, lease_key from .types import DEFAULT_LEASE_DURATION, Clock, LeaseState DEFAULT_KEY_PREFIX = "schematic:" @@ -86,17 +86,27 @@ """ ) -# Atomic check-and-decrement on `localRemainingCredits`. Returns the post-debit -# balance as a string (a Lua number reply truncates to integer, which would -# corrupt fractional credit costs); nil if there is no lease, the lease has -# expired, or there is insufficient remaining. The expiry guard compares -# against the Redis server clock, so a reserve against an expired-but-not-yet- -# evicted row during the TTL grace window is rejected. +# Atomic check-and-decrement on `localRemainingCredits`. Returns +# `{post-debit balance, charged leaseId}`, the balance as a string (a Lua +# number reply truncates to integer, which would corrupt fractional credit +# costs); nil if there is no lease, the lease has expired, or there is +# insufficient remaining. Reading the leaseId inside the same script is what +# lets the caller pin its reservation to the lease the debit actually landed +# on: the debit is not keyed by lease id, and a sibling pod can replace the +# slot's lease at any point before it. The expiry guard compares against the +# Redis server clock, so a reserve against an expired-but-not-yet-evicted row +# during the TTL grace window is rejected. +# +# The reply shape is safe for a mixed fleet: every pod EVALs its own copy of +# this script text and decodes its own reply, and the key layout and hash +# fields are untouched, so old and new pods keep sharing one lease hash. TRY_RESERVE_SCRIPT = ( LEASE_NOW_MS + """ local raw = redis.call('HGET', KEYS[1], 'localRemainingCredits') if not raw then return false end +local lease_id = redis.call('HGET', KEYS[1], 'leaseId') +if not lease_id then return false end local expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0') if expiry <= now then return false end local remaining = tonumber(raw) @@ -104,7 +114,7 @@ if remaining < requested then return false end local new_remaining = remaining - requested redis.call('HSET', KEYS[1], 'localRemainingCredits', tostring(new_remaining)) -return tostring(new_remaining) +return { tostring(new_remaining), lease_id } """ ) @@ -257,7 +267,7 @@ async def replace( ) return _to_number(result) == 1 - async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[float]: + async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[ReserveResult]: # Reject non-finite/negative debits before they reach the script: the # string form of NaN parses back to a Lua nan, slips through the `<` # comparison, and would poison the SHARED balance for every pod. @@ -268,9 +278,12 @@ async def try_reserve(self, company_id: str, credit_type_id: str, credits: float [self.hash_key(company_id, credit_type_id)], [format_amount(credits)], ) - if result is None or result is False: + # A nil reply (could not reserve) surfaces as None; success is a + # two-element multi-bulk of the post-debit balance and the charged + # lease id, both as strings. + if not result: return None - return float(to_str(result)) + return ReserveResult(balance=float(to_str(result[0])), lease_id=to_str(result[1])) async def refund( self, diff --git a/tests/conformance/test_vectors.py b/tests/conformance/test_vectors.py index 3c3693e..f438d9c 100644 --- a/tests/conformance/test_vectors.py +++ b/tests/conformance/test_vectors.py @@ -165,9 +165,13 @@ async def _op_drop_lease(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) async def _op_try_reserve(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: - balance = await h.leases.try_reserve(op["company_id"], op["credit_type_id"], op["credits"]) + reserve = await h.leases.try_reserve(op["company_id"], op["credit_type_id"], op["credits"]) if "balance" in expect: - _assert_number(balance, expect["balance"]) + _assert_number(reserve.balance if reserve is not None else None, expect["balance"]) + # Optional: the lease the debit actually landed on, which the caller must + # pin its reservation to. + if "lease_id" in expect: + assert reserve is not None and reserve.lease_id == expect["lease_id"] async def _op_refund_lease(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: diff --git a/tests/lease_support.py b/tests/lease_support.py index cb5a49f..26a19cb 100644 --- a/tests/lease_support.py +++ b/tests/lease_support.py @@ -14,7 +14,7 @@ import fakeredis.aioredis from schematic.leases import LeaseGrant, LeaseState, ReservationRecord -from schematic.leases.lease_store import LeaseStore +from schematic.leases.lease_store import LeaseStore, ReserveResult from schematic.types import ( RulesengineCheckFlagResult, RulesengineCompany, @@ -104,7 +104,9 @@ async def replace( expires_at=expires_at, ) - async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[float]: + async def try_reserve( + self, company_id: str, credit_type_id: str, credits: float + ) -> Optional[ReserveResult]: return await self._target.try_reserve(company_id, credit_type_id, credits) async def refund( @@ -148,6 +150,9 @@ 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. + self.during_extend: Optional[Any] = None async def acquire( self, @@ -195,6 +200,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( diff --git a/tests/leases/test_check_and_track.py b/tests/leases/test_check_and_track.py index e7a7b55..0fe8bbd 100644 --- a/tests/leases/test_check_and_track.py +++ b/tests/leases/test_check_and_track.py @@ -14,7 +14,7 @@ from typing import Any, Dict, List, Optional import pytest -from lease_support import ScriptedDataStream, ScriptedWireClient, VirtualClock +from lease_support import ScriptedDataStream, ScriptedWireClient, VirtualClock, make_fake_redis from schematic.client import CheckOptions, CheckResult, Reservation from schematic.leases import ( @@ -24,7 +24,11 @@ LeaseConfig, LeaseManager, LeaseStore, + RedisLeaseStore, + RedisReservationStore, ReservationRecord, + ReservationStore, + ReserveResult, check_with_lease, consume_reservation_and_build_event, ) @@ -126,7 +130,9 @@ async def __call__(self) -> CheckResult: class UnreachableLeaseStore(InMemoryLeaseStore): """A store that can be read but never debited, as an unreachable Redis is.""" - async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[float]: + async def try_reserve( + self, company_id: str, credit_type_id: str, credits: float + ) -> Optional[ReserveResult]: raise RuntimeError("redis down") @@ -150,7 +156,7 @@ class Flow: engine: FlowEngine wire: ScriptedWireClient leases: LeaseStore - reservations: InMemoryReservationStore + reservations: ReservationStore manager: LeaseManager flag_checks: RecordedFlagChecks = field(default_factory=RecordedFlagChecks) fallback: Fallback = field(default_factory=Fallback) @@ -173,13 +179,16 @@ def make_flow( *, engine: Optional[FlowEngine] = None, lease_store: Optional[LeaseStore] = None, + reservation_store: Optional[ReservationStore] = None, acquire: str = "ok", credit_balances: Optional[Dict[str, float]] = None, **datastream_kwargs: Any, ) -> Flow: engine = engine or FlowEngine() leases = lease_store if lease_store is not None else InMemoryLeaseStore(clock=clock) - reservations = InMemoryReservationStore(leases, clock=clock) + reservations = ( + reservation_store if reservation_store is not None else InMemoryReservationStore(leases, clock=clock) + ) wire = ScriptedWireClient() if acquire == "ok": wire.acquire_responses.append( @@ -512,6 +521,60 @@ async def test_a_dead_store_honors_fail_open(self, clock: VirtualClock) -> None: assert flow.engine.balance(1) == FAIL_OPEN_BALANCE +class TestLeaseReplacedMidCheck: + """The real extend window, on a shared Redis. + + The first reserve comes up short, so the flow awaits the extend, and while + that call is on the wire the slot's lease is replaced (here by the wire + stub's side effect; in production by the sweeper or a sibling pod). The + retried debit charges the successor, so the reservation has to name it. + """ + + async def test_pins_the_reservation_to_the_lease_the_retried_debit_charged( + self, frozen_clock: VirtualClock + ) -> None: + client = make_fake_redis() + leases = RedisLeaseStore(client, clock=frozen_clock) + reservations = RedisReservationStore(client, leases, clock=frozen_clock) + flow = make_flow(frozen_clock, lease_store=leases, reservation_store=reservations) + await flow.check() # draws lse_1 down to 500 + + async def swap_the_slot() -> None: + # Expire lse_1 and install lse_2 over it: replace refuses to + # displace a live lease. + await leases.drop(COMPANY["id"], CREDIT_ID) + await leases.replace( + lease_id="lse_2", + company_id=COMPANY["id"], + credit_type_id=CREDIT_ID, + granted_amount=2000, + expires_at=frozen_clock() + LEASE_DURATION, + ) + + flow.wire.during_extend = swap_the_slot + flow.wire.extend_responses.append( + {"lease": {"granted_total": 2500, "expires_at": frozen_clock() + LEASE_DURATION}} + ) + result = await flow.check(usage=90) # 900 credits, more than lse_1's 500 + + # The extend went out against the acquired lease, and the store drops + # its grant because the slot has moved on... + assert flow.wire.extend_calls[0]["lease_id"] == "lse_1" + assert result.allowed is True + assert result.reservation is not None + # ...but the debit landed on the successor that replaced it in flight. + assert result.reservation.lease_id == "lse_2" + entry = await flow.leases.get(COMPANY["id"], CREDIT_ID) + assert entry is not None and entry.lease_id == "lse_2" + assert await flow.remaining() == 1100 + + # And the settle refund lands, because it is pinned to lse_2. + outcome = await consume_reservation_and_build_event(flow.reservations, result.reservation, 50) + assert outcome.settled_locally is True + assert outcome.track.lease_id == "lse_2" + assert await flow.remaining() == 1500 + + class TestCrashWindow: async def test_the_debit_lands_before_the_record(self, clock: VirtualClock) -> None: """A crash in the gap leaks the debit; it never leaves a record with no diff --git a/tests/leases/test_crash_windows.py b/tests/leases/test_crash_windows.py index 43f9fac..9ba4cfe 100644 --- a/tests/leases/test_crash_windows.py +++ b/tests/leases/test_crash_windows.py @@ -1,25 +1,45 @@ -"""The two bounded-leak windows, on both backends. +"""The windows between the two steps of a transition, on both backends. -A crash between the two steps of a transition must strand locally held credits -(which the server reclaims at lease expiry) rather than enable a double-spend. -The debit and the claim are durable first; the record and the refund are what -may be lost. +A crash in one must strand locally held credits (which the server reclaims at +lease expiry) rather than enable a double-spend. The debit and the claim are +durable first; the record and the refund are what may be lost. + +The last window here is not a crash: the slot's lease can be replaced between +the acquire and the debit, and the reservation has to pin the lease the debit +actually charged, or its refunds are dropped. """ from __future__ import annotations -from typing import Any, Awaitable, Tuple, cast +import logging +from typing import Any, Awaitable, Optional, Tuple, cast import pytest -from lease_support import CrashingRefundLeaseStore, VirtualClock, make_fake_redis, make_reservation +from lease_support import ( + CrashingRefundLeaseStore, + ScriptedDataStream, + ScriptedEngine, + ScriptedWireClient, + VirtualClock, + make_fake_redis, + make_reservation, +) +from schematic.client import CheckOptions, CheckResult from schematic.leases import ( + CreditCheckDeps, InMemoryLeaseStore, InMemoryReservationStore, + LeaseConfig, + LeaseManager, + LeaseState, LeaseStore, RedisLeaseStore, RedisReservationStore, ReservationStore, + ReserveResult, + check_with_lease, + consume_reservation_and_build_event, ) BACKENDS = ("in_memory", "redis") @@ -70,7 +90,7 @@ async def test_debit_without_record_leaks_only_the_hold(backend: str, frozen_clo await _seed(leases, frozen_clock) # The crash: the atomic debit landed, the reservation record never did. - assert await leases.try_reserve("co_1", "ct_1", 100) == 900 + assert await leases.try_reserve("co_1", "ct_1", 100) == ReserveResult(balance=900, lease_id="lse_1") # Exactly the reserved amount is stranded, and it is invisible to the # reservation table, so no sweep can ever refund it. @@ -84,7 +104,7 @@ async def test_debit_without_record_leaks_only_the_hold(backend: str, frozen_clo async def test_debit_leak_is_reclaimed_at_lease_expiry(backend: str, frozen_clock: VirtualClock) -> None: leases, _reservations, _crash = _make_stores(backend, frozen_clock) await _seed(leases, frozen_clock) - assert await leases.try_reserve("co_1", "ct_1", 100) == 900 + assert await leases.try_reserve("co_1", "ct_1", 100) == ReserveResult(balance=900, lease_id="lse_1") frozen_clock.advance_ms(60_001) # The stale balance is never served again, and the successor installs at @@ -105,8 +125,8 @@ async def test_a_retried_check_settles_independently(backend: str, frozen_clock: leases, reservations, _crash = _make_stores(backend, frozen_clock) await _seed(leases, frozen_clock, ttl=3600) # The crashed attempt, then the retry with a fresh reservation. - assert await leases.try_reserve("co_1", "ct_1", 100) == 900 - assert await leases.try_reserve("co_1", "ct_1", 100) == 800 + assert await leases.try_reserve("co_1", "ct_1", 100) == ReserveResult(balance=900, lease_id="lse_1") + assert await leases.try_reserve("co_1", "ct_1", 100) == ReserveResult(balance=800, lease_id="lse_1") await reservations.add(make_reservation(id="res_retry", expires_at=frozen_clock() + 60)) assert await reservations.consume("res_retry", 40) == 40 @@ -123,7 +143,7 @@ async def test_a_retried_check_settles_independently(backend: str, frozen_clock: async def test_crash_before_refund_keeps_the_claim(backend: str, frozen_clock: VirtualClock) -> None: leases, reservations, crash = _make_stores(backend, frozen_clock) await _seed(leases, frozen_clock) - assert await leases.try_reserve("co_1", "ct_1", 100) == 900 + assert await leases.try_reserve("co_1", "ct_1", 100) == ReserveResult(balance=900, lease_id="lse_1") await reservations.add(make_reservation(expires_at=frozen_clock() + 60)) crash.arm() @@ -148,7 +168,7 @@ async def test_crash_before_refund_never_leaks_into_a_successor( ) -> None: leases, reservations, crash = _make_stores(backend, frozen_clock) await _seed(leases, frozen_clock) - assert await leases.try_reserve("co_1", "ct_1", 100) == 900 + assert await leases.try_reserve("co_1", "ct_1", 100) == ReserveResult(balance=900, lease_id="lse_1") await reservations.add(make_reservation(expires_at=frozen_clock() + 60)) crash.arm() @@ -178,7 +198,7 @@ async def test_crash_after_the_claim_is_reconciled_without_a_refund(frozen_clock leases = RedisLeaseStore(client, clock=frozen_clock) reservations = RedisReservationStore(client, leases, clock=frozen_clock) await _seed(leases, frozen_clock) - assert await leases.try_reserve("co_1", "ct_1", 100) == 900 + assert await leases.try_reserve("co_1", "ct_1", 100) == ReserveResult(balance=900, lease_id="lse_1") await reservations.add(make_reservation(expires_at=frozen_clock() - 0.001)) original_evalsha = client.evalsha @@ -205,3 +225,179 @@ async def crash_after_claim(sha: str, numkeys: int, *args: Any) -> Any: assert await reservations.reserved_credits("co_1", "ct_1") == 0 assert await reservations.count() == 0 assert await _balance(leases) == 900 + + +class _SwapLeaseOnReserve(LeaseStore): + """Wrapper that replaces the slot's lease with a successor right before the + first debit reaches the store. + + The deterministic form of the window between ``acquire_if_needed`` and + ``try_reserve``: the awaited extend wire call, or a sibling pod's + ``replace`` on a shared backend. The debit is not keyed by lease id, so it + lands on the successor. + """ + + def __init__(self, target: LeaseStore, successor_id: str, clock: VirtualClock) -> None: + self._target = target + self._successor_id = successor_id + self._clock = clock + self._swapped = False + + async def get(self, company_id: str, credit_type_id: str) -> Optional[LeaseState]: + return await self._target.get(company_id, credit_type_id) + + async def replace( + self, + *, + lease_id: str, + company_id: str, + credit_type_id: str, + granted_amount: float, + expires_at: float, + ) -> bool: + return await self._target.replace( + lease_id=lease_id, + company_id=company_id, + credit_type_id=credit_type_id, + granted_amount=granted_amount, + expires_at=expires_at, + ) + + async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[ReserveResult]: + if not self._swapped: + self._swapped = True + # replace refuses to displace a live lease, so the incumbent goes + # first and the successor installs into an empty slot. + await self._target.drop(company_id, credit_type_id) + await self._target.replace( + lease_id=self._successor_id, + company_id=company_id, + credit_type_id=credit_type_id, + granted_amount=1000, + expires_at=self._clock() + 60, + ) + return await self._target.try_reserve(company_id, credit_type_id, credits) + + async def refund( + self, + company_id: str, + credit_type_id: str, + credits: float, + pin_lease_id: Optional[str] = None, + ) -> None: + await self._target.refund(company_id, credit_type_id, credits, pin_lease_id) + + async def extend( + self, + company_id: str, + credit_type_id: str, + granted_total: float, + new_expires_at: Optional[float] = None, + pin_lease_id: Optional[str] = None, + ) -> None: + await self._target.extend(company_id, credit_type_id, granted_total, new_expires_at, pin_lease_id) + + async def drop(self, company_id: str, credit_type_id: str) -> None: + await self._target.drop(company_id, credit_type_id) + + +_ENTITLEMENT = { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens", +} + + +def _check_deps( + lease_store: LeaseStore, reservations: ReservationStore, clock: VirtualClock +) -> CreditCheckDeps: + """The real check flow, wired to whatever lease store is handed in.""" + + async def _noop(_body: Any) -> None: + return None + + engine = ScriptedEngine( + [ + {"value": True, "reason": "probe", "entitlement": _ENTITLEMENT}, + {"value": True, "reason": "matched", "entitlement": _ENTITLEMENT}, + ], + "inference", + ) + manager = LeaseManager( + # The seeded live lease means acquire_if_needed never hits the wire. + ScriptedWireClient(), + lease_store, + reservation_store=reservations, + config=LeaseConfig(lease_duration=300.0, reservation_ttl=60.0, lease_size=1000.0, low_water_mark=0.25), + clock=clock, + ) + return CreditCheckDeps( + datastream=ScriptedDataStream(engine, "inference", {"id": "co_1", "credit_balances": {"ct_1": 5000}}), + lease_store=lease_store, + reservations=reservations, + manager=manager, + logger=logging.getLogger("lease-swap-test"), + enqueue_flag_check=_noop, + clock=clock, + ) + + +async def _unused_fallback() -> CheckResult: + raise AssertionError("the lease path must not fall back here") + + +async def _check_over_a_swapped_lease( + backend: str, clock: VirtualClock +) -> Tuple[LeaseStore, ReservationStore, CheckResult]: + leases, reservations, _crash = _make_stores(backend, clock) + await _seed(leases, clock) + deps = _check_deps(_SwapLeaseOnReserve(leases, "lse_2", clock), reservations, clock) + result = await check_with_lease( + deps, + "inference", + {"id": "co_1"}, + None, + CheckOptions(usage=10, event_subtype="inference_tokens"), + _unused_fallback, + ) + await deps.manager._drain_background() + return leases, reservations, result + + +@pytest.mark.parametrize("backend", BACKENDS) +async def test_the_reservation_pins_the_lease_the_debit_landed_on( + backend: str, frozen_clock: VirtualClock +) -> None: + leases, reservations, result = await _check_over_a_swapped_lease(backend, frozen_clock) + + assert result.allowed is True + assert result.reservation is not None + # Pinning acquire_if_needed's lse_1 here would name a lease that was never + # charged. + assert result.reservation.lease_id == "lse_2" + entry = await leases.get("co_1", "ct_1") + assert entry is not None and entry.lease_id == "lse_2" + assert await _balance(leases) == 900 + + # Cancelling refunds the successor: refund's pin drops a refund aimed at + # any other lease, so a stale pin would leave the balance at 900. + assert await reservations.consume(result.reservation.id, 0) == 0 + assert await _balance(leases) == 1000 + + +@pytest.mark.parametrize("backend", BACKENDS) +async def test_the_settling_track_event_bills_the_lease_the_debit_landed_on( + backend: str, frozen_clock: VirtualClock +) -> None: + leases, reservations, result = await _check_over_a_swapped_lease(backend, frozen_clock) + assert result.reservation is not None + + outcome = await consume_reservation_and_build_event(reservations, result.reservation, 4) + + assert outcome.settled_locally is True + # A stale pin would bill lse_2's spend against the released lse_1, and the + # server would fall through to the grants. + assert outcome.track.lease_id == "lse_2" + # 1000 less the 100 reserved, plus the 60 unspent refunded to lse_2. + assert await _balance(leases) == 960 diff --git a/tests/leases/test_lease_store.py b/tests/leases/test_lease_store.py index 574a65d..8d0d5a2 100644 --- a/tests/leases/test_lease_store.py +++ b/tests/leases/test_lease_store.py @@ -8,7 +8,7 @@ import pytest from lease_support import VirtualClock -from schematic.leases import InMemoryLeaseStore +from schematic.leases import InMemoryLeaseStore, ReserveResult async def _seed(store: InMemoryLeaseStore, clock: VirtualClock, *, granted: float = 100, ttl: float = 60) -> None: @@ -34,9 +34,11 @@ async def test_replace_installs_at_the_full_grant(store: InMemoryLeaseStore, clo assert entry.local_remaining_credits == 100 -async def test_try_reserve_returns_the_post_debit_balance(store: InMemoryLeaseStore, clock: VirtualClock) -> None: +async def test_try_reserve_returns_the_post_debit_balance_and_charged_lease( + store: InMemoryLeaseStore, clock: VirtualClock +) -> None: await _seed(store, clock) - assert await store.try_reserve("co_1", "ct_1", 30) == 70 + assert await store.try_reserve("co_1", "ct_1", 30) == ReserveResult(balance=70, lease_id="lse_1") entry = await store.get("co_1", "ct_1") assert entry is not None and entry.local_remaining_credits == 70 @@ -65,7 +67,7 @@ async def test_try_reserve_rejects_nan_and_infinity(store: InMemoryLeaseStore, c assert await store.try_reserve("co_1", "ct_1", math.nan) is None assert await store.try_reserve("co_1", "ct_1", -10) is None assert await store.try_reserve("co_1", "ct_1", math.inf) is None - assert await store.try_reserve("co_1", "ct_1", 30) == 70 + assert await store.try_reserve("co_1", "ct_1", 30) == ReserveResult(balance=70, lease_id="lse_1") assert await store.try_reserve("co_1", "ct_1", 80) is None @@ -106,7 +108,9 @@ async def test_concurrent_try_reserves_serialize_per_slot(store: InMemoryLeaseSt store.try_reserve("co_1", "ct_1", 40), store.try_reserve("co_1", "ct_1", 40), ) - assert sorted(r for r in results if r is not None) == [20, 60] + successes = [r for r in results if r is not None] + assert sorted(r.balance for r in successes) == [20, 60] + assert [r.lease_id for r in successes] == ["lse_1", "lse_1"] entry = await store.get("co_1", "ct_1") assert entry is not None and entry.local_remaining_credits == 20 diff --git a/tests/leases/test_redis_lease_store.py b/tests/leases/test_redis_lease_store.py index d184b55..f4c913c 100644 --- a/tests/leases/test_redis_lease_store.py +++ b/tests/leases/test_redis_lease_store.py @@ -13,7 +13,7 @@ import pytest from lease_support import VirtualClock -from schematic.leases import RedisLeaseStore +from schematic.leases import RedisLeaseStore, ReserveResult from schematic.leases.redis_lease_store import LEASE_TTL_GRACE_MS @@ -106,7 +106,13 @@ async def test_try_reserve_gates_the_shared_balance(store: RedisLeaseStore, froz store.try_reserve("co_1", "ct_1", 40), store.try_reserve("co_1", "ct_1", 40), ) - assert sorted(r for r in results if r is not None) == [20, 60] + # Each success reports the post-debit balance and the lease it charged: + # one distinct balance step each, all against the one installed lease. + successes = [r for r in results if r is not None] + assert sorted(successes, key=lambda r: r.balance) == [ + ReserveResult(balance=20, lease_id="lse_1"), + ReserveResult(balance=60, lease_id="lse_1"), + ] entry = await store.get("co_1", "ct_1") assert entry is not None and entry.local_remaining_credits == 20 @@ -134,14 +140,14 @@ async def test_try_reserve_rejects_nan_before_it_reaches_the_script( assert await store.try_reserve("co_1", "ct_1", -10) is None entry = await store.get("co_1", "ct_1") assert entry is not None and entry.local_remaining_credits == 100 - assert await store.try_reserve("co_1", "ct_1", 30) == 70 + assert await store.try_reserve("co_1", "ct_1", 30) == ReserveResult(balance=70, lease_id="lse_1") async def test_fractional_credits_survive_the_round_trip( store: RedisLeaseStore, frozen_clock: VirtualClock ) -> None: await _seed(store, frozen_clock, granted=10) - assert await store.try_reserve("co_1", "ct_1", 2.5) == 7.5 + assert await store.try_reserve("co_1", "ct_1", 2.5) == ReserveResult(balance=7.5, lease_id="lse_1") await store.refund("co_1", "ct_1", 1.25) entry = await store.get("co_1", "ct_1") assert entry is not None and entry.local_remaining_credits == 8.75 @@ -293,6 +299,6 @@ async def test_a_flushed_script_cache_falls_back_to_eval( # A Redis that restarts (or is flushed) loses the cached script and answers # NOSCRIPT; the store re-sends the body rather than failing the reserve. await _seed(store, frozen_clock) - assert await store.try_reserve("co_1", "ct_1", 10) == 90 + assert await store.try_reserve("co_1", "ct_1", 10) == ReserveResult(balance=90, lease_id="lse_1") await redis_client.script_flush() - assert await store.try_reserve("co_1", "ct_1", 10) == 80 + assert await store.try_reserve("co_1", "ct_1", 10) == ReserveResult(balance=80, lease_id="lse_1")