diff --git a/README.md b/README.md index 82b9bcc..1596b8a 100644 --- a/README.md +++ b/README.md @@ -681,6 +681,10 @@ await client.identify( Or call `await client.prewarm({"id": "your-company-id"}, ["credit-type-id"])` directly. Both are no-ops in server mode, and neither raises. +Pre-warming resolves the company the way the server does: it looks the keys up +first, whatever they are named, and only when nothing matches does it read a +value carrying Schematic's `comp_` prefix as the company id. + ### When a check cannot gate A check that cannot gate, because the API is unreachable, Redis is down, or the diff --git a/conformance/README.md b/conformance/README.md index 0b0f406..d1f4a22 100644 --- a/conformance/README.md +++ b/conformance/README.md @@ -1,9 +1,9 @@ # Credit lease conformance suite -`SPEC.md` and `vectors/*.json` are copied verbatim from schematic-node -(`conformance/` on `main`), the reference implementation for client-mode credit -leases. Do not edit them here: fix or extend them in schematic-node and copy the -result back, or the SDKs stop pinning the same behavior. +`SPEC.md` and `vectors/*.json` are copied verbatim from schematic-node's +`conformance/`, the reference implementation for client-mode credit leases. Do +not edit them here: fix or extend them in schematic-node and copy the result +back, or the SDKs stop pinning the same behavior. `tests/conformance/test_vectors.py` is this repo's runner. The runner is the only language-specific piece; every SDK reimplements it and must pass the same diff --git a/conformance/SPEC.md b/conformance/SPEC.md index 8f7e51d..aed9498 100644 --- a/conformance/SPEC.md +++ b/conformance/SPEC.md @@ -109,7 +109,7 @@ Credit-metered features are gated client-side without a wire call per check. The 1. **Leases** a tranche of credits from the server per `(company_id, credit_type_id)`. The server pre-debits the company balance by the granted amount; the SDK tracks a local view of how much of the tranche remains un-reserved (`local_remaining_credits`). -2. **Reserves** `usage x consumption_rate` credits from the lease at `check()` time, atomically +2. **Reserves** `ceil(usage) x consumption_rate` credits from the lease at `check()` time, atomically (check-and-debit). A successful, engine-approved check returns a *reservation handle*. 3. **Settles** the reservation at `track()` time with the actual usage: the actually-consumed credits stay debited, the unspent slice is refunded to the lease, and a Track event bills the @@ -147,7 +147,7 @@ Leases and reservations both expire: | `company_id`, `credit_type_id` | Slot key. | | `event_subtype` | Event the settle will bill as. | | `quantity_reserved` | Caller-declared usage (event units). | -| `credits_reserved` | `quantity_reserved x consumption_rate`. | +| `credits_reserved` | `ceil(quantity_reserved) x consumption_rate` (whole event units, so the hold matches what the settle bills). | | `consumption_rate` | Rate at reservation time. | | `expires_at` | Reservation TTL deadline (sweep target). | | `eval_ctx` | Company/user keys used at check time; threaded onto the Track event. | @@ -277,13 +277,20 @@ 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. +- A caller that joins an extend already in flight must be sized too. It joins that flight only + when the flight's `additional_amount` covers its own; if the flight asked for less, it waits + that flight out, re-reads the slot the flight just moved, and sizes itself against the balance + left behind. After **two** such joins it issues **exactly one** extend of its own instead of + joining again. A joiner that silently inherits a tranche-sized ask fails its post-extend retry + with credits sitting on the server, whether that ask came from the flight it first found or + from a smaller follow-up another caller registered while it waited. The follow-up never chains: + a company whose balance cannot reach the request would otherwise spin. +- A joiner's wait is capped at the caller's **per-check timeout** when one is given. The flight + runs on the timeout of whichever call started it, which for a steady-state refresh is the + client default, so a check with a budget of its own must not sit behind it. On expiry the + joiner stops waiting and resolves to "no lease", which sends the check down its + [failure path](#failure-handling) by mode; the flight itself continues for the callers still on + it, and whatever it installs is there for the next check to read. - 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). @@ -325,7 +332,9 @@ Then: - Credit entitlement missing `credit_id`, a positive `consumption_rate`, or a resolvable `event_subtype` (caller's explicit subtype wins over the entitlement's) → fall back. - Probe error → fall back (it is a resolution step, not the gate). -6. `credit_cost = usage x consumption_rate`. +6. `credit_cost = ceil(usage) x consumption_rate`. A fraction of an event is not something the + server bills, so the hold rounds up rather than moving the local ledger by less than the Track + event will. 7. **Acquire** a lease for `(company, credit_id)`. Failure → [failure handling](#failure-handling) with reason `lease_acquire_failed`. 8. **Reserve** `credit_cost` via `try_reserve`. On refusal, opportunistically @@ -372,7 +381,8 @@ configured `on_acquire_failure` mode (default **fail-closed**): `track_with_reservation(reservation, actual_quantity)` settles a reservation: -1. `credits = actual_quantity x reservation.consumption_rate`. +1. `credits = ceil(actual_quantity) x reservation.consumption_rate`, rounded up the same way the + hold is, so the debit moves the lease by exactly what the Track event bills. 2. `consume(reservation.id, credits)`: - **Settled locally** (claim succeeded): the clamped consumed slice stays debited; the unspent slice is refunded to the lease (pinned). @@ -382,11 +392,12 @@ configured `on_acquire_failure` mode (default **fail-closed**): over. This is why `reservation_ttl_ms` should exceed the longest expected gap between `check()` and `track_with_reservation()`. 3. Either way, emit the Track event built from the **caller-held handle** (not the store): - `event = event_subtype`, `quantity = actual_quantity` (the *unclamped* actual — the server is - the source of truth for real consumption; only local bookkeeping clamps to the reserved - amount), `lease_id = reservation.lease_id` (routes the server-side consumption through the - lease's sub-ledger instead of double-debiting the pre-debited grant), plus the reservation's - `eval_ctx` company/user and any caller traits. + `event = event_subtype`, `quantity = ceil(actual_quantity)` (the *unclamped* actual, rounded + up: the server is the source of truth for real consumption and only local bookkeeping clamps + to the reserved amount, but the event's quantity has to be a whole number or the server + rejects it while processing and the usage is never billed), `lease_id = reservation.lease_id` + (routes the server-side consumption through the lease's sub-ledger instead of double-debiting + the pre-debited grant), plus the reservation's `eval_ctx` company/user and any caller traits. 4. The Track carries a deterministic idempotency key derived from the reservation id (`"lease-reservation:" + reservation.id` in Node); the server dedupes by it for 24h, so a recovery emit racing the normal emit, or an accidental double settle, collapses to one billed diff --git a/conformance/vectors/fractional-usage.json b/conformance/vectors/fractional-usage.json new file mode 100644 index 0000000..1048ddf --- /dev/null +++ b/conformance/vectors/fractional-usage.json @@ -0,0 +1,82 @@ +{ + "category": "fractional_usage", + "vectors": [ + { + "name": "fractional_usage_rounds_the_hold_and_the_settle_up", + "description": "Usage below one whole event unit: the hold is ceil(usage) x rate, the engine gates on that same cost, and a settle at the same fractional actual debits ceil(actual) x rate and bills ceil(actual) units, so the local ledger moves by exactly what the Track event bills. The reservation still records the caller-declared quantity.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 0.5, + "event_subtype": "inference_tokens", + "save_reservation_as": "r1", + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + }, + { "value": true, "reason": "ok" } + ], + "expect": { + "allowed": true, + "has_reservation": true, + "reservation": { + "lease_id": "lse_1", + "quantity_reserved": 0.5, + "credits_reserved": 10, + "consumption_rate": 10 + }, + "engine_calls": [{ "credit_balance": 5000 }, { "credit_balance": 1000, "credit_cost": 10 }] + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 990 } + }, + { + "op": "track", + "handle": "r1", + "actual_quantity": 0.5, + "expect": { + "settled_locally": true, + "track": { "event": "inference_tokens", "quantity": 1, "lease_id": "lse_1" } + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 990 } + }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 0 } } + ] + } + ] +} diff --git a/src/schematic/client.py b/src/schematic/client.py index a6f3309..72b4be7 100644 --- a/src/schematic/client.py +++ b/src/schematic/client.py @@ -6,7 +6,7 @@ import time import uuid from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Literal, Optional, Union +from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union import httpx from .base_client import AsyncBaseSchematic, BaseSchematic @@ -28,6 +28,7 @@ from .leases import ( DEFAULT_LEASE_DURATION, DEFAULT_PREWARM_RESOLVE_TIMEOUT, + SHUTDOWN_DRAIN_TIMEOUT, CreditCheckDeps, CreditsWireClient, InMemoryLeaseStore, @@ -115,7 +116,9 @@ class EventUsage: """Usage of one event subtype, for preflighting a flag check.""" event_subtype: str - quantity: int + # Any finite non-negative number. Both the REST body and the local engine + # take an integer, so a fraction rounds up at each of those boundaries. + quantity: float @dataclass @@ -129,7 +132,10 @@ class CheckFlagOptions: # They mirror the API's PreflightRequestBody. # # Quantity applied to any numeric condition met while evaluating the flag. - usage: Optional[int] = None + # Both the REST body and the local engine take an integer, so a fraction + # rounds up rather than letting the check pass on less usage than the + # action is about to record. + usage: Optional[float] = None # Usage of one specific event subtype. Preferred over `usage` when the # subtype is known, since it only moves conditions measuring that subtype. event_usage: Optional[EventUsage] = None @@ -248,6 +254,25 @@ class TrackWithReservationOptions: traits: Optional[Dict[str, Any]] = None +# Prefix Schematic's secure company ids carry, whatever key name they are +# passed under. +COMPANY_ID_PREFIX = "comp_" + + +def _schematic_id(keys: Dict[str, str], prefix: str) -> Optional[str]: + """The Schematic id hiding among a set of entity keys, recognized by its + secure-id prefix. + + The server reads keys this way once a key lookup has come up empty, so + ``{"account_id": "comp_1"}`` resolves and ``{"id": "acme"}`` does not: the + prefix decides, not the key's name. + """ + for value in keys.values(): + if isinstance(value, str) and value.startswith(prefix): + return value + return None + + def _build_preflight(options: Optional[CheckFlagOptions]) -> Optional[PreflightRequestBody]: """Build the preflight body for a flag check, or None when the caller set no preflight field.""" @@ -255,17 +280,19 @@ def _build_preflight(options: Optional[CheckFlagOptions]) -> Optional[PreflightR return None if options.usage is None and options.event_usage is None and options.credit_cost is None: return None + # The wire quantities are integers, so a fraction rounds up here the way it + # does at the engine boundary. return PreflightRequestBody( credit_cost=options.credit_cost, event_usage=( PreflightEventUsageRequestBody( event_subtype=options.event_usage.event_subtype, - quantity=options.event_usage.quantity, + quantity=_preflight_quantity(options.event_usage.quantity), ) if options.event_usage is not None else None ), - usage=options.usage, + usage=None if options.usage is None else _preflight_quantity(options.usage), ) @@ -1625,10 +1652,10 @@ async def prewarm(self, company: Dict[str, str], credit_type_ids: List[str]) -> """Acquire a lease per credit type up front, so a session's first check() does not pay the acquire round trip. - Best effort: failures are logged, never raised. When the company keys - carry no id, this fetches the company over DataStream, waiting up to - ``credit_leases.prewarm_resolve_timeout`` for it to surface, which - covers a company the server has only just ingested. + Best effort: failures are logged, never raised. The company keys are + looked up over DataStream, waiting up to + ``credit_leases.prewarm_resolve_timeout`` for the company to surface, + which covers a company the server has only just ingested. """ if self._lease_manager is None: self.logger.debug( @@ -1667,6 +1694,13 @@ async def _prewarm_one(self, company_id: str, credit_type_id: str) -> None: async def _resolve_company_id_with_wait(self, company: Dict[str, str]) -> Optional[str]: """Resolve company keys to an ID, waiting for the company to surface. + Resolved in the server's order: every supplied key/value pair is an + ordinary entity key and gets looked up first; only when nothing matches + is a value read as the company's own id, by its ``comp_`` prefix rather + than by the name of the key it sits under. An account is free to define + a key called ``id`` holding its own identifier, so the name alone + settles nothing. + identify does not push a company into the DataStream cache, since companies are only streamed on request, so this fetches (cache first, then over the socket) rather than watching an empty cache. The fetch @@ -1674,12 +1708,9 @@ async def _resolve_company_id_with_wait(self, company: Dict[str, str]) -> Option A prewarm_resolve_timeout of 0 keeps the cache lookup and skips the wait, so an already-seen company still warms. """ - company_id = company.get("id") - if company_id: - return company_id datastream = self._datastream_client if datastream is None: - return None + return _schematic_id(company, COMPANY_ID_PREFIX) # An earlier check or prewarm may already have cached this company, and # that answer costs nothing. try: @@ -1689,7 +1720,7 @@ async def _resolve_company_id_with_wait(self, company: Dict[str, str]) -> Option except Exception as e: self.logger.debug(f"prewarm: DataStream company cache lookup failed ({e})") if self._prewarm_resolve_timeout <= 0: - return None + return _schematic_id(company, COMPANY_ID_PREFIX) deadline = time.monotonic() + self._prewarm_resolve_timeout while True: try: @@ -1701,7 +1732,9 @@ async def _resolve_company_id_with_wait(self, company: Dict[str, str]) -> Option # server has yet to ingest a preceding identify. self.logger.debug(f"prewarm: DataStream company fetch failed ({e})") if time.monotonic() >= deadline: - return None + # The keys never resolved, so fall back to a comp_ value the + # way the server does once its own key lookup comes up empty. + return _schematic_id(company, COMPANY_ID_PREFIX) await asyncio.sleep(PREWARM_POLL_INTERVAL) async def _check_fallback( @@ -1931,17 +1964,26 @@ async def track_with_reservation( return quantity = _settled_quantity(actual_quantity) if reservation.mode == "server": - event = _build_reservation_track_event(reservation, quantity, options) + # Nothing local to consume: the server settles the hold by id, so + # this event is the one that records the usage. + event, settled_locally = _build_reservation_track_event(reservation, quantity, options), True else: - event = await self._settle_client_reservation(reservation, actual_quantity, quantity, options) + event, settled_locally = await self._settle_client_reservation( + reservation, actual_quantity, quantity, options, + ) await self._enqueue_event( "track", event, options=TrackOptions(idempotency_key=f"{RESERVATION_TRACK_IDEMPOTENCY_PREFIX}{reservation.id}"), ) # The settled usage counts toward the company's metrics like any other - # track event, so a locally cached company stays consistent with it. - await self._update_company_metrics(reservation.company, reservation.event_subtype, quantity) + # track event, but the cached metric moves only when this call moved + # local state with it: the server drops a duplicate event on the + # idempotency key, so bumping the metric for one would have a caller's + # retry deny its own next numeric-limit check until the stream pushes + # the real figure. + if settled_locally: + await self._update_company_metrics(reservation.company, reservation.event_subtype, quantity) async def _settle_client_reservation( self, @@ -1949,8 +1991,9 @@ async def _settle_client_reservation( actual_quantity: float, quantity: int, options: Optional[TrackWithReservationOptions], - ) -> EventBodyTrack: - """Consume a client-mode hold locally and hand back the event that bills it. + ) -> Tuple[EventBodyTrack, bool]: + """Consume a client-mode hold locally and hand back the event that bills + it, with whether the hold actually moved. The server is the source of truth for real consumption, so a settle that cannot run locally still emits: the event's idempotency key keeps @@ -1959,12 +2002,13 @@ async def _settle_client_reservation( if self._reservations is None: # The handle came from a lease-configured client, so the event # still needs its lease id and dedupe key even though this client - # holds nothing to settle. + # holds nothing to settle. The usage is new all the same, so it + # counts toward the cached metrics. self.logger.warning( "track_with_reservation: client-mode credit leases are not configured here, " "emitting an unsettled track" ) - return _build_reservation_track_event(reservation, quantity, options) + return _build_reservation_track_event(reservation, quantity, options), True try: outcome = await consume_reservation_and_build_event( self._reservations, reservation, actual_quantity, options, @@ -1974,13 +2018,13 @@ async def _settle_client_reservation( f"track_with_reservation: failed to settle reservation {reservation.id} locally ({e}), " "emitting the track anyway" ) - return _build_reservation_track_event(reservation, quantity, options) + return _build_reservation_track_event(reservation, quantity, options), False if not outcome.settled_locally: self.logger.debug( f"track_with_reservation: reservation {reservation.id} was not settled locally (swept at its " "TTL, already settled, or the store is unreachable); the track is keyed for server-side dedupe" ) - return outcome.track + return outcome.track, outcome.settled_locally async def _enqueue_event( self, @@ -2049,19 +2093,24 @@ async def shutdown(self) -> None: # wire, so a lease installed mid-shutdown is one # release_all_local_leases() can see. Both run for a shared # backend too: the work must not outlive the client. + # + # One budget across both waits, not each timeout in turn: a + # caller closing a client wants a bounded shutdown, not the sum + # of every wait inside it. + deadline = time.monotonic() + SHUTDOWN_DRAIN_TIMEOUT pending = list(self._background_tasks) for task in pending: task.cancel() if pending: await asyncio.gather(*pending, return_exceptions=True) - await self._lease_manager.drain() + await self._lease_manager.drain(deadline - time.monotonic()) if not self._lease_backend_shared: # Per-process leases have no sibling drawing on them, so # releasing hands the unspent remainder back to the company # balance now instead of at expiry. A shared lease must # survive this process's shutdown, or the release pulls the # grant out from under the pods still drawing on it. - await self._lease_manager.release_all_local_leases() + await self._lease_manager.release_all_local_leases(deadline - time.monotonic()) if self._datastream_client is not None: try: await self._datastream_client.close() diff --git a/src/schematic/datastream/datastream_client.py b/src/schematic/datastream/datastream_client.py index 2578cb9..e861c5c 100644 --- a/src/schematic/datastream/datastream_client.py +++ b/src/schematic/datastream/datastream_client.py @@ -32,6 +32,50 @@ from ..client import CheckFlagOptions +def _merged_preflight_options( + preflight: Optional[Any], + options: Optional["CheckFlagOptions"], +) -> Optional["CheckFlagOptions"]: + """The options the local engine runs with: the caller's options, with a + preflight the evaluation context carries filling what they leave unset. + + ``usage`` and ``event_usage`` move together: they ask the same question at + different granularities, so taking one from each source would preflight two + different actions. Options naming either own the pair; options naming + neither leave the context's pair alone. A credit cost on the options wins, + since it is the cost this check was priced with, and otherwise the + context's rides along; nothing here can recompute it. + """ + if preflight is None: + return options + # Deferred: the client module imports this package, so the cycle only + # closes at call time. + from ..client import CheckFlagOptions, EventUsage + + merged = CheckFlagOptions( + default_value=options.default_value if options is not None else None, + timeout=options.timeout if options is not None else None, + ) + if options is not None and (options.usage is not None or options.event_usage is not None): + merged.usage = options.usage + merged.event_usage = options.event_usage + else: + merged.usage = preflight.usage + merged.event_usage = ( + EventUsage( + event_subtype=preflight.event_usage.event_subtype, + quantity=preflight.event_usage.quantity, + ) + if preflight.event_usage is not None + else None + ) + if options is not None and options.credit_cost is not None: + merged.credit_cost = options.credit_cost + else: + merged.credit_cost = preflight.credit_cost + return merged + + _hints_cache: Dict[type, Dict[str, Any]] = {} @@ -401,8 +445,11 @@ async def check_flag( """Evaluate a flag for a company and/or user context. ``options`` carries the caller's preflight (hypothetical usage) into - the local evaluation. + the local evaluation, merged with any preflight the evaluation context + itself carries. Without the merge the same call would answer the + hypothetical over REST and the plain question here. """ + options = _merged_preflight_options(eval_ctx.preflight, options) flag = await self.get_flag(flag_key) if flag is None: raise RuntimeError(f"Flag not found: {flag_key}") diff --git a/src/schematic/datastream/rules_engine.py b/src/schematic/datastream/rules_engine.py index 6085f01..1858577 100644 --- a/src/schematic/datastream/rules_engine.py +++ b/src/schematic/datastream/rules_engine.py @@ -2,6 +2,7 @@ import json import logging +import math import re import time from pathlib import Path @@ -53,6 +54,20 @@ def _strip_none(obj: Any) -> Any: return obj +def _engine_quantity(quantity: Optional[float]) -> Optional[int]: + """The quantity as the engine takes it. + + ``usage`` and ``event_usage.quantity`` deserialize as i64 there, so a value + carrying a decimal point fails to deserialize and takes the whole check + with it. Round up, the direction the REST body takes: a preflight asks an + upper-bound question, and the check must not pass on less usage than the + action is about to record. + """ + if quantity is None: + return None + return math.ceil(quantity) + + def _engine_options(options: "CheckFlagOptions") -> Dict[str, Any]: """Build the engine's preflight options block, in the snake_case shape its serde struct expects, with unset fields dropped.""" @@ -61,11 +76,14 @@ def _engine_options(options: "CheckFlagOptions") -> Dict[str, Any]: { "credit_cost": options.credit_cost, "event_usage": ( - {"event_subtype": event_usage.event_subtype, "quantity": event_usage.quantity} + { + "event_subtype": event_usage.event_subtype, + "quantity": _engine_quantity(event_usage.quantity), + } if event_usage is not None else None ), - "usage": options.usage, + "usage": _engine_quantity(options.usage), } ) diff --git a/src/schematic/leases/__init__.py b/src/schematic/leases/__init__.py index f7e9845..e7cdf9b 100644 --- a/src/schematic/leases/__init__.py +++ b/src/schematic/leases/__init__.py @@ -27,6 +27,7 @@ DEFAULT_PREWARM_RESOLVE_TIMEOUT, DEFAULT_RESERVATION_TTL, DEFAULT_SWEEP_INTERVAL, + SHUTDOWN_DRAIN_TIMEOUT, Clock, LeaseConfig, LeaseConfigOverride, @@ -48,6 +49,7 @@ "DEFAULT_PREWARM_RESOLVE_TIMEOUT", "DEFAULT_RESERVATION_TTL", "DEFAULT_SWEEP_INTERVAL", + "SHUTDOWN_DRAIN_TIMEOUT", "InMemoryLeaseStore", "InMemoryReservationStore", "LeaseConfig", diff --git a/src/schematic/leases/check.py b/src/schematic/leases/check.py index 6915ca0..8b2bdeb 100644 --- a/src/schematic/leases/check.py +++ b/src/schematic/leases/check.py @@ -11,6 +11,7 @@ from __future__ import annotations import logging +import math import time import uuid from dataclasses import dataclass @@ -175,7 +176,11 @@ async def check_with_lease( ) return await fallback() - credit_cost = usage * consumption_rate + # Whole event units: a fraction of an event is not something the server + # bills, so the hold rounds up to what the settle will charge. Sizing it on + # the raw quantity would move the local ledger by less than the track + # event, and the two would drift apart over a session. + credit_cost = math.ceil(usage) * consumption_rate async def failure(reason: str) -> "CheckResult": result = await _handle_lease_failure( diff --git a/src/schematic/leases/lease_manager.py b/src/schematic/leases/lease_manager.py index a563730..56f79ef 100644 --- a/src/schematic/leases/lease_manager.py +++ b/src/schematic/leases/lease_manager.py @@ -30,6 +30,14 @@ logger = logging.getLogger(__name__) +# How many in-flight extends one caller will wait out before issuing its own. +# Two covers the case the single-flight was written for: the flight a caller +# joins, and the follow-up another caller registers while it was waiting. +MAX_EXTEND_JOINS = 2 + +# A wait on a shared extend that ran out the joiner's own timeout. +_JOIN_TIMED_OUT = object() + @dataclass class LeaseGrant: @@ -273,9 +281,10 @@ async def maybe_extend( 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. + server. A flight it finds on the way back is only joined if that one + covers the shortfall too; a smaller one is waited out, never inherited. """ - return await self._maybe_extend(company_id, credit_type_id, required_credits, timeout, True) + return await self._maybe_extend(company_id, credit_type_id, required_credits, timeout) async def _maybe_extend( self, @@ -283,8 +292,101 @@ async def _maybe_extend( credit_type_id: str, required_credits: Optional[float], timeout: Optional[float], - allow_follow_up: bool, ) -> Optional[LeaseState]: + # A joiner waits on someone else's wire call, which runs on whatever + # timeout ITS caller set (a background refresh uses the client + # default). So the wait is capped at this caller's own timeout: a check + # with 200ms to spend must not sit behind a 30s extend. + join_deadline = None if timeout is None else time.monotonic() + timeout + # Joins are budgeted, extends of our own are not: a caller may wait out + # flights that ask for too little, but once the budget runs out it + # issues its own single extend rather than joining again. Without the + # budget a caller could wait behind an unbounded run of other callers' + # follow-ups; without the own extend it would return a balance it + # already knows is short and fail its retry with credits on the server. + joins_left = MAX_EXTEND_JOINS + while True: + entry = await self._read_live_lease(company_id, credit_type_id) + if entry is None: + return None + resolved = self.resolve_config(credit_type_id) + if not self._needs_extend(entry, resolved, required_credits): + 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 and joins_left > 0: + joined = await self._join_within(inflight.task, join_deadline) + if joined is _JOIN_TIMED_OUT: + # The flight runs on for everybody else; we just stop + # waiting on it. Reporting no entry sends the caller down + # its fail-open/fail-closed path, which is what its timeout + # asked for. + logger.debug( + "Extend in flight for %s/%s outlasted the caller's timeout; not waiting on it", + company_id, + credit_type_id, + ) + return None + # The flight 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): + return joined + # It asked for less. Go round again to re-read the slot it just + # moved, so what we ask for next is sized against the balance + # it left rather than the one we started from. + joins_left -= 1 + continue + return await self._single_flight( + self._inflight_extend, + key, + self._recheck_and_extend( + company_id, credit_type_id, resolved, required_credits, additional_amount, timeout + ), + additional_amount, + ) + + async def _join_within( + self, + task: "asyncio.Future[Optional[LeaseState]]", + deadline: Optional[float], + ) -> Any: + """Await a flight somebody else is running, giving up at ``deadline``. + + Giving up abandons only our wait: the flight keeps running for the + callers still on it, and whatever it installs is there for our next + check to read. + """ + if deadline is None: + return await asyncio.shield(task) + remaining = deadline - time.monotonic() + if remaining <= 0: + return _JOIN_TIMED_OUT + try: + return await asyncio.wait_for(asyncio.shield(task), remaining) + except asyncio.TimeoutError: + return _JOIN_TIMED_OUT + + async def _read_live_lease(self, company_id: str, credit_type_id: str) -> Optional[LeaseState]: + """The slot's lease, or None when the read fails or the lease is absent + or expired. + + Never extend an expired lease: the server treats it as released and has + already refunded its remainder, so the only correct move is a fresh + acquire on the next check. + """ try: entry = await self._lease_store.get(company_id, credit_type_id) except Exception as err: @@ -292,54 +394,47 @@ async def _maybe_extend( return None if entry is None: return None - # Never extend an expired lease: the server treats it as released and - # has already refunded its remainder, so the only correct move is a - # fresh acquire on the next check. if entry.expires_at <= self._clock(): return None - resolved = self.resolve_config(credit_type_id) + return entry + + def _needs_extend( + self, + entry: LeaseState, + resolved: ResolvedLeaseConfig, + required_credits: Optional[float], + ) -> bool: + """Whether the slot sits low enough to warrant an extend.""" ratio = entry.local_remaining_credits / max(entry.granted_amount, 1) below_watermark = ratio <= resolved.low_water_mark below_required = required_credits is not None and entry.local_remaining_credits < required_credits - 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) + return below_watermark or below_required - key = lease_key(company_id, credit_type_id) - inflight = self._inflight_extend.get(key) - if inflight is not None: - 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, additional_amount, timeout), - additional_amount, - ) + async def _recheck_and_extend( + self, + company_id: str, + credit_type_id: str, + resolved: ResolvedLeaseConfig, + required_credits: Optional[float], + additional_amount: float, + timeout: Optional[float], + ) -> Optional[LeaseState]: + """Re-read the slot now that this flight owns it, and extend only if the + fresh row still warrants one. + + The row that decided this extend was read before the flight was + registered, so an extend that landed in that gap, clearing its own + flight on the way out, would otherwise be followed by a second extend, + under a new idempotency key, for a lease it already topped up. The + registered ``requested_additional`` stands: a joiner compares its + shortfall against that figure, so the wire body has to carry it. + """ + entry = await self._read_live_lease(company_id, credit_type_id) + if entry is None: + return None + if not self._needs_extend(entry, resolved, required_credits): + return entry + return await self._extend(entry, resolved, additional_amount, timeout) async def _extend( self, @@ -385,14 +480,18 @@ async def run() -> None: self._spawn(run()) - async def release_all_local_leases(self) -> None: + async def release_all_local_leases(self, timeout: Optional[float] = None) -> None: """Release every live lease this process exclusively holds. Only a per-process store answers ``list_leases``; a shared backend returns ``None`` and is skipped, since sibling pods still draw on those leases. Expired leases are skipped too: the server already swept them. Best-effort, with failures falling back to server-side expiry. + + Bounded by ``timeout``, so a store or wire call that never lands cannot + hold a closing client open; whatever is abandoned expires server-side. """ + budget = SHUTDOWN_DRAIN_TIMEOUT if timeout is None else timeout try: entries = self._lease_store.list_leases() except Exception as err: @@ -401,19 +500,30 @@ async def release_all_local_leases(self) -> None: if not entries: return now = self._clock() - for entry in entries: - if entry.expires_at <= now: - continue - try: - await self._wire.release(entry.lease_id) - await self._lease_store.drop(entry.company_id, entry.credit_type_id) - logger.debug("Released credit lease %s on close", entry.lease_id) - except Exception as err: - logger.warning( - "Failed to release credit lease %s on close (it will expire server-side): %s", - entry.lease_id, - err, - ) + live = [entry for entry in entries if entry.expires_at > now] + if not live: + return + releases = asyncio.gather(*(self._release_local_lease(entry) for entry in live)) + try: + await asyncio.wait_for(releases, budget) + except asyncio.TimeoutError: + logger.warning( + "Timed out after %ss releasing credit leases on close; " + "any still held will be released by server-side expiry", + budget, + ) + + async def _release_local_lease(self, entry: LeaseState) -> None: + try: + await self._wire.release(entry.lease_id) + await self._lease_store.drop(entry.company_id, entry.credit_type_id) + logger.debug("Released credit lease %s on close", entry.lease_id) + except Exception as err: + logger.warning( + "Failed to release credit lease %s on close (it will expire server-side): %s", + entry.lease_id, + err, + ) def start_sweep(self) -> None: """Run the expired-reservation sweep on an interval. Safe to call twice.""" @@ -501,20 +611,21 @@ async def _drain_background(self) -> None: while self._background: await asyncio.gather(*list(self._background), return_exceptions=True) - async def drain(self) -> None: + async def drain(self, timeout: Optional[float] = None) -> None: """Wait out in-flight lease work, so a close can release what it installed. - Bounded: whatever has not landed by ``SHUTDOWN_DRAIN_TIMEOUT`` is - cancelled rather than stalling the caller's shutdown, and a grant the - server issued for it falls back to server-side expiry. + Bounded: whatever has not landed by ``timeout`` is cancelled rather + than stalling the caller's shutdown, and a grant the server issued for + it falls back to server-side expiry. """ + budget = SHUTDOWN_DRAIN_TIMEOUT if timeout is None else timeout try: - await asyncio.wait_for(self._drain_background(), SHUTDOWN_DRAIN_TIMEOUT) + await asyncio.wait_for(self._drain_background(), budget) except asyncio.TimeoutError: logger.warning( "Timed out after %ss draining in-flight credit lease work; " "any credits it holds will be released by server-side expiry", - SHUTDOWN_DRAIN_TIMEOUT, + budget, ) diff --git a/src/schematic/leases/redis_reservation_store.py b/src/schematic/leases/redis_reservation_store.py index be3c3e8..0ed9473 100644 --- a/src/schematic/leases/redis_reservation_store.py +++ b/src/schematic/leases/redis_reservation_store.py @@ -88,29 +88,62 @@ def _index_key(self) -> str: def _by_credit_key(self, company_id: str, credit_type_id: str) -> str: return f"{self._key_prefix}{RES_BYCREDIT_NAMESPACE}{company_id}:{credit_type_id}" + def _transaction(self) -> Optional[Any]: + """A MULTI/EXEC pipeline, or None when this client cannot open one. + + Probed by calling it, because having the attribute is not the same as + honouring the argument: ``redis.asyncio.cluster.RedisCluster`` carries + ``pipeline`` and raises on ``transaction=True`` (before any I/O), and a + client shim may not take the keyword at all. Both resolve to the + sequential path rather than failing the check that is holding the + credits. + """ + pipeline = getattr(self._client, "pipeline", None) + if pipeline is None: + return None + try: + return pipeline(transaction=True) + except Exception: + return None + async def add(self, reservation: ReservationRecord) -> None: expires_ms = int(to_epoch_ms(reservation.expires_at)) hash_key = self._hash_key(reservation.id) - # The hash goes out first so the reservation exists before anything - # references it. These are independent single-key ops rather than one - # multi-key script: a partial failure at worst leaves an un-indexed - # reservation that the TTL reaps, never a double-spend. - await self._client.hset( - hash_key, - mapping={ - "id": reservation.id, - "leaseId": reservation.lease_id, - "companyId": reservation.company_id, - "creditTypeId": reservation.credit_type_id, - "eventSubtype": reservation.event_subtype, - "quantityReserved": format_amount(reservation.quantity_reserved), - "creditsReserved": format_amount(reservation.credits_reserved), - "consumptionRate": format_amount(reservation.consumption_rate), - "expiresAt": str(expires_ms), - "evalCtx": _encode_eval_ctx(reservation), - }, - ) - await self._client.pexpireat(hash_key, expires_ms + RES_TTL_GRACE_MS) + fields = { + "id": reservation.id, + "leaseId": reservation.lease_id, + "companyId": reservation.company_id, + "creditTypeId": reservation.credit_type_id, + "eventSubtype": reservation.event_subtype, + "quantityReserved": format_amount(reservation.quantity_reserved), + "creditsReserved": format_amount(reservation.credits_reserved), + "consumptionRate": format_amount(reservation.consumption_rate), + "expiresAt": str(expires_ms), + "evalCtx": _encode_eval_ctx(reservation), + } + ttl_at = expires_ms + RES_TTL_GRACE_MS + # The hash and its expiry go out as one MULTI/EXEC. Written separately, + # a crash in the gap leaves a reservation row with no TTL: once the + # sweeper drops its index entry, nothing points at the row and nothing + # reaps it, so it sits in Redis for good. Same commands, same key, same + # fields as before, so what other SDKs read is unchanged, and both + # commands touch the one key, so this is Cluster-safe in principle; + # redis-py's cluster client refuses MULTI all the same, and falls back + # to the sequential path below. + pipe = self._transaction() + if pipe is not None: + pipe.hset(hash_key, mapping=fields) + pipe.pexpireat(hash_key, ttl_at) + await pipe.execute() + else: + await self._client.hset(hash_key, mapping=fields) + await self._client.pexpireat(hash_key, ttl_at) + # The two indexes (expiry zset for the sweeper, per-tenant hash for + # reserved_credits) only depend on the hash existing. They stay outside + # the transaction because their keys hash to other slots. A partial + # failure here at worst leaves an un-indexed reservation that the TTL + # reaps (its slice reclaimed when the lease expires), never a + # double-spend. member = _encode_member(reservation.company_id, reservation.credit_type_id, reservation.id) await self._client.zadd(self._index_key(), {member: expires_ms}) await self._client.hset( @@ -137,10 +170,14 @@ async def consume(self, reservation_id: str, credits_consumed: float) -> Optiona # Index cleanup, single-key ops. The per-tenant hash loses the slice # BEFORE the refund below, so the lease (local remaining plus this - # hash) never transiently double-counts it. + # hash) never transiently double-counts it. The per-tenant field goes + # first and the expiry index second, because the index is what the + # sweeper would reach a surviving field through: dropping the index + # first and then failing on the field would inflate reserved_credits + # for that tenant forever. member = _encode_member(company_id, credit_type_id, reservation_id) - await _ignore_errors(self._client.zrem(self._index_key(), member)) await _ignore_errors(self._client.hdel(self._by_credit_key(company_id, credit_type_id), reservation_id)) + await _ignore_errors(self._client.zrem(self._index_key(), member)) consumed = clamp_consumption(credits_consumed, reserved) refund = reserved - consumed diff --git a/src/schematic/leases/track.py b/src/schematic/leases/track.py index 76c68ab..0cf3aac 100644 --- a/src/schematic/leases/track.py +++ b/src/schematic/leases/track.py @@ -39,7 +39,10 @@ async def consume_reservation_and_build_event( options: Optional["TrackWithReservationOptions"] = None, ) -> ReservationConsumeResult: """Settle a hold against its lease and build the track event that bills it.""" - credits = actual_quantity * reservation.consumption_rate + # Rounded up for the same reason the hold is (see ``check_with_lease``): + # the debit has to move the local ledger by exactly what the track event + # bills. + credits = math.ceil(actual_quantity) * reservation.consumption_rate consumed = await reservations.consume(reservation.id, credits) return ReservationConsumeResult( track=build_reservation_track_event(reservation, settled_quantity(actual_quantity), options), diff --git a/tests/custom/test_client.py b/tests/custom/test_client.py index 7b842a1..981c221 100644 --- a/tests/custom/test_client.py +++ b/tests/custom/test_client.py @@ -1642,6 +1642,26 @@ def test_event_usage_and_credit_cost_are_forwarded_as_preflight(self): ), ) + def test_a_fractional_preflight_quantity_rounds_up_on_the_wire(self): + # The options take any finite quantity; the REST body's usage is an + # integer, so a fraction rounds up rather than letting the check pass + # on less usage than the action is about to record. + self.schematic.check_flag( + "inference", + company={"id": "co_1"}, + options=CheckFlagOptions( + usage=0.5, event_usage=EventUsage(event_subtype="inference_tokens", quantity=2.5), + ), + ) + preflight = self.schematic.features.check_flag.call_args.kwargs["preflight"] + self.assertEqual( + preflight, + PreflightRequestBody( + usage=1, + event_usage=PreflightEventUsageRequestBody(event_subtype="inference_tokens", quantity=3), + ), + ) + def test_preflighted_check_neither_reads_nor_writes_the_cache(self): company = {"id": "co_1"} options = CheckFlagOptions(usage=5) @@ -2698,6 +2718,29 @@ async def test_track_with_reservation_settles_the_local_hold(self): finally: await self._drain(client) + async def test_track_with_reservation_moves_the_cached_metric_only_on_the_settle(self): + client = _async_lease_client() + client._datastream_client = _lease_datastream([LEASE_PROBE, LEASE_GATE]) + try: + result = await self._check(client) + assert result.reservation is not None + with patch.object(client.event_buffer, "push", new=AsyncMock()) as mock_push: + await client.track_with_reservation(result.reservation, 20) + client._datastream_client.update_company_metrics.assert_awaited_once_with( + {"id": "co_1"}, "inference_tokens", 20, + ) + + # The hold is already consumed, so this settle changes nothing + # locally and the server drops the event on its idempotency + # key. Bumping the metric again would deny the company's next + # numeric-limit check on usage nobody recorded. + await client.track_with_reservation(result.reservation, 20) + + assert mock_push.await_count == 2 + assert client._datastream_client.update_company_metrics.await_count == 1 + finally: + await self._drain(client) + async def test_track_with_reservation_emits_even_when_the_settle_raises(self): client = _async_lease_client() client._datastream_client = _lease_datastream([LEASE_PROBE, LEASE_GATE]) @@ -2800,6 +2843,43 @@ async def test_prewarm_with_no_wait_gives_up_on_an_uncached_company(self): finally: await self._drain(client) + async def test_prewarm_resolves_an_account_defined_id_key_through_the_cache(self): + # The account's own identifier happens to live under a key named `id`. + # It is an ordinary entity key, so the lookup decides. + client = _async_lease_client( + credit_leases=CreditLeaseConfig(default_lease_size=1000.0, prewarm_resolve_timeout=0) + ) + client._datastream_client = _lease_datastream([], company_cached=True) + try: + await client.prewarm({"id": "acme"}, ["bilcr_inference"]) + client.credits.acquire_credit_lease.assert_awaited_once() + assert client.credits.acquire_credit_lease.call_args.kwargs["company_id"] == "co_1" + finally: + await self._drain(client) + + async def test_prewarm_falls_back_to_a_comp_prefixed_value_when_the_keys_miss(self): + client = _async_lease_client( + credit_leases=CreditLeaseConfig(default_lease_size=1000.0, prewarm_resolve_timeout=0) + ) + client._datastream_client = _lease_datastream([]) + try: + await client.prewarm({"account_id": "comp_1"}, ["bilcr_inference"]) + client.credits.acquire_credit_lease.assert_awaited_once() + assert client.credits.acquire_credit_lease.call_args.kwargs["company_id"] == "comp_1" + finally: + await self._drain(client) + + async def test_prewarm_resolves_nothing_when_the_keys_miss_and_carry_no_schematic_id(self): + client = _async_lease_client( + credit_leases=CreditLeaseConfig(default_lease_size=1000.0, prewarm_resolve_timeout=0) + ) + client._datastream_client = _lease_datastream([]) + try: + await client.prewarm({"id": "acme"}, ["bilcr_inference"]) + client.credits.acquire_credit_lease.assert_not_awaited() + finally: + await self._drain(client) + async def test_prewarm_is_a_no_op_in_server_mode(self): client = _async_server_client() try: @@ -2957,6 +3037,30 @@ async def test_prewarm_started_during_shutdown_acquires_nothing(self): finally: await self._drain(client) + async def test_shutdown_returns_within_the_budget_when_a_release_never_lands(self): + client = _async_lease_client() + await client._lease_store.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="bilcr_inference", + granted_amount=1000, + expires_at=time.time() + 300, + ) + + async def never(*args, **kwargs): + await asyncio.Event().wait() + + client.credits.release_credit_lease = AsyncMock(side_effect=never) + + with patch("schematic.client.SHUTDOWN_DRAIN_TIMEOUT", 0.05): + started = time.monotonic() + await client.shutdown() + + # The drain and the release share one budget, so a wire call that never + # lands cannot hold a closing client open. + assert time.monotonic() - started < 1 + + async def test_shutdown_leaves_a_shared_lease_for_the_pods_still_drawing_on_it(self): redis_client = make_fake_redis() client = _async_lease_client( diff --git a/tests/datastream/test_datastream_client.py b/tests/datastream/test_datastream_client.py index 5764b63..6cf0195 100644 --- a/tests/datastream/test_datastream_client.py +++ b/tests/datastream/test_datastream_client.py @@ -1177,3 +1177,92 @@ async def test_check_flag_returns_default_on_user_key_conflict(self, client: Dat assert result.reason == "key conflict" assert result.flag_key == "user-conflict-flag" assert result.err is not None + + +class TestDataStreamClientEvalContextPreflight: + """A preflight on the evaluation context has to reach the local engine, or + the same call answers the hypothetical over REST and the plain question + here.""" + + async def _client(self, logger: logging.Logger) -> tuple[DataStreamClient, MagicMock]: + cache = MockCacheProvider() + client = DataStreamClient(DataStreamClientOptions( + api_key="test-key", + logger=logger, + replicator_mode=True, + company_cache=cache, + company_lookup_cache=cache, + user_cache=cache, + user_lookup_cache=cache, + flag_cache=cache, + )) + engine = MagicMock() + engine.is_initialized.return_value = True + engine.check_flag.return_value = RulesengineCheckFlagResult( + value=True, reason="match", flag_key="pf-flag", + ) + client._rules_engine = engine + await client._handle_message(DataStreamResp( + data={ + "key": "pf-flag", "id": "f1", "default_value": True, "rules": [], + "account_id": "acc_1", "environment_id": "env_1", + }, + entity_type=EntityType.FLAG.value, + message_type=MessageType.FULL.value, + )) + return client, engine + + async def test_hands_the_engine_a_preflight_the_eval_context_carries( + self, logger: logging.Logger + ) -> None: + from schematic.types import PreflightRequestBody + + client, engine = await self._client(logger) + + await client.check_flag( + CheckFlagRequestBody(preflight=PreflightRequestBody(usage=7)), "pf-flag", + ) + + options = engine.check_flag.call_args.args[3] + assert options is not None + assert options.usage == 7 + assert options.event_usage is None + + async def test_the_options_usage_knobs_replace_the_eval_contexts( + self, logger: logging.Logger + ) -> None: + from schematic.client import CheckFlagOptions, EventUsage + from schematic.types import PreflightRequestBody + + client, engine = await self._client(logger) + + await client.check_flag( + CheckFlagRequestBody( + preflight=PreflightRequestBody(usage=7, credit_cost={"credit-1": 20}), + ), + "pf-flag", + options=CheckFlagOptions(event_usage=EventUsage(event_subtype="tokens", quantity=9)), + ) + + options = engine.check_flag.call_args.args[3] + assert options is not None + # The pair moves as one, so the context's usage goes with it; the + # credit cost it carried rides along, since the options named none. + assert options.usage is None + assert options.event_usage == EventUsage(event_subtype="tokens", quantity=9) + assert options.credit_cost == {"credit-1": 20} + + async def test_an_options_credit_cost_wins(self, logger: logging.Logger) -> None: + from schematic.client import CheckFlagOptions + from schematic.types import PreflightRequestBody + + client, engine = await self._client(logger) + + await client.check_flag( + CheckFlagRequestBody(preflight=PreflightRequestBody(credit_cost={"credit-1": 20})), + "pf-flag", + options=CheckFlagOptions(credit_cost={"credit-1": 5}), + ) + + options = engine.check_flag.call_args.args[3] + assert options is not None and options.credit_cost == {"credit-1": 5} diff --git a/tests/datastream/test_rules_engine.py b/tests/datastream/test_rules_engine.py index d20019f..3b64e89 100644 --- a/tests/datastream/test_rules_engine.py +++ b/tests/datastream/test_rules_engine.py @@ -327,7 +327,7 @@ async def engine(self) -> RulesEngineClient: await e.initialize() return e - def _metered_company(self) -> RulesengineCompany: + def _metered_company(self, value: int = 95) -> RulesengineCompany: company_id = "co_metered" company_condition = RulesengineCondition( id="cond_company", @@ -370,7 +370,7 @@ def _metered_company(self) -> RulesengineCompany: event_subtype="api-calls", period="current_month", month_reset="billing_cycle", - value=95, + value=value, created_at="2023-01-01T00:00:00Z", ) return RulesengineCompany( @@ -401,6 +401,32 @@ async def test_usage_that_crosses_the_limit_denies(self, engine: RulesEngineClie result = engine.check_flag(self._flag(), self._metered_company(), None, CheckFlagOptions(usage=10)) assert result.value is False + async def test_a_fractional_usage_counts_as_one_whole_unit(self, engine: RulesEngineClient) -> None: + from schematic.client import CheckFlagOptions + + # One call short of the limit, so rounding the fraction up is what + # decides the verdict. Sent as 0.5 the envelope fails to deserialize + # and takes the whole check with it. + assert engine.check_flag(self._flag(), self._metered_company(99)).value is True + + result = engine.check_flag( + self._flag(), self._metered_company(99), None, CheckFlagOptions(usage=0.5), + ) + assert result.value is False + + async def test_a_fractional_event_usage_counts_as_one_whole_unit( + self, engine: RulesEngineClient, + ) -> None: + from schematic.client import CheckFlagOptions, EventUsage + + result = engine.check_flag( + self._flag(), + self._metered_company(99), + None, + CheckFlagOptions(event_usage=EventUsage(event_subtype="api-calls", quantity=0.5)), + ) + assert result.value is False + async def test_event_usage_for_another_subtype_leaves_the_verdict_alone( self, engine: RulesEngineClient, ) -> None: diff --git a/tests/leases/test_lease_manager.py b/tests/leases/test_lease_manager.py index 1f848c8..d570280 100644 --- a/tests/leases/test_lease_manager.py +++ b/tests/leases/test_lease_manager.py @@ -7,6 +7,8 @@ from __future__ import annotations import asyncio +import logging +import time from typing import Any, List, Optional import pytest @@ -167,6 +169,37 @@ 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_stale_trigger_sends_nothing_once_the_previous_extend_landed( + clock: VirtualClock, monkeypatch: Any +) -> None: + manager, store, wire = _make_manager(clock) + await _drawn_down_lease(store, clock) + stale = await store.get("co_1", "ct_1") + + wire.extend_responses.append({"lease": {"granted_total": 2000, "expires_at": clock() + 600}}) + await manager.maybe_extend("co_1", "ct_1") + assert len(wire.extend_calls) == 1 + + # The second trigger reads the slot as it was before that extend landed: + # its own flight is gone, so nothing stops it reaching the wire but the + # re-read the flight registration now makes. + live = store.get + reads = 0 + + async def staged_get(company_id: str, credit_type_id: str) -> Optional[LeaseState]: + nonlocal reads + reads += 1 + return stale if reads == 1 else await live(company_id, credit_type_id) + + monkeypatch.setattr(store, "get", staged_get) + + await manager.maybe_extend("co_1", "ct_1") + + assert len(wire.extend_calls) == 1 + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.granted_amount == 2000 + + 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 @@ -241,6 +274,53 @@ async def test_two_watermark_joiners_share_the_one_wire_call(clock: VirtualClock assert [entry.local_remaining_credits for entry in results if entry] == [1200] * 3 +async def test_a_follow_up_does_not_inherit_another_callers_smaller_follow_up( + clock: VirtualClock, +) -> None: + # Two checks join one 10000 refill, both needing more than it asked for. + # C's follow-up (2000) registers first; A still needs 18000. Taking C's + # result would send A's retry back to a lease it already knows is short, + # with the credits sitting on the server. + manager, store, wire = _make_manager(clock) + await store.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=10_000, + expires_at=clock() + 300, + ) + await store.try_reserve("co_1", "ct_1", 10_000) + + arrived, release = wire.hold_extend() + wire.extend_responses.append({"lease": {"granted_total": 20_000, "expires_at": clock() + 600}}) + wire.extend_responses.append({"lease": {"granted_total": 22_000, "expires_at": clock() + 600}}) + wire.extend_responses.append({"lease": {"granted_total": 40_000, "expires_at": clock() + 600}}) + + refill = asyncio.ensure_future(manager.maybe_extend("co_1", "ct_1", 10_000)) + await arrived.wait() + assert wire.extend_calls[0]["additional_amount"] == 10_000 + + # C joins first, so its follow-up is the one in flight when A resumes. + joiner_c = asyncio.ensure_future(manager.maybe_extend("co_1", "ct_1", 12_000)) + await _settle() + joiner_a = asyncio.ensure_future(manager.maybe_extend("co_1", "ct_1", 28_000)) + await _settle() + assert len(wire.extend_calls) == 1 + + release.set() + await refill + entry_c = await joiner_c + entry_a = await joiner_a + + # C topped up its 2000 shortfall; A then asked for its own, sized against + # the 12000 C left behind rather than inheriting C's ask. + assert len(wire.extend_calls) == 3 + assert wire.extend_calls[1]["additional_amount"] == 2_000 + assert wire.extend_calls[2]["additional_amount"] == 16_000 + assert entry_c is not None and entry_c.local_remaining_credits >= 12_000 + assert entry_a is not None and entry_a.local_remaining_credits >= 28_000 + + 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 @@ -265,6 +345,37 @@ async def test_the_follow_up_never_chains(clock: VirtualClock) -> None: assert joined is not None and joined.local_remaining_credits == 2200 +async def test_a_joiner_gives_up_on_a_flight_that_outlasts_its_own_timeout( + clock: VirtualClock, +) -> None: + # The flight runs on whatever timeout started it (a background refresh uses + # the client default). A check with 50ms to spend must not sit behind it: + # it gives up, takes its fail-open/fail-closed path, and leaves the flight + # running for everyone else. + 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}}) + + flight = asyncio.ensure_future(manager.maybe_extend("co_1", "ct_1")) + await arrived.wait() + assert len(wire.extend_calls) == 1 + + started = time.monotonic() + impatient = await manager.maybe_extend("co_1", "ct_1", 900, 0.05) + waited = time.monotonic() - started + + assert impatient is None + assert waited < 1 + # No second wire call: the joiner abandoned its wait, it did not race + # another extend onto the lease. + assert len(wire.extend_calls) == 1 + + release.set() + entry = await flight + assert entry is not None and entry.granted_amount == 2000 + + 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. @@ -359,6 +470,32 @@ async def test_release_all_releases_live_and_skips_expired(clock: VirtualClock) assert await store.get("co_2", "ct_1") is not None +async def test_release_all_gives_up_on_a_release_that_never_lands( + clock: VirtualClock, caplog: Any +) -> None: + manager, store, wire = _make_manager(clock) + await store.replace( + lease_id="lse_live", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=clock() + 60, + ) + + async def never(lease_id: str) -> None: + await asyncio.Event().wait() + + wire.release = never # type: ignore[assignment] + + started = time.monotonic() + with caplog.at_level(logging.WARNING, logger="schematic.leases.lease_manager"): + await manager.release_all_local_leases(0.05) + + # A shutdown that hangs is worse than a hold the server expires. + assert time.monotonic() - started < 1 + assert any("releasing credit leases on close" in record.getMessage() for record in caplog.records) + + async def test_release_all_skips_a_shared_store(clock: VirtualClock) -> None: # A shared backend cannot enumerate: sibling pods still draw on its leases. client = make_fake_redis() diff --git a/tests/leases/test_redis_reservation_store.py b/tests/leases/test_redis_reservation_store.py index c6271a4..4034877 100644 --- a/tests/leases/test_redis_reservation_store.py +++ b/tests/leases/test_redis_reservation_store.py @@ -50,6 +50,60 @@ async def test_add_round_trips_and_indexes( assert await reservations.count() == 1 +async def test_add_writes_the_hash_and_its_ttl_in_one_transaction( + redis_client: Any, reservations: RedisReservationStore, frozen_clock: VirtualClock +) -> None: + transactions: list[list[str]] = [] + original = redis_client.pipeline + + def recording_pipeline(*args: Any, **kwargs: Any) -> Any: + pipe = original(*args, **kwargs) + execute = pipe.execute + + async def record(*call_args: Any, **call_kwargs: Any) -> Any: + transactions.append([str(queued[0][0]) for queued in pipe.command_stack]) + return await execute(*call_args, **call_kwargs) + + pipe.execute = record + return pipe + + redis_client.pipeline = recording_pipeline + try: + await reservations.add(make_reservation(expires_at=frozen_clock() + 60)) + finally: + redis_client.pipeline = original + + # Written separately, a crash between the two leaves a row that never + # expires and that nothing points at once the sweeper drops its index entry. + assert transactions == [["HSET", "PEXPIREAT"]] + fetched = await reservations.get("res_1") + assert fetched is not None and fetched.credits_reserved == 100 + + +async def test_add_falls_back_when_the_client_refuses_a_transaction( + redis_client: Any, reservations: RedisReservationStore, frozen_clock: VirtualClock +) -> None: + # redis-py's cluster client carries `pipeline` and raises on + # `transaction=True`, so having the attribute settles nothing. Failing here + # would fail the check that is already holding the credits. + original = redis_client.pipeline + + def refuses(*args: Any, **kwargs: Any) -> Any: + raise RuntimeError("transaction is deprecated in cluster mode") + + redis_client.pipeline = refuses + try: + await reservations.add(make_reservation(expires_at=frozen_clock() + 60)) + finally: + redis_client.pipeline = original + + fetched = await reservations.get("res_1") + assert fetched is not None and fetched.credits_reserved == 100 + assert await reservations.count() == 1 + # The TTL still landed, on the sequential path. + assert await redis_client.pttl("schematic:credit-reservation:res_1") > 0 + + async def test_consume_refunds_the_unspent_slice( leases: RedisLeaseStore, reservations: RedisReservationStore, frozen_clock: VirtualClock ) -> None: