diff --git a/.fernignore b/.fernignore index d201894..917bcd7 100644 --- a/.fernignore +++ b/.fernignore @@ -6,8 +6,10 @@ CONTRIBUTING.md LICENSE README.md WASM_VERSION +conformance/ custom.gemspec.rb lib/schematic/cache.rb +lib/schematic/credits/leases/ lib/schematic/datastream/ lib/schematic/event_buffer.rb lib/schematic/logger.rb @@ -18,7 +20,11 @@ lib/schematic/wasm/ lib/schematic/webhook_verification.rb lib/schematichq.rb scripts/ +test/conformance_test.rb +test/credits_test.rb test/custom.test.rb +test/lease_support.rb +test/rules_engine_clock_test.rb testapp/ .fern/replay.lock .fern/replay.yml diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 65f57c4..0c65e84 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -37,6 +37,10 @@ jobs: bundler-cache: true - name: Download WASM binary run: ./scripts/download-wasm.sh + # Not `rake test`: the Rakefile is generated and not fernignored, so a + # regeneration could silently drop the glob that picks these up. + - name: Run test suite + run: bundle exec ruby -I lib -I test -e 'Dir["test/**/{test_*,*_test}.rb"].sort.each { |f| require File.expand_path(f) }' - run: bundle exec ruby -I lib -I test test/custom.test.rb publish: diff --git a/README.md b/README.md index 99e7464..8954723 100644 --- a/README.md +++ b/README.md @@ -363,6 +363,147 @@ end client.close ``` +### Credit Leases and Reservations + +For features metered by credit burndown (e.g. inference tokens), `check` reserves credits for the work about to run and `track_with_reservation` settles the reservation with actual usage. The SDK gates in one of two modes: + +- **Client mode** acquires a **lease**, a tranche of credits held against the company's balance, and carves a per-request **reservation** out of it locally, so a check needs no API call. It requires [DataStream](#datastream) (or [Replicator Mode](#replicator-mode)) and, across multiple processes, a shared Redis so every instance gates against the same lease. +- **Server mode** makes one `check-and-reserve` API call per check. No lease, no Redis, no local state. + +`credit_leases[:mode]` defaults to `:auto`: client when DataStream is enabled, server otherwise. Client mode suits high-throughput gating; server mode suits low-volume checks and operations that run for seconds. + +Durations are in milliseconds, as in the other Schematic SDKs. + +#### Setup + +The `redis` gem is not a dependency; pass a client you already have. + +```ruby +require "redis" +require "schematichq" + +redis_client = Redis.new(url: "redis://localhost:6379") + +client = Schematic::SchematicClient.new( + api_key: ENV["SCHEMATIC_API_KEY"], + use_data_stream: true, + datastream_options: { + redis_client: redis_client # also backs lease and reservation state + }, + credit_leases: { + default_lease_size: 10_000, # credits requested per lease + default_lease_duration: 5 * 60 * 1000, # lease lifetime (ms) + default_reservation_ttl: 60 * 1000 # how long a reservation is held if no track settles it (ms) + } +) +``` + +Set `credit_leases[:redis_client]` to keep lease state in a different Redis than the DataStream cache. + +Server mode needs only a TTL: + +```ruby +client = Schematic::SchematicClient.new( + api_key: ENV["SCHEMATIC_API_KEY"], + credit_leases: { + default_reservation_ttl: 60 * 1000 # how long the server holds the credits if no track settles them (ms, max 1 hour) + } +) +``` + +Only `mode` and `default_reservation_ttl` apply in server mode; the SDK warns at startup if a client-only option is set. + +#### Checking and tracking + +```ruby +# Reserve up to max_tokens for this operation. +result = client.check( + "inference", + company: { "id" => "your-company-id" }, + usage: max_tokens, # upper bound for this operation + event_subtype: "inference_tokens" # the metered event +) +raise "credit balance exceeded" unless result.allowed? + +inference = run_inference(...) + +# Report actual usage; the unused slice of the reservation is refunded. +if result.reservation + client.track_with_reservation(result.reservation, inference.tokens_used) +else + client.track({ + event: "inference_tokens", + company: { "id" => "your-company-id" }, + quantity: inference.tokens_used + }) +end +``` + +A check can allow without reserving credits (the feature is not credit-metered, `usage` is 0, or the check failed open), and that usage still has to be tracked. + +Without `credit_leases` configured, or without a `usage`, `check` falls through to a plain flag check with no reservation. `usage` is still sent as a preflight, locally or to the API, so the verdict accounts for what the call is about to spend. Preflighted verdicts are not cached. + +`usage` may be fractional. A client-mode reservation is sized in whole event units, `ceil(usage) x consumption_rate`, and a settle debits the lease by `ceil(actual) x consumption_rate`, so the local ledger moves by exactly what the track event bills. The reservation still records the fractional quantity the caller declared. Everything on the wire rounds up with the debit: the track event's quantity, a server-mode reservation, and the preflight quantity, whether it goes to the API or the local engine. + +`result.entitlement` is a symbol-keyed hash with the field names of `Schematic::Types::FeatureEntitlement`, the same in both modes. + +An unsettled reservation expires after `default_reservation_ttl` and its credits return to the lease. A late settle still bills the usage (the track event carries a deterministic idempotency key, so it never double-bills) but does not re-debit the local lease, so set `default_reservation_ttl` above the longest expected gap between `check` and `track_with_reservation`. + +#### Pre-warming + +Pre-warm leases when the user is identified, so a session's first check does not wait on a lease acquire: + +```ruby +client.identify( + { + keys: { "user_id" => "your-user-id" }, + company: { keys: { "id" => "your-company-id" } } + }, + prewarm: ["credit-type-id"] # credit type IDs to acquire leases for +) +``` + +Or call `client.prewarm(credit_type_ids, company:)` directly. Both are no-ops in server mode. + +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. + +#### Failure behavior + +A check that cannot be gated (API unreachable, Redis down, lease exhausted) fails closed by default. Override per check: + +```ruby +result = client.check( + "inference", + company: { "id" => "your-company-id" }, + usage: max_tokens, + event_subtype: "inference_tokens", + on_acquire_failure: :fail_open +) +``` + +In server mode, a check that times out after the server has already reserved leaves those credits reserved until the TTL expires, so keep `default_reservation_ttl` short there. + +In client mode, `:fail_open` still evaluates the flag's rules with the credit balance assumed sufficient, so plan targeting and all non-credit conditions apply and only the credit gate is bypassed. In server mode it returns the flag's default value, which is `false` unless you pass `default_value` or configure a `flag_defaults` entry. + +`default_value` also applies when a check falls back to a plain flag check and that check fails. Pass a boolean or a callable. + +`check` accepts `timeout_ms`, but no timeout is configurable yet: the generated HTTP transport ignores per-request timeouts and the client exposes no setting. The option is carried on the request so it takes effect once the transport reads it. + +#### Configuration options + +| Option | Type | Default | Description | +|---|---|---|---| +| `mode` | `:client`, `:server`, `:auto` | `:auto` | Where credits are reserved; `:auto` picks client when DataStream is enabled, server otherwise | +| `default_reservation_ttl` | `Integer` | 60000 (60 seconds) | How long an unsettled reservation is held (ms) | +| `default_lease_duration` | `Integer` | 300000 (5 minutes) | (client mode) Lease lifetime (ms) | +| `default_lease_size` | `Numeric` | 10000 | (client mode) Credits requested per lease acquire or extend | +| `low_water_mark` | `Float` | 0.25 | (client mode) Extend in the background when the lease balance dips below this fraction | +| `sweep_interval_ms` | `Integer` | 1000 | (client mode) Sweep interval for expired reservations (ms) | +| `prewarm_resolve_timeout_ms` | `Integer` | 5000 | (client mode) How long `prewarm` waits for a freshly identified company to surface (ms); 0 resolves from the DataStream cache only | +| `redis_client` | Redis client | `datastream_options[:redis_client]` | (client mode) Redis client for lease and reservation state | +| `redis_key_prefix` | `String` | `datastream_options[:redis_key_prefix]` | (client mode) Key prefix for lease and reservation keys | +| `overrides` | `Hash` | none | (client mode) Per-credit-type overrides of the above, keyed by credit type ID | + ### Other API operations The Schematic API supports many operations beyond these, accessible via the API modules on the client: `accounts`, `billing`, `companies`, `credits`, `entitlements`, `events`, `features`, and `plans`. diff --git a/conformance/SPEC.md b/conformance/SPEC.md new file mode 100644 index 0000000..aed9498 --- /dev/null +++ b/conformance/SPEC.md @@ -0,0 +1,480 @@ +# Credit lease & reservation semantics — conformance spec + +This document specifies the client-side credit lease/reservation semantics implemented by the +Schematic Node SDK (the reference implementation), in enough detail to reimplement them in another +language without reading the Node source. The machine-readable test vectors in +`conformance/vectors/*.json` pin the observable behavior; this spec explains the model, the +configuration knobs, and the invariants that cannot be expressed as deterministic vectors. + +Where this document and the vectors disagree, the vectors win — they are generated from the +reference implementation's behavior. + +- [Vector format](#vector-format) +- [Model overview](#model-overview) +- [State](#state) +- [Store operations](#store-operations) +- [Lease manager](#lease-manager) +- [Check flow](#check-flow) +- [Track / settle flow](#track--settle-flow) +- [Configuration knobs](#configuration-knobs) +- [Bounded-leak contract](#bounded-leak-contract) +- [Invariants not expressible as vectors](#invariants-not-expressible-as-vectors) + +## Vector format + +Each file in `conformance/vectors/` is a JSON document: + +```json +{ + "category": "reservation_lifecycle", + "vectors": [ + { + "name": "unique_snake_case_name", + "description": "What this vector pins and why.", + "backends": ["in_memory", "redis"], + "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": 60000 } ] + }, + "operations": [ + { "op": "try_reserve", "company_id": "co_1", "credit_type_id": "ct_1", "credits": 100, "expect": { "balance": 900 } } + ] + } + ] +} +``` + +Rules: + +- All keys are `snake_case`. Vectors are plain JSON — no language-specific types. +- **Virtual clock.** The run starts at a fixed virtual instant `t0`. Every `*_at_ms` field is an + offset in milliseconds from `t0` (an absolute position on the virtual timeline, not relative to + the current operation). The `advance_clock` operation moves the clock forward; nothing else does. + Runners must execute vectors against a controllable clock (no wall time). +- `backends` restricts which store backends the vector runs against; when omitted, the vector must + pass against every backend the SDK ships (in-memory and Redis for Node). +- `given.leases` are installed via the store's `replace` operation at `t0` (each install must + return "written"). +- Assertions are attached per-operation via `expect`. Final-state assertions are expressed as + trailing read operations (`get_lease`, `reserved_credits`, `reservation_count`). +- `expect.balance` / `expect.consumed` use JSON `null` for the "no / refused" result. +- Reservation ids created by `check` operations are random; the vector names them via + `save_reservation_as` and later operations reference them with `handle`. + +### Operations + +Store-level (exercise the lease store and reservation store directly): + +| op | fields | expect | +| --- | --- | --- | +| `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`), `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`? | +| `add_reservation` | `id`, `lease_id`, `company_id`, `credit_type_id`, `event_subtype`, `quantity_reserved`, `credits_reserved`, `consumption_rate`, `expires_at_ms` | — | +| `consume_reservation` | `id` or `handle`, `credits`, `crash_before_refund`? (bool, one-shot) | `consumed` (number or `null`), `throws`? | +| `get_reservation` | `id` or `handle` | `exists` | +| `reserved_credits` | `company_id`, `credit_type_id` | `total` | +| `reservation_count` | — | `count` | +| `sweep_expired` | — | `swept` | + +Manager-level (exercise the lease manager with a scripted wire client): + +| op | fields | expect | +| --- | --- | --- | +| `acquire_if_needed` | `company_id`, `credit_type_id`, `server`? ( `{ "lease": {...} }` or `{ "error": "..." }` ), `install_during_wire`? (lease installed into the store while the wire call is in flight, emulating a sibling pod winning the race) | `lease_id` (or `null`), `wire_acquires` (cumulative count), `last_acquire_requested_amount`?, `released_lease_ids` (cumulative) | +| `maybe_extend` | `company_id`, `credit_type_id`, `required_credits`?, `server`? ( `{ "lease": { "granted_total", "expires_at_ms" } }` or `{ "error": "..." }` ) | `wire_extends` (cumulative count), `last_extend_additional_amount`?, `last_extend_lease_id`? | +| `release_all_local_leases` | — (in-memory backend only) | `released_lease_ids`, `remaining_slots` | + +Flow-level (exercise the full check/track orchestration with a scripted rules engine): + +| op | fields | expect | +| --- | --- | --- | +| `check` | `flag_key`, `company` (`{ id, credit_balances }`), `usage`, `event_subtype`?, `on_acquire_failure`?, `engine` (array of scripted engine results, consumed in call order), `server`? (as above), `save_reservation_as`? | `allowed`, `reason`?, `err`?, `has_reservation`, `reservation`? (field subset), `fallback_called`?, `engine_calls`? (per-call `{ credit_balance, credit_cost?, event_usage? }`; `credit_balance` may be the string `"max_safe_integer"`) | +| `track` | `handle`, `actual_quantity` | `settled_locally`, `track` (`{ event, quantity, lease_id }`) | + +A scripted engine result is `{ "value": bool, "reason"?: string, "entitlement"?: { "value_type", +"credit_id"?, "consumption_rate"?, "event_subtype"?, "feature_id"?, "feature_key"? } }`. The engine +is an oracle: the vectors pin the *orchestration around* the rules engine (what it is called with, +and what the SDK does with its answer), not the engine itself — the engine is shared WASM across +SDKs and has its own tests. + +## Model overview + +Credit-metered features are gated client-side without a wire call per check. The SDK: + +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** `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 + server (the server is the source of truth for real consumption). + +Everything client-side is *local bookkeeping against the leased tranche*. The server reconciles: +an expired lease's unspent remainder is refunded to the company balance server-side, and Track +events (keyed by `lease_id`) drive the authoritative consumption. + +Leases and reservations both expire: + +- A **lease** past its expiry must be treated as *released* — its local balance is stale (the + server already refunded the remainder) and must never serve another reserve or be extended. +- A **reservation** past its TTL is swept: removed from the table and its full hold refunded to + the lease. Work that finishes after the sweep still bills the server (recovery emit) but does + not re-debit the local lease. + +## State + +**Lease slot** — at most one lease per `(company_id, credit_type_id)` key: + +| field | meaning | +| --- | --- | +| `lease_id` | Server-issued id. | +| `granted_amount` | Server-authoritative total granted to this lease (grows on extend). | +| `local_remaining_credits` | Local view: granted minus outstanding holds/consumption. Initialized to `granted_amount` on install. | +| `expires_at` | Expiry instant. Past it the lease is dead (see above). | + +**Reservation** — keyed by a unique id: + +| field | meaning | +| --- | --- | +| `id` | Unique (UUID in Node). | +| `lease_id` | The lease the hold was carved from. Pins refunds. | +| `company_id`, `credit_type_id` | Slot key. | +| `event_subtype` | Event the settle will bill as. | +| `quantity_reserved` | Caller-declared usage (event units). | +| `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. | + +## Store operations + +These are the primitives both store backends (per-process in-memory; shared Redis) must implement +with identical observable semantics. Each mutation must be atomic per slot/reservation (see +[Invariants](#invariants-not-expressible-as-vectors)). + +### `replace(lease)` — install-if-not-live + +Install a fresh lease with `local_remaining_credits = granted_amount`, **only if** the slot is +empty or the existing lease is expired *and carries a different `lease_id`*. If a *live* lease +occupies the slot — even with a different `lease_id` (a sibling pod won the race) — leave it +untouched (its already-debited balance wins) and report "kept". If an *expired* lease with the +**same** `lease_id` occupies the slot (a stale acquire response for a lease the idempotent server +also handed to a racing sibling, which may since have extended it), do not rewrite it either: +rewriting would reset `local_remaining_credits` to the full grant and erase debits whose +reservations are still open. Reconcile it like `extend` instead — granted to the incoming total +(lower/equal totals are no-ops), expiry only forward, balance untouched — and report "kept". +Returns written/kept so the caller can run the redundant-lease release logic (see manager). + +### `try_reserve(company, credit, credits)` — atomic check-and-debit + +- Reject (return `null`, touch nothing) if: no lease in the slot; the lease is **expired**; the + remaining balance is `< credits`; or `credits` is not a finite non-negative number (NaN must + 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) **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?)` + +Add credits back to `local_remaining_credits`, **clamped at `granted_amount`**. No-op if +`credits <= 0` or no lease is in the slot. When `pin_lease_id` is given, the refund applies +**only if** the slot still holds that lease: a hold carved out of expired lease A must never +inflate successor lease B — A's remainder (including this slice) was already refunded to the +company balance server-side when A expired, so crediting B would double-count. + +### `extend(company, credit, granted_total, new_expires_at?, pin_lease_id?)` — reconcile to total + +After a remote extend, reconcile the slot to the **server-authoritative total**: + +- Compute `delta = granted_total - stored granted_amount` **atomically against the currently + stored total** — never from a caller-held pre-wire-call read (two pods extending concurrently + from the same stale read would each apply a delta and mint phantom credits). If `delta > 0`, + set `granted_amount = granted_total` and add `delta` to `local_remaining_credits`. If + `delta <= 0` (a total a sibling already applied, or a stale lower total) it is a **no-op** — + applies converge in any order. +- Expiry only ever moves **forward**: `new_expires_at` is applied only if later than the stored + expiry, so an out-of-order apply cannot shorten a lease a sibling just extended. +- When `pin_lease_id` is given and the slot holds a different lease, drop the whole extend + (credits and expiry): the server granted the extension to the pinned lease; crediting a + successor would mint credits the server refunds with the pinned lease at its expiry. +- No-op if the slot is empty. + +### `drop(company, credit)` + +Remove the slot entry (after a remote release). Plain delete. + +### Reservation table: `add`, `get`, `consume`, `reserved_credits`, `sweep_expired` + +- `add(reservation)` — register. Idempotent on id. `add` does NOT debit the lease; the debit + already happened in `try_reserve` (see [ordering](#bounded-leak-contract)). +- `consume(id, credits_consumed)` — **exactly-once claim**: atomically remove the reservation + from the table; if it was already gone (swept, or consumed by a racing caller) return `null` + and touch nothing. On a successful claim, clamp `credits_consumed` to + `[0, credits_reserved]`, refund `credits_reserved - clamped` to the lease (pinned to the + reservation's `lease_id`), and return the clamped figure. The claim and the refund are two + steps; the claim is the arbiter (see bounded-leak contract). +- `reserved_credits(company, credit)` — sum of `credits_reserved` across open reservations for + the slot. A reservation counts iff it is still in the table, so + `local_remaining_credits + reserved_credits` stays exact between operations. +- `sweep_expired(now)` — remove every reservation with `expires_at <= now` and refund its full + hold to its lease (pinned to its `lease_id`; a stale-lease hold is dropped, not refunded). + Returns the number swept. Runs on a background interval (`sweep_interval_ms`) in production; + vectors call it explicitly. + +## Lease manager + +Owns the lease lifecycle against the server wire API (`acquire`, `extend`, `release`). + +### Acquire (`acquire_if_needed`) + +- If the slot holds a **live** lease, return it — no wire call. +- Otherwise call the server: `requested_amount = lease_size`, `expires_at = now + + lease_duration_ms`. An expired local entry is left in place for `replace` to overwrite + atomically (deleting it first would open a race window against sibling pods; every reader + re-guards on expiry anyway). +- On response, `replace` the slot. If `replace` kept an existing lease (a sibling won, or the + slot's expired row was reconciled in place): + - If the installed lease has a **different id** than the one the server handed us, ours is a + redundant hold nobody will draw on — release it (fire-and-forget; a failed release falls + back to server-side lease expiry). + - If the ids are the **same** (the server is idempotent for an active slot and handed the + racing acquire the sibling's lease back), release **nothing** — releasing would pull the + shared lease out from under every sibling. + - If the slot reads empty (expired in the gap), also release nothing. + - Either way, return whatever the slot now holds. +- Wire or store failure: return "no lease" (never throw) — the caller routes it through + fail-open/fail-closed. +- Per-process single-flight per slot: concurrent callers share one in-flight wire call + (best-effort; duplicates are absorbed by the idempotent server + `replace`). + +### Extend (`maybe_extend`) + +Triggered when EITHER: + +- `local_remaining_credits / max(granted_amount, 1) <= low_water_mark` (steady-state refresh), or +- the caller passes `required_credits` and `local_remaining_credits < required_credits` (a check + just failed a reserve of that size — extend opportunistically). + +Rules: + +- **Never extend an expired lease** — the server treats it as released; the right move is a fresh + acquire on the next check. +- Wire body: `additional_amount = max(lease_size, required_credits - local_remaining_credits)`. + 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. 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). +- Per-process single-flight per slot, kept separate from acquire's (an in-flight extend must not + satisfy an acquire, or vice versa). + +### Release on close (`release_all_local_leases`) + +Only for a **per-process (in-memory) store**, whose leases are exclusively this process's: +release every live lease over the wire (returning the unspent remainder to the company balance +immediately) and drop it locally; **skip expired** leases (already swept server-side). A shared +(Redis) store must never do this — sibling pods still draw on those leases. Best-effort: +failures fall back to server-side expiry. + +## Check flow + +`check(eval_ctx, flag_key, { usage, event_subtype?, on_acquire_failure?, ... })` — the +lease-gated feature check. Fallback = the plain (non-lease) flag check, which has its own +degradation story; when the flow "falls back", no reservation is issued and no lease state is +touched beyond what already happened. + +Guards, in order: + +1. `usage` missing → plain check (lease path not requested). +2. `usage` not a finite non-negative number → resolve **statically** by `on_acquire_failure` + (deny for fail-closed; blanket allow for fail-open, reason `invalid_usage`). The value must + never reach the stores. +3. `usage == 0` → nothing to reserve; fall back to the plain check (no 0-credit reservation). +4. No datastream / cached flag / resolvable company (or named user) → fall back. + +Then: + +5. **Entitlement probe.** Run the rules engine once against the company's *real* balance — no + substitution, no credit-cost preflight (a preflight against the lease-depleted server balance + could fail the credit condition and hide the entitlement being probed for). Read the matched + entitlement's shape: + - Not credit-metered (`value_type != "credit"`: boolean/override grant, numeric allocation, + unlimited, or not entitled) → **fall back**, no lease traffic at all. + - 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 = 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 + `maybe_extend(required_credits = credit_cost)` (awaited) and retry the reserve **once**. + 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. 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 + read race), passing `credit_cost = { credit_id: credit_cost }` so the engine evaluates + `pre_reservation - credit_cost >= 0` — the same arithmetic `try_reserve` just enforced, plus + every non-credit rule (plan targeting, overrides). + - Engine **allows** → keep the hold; return `{ allowed: true, reservation }`. Fire-and-forget + a watermark-driven `maybe_extend`. + - Engine **denies** → cancel the reservation (claim + full refund) and return + `{ allowed: false }` with the engine's reason. + - Engine **errors** → cancel the reservation and resolve **statically** by mode (the engine + itself is down, so no fail-open re-evaluation is possible). + +### Failure handling + +Every can't-gate outcome (acquire failed, store unreachable, lease exhausted) funnels through the +configured `on_acquire_failure` mode (default **fail-closed**): + +- **fail-closed** → `{ allowed: false }`, reason = the failure reason. No reservation. +- **fail-open** → *err on the side of assuming the credits are there*, **not** blanket allow: + re-run the engine with the credit balance substituted to an effectively unlimited value + (`MAX_SAFE_INTEGER` in Node) and the caller's usage preflight threaded through. Plan + targeting, overrides, and every non-credit condition still apply — a company that is not + entitled stays **denied** even with the lease backend down. No reservation is issued either + way; `err` carries the failure reason. Only if that evaluation itself errors does the SDK + fall back to a blanket allow. + +## Track / settle flow + +`track_with_reservation(reservation, actual_quantity)` settles a reservation: + +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). + - **Not settled** (`null`: already swept after TTL, already consumed, or store unreachable): + local lease state is untouched — if the sweeper already refunded the full hold, nothing + re-debits the consumed slice, so the local balance reads **high** until the lease rolls + 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 = 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 + event across pods and restarts. +5. Guard: a non-finite or negative `actual_quantity` skips the settle entirely (no store call, no + event) — the untouched reservation expires at its TTL and the sweeper refunds the full hold. + +## Configuration knobs + +| knob (vector key) | Node name | default | meaning | +| --- | --- | --- | --- | +| `lease_duration_ms` | `defaultLeaseDuration` | 300 000 (5 min) | Lease lifetime requested at acquire/extend (`expires_at = now + duration`). | +| `reservation_ttl_ms` | `defaultReservationTTL` | 60 000 (60 s) | Reservation lifetime; the sweep deadline. Size above the longest expected check→track gap. | +| `lease_size` | `defaultLeaseSize` | 10 000 | Credits requested per acquire, and the minimum extend tranche. | +| `low_water_mark` | `lowWaterMark` | 0.25 | Remaining/granted ratio at or below which a background extend is kicked off. | +| `sweep_interval_ms` | `sweepIntervalMs` | 1 000 | Expired-reservation sweep cadence. | +| — | `onAcquireFailure` | `fail-closed` | Per-check failure mode (see check flow). | + +Per-credit-type overrides of the first four are supported (keyed by credit type id); resolution is +override → client config → default. + +## Bounded-leak contract + +The flow deliberately orders its two-step transitions so that a process crash between steps leaks +*locally held credits* (which the server reclaims at lease expiry) rather than enabling a +double-spend. The invariant direction is always: **the debit/claim is durable first; the +record/refund may be lost.** + +| # | crash window | what leaks | bound | reclaimed by | must NOT happen | +| --- | --- | --- | --- | --- | --- | +| 1 | after `try_reserve` (debit), before `add` (record) | the debited hold — invisible to the reservation table, so the sweeper can never refund it | `credits_reserved` of that one check | lease expiry: the expired balance is never served again, and the server refunds the whole grant; the successor lease installs at full grant | a reservation record without a debit (a later consume would refund credits never held → double-spend). Vectors pin that the debit lands strictly before the record. | +| 2 | inside `consume`: after the claim, before the refund | the unspent slice of that reservation | `credits_reserved` of that one reservation | lease expiry (same mechanism) | a double refund: the claim is exactly-once, so a retried settle or a sweeper finds nothing to claim and refunds nothing | +| 3 | (Redis only) after the claim, before index cleanup | nothing (bookkeeping only): the per-slot reserved-credits index transiently over-counts | one index field | the sweeper reconciles the orphaned index entry — **without refunding** (without the claimed record, exactly-once cannot be arbitrated across racing sweepers) | a refund driven by an index entry alone | + +Additional pinned properties: + +- A leak never survives its lease: after lease expiry the stale balance is refused + (`try_reserve → null`) and a successor lease restores the full grant. +- A retried check after a window-1 crash settles independently: its own slice refunds exactly + once; the leaked slice never refunds. +- A late retried settle after a window-2 crash (even after a successor lease is installed) + refunds nothing into the successor. + +## Invariants not expressible as vectors + +These hold in the reference implementation but need concurrency, wall clocks, or non-JSON values +to demonstrate; ports must uphold them and should test them natively. + +1. **Per-slot atomicity.** `replace`, `try_reserve`, `refund`, `extend` are atomic per lease + slot; `consume`'s claim is atomic per reservation. In-memory: per-key mutual exclusion. Redis: + single-key Lua scripts (single-key keeps them Redis-Cluster-safe; the refund to the lease hash + is deliberately a separate single-key step, never a multi-key script — see leak window 2/3). +2. **Server-clock expiry (shared backend).** With a shared store, lease expiry must be decided + against the *store's* clock (Redis `TIME`), not the calling process's — pods with skewed + clocks must agree on liveness. Backend rows carry a TTL grace window past `expires_at` + (60 s lease / 30 s reservation in Node) so the sweeper can still read them; expired-but-not- + evicted rows must still refuse reserves. +3. **NaN/precision guards.** Non-finite or negative amounts are rejected at every boundary + (`usage`, `try_reserve`, `actual_quantity`) — JSON cannot encode NaN, so vectors only cover + the negative case. Fractional credit amounts are legal throughout (rates like 0.1); Redis + stores balances as strings to avoid integer truncation. +4. **Single-flight.** Per-process, per-slot single-flight for acquire and for extend, tracked + separately. Best-effort only: duplicate wire calls are safe (idempotent server + keep-first + `replace` + reconcile-to-total `extend`). An extend flight carries the `additional_amount` it + asked for: a joiner whose required shortfall exceeds that figure waits the flight out and + then issues exactly one further extend for the remaining shortfall, while a joiner the flight + already covers — every watermark-driven one, the common case — issues nothing and shares the + single wire call. +5. **Concurrent cross-pod extends converge.** Two pods extending from the same stale read must + not double-count — guaranteed by reconcile-to-total computed inside the store (the sequential + out-of-order-totals vector pins the arithmetic; the concurrent schedule needs a race). +6. **Idempotent billing.** The Track idempotency key is deterministic from the reservation id; + the server dedupes for 24h. Double settles and recovery emits collapse to one billed event. +7. **Background sweep loop.** `start_sweep`/`stop` run `sweep_expired` on an interval; timers + must not keep the process alive. Vectors call `sweep_expired` explicitly instead. +8. **Fire-and-forget never rejects.** `acquire_if_needed`, `maybe_extend`, and release paths + resolve (to "no lease") on failure rather than rejecting — they are often unawaited. +9. **Offline/unconfigured degradation.** Lease config absent → `check` is a plain flag check; + `track_with_reservation` on an unconfigured client still emits the billing event with the + `lease_id` and idempotency key intact. diff --git a/conformance/vectors/check-flow.json b/conformance/vectors/check-flow.json new file mode 100644 index 0000000..eff9990 --- /dev/null +++ b/conformance/vectors/check-flow.json @@ -0,0 +1,527 @@ +{ + "category": "check_flow", + "vectors": [ + { + "name": "check_happy_path_gates_on_pre_reservation_balance", + "description": "A lease-bearing check: probe against the real balance (no substitution, no credit cost), reserve usage x rate from the lease, then gate the engine on the PRE-reservation local balance with credit_cost — the same arithmetic the atomic reserve just enforced. The hold sticks only because the engine allowed.", + "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": 10, + "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", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10 + }, + "engine_calls": [{ "credit_balance": 5000 }, { "credit_balance": 1000, "credit_cost": 100 }] + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 100 } } + ] + }, + { + "name": "check_denied_by_engine_cancels_the_hold", + "description": "When the gate evaluation denies, the reservation made before the eval is cancelled: claimed and fully refunded, leaving no hold and no reservation.", + "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": 10, + "event_subtype": "inference_tokens", + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + }, + { "value": false, "reason": "denied_by_targeting" } + ], + "expect": { "allowed": false, "reason": "denied_by_targeting", "has_reservation": false } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 0 } }, + { "op": "reservation_count", "expect": { "count": 0 } } + ] + }, + { + "name": "check_acquire_failure_fail_closed_denies", + "description": "fail-closed (the default): when no lease can be acquired the check denies outright with the failure reason; no reservation is issued.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "on_acquire_failure": "fail-closed", + "server": { "acquire": { "error": "wire down" } }, + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + } + ], + "expect": { + "allowed": false, + "reason": "lease_acquire_failed", + "err": "lease_acquire_failed", + "has_reservation": false + } + }, + { "op": "reservation_count", "expect": { "count": 0 } } + ] + }, + { + "name": "check_acquire_failure_fail_open_reevaluates", + "description": "fail-open is NOT blanket allow: the engine re-runs with the credit balance substituted to an effectively unlimited value and the caller's usage preflight threaded through, so non-credit rules still apply. Here they pass, so the check allows — with the failure recorded in err and no reservation.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "on_acquire_failure": "fail-open", + "server": { "acquire": { "error": "wire down" } }, + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + }, + { "value": true, "reason": "evaluated" } + ], + "expect": { + "allowed": true, + "reason": "evaluated (lease_acquire_failed_fail_open)", + "err": "lease_acquire_failed", + "has_reservation": false, + "engine_calls": [ + { "credit_balance": 5000 }, + { + "credit_balance": "max_safe_integer", + "event_usage": { "event_subtype": "inference_tokens", "quantity": 10 } + } + ] + } + }, + { "op": "reservation_count", "expect": { "count": 0 } } + ] + }, + { + "name": "check_fail_open_still_denies_when_rules_deny", + "description": "fail-open with a denying rules evaluation stays denied: substituting an unlimited balance only bypasses the credit gate, never plan targeting or overrides.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "on_acquire_failure": "fail-open", + "server": { "acquire": { "error": "wire down" } }, + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + }, + { "value": false, "reason": "not_targeted" } + ], + "expect": { + "allowed": false, + "reason": "not_targeted (lease_acquire_failed_fail_open)", + "err": "lease_acquire_failed", + "has_reservation": false + } + } + ] + }, + { + "name": "check_insufficient_lease_extends_and_retries", + "description": "A reserve refusal triggers an awaited opportunistic extend sized to cover the request (required_credits = credit_cost), then exactly one reserve retry. On success the check proceeds normally, gating on the post-extend pre-reservation balance.", + "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": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 950, + "expect": { "balance": 50 } + }, + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "server": { "extend": { "lease": { "granted_total": 2000, "expires_at_ms": 600000 } } }, + "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, + "wire_extends": 1, + "last_extend_additional_amount": 1000, + "engine_calls": [{ "credit_balance": 5000 }, { "credit_balance": 1050, "credit_cost": 100 }] + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 2000, "local_remaining_credits": 950 } + } + ] + }, + { + "name": "check_insufficient_lease_after_failed_extend_resolves_by_mode", + "description": "When the opportunistic extend fails and the retry is still refused, the check resolves through the failure mode (fail-closed here) with reason insufficient_lease_balance, leaving the lease balance untouched.", + "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": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 950, + "expect": { "balance": 50 } + }, + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "on_acquire_failure": "fail-closed", + "server": { "extend": { "error": "wire down" } }, + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + } + ], + "expect": { + "allowed": false, + "reason": "insufficient_lease_balance", + "err": "insufficient_lease_balance", + "has_reservation": false + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 50 } + }, + { "op": "reservation_count", "expect": { "count": 0 } } + ] + }, + { + "name": "check_falls_back_without_usable_credit_entitlement", + "description": "A non-credit matched entitlement (boolean/override/numeric/unlimited/not entitled) or an incomplete credit entitlement (no positive consumption rate) defers to the plain check: no lease traffic, no reservation.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": {} }, + "usage": 10, + "event_subtype": "inference_tokens", + "engine": [{ "value": true, "reason": "probe", "entitlement": { "value_type": "boolean" } }], + "expect": { "fallback_called": true, "reason": "fallback", "has_reservation": false } + }, + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 0, + "event_subtype": "inference_tokens" + } + } + ], + "expect": { "fallback_called": true, "reason": "fallback", "has_reservation": false } + }, + { "op": "reservation_count", "expect": { "count": 0 } } + ] + }, + { + "name": "check_zero_usage_falls_back", + "description": "usage = 0 means nothing to reserve: the check defers to the plain (preflight-threaded) check instead of issuing a no-op 0-credit reservation.", + "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, + "event_subtype": "inference_tokens", + "engine": [], + "expect": { "fallback_called": true, "reason": "fallback", "has_reservation": false } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + } + ] + }, + { + "name": "check_invalid_usage_resolves_statically_by_mode", + "description": "A negative (or non-finite) usage must never reach the stores; the check resolves statically by mode without any engine evaluation: deny for fail-closed, blanket allow for fail-open, reason invalid_usage either way.", + "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": -5, + "event_subtype": "inference_tokens", + "on_acquire_failure": "fail-closed", + "engine": [], + "expect": { + "allowed": false, + "reason": "invalid_usage", + "err": "invalid_usage", + "has_reservation": false, + "engine_calls": [] + } + }, + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": -5, + "event_subtype": "inference_tokens", + "on_acquire_failure": "fail-open", + "engine": [], + "expect": { + "allowed": true, + "reason": "invalid_usage_fail_open", + "err": "invalid_usage", + "has_reservation": false, + "engine_calls": [] + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + } + ] + } + ] +} diff --git a/conformance/vectors/crash-windows.json b/conformance/vectors/crash-windows.json new file mode 100644 index 0000000..c8ee05b --- /dev/null +++ b/conformance/vectors/crash-windows.json @@ -0,0 +1,271 @@ +{ + "category": "crash_window", + "vectors": [ + { + "name": "debit_without_record_leaks_bounded", + "description": "Crash window 1 (debit-then-add): the atomic debit landed but the reservation record never did. The leak is exactly the reserved amount; the sweeper can never refund a hold that was never recorded.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 0 } }, + { "op": "advance_clock", "ms": 55000 }, + { "op": "sweep_expired", "expect": { "swept": 0 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + } + ] + }, + { + "name": "debit_leak_reclaimed_at_lease_expiry", + "description": "Crash window 1 recovery: the leaked balance is never served after lease expiry, and the successor lease installs at the full grant — the leak does not outlive the lease.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 1, + "expect": { "balance": null } + }, + { + "op": "replace_lease", + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 180000, + "expect": { "written": true } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + } + ] + }, + { + "name": "retried_check_after_debit_leak_settles_once", + "description": "A retry after a window-1 crash is a fresh check with a fresh reservation: its unspent slice refunds exactly once; the leaked slice never refunds — not on a repeat consume, not on a sweep.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 3600000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 800 } + }, + { + "op": "add_reservation", + "id": "res_retry", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { "op": "consume_reservation", "id": "res_retry", "credits": 40, "expect": { "consumed": 40 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 860 } + }, + { "op": "consume_reservation", "id": "res_retry", "credits": 40, "expect": { "consumed": null } }, + { "op": "advance_clock", "ms": 70000 }, + { "op": "sweep_expired", "expect": { "swept": 0 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 860 } + } + ] + }, + { + "name": "crash_before_refund_claim_is_durable", + "description": "Crash window 2 (consume-then-refund): the claim survives the crash, so the reservation is gone everywhere and nothing can double-spend; the unspent slice's refund is lost, bounded by credits_reserved. A retried settle neither re-claims nor double-refunds, and the sweeper cannot refund a claimed reservation.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 3600000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { + "op": "consume_reservation", + "id": "res_1", + "credits": 30, + "crash_before_refund": true, + "expect": { "throws": true } + }, + { "op": "get_reservation", "id": "res_1", "expect": { "exists": false } }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 0 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + }, + { "op": "consume_reservation", "id": "res_1", "credits": 30, "expect": { "consumed": null } }, + { "op": "advance_clock", "ms": 70000 }, + { "op": "sweep_expired", "expect": { "swept": 0 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + } + ] + }, + { + "name": "crash_before_refund_reclaimed_at_lease_expiry", + "description": "Crash window 2 recovery: after the lease expires and a successor takes the slot at full grant, a very late retried settle of the crashed reservation must not leak the lost refund into the successor.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { + "op": "consume_reservation", + "id": "res_1", + "credits": 30, + "crash_before_refund": true, + "expect": { "throws": true } + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 1, + "expect": { "balance": null } + }, + { + "op": "replace_lease", + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 180000, + "expect": { "written": true } + }, + { "op": "consume_reservation", "id": "res_1", "credits": 0, "expect": { "consumed": null } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_2", "local_remaining_credits": 1000 } + } + ] + } + ] +} diff --git a/conformance/vectors/expiry.json b/conformance/vectors/expiry.json new file mode 100644 index 0000000..e1f7e2c --- /dev/null +++ b/conformance/vectors/expiry.json @@ -0,0 +1,198 @@ +{ + "category": "expiry", + "vectors": [ + { + "name": "expired_lease_never_serves_reserves", + "description": "Past its expiry a lease's balance is stale (the server refunded the grant): reserves are refused even with ample local balance.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 1, + "expect": { "balance": null } + } + ] + }, + { + "name": "successor_after_expiry_restores_full_grant", + "description": "A successor lease installed over an expired slot starts at its full grant — nothing from the expired lease (debits or leaks) carries over.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 700, + "expect": { "balance": 300 } + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "replace_lease", + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 180000, + "expect": { "written": true } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_2", "local_remaining_credits": 1000 } + } + ] + }, + { + "name": "sweep_refunds_expired_holds_only", + "description": "The sweeper refunds an expired reservation's full hold to its lease and leaves live reservations untouched.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 300, + "expect": { "balance": 700 } + }, + { + "op": "add_reservation", + "id": "res_short", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 10000 + }, + { + "op": "add_reservation", + "id": "res_long", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 20, + "credits_reserved": 200, + "consumption_rate": 10, + "expires_at_ms": 200000 + }, + { "op": "sweep_expired", "expect": { "swept": 0 } }, + { "op": "advance_clock", "ms": 10001 }, + { "op": "sweep_expired", "expect": { "swept": 1 } }, + { "op": "get_reservation", "id": "res_short", "expect": { "exists": false } }, + { "op": "get_reservation", "id": "res_long", "expect": { "exists": true } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 800 } + }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 200 } } + ] + }, + { + "name": "sweep_of_stale_lease_hold_does_not_inflate_successor", + "description": "Sweeping (or consuming) a reservation carved from an expired lease refunds nothing into the successor lease occupying the slot: the hold is pinned to its originating lease_id.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_stale", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 90000 + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "replace_lease", + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000, + "expect": { "written": true } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 200, + "expect": { "balance": 800 } + }, + { "op": "advance_clock", "ms": 30000 }, + { "op": "sweep_expired", "expect": { "swept": 1 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_2", "local_remaining_credits": 800 } + } + ] + } + ] +} 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/conformance/vectors/lease-lifecycle.json b/conformance/vectors/lease-lifecycle.json new file mode 100644 index 0000000..c9f5db1 --- /dev/null +++ b/conformance/vectors/lease-lifecycle.json @@ -0,0 +1,497 @@ +{ + "category": "lease_lifecycle", + "vectors": [ + { + "name": "replace_installs_full_grant", + "description": "A fresh install initializes local_remaining_credits to the full granted amount.", + "operations": [ + { + "op": "replace_lease", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000, + "expect": { "written": true } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { + "exists": true, + "lease_id": "lse_1", + "granted_amount": 1000, + "local_remaining_credits": 1000 + } + } + ] + }, + { + "name": "replace_keeps_live_lease_even_with_different_id", + "description": "A live lease occupying the slot wins over any replace — even one carrying a different lease_id (a sibling raced this acquire). Its already-debited balance is preserved.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 400, + "expect": { "balance": 600 } + }, + { + "op": "replace_lease", + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 5000, + "expires_at_ms": 120000, + "expect": { "written": false } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { + "exists": true, + "lease_id": "lse_1", + "granted_amount": 1000, + "local_remaining_credits": 600 + } + } + ] + }, + { + "name": "replace_reconciles_expired_slot_with_same_id", + "description": "A stale acquire response for the SAME lease landing over its own expired local row must not reinstall it: that would reset local_remaining_credits and erase debits whose reservations are still open. The row is reconciled like an extend (granted to total, expiry forward, balance untouched) and reported as kept.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 400, + "expect": { "balance": 600 } + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "replace_lease", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1500, + "expires_at_ms": 120000, + "expect": { "written": false } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { + "exists": true, + "lease_id": "lse_1", + "granted_amount": 1500, + "local_remaining_credits": 1100 + } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 1000 } + } + ] + }, + { + "name": "replace_overwrites_expired_lease", + "description": "An expired lease does not block the slot: replace overwrites it atomically and restores the full new grant.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 400, + "expect": { "balance": 600 } + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "replace_lease", + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 180000, + "expect": { "written": true } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_2", "local_remaining_credits": 1000 } + } + ] + }, + { + "name": "try_reserve_insufficient_and_boundary", + "description": "A reserve larger than the remaining balance is refused and touches nothing; reserving down to exactly zero is allowed; negative amounts are always refused.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 100, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 101, + "expect": { "balance": null } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 100 } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": -1, + "expect": { "balance": null } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 0 } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 1, + "expect": { "balance": null } + }, + { + "op": "try_reserve", + "company_id": "co_9", + "credit_type_id": "ct_1", + "credits": 1, + "expect": { "balance": null } + } + ] + }, + { + "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.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { "op": "refund_lease", "company_id": "co_1", "credit_type_id": "ct_1", "credits": 500 }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + } + ] + }, + { + "name": "refund_pinned_to_lease_id_dropped_on_successor", + "description": "A refund pinned to an expired lease's id must not inflate the successor lease occupying the slot — the expired lease's remainder was already returned to the company balance server-side. An unpinned refund still applies.", + "given": { + "leases": [ + { + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 200, + "expect": { "balance": 800 } + }, + { + "op": "refund_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "pin_lease_id": "lse_1" + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 800 } + }, + { + "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, "local_remaining_credits": 900 } + } + ] + }, + { + "name": "extend_reconciles_to_total_and_converges", + "description": "Extend applies the server-authoritative TOTAL: the delta is computed against the stored total, so a repeated or stale-lower total is a no-op and out-of-order applies converge without minting credits.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 800, + "expect": { "balance": 200 } + }, + { + "op": "extend_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_total": 3000, + "expires_at_ms": 300000 + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 3000, "local_remaining_credits": 2200 } + }, + { + "op": "extend_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_total": 3000, + "expires_at_ms": 300000 + }, + { + "op": "extend_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_total": 2000, + "expires_at_ms": 300000 + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 3000, "local_remaining_credits": 2200 } + } + ] + }, + { + "name": "extend_expiry_only_moves_forward", + "description": "An extend carrying an earlier expiry must not shorten the lease: after an extend to a later expiry, a stale out-of-order apply with an earlier expiry leaves the lease live past that earlier instant.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "extend_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_total": 2000, + "expires_at_ms": 120000 + }, + { + "op": "extend_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_total": 2000, + "expires_at_ms": 30000 + }, + { "op": "advance_clock", "ms": 90000 }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 1900 } + } + ] + }, + { + "name": "extend_pinned_to_lease_id_dropped_on_successor", + "description": "An extend pinned to a lease the slot no longer holds is dropped entirely — crediting the successor would mint credits the server granted to the expired lease.", + "given": { + "leases": [ + { + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "extend_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_total": 5000, + "expires_at_ms": 600000, + "pin_lease_id": "lse_1" + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 1000, "local_remaining_credits": 1000 } + }, + { + "op": "extend_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_total": 2000, + "expires_at_ms": 600000, + "pin_lease_id": "lse_2" + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 2000, "local_remaining_credits": 2000 } + } + ] + } + ] +} diff --git a/conformance/vectors/lease-manager.json b/conformance/vectors/lease-manager.json new file mode 100644 index 0000000..79ab1af --- /dev/null +++ b/conformance/vectors/lease-manager.json @@ -0,0 +1,417 @@ +{ + "category": "lease_manager", + "vectors": [ + { + "name": "acquire_installs_tranche_and_reuses_live_lease", + "description": "First acquire requests lease_size from the server and installs the response at full grant; a second acquire while the lease is live makes no wire call.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "acquire_if_needed", + "company_id": "co_1", + "credit_type_id": "ct_1", + "server": { + "lease": { "lease_id": "lse_1", "granted_amount": 1000, "expires_at_ms": 300000 } + }, + "expect": { + "lease_id": "lse_1", + "wire_acquires": 1, + "last_acquire_requested_amount": 1000, + "released_lease_ids": [] + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_1", "local_remaining_credits": 1000 } + }, + { + "op": "acquire_if_needed", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "lease_id": "lse_1", "wire_acquires": 1 } + } + ] + }, + { + "name": "acquire_replaces_expired_slot_without_release", + "description": "An expired slot triggers a fresh acquire that supplants the stale entry in place; the redundant-lease release path must not fire (replace wrote, it did not keep a live lease).", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_stale", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { "op": "advance_clock", "ms": 60001 }, + { + "op": "acquire_if_needed", + "company_id": "co_1", + "credit_type_id": "ct_1", + "server": { + "lease": { "lease_id": "lse_fresh", "granted_amount": 1000, "expires_at_ms": 360000 } + }, + "expect": { "lease_id": "lse_fresh", "wire_acquires": 1, "released_lease_ids": [] } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_fresh", "local_remaining_credits": 1000 } + } + ] + }, + { + "name": "lost_acquire_race_different_id_releases_redundant_lease", + "description": "A sibling installs a live lease while this acquire's wire call is in flight. The installed lease (with its debited balance) wins; the lease the server minted for the loser is redundant and gets released so it is not orphaned against the company balance.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "acquire_if_needed", + "company_id": "co_1", + "credit_type_id": "ct_1", + "server": { + "lease": { "lease_id": "lse_loser", "granted_amount": 1000, "expires_at_ms": 300000 } + }, + "install_during_wire": { + "lease_id": "lse_winner", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + }, + "expect": { "lease_id": "lse_winner", "wire_acquires": 1, "released_lease_ids": ["lse_loser"] } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_winner" } + } + ] + }, + { + "name": "lost_acquire_race_same_id_releases_nothing", + "description": "The server is idempotent for an active slot: a racing acquire is handed back the SAME lease the sibling installed. There is nothing to release — releasing would pull the shared lease out from under every sibling.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "acquire_if_needed", + "company_id": "co_1", + "credit_type_id": "ct_1", + "server": { + "lease": { "lease_id": "lse_shared", "granted_amount": 1000, "expires_at_ms": 300000 } + }, + "install_during_wire": { + "lease_id": "lse_shared", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + }, + "expect": { "lease_id": "lse_shared", "wire_acquires": 1, "released_lease_ids": [] } + } + ] + }, + { + "name": "extend_triggered_at_low_water_mark_requests_tranche", + "description": "At or below the low-water-mark ratio a steady-state extend fires, requesting the configured tranche (lease_size) and reconciling the local row to the server total.", + "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": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 800, + "expect": { "balance": 200 } + }, + { + "op": "maybe_extend", + "company_id": "co_1", + "credit_type_id": "ct_1", + "server": { "lease": { "granted_total": 2000, "expires_at_ms": 600000 } }, + "expect": { + "wire_extends": 1, + "last_extend_additional_amount": 1000, + "last_extend_lease_id": "lse_1" + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 2000, "local_remaining_credits": 1200 } + } + ] + }, + { + "name": "extend_triggered_by_required_credits_above_watermark", + "description": "Above the watermark no steady-state extend fires; a required_credits hint larger than the local remaining triggers one anyway (a check just failed a reserve of that size).", + "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": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "maybe_extend", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "wire_extends": 0 } + }, + { + "op": "maybe_extend", + "company_id": "co_1", + "credit_type_id": "ct_1", + "required_credits": 1500, + "server": { "lease": { "granted_total": 2000, "expires_at_ms": 600000 } }, + "expect": { "wire_extends": 1, "last_extend_additional_amount": 1000 } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 2000, "local_remaining_credits": 1900 } + } + ] + }, + { + "name": "extend_sized_to_shortfall_when_larger_than_tranche", + "description": "additional_amount = max(lease_size, required_credits - local_remaining): a single request larger than remaining + tranche must extend by the shortfall, or its post-extend retry would fail forever regardless of server balance.", + "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": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "maybe_extend", + "company_id": "co_1", + "credit_type_id": "ct_1", + "required_credits": 5000, + "server": { "lease": { "granted_total": 5100, "expires_at_ms": 600000 } }, + "expect": { "wire_extends": 1, "last_extend_additional_amount": 4100 } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 5100, "local_remaining_credits": 5000 } + } + ] + }, + { + "name": "never_extend_an_expired_lease", + "description": "An expired lease is released as far as the server is concerned — the only correct move is a fresh acquire, never an extend, no matter how depleted the balance.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_old", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 900, + "expect": { "balance": 100 } + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "maybe_extend", + "company_id": "co_1", + "credit_type_id": "ct_1", + "required_credits": 1500, + "expect": { "wire_extends": 0 } + } + ] + }, + { + "name": "wire_failures_resolve_to_no_lease_without_state_changes", + "description": "A failed acquire yields no lease and installs nothing; a failed extend leaves the local row untouched. Neither throws (both are routed through fail-open/fail-closed by callers).", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "acquire_if_needed", + "company_id": "co_1", + "credit_type_id": "ct_1", + "server": { "error": "wire down" }, + "expect": { "lease_id": null, "wire_acquires": 1 } + }, + { "op": "get_lease", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "exists": false } }, + { + "op": "replace_lease", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000, + "expect": { "written": true } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 800, + "expect": { "balance": 200 } + }, + { + "op": "maybe_extend", + "company_id": "co_1", + "credit_type_id": "ct_1", + "server": { "error": "wire down" }, + "expect": { "wire_extends": 1 } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 1000, "local_remaining_credits": 200 } + } + ] + }, + { + "name": "release_all_releases_live_and_skips_expired", + "description": "On close, a per-process store releases its live leases over the wire (returning remainders immediately) and drops them locally; expired leases are skipped — the server already swept them. Only valid for an exclusively-owned (in-memory) store; a shared backend must never enumerate-and-release.", + "backends": ["in_memory"], + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_live", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + }, + { + "lease_id": "lse_expired", + "company_id": "co_2", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 30000 + } + ] + }, + "operations": [ + { "op": "advance_clock", "ms": 30001 }, + { + "op": "release_all_local_leases", + "expect": { "released_lease_ids": ["lse_live"] } + }, + { "op": "get_lease", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "exists": false } } + ] + } + ] +} diff --git a/conformance/vectors/reservation-lifecycle.json b/conformance/vectors/reservation-lifecycle.json new file mode 100644 index 0000000..eea6b50 --- /dev/null +++ b/conformance/vectors/reservation-lifecycle.json @@ -0,0 +1,353 @@ +{ + "category": "reservation_lifecycle", + "vectors": [ + { + "name": "consume_exact_usage_no_refund", + "description": "Consuming exactly the reserved amount removes the reservation and refunds nothing.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { "op": "consume_reservation", "id": "res_1", "credits": 100, "expect": { "consumed": 100 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + }, + { "op": "get_reservation", "id": "res_1", "expect": { "exists": false } } + ] + }, + { + "name": "consume_under_reserved_refunds_unspent", + "description": "Consuming less than reserved refunds the unspent slice to the lease in the same step.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { "op": "consume_reservation", "id": "res_1", "credits": 30, "expect": { "consumed": 30 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 970 } + } + ] + }, + { + "name": "consume_over_reserved_clamps_to_hold", + "description": "Local consumption is clamped to credits_reserved: over-use consumes the full hold, refunds nothing, and never debits the lease beyond the reservation. (The billed Track quantity is NOT clamped — see track-settle vectors.)", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { "op": "consume_reservation", "id": "res_1", "credits": 999, "expect": { "consumed": 100 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + } + ] + }, + { + "name": "consume_zero_cancels_with_full_refund", + "description": "Consuming 0 credits acts as a cancel: the full hold is refunded. Negative consumption clamps to 0 the same way.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { "op": "consume_reservation", "id": "res_1", "credits": 0, "expect": { "consumed": 0 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + } + ] + }, + { + "name": "consume_is_exactly_once", + "description": "A missing reservation and a second consume of the same id both return null and refund nothing — the claim is the exactly-once arbiter.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { "op": "consume_reservation", "id": "res_missing", "credits": 10, "expect": { "consumed": null } }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { "op": "consume_reservation", "id": "res_1", "credits": 30, "expect": { "consumed": 30 } }, + { "op": "consume_reservation", "id": "res_1", "credits": 30, "expect": { "consumed": null } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 970 } + } + ] + }, + { + "name": "reserved_credits_sums_open_holds_per_slot", + "description": "reserved_credits sums credits_reserved across open reservations for the exact (company, credit) slot only, and a hold stops counting the moment it is consumed.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 350, + "expect": { "balance": 650 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { + "op": "add_reservation", + "id": "res_2", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 25, + "credits_reserved": 250, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { + "op": "add_reservation", + "id": "res_other_credit", + "lease_id": "lse_9", + "company_id": "co_1", + "credit_type_id": "ct_2", + "event_subtype": "inference_tokens", + "quantity_reserved": 99, + "credits_reserved": 999, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { + "op": "add_reservation", + "id": "res_other_company", + "lease_id": "lse_8", + "company_id": "co_2", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 99, + "credits_reserved": 999, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { + "op": "reserved_credits", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "total": 350 } + }, + { + "op": "reserved_credits", + "company_id": "co_1", + "credit_type_id": "ct_2", + "expect": { "total": 999 } + }, + { "op": "reserved_credits", "company_id": "co_9", "credit_type_id": "ct_1", "expect": { "total": 0 } }, + { "op": "consume_reservation", "id": "res_1", "credits": 40, "expect": { "consumed": 40 } }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 250 } } + ] + }, + { + "name": "fractional_credit_amounts_are_exact", + "description": "Fractional consumption rates produce fractional holds; reserve, consume, and refund arithmetic must not truncate.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 10, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 2.5, + "expect": { "balance": 7.5 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 25, + "credits_reserved": 2.5, + "consumption_rate": 0.1, + "expires_at_ms": 60000 + }, + { "op": "consume_reservation", "id": "res_1", "credits": 1.5, "expect": { "consumed": 1.5 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 8.5 } + } + ] + } + ] +} diff --git a/conformance/vectors/track-settle.json b/conformance/vectors/track-settle.json new file mode 100644 index 0000000..1da71a0 --- /dev/null +++ b/conformance/vectors/track-settle.json @@ -0,0 +1,194 @@ +{ + "category": "track_settle", + "vectors": [ + { + "name": "track_underuse_settles_and_refunds_unspent", + "description": "Settling with less than the reserved usage consumes actual x rate, refunds the unspent slice to the lease, and emits a Track billing the ACTUAL quantity keyed to the lease.", + "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": 10, + "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 } + }, + { + "op": "track", + "handle": "r1", + "actual_quantity": 4, + "expect": { + "settled_locally": true, + "track": { "event": "inference_tokens", "quantity": 4, "lease_id": "lse_1" } + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 960 } + }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 0 } } + ] + }, + { + "name": "track_overuse_bills_actual_but_clamps_local_debit", + "description": "Actual usage above the reservation: the LOCAL settle clamps consumption to the reserved hold (the lease is never debited past the reservation), but the Track event bills the unclamped actual quantity — the server is the source of truth for real consumption.", + "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": 10, + "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 } + }, + { + "op": "track", + "handle": "r1", + "actual_quantity": 25, + "expect": { + "settled_locally": true, + "track": { "event": "inference_tokens", "quantity": 25, "lease_id": "lse_1" } + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + } + ] + }, + { + "name": "track_after_sweep_is_a_recovery_emit", + "description": "Work outliving the reservation TTL: the sweeper already refunded the full hold, so the late settle does not touch the lease (the local balance reads high until rollover) — but the Track is still emitted so the server bills the actual usage. Server-side idempotency (a deterministic key derived from the reservation id) is what keeps a racing normal emit from double-billing.", + "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": 10, + "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 } + }, + { "op": "advance_clock", "ms": 60001 }, + { "op": "sweep_expired", "expect": { "swept": 1 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + }, + { + "op": "track", + "handle": "r1", + "actual_quantity": 4, + "expect": { + "settled_locally": false, + "track": { "event": "inference_tokens", "quantity": 4, "lease_id": "lse_1" } + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + } + ] + } + ] +} diff --git a/custom.gemspec.rb b/custom.gemspec.rb index 853485b..441f5e5 100644 --- a/custom.gemspec.rb +++ b/custom.gemspec.rb @@ -3,6 +3,12 @@ def add_custom_gemspec_data(spec) spec.homepage = "https://github.com/SchematicHQ/schematic-ruby" + # Cross-SDK test fixtures, read only by test/conformance_test.rb, which the + # generated gemspec already drops with the rest of test/. Rejected here rather + # than there because that file is generated; npm's allowlist keeps them out of + # schematic-node the same way. + spec.files.reject! { |f| f.start_with?("conformance/") } + # The WASM binary is gitignored (downloaded at build time via scripts/download-wasm.sh) # but must be included in the published gem. Append it to the file list. wasm_file = "lib/schematic/wasm/rulesengine.wasm" diff --git a/lib/schematic/credits/leases/check.rb b/lib/schematic/credits/leases/check.rb new file mode 100644 index 0000000..15adc54 --- /dev/null +++ b/lib/schematic/credits/leases/check.rb @@ -0,0 +1,477 @@ +# frozen_string_literal: true + +require "securerandom" + +module Schematic + module Credits + module Leases + # Everything a lease-bearing check needs. enqueue_flag_check_event reports + # a flag_check event for a check the lease path resolved itself, mirroring + # the plain check paths so lease-gated checks stay visible to flag-check + # analytics and company last-seen. Fallback exits do not call it: the + # plain check they delegate to enqueues its own. + CheckDeps = Struct.new(:lease_store, :reservations, :manager, :datastream, :logger, :clock, + :enqueue_flag_check_event, keyword_init: true) + + # Read a field from an engine payload, which carries camelCase keys inside + # the entitlement even though the top-level result is normalized. + def self.field(hash, *names) + return nil if hash.nil? + + names.each do |name| + value = hash[name] || hash[name.to_s] + return value unless value.nil? + end + nil + end + + # The preflight envelope for a client-side rules evaluation. With an + # event_subtype the quantity goes out as the event_usage pair so the + # engine matches it to that subtype's condition; without one it goes out + # as the generic usage knob. Public so the plain check path can thread the + # same preflight through when the lease path cannot run. + def self.build_preflight_options(options) + usage = options[:usage] + return nil if usage.nil? + + # The preflight quantity is an integer on both seams (the engine + # envelope and the API's preflight body), and it asks an upper-bound + # question, so a fractional usage rounds up rather than gating on less + # usage than the operation is about to record. + quantity = Leases.wire_quantity(usage) + # A zero usage has no effect server-side, so sending a preflight for one + # would only cost the check its flag cache. A zero credit_cost would be + # different, saying free rather than absent, but this never emits one. + return nil if quantity.zero? + + if options[:event_subtype] + { event_usage: { event_subtype: options[:event_subtype], quantity: quantity } } + else + { usage: quantity } + end + end + + # Drive a single lease-gated check. + # + # 1. Probe the engine once against the company's real balance, with no + # substitution and no preflight, and read the matched entitlement. A + # non-credit entitlement means there is nothing to lease, so defer to + # the plain check. + # 2. Acquire (or reuse) a lease for (company, credit id). + # 3. Reserve ceil(usage) x consumption_rate from it, atomically. + # 4. Re-run the engine against a company snapshot whose balance for that + # credit is the PRE-reservation local balance, with credit_cost set, so + # the engine evaluates the same arithmetic try_reserve just enforced. + # The hold only sticks if the engine allows. + def self.check_with_lease(deps, key, eval_ctx, options, &fallback) + Check.new(deps, key, eval_ctx, options, fallback).run + end + + # The check flow, as an object so its steps can pass state without + # threading a dozen arguments through every helper. + class Check + def initialize(deps, key, eval_ctx, options, fallback) + @deps = deps + @key = key + @eval_ctx = eval_ctx || {} + @options = options || {} + @fallback = fallback + @logger = deps.logger + @clock = deps.clock || DEFAULT_CLOCK + @on_failure = Leases.resolve_failure_mode(@options[:on_acquire_failure], @logger) + # Fixed now, not at each wait: a deadline taken when a join starts + # would let a check spend its acquire and reserve time and then its + # whole timeout again behind someone else's extend. + @deadline = Leases.join_deadline(request_options) + end + + def run + guard = check_guards + return guard if guard + + resolved = resolve_entitlement + return @fallback.call if resolved.nil? + + @credit_id, @consumption_rate, @event_subtype = resolved + # 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. + # Rounded through wire_quantity, the one the Track event's quantity + # goes through, so the hold, the local debit and the billed figure + # cannot disagree over a float that is a hair above a whole unit. + @credit_cost = Leases.wire_quantity(@options[:usage]) * @consumption_rate + + lease = @deps.manager.acquire_if_needed(@company[:id], @credit_id, request_options, deadline: @deadline) + return failure("lease_acquire_failed") if lease.nil? + + reserve = reserve_credits + return reserve if reserve.is_a?(CheckResult) + + reservation = register_reservation(reserve) + persisted = persist(reservation, reserve) + return persisted if persisted.is_a?(CheckResult) + + gate(reservation, reserve) + end + + private + + # Guards, in order: a malformed usage never reaches the stores; zero + # usage has nothing to reserve; and without a datastream, a cached flag, + # or a resolvable company there is no local evaluation to gate with. + def check_guards + usage = @options[:usage] + # NaN slips through every numeric comparison, so a single NaN debit + # would poison the (possibly shared) lease balance into approving + # every later reserve. The stores guard too, but resolve it here + # through the caller's failure contract rather than letting it surface + # as an opaque reserve failure. + unless Leases.valid_quantity?(usage) + @logger.error( + "Lease check: invalid usage #{usage.inspect} for flag #{@key}, must be a finite non-negative number" + ) + return emit(static_failure_result("invalid_usage", nil)) + end + + if usage.zero? + @logger.debug("Lease check: usage is 0 for flag #{@key}, nothing to reserve, using plain check") + return @fallback.call + end + + return @fallback.call if datastream_unavailable? + return @fallback.call if load_flag.nil? + return @fallback.call unless entities_resolved? + + nil + end + + def datastream_unavailable? + return false if @deps.datastream + + @logger.debug("Credit-lease check requested without datastream, falling back to plain check") + true + end + + def load_flag + @flag = begin + @deps.datastream.get_flag(@key) + rescue StandardError => e + @logger.warn("Lease check: failed to load flag #{@key}: #{e.message}") + nil + end + @logger.debug("Lease check: no cached flag for #{@key}, falling back") if @flag.nil? + @flag + end + + # Resolve company and user the way a plain datastream check does. An + # evaluation with a missing entity is not an option: a nil user would + # silently skip user-targeted rules and overrides, so a named entity + # that cannot be resolved falls back to the plain check, which has its + # own degradation story. + def entities_resolved? + company_keys = @eval_ctx[:company] || @eval_ctx["company"] + if company_keys.nil? || company_keys.empty? + @logger.debug("Lease check: no company on eval context, falling back") + return false + end + @company = fetch_entity("company") { @deps.datastream.get_company(company_keys) } + return false if @company.nil? + + user_keys = @eval_ctx[:user] || @eval_ctx["user"] + return true if user_keys.nil? || user_keys.empty? + + @user = fetch_entity("user") { @deps.datastream.get_user(user_keys) } + !@user.nil? + end + + def fetch_entity(kind) + yield + rescue StandardError => e + @logger.debug("Lease check: #{kind} fetch failed (#{e.message}), falling back") + nil + end + + # One probe against the real balance surfaces the matched entitlement, + # which names the credit directly and lets a non-credit grant skip the + # lease round trip entirely. The probe omits preflight on purpose: + # charging a cost against the lease-depleted server balance could fail + # the credit condition, drop the engine to a lower-priority rule, and + # hide the very entitlement being identified. + def resolve_entitlement + probe = probe_entitlement + return nil if probe.nil? + + entitlement = probe[:entitlement] + value_type = Leases.field(entitlement, :valueType, :value_type) + unless value_type == "credit" + # A boolean or override grant, a numeric allocation, unlimited, or + # simply not entitled. The feature resolves without drawing a + # credit, so skip the lease round-trip and let the plain check, + # which is preflight-aware, decide. + @logger.debug( + "Lease check: flag #{@key} matched a non-credit entitlement " \ + "(value_type=#{value_type || ""}), falling back to plain check, no reservation" + ) + return nil + end + + credit_id = Leases.field(entitlement, :creditId, :credit_id) + consumption_rate = (Leases.field(entitlement, :consumptionRate, :consumption_rate) || 0).to_f + # The caller's explicit subtype wins; otherwise the entitlement names + # the metered event. The reservation settles into a track event named + # by this subtype, so a credit entitlement with neither a resolvable + # subtype nor a positive rate cannot be billed and is ungateable. + subtype = @options[:event_subtype] || Leases.field(entitlement, :eventSubtype, :event_subtype) + if credit_id.nil? || consumption_rate <= 0 || subtype.nil? + @logger.debug( + "Lease check: flag #{@key} credit entitlement is incomplete " \ + "(credit_id=#{credit_id || ""}, consumption_rate=#{consumption_rate}, " \ + "subtype=#{subtype || ""}), falling back" + ) + return nil + end + + [credit_id, consumption_rate, subtype] + end + + def probe_entitlement + evaluate(@company, nil) + rescue StandardError => e + # The probe is a resolution step, not the gate, so a failure means the + # credit could not be resolved. Defer to the plain check rather than + # hard-denying. No reservation exists yet, so nothing to cancel. + @logger.warn("Lease check: entitlement probe failed for flag #{@key} (#{e.message}), falling back") + nil + end + + # try_reserve is the atomic gate: check and debit in one step, returning + # the post-debit balance (so the pre-debit figure follows without a + # second store read) AND the id of the lease it charged. That id, not + # the acquired one, is what the reservation is pinned to: the debit is + # not keyed by lease, so the slot's lease may have been replaced since + # the acquire, and the window spans the extend awaited below. + def reserve_credits + reserve = @deps.lease_store.try_reserve(@company[:id], @credit_id, @credit_cost) + if reserve.nil? + # The lease has less than credit_cost left locally. Passing + # credit_cost extends even when the ratio is still above the low + # water mark, which a single large request needs. + @deps.manager.maybe_extend_in_background(@company[:id], @credit_id, @credit_cost, + request_options, deadline: @deadline)&.join + reserve = @deps.lease_store.try_reserve(@company[:id], @credit_id, @credit_cost) + end + return failure("insufficient_lease_balance") if reserve.nil? + + reserve + rescue StandardError => e + @logger.error("Lease check: reserve against #{@company[:id]}/#{@credit_id} failed: #{e.message}") + failure("lease_store_error") + end + + def register_reservation(reserve) + resolved = @deps.manager.resolve_config(@credit_id) + Reservation.new( + id: SecureRandom.uuid, + # The lease the debit actually landed on, which may not be the one + # the acquire 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: @company[:id], + credit_type_id: @credit_id, + event_subtype: @event_subtype, + quantity_reserved: @options[:usage], + credits_reserved: @credit_cost, + consumption_rate: @consumption_rate, + expires_at: @clock.call + (resolved.reservation_ttl_ms / 1000.0), + eval_ctx: @eval_ctx + ) + end + + # Recorded between the debit and the gate, so a crash leaves a sweepable + # reservation rather than credits stranded until the lease expires. The + # window it cannot cover is the gap before this call, which has no I/O + # in it and leaks at most credit_cost until that expiry. + def persist(reservation, reserve) + @deps.reservations.add(reservation) + nil + rescue StandardError => e + @logger.error("Lease check: failed to persist reservation #{reservation.id}: #{e.message}") + undo_debit(reservation, reserve) + failure("lease_store_error") + end + + # Undo the local debit so the credits are not stranded until lease + # expiry. consume claims whatever slice of the add made it to the store + # and refunds it; if nothing was persisted, refund the debit directly. + # Both are pinned to the lease the debit landed on. If even the undo + # fails, accept the bounded leak: the slice is reclaimed when the lease + # expires server-side, which beats risking a double refund. + def undo_debit(reservation, reserve) + undone = @deps.reservations.consume(reservation.id, 0) + @deps.lease_store.refund(@company[:id], @credit_id, @credit_cost, reserve.lease_id) if undone.nil? + rescue StandardError => e + @logger.warn( + "Lease check: could not undo local debit for #{reservation.id} (#{e.message}); " \ + "the slice is reclaimed at lease expiry" + ) + end + + # The engine gates against the lease's local view, not the server's + # balance, so it re-checks the arithmetic try_reserve just enforced plus + # every non-credit rule. The pre-reservation figure is the atomic + # reserve's own post-debit balance plus what it debited, so it is exact + # as of the debit with no read race. + def gate(reservation, reserve) + pre_reservation = reserve.balance + @credit_cost + substituted = substitute_credit_balance(@company, @credit_id, pre_reservation) + begin + result = evaluate(substituted, { credit_cost: { @credit_id => @credit_cost } }) + rescue StandardError => e + @logger.error("Lease check: rules engine evaluation failed: #{e.message}") + # Cancel the hold, then resolve the mode statically: the engine + # itself just failed, so a fail-open re-evaluation is impossible. + cancel_reservation(reservation) + return emit(static_failure_result("wasm_error: #{e.message}", @flag), engine_ids(nil)) + end + + ids = engine_ids(result) + # A nil verdict denies here rather than standing in the caller's + # default, which is what the plain check path does with the same nil. + # The difference is deliberate: this branch holds credits, and the + # default exists to answer a flag nothing evaluated, not to release a + # hold the engine declined to approve. + unless result[:value] + cancel_reservation(reservation) + return emit( + CheckResult.new( + allowed: false, value: false, reason: result[:reason] || "denied_by_engine", + entitlement: result[:entitlement], flag_key: result[:flag_key] || @key, flag_id: result[:flag_id] + ), ids + ) + end + + # The engine allowed against the substituted lease balance, so the + # hold stays. No rule-match disambiguation is needed: the probe + # already established that this company's matched entitlement is the + # credit one, so an override or boolean grant would have skipped the + # reserve path. Bumping only the credit balance cannot make a + # different rule match here. + + # Fire and forget the low-water-mark refresh now that we have debited. + @deps.manager.maybe_extend_in_background(@company[:id], @credit_id) + + emit( + CheckResult.new( + allowed: true, value: true, reservation: reservation, + reason: result[:reason] || "lease_reserved", entitlement: result[:entitlement], + flag_key: result[:flag_key] || @key, flag_id: result[:flag_id] + ), ids + ) + end + + # Best-effort cancel: claims the record and refunds its full hold. + def cancel_reservation(reservation) + @deps.reservations.consume(reservation.id, 0) + rescue StandardError => e + @logger.warn( + "Lease check: failed to cancel reservation #{reservation.id} (#{e.message}); " \ + "its hold is reclaimed by the sweeper or at lease expiry" + ) + end + + # Every can't-gate outcome funnels through here. fail-open means assume + # the credits are there, NOT skip evaluation: the engine still runs with + # the balance substituted to an effectively unlimited value, so a company + # that is not entitled stays denied even with the lease backend down. + # Only if that evaluation itself fails does this fall back to a blanket + # allow. + def failure(reason) + result = + if @on_failure == :fail_closed + static_failure_result(reason, @flag) + else + fail_open_result(reason) + end + emit(result, { company_id: @company&.dig(:id), user_id: @user&.dig(:id) }) + end + + def fail_open_result(reason) + substituted = substitute_credit_balance(@company, @credit_id, FAIL_OPEN_BALANCE) + result = evaluate(substituted, Leases.build_preflight_options(@options)) + CheckResult.new( + allowed: result[:value], value: result[:value], + reason: "#{result[:reason] || "evaluated"} (#{reason}_fail_open)", + entitlement: result[:entitlement], flag_key: result[:flag_key] || @key, + flag_id: result[:flag_id] || @flag&.dig(:id), error: reason + ) + rescue StandardError => e + @logger.warn("Lease check: fail-open evaluation failed (#{e.message}); allowing") + static_failure_result(reason, @flag) + end + + # A mode resolved without an engine evaluation: deny for fail-closed, + # blanket allow for fail-open. Used when the engine itself is the thing + # that failed, and as the fallback when a fail-open evaluation errors. + def static_failure_result(reason, flag) + if @on_failure == :fail_closed + CheckResult.new(allowed: false, value: false, reason: reason, flag_key: @key, + flag_id: flag&.dig(:id), error: reason) + else + CheckResult.new(allowed: true, value: true, reason: "#{reason}_fail_open", flag_key: @key, + flag_id: flag&.dig(:id), error: reason) + end + end + + def evaluate(company, options) + @deps.datastream.check_flag_with_options(@flag, company, @user, options) + end + + def substitute_credit_balance(company, credit_id, balance) + substituted = company.dup + balances = (company[:credit_balances] || company["credit_balances"] || {}).dup + # The cache symbolizes keys, so a balance may be filed under either + # spelling. Replace whichever is there so the engine sees one value. + balances.delete(credit_id.to_sym) + balances.delete(credit_id.to_s) + balances[credit_id] = balance + substituted[:credit_balances] = balances + substituted + end + + def engine_ids(result) + { + company_id: (result && result[:company_id]) || @company&.dig(:id), + user_id: (result && result[:user_id]) || @user&.dig(:id), + rule_id: result && result[:rule_id] + } + end + + # Thread the caller's per-check timeout to the lease wire calls the same + # way the fallback path threads it to a plain check. + def request_options + return {} if @options[:timeout_ms].nil? + + { timeout_in_seconds: @options[:timeout_ms] / 1000.0 } + end + + def emit(result, ids = {}) + @deps.enqueue_flag_check_event&.call( + flag_key: result.flag_key, + value: result.value, + reason: result.reason, + error: result.error, + flag_id: result.flag_id, + company_id: ids[:company_id], + user_id: ids[:user_id], + rule_id: ids[:rule_id], + req_company: @eval_ctx[:company] || @eval_ctx["company"], + req_user: @eval_ctx[:user] || @eval_ctx["user"] + ) + result + end + end + end + end +end diff --git a/lib/schematic/credits/leases/lease_manager.rb b/lib/schematic/credits/leases/lease_manager.rb new file mode 100644 index 0000000..7b75c30 --- /dev/null +++ b/lib/schematic/credits/leases/lease_manager.rb @@ -0,0 +1,565 @@ +# frozen_string_literal: true + +require "securerandom" + +module Schematic + module Credits + module Leases + # The monotonic-clock millisecond deadline a caller's timeout sets for + # waiting on a shared flight, or nil for no cap. Taken once, at check + # start, so every wait the check makes draws on the same budget. + def self.join_deadline(request_options) + seconds = request_options.is_a?(Hash) ? request_options[:timeout_in_seconds] : nil + return nil unless seconds.is_a?(Numeric) && seconds.to_f.finite? + + (Process.clock_gettime(Process::CLOCK_MONOTONIC) * 1000) + (seconds * 1000) + end + + # One wire call in flight for a slot, plus the figure it asked for, which + # is what a joiner compares its own shortfall against. Threads that arrive + # while it runs wait on it instead of issuing a second call. + class Flight + attr_reader :requested_additional + + def initialize(requested_additional = nil) + @requested_additional = requested_additional + @mutex = Mutex.new + @condition = ConditionVariable.new + @done = false + @value = nil + end + + def complete(value) + @mutex.synchronize do + @value = value + @done = true + @condition.broadcast + end + end + + # Wait for the flight to land. A timeout_ms of nil waits indefinitely; + # a shutdown passes its remaining budget so a stalled wire call cannot + # hold the close open. + def wait(timeout_ms = nil) + deadline = timeout_ms ? monotonic_ms + timeout_ms : nil + @mutex.synchronize do + until @done + if deadline + remaining = deadline - monotonic_ms + break if remaining <= 0 + + @condition.wait(@mutex, remaining / 1000.0) + else + @condition.wait(@mutex) + end + end + @value + end + end + + def done? + @mutex.synchronize { @done } + end + + private + + def monotonic_ms + Process.clock_gettime(Process::CLOCK_MONOTONIC) * 1000 + end + end + + # Owns the lifecycle of credit leases for a single client: acquire on + # first use or after expiry, extend when the local view dips below the low + # water mark, release on close. + # + # Acquire and extend each get their own single-flight map, kept separate + # so an in-flight extend can never satisfy an acquire. Best-effort: a + # caller racing ahead of the registration can still duplicate a wire call, + # which is safe, because the server is idempotent for an active slot, + # replace keeps the first live lease, and extend reconciles to a total. + class LeaseManager + # A wait on a shared extend that ran out the joiner's own timeout. + JOIN_TIMED_OUT = Object.new.freeze + + def initialize(wire_client:, lease_store:, logger:, config: {}, clock: DEFAULT_CLOCK) + @wire = wire_client + @lease_store = lease_store + @logger = logger + @config = config || {} + @clock = clock + @flight_mutex = Mutex.new + @inflight_acquire = {} + @inflight_extend = {} + # Lease work nobody joins: the redundant release a lost acquire race + # issues, and the background extends callers fire and forget. drain + # waits these out so a close releases what they installed. + @background = [] + @stopped = false + end + + def resolve_config(credit_type_id) + Leases.resolve_config(@config, credit_type_id) + end + + # Return the slot's lease, acquiring one (or replacing an expired one) + # if none is live. Never raises: a wire or store failure is logged and + # reported as nil, so callers route it through their fail-open or + # fail-closed handling. + # + # deadline is a monotonic-clock millisecond cap on waiting for another + # caller's acquire, fixed once when the check started so the waits a + # check makes cannot each restart its timeout. Without one the cap is + # derived from request_options' timeout as of this call. + def acquire_if_needed(company_id, credit_type_id, request_options = nil, deadline: nil) + # Past stop the drain has run or is running, so a lease acquired now + # is one nothing is left to release. + return log_stopped("acquire", company_id, credit_type_id) if @stopped + + begin + existing = @lease_store.get(company_id, credit_type_id) + rescue StandardError => e + @logger.error("Failed to read lease store for #{company_id}/#{credit_type_id}: #{e.message}") + return nil + end + return existing if existing && !existing.expired?(@clock.call) + + # An expired or absent slot is left for replace to overwrite, which + # guards on expiry and writes atomically. Dropping it first would be a + # separate op that can interleave between a sibling's read and its + # replace, clobbering the lease that sibling just installed. + + # Check again: stop may have landed during the store read. + return log_stopped("acquire", company_id, credit_type_id) if @stopped + + key = Leases.lease_key(company_id, credit_type_id) + flight, leader = enlist(@inflight_acquire, key) { Flight.new } + return log_stopped("acquire", company_id, credit_type_id) if flight.nil? + return join_acquire(flight, deadline || join_deadline(request_options), company_id, credit_type_id) unless leader + + begin + result = acquire(company_id, credit_type_id, request_options) + ensure + @flight_mutex.synchronize { @inflight_acquire.delete(key) } + flight.complete(result) + end + result + end + + # Kick off an extend when one is due, on a background thread a caller can + # fire and forget. Returns the thread, or nil when nothing is due; the + # check flow joins it when the extend covers a reserve it just failed. + # + # Due means at or below the low-water-mark ratio, or below a named + # required_credits: one large check should not wait for the next + # sub-watermark check to top the lease up. deadline caps a wait on + # another caller's extend, as for acquire_if_needed. Never raises. + def maybe_extend_in_background(company_id, credit_type_id, required_credits = nil, request_options = nil, + deadline: nil) + if @stopped + # Extending past stop re-holds credits on a lease the close is about + # to release, or has already released. + log_stopped("extend", company_id, credit_type_id) + return nil + end + # Decided here rather than inside the thread: most checks sit nowhere + # near the water mark, and spawning one thread per check to learn that + # costs far more than the store read that answers it. + return nil unless extend_due?(company_id, credit_type_id, required_credits) + # An extend already on the wire is fetching the credits a watermark + # refresh wants, so a thread for it would find the flight and exit. + # A caller naming required_credits still spawns: it has a reserve to + # retry, and may need more than the flight asked for. + return nil if required_credits.nil? && extend_in_flight?(company_id, credit_type_id) + + # The whole call is tracked, not just the wire call inside it: callers + # drop the thread on the floor, so between the store read and the + # extend there would otherwise be a window where a drain sees nothing + # pending. + track do + extend_if_needed(company_id, credit_type_id, required_credits, request_options, deadline) + end + end + + # Refuse new lease work. Idempotent, and paired with drain: stopping + # first is what makes the drain terminate, since nothing can enqueue + # behind it. Taken under the flight lock that enlist reads it under, so + # once stop returns no further flight can register and the drain that + # follows sees every flight there will ever be. + def stop + @flight_mutex.synchronize { @stopped = true } + nil + end + + # Wait out lease work already on the wire, so a close releases what that + # work installs instead of orphaning it. Bounded: whatever has not + # landed by the deadline is abandoned rather than stalling the caller's + # shutdown, and the credits it holds fall back to server-side expiry. + def drain(timeout_ms = SHUTDOWN_DRAIN_TIMEOUT_MS) + deadline = monotonic_ms + timeout_ms + loop do + pending_threads, pending_flights = pending_work + return if pending_threads.empty? && pending_flights.empty? + + remaining = deadline - monotonic_ms + if remaining <= 0 + @logger.warn( + "Timed out after #{timeout_ms}ms draining in-flight credit lease work; " \ + "any credits it holds will be released by server-side expiry" + ) + return + end + # Recomputed per thread, not once for the round: a shared budget + # spent thread by thread would let N stalled threads wait N times + # the timeout a caller asked close to take. + pending_threads.each do |thread| + left = deadline - monotonic_ms + break if left <= 0 + + thread.join(left / 1000.0) + end + pending_flights.each { |flight| flight.wait(deadline - monotonic_ms) } + # Settling one round can enqueue another (an acquire that loses its + # race fires a release), so keep going until nothing is left. + end + end + + # Release every live lease held in the store, returning its unspent + # remainder now rather than at expiry. ONLY safe per-process: a shared + # store's siblings are still drawing on those leases, which the list + # capability check excludes. + # + # timeout_ms bounds the whole pass, because each release is a + # synchronous round trip and a slow API would otherwise stretch close by + # one timeout per slot. What is left expires server-side, the same + # outcome a failed release already has. + def release_all_local_leases(timeout_ms = nil) + return nil unless @lease_store.respond_to?(:list) + + # A failed listing is logged and left to server-side expiry: raising + # here would skip the rest of the caller's close. + entries = begin + @lease_store.list + rescue StandardError => e + @logger.warn("Failed to list credit leases on close (they will expire server-side): #{e.message}") + nil + end + return nil if entries.nil? || entries.empty? + + deadline = timeout_ms.nil? ? nil : monotonic_ms + timeout_ms + entries.each_with_index do |entry, index| + # Skip expired leases: the server already swept and refunded them. + next if entry.expired?(@clock.call) + + if deadline && monotonic_ms >= deadline + left = entries[index..].count { |remaining| !remaining.expired?(@clock.call) } + @logger.warn( + "Timed out after #{timeout_ms.round}ms releasing credit leases on close; " \ + "#{left} left to server-side expiry" + ) + break + end + + begin + @wire.release(lease_id: entry.lease_id) + @lease_store.drop(entry.company_id, entry.credit_type_id) + @logger.debug("Released credit lease #{entry.lease_id} on close") + rescue StandardError => e + @logger.warn( + "Failed to release credit lease #{entry.lease_id} on close " \ + "(it will expire server-side): #{e.message}" + ) + end + end + nil + end + + private + + def acquire(company_id, credit_type_id, request_options) + resolved = resolve_config(credit_type_id) + grant = @wire.acquire( + company_id: company_id, + credit_type_id: credit_type_id, + requested_amount: resolved.lease_size, + expires_at: @clock.call + (resolved.lease_duration_ms / 1000.0), + request_options: request_options || {} + ) + wrote = @lease_store.replace(LeaseEntry.new( + lease_id: grant.lease_id, + company_id: grant.company_id, + credit_type_id: grant.credit_type_id, + granted_amount: grant.granted_amount, + expires_at: grant.expires_at + )) + return @lease_store.get(company_id, credit_type_id) if wrote + + settle_lost_race(company_id, credit_type_id, grant) + rescue StandardError => e + @logger.error("Failed to acquire credit lease for #{company_id}/#{credit_type_id}: #{e.message}") + nil + end + + # A sibling installed a live lease first, so replace kept theirs. The + # server is idempotent for an active slot, so ours is normally the SAME + # lease and releasing it would refund a lease every process is still + # reserving against. Only a DIFFERENT lease is a redundant hold worth + # releasing. An empty slot (it expired in the gap) is skipped too: that + # id may well be what the next acquire is handed back. + def settle_lost_race(company_id, credit_type_id, grant) + current = @lease_store.get(company_id, credit_type_id) + if current && current.lease_id != grant.lease_id + @logger.debug( + "Lost acquire race for #{company_id}/#{credit_type_id}; releasing redundant lease #{grant.lease_id}" + ) + # Fire and forget: a failed release just falls back to lease expiry. + track do + @wire.release(lease_id: grant.lease_id) + rescue StandardError => e + @logger.warn("Failed to release redundant credit lease #{grant.lease_id}: #{e.message}") + end + else + @logger.debug( + "Lost acquire race for #{company_id}/#{credit_type_id}; " \ + "server returned the installed lease #{grant.lease_id}, nothing to release" + ) + end + current + end + + # Cheap read-only answer to "would extend_if_needed do anything". It can + # go stale between here and the thread, which is harmless: the thread + # re-reads and re-decides under the flight. + def extend_due?(company_id, credit_type_id, required_credits) + entry = @lease_store.get(company_id, credit_type_id) + return false if entry.nil? || entry.expired?(@clock.call) + + needs_extend?(entry, resolve_config(credit_type_id), required_credits) + rescue StandardError => e + # No extend without a reading: a store outage answering true would + # spawn a thread per check, and each would only fail the same read. + @logger.debug("Failed to read lease store for #{company_id}/#{credit_type_id}: #{e.message}") + false + end + + # 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 extend of its own it would return a balance it + # already knows is short and fail its retry with credits on the server. + def extend_if_needed(company_id, credit_type_id, required_credits, request_options, deadline = nil) + # 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 deadline: a + # check with 200ms to spend must not sit behind a 30s extend. A check + # passes the deadline it fixed when it started, since one taken now + # would grant a fresh timeout on top of the acquire and reserve. + deadline ||= join_deadline(request_options) + joins_left = MAX_EXTEND_JOINS + loop do + entry = read_live_lease(company_id, credit_type_id) + return nil if entry.nil? + + resolved = resolve_config(credit_type_id) + return entry unless needs_extend?(entry, resolved, required_credits) + + # Sized to cover the request that triggered it, or a check needing + # more than the tranche would fail its post-extend retry forever + # against an ample server balance. Sized here rather than at the + # wire call so the flight below and the request body provably carry + # the same number for a joiner to compare against. + shortfall = required_credits ? required_credits - entry.local_remaining_credits : 0 + additional_amount = [resolved.lease_size, shortfall].max + + key = Leases.lease_key(company_id, credit_type_id) + # With the budget spent we take the slot ourselves rather than + # joining a flight that has already proved too small. + flight, leader = enlist(@inflight_extend, key, force: joins_left <= 0) { Flight.new(additional_amount) } + return log_stopped("extend", company_id, credit_type_id) if flight.nil? + return run_extend(key, flight, entry, resolved, additional_amount, request_options) if leader + + # A watermark refresh that finds a flight already running has + # nothing to wait for: the credits it wants are the ones that call + # is fetching. Only a caller naming required_credits waits, because + # it has a reserve to retry. + return nil if required_credits.nil? + + joined = join_within(flight, deadline) + if joined.equal?(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 or + # fail-closed path, which is what its timeout asked for. + @logger.debug( + "Extend in flight for #{company_id}/#{credit_type_id} outlasted the caller's timeout; " \ + "not waiting on it" + ) + return nil + end + # The flight already asked for at least what we need, which covers + # every watermark-driven joiner and any check the tranche covers. + # One wire call serves all of them, which is the point. + return joined if additional_amount <= flight.requested_additional + + # 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 + end + end + + # Wait on another caller's acquire, reporting nil once the deadline + # passes. The acquire runs on for whoever else is on it, and what it + # installs is there for the next check to read. + def join_acquire(flight, deadline, company_id, credit_type_id) + joined = join_within(flight, deadline) + return joined unless joined.equal?(JOIN_TIMED_OUT) + + @logger.debug( + "Acquire in flight for #{company_id}/#{credit_type_id} outlasted the caller's timeout; not waiting on it" + ) + nil + end + + # Read the slot, reporting nil when the read fails or the lease is + # absent or expired. An expired lease is never extended: the server + # treats it as released, with its remainder already refunded to the + # company balance, so the right move is the fresh acquire the next check + # performs. + def read_live_lease(company_id, credit_type_id) + entry = @lease_store.get(company_id, credit_type_id) + return nil if entry.nil? || entry.expired?(@clock.call) + + entry + rescue StandardError => e + @logger.warn("Failed to read lease store for #{company_id}/#{credit_type_id}: #{e.message}") + nil + end + + # Whether the entry sits low enough to warrant an extend. + def needs_extend?(entry, resolved, required_credits) + ratio = entry.local_remaining_credits / [entry.granted_amount, 1].max + below_watermark = ratio <= resolved.low_water_mark + below_required = !required_credits.nil? && entry.local_remaining_credits < required_credits + below_watermark || below_required + end + + # When a joiner's wait on a shared flight runs out, or nil for no cap. + def join_deadline(request_options) + Leases.join_deadline(request_options) + end + + # Wait out a flight somebody else is running, giving up at the 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. + def join_within(flight, deadline) + return flight.wait if deadline.nil? + + remaining = deadline - monotonic_ms + return JOIN_TIMED_OUT if remaining <= 0 + + joined = flight.wait(remaining) + flight.done? ? joined : JOIN_TIMED_OUT + end + + # Run the extend as the slot's in-flight one. The cleanup is + # identity-guarded rather than an unconditional delete: a joiner whose + # shortfall outran this flight registers a follow-up for the same key, + # and this flight must not evict it. + def run_extend(key, flight, entry, resolved, additional_amount, request_options) + result = extend(entry, resolved, additional_amount, request_options) + result + ensure + @flight_mutex.synchronize { @inflight_extend.delete(key) if @inflight_extend[key].equal?(flight) } + flight.complete(result) + end + + def extend(entry, resolved, additional_amount, request_options) + grant = @wire.extend( + lease_id: entry.lease_id, + additional_amount: additional_amount, + expires_at: @clock.call + (resolved.lease_duration_ms / 1000.0), + # Minted once per extend, outside the wire call, so the transport's + # retries resend the same key: a retry after a lost 2xx is handed + # the lease as it stands instead of growing it a second time. + idempotency_key: SecureRandom.uuid, + request_options: request_options || {} + ) + # Reconciled to the server's TOTAL, with the store computing the delta + # against its own current figure rather than the read above: two + # sibling processes extending one shared lease would each apply a + # stale-read delta and mint phantom credits. Pinned to the lease the + # server extended, so an expiry mid-call drops the delta instead of + # minting it onto the successor. + @lease_store.extend(entry.company_id, entry.credit_type_id, grant.granted_amount, grant.expires_at, + entry.lease_id) + @logger.debug( + "Extended credit lease #{entry.lease_id} to #{grant.granted_amount} " \ + "(was #{entry.granted_amount} at last read)" + ) + @lease_store.get(entry.company_id, entry.credit_type_id) + rescue StandardError => e + @logger.warn("Failed to extend credit lease #{entry.lease_id}: #{e.message}") + nil + end + + # Register this caller against the slot's in-flight call, or become the + # one that makes it. Returns [flight, leader], or [nil, false] once the + # manager is stopped: the stopped read happens here, under the same lock + # stop takes, because an unguarded read outside it can be overtaken by a + # close whose drain then never sees the flight this would have made. + def enlist(flights, key, force: false) + @flight_mutex.synchronize do + next [nil, false] if @stopped + + existing = flights[key] + next [existing, false] if existing && !force + + flight = yield + flights[key] = flight + [flight, true] + end + end + + # Whether the slot already has an extend on the wire. + def extend_in_flight?(company_id, credit_type_id) + @flight_mutex.synchronize { !@inflight_extend[Leases.lease_key(company_id, credit_type_id)].nil? } + end + + # Hold a reference to work nobody joins so drain can wait it out. + def track(&block) + thread = Thread.new do + block.call + rescue StandardError => e + @logger.warn("Background credit lease work failed: #{e.message}") + nil + end + thread.abort_on_exception = false + @flight_mutex.synchronize do + # Finished work is dropped as new work arrives, so the list tracks + # what is still pending rather than growing for the process's life. + @background.select!(&:alive?) + @background << thread + end + thread + end + + def pending_work + @flight_mutex.synchronize do + @background.select!(&:alive?) + [@background.dup, (@inflight_acquire.values + @inflight_extend.values).reject(&:done?)] + end + end + + def log_stopped(action, company_id, credit_type_id) + @logger.debug("Lease manager is stopped; skipping #{action} for #{company_id}/#{credit_type_id}") + nil + end + + def monotonic_ms + Process.clock_gettime(Process::CLOCK_MONOTONIC) * 1000 + end + end + end + end +end diff --git a/lib/schematic/credits/leases/lease_store.rb b/lib/schematic/credits/leases/lease_store.rb new file mode 100644 index 0000000..a80ecdb --- /dev/null +++ b/lib/schematic/credits/leases/lease_store.rb @@ -0,0 +1,240 @@ +# frozen_string_literal: true + +module Schematic + module Credits + module Leases + def self.lease_key(company_id, credit_type_id) + "#{company_id}:#{credit_type_id}" + end + + # Per-process lease store, keyed by ":". + # + # Holds the lease id, the server-authoritative granted total, the expiry, + # and the SDK's local view of local_remaining_credits: the portion of the + # lease not yet carved out by an open reservation. A per-key mutex keeps + # reserve, refund, extend, and replace atomic per slot. + # + # The store never talks to the API. LeaseManager drives acquire, extend, + # and release over the wire and uses these methods to mirror remote state. + # + # For cross-process deployments use RedisLeaseStore instead: both answer + # the same calls with the same semantics. + class LeaseStore + def initialize(clock: DEFAULT_CLOCK) + @clock = clock + @leases = {} + @locks = {} + # Guards the two hashes themselves. Always taken inside a key lock, + # never the other way round, so the pair cannot deadlock. + @table_mutex = Mutex.new + end + + # Snapshot of the current entry, or nil when no lease occupies the slot. + # A copy, so a caller cannot mutate store state by holding the result. + def get(company_id, credit_type_id) + key = Leases.lease_key(company_id, credit_type_id) + with_lock_if_present(key) { read(key)&.dup } + end + + # Install a fresh lease for the slot, but only if no live lease already + # occupies it. + # + # A live lease is left untouched even when its lease_id differs (a + # sibling process won the race), so its already-debited + # local_remaining_credits wins. An expired row carrying the SAME + # lease_id is not rewritten either: rewriting would reset + # local_remaining_credits to the full grant and erase debits whose + # reservations are still open, and the idempotent server hands a racing + # acquire that same active lease back. Such a row is reconciled like an + # extend instead: granted to the incoming total, expiry only forward, + # balance untouched. + # + # Returns true when a fresh row was written, false when an existing + # lease was kept or reconciled. + def replace(entry) + key = Leases.lease_key(entry.company_id, entry.credit_type_id) + with_lock(key) do + existing = read(key) + if existing && !existing.expired?(@clock.call) + false + elsif existing && existing.lease_id == entry.lease_id + reconcile(existing, entry.granted_amount, entry.expires_at) + false + else + write(key, LeaseEntry.new( + lease_id: entry.lease_id, + company_id: entry.company_id, + credit_type_id: entry.credit_type_id, + granted_amount: entry.granted_amount, + expires_at: entry.expires_at + )) + true + end + end + end + + # Reconcile the slot to the server-authoritative granted total after a + # remote extend, crediting the difference to local_remaining_credits. + # + # The delta is computed here against the CURRENT stored total, never by + # the caller from a pre-wire-call read: two writers extending the same + # lease concurrently would each apply a delta against the same stale + # read and mint phantom credits. Reconciling to the absolute total + # converges in any order, and the expiry only ever moves forward so an + # out-of-order apply cannot shorten a lease a sibling just extended. + # + # When pin_lease_id is given the extend applies only if the slot still + # holds that lease: the server extended lease A, so its credits must not + # land on a successor B that replaced A after it expired mid-extend. + def extend(company_id, credit_type_id, granted_amount, new_expires_at = nil, pin_lease_id = nil) + key = Leases.lease_key(company_id, credit_type_id) + with_lock_if_present(key) do + entry = read(key) + next if entry.nil? + next if pin_lease_id && entry.lease_id != pin_lease_id + + reconcile(entry, granted_amount, new_expires_at) + end + nil + end + + # Drop the slot entry, after a remote release. An explicit drop is the + # only thing that removes one: a lease that merely expired stays here, + # readable, until it is dropped or replaced, which is what the spec + # requires, so every path re-guards on expiry rather than trusting + # presence. The slot's mutex goes with the entry, so a long-lived + # process does not accumulate one per slot it has ever leased. + def drop(company_id, credit_type_id) + key = Leases.lease_key(company_id, credit_type_id) + with_lock(key) do + @table_mutex.synchronize do + @leases.delete(key) + @locks.delete(key) + end + end + nil + end + + # Atomically check and debit `credits` from the lease's remaining + # balance. Returns a ReserveResult carrying the post-debit balance and + # the lease the debit landed on, or nil when there is no live lease, the + # balance is short, or `credits` is not a finite non-negative number. + # + # The debit is not keyed by lease id: it charges whichever lease holds + # the slot at that moment, which need not be the one the caller's + # acquire returned. The caller must therefore pin its reservation to the + # returned lease_id, since the settle refund, the sweep refund, and the + # track event's lease_id all have to name the lease that was charged. + def try_reserve(company_id, credit_type_id, credits) + # NaN passes every comparison below, and a NaN balance would approve + # every later reserve, so reject it before it reaches the arithmetic. + return nil unless Leases.valid_quantity?(credits) + + key = Leases.lease_key(company_id, credit_type_id) + with_lock_if_present(key) do + entry = read(key) + next nil if entry.nil? + # An expired lease is released server-side and its grant refunded to + # the company balance, so its local balance is stale. + next nil if entry.expired?(@clock.call) + next nil if entry.local_remaining_credits < credits + + entry.local_remaining_credits -= credits + # The lease id is read under the same lock as the debit: a read + # afterwards could name a lease that replaced this one in between. + ReserveResult.new(entry.local_remaining_credits, entry.lease_id) + end + end + + # Refund credits to the slot's lease balance, clamped at granted_amount. + # + # When pin_lease_id is given the refund applies only if the slot still + # holds that lease: a hold carved out of expired lease A must never + # inflate a successor B, because A's unspent remainder was already + # returned to the company balance server-side when A expired. + def refund(company_id, credit_type_id, credits, pin_lease_id = nil) + return nil if credits.nil? || credits <= 0 + + key = Leases.lease_key(company_id, credit_type_id) + with_lock_if_present(key) do + entry = read(key) + next if entry.nil? + next if pin_lease_id && entry.lease_id != pin_lease_id + + entry.local_remaining_credits = [entry.local_remaining_credits + credits, entry.granted_amount].min + end + nil + end + + # Snapshot of every entry. Only the per-process store implements this: + # close releases the leases this process exclusively holds, and a shared + # backend must never enumerate, since sibling processes may still be + # drawing on those leases. + def list + keys = @table_mutex.synchronize { @leases.keys } + keys.filter_map { |key| with_lock(key) { read(key)&.dup } } + end + + private + + # Granted to the incoming total, expiry only forward, balance untouched + # except for the credited difference. Shared by replace's same-lease + # path and extend, which reconcile identically. + def reconcile(entry, granted_amount, new_expires_at) + add = granted_amount.to_f - entry.granted_amount + if add.positive? + entry.granted_amount = granted_amount.to_f + entry.local_remaining_credits += add + end + return unless new_expires_at && new_expires_at.to_f > entry.expires_at.to_f + + entry.expires_at = new_expires_at + end + + def read(key) + @table_mutex.synchronize { @leases[key] } + end + + def write(key, entry) + @table_mutex.synchronize { @leases[key] = entry } + end + + # A read of a slot nothing has leased must not leave a mutex behind, or + # pruning would only half work: a process checking flags for companies + # that never lease would still collect one per slot it asked about. With + # no lock registered and no lease stored there is nothing to serialize + # against, so answer without taking one. A writer landing in that window + # is the same race as reading a moment earlier. + def with_lock_if_present(key, &) + present = @table_mutex.synchronize { @locks.key?(key) || @leases.key?(key) } + return nil unless present + + with_lock(key, &) + end + + # Serialize on the slot's mutex, re-checking after the acquire that it + # is still the registered one. + # + # Pruning is what makes the re-check necessary. A drop deletes the mutex + # while holding it, so a thread that was already blocked on it wakes + # owning an object the table no longer knows about, while a thread + # arriving afterwards takes a fresh mutex for the same slot. Without the + # check those two would run side by side on one slot. The waiter sees + # its mutex is no longer the registered one and retries against the + # current one instead. Every retry follows a drop, which happens once + # per lease, so the loop cannot spin. + def with_lock(key, &block) + loop do + lock = @table_mutex.synchronize { @locks[key] ||= Mutex.new } + stale = false + result = lock.synchronize do + stale = @table_mutex.synchronize { !@locks[key].equal?(lock) } + block.call unless stale + end + return result unless stale + end + end + end + end + end +end diff --git a/lib/schematic/credits/leases/redis_lease_store.rb b/lib/schematic/credits/leases/redis_lease_store.rb new file mode 100644 index 0000000..4a04968 --- /dev/null +++ b/lib/schematic/credits/leases/redis_lease_store.rb @@ -0,0 +1,358 @@ +# frozen_string_literal: true + +require "digest" + +module Schematic + module Credits + module Leases + # Minimal interface describing the Redis client methods the lease and + # reservation stores use. Compatible with the "redis" gem's client, which + # is deliberately not a dependency: the caller injects a client it already + # has, the same way RedisCacheProvider takes one. + module RedisLeaseClientInterface + def evalsha(sha, keys:, argv:) + raise NotImplementedError + end + + def eval(script, keys:, argv:) + raise NotImplementedError + end + + def hgetall(key) + raise NotImplementedError + end + + # Called as hset(key, field1, value1, field2, value2, ...). + def hset(key, *pairs) + raise NotImplementedError + end + + def hdel(key, field) + raise NotImplementedError + end + + def del(*keys) + raise NotImplementedError + end + + def pexpireat(key, millis) + raise NotImplementedError + end + + def zadd(key, score, member) + raise NotImplementedError + end + + def zrem(key, member) + raise NotImplementedError + end + + def zrangebyscore(key, min, max, limit: nil) + raise NotImplementedError + end + + def zcard(key) + raise NotImplementedError + end + end + + # A Lua script addressed by its SHA. EVALSHA first so the script body is + # not resent on every call; a Redis that has never seen it (a restart, or + # a node this process has not talked to) answers NOSCRIPT and the full + # body goes out once to load it. + class Script + attr_reader :source, :sha + + def initialize(source) + @source = source + @sha = Digest::SHA1.hexdigest(source) + end + + def call(client, keys:, argv:) + client.evalsha(@sha, keys: keys, argv: argv) + rescue StandardError => e + raise unless e.message.to_s.include?("NOSCRIPT") + + client.eval(@source, keys: keys, argv: argv) + end + end + + # Redis-backed lease store. One hash per (company_id, credit_type_id) + # slot, mutated by single-key Lua scripts so the store stays correct on + # standalone and clustered Redis alike. + # + # The scripts below are byte-identical to the ones the Node, Go, and + # Python SDKs ship, and the key layout matches, so a mixed-language fleet + # shares one lease per slot. Do not re-derive them. + class RedisLeaseStore + DEFAULT_KEY_PREFIX = "schematic:" + LEASE_KEY_NAMESPACE = "credit-lease:" + # How long after the declared expiry the Redis row is kept before + # auto-eviction. Gives the sweeper a window to refund expired + # reservations before the underlying lease state disappears. + LEASE_TTL_GRACE_MS = 60_000 + + # Every script below touches exactly ONE key (the lease hash), which + # keeps them safe under Redis Cluster, where a multi-key script spanning + # slots raises CROSSSLOT. Only the lease hash needs atomic mutation; + # cross-key bookkeeping uses ordinary single-key commands. + # + # Expiry is decided against the Redis server's clock (redis.call('TIME')), + # not the calling process's: with many processes sharing one lease, local + # clock skew would let them disagree on whether the lease is live. + + # Atomic replace. Writes the lease hash only when the slot is empty or + # the existing lease has expired. Returns 1 on write, 0 when a LIVE + # lease already occupies the slot, even one with a different leaseId + # installed by a sibling that raced this acquire. An expired row with + # the SAME leaseId is reconciled like an extend instead of rewritten, + # since rewriting would reset the balance and erase debits whose + # reservations are still open. + REPLACE_SCRIPT = Script.new(<<~LUA) + + redis.replicate_commands() + local t = redis.call('TIME') + local now = (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000) + + local existing_id = redis.call('HGET', KEYS[1], 'leaseId') + local existing_expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0') + local new_id = ARGV[1] + local new_granted = ARGV[2] + local new_expiry = tonumber(ARGV[3]) + local grace = tonumber(ARGV[4]) + + if existing_id and existing_expiry > now then + return 0 + end + + if existing_id == new_id then + local granted = tonumber(redis.call('HGET', KEYS[1], 'grantedAmount') or '0') + local add = tonumber(new_granted) - granted + if add > 0 then + local remaining = tonumber(redis.call('HGET', KEYS[1], 'localRemainingCredits') or '0') + redis.call('HSET', KEYS[1], + 'grantedAmount', new_granted, + 'localRemainingCredits', tostring(remaining + add)) + end + if new_expiry > existing_expiry then + redis.call('HSET', KEYS[1], 'expiresAt', ARGV[3]) + redis.call('PEXPIREAT', KEYS[1], new_expiry + grace) + end + return 0 + end + + redis.call('DEL', KEYS[1]) + redis.call('HSET', KEYS[1], + 'leaseId', new_id, + 'companyId', ARGV[5], + 'creditTypeId', ARGV[6], + 'grantedAmount', new_granted, + 'localRemainingCredits', new_granted, + 'expiresAt', ARGV[3]) + redis.call('PEXPIREAT', KEYS[1], new_expiry + grace) + return 1 + LUA + + # Atomic check and decrement on localRemainingCredits. Returns + # [post-debit balance, charged leaseId] on success, with the balance as + # a string because a Lua number reply truncates to integer and would + # corrupt fractional credit costs; a nil reply when there is no lease, + # the lease has expired, or the remaining balance is short. Returning + # the lease id read inside the same script is what lets the caller pin + # its reservation to the lease the debit actually landed on. + TRY_RESERVE_SCRIPT = Script.new(<<~LUA) + + redis.replicate_commands() + local t = redis.call('TIME') + local now = (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000) + + 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) + local requested = tonumber(ARGV[1]) + 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), lease_id } + LUA + + # Refund credits, clamped at grantedAmount. ARGV[2], when non-empty, + # pins the refund to a leaseId: if the slot now holds a different lease, + # the refund is dropped, because the expired lease's unspent remainder + # was already returned to the company balance server-side. + REFUND_SCRIPT = Script.new(<<~LUA) + + local raw_remaining = redis.call('HGET', KEYS[1], 'localRemainingCredits') + if not raw_remaining then return 0 end + local required_lease = ARGV[2] + if required_lease and required_lease ~= '' then + local current_lease = redis.call('HGET', KEYS[1], 'leaseId') + if current_lease ~= required_lease then return 0 end + end + local remaining = tonumber(raw_remaining) + local granted = tonumber(redis.call('HGET', KEYS[1], 'grantedAmount') or '0') + local refund = tonumber(ARGV[1]) + local new_balance = remaining + refund + if new_balance > granted then new_balance = granted end + redis.call('HSET', KEYS[1], 'localRemainingCredits', tostring(new_balance)) + return 1 + LUA + + # Reconcile the lease to the server-authoritative grantedAmount total, + # crediting the difference to localRemainingCredits. The delta is + # computed HERE, atomically against the hash's current total, so two + # processes extending the same lease concurrently converge instead of + # minting phantom credits. Expiry only moves forward. ARGV[4] pins the + # extend to a leaseId, mirroring the pin on the refund. + EXTEND_SCRIPT = Script.new(<<~LUA) + + local raw_granted = redis.call('HGET', KEYS[1], 'grantedAmount') + if not raw_granted then return 0 end + local required_lease = ARGV[4] + if required_lease and required_lease ~= '' then + local current_lease = redis.call('HGET', KEYS[1], 'leaseId') + if current_lease ~= required_lease then return 0 end + end + local granted = tonumber(raw_granted) + local target = tonumber(ARGV[1]) + local add = target - granted + if add > 0 then + local remaining = tonumber(redis.call('HGET', KEYS[1], 'localRemainingCredits') or '0') + redis.call('HSET', KEYS[1], + 'grantedAmount', tostring(target), + 'localRemainingCredits', tostring(remaining + add)) + end + local new_expiry = tonumber(ARGV[2]) + local grace = tonumber(ARGV[3]) + local current_expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0') + if new_expiry > current_expiry then + redis.call('HSET', KEYS[1], 'expiresAt', ARGV[2]) + redis.call('PEXPIREAT', KEYS[1], new_expiry + grace) + end + return 1 + LUA + + def initialize(client:, key_prefix: nil, default_lease_duration_ms: DEFAULT_LEASE_DURATION_MS, + clock: DEFAULT_CLOCK) + @client = client + @key_prefix = key_prefix || DEFAULT_KEY_PREFIX + # Defensive fallback for a direct caller that extends without naming a + # new expiry. The lease manager always passes one. + @default_lease_duration_ms = default_lease_duration_ms + @clock = clock + end + + # Public so the reservation store can target the same lease hash. + def hash_key(company_id, credit_type_id) + "#{@key_prefix}#{LEASE_KEY_NAMESPACE}#{Leases.lease_key(company_id, credit_type_id)}" + end + + def get(company_id, credit_type_id) + raw = @client.hgetall(hash_key(company_id, credit_type_id)) + return nil if raw.nil? || raw["leaseId"].nil? + + decode_entry(raw) + end + + # rubocop:disable Naming/PredicateMethod + # replace is the store method name every Schematic SDK shares; its + # boolean says whether a fresh row was written or an existing lease kept. + def replace(entry) + # No process clock here: the script reads now from the Redis server + # via TIME, so every process agrees on expiry. + result = REPLACE_SCRIPT.call( + @client, + keys: [hash_key(entry.company_id, entry.credit_type_id)], + argv: [ + entry.lease_id, + num(entry.granted_amount), + millis(entry.expires_at).to_s, + LEASE_TTL_GRACE_MS.to_s, + entry.company_id, + entry.credit_type_id + ] + ) + result.to_i == 1 + end + # rubocop:enable Naming/PredicateMethod + + def extend(company_id, credit_type_id, granted_amount, new_expires_at = nil, pin_lease_id = nil) + expiry = new_expires_at ? millis(new_expires_at) : millis(@clock.call) + @default_lease_duration_ms + # granted_amount is the server-authoritative TOTAL; the script computes + # the credit delta atomically against the stored total. An empty string + # disables the lease pin, since Lua has no nil ARGV. + EXTEND_SCRIPT.call( + @client, + keys: [hash_key(company_id, credit_type_id)], + argv: [num(granted_amount), expiry.to_s, LEASE_TTL_GRACE_MS.to_s, pin_lease_id.to_s] + ) + nil + end + + def drop(company_id, credit_type_id) + # A plain single-key delete: there is no secondary index to keep in sync. + @client.del(hash_key(company_id, credit_type_id)) + nil + end + + def try_reserve(company_id, credit_type_id, credits) + # Reject non-finite or negative debits before they reach the script: + # NaN.to_s parses back to a Lua nan, slips through the comparison, and + # would poison the SHARED balance for every process. + return nil unless Leases.valid_quantity?(credits) + + result = TRY_RESERVE_SCRIPT.call( + @client, + keys: [hash_key(company_id, credit_type_id)], + # Only the requested amount: now comes from the Redis server clock. + argv: [num(credits)] + ) + return nil if result.nil? || result == false + + ReserveResult.new(result[0].to_f, result[1].to_s) + end + + def refund(company_id, credit_type_id, credits, pin_lease_id = nil) + return nil if credits.nil? || credits <= 0 + + REFUND_SCRIPT.call( + @client, + keys: [hash_key(company_id, credit_type_id)], + # An empty string disables the lease pin, since Lua has no nil ARGV. + argv: [num(credits), pin_lease_id.to_s] + ) + nil + end + + private + + # Lua's tonumber reads a decimal string, and Ruby's to_s does not always + # produce one: a Rational quantity renders as "1/2", which the script + # reads as nil and treats as a zero. Integers are already safe and stay + # exact, so only the rest is forced through Float. + def num(value) + value.is_a?(Integer) ? value.to_s : value.to_f.to_s + end + + def millis(time) + (time.to_f * 1000).round + end + + def decode_entry(raw) + LeaseEntry.new( + lease_id: raw["leaseId"], + company_id: raw["companyId"], + credit_type_id: raw["creditTypeId"], + granted_amount: (raw["grantedAmount"] || 0).to_f, + local_remaining_credits: (raw["localRemainingCredits"] || 0).to_f, + expires_at: Time.at((raw["expiresAt"] || 0).to_f / 1000.0) + ) + end + end + end + end +end diff --git a/lib/schematic/credits/leases/redis_reservation_store.rb b/lib/schematic/credits/leases/redis_reservation_store.rb new file mode 100644 index 0000000..7beb8e3 --- /dev/null +++ b/lib/schematic/credits/leases/redis_reservation_store.rb @@ -0,0 +1,326 @@ +# frozen_string_literal: true + +require "json" + +module Schematic + module Credits + module Leases + # Redis-backed reservation table. Each reservation is a hash, indexed by + # expires_at in a sorted set so the sweeper can pop expired entries in + # O(log n), and by (company, credit) so a balance display can sum a + # tenant's open holds. + # + # Every mutation is a single-key operation or single-key Lua, so the store + # is correct on standalone and clustered Redis alike: the unspent-slice + # refund is delegated to the lease store rather than reaching across to + # the lease hash inside a multi-key script. + class RedisReservationStore + DEFAULT_KEY_PREFIX = "schematic:" + RES_KEY_NAMESPACE = "credit-reservation:" + # Sorted set scoring open reservations by expires_at. Members encode the + # full (company, credit, id) tuple so the sweeper can clean the + # per-tenant hash even after the reservation hash has TTL-evicted, at + # which point the claim returns nil and cannot report company or credit; + # otherwise the orphaned field would permanently inflate reserved_credits. + RES_INDEX_KEY = "credit-reservations:byExpiry" + # Per-(company, credit) index of open reservations, one hash of + # reservationId to creditsReserved. reserved_credits then reads a whole + # tenant's holds with ONE HGETALL (single key, Cluster-safe) and sums the + # values. The hash also IS the source of truth for the sum: a field + # exists iff its reservation is open and unrefunded. + RES_BYCREDIT_NAMESPACE = "credit-reservations:byCredit:" + # Buffer past expires_at before Redis auto-evicts the row, so the + # sweeper has a window to refund. + RES_TTL_GRACE_MS = 30_000 + # The delimiter is absent from Schematic ids and from the UUID + # reservation id. + MEMBER_DELIMITER = "|" + # Page size for the sweeper's ZRANGEBYSCORE. Without a limit, a backlog + # of expired holds would come back as one giant reply on every tick. + SWEEP_BATCH_SIZE = 256 + # Upper bound on pages per tick. Keeps one sweep's work bounded and + # guards against an endless loop if a zrem persistently fails. Anything + # left over is picked up next tick. + MAX_SWEEP_BATCHES = 16 + + # Atomic claim: read the reservation hash and delete it in one step, + # returning its fields, or nil if it was already gone. One key, so it is + # Cluster-safe. The atomic read-then-delete is what makes consume + # exactly-once: of two racing callers only one gets the fields back and + # proceeds to refund. The refund is a separate single-key op; a crash in + # the gap leaves the unspent slice held on the lease until the lease + # itself expires, never double-refunded. + CLAIM_SCRIPT = Script.new(<<~LUA) + + local raw = redis.call('HGETALL', KEYS[1]) + if #raw == 0 then return nil end + redis.call('DEL', KEYS[1]) + return raw + LUA + + def initialize(client:, lease_store:, sweep_interval_ms: DEFAULT_SWEEP_INTERVAL_MS, key_prefix: nil, + clock: DEFAULT_CLOCK, logger: nil) + @client = client + @lease_store = lease_store + @sweep_interval_ms = sweep_interval_ms + @key_prefix = key_prefix || DEFAULT_KEY_PREFIX + @clock = clock + @logger = logger + @mutex = Mutex.new + @sweep_thread = nil + @stopped = false + end + + def add(reservation) + expires_ms = millis(reservation.expires_at) + key = hash_key(reservation.id) + fields = [ + "id", reservation.id, + "leaseId", reservation.lease_id, + "companyId", reservation.company_id, + "creditTypeId", reservation.credit_type_id, + "eventSubtype", reservation.event_subtype.to_s, + "quantityReserved", reservation.quantity_reserved.to_s, + "creditsReserved", reservation.credits_reserved.to_s, + "consumptionRate", reservation.consumption_rate.to_s, + "expiresAt", expires_ms.to_s, + "evalCtx", JSON.generate(reservation.eval_ctx) + ] + 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, so what other SDKs read is unchanged, and + # both commands touch the one key, so this is Cluster-safe. + if @client.respond_to?(:multi) + @client.multi do |tx| + tx.hset(key, *fields) + tx.pexpireat(key, ttl_at) + end + else + @client.hset(key, *fields) + @client.pexpireat(key, ttl_at) + end + # The two indexes (expiry zset for the sweeper, per-tenant hash for + # reserved_credits) only depend on the hash existing, so they stay + # outside the transaction: their keys hash to other slots. A partial + # failure here at worst leaves an un-indexed reservation that the TTL + # reaps, never a double spend. + @client.zadd(index_key, expires_ms, encode_member(reservation.company_id, reservation.credit_type_id, + reservation.id)) + @client.hset(by_credit_key(reservation.company_id, reservation.credit_type_id), reservation.id, + reservation.credits_reserved.to_s) + nil + end + + def get(id) + raw = @client.hgetall(hash_key(id)) + return nil if raw.nil? || raw["id"].nil? + + decode_reservation(raw) + end + + # Sum a tenant's open holds with a single HGETALL on the per-tenant + # index hash: one round trip, one key, no per-reservation fan-out. A + # field is present iff its reservation is open and unrefunded, so the + # sum is exact. + def reserved_credits(company_id, credit_type_id) + by_credit = begin + @client.hgetall(by_credit_key(company_id, credit_type_id)) + rescue StandardError + {} + end + (by_credit || {}).values.sum(&:to_f) + end + + def consume(id, credits_consumed) + # Atomically claim the reservation hash. Only one caller wins; a + # duplicate or racing consume gets nil. + claimed = CLAIM_SCRIPT.call(@client, keys: [hash_key(id)], argv: []) + raw = decode_raw_array(claimed) + return nil if raw.nil? || raw["id"].nil? + + company_id = raw["companyId"] + credit_type_id = raw["creditTypeId"] + reserved = raw["creditsReserved"].to_f + + # Index cleanup, single-key ops. The credits leave the per-tenant hash + # BEFORE the refund below so the lease (local remaining plus this + # hash) never transiently double-counts the slice, and before the + # expiry index because that index is the only way the sweeper reaches + # a surviving field: dropping the index first and then failing on the + # field would inflate reserved_credits for that tenant forever. + swallow { @client.hdel(by_credit_key(company_id, credit_type_id), id) } + swallow { @client.zrem(index_key, encode_member(company_id, credit_type_id, id)) } + + consumed = credits_consumed.clamp(0, reserved) + refund = reserved - consumed + # A hold that cannot name its lease is not refundable: the refund + # script reads an empty pin as no pin and would credit whichever lease + # holds the slot now, possibly a successor. The slice comes back when + # the lease expires instead, the same call the per-process store makes. + if refund.positive? && !raw["leaseId"].to_s.empty? + # Delegated to the lease store, which owns the lease hash, so the + # cross-key write stays out of a single Lua script. Pinned to the + # reservation's leaseId so a hold carved out of an expired lease + # cannot inflate a successor lease's balance. + @lease_store.refund(company_id, credit_type_id, refund, raw["leaseId"]) + end + consumed + end + + def start_sweep + @mutex.synchronize do + # alive?, not presence: a process that forks after start (Puma or + # Unicorn with preload) hands the child a thread object whose thread + # did not survive the fork, and the child would never sweep. + return if @sweep_thread&.alive? || @stopped + + interval = @sweep_interval_ms.to_f / 1000.0 + @sweep_thread = Thread.new do + loop do + sleep(interval) + break if @stopped + + begin + sweep_expired + rescue StandardError => e + @logger&.debug("Reservation sweep failed: #{e.message}") + end + end + end + @sweep_thread.abort_on_exception = false + end + end + + def sweep_expired(now = nil) + cutoff = millis(now || @clock.call) + swept = 0 + # Page through expired members rather than fetching them all at once. + # Each processed member is removed below, so re-reading at offset 0 + # advances through the backlog. + MAX_SWEEP_BATCHES.times do + expired = @client.zrangebyscore(index_key, 0, cutoff, limit: [0, SWEEP_BATCH_SIZE]) || [] + break if expired.empty? + + expired.each { |member| swept += 1 if member_refunded?(member) } + break if expired.size < SWEEP_BATCH_SIZE + end + swept + end + + def stop + thread = @mutex.synchronize do + @stopped = true + thread = @sweep_thread + @sweep_thread = nil + thread + end + return nil if thread.nil? + + # A sweep past its claim has deleted the reservation but not yet + # refunded the lease, and killing it there strands the unspent slice + # until the lease expires. Give it a bounded moment to land that + # refund, then kill: the loop checks @stopped as soon as its sleep + # ends, so a parked sweeper costs at most this wait. + thread.kill unless thread.join(SWEEP_STOP_JOIN_MS / 1000.0) + nil + end + + def size + @client.zcard(index_key) + rescue StandardError + 0 + end + + private + + def member_refunded?(member) + decoded = decode_member(member) + if decoded.nil? + # Nothing but add writes these, so this is belt and braces. Drop it + # so an unparseable member cannot wedge the sweeper. + swallow { @client.zrem(index_key, member) } + return false + end + company_id, credit_type_id, id = decoded + refunded = consume(id, 0) + # Always drop the member just read. On the success path consume + # already removed it, so this is an idempotent no-op; it also covers + # the hash-evicted path below. + swallow { @client.zrem(index_key, member) } + return true unless refunded.nil? + + # The claim found no reservation hash. Either a racing track already + # consumed it and reconciled the byCredit field, or the hash + # TTL-evicted before the sweeper reached it, orphaning that field. + # Reconcile it so reserved_credits cannot keep summing an evicted + # hold. The unspent slice is NOT refunded here: without the hash the + # claim cannot arbitrate exactly-once across racing sweepers, so the + # slice is reclaimed when the lease itself expires server-side. + swallow { @client.hdel(by_credit_key(company_id, credit_type_id), id) } + false + end + + def hash_key(id) + "#{@key_prefix}#{RES_KEY_NAMESPACE}#{id}" + end + + def index_key + "#{@key_prefix}#{RES_INDEX_KEY}" + end + + def by_credit_key(company_id, credit_type_id) + "#{@key_prefix}#{RES_BYCREDIT_NAMESPACE}#{company_id}:#{credit_type_id}" + end + + def encode_member(company_id, credit_type_id, id) + [company_id, credit_type_id, id].join(MEMBER_DELIMITER) + end + + def decode_member(member) + parts = member.to_s.split(MEMBER_DELIMITER) + return nil unless parts.size == 3 + + parts + end + + def millis(time) + (time.to_f * 1000).round + end + + def swallow + yield + rescue StandardError + nil + end + + # Decode a flat [field, value, field, value, ...] HGETALL array, which + # is what the claim script returns. + def decode_raw_array(raw) + return nil unless raw.is_a?(Array) && !raw.empty? + + out = {} + raw.each_slice(2) { |field, value| out[field.to_s] = value.to_s if value } + out + end + + def decode_reservation(raw) + Reservation.new( + id: raw["id"], + lease_id: raw["leaseId"], + company_id: raw["companyId"], + credit_type_id: raw["creditTypeId"], + event_subtype: raw["eventSubtype"], + quantity_reserved: raw["quantityReserved"].to_f, + credits_reserved: raw["creditsReserved"].to_f, + consumption_rate: raw["consumptionRate"].to_f, + expires_at: Time.at(raw["expiresAt"].to_f / 1000.0), + eval_ctx: raw["evalCtx"] ? JSON.parse(raw["evalCtx"], symbolize_names: true) : {} + ) + end + end + end + end +end diff --git a/lib/schematic/credits/leases/reservation_store.rb b/lib/schematic/credits/leases/reservation_store.rb new file mode 100644 index 0000000..cb8da62 --- /dev/null +++ b/lib/schematic/credits/leases/reservation_store.rb @@ -0,0 +1,161 @@ +# frozen_string_literal: true + +module Schematic + module Credits + module Leases + # Per-process reservation table, paired with a sweep loop that returns + # expired reservations to their underlying lease. + # + # For cross-process deployments use RedisReservationStore instead: both + # answer the same calls with the same semantics. + class ReservationStore + def initialize(lease_store, sweep_interval_ms = DEFAULT_SWEEP_INTERVAL_MS, clock: DEFAULT_CLOCK, + logger: nil) + @lease_store = lease_store + @sweep_interval_ms = sweep_interval_ms + @clock = clock + @logger = logger + @reservations = {} + @mutex = Mutex.new + @sweep_thread = nil + @stopped = false + end + + # Register a new reservation. Idempotent on id. It does NOT debit the + # lease: the debit already happened in try_reserve, and recording the + # hold after the debit is what bounds a crash between them to leaked + # credits rather than a double spend. + def add(reservation) + @mutex.synchronize { @reservations[reservation.id] = reservation } + nil + end + + def get(id) + @mutex.synchronize { @reservations[id] } + end + + # Sum the open reservations for a (company, credit). A reservation + # counts while it is in the table: its credits stay carved out of the + # lease's local_remaining_credits until consume or sweep_expired removes + # it and refunds the unspent remainder in the same step, so + # local_remaining_credits + reserved_credits stays exact. + def reserved_credits(company_id, credit_type_id) + @mutex.synchronize do + @reservations.each_value.sum do |reservation| + if reservation.company_id == company_id && reservation.credit_type_id == credit_type_id + reservation.credits_reserved + else + 0 + end + end + end + end + + # Claim a reservation exactly once and settle it: the clamped consumed + # slice stays debited and the unspent remainder is refunded to the + # lease. Returns the clamped figure, or nil when the reservation was + # already gone (swept at its TTL, or claimed by a racing caller). + # + # The claim is the arbiter. A crash between the claim and the refund + # leaks the unspent slice until the lease expires, which is the safe + # direction: a second claim finds nothing and refunds nothing. + def consume(id, credits_consumed) + reservation = @mutex.synchronize { @reservations.delete(id) } + return nil if reservation.nil? + + actual = credits_consumed.clamp(0, reservation.credits_reserved) + refund = reservation.credits_reserved - actual + if refund.positive? && refundable?(reservation) + # Pinned to the originating lease: if that lease has expired and a + # successor holds the slot, the refund is dropped, because the + # expired lease's remainder was already returned server-side. + @lease_store.refund(reservation.company_id, reservation.credit_type_id, refund, reservation.lease_id) + end + actual + end + + # Start the background sweep loop. Safe to call repeatedly. + def start_sweep + @mutex.synchronize do + # alive?, not presence: a process that forks after start (Puma or + # Unicorn with preload) hands the child a thread object whose thread + # did not survive the fork, and the child would never sweep. + return if @sweep_thread&.alive? || @stopped + + interval = @sweep_interval_ms.to_f / 1000.0 + @sweep_thread = Thread.new do + loop do + sleep(interval) + break if @stopped + + begin + sweep_expired + rescue StandardError => e + # Keep the loop alive: a store blip must not stop the sweeper. + @logger&.debug("Reservation sweep failed: #{e.message}") + end + end + end + @sweep_thread.abort_on_exception = false + end + end + + # Remove every reservation past its TTL and refund its full hold to its + # lease. Returns the number swept. + def sweep_expired(now = nil) + now ||= @clock.call + expired = @mutex.synchronize do + due = @reservations.each_value.select { |r| r.expires_at.to_f <= now.to_f } + due.each { |r| @reservations.delete(r.id) } + due + end + expired.each do |reservation| + next unless refundable?(reservation) + + @lease_store.refund( + reservation.company_id, + reservation.credit_type_id, + reservation.credits_reserved, + reservation.lease_id + ) + end + expired.size + end + + def stop + thread = @mutex.synchronize do + @stopped = true + thread = @sweep_thread + @sweep_thread = nil + thread + end + return nil if thread.nil? + + # A sweep past its claim has deleted the reservation but not yet + # refunded the lease, and killing it there strands the unspent slice + # until the lease expires. Give it a bounded moment to land that + # refund, then kill: the loop checks @stopped as soon as its sleep + # ends, so a parked sweeper costs at most this wait. + thread.kill unless thread.join(SWEEP_STOP_JOIN_MS / 1000.0) + nil + end + + def size + @mutex.synchronize { @reservations.size } + end + + private + + # A hold that cannot name the lease it came out of is not refundable: + # crediting whichever lease holds the slot now could inflate a successor + # whose grant the server issued whole, and the slice comes back when the + # lease expires anyway. Decided here rather than left to the lease + # store, because the two lease stores read a missing pin differently, + # and a fleet mixing SDKs on one Redis must agree. + def refundable?(reservation) + !reservation.lease_id.nil? && !reservation.lease_id.to_s.empty? + end + end + end + end +end diff --git a/lib/schematic/credits/leases/server_check.rb b/lib/schematic/credits/leases/server_check.rb new file mode 100644 index 0000000..0b97cc3 --- /dev/null +++ b/lib/schematic/credits/leases/server_check.rb @@ -0,0 +1,237 @@ +# frozen_string_literal: true + +require "securerandom" +require "time" + +module Schematic + module Credits + module Leases + # Everything a server-mode check needs. reservation_ttl_ms is how far out + # the hold's expires_at is set, and default_value resolves the caller's + # default for the flag, which the fail-open branch returns because there + # is no local engine to re-run. + ServerCheckDeps = Struct.new(:features, :credits, :logger, :reservation_ttl_ms, :default_value, :clock, + keyword_init: true) + + # Mirrors the reason the API returns on a 200 with value false for the + # same denial, so a caller matching on reason has one string to match + # either way. + INSUFFICIENT_CREDITS_REASON = "Insufficient credits" + + # Drive a single check with usage set, in server mode. + # + # One check-and-reserve call does everything the client path spreads + # across a lease acquire, a local reserve, and a rules evaluation: the + # server evaluates the flag against the company's real balance, applies + # the preflight cost, and takes the hold in the same round trip. There is + # no lease, no local store, and no rules engine involved. + # + # The failure contract differs from client mode in one place. fail-open + # there means re-run the engine with the credit balance assumed + # sufficient, so plan targeting and every non-credit condition still + # apply. Server mode has no local engine to re-run, since the call that + # would have answered is the one that failed, so fail-open returns the + # caller's default value instead. fail-closed denies, same as client mode. + # + # No flag_check event is enqueued here: the server logs the flag check for + # check-and-reserve itself, the same way the REST check path does. + def self.check_with_server_reservation(deps, key, eval_ctx, options, &fallback) + ServerCheck.new(deps, key, eval_ctx, options, fallback).run + end + + class ServerCheck + def initialize(deps, key, eval_ctx, options, fallback) + @deps = deps + @key = key + @eval_ctx = eval_ctx || {} + @options = options || {} + @fallback = fallback + @logger = deps.logger + @clock = deps.clock || DEFAULT_CLOCK + @on_failure = Leases.resolve_failure_mode(@options[:on_acquire_failure], @logger) + # One key for this check, minted before the call rather than per + # attempt: check-and-reserve takes a hold, so a 502 arriving after the + # API committed one would otherwise have the retry take a second and + # park the first until its TTL. Sharing the key across attempts makes + # the server return the hold it already took. + @idempotency_key = SecureRandom.uuid + end + + def run + usage = @options[:usage] + # The same guard as the client path: a malformed usage must never + # reach the wire. NaN slips through every numeric comparison, so the + # server would size a hold off a value no comparison can reject. + unless Leases.valid_quantity?(usage) + @logger.error( + "Server reservation: invalid usage #{usage.inspect} for flag #{@key}, " \ + "must be a finite non-negative number" + ) + return failure_result("invalid_usage") + end + + if usage.zero? + @logger.debug("Server reservation: usage is 0 for flag #{@key}, nothing to reserve, using plain check") + return @fallback.call + end + + data = call_api + return data if data.is_a?(CheckResult) + + build_result(data) + end + + private + + def call_api + response = @deps.features.check_and_reserve_flag(request_options: request_options, **request_body) + response.data + rescue StandardError => e + # A 402 is the server's definitive answer, not a can't-gate: it knows + # the credits are not there. Deny regardless of the failure mode, + # since failing open here would hand out credit the balance cannot + # cover. check-and-reserve itself answers 200 with value false for + # insufficient credits; this is defensive. + return payment_required_result(e) if payment_required?(e) + + @logger.error("Server reservation: check-and-reserve for flag #{@key} failed: #{e.message}") + failure_result("server_reservation_failed") + end + + def request_body + # The request body's quantity is an integer, so a fractional usage + # would truncate and the server would size the hold below the work + # about to run. Round up, and size the preflight from the same number + # so the flag is evaluated against the quantity actually held. + quantity = Leases.wire_quantity(@options[:usage]) + body = { + key: @key, + quantity: quantity, + expires_at: (@clock.call + (@deps.reservation_ttl_ms / 1000.0)).utc.iso8601 + } + company = @eval_ctx[:company] || @eval_ctx["company"] + user = @eval_ctx[:user] || @eval_ctx["user"] + body[:company] = company if company && !company.empty? + body[:user] = user if user && !user.empty? + preflight = Leases.build_preflight_options(@options.merge(usage: quantity)) + body[:preflight] = preflight if preflight + body[:idempotency_key] = @idempotency_key + body + end + + # Only the timeout is set here. The call keeps the client's default + # retry policy, which the idempotency key makes safe. + def request_options + return {} if @options[:timeout_ms].nil? + + { timeout_in_seconds: @options[:timeout_ms] / 1000.0 } + end + + def build_result(data) + base = CheckResult.new( + allowed: data.value, value: data.value, reason: data.reason, entitlement: data.entitlement, + flag_key: data.flag || @key, flag_id: data.flag_id, error: data.error + ) + held = data.reservation + # No reservation comes back when the flag denied, the credits were + # insufficient (a 200 with value false), or the feature is not + # credit-metered. Nothing was held, so there is nothing to release. + return base if !data.value || held.nil? + + # The settling track event is named by the event subtype; the caller's + # explicit one wins, otherwise the server names it on the hold. With + # neither, the hold could never be settled. + event_subtype = @options[:event_subtype] || held.event_subtype + return release_unsettleable(held, base) if event_subtype.nil? || event_subtype.empty? + + CheckResult.new( + allowed: true, value: true, reason: data.reason, entitlement: data.entitlement, + flag_key: data.flag || @key, flag_id: data.flag_id, + reservation: reservation_from(held, event_subtype) + ) + end + + def reservation_from(held, event_subtype) + Reservation.new( + id: held.id, + # No lease exists in server mode; mirror the id so the field stays + # populated and a handle round-trips through code that reads it. + lease_id: held.id, + mode: :server, + company_id: held.company_id, + credit_type_id: held.credit_type_id, + event_subtype: event_subtype, + quantity_reserved: held.quantity_reserved, + credits_reserved: held.credits_reserved, + consumption_rate: held.consumption_rate, + expires_at: parse_time(held.expires_at), + eval_ctx: @eval_ctx + ) + end + + def release_unsettleable(held, base) + @logger.error( + "Server reservation: reservation #{held.id} for flag #{@key} has no event subtype, " \ + "releasing, it could never be settled" + ) + begin + @deps.credits.release_credit_reservation(reservation_id: held.id) + rescue StandardError => e + @logger.warn( + "Server reservation: failed to release #{held.id} (#{e.message}); its hold is refunded when it expires" + ) + end + return failure_result("missing_event_subtype") if @on_failure == :fail_closed + + # Fail-open means assume the credits are there, and the server has + # already evaluated the flag and allowed this check. Only the settle + # is impossible, so keep the server's verdict rather than falling back + # to the caller's default, which could deny what the server allowed. + CheckResult.new( + allowed: base.allowed, value: base.value, reason: base.reason, entitlement: base.entitlement, + flag_key: base.flag_key, flag_id: base.flag_id, error: "missing_event_subtype" + ) + end + + # Resolve a can't-gate outcome. fail-closed denies; fail-open returns + # the caller's default value, since there is no local engine to + # re-evaluate with an assumed-sufficient balance the way client mode + # does. + def failure_result(reason) + return CheckResult.new(allowed: false, value: false, reason: reason, flag_key: @key, error: reason) if @on_failure == :fail_closed + + value = @deps.default_value.call + CheckResult.new(allowed: value, value: value, reason: "#{reason}_fail_open", flag_key: @key, error: reason) + end + + # The generated client has no 402-specific error class, so a payment + # required arrives as a ClientError carrying the status code. + def payment_required?(error) + error.respond_to?(:code) && error.code.to_i == 402 + end + + def payment_required_result(error) + CheckResult.new( + allowed: false, value: false, reason: INSUFFICIENT_CREDITS_REASON, flag_key: @key, + error: api_error_message(error) + ) + end + + # A response error's message is the raw body, which carries the API's + # own error string when it is JSON. + def api_error_message(error) + parsed = JSON.parse(error.message.to_s) + parsed.is_a?(Hash) && parsed["error"] ? parsed["error"] : error.message + rescue StandardError + error.message + end + + def parse_time(value) + Leases.parse_api_time(value) + rescue StandardError + @clock.call + end + end + end + end +end diff --git a/lib/schematic/credits/leases/track.rb b/lib/schematic/credits/leases/track.rb new file mode 100644 index 0000000..dfd547f --- /dev/null +++ b/lib/schematic/credits/leases/track.rb @@ -0,0 +1,66 @@ +# frozen_string_literal: true + +module Schematic + module Credits + module Leases + # What a settle did locally, and what it owes the server. + # + # settled_locally is true when the hold was still open and this call + # debited the consumed slice and refunded the rest. False when it had + # already been swept at its TTL, already settled, or the store was + # unreachable: the lease balance was not touched here, so it reads high + # until the lease rolls over, and the event is a recovery emit. + SettleOutcome = Struct.new(:track, :settled_locally, keyword_init: true) + + # Consume a reservation against its lease and build the event that bills + # it. + # + # The event is built from the caller-held handle rather than the store, so + # the usage is still billed once the hold has been swept. Only the local + # bookkeeping clamps to the reserved amount; the event carries the + # unclamped actual, because the server is the source of truth for real + # consumption. + def self.consume_reservation_and_build_event(reservations, reservation, actual_quantity, traits: nil) + # Rounded up for the same reason the hold is (see check_with_lease), and + # through the same helper the event's quantity uses: the debit has to + # move the local ledger by exactly what the Track event bills. + consumed = reservations.consume( + reservation.id, Leases.wire_quantity(actual_quantity) * reservation.consumption_rate + ) + SettleOutcome.new( + track: build_reservation_track_event(reservation, actual_quantity, traits: traits), + settled_locally: !consumed.nil? + ) + end + + # Build the track event for a reservation from the handle alone, with no + # store access, so the client can still bill the usage when the local + # settle fails against an unreachable store. + def self.build_reservation_track_event(reservation, actual_quantity, traits: nil) + # Whole event units, the same rounding the hold and the settle debit + # use: the API rejects a non-integer quantity during processing, and + # the local ledger has to move by what this event bills. + body = { event: reservation.event_subtype, quantity: Leases.wire_quantity(actual_quantity) } + if reservation.server_mode? + # The hold lives on the server, so the event settles it by id. Never + # send lease_id too: the server prefers it when both are set, and + # there is no lease here for it to route through. + body[:reservation_id] = reservation.id + else + # Routes the server-side credit consumption through the lease's + # sub-ledger instead of decrementing the grant again, which the + # acquire already pre-debited. Without this the grant double-debits + # and eventually starves redemptions mid-session. + body[:lease_id] = reservation.lease_id + end + eval_ctx = reservation.eval_ctx || {} + company = eval_ctx[:company] || eval_ctx["company"] + user = eval_ctx[:user] || eval_ctx["user"] + body[:company] = company if company && !company.empty? + body[:user] = user if user && !user.empty? + body[:traits] = traits if traits + body + end + end + end +end diff --git a/lib/schematic/credits/leases/types.rb b/lib/schematic/credits/leases/types.rb new file mode 100644 index 0000000..92f9823 --- /dev/null +++ b/lib/schematic/credits/leases/types.rb @@ -0,0 +1,329 @@ +# frozen_string_literal: true + +require "time" + +module Schematic + module Credits + # Client-side credit leases, reservations, and preflight checks. + # + # A credit-metered feature is gated without a wire call per check: the SDK + # leases a tranche of credits per (company, credit type), carves a + # reservation out of it at check time, and settles the reservation against + # the actual usage when the work finishes. See conformance/SPEC.md at the + # repo root for the full model; the vectors there pin the semantics every + # Schematic SDK shares. + module Leases + # Durations are milliseconds throughout, matching the names and units the + # other SDKs use so one set of numbers configures a mixed fleet. + + DEFAULT_LEASE_DURATION_MS = 5 * 60 * 1000 + DEFAULT_RESERVATION_TTL_MS = 60 * 1000 + # The API rejects a hold whose expires_at is more than an hour out, so a + # larger reservation TTL would fail every server-mode check. The SDK + # clamps to this instead. + MAX_RESERVATION_TTL_MS = 60 * 60 * 1000 + # The API measures that hour against its own clock while the SDK computes + # expires_at against the caller's, so a client running ahead would be + # rejected at exactly the cap. Hold this much back from it. + RESERVATION_TTL_SKEW_ALLOWANCE_MS = 60 * 1000 + DEFAULT_LEASE_SIZE = 10_000 + DEFAULT_LOW_WATER_MARK = 0.25 + DEFAULT_SWEEP_INTERVAL_MS = 1000 + # How long prewarm is willing to wait for a freshly identified company to + # surface in the datastream cache before giving up. Long enough to cover + # the buffer-flush, server-ingest, datastream-push round trip for a new + # company; short enough that a misconfigured caller does not hang. + DEFAULT_PREWARM_RESOLVE_TIMEOUT_MS = 5000 + DEFAULT_PREWARM_POLL_INTERVAL_MS = 100 + # How long close waits for in-flight lease work to land before giving up + # on it. Bounded on purpose: a shutdown that hangs is worse than a hold + # the server expires at DEFAULT_LEASE_DURATION_MS. + SHUTDOWN_DRAIN_TIMEOUT_MS = 5000 + + # 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 + + # Balance substituted for a fail-open evaluation: large enough that the + # credit gate always passes, and the same figure the other SDKs use so a + # shared vector can name it. + FAIL_OPEN_BALANCE = (2**53) - 1 + + # Where a credit hold lives for a check that passes usage. + # :client - local leases over DataStream. + # :server - one check-and-reserve API call per check. + # :auto - client when DataStream is enabled, server otherwise. + MODES = %i[client server auto].freeze + + # What to do when a lease cannot be acquired or reserved against. + # :fail_open - re-run the engine with the credit balance assumed + # sufficient, so non-credit rules still apply. + # :fail_closed - deny, because the gate cannot gate. + FAILURE_MODES = %i[fail_open fail_closed].freeze + + # Reads the current time. Every store and the lease manager take one so + # tests and the conformance runner can drive a virtual clock instead of + # wall time. + DEFAULT_CLOCK = -> { Time.now } + + # Convert a caller-supplied mode to a symbol, accepting the hyphenated + # spellings the other SDKs use ("fail-open") so one config shape travels. + def self.normalize_symbol(value) + return nil if value.nil? + + value.to_s.tr("-", "_").downcase.to_sym + end + + # How long stop waits for the sweeper to finish what it is doing before + # killing it. Long enough for a sweep already past its claim to land its + # refund, short enough that close stays prompt. + SWEEP_STOP_JOIN_MS = 100 + + # Resolve a caller's on_acquire_failure, or fail closed. + # + # An unrecognized value must not read as fail-open: the branches all test + # for :fail_closed, so a typo would quietly turn the safe default into the + # permissive one on every check. Name the value in the warning, since a + # silent downgrade of a gate is the thing worth telling someone about. + def self.resolve_failure_mode(value, logger = nil) + return :fail_closed if value.nil? + + mode = normalize_symbol(value) + return mode if FAILURE_MODES.include?(mode) + + logger&.warn( + "Unrecognized on_acquire_failure #{value.inspect}; expected one of " \ + "#{FAILURE_MODES.join(" or ")}. Failing closed." + ) + :fail_closed + end + + # Whether a caller-supplied quantity can size a credit hold. NaN is the + # dangerous case: it slips through every numeric comparison, and a NaN + # balance would approve every later reserve on a possibly shared lease. + def self.valid_quantity?(value) + value.is_a?(Numeric) && !value.to_f.nan? && value.to_f.finite? && value >= 0 + end + + # The local view of the one lease a (company, credit type) slot holds. + class LeaseEntry + attr_accessor :lease_id, :company_id, :credit_type_id, :granted_amount, + :local_remaining_credits, :expires_at + + def initialize(lease_id:, company_id:, credit_type_id:, granted_amount:, expires_at:, + local_remaining_credits: nil) + @lease_id = lease_id + @company_id = company_id + @credit_type_id = credit_type_id + @granted_amount = granted_amount.to_f + @local_remaining_credits = (local_remaining_credits || granted_amount).to_f + @expires_at = expires_at + end + + def dup + LeaseEntry.new( + lease_id: @lease_id, + company_id: @company_id, + credit_type_id: @credit_type_id, + granted_amount: @granted_amount, + local_remaining_credits: @local_remaining_credits, + expires_at: @expires_at + ) + end + + def expired?(now) + @expires_at.to_f <= now.to_f + end + end + + # The post-debit balance a successful try_reserve returns, plus the id of + # the lease the credits actually came out of. The debit is not keyed by + # lease id, so the caller pins its reservation to this id and never to the + # one its acquire handed back. + ReserveResult = Struct.new(:balance, :lease_id) + + # One credit hold carved out of a lease by a check. Returned to the caller + # from check and handed back to track_with_reservation. + class Reservation + attr_reader :id, :lease_id, :mode, :company_id, :credit_type_id, :event_subtype, + :quantity_reserved, :credits_reserved, :consumption_rate, :expires_at, :eval_ctx + + def initialize(id:, lease_id:, company_id:, credit_type_id:, event_subtype:, + quantity_reserved:, credits_reserved:, consumption_rate:, expires_at:, + eval_ctx: {}, mode: nil) + @id = id + @lease_id = lease_id + # :server means the API holds the credits and the settling track event + # routes by reservation_id; nil (or :client) means the hold is a local + # carve-out of a lease. + @mode = mode + @company_id = company_id + @credit_type_id = credit_type_id + @event_subtype = event_subtype + @quantity_reserved = quantity_reserved.to_f + @credits_reserved = credits_reserved.to_f + @consumption_rate = consumption_rate.to_f + @expires_at = expires_at + @eval_ctx = eval_ctx || {} + end + + def server_mode? + @mode == :server + end + + def to_h + { + id: @id, + lease_id: @lease_id, + mode: @mode, + company_id: @company_id, + credit_type_id: @credit_type_id, + event_subtype: @event_subtype, + quantity_reserved: @quantity_reserved, + credits_reserved: @credits_reserved, + consumption_rate: @consumption_rate, + expires_at: @expires_at, + eval_ctx: @eval_ctx + } + end + end + + # Cast a usage onto the integer the wire carries. A hold can be sized from + # a fractional usage, but every quantity field on the API (the + # check-and-reserve ask, the preflight envelope, a track event) is an + # integer, and the generated models truncate a float onto it. A preflight + # asks an upper-bound question and a settle must not bill a partial unit + # as none, so a fraction rounds up in both directions. + # Float noise is shaved off before the rounding: (0.1 + 0.2) * 10 is + # 3.0000000000000004, and a bare ceil would bill that as 4. + ROUND_UP_TOLERANCE = 1e-9 + + def self.wire_quantity(value) + return value unless value.is_a?(Numeric) + return value if value.is_a?(Integer) + return value unless value.finite? + + shaved = value - ROUND_UP_TOLERANCE + shaved.positive? ? shaved.ceil : value.ceil + end + + # One shape for the matched entitlement whichever mode produced it. + # + # The WASM engine hands back a camelCase hash and the API hands back a + # generated model, so without this a caller reading result.entitlement + # would need one accessor for client mode and another for server mode. + # Both become a snake_case, symbol-keyed Hash matching the field names on + # Schematic::Types::FeatureEntitlement, and metric_reset_at is parsed to a + # Time so a caller can compare it without knowing which mode it came from. + def self.normalize_entitlement(raw) + return nil if raw.nil? + + hash = raw.is_a?(Hash) ? raw : entitlement_to_h(raw) + return nil if hash.nil? + + normalized = deep_snake_case(hash) + reset_at = normalized[:metric_reset_at] + normalized[:metric_reset_at] = parse_reset_at(reset_at) unless reset_at.nil? + normalized + end + + def self.entitlement_to_h(raw) + return raw.to_h if raw.respond_to?(:to_h) + + nil + end + private_class_method :entitlement_to_h + + def self.parse_reset_at(value) + parse_api_time(value) + rescue StandardError + value + end + private_class_method :parse_reset_at + + # The API's timestamps are UTC, but Time.iso8601 reads one that carries + # no offset as host-local time, which would shift a lease's expiry by the + # host's UTC offset. A timestamp with no zone is read as UTC instead. + def self.parse_api_time(value) + return value.getutc if value.is_a?(Time) + + text = value.to_s + text += "Z" if text.include?("T") && !text.match?(/(?:Z|[+-]\d{2}(?::?\d{2})?)\z/i) + Time.iso8601(text).utc + end + + def self.deep_snake_case(value) + case value + when Hash + value.each_with_object({}) do |(key, inner), out| + out[key.to_s.gsub(/([a-z\d])([A-Z])/, '\1_\2').downcase.to_sym] = deep_snake_case(inner) + end + when Array + value.map { |inner| deep_snake_case(inner) } + else + value + end + end + private_class_method :deep_snake_case + + # What a lease-aware check decided. `allowed` is what the caller gates on; + # `reservation` is present only when a hold was taken. + class CheckResult + attr_reader :allowed, :value, :reservation, :reason, :entitlement, :flag_key, :flag_id, :error + + def initialize(allowed:, value:, reason:, flag_key:, reservation: nil, entitlement: nil, + flag_id: nil, error: nil) + @allowed = allowed + @value = value + @reservation = reservation + @reason = reason + @entitlement = Leases.normalize_entitlement(entitlement) + @flag_key = flag_key + @flag_id = flag_id + @error = error + end + + def allowed? + @allowed + end + + def to_h + { + allowed: @allowed, + value: @value, + reservation: @reservation&.to_h, + reason: @reason, + entitlement: @entitlement, + flag_key: @flag_key, + flag_id: @flag_id, + error: @error + }.compact + end + end + + # The four resolvable knobs for one credit type, after overrides and + # defaults. + ResolvedLeaseConfig = Struct.new(:lease_duration_ms, :reservation_ttl_ms, :lease_size, :low_water_mark, + keyword_init: true) + + # Resolve the knobs for one credit type: the credit type's override wins, + # then the client-wide config, then the default. + def self.resolve_config(config, credit_type_id) + config ||= {} + overrides = config[:overrides] || {} + override = overrides[credit_type_id] || overrides[credit_type_id.to_s] || + overrides[credit_type_id.to_sym] || {} + ResolvedLeaseConfig.new( + lease_duration_ms: override[:default_lease_duration] || config[:default_lease_duration] || + DEFAULT_LEASE_DURATION_MS, + reservation_ttl_ms: override[:default_reservation_ttl] || config[:default_reservation_ttl] || + DEFAULT_RESERVATION_TTL_MS, + lease_size: override[:default_lease_size] || config[:default_lease_size] || DEFAULT_LEASE_SIZE, + low_water_mark: override[:low_water_mark] || config[:low_water_mark] || DEFAULT_LOW_WATER_MARK + ) + end + end + end +end diff --git a/lib/schematic/credits/leases/wire_client.rb b/lib/schematic/credits/leases/wire_client.rb new file mode 100644 index 0000000..a1a489b --- /dev/null +++ b/lib/schematic/credits/leases/wire_client.rb @@ -0,0 +1,83 @@ +# frozen_string_literal: true + +require "securerandom" +require "time" + +module Schematic + module Credits + module Leases + # What the server says a lease is after an acquire or an extend. It is + # also what installs a lease into a store: the local balance is derived, + # never supplied. + LeaseGrant = Struct.new(:lease_id, :company_id, :credit_type_id, :granted_amount, :expires_at, + keyword_init: true) + + # The lease lifecycle over the generated API client. Separated from + # LeaseManager so tests and the conformance runner can script the wire. + # + # Both calls take the client's default retry policy. Acquire is safe to + # retry because the server hands back the slot's existing active lease + # rather than opening a second one. An extend is an increment, so a retry + # after a lost response would grant the tranche twice; the request carries + # an idempotency key minted once per extend to collapse them. + class ApiWireClient + def initialize(credits_client:) + @credits = credits_client + end + + def acquire(company_id:, credit_type_id:, requested_amount:, expires_at:, request_options: {}) + response = @credits.acquire_credit_lease( + request_options: request_options, + company_id: company_id, + credit_type_id: credit_type_id, + # Rounded up, not to nearest: the API takes a whole number, and a + # shortfall of 10.4 asked for as 10 leaves the retried reserve short + # by the same fraction every time. + requested_amount: Leases.wire_quantity(requested_amount), + expires_at: expires_at.utc.iso8601 + ) + grant_from(response) + end + + def extend(lease_id:, additional_amount:, expires_at:, idempotency_key: nil, request_options: {}) + body = { + lease_id: lease_id, + additional_amount: Leases.wire_quantity(additional_amount), + expires_at: expires_at.utc.iso8601, + # The key is minted once here, outside the retry loop, so every + # attempt of this extend carries the same one and the server folds + # them into a single increment. A caller that already minted one + # keeps it: the manager does, so its single-flight followers settle + # against the same key. + idempotency_key: idempotency_key || SecureRandom.uuid + } + grant_from(@credits.extend_credit_lease(request_options: request_options, **body)) + end + + def release(lease_id:, request_options: {}) + @credits.release_credit_lease(request_options: request_options, lease_id: lease_id) + nil + end + + private + + def grant_from(response) + data = response&.data + raise "credit lease response carried no data" if data.nil? + + LeaseGrant.new( + lease_id: data.id, + company_id: data.company_id, + credit_type_id: data.credit_type_id, + granted_amount: data.granted_amount.to_f, + expires_at: parse_time(data.expires_at) + ) + end + + def parse_time(value) + Leases.parse_api_time(value) + end + end + end + end +end diff --git a/lib/schematic/datastream/client.rb b/lib/schematic/datastream/client.rb index 61b832c..ef6c76e 100644 --- a/lib/schematic/datastream/client.rb +++ b/lib/schematic/datastream/client.rb @@ -146,7 +146,11 @@ def close # Raises DataStream::EvaluationError when the flag cannot be evaluated # locally (e.g. flag not in cache, rules engine unavailable), so the caller # can fall back to the API path. - def check_flag(eval_ctx, flag_key) + # options carries the caller's preflight (a simulated usage, an + # event-scoped usage, or a per-credit cost) through to the engine, so a + # client-side evaluation gates on the post-call balance. Nil evaluates the + # flag as it stands. + def check_flag(eval_ctx, flag_key, options = nil) flag = get_flag(flag_key) raise EvaluationError, "flag '#{flag_key}' not found in cache" unless flag @@ -172,12 +176,22 @@ def check_flag(eval_ctx, flag_key) end @logger.debug("Evaluating flag with rules engine: flag=#{flag_key}, company=#{company&.dig(:id)}, user=#{user&.dig(:id)}") - result = @rules_engine.check_flag(flag, company, user) + result = @rules_engine.check_flag_with_options(flag, company, user, options) @logger.debug("Rules engine evaluation result: value=#{result[:value]}, reason=#{result[:reason]}") result[:flag_key] = flag_key result end + # Evaluate a flag against an explicit company and user with preflight + # options, skipping the cache lookups check_flag does. The credit-lease + # path has already resolved both entities and needs to gate against a + # substituted credit balance, which only it can build. + def check_flag_with_options(flag, company, user, options) + raise EvaluationError, "rules engine not available" unless @rules_engine&.initialized? + + @rules_engine.check_flag_with_options(flag, company, user, options) + end + def update_company_metrics(company_keys, event_name, quantity) company = @company_cache.get_by_keys(company_keys) return unless company @@ -231,6 +245,10 @@ def get_company(keys) return cached if cached return nil if @replicator_mode + # A request sent over a closed socket is dropped without an error, so + # waiting on it just burns the resource timeout. Answer now and let the + # caller fall back. + return nil unless connected? request_entity(ENTITY_TYPE_COMPANY, keys, @pending_companies) do |_data| @company_cache.get_by_keys(keys) @@ -242,6 +260,7 @@ def get_user(keys) return cached if cached return nil if @replicator_mode + return nil unless connected? request_entity(ENTITY_TYPE_USER, keys, @pending_users) do |_data| @user_cache.get_by_keys(keys) @@ -513,6 +532,17 @@ def check_replicator_health @logger.info("Replicator is no longer ready") if was_ready @logger.warn("Replicator health check error: #{e.message}") end + + # The cached company for these keys, without the fetch get_company falls + # back to. prewarm polls this to resolve a company id from secondary keys. + def get_cached_company(keys) + @company_cache.get_by_keys(keys) + end + + # Cache-first flag and entity lookups, declared public after their + # definitions. The credit-lease check path resolves the flag, the company, + # and the user itself before evaluating, so it needs all three. + public :get_flag, :get_company, :get_user, :get_cached_company end end end diff --git a/lib/schematic/rules_engine.rb b/lib/schematic/rules_engine.rb index 97c5cfd..679592e 100644 --- a/lib/schematic/rules_engine.rb +++ b/lib/schematic/rules_engine.rb @@ -73,12 +73,26 @@ def initialized? end def check_flag(flag, company = nil, user = nil) + check_flag_with_options(flag, company, user, nil) + end + + # Evaluate a flag with preflight options: a simulated usage, an event-scoped + # usage, or a pre-computed per-credit cost. The engine applies them to the + # conditions they match without mutating any state, so a caller can ask + # "would this call still be allowed after it lands". + # + # options is a hash with any of :credit_cost (credit id to cost), + # :usage, and :event_usage ({ event_subtype:, quantity: }). The engine reads + # them snake_case in both directions, so they go on the envelope as given. + def check_flag_with_options(flag, company = nil, user = nil, options = nil) raise "WASM rules engine not initialized" unless @initialized # Build combined JSON envelope (same format as Python/C#) envelope = { flag: strip_nulls(flag) } envelope[:company] = strip_nulls(company) if company envelope[:user] = strip_nulls(user) if user + # Omitted entirely when empty, so the engine uses its own defaults. + envelope[:options] = strip_nulls(engine_options(options)) if options && !options.empty? json_bytes = JSON.generate(envelope).encode("UTF-8") @@ -123,6 +137,29 @@ def version_key private + # The options as the engine takes them. usage and event_usage.quantity + # deserialize as i64 there, so a value with a decimal point fails the whole + # check. Rounded through the helper every other wire field uses, so a caller + # passing a preflight straight to check_flag_with_entitlement is asked the + # same question a lease check asks, and an infinite usage is handed over + # rather than raising out of ceil. + def engine_options(options) + out = options.dup + usage_key = out.key?(:usage) ? :usage : "usage" + out[usage_key] = Credits::Leases.wire_quantity(out[usage_key]) if out[usage_key].is_a?(Numeric) + event_key = out.key?(:event_usage) ? :event_usage : "event_usage" + event_usage = out[event_key] + return out unless event_usage.is_a?(Hash) + + quantity_key = event_usage.key?(:quantity) ? :quantity : "quantity" + return out unless event_usage[quantity_key].is_a?(Numeric) + + out[event_key] = event_usage.merge( + quantity_key => Credits::Leases.wire_quantity(event_usage[quantity_key]) + ) + out + end + def export_func(name) exp = @instance.export(name) raise "WASM export '#{name}' not found" unless exp diff --git a/lib/schematic/schematic_client.rb b/lib/schematic/schematic_client.rb index fd74b7f..65fddd2 100644 --- a/lib/schematic/schematic_client.rb +++ b/lib/schematic/schematic_client.rb @@ -59,6 +59,24 @@ class SchematicClient # Optional event metadata accepted via the `options:` keyword on track/identify. # identify only honors :idempotency_key; track also honors :sent_at, # :trusted_client_clock, and :backfill. Fields are only sent when set. + # Namespaces the idempotency key on the track event a reservation settles + # into. Deterministic per reservation, so a recovery emit (the work outlived + # the local reservation TTL) and an accidental double settle collapse to one + # billed event: the pipeline drops duplicates for 24h before any credit + # consumption runs. + RESERVATION_TRACK_IDEMPOTENCY_PREFIX = "lease-reservation:" + + # The prefix Schematic's secure company ids carry, whatever key name they + # are passed under. + COMPANY_ID_PREFIX = "comp_" + + # Knobs that only steer the local lease plumbing, which server mode never + # builds. Setting one there does nothing, so the client says so at startup. + CLIENT_ONLY_LEASE_OPTIONS = %i[ + default_lease_duration default_lease_size low_water_mark sweep_interval_ms + redis_client redis_key_prefix prewarm_resolve_timeout_ms overrides + ].freeze + TRACK_OPTION_KEYS = %i[idempotency_key sent_at trusted_client_clock backfill].freeze IDENTIFY_OPTION_KEYS = %i[idempotency_key].freeze @@ -72,6 +90,7 @@ def initialize( event_capture_base_url: nil, use_data_stream: false, datastream_options: {}, + credit_leases: nil, logger: nil, log_level: :warn ) @@ -89,6 +108,13 @@ def initialize( end @offline = offline + # Validated here rather than alongside the rest of the lease setup below: + # a rejected knob raises out of the constructor, so the caller never gets + # a client to close, and anything already running by then (the event + # buffer's flush thread, the DataStream socket) is leaked for the life of + # the process. Nothing has started yet at this point. + validated_leases = validated_credit_lease_config(credit_leases) + # Initialize Fern-generated API client @api_client = if @offline nil @@ -118,6 +144,25 @@ def initialize( @rules_engine = nil setup_datastream(datastream_options) if use_data_stream && !@offline + # Credit lease + reservation plumbing, if the caller opted in. + @credit_lease_config = nil + @credit_lease_mode = nil + @credit_lease_manager = nil + @lease_store = nil + @reservations = nil + # True when lease state lives in a shared backend that sibling processes + # may also be drawing on. close must then NOT release leases. + @lease_backend_shared = false + @server_reservation_ttl_ms = Credits::Leases::DEFAULT_RESERVATION_TTL_MS + @prewarm_resolve_timeout_ms = Credits::Leases::DEFAULT_PREWARM_RESOLVE_TIMEOUT_MS + # Prewarms identify spawned and nobody joins. close waits them out: an + # acquire that lands after the release installs a lease nothing releases, + # and its credits stay held until the server expires them. + @pending_prewarms = [] + @pending_prewarms_mutex = Mutex.new + @closing = false + setup_credit_leases(validated_leases, datastream_options) if credit_leases + # Register shutdown hook to ensure graceful cleanup on process exit at_exit { close } end @@ -128,11 +173,19 @@ def check_flag(flag_key, company: nil, user: nil) check_flag_with_entitlement(flag_key, company: company, user: user).value end - def check_flag_with_entitlement(flag_key, company: nil, user: nil) + # default_value overrides the registered flag default on every path that + # cannot answer from the flag itself: offline, an API error, or a response + # with no value. Leaving it nil keeps the registered default. timeout_ms is + # threaded to the API call, but see the note on check: the generated + # transport does not yet apply a per-request timeout. + def check_flag_with_entitlement(flag_key, company: nil, user: nil, preflight: nil, default_value: nil, + timeout_ms: nil) + get_default = -> { resolve_default_value(flag_key, default_value) } + # Offline mode if @offline return CheckFlagResponse.new( - value: get_flag_default(flag_key), + value: get_default.call, flag_key: flag_key, reason: "offline mode" ) @@ -142,8 +195,22 @@ def check_flag_with_entitlement(flag_key, company: nil, user: nil) if @datastream_client&.connected? begin eval_ctx = build_eval_context(company, user) - result = @datastream_client.check_flag(eval_ctx, flag_key) - + # Only widen the call when there is something to pass: a DataStream + # double written against the two-argument form still works, and the + # preflight envelope reaches the engine when a lease check supplies it. + result = if preflight.nil? + @datastream_client.check_flag(eval_ctx, flag_key) + else + @datastream_client.check_flag(eval_ctx, flag_key, preflight) + end + + # A nil value is the engine declining to answer, not a false. The + # caller's default stands in, falling back to the registered one, + # which is what CheckFlagResponse's own coercion to false would + # otherwise hide. Resolved the same way the offline and API paths + # resolve it, or one call would honour default_value and another + # would not. + result[:value] = get_default.call if result[:value].nil? response = CheckFlagResponse.new(result) enqueue_flag_check_event(flag_key, response, company, user) return response @@ -155,11 +222,12 @@ def check_flag_with_entitlement(flag_key, company: nil, user: nil) end # API path with caching - check_flag_via_api(flag_key, company, user) + check_flag_via_api(flag_key, company, user, preflight: preflight, timeout_ms: timeout_ms, + get_default: get_default) rescue StandardError => e @logger.error("check_flag_with_entitlement error for '#{flag_key}': #{e.message}") CheckFlagResponse.new( - value: get_flag_default(flag_key), + value: get_default.call, flag_key: flag_key, reason: "error: #{e.message}" ) @@ -255,29 +323,216 @@ def check_flags(company: nil, user: nil, keys: nil) end end - # --- Event Submission --- + # --- Credit-aware Flag Checking --- + + # Credit-aware feature check. With credit_leases configured and a usage + # passed (optionally qualified by an event_subtype), this gates the check + # against the company's credit balance and returns a reservation handle on + # success. Hand that handle to track_with_reservation when the work + # completes. + # + # In client mode (DataStream enabled) the hold is carved out of a local + # lease and the flag is evaluated by the WASM engine. In server mode it is a + # single check-and-reserve API call that evaluates the flag and takes the + # hold server-side. credit_leases[:mode] picks; the default, :auto, uses + # client mode when DataStream is enabled and server mode otherwise. + # + # Without credit_leases (or without a usage) this falls through to a plain + # flag check and returns a result with no reservation. The caller's + # preflight is still threaded through that plain check, so any client-side + # evaluation path gates on the post-call balance, just without a + # reservation, and the REST path sends the preflight too. default_value + # governs that fallback too, so a check that cannot reach the credit path + # still answers the way the caller asked. + # + # timeout_ms is carried on the API request but not yet applied: the + # generated transport takes its timeout from the construction of its HTTP + # client, and nothing exposes that, so no per-check or client-level timeout + # is configurable today. + def check(flag_key, company: nil, user: nil, usage: nil, event_subtype: nil, on_acquire_failure: nil, + default_value: nil, timeout_ms: nil) + options = { + usage: usage, + event_subtype: event_subtype, + on_acquire_failure: on_acquire_failure, + default_value: default_value, + timeout_ms: timeout_ms + } + eval_ctx = build_eval_context(company, user) + fallback = -> { plain_check_result(flag_key, company, user, options) } + + mode = effective_lease_mode + # With gating configured, the lease and server paths own a malformed + # usage and refuse it by the caller's failure mode. With no gating there + # is nothing to refuse: the value would only reach the preflight builder, + # which the engine and the REST body both take as an integer, so drop it + # and ask the plain question. Ruby has no type to catch it at the boundary + # the way the other SDKs do. + if mode.nil? && !usage.nil? && !Credits::Leases.valid_quantity?(usage) + @logger.warn( + "check: invalid usage #{usage.inspect} for flag #{flag_key}, must be a finite non-negative " \ + "number; continuing without one" + ) + options[:usage] = nil + usage = nil + end + return fallback.call if usage.nil? || mode.nil? + + if mode == :server + return Credits::Leases.check_with_server_reservation( + Credits::Leases::ServerCheckDeps.new( + features: features, credits: credits, logger: @logger, + reservation_ttl_ms: @server_reservation_ttl_ms, + default_value: -> { resolve_default_value(flag_key, default_value) } + ), + flag_key, eval_ctx, options, &fallback + ) + end - def identify(body, options: nil) + # Client mode without the local plumbing (mode: :client and no DataStream) + # keeps the old behavior: a plain, ungated flag check. + return fallback.call unless @credit_lease_manager && @lease_store && @reservations + + Credits::Leases.check_with_lease( + Credits::Leases::CheckDeps.new( + lease_store: @lease_store, reservations: @reservations, manager: @credit_lease_manager, + datastream: @datastream_client, logger: @logger, + # Lease-path checks must stay visible to flag-check analytics and + # company last-seen, the same as every plain check path. + enqueue_flag_check_event: ->(body) { enqueue_lease_flag_check_event(body) } + ), + flag_key, eval_ctx, options, &fallback + ) + end + + # Consume a reservation issued by check. Refunds the unused slice back to + # the lease's local balance and enqueues a track event with the actual + # quantity; the server-side event processor consumes + # actual_quantity x consumption_rate from the company's real credit balance. + # + # A server-mode handle has no local hold to refund: the track event carries + # the reservation id, and the server settles the hold when it processes the + # event. + # + # If the work outlived the reservation's TTL and the sweeper already + # returned the hold to the lease, the local refund has happened but the + # usage must still be billed, so the track is emitted anyway as a recovery + # emit. Double billing is prevented server-side: the track carries a + # deterministic idempotency key derived from the reservation id, and the + # events pipeline drops duplicates for 24h before any credit consumption + # runs. So a recovery emit racing the normal emit, or an accidental second + # settle, collapses to a single billed event, across processes and restarts. + def track_with_reservation(reservation, actual_quantity, traits: nil) return if @offline - @event_buffer.push(build_event("identify", body, options, IDENTIFY_OPTION_KEYS)) - rescue StandardError => e - @logger.error("Error sending identify event: #{e.message}") + # check allows without a hold in several ordinary cases: the feature is + # not credit-metered, the check failed open, usage was 0, or credit leases + # are not configured. Callers pass result.reservation straight through, so + # take the nil and tell them how to bill the usage instead of raising on a + # settle that has nothing to settle. + if reservation.nil? + @logger.error( + "track_with_reservation called without a reservation: the check allowed without taking a hold, " \ + "so there is nothing to settle. Report the usage with track instead." + ) + return + end + + # Mirror the check-path usage guard: a non-finite quantity must reach + # neither the store (clamping against NaN claims the reservation with NO + # refund of the unspent slice) nor the billing event, and a negative one + # would bill negative usage. Skipping the settle leaves the reservation to + # expire at its TTL, where the sweeper refunds the full hold, so no + # credits are lost and nothing bogus is billed. + unless Credits::Leases.valid_quantity?(actual_quantity) + @logger.error( + "track_with_reservation: invalid actual_quantity #{actual_quantity.inspect} for reservation " \ + "#{reservation.id}, must be a finite non-negative number; skipping settle " \ + "(the hold is refunded at its TTL)" + ) + return + end + + settle_reservation(reservation, actual_quantity, traits) end - def track(body, options: nil) + # Pre-warm a credit lease for each given credit type id, so the first check + # against it does not pay the acquire round trip. Failures are logged, never + # raised. + # + # When the company carries only secondary keys (no id), prewarm actively + # fetches it over the datastream, waiting up to + # credit_leases[:prewarm_resolve_timeout_ms], which both resolves the id and + # warms the cache so the first check hits the lease path. + def prewarm(credit_type_ids, company: nil) + if @credit_lease_manager.nil? || @lease_store.nil? + @logger.debug( + effective_lease_mode == :server ? "prewarm is a no-op in server mode, there is no local lease to warm" : "prewarm called but credit_leases is not configured" + ) + return + end + if company.nil? || company.empty? + @logger.debug("prewarm requires a company") + return + end + # Documented as never raising, and a caller reading ids out of config can + # hand over nil or an empty list without meaning to. + if credit_type_ids.nil? || credit_type_ids.empty? + @logger.debug("prewarm requires at least one credit type id") + return + end + if @closing + # close only waits out the prewarms it spawned; a caller invoking + # prewarm directly would otherwise install a lease after the release has + # already listed the store. + @logger.debug("prewarm: client is closing, skipping acquire") + return + end + + company_id = resolve_company_id_with_wait(company) + if company_id.nil? + @logger.debug( + "prewarm: company not resolved within #{@prewarm_resolve_timeout_ms}ms for keys #{company} " \ + "(first check will acquire)" + ) + return + end + + credit_type_ids.each do |credit_type_id| + @credit_lease_manager.acquire_if_needed(company_id, credit_type_id) + rescue StandardError => e + @logger.warn("prewarm: failed to acquire lease for #{credit_type_id}: #{e.message}") + end + nil + end + + # --- Event Submission --- + + # prewarm names credit type ids to acquire leases for in the background once + # the identify event is enqueued. Failures never surface to the caller, and + # it is a no-op unless credit_leases is configured. + def identify(body, options: nil, prewarm: nil) return if @offline - @event_buffer.push(build_event("track", body, options, TRACK_OPTION_KEYS)) + begin + @event_buffer.push(build_event("identify", body, options, IDENTIFY_OPTION_KEYS)) + rescue StandardError => e + @logger.error("Error sending identify event: #{e.message}") + end - # Update company metrics locally if DataStream is active and connected - if @datastream_client&.connected? && body[:company] - event_name = body[:event] || body["event"] - quantity = body[:quantity] || body["quantity"] || 1 - @datastream_client.update_company_metrics(body[:company], event_name, quantity) + return if prewarm.nil? || prewarm.empty? + + begin + prewarm_after_identify(body, prewarm) + rescue StandardError => e + # identify never raises into its caller, and a prewarm is the least of + # the reasons it should start. + @logger.warn("identify prewarm setup failed: #{e.message}") end - rescue StandardError => e - @logger.error("Error sending track event: #{e.message}") + end + + def track(body, options: nil) + emit_track(body, options, update_metrics: true) end # --- Flag Defaults --- @@ -370,12 +625,27 @@ def webhooks # --- Lifecycle --- + # Credit leases: with the per-process in-memory backend this process is the + # only holder of its leases, so they are released here (best-effort), which + # returns their unspent remainder to the company balance immediately instead + # of waiting out the lease expiry. With a shared backend, leases are + # deliberately NOT released: one row per company and credit is shared across + # every SDK instance pointed at that backend, so a single process shutting + # down must not release a lease its siblings are still drawing on. Shared + # leases reclaim themselves by expiring or being fully consumed. + # + # Lease work already in flight is waited out, bounded, before the release, + # so an acquire that lands mid-shutdown is one the release can see. def close return if @closed @closed = true - @event_buffer.stop + @closing = true + shut_down_credit_leases + # DataStream first, then the buffer: its stop flushes, and a flush is the + # last thing that should still be running. @datastream_client&.close + @event_buffer.stop @flag_check_cache_providers.each { |c| c.stop if c.respond_to?(:stop) } @logger.debug("SchematicClient closed") end @@ -393,14 +663,20 @@ def coerce_cached_response(cached) CheckFlagResponse.new(cached) end - def check_flag_via_api(flag_key, company, user) - # Check cache + def check_flag_via_api(flag_key, company, user, preflight: nil, timeout_ms: nil, get_default: nil) + get_default ||= -> { get_flag_default(flag_key) } cache_key = build_cache_key(flag_key, company, user) - @flag_check_cache_providers.each do |provider| - cached = coerce_cached_response(provider.get(cache_key)) - if cached - @logger.debug("Flag '#{flag_key}' found in cache (value=#{cached.value})") - return cached + # The cache is keyed by flag, company and user, so a preflighted check and + # a plain one collide on one entry while asking different questions ("is + # this allowed after the action" versus "is it allowed now"). A + # preflighted check therefore neither reads the cache nor writes to it. + if preflight.nil? + @flag_check_cache_providers.each do |provider| + cached = coerce_cached_response(provider.get(cache_key)) + if cached + @logger.debug("Flag '#{flag_key}' found in cache (value=#{cached.value})") + return cached + end end end @@ -410,11 +686,19 @@ def check_flag_via_api(flag_key, company, user) eval_body = {} eval_body[:company] = company if company&.any? eval_body[:user] = user if user&.any? + eval_body[:preflight] = preflight if preflight - api_response = @api_client.features.check_flag(key: flag_key, **eval_body) + api_response = @api_client.features.check_flag( + request_options: api_request_options(timeout_ms), key: flag_key, **eval_body + ) data = api_response.data @logger.debug("API returned flag '#{flag_key}' value=#{data.value}, reason=#{data.reason}") + if data.value.nil? + @logger.debug("No value returned from feature flag API for flag '#{flag_key}', falling back to default") + return CheckFlagResponse.new(value: get_default.call, flag_key: flag_key, reason: "flag default") + end + response = CheckFlagResponse.new( value: data.value, flag_key: data.flag, @@ -433,22 +717,35 @@ def check_flag_via_api(flag_key, company, user) feature_usage_reset_at: data.respond_to?(:feature_usage_reset_at) ? data.feature_usage_reset_at : nil ) - # Cache the response - @flag_check_cache_providers.each do |provider| - provider.set(cache_key, response) + # Cache the response, unless the verdict was preflighted: it answers a + # question a later plain check is not asking. + if preflight.nil? + @flag_check_cache_providers.each do |provider| + provider.set(cache_key, response) + end end response rescue StandardError => e @logger.error("API flag check failed for '#{flag_key}': #{e.message}") CheckFlagResponse.new( - value: get_flag_default(flag_key), + value: get_default.call, flag_key: flag_key, reason: "error: #{e.message}" ) end end + # The generated transport reads its timeout from its own construction, not + # from a request, so timeout_in_seconds is carried but not yet honored. + # Sending it anyway means a per-check timeout starts working the moment the + # transport does, without another change here. + def api_request_options(timeout_ms) + return {} if timeout_ms.nil? + + { timeout_in_seconds: timeout_ms / 1000.0 } + end + def get_flag_default(flag_key) @flag_defaults_mutex.synchronize do @flag_defaults.fetch(flag_key, false) @@ -557,6 +854,506 @@ def enqueue_flag_check_event(flag_key, response, company, user) }) end + # --- Credit Leases --- + + # Normalize and validate before the constructor starts anything, so a bad + # knob is a clean raise rather than a raise on top of a running flush thread + # and an open socket. Offline is left alone: it starts neither, and the + # warning below already says lease gating is off. + def validated_credit_lease_config(config) + return nil if config.nil? || @offline + + normalize_credit_lease_config(config).tap { |normalized| validate_credit_lease_config(normalized) } + end + + def setup_credit_leases(config, datastream_options) + if @offline + @logger.warn( + "credit_leases is configured but the client is in offline mode; lease-gated checks are disabled " \ + "and check will return flag defaults with no credit gating." + ) + return + end + + # Already normalized and validated in the constructor. + @credit_lease_config = config + @credit_lease_mode = config[:mode] || :auto + resolve_server_reservation_ttl(config) + warn_about_mode(config) + return unless credit_lease_mode_uses_leases? + + build_lease_plumbing(config, datastream_options) + end + + # Reject a knob that cannot mean anything, at construction, where the stack + # still points at the caller. Left to run, a zero sweep interval spins a + # thread flat out, a non-positive lease size or duration acquires a lease + # nothing can reserve against, and a water mark outside (0, 1) either never + # extends or extends on every check. + def validate_credit_lease_config(config) + %i[default_lease_duration default_reservation_ttl default_lease_size sweep_interval_ms + prewarm_resolve_timeout_ms].each do |knob| + # prewarm_resolve_timeout_ms documents 0 as cache-only, so it alone may + # be zero. + validate_positive_number(config, knob, allow_zero: knob == :prewarm_resolve_timeout_ms) + end + validate_low_water_mark(config) + (config[:overrides] || {}).each_value { |override| validate_credit_lease_config(override) } + nil + end + + def validate_positive_number(config, knob, allow_zero: false) + value = config[knob] + return if value.nil? + + # finite? rejects NaN and both infinities, neither of which can size a + # lease, a sweep, or a timeout. + valid = value.is_a?(Numeric) && value.to_f.finite? && (allow_zero ? value >= 0 : value.positive?) + return if valid + + raise ArgumentError, + "credit_leases[:#{knob}] must be a finite #{allow_zero ? "non-negative" : "positive"} number, " \ + "got #{value.inspect}" + end + + def validate_low_water_mark(config) + value = config[:low_water_mark] + return if value.nil? + return if value.is_a?(Numeric) && value.to_f.finite? && value.positive? && value < 1 + + raise ArgumentError, "credit_leases[:low_water_mark] must be a number between 0 and 1, got #{value.inspect}" + end + + # Accept the hyphenated spellings the other SDKs use for the two enum-ish + # knobs, so one config shape travels across a mixed fleet. + def normalize_credit_lease_config(config) + normalized = config.transform_keys(&:to_sym) + normalized[:mode] = resolve_credit_lease_mode(normalized[:mode]) + overrides = normalized[:overrides] + return normalized unless overrides.is_a?(Hash) + + # A credit type id is a string, but { "ct_1": {} } in Ruby is the symbol + # :ct_1, and the knobs inside are read by symbol. Settle both spellings + # here: read the wrong way round, an override silently falls back to the + # client-wide defaults. + normalized[:overrides] = overrides.each_with_object({}) do |(credit_type_id, knobs), out| + out[credit_type_id.to_s] = knobs.is_a?(Hash) ? knobs.transform_keys(&:to_sym) : knobs + end + normalized + end + + # An unrecognized mode must not read as :auto in silence: the branches test + # for :client and :server and everything else falls through, so a typo would + # quietly pick a mode the caller did not ask for. Name the value, then use + # the documented default. + def resolve_credit_lease_mode(value) + return :auto if value.nil? + + mode = Credits::Leases.normalize_symbol(value) + return mode if Credits::Leases::MODES.include?(mode) + + @logger.warn( + "Unrecognized credit_leases[:mode] #{value.inspect}; expected one of " \ + "#{Credits::Leases::MODES.join(", ")}. Using :auto." + ) + :auto + end + + # The API refuses a hold expiring more than an hour after its own clock, and + # this TTL is applied to the caller's, so clamp a step below the cap to leave + # room for skew. Only server mode sends the value to the API: in client mode + # it sizes the local sweep, so clamping there would shorten holds for no + # reason and the warning would be untrue. + def resolve_server_reservation_ttl(config) + configured = config[:default_reservation_ttl] || Credits::Leases::DEFAULT_RESERVATION_TTL_MS + max_ttl = Credits::Leases::MAX_RESERVATION_TTL_MS - Credits::Leases::RESERVATION_TTL_SKEW_ALLOWANCE_MS + @server_reservation_ttl_ms = @credit_lease_mode == :client ? configured : [configured, max_ttl].min + return unless @credit_lease_mode != :client && configured > max_ttl + + @logger.warn( + "credit_leases[:default_reservation_ttl] of #{configured}ms is longer than the API will hold credits " \ + "for; server-mode holds will be clamped to #{max_ttl}ms (the " \ + "#{Credits::Leases::MAX_RESERVATION_TTL_MS}ms maximum, less " \ + "#{Credits::Leases::RESERVATION_TTL_SKEW_ALLOWANCE_MS}ms of room for clock skew)." + ) + end + + def warn_about_mode(config) + # Server mode holds credits over the API, so none of the local lease + # plumbing is built and options that only steer it would silently do + # nothing. Say so once, at startup. :auto with no DataStream lands in + # server mode too, and is the likelier way to get here. + if @credit_lease_mode == :server || (@credit_lease_mode == :auto && @datastream_client.nil?) + client_only = CLIENT_ONLY_LEASE_OPTIONS.reject { |name| config[name].nil? } + if client_only.any? + @logger.warn( + "credit_leases resolves to server mode, so #{client_only.join(", ")} will be ignored: " \ + "those options only apply to client mode (local leases over DataStream)." + ) + end + end + + # :auto with no DataStream is the server-mode default, not a + # misconfiguration: check-and-reserve gates over the API instead. :client + # without DataStream is the degraded path, where every check falls back to + # a plain flag check with usage ignored, so it warns. + if @credit_lease_mode == :auto && @datastream_client.nil? + @logger.info( + "credit_leases is configured and DataStream is not enabled; credit reservations will run in server " \ + "mode (one check-and-reserve API call per check). Set use_data_stream: true (or replicator mode) " \ + "for client-side leases." + ) + end + return unless @credit_lease_mode == :client && @datastream_client.nil? + + @logger.warn( + "credit_leases is configured but DataStream is not enabled; check will fall back to plain flag checks " \ + "with NO credit gating (usage is ignored). Set use_data_stream: true (or replicator mode) to enable " \ + "lease-gated checks." + ) + end + + def build_lease_plumbing(config, datastream_options) + sweep_ms = config[:sweep_interval_ms] || Credits::Leases::DEFAULT_SWEEP_INTERVAL_MS + # Lease and reservation state belongs in a shared cache so gating holds + # across horizontally scaled processes. Prefer an explicit client, but + # otherwise reuse the one the DataStream cache is already configured with, + # so an existing Redis setup backs leases automatically. Same for the key + # prefix. + redis_client = config[:redis_client] || datastream_options[:redis_client] + key_prefix = config[:redis_key_prefix] || datastream_options[:redis_key_prefix] + + if redis_client + # Shared-state backend: the lease balance and the reservation table live + # in Redis, and the Lua-driven reserve and consume paths give atomic + # cross-process gating without a separate lock service. + @lease_backend_shared = true + @lease_store = Credits::Leases::RedisLeaseStore.new( + client: redis_client, key_prefix: key_prefix, + default_lease_duration_ms: config[:default_lease_duration] || Credits::Leases::DEFAULT_LEASE_DURATION_MS + ) + @reservations = Credits::Leases::RedisReservationStore.new( + client: redis_client, lease_store: @lease_store, sweep_interval_ms: sweep_ms, + key_prefix: key_prefix, logger: @logger + ) + else + # No shared backend configured. In a horizontally scaled deployment each + # process then acquires and gates against its own leases, which defeats + # the cross-process over-spend protection that is the point of leasing, + # so warn rather than degrade silently. + @logger.warn( + "credit_leases is enabled without a shared Redis backend; lease and reservation state will be kept " \ + "per-process. Configure datastream_options[:redis_client] (or credit_leases[:redis_client]) so " \ + "leases gate correctly across multiple SDK instances." + ) + @lease_store = Credits::Leases::LeaseStore.new + @reservations = Credits::Leases::ReservationStore.new(@lease_store, sweep_ms, logger: @logger) + end + + @reservations.start_sweep + @credit_lease_manager = Credits::Leases::LeaseManager.new( + wire_client: Credits::Leases::ApiWireClient.new(credits_client: credits), + lease_store: @lease_store, + logger: @logger, + config: config + ) + @prewarm_resolve_timeout_ms = + config[:prewarm_resolve_timeout_ms] || Credits::Leases::DEFAULT_PREWARM_RESOLVE_TIMEOUT_MS + end + + # Whether the configured mode wants the local lease plumbing. Read during + # construction, after the DataStream client has been wired, so :auto can + # resolve against it. + def credit_lease_mode_uses_leases? + return false if @credit_lease_mode.nil? || @credit_lease_mode == :server + return true if @credit_lease_mode == :client + + !@datastream_client.nil? + end + + # Which reservation mode a check with usage resolves to right now. Nil means + # no credit gating at all: credit_leases is not configured, or the client is + # offline. + # + # :auto is resolved per check rather than once at startup, so a DataStream + # that failed to start after construction falls to server mode instead of + # silently dropping every check to a plain, ungated flag check. + def effective_lease_mode + return nil if @credit_lease_mode.nil? || @offline + return :server if @credit_lease_mode == :server + return :client if @credit_lease_mode == :client + + plumbing_ready = !@credit_lease_manager.nil? && !@lease_store.nil? && !@reservations.nil? + @datastream_client && plumbing_ready ? :client : :server + end + + # The plain (non-lease) check the credit paths fall back to, with the + # caller's preflight threaded through so a client-side evaluation still + # gates on the post-call balance. + def plain_check_result(flag_key, company, user, options) + # The caller's default_value governs this path too. Without it a check + # that falls back and then fails would answer with the registered flag + # default, denying where the caller asked to allow. + response = check_flag_with_entitlement( + flag_key, company: company, user: user, + preflight: Credits::Leases.build_preflight_options(options), + default_value: options[:default_value], timeout_ms: options[:timeout_ms] + ) + value = response.value + Credits::Leases::CheckResult.new( + allowed: value, value: value, reason: response.reason, entitlement: response.entitlement, + flag_key: response.flag_key || flag_key, flag_id: response.flag_id, error: response.error + ) + end + + def resolve_default_value(flag_key, default_value) + return get_flag_default(flag_key) if default_value.nil? + return default_value.call if default_value.respond_to?(:call) + + default_value + end + + def enqueue_lease_flag_check_event(body) + payload = { + flag_key: body[:flag_key], + value: body[:value], + reason: body[:reason] + } + payload[:error] = body[:error] if body[:error] + payload[:flag_id] = body[:flag_id] if body[:flag_id] + payload[:rule_id] = body[:rule_id] if body[:rule_id] + payload[:company_id] = body[:company_id] if body[:company_id] + payload[:user_id] = body[:user_id] if body[:user_id] + payload[:company] = body[:req_company] if body[:req_company]&.any? + payload[:user] = body[:req_user] if body[:req_user]&.any? + + @event_buffer.push({ event_type: "flag_check", body: payload, sent_at: Time.now.utc.iso8601 }) + rescue StandardError => e + @logger.error("Error enqueueing flag_check event: #{e.message}") + end + + def settle_reservation(reservation, actual_quantity, traits) + idempotency_key = "#{RESERVATION_TRACK_IDEMPOTENCY_PREFIX}#{reservation.id}" + # Server mode: the hold lives on the server and settles by id, so there is + # nothing local to consume or refund. Just emit the track. + if reservation.server_mode? || @reservations.nil? + @logger.warn("track_with_reservation called but credit_leases is not configured; emitting unsettled track") if @reservations.nil? && !reservation.server_mode? + # Without a local store there is nothing to settle against, but the + # billing event must still carry the lease id (the handle was issued by + # a lease-configured client, and dropping it would double-debit the + # grant) and the deterministic idempotency key. + track(Credits::Leases.build_reservation_track_event(reservation, actual_quantity, traits: traits), + options: { idempotency_key: idempotency_key }) + return + end + + outcome = begin + Credits::Leases.consume_reservation_and_build_event(@reservations, reservation, actual_quantity, + traits: traits) + rescue StandardError => e + # The local settle failed, most likely an unreachable Redis. The usage + # still has to be billed: build the track from the caller-held handle + # and emit it anyway. The unsettled local hold is reclaimed by the + # sweeper at its TTL or at lease expiry, and the idempotency key keeps a + # retried settle from double billing. + @logger.warn( + "track_with_reservation: failed to settle reservation #{reservation.id} locally (#{e.message}), " \ + "emitting track anyway" + ) + Credits::Leases::SettleOutcome.new( + track: Credits::Leases.build_reservation_track_event(reservation, actual_quantity, traits: traits), + settled_locally: false + ) + end + + unless outcome.settled_locally + @logger.debug( + "track_with_reservation: reservation #{reservation.id} was not settled locally (expired or swept, " \ + "already settled, or store unreachable), emitting track keyed for idempotent server-side dedupe" + ) + end + # The cached company metric moves only when this call moved local state + # with it: the server drops a duplicate event on the 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. + emit_track(outcome.track, { idempotency_key: idempotency_key }, update_metrics: outcome.settled_locally) + end + + # Enqueue a track event, optimistically bumping the cached company metric + # with it unless the caller says not to. The bump is a local prediction of + # what the stream will push back, so it belongs only to an event that + # records usage the server has not already counted. + def emit_track(body, options, update_metrics:) + return if @offline + + @event_buffer.push(build_event("track", body, options, TRACK_OPTION_KEYS)) + + if update_metrics && @datastream_client&.connected? && body[:company] + event_name = body[:event] || body["event"] + quantity = body[:quantity] || body["quantity"] || 1 + @datastream_client.update_company_metrics(body[:company], event_name, quantity) + end + rescue StandardError => e + @logger.error("Error sending track event: #{e.message}") + end + + def prewarm_after_identify(body, credit_type_ids) + # A thread that would only log "no-op in server mode" is still a thread, + # and close still has to wait it out. Decide before spawning one. + return nil if @credit_lease_manager.nil? || @lease_store.nil? + + company = identify_company_keys(body) + + thread = Thread.new do + # Force a flush so the server processes the identify as soon as + # possible. Without it the company may sit in the local buffer for up to + # the flush interval before the server even sees it, and prewarm's + # bounded poll would just be waiting on us. It runs here rather than on + # the caller's thread because a flush is an HTTP post with retries, and + # identify with a prewarm must stay the buffer push that identify + # without one is. The ordering the poll needs still holds: the flush and + # the poll are the same thread, in that order. close waits on this + # thread, so a shutdown still covers the flush. + begin + @event_buffer.flush + rescue StandardError => e + @logger.debug("identify flush before prewarm failed: #{e.message}") + end + prewarm(credit_type_ids, company: company) + rescue StandardError => e + @logger.warn("identify prewarm failed: #{e.message}") + end + thread.abort_on_exception = false + @pending_prewarms_mutex.synchronize do + @pending_prewarms.select!(&:alive?) + @pending_prewarms << thread + end + nil + end + + # The Schematic company id for a set of entity keys, 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. + # + # On a cache miss this actively fetches the company over the datastream, + # warming the cache as a side effect. identify does not push a company into + # that cache: companies are only streamed in response to a request. So this + # fetches rather than passively polling the cache, which would watch an + # empty cache until it times out. Fetching also primes the cache so the + # first real check hits the lease path instead of falling back. + def resolve_company_id_with_wait(company) + return schematic_company_id(company) if @datastream_client.nil? + + cached = @datastream_client.get_cached_company(company) + cached_id = cached && (cached[:id] || cached["id"]) + return cached_id if cached_id + # A zero or negative timeout means cache-only: answer from what the + # DataStream already holds and never fetch or poll. A prewarm still + # acquires when an earlier check warmed the company. + return schematic_company_id(company) if @prewarm_resolve_timeout_ms <= 0 + + deadline = monotonic_ms + @prewarm_resolve_timeout_ms + loop do + # A company resolved for a client that is shutting down warms nothing, + # and close would be waiting out the rest of this poll. + return nil if @closing + + if @datastream_client.connected? + resolved_id = fetch_company_id_within(company, deadline) + return resolved_id if resolved_id + else + # A request sent over a closed socket is dropped without an error, so + # the fetch would just sit out its own timeout. Poll the cache until + # the socket returns or the budget runs out. + cached = @datastream_client.get_cached_company(company) + cached_id = cached && (cached[:id] || cached["id"]) + return cached_id if cached_id + end + # 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_company_id(company) if monotonic_ms >= deadline + + sleep([Credits::Leases::DEFAULT_PREWARM_POLL_INTERVAL_MS, deadline - monotonic_ms].min / 1000.0) + end + end + + # 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. + def schematic_company_id(keys) + return nil unless keys.is_a?(Hash) + + keys.values.find { |v| v.is_a?(String) && v.start_with?(COMPANY_ID_PREFIX) } + end + + # A caller can hand identify any shape, so read the company keys without + # assuming one. + def identify_company_keys(body) + return nil unless body.is_a?(Hash) + + company = body[:company] || body["company"] + return nil unless company.is_a?(Hash) + + company[:keys] || company["keys"] + end + + # get_company waits on the DataStream's own resource timeout, which is many + # times the prewarm budget, so it runs on a thread joined to what is left. + # An abandoned fetch is left to finish: it still warms the cache for the + # next poll or the first real check. + def fetch_company_id_within(company, deadline) + remaining = deadline - monotonic_ms + return nil if remaining <= 0 + + fetch = Thread.new { @datastream_client.get_company(company) } + fetch.abort_on_exception = false + resolved = fetch.join(remaining / 1000.0)&.value + resolved && (resolved[:id] || resolved["id"]) + rescue StandardError => e + @logger.debug("prewarm: datastream company fetch failed (#{e.message})") + nil + end + + def shut_down_credit_leases + @reservations&.stop + return if @credit_lease_manager.nil? + + # Refuse new lease work first, so the waits below are waiting on work that + # is already unwinding rather than work still starting. Both steps run for + # a shared backend too: the work must not outlive the client, even where + # there is nothing to release. + @credit_lease_manager.stop + # 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 = monotonic_ms + Credits::Leases::SHUTDOWN_DRAIN_TIMEOUT_MS + prewarms = @pending_prewarms_mutex.synchronize { @pending_prewarms.dup } + prewarms.each do |thread| + remaining = deadline - monotonic_ms + break if remaining <= 0 + + thread.join(remaining / 1000.0) + end + if prewarms.any?(&:alive?) + @logger.warn( + "Timed out after #{Credits::Leases::SHUTDOWN_DRAIN_TIMEOUT_MS}ms waiting for in-flight prewarms on close" + ) + end + @credit_lease_manager.drain([deadline - monotonic_ms, 0].max) + return if @lease_backend_shared + + # The releases share the shutdown budget too, so a slow API cannot stretch + # close past what the caller was promised. + @credit_lease_manager.release_all_local_leases([deadline - monotonic_ms, 0].max) + end + + def monotonic_ms + Process.clock_gettime(Process::CLOCK_MONOTONIC) * 1000 + end + def setup_datastream(options) @rules_engine = RulesEngine.new(logger: @logger) diff --git a/lib/schematichq.rb b/lib/schematichq.rb index 13538df..2432533 100644 --- a/lib/schematichq.rb +++ b/lib/schematichq.rb @@ -14,4 +14,14 @@ require_relative "schematic/datastream/resource_cache" require_relative "schematic/datastream/websocket_client" require_relative "schematic/datastream/client" +require_relative "schematic/credits/leases/types" +require_relative "schematic/credits/leases/lease_store" +require_relative "schematic/credits/leases/reservation_store" +require_relative "schematic/credits/leases/redis_lease_store" +require_relative "schematic/credits/leases/redis_reservation_store" +require_relative "schematic/credits/leases/wire_client" +require_relative "schematic/credits/leases/lease_manager" +require_relative "schematic/credits/leases/check" +require_relative "schematic/credits/leases/server_check" +require_relative "schematic/credits/leases/track" require_relative "schematic/schematic_client" diff --git a/test/conformance_test.rb b/test/conformance_test.rb new file mode 100644 index 0000000..9769279 --- /dev/null +++ b/test/conformance_test.rb @@ -0,0 +1,410 @@ +# frozen_string_literal: true + +require "minitest/autorun" +require "json" +require_relative "../lib/schematichq" +require_relative "lease_support" + +# Runs the language-agnostic conformance vectors against this SDK. +# +# The vectors and the semantics they pin live in conformance/ at the repo root, +# copied verbatim from schematic-node (the reference implementation). This +# runner is the only language-specific piece; every port reimplements it and +# must pass the same vectors, against every store backend it ships. +class ConformanceTest < Minitest::Test + Leases = Schematic::Credits::Leases + + VECTOR_DIR = File.expand_path("../conformance/vectors", __dir__) + BACKENDS = %w[in_memory redis].freeze + + # One vector's stores, manager, clock, and reservation handles. + class Harness + attr_reader :clock, :leases, :reservations, :crash, :handles, :redis + + def initialize(backend, config) + @config = config + @clock = LeaseSupport::VirtualClock.new + @handles = {} + if backend == "redis" + @redis = LeaseSupport::FakeRedis.new(@clock) + @leases = Leases::RedisLeaseStore.new(client: @redis, clock: @clock.to_proc) + @crash = LeaseSupport::CrashingRefund.new(@leases) + @reservations = Leases::RedisReservationStore.new(client: @redis, lease_store: @crash, + clock: @clock.to_proc) + else + @leases = Leases::LeaseStore.new(clock: @clock.to_proc) + @crash = LeaseSupport::CrashingRefund.new(@leases) + @reservations = Leases::ReservationStore.new(@crash, Leases::DEFAULT_SWEEP_INTERVAL_MS, + clock: @clock.to_proc) + end + end + + # Built on first use so a pure store vector never spins one up. + def wire + @wire ||= LeaseSupport::ScriptedWireClient.new(@clock) + end + + def unconsumed_wire_scripts + @wire ? @wire.pending_scripts : [] + end + + def manager + @manager ||= Leases::LeaseManager.new( + wire_client: wire, + lease_store: @leases, + logger: Schematic::ConsoleLogger.new(level: :error), + config: @config, + clock: @clock.to_proc + ) + end + + def stop + @reservations&.stop + @manager&.stop + end + + def reservation_id(step) + handle = step["handle"] + return @handles.fetch(handle).id if handle + + step.fetch("id") + end + end + + # Build one test method per (backend, category, vector) so a failure names the + # vector it came from. + Dir[File.join(VECTOR_DIR, "*.json")].each do |path| + document = JSON.parse(File.read(path)) + category = document["category"] + document["vectors"].each do |vector| + BACKENDS.each do |backend| + next if vector["backends"] && !vector["backends"].include?(backend) + + define_method("test_#{backend}_#{category}_#{vector["name"]}") do + run_vector(backend, vector) + end + end + end + end + + private + + def run_vector(backend, vector) + given = vector["given"] || {} + harness = Harness.new(backend, config_from(given["config"])) + (given["leases"] || []).each do |lease| + wrote = harness.leases.replace(lease_from(harness, lease)) + + assert wrote, "given.leases must install: #{lease["lease_id"]}" + end + (vector["operations"] || []).each { |step| run_operation(harness, step) } + # A scripted response nobody consumed means the run took a different path + # through the wire than the vector describes, which no per-step assertion + # would notice. + assert_empty harness.unconsumed_wire_scripts, "scripted wire responses left unconsumed" + ensure + harness&.stop + end + + def config_from(config) + config ||= {} + { + default_lease_duration: config["lease_duration_ms"], + default_reservation_ttl: config["reservation_ttl_ms"], + default_lease_size: config["lease_size"], + low_water_mark: config["low_water_mark"] + }.compact + end + + def lease_from(harness, spec) + Leases::LeaseEntry.new( + lease_id: spec["lease_id"], + company_id: spec["company_id"], + credit_type_id: spec["credit_type_id"], + granted_amount: spec["granted_amount"], + expires_at: harness.clock.at_ms(spec["expires_at_ms"]) + ) + end + + def run_operation(harness, step) + handler = "op_#{step["op"]}" + raise "unknown conformance step: #{step["op"]}" unless respond_to?(handler, true) + + send(handler, harness, step, step["expect"] || {}) + end + + # --- store-level ops ------------------------------------------------------- + + def op_advance_clock(harness, step, _expect) + harness.clock.advance_ms(step["ms"]) + end + + def op_replace_lease(harness, step, expect) + wrote = harness.leases.replace(lease_from(harness, step)) + assert_equal expect["written"], wrote if expect.key?("written") + end + + def op_drop_lease(harness, step, _expect) + harness.leases.drop(step["company_id"], step["credit_type_id"]) + end + + def op_try_reserve(harness, step, expect) + result = harness.leases.try_reserve(step["company_id"], step["credit_type_id"], step["credits"]) + + assert_nullable_number(expect, "balance", result&.balance) + return unless expect.key?("lease_id") + + if expect["lease_id"].nil? + assert_nil result + else + # The charged lease is what a caller pins its reservation to, so a vector + # naming one is checking the pin, not just the arithmetic. + assert_equal expect["lease_id"], result.lease_id + end + end + + def op_refund_lease(harness, step, _expect) + harness.leases.refund(step["company_id"], step["credit_type_id"], step["credits"], step["pin_lease_id"]) + end + + def op_extend_lease(harness, step, _expect) + expires_at = step["expires_at_ms"] ? harness.clock.at_ms(step["expires_at_ms"]) : nil + harness.leases.extend(step["company_id"], step["credit_type_id"], step["granted_total"], expires_at, + step["pin_lease_id"]) + end + + def op_get_lease(harness, step, expect) + entry = harness.leases.get(step["company_id"], step["credit_type_id"]) + assert_equal expect["exists"], !entry.nil? if expect.key?("exists") + if expect.key?("lease_id") + expect["lease_id"].nil? ? assert_nil(entry) : assert_equal(expect["lease_id"], entry.lease_id) + end + assert_in_delta expect["granted_amount"], entry.granted_amount if expect.key?("granted_amount") + return unless expect.key?("local_remaining_credits") + + assert_in_delta expect["local_remaining_credits"], entry.local_remaining_credits + end + + def op_add_reservation(harness, step, _expect) + harness.reservations.add(Leases::Reservation.new( + id: step["id"], + lease_id: step["lease_id"], + company_id: step["company_id"], + credit_type_id: step["credit_type_id"], + event_subtype: step["event_subtype"], + quantity_reserved: step["quantity_reserved"], + credits_reserved: step["credits_reserved"], + consumption_rate: step["consumption_rate"], + expires_at: harness.clock.at_ms(step["expires_at_ms"]), + eval_ctx: { company: { id: step["company_id"] } } + )) + end + + def op_consume_reservation(harness, step, expect) + id = harness.reservation_id(step) + if step["crash_before_refund"] + harness.crash.arm + assert_raises(LeaseSupport::CrashingRefund::SimulatedCrash) do + harness.reservations.consume(id, step["credits"]) + end + assert expect["throws"] + return + end + + assert_nullable_number(expect, "consumed", harness.reservations.consume(id, step["credits"])) + end + + def op_get_reservation(harness, step, expect) + reservation = harness.reservations.get(harness.reservation_id(step)) + assert_equal expect["exists"], !reservation.nil? if expect.key?("exists") + end + + def op_reserved_credits(harness, step, expect) + total = harness.reservations.reserved_credits(step["company_id"], step["credit_type_id"]) + + assert_in_delta expect["total"], total + end + + def op_reservation_count(harness, _step, expect) + assert_equal expect["count"], harness.reservations.size + end + + def op_sweep_expired(harness, _step, expect) + swept = harness.reservations.sweep_expired + assert_equal expect["swept"], swept if expect.key?("swept") + end + + # --- manager-level ops ----------------------------------------------------- + + def op_acquire_if_needed(harness, step, expect) + harness.wire.queue_acquire(step["server"]) if step["server"] + if (install = step["install_during_wire"]) + harness.wire.during_acquire = -> { harness.leases.replace(lease_from(harness, install)) } + end + entry = harness.manager.acquire_if_needed(step["company_id"], step["credit_type_id"]) + harness.manager.drain + + if expect.key?("lease_id") + expect["lease_id"].nil? ? assert_nil(entry) : assert_equal(expect["lease_id"], entry&.lease_id) + end + assert_equal expect["wire_acquires"], harness.wire.acquire_calls.size if expect.key?("wire_acquires") + assert_in_delta expect["last_acquire_requested_amount"], harness.wire.acquire_calls.last.requested_amount if expect.key?("last_acquire_requested_amount") + return unless expect.key?("released_lease_ids") + + assert_equal expect["released_lease_ids"], harness.wire.release_calls + end + + def op_maybe_extend(harness, step, expect) + harness.wire.queue_extend(step["server"]) if step["server"] + harness.manager.maybe_extend_in_background(step["company_id"], step["credit_type_id"], step["required_credits"]) + harness.manager.drain + + assert_equal expect["wire_extends"], harness.wire.extend_calls.size if expect.key?("wire_extends") + assert_in_delta expect["last_extend_additional_amount"], harness.wire.extend_calls.last.additional_amount if expect.key?("last_extend_additional_amount") + return unless expect.key?("last_extend_lease_id") + + assert_equal expect["last_extend_lease_id"], harness.wire.extend_calls.last.lease_id + end + + def op_release_all_local_leases(harness, _step, expect) + harness.manager.release_all_local_leases + assert_equal expect["released_lease_ids"], harness.wire.release_calls if expect.key?("released_lease_ids") + return unless expect.key?("remaining_slots") + + assert_equal expect["remaining_slots"], harness.leases.list.size + end + + # --- flow-level ops -------------------------------------------------------- + + def op_check(harness, step, expect) + flag_key = step["flag_key"] || "flag" + spec = step["company"] || { "id" => "co_1" } + if (script = step["server"]) + harness.wire.queue_acquire(script["acquire"]) if script["acquire"] + harness.wire.queue_extend(script["extend"]) if script["extend"] + end + datastream = LeaseSupport::ScriptedDataStream.new( + flag_key: flag_key, + company: { id: spec["id"], credit_balances: symbolize_balances(spec["credit_balances"]) }, + results: step["engine"] || [] + ) + + fell_back = false + fallback = -> do + fell_back = true + Leases::CheckResult.new(allowed: true, value: true, reason: "fallback", flag_key: flag_key) + end + result = Leases.check_with_lease( + Leases::CheckDeps.new( + lease_store: harness.leases, reservations: harness.reservations, manager: harness.manager, + datastream: datastream, logger: Schematic::ConsoleLogger.new(level: :error), clock: harness.clock.to_proc + ), + flag_key, + { company: { id: spec["id"] } }, + { usage: step["usage"], event_subtype: step["event_subtype"], + on_acquire_failure: step["on_acquire_failure"] }, + &fallback + ) + harness.manager.drain + + assert_check_result(expect, result, fell_back) + assert_engine_calls(step, expect, datastream.calls, spec) + # SPEC.md makes expect.engine_calls optional, so an absent one asserts + # nothing. The scripted engine results are not optional: one left over means + # the check reached the engine fewer times than the vector describes, which + # no other assertion here would notice. An extra call already raises. + assert_empty datastream.pending_results, "scripted engine results left unconsumed" + assert_equal expect["wire_extends"], harness.wire.extend_calls.size if expect.key?("wire_extends") + assert_in_delta expect["last_extend_additional_amount"], harness.wire.extend_calls.last.additional_amount if expect.key?("last_extend_additional_amount") + return unless step["save_reservation_as"] && result.reservation + + harness.handles[step["save_reservation_as"]] = result.reservation + end + + def op_track(harness, step, expect) + reservation = harness.handles.fetch(step["handle"]) + outcome = Leases.consume_reservation_and_build_event(harness.reservations, reservation, step["actual_quantity"]) + assert_equal expect["settled_locally"], outcome.settled_locally if expect.key?("settled_locally") + return unless expect.key?("track") + + want = expect["track"] + assert_equal want["event"], outcome.track[:event] if want.key?("event") + assert_in_delta want["quantity"], outcome.track[:quantity] if want.key?("quantity") + assert_equal want["lease_id"], outcome.track[:lease_id] if want.key?("lease_id") + end + + def assert_check_result(expect, result, fell_back) + assert_equal expect["allowed"], result.allowed if expect.key?("allowed") + assert_equal expect["reason"], result.reason if expect.key?("reason") + assert_equal expect["err"], result.error if expect.key?("err") + assert_equal expect["has_reservation"], !result.reservation.nil? if expect.key?("has_reservation") + assert_equal expect["fallback_called"], fell_back if expect.key?("fallback_called") + return unless expect.key?("reservation") + + want = expect["reservation"] + reservation = result.reservation + + refute_nil reservation + assert_equal want["lease_id"], reservation.lease_id if want.key?("lease_id") + assert_equal want["credit_type_id"], reservation.credit_type_id if want.key?("credit_type_id") + assert_equal want["event_subtype"], reservation.event_subtype if want.key?("event_subtype") + assert_in_delta want["quantity_reserved"], reservation.quantity_reserved if want.key?("quantity_reserved") + assert_in_delta want["credits_reserved"], reservation.credits_reserved if want.key?("credits_reserved") + assert_in_delta want["consumption_rate"], reservation.consumption_rate if want.key?("consumption_rate") + end + + def assert_engine_calls(step, expect, calls, spec) + return unless expect.key?("engine_calls") + + want = expect["engine_calls"] + + assert_equal want.size, calls.size + credit_id = credit_id_from(step, spec) + want.each_with_index do |expected, index| + got = calls[index] + if expected.key?("credit_balance") + refute_empty credit_id.to_s, "engine_calls needs a credit id to assert a balance against" + assert_in_delta expected_balance(expected["credit_balance"]), got.credit_balances[credit_id] + end + assert_in_delta expected["credit_cost"], got.options[:credit_cost][credit_id] if expected.key?("credit_cost") + if expected.key?("event_usage") + assert_equal expected["event_usage"]["event_subtype"], got.options[:event_usage][:event_subtype] + assert_in_delta expected["event_usage"]["quantity"], got.options[:event_usage][:quantity] + end + assert_in_delta expected["usage"], got.options[:usage] if expected.key?("usage") + end + end + + # A balance expectation is a number, or the name of the effectively unlimited + # figure the fail-open substitution uses. + def expected_balance(raw) + raw == "max_safe_integer" ? Leases::FAIL_OPEN_BALANCE : raw + end + + # The credit a vector's balance and cost expectations are about: the one the + # scripted entitlement meters, or the company's only balance when no + # entitlement names one. + def credit_id_from(step, spec) + named = (step["engine"] || []).filter_map { |result| result.dig("entitlement", "credit_id") }.first + named || (spec["credit_balances"] || {}).keys.first + end + + def symbolize_balances(balances) + (balances || {}).each_with_object({}) { |(id, value), out| out[id] = value } + end + + # --- assertions ------------------------------------------------------------ + + # Compare a result that may be the refused outcome, which the vectors write as + # JSON null and never as a figure. + def assert_nullable_number(expect, key, value) + return unless expect.key?(key) + + if expect[key].nil? + assert_nil value, "expected #{key} to be refused, got #{value}" + else + refute_nil value, "expected #{key} to be #{expect[key]}, got the refused result" + assert_in_delta expect[key], value + end + end +end diff --git a/test/credits_test.rb b/test/credits_test.rb new file mode 100644 index 0000000..f214b24 --- /dev/null +++ b/test/credits_test.rb @@ -0,0 +1,2220 @@ +# frozen_string_literal: true + +require "minitest/autorun" +require "webmock/minitest" + +# Requiring WebMock blocks outbound HTTP for the whole suite, and the default +# rake run also covers a raw-client test that talks to a throwaway server on +# loopback. Leave that alone; every stub here targets a remote host. +WebMock.disable_net_connect!(allow_localhost: true) +require "digest" +require_relative "../lib/schematichq" +require_relative "lease_support" + +# Unit coverage for credit leases, reservations, and preflight checks. The +# conformance vectors in test/conformance_test.rb pin the cross-SDK semantics; +# these cover what a vector cannot: concurrency, wire shapes, client wiring, and +# the Redis key layout. +module CreditTestHelpers + Leases = Schematic::Credits::Leases + + def silent_logger + Schematic::ConsoleLogger.new(level: :error) + end + + # Captures warnings so a test can assert the SDK said something about a + # misconfiguration rather than correcting it in silence. + class RecordingLogger + attr_reader :warnings + + def initialize + @warnings = [] + end + + def warn(message) + @warnings << message + end + + def error(_message); end + def debug(_message); end + def info(_message); end + end + + def clock + @clock ||= LeaseSupport::VirtualClock.new + end + + def lease_entry(lease_id: "lse_1", company_id: "co_1", credit_type_id: "ct_1", granted_amount: 1000, + expires_at_ms: 300_000) + Leases::LeaseEntry.new( + lease_id: lease_id, company_id: company_id, credit_type_id: credit_type_id, + granted_amount: granted_amount, expires_at: clock.at_ms(expires_at_ms) + ) + end + + UUID_SHAPE = /\A\h{8}-\h{4}-\h{4}-\h{4}-\h{12}\z/ + + # Collects the idempotency key off every attempt WebMock sees, so a test can + # tell "one key, several attempts" from "a key per attempt". + def record_idempotency_keys(stub, sink) + stub.with do |req| + sink << JSON.parse(req.body)["idempotency_key"] + true + end + end + + def reservation(id: "res_1", lease_id: "lse_1", credits_reserved: 100, expires_at_ms: 60_000, + consumption_rate: 10, mode: nil) + Leases::Reservation.new( + id: id, lease_id: lease_id, mode: mode, company_id: "co_1", credit_type_id: "ct_1", + event_subtype: "inference_tokens", quantity_reserved: credits_reserved / consumption_rate, + credits_reserved: credits_reserved, consumption_rate: consumption_rate, + expires_at: clock.at_ms(expires_at_ms), eval_ctx: { company: { "id" => "co_1" } } + ) + end +end + +# A Lua change that is not mirrored into a real Redis would split a +# mixed-language fleet's lease in two, so the shipped script text is pinned +# against the reference implementation's byte for byte. +class LeaseScriptIdentityTest < Minitest::Test + include CreditTestHelpers + + # SHA1 of each script as schematic-node ships it. Recompute from that repo, + # never from this one, if the scripts ever legitimately change. + REFERENCE_SHAS = { + "REPLACE_SCRIPT" => "96e6830dd5664bee5700e30d64f04b1726b37c3f", + "TRY_RESERVE_SCRIPT" => "53d744e6e38eca7fc26462fd7e8b5e70bbdfb642", + "REFUND_SCRIPT" => "43873b5f916a0ec012d9377af400e1ecc0b3ec0d", + "EXTEND_SCRIPT" => "842913369460707b70475b6853d996687878b7eb" + }.freeze + + def test_lease_scripts_match_the_reference_implementation + REFERENCE_SHAS.each do |name, sha| + assert_equal sha, Leases::RedisLeaseStore.const_get(name).sha, "#{name} diverged from schematic-node" + end + end + + def test_claim_script_matches_the_reference_implementation + assert_equal "b89cdd478995f79d7fd6e5534d434ca08d214c6c", Leases::RedisReservationStore::CLAIM_SCRIPT.sha + end + + # A script is addressed by its SHA so the body is not resent on every call; a + # Redis that has never seen it answers NOSCRIPT and the body goes out once. + def test_script_falls_back_to_eval_when_redis_has_not_loaded_it + redis = LeaseSupport::FakeRedis.new(clock) + store = Leases::RedisLeaseStore.new(client: redis, clock: clock.to_proc) + + assert store.replace(lease_entry) + # The second call finds the script cached under its SHA. + refute store.replace(lease_entry(lease_id: "lse_2")) + end +end + +# Every integer field on the wire goes through wire_quantity, so its rounding +# is what decides what a caller is billed. +class WireQuantityTest < Minitest::Test + include CreditTestHelpers + + def test_a_partial_unit_rounds_up + assert_equal 2, Leases.wire_quantity(1.2) + assert_equal 1, Leases.wire_quantity(0.0001) + end + + # (0.1 + 0.2) * 10 is 3.0000000000000004, and a bare ceil would bill it as 4. + def test_float_noise_does_not_cost_a_whole_unit + assert_equal 3, Leases.wire_quantity((0.1 + 0.2) * 10) + assert_equal 1, Leases.wire_quantity(1.0000000000000002) + assert_equal 10, Leases.wire_quantity(10.0) + end + + def test_integers_and_non_numbers_pass_through + assert_equal 7, Leases.wire_quantity(7) + assert_equal 0, Leases.wire_quantity(0) + # Not finite, so it is handed back untouched rather than raising in ceil. + assert_predicate Leases.wire_quantity(Float::INFINITY), :infinite? + assert_predicate Leases.wire_quantity(Float::NAN), :nan? + end +end + +# The engine deserializes usage and event_usage.quantity as i64, so a value +# with a decimal point fails the whole check. +class EngineQuantityTest < Minitest::Test + def options + Schematic::RulesEngine.new.send( + :engine_options, + { credit_cost: { "ct_1" => 2.5 }, usage: 0.5, event_usage: { event_subtype: "tokens", quantity: 1.2 } } + ) + end + + def test_a_fractional_preflight_quantity_rounds_up + assert_equal 1, options[:usage] + assert_equal 2, options[:event_usage][:quantity] + end + + # The cost is credits, not event units, and the engine takes it as a float. + def test_the_credit_cost_is_left_alone + assert_in_delta 2.5, options[:credit_cost]["ct_1"] + end +end + +class LeaseStoreTest < Minitest::Test + include CreditTestHelpers + + def setup + @store = Leases::LeaseStore.new(clock: clock.to_proc) + end + + def test_try_reserve_rejects_a_nan_debit + @store.replace(lease_entry) + + assert_nil @store.try_reserve("co_1", "ct_1", Float::NAN) + assert_in_delta 1000, @store.get("co_1", "ct_1").local_remaining_credits + end + + def test_try_reserve_rejects_an_infinite_debit + @store.replace(lease_entry) + + assert_nil @store.try_reserve("co_1", "ct_1", Float::INFINITY) + end + + def test_concurrent_reserves_never_oversell_the_lease + @store.replace(lease_entry(granted_amount: 100)) + threads = 20.times.map do + Thread.new { @store.try_reserve("co_1", "ct_1", 10) } + end + granted = threads.map(&:value).compact + + assert_equal 10, granted.size + assert_in_delta 0, @store.get("co_1", "ct_1").local_remaining_credits + end + + def test_get_returns_a_copy_a_caller_cannot_mutate_the_store_through + @store.replace(lease_entry) + entry = @store.get("co_1", "ct_1") + entry.local_remaining_credits = 0 + + assert_in_delta 1000, @store.get("co_1", "ct_1").local_remaining_credits + end + + def test_list_enumerates_every_slot + @store.replace(lease_entry) + @store.replace(lease_entry(lease_id: "lse_2", credit_type_id: "ct_2")) + + assert_equal %w[lse_1 lse_2], @store.list.map(&:lease_id).sort + end + + # A process checking flags for companies that never lease must not collect a + # mutex per slot it merely asked about, or pruning only half works. + def test_reading_a_slot_with_no_lease_registers_no_lock + assert_nil @store.get("co_missing", "ct_1") + assert_nil @store.try_reserve("co_missing", "ct_1", 10) + assert_nil @store.refund("co_missing", "ct_1", 10) + + assert_empty @store.instance_variable_get(:@locks) + end + + # A long-lived process leases for many companies, and a mutex per slot it has + # ever touched is a leak. + def test_dropping_a_lease_prunes_its_lock + @store.replace(lease_entry) + @store.try_reserve("co_1", "ct_1", 10) + + @store.drop("co_1", "ct_1") + + assert_empty @store.instance_variable_get(:@locks) + end + + # Pruning under the slot's own mutex strands anyone already blocked on it, so + # a waiter has to notice and retry against the mutex the table now holds. + # Without that, the waiter and a newcomer would run side by side on one slot. + def test_a_waiter_on_a_pruned_lock_retries_against_the_current_one + locks = @store.instance_variable_get(:@locks) + table = @store.instance_variable_get(:@table_mutex) + held = Queue.new + release = Queue.new + order = [] + + holder = Thread.new do + @store.send(:with_lock, "co_1:ct_1") do + held << true + release.pop + # Exactly what drop does: prune while holding the very mutex it prunes. + table.synchronize { locks.delete("co_1:ct_1") } + order << :holder + end + end + held.pop + waiter = Thread.new { @store.send(:with_lock, "co_1:ct_1") { order << :waiter } } + # Let the waiter block on the mutex that is about to be pruned. + sleep 0.05 + release << true + [holder, waiter].each(&:join) + + assert_equal %i[holder waiter], order + # The waiter re-registered a mutex for the slot, which it could only do by + # noticing the one it woke holding was stale. + refute_empty locks + end +end + +class ReservationStoreTest < Minitest::Test + include CreditTestHelpers + + def setup + @leases = Leases::LeaseStore.new(clock: clock.to_proc) + @leases.replace(lease_entry) + @store = Leases::ReservationStore.new(@leases, 1000, clock: clock.to_proc) + end + + def teardown + @store.stop + end + + # The claim is the arbiter, so of two racing settles only one refunds. + def test_consume_is_exactly_once + @leases.try_reserve("co_1", "ct_1", 100) + @store.add(reservation) + + assert_in_delta 0, @store.consume("res_1", 0) + assert_nil @store.consume("res_1", 0) + assert_in_delta 1000, @leases.get("co_1", "ct_1").local_remaining_credits + end + + def test_sweep_refunds_only_expired_holds + @leases.try_reserve("co_1", "ct_1", 200) + @store.add(reservation(id: "res_1", credits_reserved: 100, expires_at_ms: 10_000)) + @store.add(reservation(id: "res_2", credits_reserved: 100, expires_at_ms: 90_000)) + clock.advance_ms(60_000) + + assert_equal 1, @store.sweep_expired + assert_in_delta 900, @leases.get("co_1", "ct_1").local_remaining_credits + assert_in_delta 100, @store.reserved_credits("co_1", "ct_1") + end + + # The sweeper must not keep the process alive, and stopping it twice is safe. + def test_start_and_stop_sweep_are_idempotent + @store.start_sweep + @store.start_sweep + @store.stop + @store.stop + end + + # A fork (Puma or Unicorn with preload) hands the child a thread object whose + # thread did not survive, and holds would then pile up unswept in the child. + def test_a_sweeper_that_did_not_survive_a_fork_starts_again + dead = Thread.new { nil } + dead.join + @store.instance_variable_set(:@sweep_thread, dead) + + @store.start_sweep + + assert_predicate @store.instance_variable_get(:@sweep_thread), :alive? + ensure + @store.stop + end +end + +# The sweeper deletes a reservation and then refunds its lease. Killing it +# between the two strands the unspent slice until the lease expires, so stop +# gives it a bounded moment to land the refund first. +class ReservationSweepStopTest < Minitest::Test + include CreditTestHelpers + + def test_stop_lets_a_sweep_in_progress_land_its_refund + refunding = Queue.new + refunded = [] + leases = Leases::LeaseStore.new(clock: clock.to_proc) + leases.replace(lease_entry) + slow = Object.new + slow.define_singleton_method(:refund) do |company_id, credit_type_id, credits, pin_lease_id = nil| + refunding << true + # The window a kill would land in. + sleep 0.05 + refunded << credits + leases.refund(company_id, credit_type_id, credits, pin_lease_id) + end + store = Leases::ReservationStore.new(slow, 10, clock: clock.to_proc, logger: silent_logger) + store.add(reservation(expires_at_ms: 0)) + store.start_sweep + refunding.pop + + started = Process.clock_gettime(Process::CLOCK_MONOTONIC) + store.stop + elapsed_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started) * 1000 + + assert_equal [100.0], refunded + assert_in_delta 1000, leases.get("co_1", "ct_1").local_remaining_credits + # And still returns promptly: the join is bounded, not open-ended. + assert_operator elapsed_ms, :<, 400 + end + + # Nothing to wait for is the common case, and it must not cost the join + # budget. + def test_stop_returns_at_once_when_no_sweep_is_running + store = Leases::ReservationStore.new(Leases::LeaseStore.new(clock: clock.to_proc), 1000, + clock: clock.to_proc, logger: silent_logger) + + started = Process.clock_gettime(Process::CLOCK_MONOTONIC) + store.stop + elapsed_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started) * 1000 + + assert_operator elapsed_ms, :<, 50 + end +end + +# EVALSHA is retried with the full script body on NOSCRIPT alone. Any other +# error is the script's own: replaying it could debit a second time. +class ScriptFallbackTest < Minitest::Test + include CreditTestHelpers + + def test_a_noscript_error_reloads_the_script + script = Leases::Script.new("return 1") + client = Object.new + def client.evalsha(_sha, **) = raise("NOSCRIPT No matching script") + def client.eval(_source, **) = "loaded" + + assert_equal "loaded", script.call(client, keys: [], argv: []) + end + + def test_any_other_error_is_raised_rather_than_retried + script = Leases::Script.new("return 1") + client = Object.new + def client.evalsha(_sha, **) = raise("ERR Lua script attempted to access nonexistent key") + def client.eval(_source, **) = raise("a retry could debit twice") + + error = assert_raises(RuntimeError) { script.call(client, keys: [], argv: []) } + + assert_match(/nonexistent key/, error.message) + end +end + +class RedisStoreLayoutTest < Minitest::Test + include CreditTestHelpers + + def setup + @redis = LeaseSupport::FakeRedis.new(clock) + @leases = Leases::RedisLeaseStore.new(client: @redis, key_prefix: "acme:", clock: clock.to_proc) + @reservations = Leases::RedisReservationStore.new(client: @redis, lease_store: @leases, key_prefix: "acme:", + clock: clock.to_proc) + end + + # Lua's tonumber reads a decimal string, and a Rational renders as "1/2", + # which the script reads as nil and treats as a zero. + def test_a_rational_amount_reaches_the_script_as_a_decimal + @leases.replace(lease_entry(granted_amount: Rational(3, 2))) + + assert_in_delta 1.5, @leases.get("co_1", "ct_1").granted_amount + assert_in_delta 1.5, @leases.get("co_1", "ct_1").local_remaining_credits + refute_nil @leases.try_reserve("co_1", "ct_1", Rational(1, 2)) + assert_in_delta 1.0, @leases.get("co_1", "ct_1").local_remaining_credits + end + + # The key layout is what lets a mixed-language fleet share one lease per slot. + def test_lease_hash_key_layout + assert_equal "acme:credit-lease:co_1:ct_1", @leases.hash_key("co_1", "ct_1") + end + + def test_reservation_add_writes_the_hash_and_both_indexes + @leases.replace(lease_entry) + @leases.try_reserve("co_1", "ct_1", 100) + @reservations.add(reservation) + + assert_equal "lse_1", @redis.hgetall("acme:credit-reservation:res_1")["leaseId"] + assert_equal ["co_1|ct_1|res_1"], @redis.zrangebyscore("acme:credit-reservations:byExpiry", 0, 1e18) + assert_in_delta 100, @reservations.reserved_credits("co_1", "ct_1") + end + + # 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. + def test_reservation_add_writes_the_hash_and_its_ttl_in_one_transaction + @leases.replace(lease_entry) + @leases.try_reserve("co_1", "ct_1", 100) + @reservations.add(reservation) + + assert_equal [%i[hset pexpireat]], @redis.transactions + assert_equal "res_1", @reservations.get("res_1").id + end + + # The row outlives its declared expiry by the grace window so the sweeper can + # still read it, but an expired-but-not-evicted lease must refuse a reserve. + def test_expired_lease_within_the_grace_window_refuses_a_reserve + @leases.replace(lease_entry(expires_at_ms: 1000)) + clock.advance_ms(2000) + + assert_nil @leases.try_reserve("co_1", "ct_1", 1) + refute_nil @leases.get("co_1", "ct_1") + end + + def test_lease_row_is_evicted_after_the_grace_window + @leases.replace(lease_entry(expires_at_ms: 1000)) + clock.advance_ms(1000 + Leases::RedisLeaseStore::LEASE_TTL_GRACE_MS + 1) + + assert_nil @leases.get("co_1", "ct_1") + end + + # Without the reservation hash the claim cannot arbitrate exactly-once across + # racing sweepers, so the orphaned index field is reconciled, not refunded. + def test_sweep_reconciles_an_evicted_hold_without_refunding_it + @leases.replace(lease_entry) + @leases.try_reserve("co_1", "ct_1", 100) + @reservations.add(reservation(expires_at_ms: 1000)) + clock.advance_ms(1000 + Leases::RedisReservationStore::RES_TTL_GRACE_MS + 1) + + assert_equal 0, @reservations.sweep_expired + assert_in_delta 0, @reservations.reserved_credits("co_1", "ct_1") + assert_in_delta 900, @leases.get("co_1", "ct_1").local_remaining_credits + end + + # The two lease stores read an empty pin differently: the Lua takes it as no + # pin and credits whichever lease holds the slot, while the per-process store + # takes it as a pin nothing matches. Both reservation stores decide it + # instead, and both decline, so a hold that cannot name its lease never + # lands on a successor. + def test_neither_backend_refunds_a_hold_that_cannot_name_its_lease + memory_leases = Leases::LeaseStore.new(clock: clock.to_proc) + memory = Leases::ReservationStore.new(memory_leases, clock: clock.to_proc) + [[@reservations, @leases], [memory, memory_leases]].each do |store, leases| + [nil, ""].each do |orphan_id| + leases.drop("co_1", "ct_1") + leases.replace(lease_entry(lease_id: "lse_successor")) + leases.try_reserve("co_1", "ct_1", 100) + # Redis stores every field as a string, so a missing id reads back empty. + next if orphan_id.nil? && store.equal?(@reservations) + + store.add(reservation(lease_id: orphan_id)) + store.consume("res_1", 0) + + # 900, not 1000: the successor's balance is untouched. + assert_in_delta 900, leases.get("co_1", "ct_1").local_remaining_credits + end + end + end + + def test_the_sweep_skips_refunding_a_hold_that_cannot_name_its_lease + leases = Leases::LeaseStore.new(clock: clock.to_proc) + store = Leases::ReservationStore.new(leases, clock: clock.to_proc) + leases.replace(lease_entry(lease_id: "lse_successor")) + leases.try_reserve("co_1", "ct_1", 100) + store.add(reservation(lease_id: nil, expires_at_ms: 1000)) + clock.advance_ms(2000) + + assert_equal 1, store.sweep_expired + assert_in_delta 900, leases.get("co_1", "ct_1").local_remaining_credits + end + + # A shared store is never enumerated: sibling processes may still be drawing + # on those leases, so close must not be able to release them. + def test_redis_store_does_not_expose_list + refute_respond_to @leases, :list + end +end + +class LeaseManagerTest < Minitest::Test + include CreditTestHelpers + + def setup + @wire = LeaseSupport::ScriptedWireClient.new(clock) + @leases = Leases::LeaseStore.new(clock: clock.to_proc) + @manager = Leases::LeaseManager.new( + wire_client: @wire, lease_store: @leases, logger: silent_logger, + config: { default_lease_size: 1000, default_lease_duration: 300_000, low_water_mark: 0.25 }, + clock: clock.to_proc + ) + end + + def teardown + @manager.stop + end + + def acquire_script(lease_id: "lse_1", granted_amount: 1000, expires_at_ms: 300_000) + { "lease" => { "lease_id" => lease_id, "granted_amount" => granted_amount, + "expires_at_ms" => expires_at_ms } } + end + + # Callers arriving while a wire call is in flight share it rather than opening + # a second lease. + def test_concurrent_acquires_share_one_wire_call + started = Queue.new + release = Queue.new + @wire.queue_acquire(acquire_script) + @wire.during_acquire = -> do + started << true + release.pop + end + + first = Thread.new { @manager.acquire_if_needed("co_1", "ct_1") } + started.pop + joiner = Thread.new { @manager.acquire_if_needed("co_1", "ct_1") } + # Give the joiner time to register against the flight before it lands. + sleep 0.05 + release << true + + assert_equal "lse_1", first.value.lease_id + assert_equal "lse_1", joiner.value.lease_id + assert_equal 1, @wire.acquire_calls.size + end + + # A store outage answering "due" would spawn a thread per check, and each + # would only fail the same read again. + def test_no_thread_is_spawned_when_the_store_cannot_be_read + broken = Object.new + def broken.get(_company_id, _credit_type_id) = raise("redis down") + manager = Leases::LeaseManager.new(wire_client: @wire, lease_store: broken, logger: silent_logger, + clock: clock.to_proc) + + assert_nil manager.maybe_extend_in_background("co_1", "ct_1") + assert_empty @wire.extend_calls + ensure + manager&.stop + end + + # Most checks sit nowhere near the water mark, so learning that must not cost + # a thread each. + def test_no_thread_is_spawned_when_no_extend_is_due + @leases.replace(lease_entry(granted_amount: 1000)) + + assert_nil @manager.maybe_extend_in_background("co_1", "ct_1") + assert_empty @wire.extend_calls + end + + # Watermark refreshes arriving during one slow extend have nothing to wait + # for: the credits they want are the ones that call is fetching. + def test_watermark_refreshes_during_one_extend_make_one_wire_call + @leases.replace(lease_entry(granted_amount: 1000)) + # Draw the slot down under the water mark. + @leases.try_reserve("co_1", "ct_1", 900) + started = Queue.new + release = Queue.new + @wire.queue_extend({ "lease" => { "lease_id" => "lse_1", "granted_amount" => 2000, + "expires_at_ms" => 300_000 } }) + @wire.during_extend = -> do + started << true + release.pop + end + + first = @manager.maybe_extend_in_background("co_1", "ct_1") + started.pop + followers = Array.new(10) { @manager.maybe_extend_in_background("co_1", "ct_1") } + release << true + ([first] + followers).compact.each(&:join) + + assert_equal 1, @wire.extend_calls.size + # No thread each: the credits they want are the ones the flight is already + # fetching, so there is nothing for a thread to do but find it and exit. + assert_empty followers.compact + end + + # A caller naming required_credits still spawns while a flight runs: it has a + # reserve to retry, and may need more than that flight asked for. + def test_a_check_naming_its_shortfall_still_spawns_during_one_extend + @leases.replace(lease_entry(granted_amount: 1000)) + @leases.try_reserve("co_1", "ct_1", 900) + started = Queue.new + release = Queue.new + 2.times do + @wire.queue_extend({ "lease" => { "lease_id" => "lse_1", "granted_total" => 20_000, + "expires_at_ms" => 300_000 } }) + end + @wire.during_extend = -> do + started << true + release.pop + end + + first = @manager.maybe_extend_in_background("co_1", "ct_1") + started.pop + joiner = @manager.maybe_extend_in_background("co_1", "ct_1", 5000) + release << true + + refute_nil joiner + [first, joiner].compact.each(&:join) + end + + # A follow-up that joins a flight asking for less than it needs, and finds + # another equally small follow-up on the way round, must still end with an + # extend of its own: returning the small one leaves its retry short with + # credits sitting on the server. + def test_a_follow_up_waits_out_a_too_small_flight_and_then_extends_itself + @leases.replace(lease_entry(granted_amount: 1000)) + @leases.try_reserve("co_1", "ct_1", 1000) + key = Leases.lease_key("co_1", "ct_1") + inflight = @manager.instance_variable_get(:@inflight_extend) + inflight[key] = Leases::Flight.new(2000) + # The second read is the caller waking from that flight: by then another + # caller's follow-up, just as small, holds the slot. + reads = 0 + leases = @leases + @leases.define_singleton_method(:get) do |company_id, credit_type_id| + reads += 1 + if reads == 2 + landed = Leases::Flight.new(2000) + landed.complete(leases.get("co_1", "ct_1")) + inflight[key] = landed + end + super(company_id, credit_type_id) + end + @wire.queue_extend({ "lease" => { "lease_id" => "lse_1", "granted_total" => 19_000, + "expires_at_ms" => 600_000 } }) + + waiter = Thread.new { @manager.send(:extend_if_needed, "co_1", "ct_1", 18_000, nil) } + inflight[key].complete(@leases.get("co_1", "ct_1")) + entry = waiter.value + + assert_equal 1, @wire.extend_calls.size + assert_in_delta 18_000, @wire.extend_calls.first.additional_amount + assert_in_delta 18_000, entry.local_remaining_credits + end + + # A joiner waits on a call started by someone else, on someone else's timeout. + # A check with a budget of its own must not sit behind it. + def test_a_joiner_gives_up_on_a_shared_extend_at_its_own_timeout + @leases.replace(lease_entry(granted_amount: 1000)) + @leases.try_reserve("co_1", "ct_1", 1000) + key = Leases.lease_key("co_1", "ct_1") + flight = Leases::Flight.new(2000) + @manager.instance_variable_get(:@inflight_extend)[key] = flight + + started = Process.clock_gettime(Process::CLOCK_MONOTONIC) + entry = @manager.send(:extend_if_needed, "co_1", "ct_1", 18_000, { timeout_in_seconds: 0.05 }) + elapsed_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started) * 1000 + + # No entry, so the check takes its failure path by mode, and the flight is + # left running for whoever else is on it. + assert_nil entry + assert_operator elapsed_ms, :<, 1000 + refute_predicate flight, :done? + assert_empty @wire.extend_calls + end + + # The deadline a check fixed at its start caps the join, not one taken at the + # join: time already spent on the acquire and reserve is not given back. + def test_a_joiner_spends_the_deadline_its_check_started_with + @leases.replace(lease_entry(granted_amount: 1000)) + @leases.try_reserve("co_1", "ct_1", 1000) + key = Leases.lease_key("co_1", "ct_1") + flight = Leases::Flight.new(20_000) + @manager.instance_variable_get(:@inflight_extend)[key] = flight + spent = Leases.join_deadline({ timeout_in_seconds: 0.05 }) - 50 + + started = Process.clock_gettime(Process::CLOCK_MONOTONIC) + entry = @manager.send(:extend_if_needed, "co_1", "ct_1", 18_000, { timeout_in_seconds: 5 }, spent) + elapsed_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started) * 1000 + + assert_nil entry + assert_operator elapsed_ms, :<, 1000 + refute_predicate flight, :done? + end + + # An acquire joiner is capped the same way: someone else's acquire runs on + # someone else's timeout. + def test_an_acquire_joiner_gives_up_at_its_own_deadline + key = Leases.lease_key("co_1", "ct_1") + flight = Leases::Flight.new + @manager.instance_variable_get(:@inflight_acquire)[key] = flight + + started = Process.clock_gettime(Process::CLOCK_MONOTONIC) + entry = @manager.acquire_if_needed("co_1", "ct_1", { timeout_in_seconds: 0.05 }) + elapsed_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started) * 1000 + + assert_nil entry + assert_operator elapsed_ms, :<, 1000 + refute_predicate flight, :done? + assert_empty @wire.acquire_calls + end + + # stop and enlist share the flight lock, so once stop returns no acquire can + # register behind the drain that follows it. + def test_an_acquire_cannot_register_once_the_manager_is_stopped + @manager.stop + @wire.queue_acquire(acquire_script) + + assert_nil @manager.acquire_if_needed("co_1", "ct_1") + assert_empty @wire.acquire_calls + end + + # A stop landing mid-acquire leaves the lease installed, for the drain that + # follows or for server-side expiry. Releasing it here would refund a lease + # sibling processes sharing the backend are still reserving against. + def test_a_lease_acquired_during_a_stop_is_left_for_the_drain + @wire.queue_acquire(acquire_script) + manager = @manager + @wire.during_acquire = -> { manager.stop } + + @manager.acquire_if_needed("co_1", "ct_1") + + assert_empty @wire.release_calls + assert_equal "lse_1", @leases.get("co_1", "ct_1").lease_id + end + + # The drain budget is what a caller asked close to take in total. Spending it + # thread by thread would multiply it by however many are stalled. + def test_drain_bounds_the_total_wait_across_stalled_threads + release = Queue.new + 5.times { @manager.send(:track) { release.pop } } + + started = Process.clock_gettime(Process::CLOCK_MONOTONIC) + @manager.drain(300) + elapsed_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started) * 1000 + + # One budget, not five. The ceiling is loose enough for scheduling noise and + # far below the 1500ms a per-thread budget would cost. + assert_operator elapsed_ms, :<, 900 + ensure + 5.times { release << true } + end + + def test_acquire_returns_nil_rather_than_raising_when_the_wire_fails + @wire.queue_acquire({ "error" => "wire down" }) + + assert_nil @manager.acquire_if_needed("co_1", "ct_1") + end + + def test_acquire_is_refused_once_the_manager_is_stopped + @manager.stop + @wire.queue_acquire(acquire_script) + + assert_nil @manager.acquire_if_needed("co_1", "ct_1") + assert_empty @wire.acquire_calls + end + + # The server treats an expired lease as released and has already refunded its + # remainder, so extending it would resurrect a stale local row. + def test_expired_leases_are_never_extended + @leases.replace(lease_entry(granted_amount: 1000, expires_at_ms: 1000)) + @leases.try_reserve("co_1", "ct_1", 900) + clock.advance_ms(2000) + + @manager.maybe_extend_in_background("co_1", "ct_1")&.join + + assert_empty @wire.extend_calls + end + + # An extend is an increment, so every attempt of one extend must carry the + # same key while a later extend gets its own. + def test_each_extend_mints_a_fresh_idempotency_key + @leases.replace(lease_entry(granted_amount: 1000)) + @leases.try_reserve("co_1", "ct_1", 900) + @wire.queue_extend({ "lease" => { "granted_total" => 2000, "expires_at_ms" => 600_000 } }) + @manager.maybe_extend_in_background("co_1", "ct_1")&.join + + @leases.try_reserve("co_1", "ct_1", 1000) + @wire.queue_extend({ "lease" => { "granted_total" => 3000, "expires_at_ms" => 900_000 } }) + @manager.maybe_extend_in_background("co_1", "ct_1")&.join + + keys = @wire.extend_calls.map(&:idempotency_key) + + assert_equal 2, keys.compact.uniq.size + end + + def test_release_all_local_leases_skips_expired_and_drops_released + @leases.replace(lease_entry(lease_id: "lse_live")) + @leases.replace(lease_entry(lease_id: "lse_dead", credit_type_id: "ct_2", expires_at_ms: 1000)) + clock.advance_ms(2000) + + @manager.release_all_local_leases + + assert_equal [], @wire.release_calls - ["lse_live"] + assert_equal ["lse_dead"], @leases.list.map(&:lease_id) + end + + # The close calls this before it shuts the DataStream and the event buffer, so + # a store that cannot list must not raise out of it. + def test_release_all_local_leases_survives_a_failed_listing + @leases.define_singleton_method(:list) { raise "store down" } + + assert_nil @manager.release_all_local_leases + assert_empty @wire.release_calls + end + + # A shared store never reports its contents, so there is nothing for the close + # to release out from under a sibling process. + def test_release_all_local_leases_is_a_no_op_without_an_enumerable_store + shared = Leases::RedisLeaseStore.new(client: LeaseSupport::FakeRedis.new(clock), clock: clock.to_proc) + manager = Leases::LeaseManager.new(wire_client: @wire, lease_store: shared, logger: silent_logger, + config: {}, clock: clock.to_proc) + shared.replace(lease_entry) + + manager.release_all_local_leases + + assert_empty @wire.release_calls + end + + def test_drain_returns_once_background_work_has_landed + @leases.replace(lease_entry(granted_amount: 1000)) + @leases.try_reserve("co_1", "ct_1", 900) + @wire.queue_extend({ "lease" => { "granted_total" => 2000, "expires_at_ms" => 600_000 } }) + @manager.maybe_extend_in_background("co_1", "ct_1") + + @manager.drain + + assert_equal 1, @wire.extend_calls.size + assert_in_delta 1100, @leases.get("co_1", "ct_1").local_remaining_credits + end +end + +class WireClientTest < Minitest::Test + include CreditTestHelpers + + LEASE_BODY = { + "data" => { + "id" => "lse_1", "company_id" => "co_1", "credit_type_id" => "ct_1", + "granted_amount" => 1000, "tracked_amount" => 0, + "expires_at" => "2026-01-01T00:05:00Z", "created_at" => "2026-01-01T00:00:00Z", + "updated_at" => "2026-01-01T00:00:00Z" + }, + "params" => {} + }.freeze + + def setup + @client = Schematic::Client.new(api_key: "sch_test", base_url: "https://api.schematichq.test") + @wire = Leases::ApiWireClient.new(credits_client: @client.credits) + end + + def test_acquire_sends_the_requested_amount_and_expiry + stub_request(:post, "https://api.schematichq.test/billing/credits/lease") + .to_return(status: 200, body: JSON.generate(LEASE_BODY), headers: { "Content-Type" => "application/json" }) + + grant = @wire.acquire(company_id: "co_1", credit_type_id: "ct_1", requested_amount: 1000, + expires_at: Time.utc(2026, 1, 1, 0, 5)) + + assert_equal "lse_1", grant.lease_id + assert_in_delta 1000, grant.granted_amount + assert_equal Time.utc(2026, 1, 1, 0, 5), grant.expires_at + assert_requested(:post, "https://api.schematichq.test/billing/credits/lease") do |req| + body = JSON.parse(req.body) + body["company_id"] == "co_1" && body["requested_amount"] == 1000 && + body["expires_at"] == "2026-01-01T00:05:00Z" + end + end + + # Read as host-local time, a zone-less expiry would move by the host's UTC + # offset, so a lease could look live hours after the server expired it. + def test_an_expiry_without_a_zone_is_read_as_utc + body = LEASE_BODY.merge("data" => LEASE_BODY["data"].merge("expires_at" => "2026-01-01T00:05:00")) + stub_request(:post, "https://api.schematichq.test/billing/credits/lease") + .to_return(status: 200, body: JSON.generate(body), headers: { "Content-Type" => "application/json" }) + + grant = with_tz("America/New_York") do + @wire.acquire(company_id: "co_1", credit_type_id: "ct_1", requested_amount: 1000, + expires_at: Time.utc(2026, 1, 1, 0, 5)) + end + + assert_equal Time.utc(2026, 1, 1, 0, 5), grant.expires_at + end + + def test_an_expiry_with_an_offset_keeps_it + with_tz("America/New_York") do + assert_equal Time.utc(2026, 1, 1, 0, 5), Leases.parse_api_time("2026-01-01T02:05:00+02:00") + assert_equal Time.utc(2026, 1, 1, 0, 5), Leases.parse_api_time("2026-01-01T00:05:00Z") + end + end + + def with_tz(zone) + previous = ENV.fetch("TZ", nil) + ENV["TZ"] = zone + yield + ensure + ENV["TZ"] = previous + end + + def test_extend_sends_the_additional_amount + stub_request(:put, "https://api.schematichq.test/billing/credits/lease/lse_1/extend") + .to_return(status: 200, body: JSON.generate(LEASE_BODY), headers: { "Content-Type" => "application/json" }) + + @wire.extend(lease_id: "lse_1", additional_amount: 500, expires_at: Time.utc(2026, 1, 1, 0, 5), + idempotency_key: "key-1") + + assert_requested(:put, "https://api.schematichq.test/billing/credits/lease/lse_1/extend") do |req| + body = JSON.parse(req.body) + body["additional_amount"] == 500 && body["idempotency_key"] == "key-1" + end + end + + # A shortfall of 10.4 asked for as 10 leaves the retried reserve short by the + # same fraction every time, so the ask rounds up. + def test_a_fractional_amount_rounds_up_on_both_calls + stub_request(:post, "https://api.schematichq.test/billing/credits/lease") + .to_return(status: 200, body: JSON.generate(LEASE_BODY), headers: { "Content-Type" => "application/json" }) + stub_request(:put, "https://api.schematichq.test/billing/credits/lease/lse_1/extend") + .to_return(status: 200, body: JSON.generate(LEASE_BODY), headers: { "Content-Type" => "application/json" }) + + @wire.acquire(company_id: "co_1", credit_type_id: "ct_1", requested_amount: 10.4, + expires_at: Time.utc(2026, 1, 1, 0, 5)) + @wire.extend(lease_id: "lse_1", additional_amount: 10.4, expires_at: Time.utc(2026, 1, 1, 0, 5)) + + assert_requested(:post, "https://api.schematichq.test/billing/credits/lease") do |req| + JSON.parse(req.body)["requested_amount"] == 11 + end + assert_requested(:put, "https://api.schematichq.test/billing/credits/lease/lse_1/extend") do |req| + JSON.parse(req.body)["additional_amount"] == 11 + end + end + + # An extend is an increment, so the server needs a key to fold a retried one + # back into a single grant. The manager supplies one; a caller that does not + # still gets a key rather than an unguarded increment. + def test_extend_mints_a_key_when_the_caller_omits_one + keys = [] + record_idempotency_keys( + stub_request(:put, "https://api.schematichq.test/billing/credits/lease/lse_1/extend"), keys + ).to_return(status: 200, body: JSON.generate(LEASE_BODY), headers: { "Content-Type" => "application/json" }) + + 2.times do + @wire.extend(lease_id: "lse_1", additional_amount: 500, expires_at: Time.utc(2026, 1, 1, 0, 5)) + end + + assert_equal 2, keys.size + keys.each { |key| assert_match UUID_SHAPE, key } + refute_equal keys[0], keys[1], "a separate extend must not reuse the previous call's key" + end + + # The key is minted outside the client's retry loop, so a 500 that the client + # retries carries the same key and the server grants the tranche once. + def test_one_key_rides_every_attempt_of_a_single_extend + keys = [] + record_idempotency_keys( + stub_request(:put, "https://api.schematichq.test/billing/credits/lease/lse_1/extend"), keys + ).to_return( + { status: 500, body: JSON.generate({ "error" => "boom" }), headers: { "Content-Type" => "application/json" } }, + { status: 200, body: JSON.generate(LEASE_BODY), headers: { "Content-Type" => "application/json" } } + ) + + grant = @wire.extend(lease_id: "lse_1", additional_amount: 500, expires_at: Time.utc(2026, 1, 1, 0, 5)) + + assert_equal "lse_1", grant.lease_id + assert_equal 2, keys.size + assert_equal keys[0], keys[1] + end +end + +class CheckFlowTest < Minitest::Test + include CreditTestHelpers + + CREDIT_ENTITLEMENT = { + "value_type" => "credit", "credit_id" => "ct_1", "consumption_rate" => 10, + "event_subtype" => "inference_tokens" + }.freeze + + def setup + @wire = LeaseSupport::ScriptedWireClient.new(clock) + @leases = Leases::LeaseStore.new(clock: clock.to_proc) + @reservations = Leases::ReservationStore.new(@leases, 1000, clock: clock.to_proc) + @manager = Leases::LeaseManager.new( + wire_client: @wire, lease_store: @leases, logger: silent_logger, + config: { default_lease_size: 1000, default_reservation_ttl: 60_000 }, clock: clock.to_proc + ) + @events = [] + end + + def teardown + @reservations.stop + @manager.stop + end + + def run_check(results, usage: 10, event_subtype: "inference_tokens", on_acquire_failure: nil, + balances: { "ct_1" => 5000 }, reservations: @reservations) + datastream = LeaseSupport::ScriptedDataStream.new( + flag_key: "inference", company: { id: "co_1", credit_balances: balances }, results: results + ) + @datastream = datastream + @fell_back = false + Leases.check_with_lease( + Leases::CheckDeps.new( + lease_store: @leases, reservations: reservations, manager: @manager, datastream: datastream, + logger: silent_logger, clock: clock.to_proc, + enqueue_flag_check_event: ->(body) { @events << body } + ), + "inference", { company: { "id" => "co_1" } }, + { usage: usage, event_subtype: event_subtype, on_acquire_failure: on_acquire_failure } + ) do + @fell_back = true + Leases::CheckResult.new(allowed: true, value: true, reason: "fallback", flag_key: "inference") + end + end + + def probe(entitlement = CREDIT_ENTITLEMENT) + { "value" => true, "reason" => "probe", "entitlement" => entitlement } + end + + # The WASM engine hands back camelCase keys, which a caller must not have to + # read differently from the server-mode model. + def test_the_engine_entitlement_is_normalized_to_snake_case + entitlement = CREDIT_ENTITLEMENT.merge("metric_reset_at" => "2026-02-01T00:00:00Z") + @leases.replace(lease_entry) + result = run_check([probe(entitlement), + { "value" => true, "reason" => "ok", "entitlement" => entitlement }]) + + assert_predicate result, :allowed? + assert_equal "inference", result.entitlement[:feature_key] + assert_equal "credit", result.entitlement[:value_type] + assert_in_delta 10, result.entitlement[:consumption_rate] + assert_equal Time.utc(2026, 2, 1), result.entitlement[:metric_reset_at] + end + + # Every branch tests for :fail_closed, so a typo would quietly read as + # fail-open and turn a gate permissive. + def test_a_misspelled_failure_mode_fails_closed_and_says_so + logger = RecordingLogger.new + @wire.queue_acquire({ "error" => "wire down" }) + datastream = LeaseSupport::ScriptedDataStream.new( + flag_key: "inference", company: { id: "co_1", credit_balances: { "ct_1" => 5000 } }, + results: [probe, { "value" => true, "reason" => "ok" }] + ) + result = Leases.check_with_lease( + Leases::CheckDeps.new( + lease_store: @leases, reservations: @reservations, manager: @manager, datastream: datastream, + logger: logger, clock: clock.to_proc, enqueue_flag_check_event: ->(body) { @events << body } + ), + "inference", { company: { "id" => "co_1" } }, + { usage: 10, event_subtype: "inference_tokens", on_acquire_failure: :fail_closd } + ) { Leases::CheckResult.new(allowed: true, value: true, reason: "fallback", flag_key: "inference") } + + refute_predicate result, :allowed? + assert(logger.warnings.any? { |w| w.include?("fail_closd") }) + end + + # A fraction of an event is not something the server bills, so the hold is + # sized in whole event units. The reservation still records what the caller + # declared. + def test_a_fractional_usage_rounds_the_hold_up_to_a_whole_unit + @leases.replace(lease_entry) + result = run_check([probe, { "value" => true, "reason" => "ok" }], usage: 0.5) + + assert_predicate result, :allowed? + assert_in_delta 10, result.reservation.credits_reserved + assert_in_delta 0.5, result.reservation.quantity_reserved + assert_in_delta 990, @leases.get("co_1", "ct_1").local_remaining_credits + end + + # A boolean or override grant resolves without drawing a credit, so it must + # not cost a lease acquire and a reserve-then-cancel. + def test_a_non_credit_entitlement_falls_back_without_touching_the_wire + @leases.replace(lease_entry) + result = run_check([probe({ "value_type" => "boolean" })]) + + assert @fell_back + assert_nil result.reservation + assert_empty @wire.acquire_calls + assert_in_delta 1000, @leases.get("co_1", "ct_1").local_remaining_credits + end + + def test_an_incomplete_credit_entitlement_falls_back + @leases.replace(lease_entry) + run_check([probe({ "value_type" => "credit", "credit_id" => "ct_1", "consumption_rate" => 0 })]) + + assert @fell_back + end + + # The reservation is pinned to the lease try_reserve charged, not the one the + # acquire handed back, so the settle refund names the lease actually debited. + def test_the_reservation_pins_the_charged_lease + @leases.replace(lease_entry(lease_id: "lse_current")) + result = run_check([probe, { "value" => true, "reason" => "ok" }]) + + assert_predicate result, :allowed? + assert_equal "lse_current", result.reservation.lease_id + end + + # A lease-path check resolves itself, so it has to enqueue the flag_check + # event the plain paths would have. + def test_a_lease_path_check_reports_a_flag_check_event + @leases.replace(lease_entry) + run_check([probe, { "value" => true, "reason" => "ok" }]) + + assert_equal 1, @events.size + assert_equal "inference", @events.first[:flag_key] + assert @events.first[:value] + end + + def test_a_fallback_exit_reports_no_flag_check_event + @leases.replace(lease_entry) + run_check([probe({ "value_type" => "boolean" })]) + + assert_empty @events + end + + # Fail-open is not a blanket allow: the engine still runs, so a company that + # is not entitled stays denied with the lease backend down. + def test_fail_open_still_denies_an_unentitled_company + @wire.queue_acquire({ "error" => "wire down" }) + result = run_check([probe, { "value" => false, "reason" => "not_entitled" }], on_acquire_failure: :fail_open) + + refute_predicate result, :allowed? + assert_equal "lease_acquire_failed", result.error + end + + def test_fail_open_accepts_the_hyphenated_spelling + @wire.queue_acquire({ "error" => "wire down" }) + result = run_check([probe, { "value" => true, "reason" => "evaluated" }], on_acquire_failure: "fail-open") + + assert_predicate result, :allowed? + end + + # An engine that is itself down leaves no evaluation to fail open with. + def test_an_engine_failure_during_the_gate_cancels_the_hold + @leases.replace(lease_entry) + datastream = LeaseSupport::ScriptedDataStream.new( + flag_key: "inference", company: { id: "co_1", credit_balances: { "ct_1" => 5000 } }, + results: [probe] + ) + result = Leases.check_with_lease( + Leases::CheckDeps.new( + lease_store: @leases, reservations: @reservations, manager: @manager, datastream: datastream, + logger: silent_logger, clock: clock.to_proc, enqueue_flag_check_event: ->(body) { @events << body } + ), + "inference", { company: { "id" => "co_1" } }, + { usage: 10, event_subtype: "inference_tokens" } + ) { Leases::CheckResult.new(allowed: true, value: true, reason: "fallback", flag_key: "inference") } + + refute_predicate result, :allowed? + assert_in_delta 1000, @leases.get("co_1", "ct_1").local_remaining_credits + assert_equal 0, @reservations.size + end + + # A store that cannot record the hold must not leave the debit stranded. + def test_a_failed_reservation_persist_undoes_the_debit + @leases.replace(lease_entry) + failing = Object.new + def failing.add(_reservation) = raise("redis down") + def failing.consume(_id, _credits) = nil + + result = run_check([probe, { "value" => true, "reason" => "ok" }], reservations: failing) + + refute_predicate result, :allowed? + assert_equal "lease_store_error", result.error + assert_in_delta 1000, @leases.get("co_1", "ct_1").local_remaining_credits + end + + # A NaN debit would slip through every comparison and poison a possibly shared + # lease balance into approving everything. + def test_a_nan_usage_never_reaches_the_store + @leases.replace(lease_entry) + result = run_check([], usage: Float::NAN) + + refute_predicate result, :allowed? + assert_equal "invalid_usage", result.error + assert_in_delta 1000, @leases.get("co_1", "ct_1").local_remaining_credits + end + + def test_zero_usage_falls_back_without_issuing_a_handle + @leases.replace(lease_entry) + result = run_check([], usage: 0) + + assert @fell_back + assert_nil result.reservation + end +end + +class TrackSettleTest < Minitest::Test + include CreditTestHelpers + + def setup + @leases = Leases::LeaseStore.new(clock: clock.to_proc) + @leases.replace(lease_entry) + @reservations = Leases::ReservationStore.new(@leases, 1000, clock: clock.to_proc) + end + + def teardown + @reservations.stop + end + + def test_a_client_mode_settle_routes_through_the_lease_sub_ledger + @leases.try_reserve("co_1", "ct_1", 100) + @reservations.add(reservation) + + outcome = Leases.consume_reservation_and_build_event(@reservations, reservation, 4) + + assert outcome.settled_locally + assert_equal "lse_1", outcome.track[:lease_id] + assert_nil outcome.track[:reservation_id] + assert_in_delta 4, outcome.track[:quantity] + assert_in_delta 960, @leases.get("co_1", "ct_1").local_remaining_credits + end + + # The server prefers a lease id when both are set, and a server-mode hold has + # no lease for it to route through. + def test_a_server_mode_settle_names_the_reservation_and_not_a_lease + body = Leases.build_reservation_track_event(reservation(mode: :server), 4) + + assert_equal "res_1", body[:reservation_id] + assert_nil body[:lease_id] + end + + # The event's quantity is an integer on the wire, so a partial unit bills as a + # whole one, and the debit rounds up with it: the local ledger moves by + # exactly what the event costs. + def test_a_fractional_settle_bills_and_debits_whole_units + @leases.try_reserve("co_1", "ct_1", 100) + @reservations.add(reservation) + + outcome = Leases.consume_reservation_and_build_event(@reservations, reservation, 1.2) + + assert_equal 2, outcome.track[:quantity] + assert_in_delta 980, @leases.get("co_1", "ct_1").local_remaining_credits + end + + # A quantity a hair above a whole unit bills as that unit, so the debit has to + # shave the same float noise. A bare ceil would debit for 4 and bill for 3. + def test_float_noise_debits_and_bills_the_same_whole_units + @leases.try_reserve("co_1", "ct_1", 100) + @reservations.add(reservation) + + outcome = Leases.consume_reservation_and_build_event(@reservations, reservation, (0.1 + 0.2) * 10) + + assert_equal 3, outcome.track[:quantity] + assert_in_delta 970, @leases.get("co_1", "ct_1").local_remaining_credits + end + + # A hold swept at its TTL still has to bill: the event is built from the + # caller-held handle, not the store. + def test_a_settle_after_the_sweep_still_bills_as_a_recovery_emit + @leases.try_reserve("co_1", "ct_1", 100) + @reservations.add(reservation(expires_at_ms: 10_000)) + clock.advance_ms(60_000) + @reservations.sweep_expired + + outcome = Leases.consume_reservation_and_build_event(@reservations, reservation, 4) + + refute outcome.settled_locally + assert_equal "inference_tokens", outcome.track[:event] + assert_in_delta 1000, @leases.get("co_1", "ct_1").local_remaining_credits + end + + def test_the_event_carries_the_evaluation_context_and_caller_traits + body = Leases.build_reservation_track_event(reservation, 4, traits: { "model" => "sonnet" }) + + assert_equal({ "id" => "co_1" }, body[:company]) + assert_equal({ "model" => "sonnet" }, body[:traits]) + end +end + +class ServerCheckTest < Minitest::Test + include CreditTestHelpers + + RESERVATION_BODY = { + "id" => "res_1", "company_id" => "co_1", "credit_type_id" => "ct_1", "consumption_rate" => 10, + "credits_reserved" => 100, "quantity_reserved" => 10, "event_subtype" => "inference_tokens", + "expires_at" => "2026-01-01T00:01:00Z" + }.freeze + + def setup + @client = Schematic::Client.new(api_key: "sch_test", base_url: "https://api.schematichq.test") + @deps = Leases::ServerCheckDeps.new( + features: @client.features, credits: @client.credits, logger: silent_logger, + reservation_ttl_ms: 60_000, default_value: -> { false }, clock: clock.to_proc + ) + end + + def stub_check_and_reserve(status: 200, body: nil) + stub_request(:post, "https://api.schematichq.test/flags/inference/check-and-reserve") + .to_return(status: status, body: JSON.generate(body), headers: { "Content-Type" => "application/json" }) + end + + def run_server_check(options = {}) + Leases.check_with_server_reservation( + @deps, "inference", { company: { "id" => "co_1" } }, + { usage: 10, event_subtype: "inference_tokens" }.merge(options) + ) { Leases::CheckResult.new(allowed: true, value: true, reason: "fallback", flag_key: "inference") } + end + + def test_a_granted_hold_becomes_a_server_mode_handle + stub_check_and_reserve(body: { "data" => { "flag" => "inference", "value" => true, + "reason" => "ok", "reservation" => RESERVATION_BODY } }) + + result = run_server_check + + assert_predicate result, :allowed? + assert_equal :server, result.reservation.mode + assert_equal "res_1", result.reservation.lease_id + assert_requested(:post, "https://api.schematichq.test/flags/inference/check-and-reserve") do |req| + body = JSON.parse(req.body) + body["quantity"] == 10 && body["expires_at"] == "2026-01-01T00:01:00Z" && + body.dig("preflight", "event_usage", "quantity") == 10 + end + end + + def test_a_denied_check_returns_no_handle + stub_check_and_reserve(body: { "data" => { "flag" => "inference", "value" => false, + "reason" => "Insufficient credits" } }) + + result = run_server_check + + refute_predicate result, :allowed? + assert_nil result.reservation + end + + # A 402 is the server's definitive answer, not a can't-gate, so it denies even + # under fail-open: failing open would hand out credit the balance cannot cover. + def test_a_402_denies_even_when_configured_to_fail_open + stub_check_and_reserve(status: 402, body: { "error" => "insufficient credits" }) + + result = run_server_check(on_acquire_failure: :fail_open) + + refute_predicate result, :allowed? + assert_equal Leases::INSUFFICIENT_CREDITS_REASON, result.reason + end + + # Server mode has no local engine to re-run, so fail-open returns the caller's + # default. + def test_fail_open_returns_the_caller_default_when_the_call_fails + stub_check_and_reserve(status: 500, body: { "error" => "boom" }) + @deps.default_value = -> { true } + + result = run_server_check(on_acquire_failure: :fail_open) + + assert_predicate result, :allowed? + assert_equal "server_reservation_failed", result.error + end + + def test_fail_closed_denies_when_the_call_fails + stub_check_and_reserve(status: 500, body: { "error" => "boom" }) + + result = run_server_check + + refute_predicate result, :allowed? + assert_equal "server_reservation_failed", result.error + end + + # A hold with no subtype could never be settled, so it is released now rather + # than parking credits until its TTL. + def test_a_hold_with_no_event_subtype_is_released + stub_check_and_reserve(body: { "data" => { "flag" => "inference", "value" => true, "reason" => "ok", + "reservation" => RESERVATION_BODY.merge("event_subtype" => nil) } }) + release = stub_request(:put, "https://api.schematichq.test/billing/credits/reservations/res_1/release") + .to_return(status: 200, body: "{}", headers: { "Content-Type" => "application/json" }) + + result = run_server_check(event_subtype: nil) + + refute_predicate result, :allowed? + assert_equal "missing_event_subtype", result.error + assert_requested release + end + + # The request body's quantity is an integer, so a fractional usage would + # truncate and the server would size the hold below the work about to run. + def test_a_fractional_usage_rounds_up_on_the_wire + stub_check_and_reserve(body: { "data" => { "flag" => "inference", "value" => true, + "reason" => "ok", "reservation" => RESERVATION_BODY } }) + + run_server_check(usage: 1.5) + + assert_requested(:post, "https://api.schematichq.test/flags/inference/check-and-reserve") do |req| + body = JSON.parse(req.body) + body["quantity"] == 2 && body.dig("preflight", "event_usage", "quantity") == 2 + end + end + + # Server mode gets its entitlement as a generated model and client mode as a + # camelCase engine hash. A caller reading result.entitlement must not need to + # know which one answered. + def test_the_entitlement_is_normalized_to_one_shape + stub_check_and_reserve( + body: { "data" => { "flag" => "inference", "value" => true, "reason" => "ok", + "reservation" => RESERVATION_BODY, + "entitlement" => { "feature_id" => "feat_1", "feature_key" => "inference", + "value_type" => "credit", "credit_id" => "ct_1", + "consumption_rate" => 10, + "metric_reset_at" => "2026-02-01T00:00:00Z" } } } + ) + + result = run_server_check + + assert_instance_of Hash, result.entitlement + assert_equal "inference", result.entitlement[:feature_key] + assert_equal "credit", result.entitlement[:value_type] + assert_equal Time.utc(2026, 2, 1), result.entitlement[:metric_reset_at] + end + + def test_a_misspelled_failure_mode_fails_closed_and_says_so + logger = RecordingLogger.new + @deps.logger = logger + stub_check_and_reserve(status: 500, body: { "error" => "boom" }) + @deps.default_value = -> { true } + + result = run_server_check(on_acquire_failure: "fail-opne") + + refute_predicate result, :allowed? + assert(logger.warnings.any? { |w| w.include?("fail-opne") }) + end + + def test_each_check_sends_its_own_idempotency_key + keys = [] + record_idempotency_keys( + stub_request(:post, "https://api.schematichq.test/flags/inference/check-and-reserve"), keys + ).to_return( + status: 200, + body: JSON.generate({ "data" => { "flag" => "inference", "value" => true, "reason" => "ok", + "reservation" => RESERVATION_BODY } }), + headers: { "Content-Type" => "application/json" } + ) + + 2.times { run_server_check } + + assert_equal 2, keys.size + keys.each { |key| assert_match UUID_SHAPE, key } + refute_equal keys[0], keys[1], "a separate check must not reuse the previous call's key" + end + + # check-and-reserve takes a hold, so a 500 the client retries would otherwise + # take a second one. The key is minted per check, not per attempt, so the + # server hands back the hold it already took and the check still resolves. + def test_a_retried_check_keeps_one_key_and_still_yields_the_hold + keys = [] + record_idempotency_keys( + stub_request(:post, "https://api.schematichq.test/flags/inference/check-and-reserve"), keys + ).to_return( + { status: 500, body: JSON.generate({ "error" => "boom" }), headers: { "Content-Type" => "application/json" } }, + { status: 200, + body: JSON.generate({ "data" => { "flag" => "inference", "value" => true, "reason" => "ok", + "reservation" => RESERVATION_BODY } }), + headers: { "Content-Type" => "application/json" } } + ) + + result = run_server_check + + assert_predicate result, :allowed? + assert_equal "res_1", result.reservation.id + assert_equal 2, keys.size + assert_equal keys[0], keys[1] + end +end + +class CreditLeaseClientWiringTest < Minitest::Test + include CreditTestHelpers + + def build_client(credit_leases:, **) + Schematic::SchematicClient.new( + api_key: "sch_test", base_url: "https://api.schematichq.test", + credit_leases: credit_leases, logger: silent_logger, ** + ) + end + + # Without DataStream, auto resolves to server mode rather than dropping every + # check to a plain, ungated flag check. + def test_auto_without_datastream_resolves_to_server_mode + client = build_client(credit_leases: { default_reservation_ttl: 60_000 }) + + assert_equal :server, client.send(:effective_lease_mode) + assert_nil client.instance_variable_get(:@lease_store) + ensure + client&.close + end + + # A DataStream whose start fails after construction clears the client, so auto + # has to be resolved per check or every check would drop to a plain, ungated + # flag check. + def test_auto_resolves_per_check_and_falls_to_server_when_the_datastream_goes_away + client = build_client(credit_leases: { mode: :auto }) + client.instance_variable_set(:@datastream_client, Object.new) + %i[@credit_lease_manager @lease_store @reservations].each do |ivar| + client.instance_variable_set(ivar, Object.new) + end + + assert_equal :client, client.send(:effective_lease_mode) + + client.instance_variable_set(:@datastream_client, nil) + + assert_equal :server, client.send(:effective_lease_mode) + ensure + %i[@credit_lease_manager @lease_store @reservations].each do |ivar| + client&.instance_variable_set(ivar, nil) + end + client&.close + end + + # The API refuses a hold expiring more than an hour out, measured against its + # own clock, so a server-mode TTL is clamped a step below the cap. + def test_a_server_mode_reservation_ttl_is_clamped_below_the_api_cap + client = build_client(credit_leases: { mode: :server, default_reservation_ttl: 2 * 60 * 60 * 1000 }) + + assert_equal Leases::MAX_RESERVATION_TTL_MS - Leases::RESERVATION_TTL_SKEW_ALLOWANCE_MS, + client.instance_variable_get(:@server_reservation_ttl_ms) + ensure + client&.close + end + + # In client mode the TTL sizes the local sweep, so clamping it would shorten + # holds for no reason. + def test_a_client_mode_reservation_ttl_is_left_alone + client = build_client(credit_leases: { mode: :client, default_reservation_ttl: 2 * 60 * 60 * 1000 }) + + assert_equal 2 * 60 * 60 * 1000, client.instance_variable_get(:@server_reservation_ttl_ms) + ensure + client&.close + end + + # An existing Redis setup backs leases automatically, with no second client to + # wire up. + def test_lease_state_reuses_the_datastream_redis_client + redis = LeaseSupport::FakeRedis.new(clock) + client = build_client( + credit_leases: { mode: :client }, use_data_stream: false, + datastream_options: { redis_client: redis, redis_key_prefix: "acme:" } + ) + + assert_instance_of Leases::RedisLeaseStore, client.instance_variable_get(:@lease_store) + assert client.instance_variable_get(:@lease_backend_shared) + ensure + client&.close + end + + # A shared lease is one row per company and credit across every instance, so a + # single process shutting down must not release it. + def test_close_does_not_release_leases_held_in_a_shared_backend + redis = LeaseSupport::FakeRedis.new(clock) + client = build_client( + credit_leases: { mode: :client }, datastream_options: { redis_client: redis } + ) + store = client.instance_variable_get(:@lease_store) + store.replace(lease_entry(expires_at_ms: 10 * 60 * 1000)) + released = stub_request(:put, %r{/billing/credits/lease/.*/release}).to_return(status: 200, body: "{}") + + client.close + + assert_not_requested released + end + + def test_close_releases_leases_held_per_process + client = build_client(credit_leases: { mode: :client }) + client.instance_variable_get(:@lease_store).replace( + Leases::LeaseEntry.new(lease_id: "lse_1", company_id: "co_1", credit_type_id: "ct_1", + granted_amount: 1000, expires_at: Time.now + 600) + ) + released = stub_request(:put, "https://api.schematichq.test/billing/credits/lease/lse_1/release") + .to_return(status: 200, body: "{}", headers: { "Content-Type" => "application/json" }) + + client.close + + assert_requested released + end + + # A check with a usage that falls back to a plain check and then cannot reach + # the API must answer with the caller's default, not the registered one. + def test_a_fallback_check_that_fails_honors_the_caller_default + client = Schematic::SchematicClient.new(api_key: "sch_test", base_url: "https://api.schematichq.test", + logger: silent_logger, flag_defaults: { "inference" => false }) + stub_request(:post, "https://api.schematichq.test/flags/inference/check") + .to_return(status: 500, body: JSON.generate({ "error" => "boom" }), + headers: { "Content-Type" => "application/json" }) + + result = client.check("inference", company: { "id" => "co_1" }, usage: 10, default_value: true) + + assert_predicate result, :allowed? + ensure + client&.close + end + + # A callable default is resolved at check time, the same as on the credit + # paths. + def test_a_callable_default_is_resolved_on_the_fallback_path + client = Schematic::SchematicClient.new(api_key: "sch_test", base_url: "https://api.schematichq.test", + logger: silent_logger) + stub_request(:post, "https://api.schematichq.test/flags/inference/check") + .to_return(status: 500, body: "{}", headers: { "Content-Type" => "application/json" }) + + result = client.check("inference", company: { "id" => "co_1" }, usage: 10, default_value: -> { true }) + + assert_predicate result, :allowed? + ensure + client&.close + end + + # The cached metric is a local prediction of what the stream will push back, + # so it belongs to an event that moved local state. A recovery emit dedupes + # server-side, and bumping the metric for one would deny the caller's next + # numeric-limit check until the stream corrected it. + def test_a_settle_that_did_not_land_locally_leaves_the_cached_metric_alone + client = build_client(credit_leases: { mode: :client }) + stub_request(:post, "https://c.schematichq.com/batch").to_return(status: 200, body: "") + metrics = [] + datastream = Object.new + datastream.define_singleton_method(:update_company_metrics) do |company, event, quantity| + metrics << [company, event, quantity] + end + def datastream.connected? = true + def datastream.close = nil + client.instance_variable_set(:@datastream_client, datastream) + leases = client.instance_variable_get(:@lease_store) + leases.replace(lease_entry) + leases.try_reserve("co_1", "ct_1", 100) + + # Never added to the store: swept at its TTL, or already settled. + client.track_with_reservation(reservation, 4) + + assert_empty metrics + + client.instance_variable_get(:@reservations).add(reservation) + client.track_with_reservation(reservation, 4) + + assert_equal [[{ "id" => "co_1" }, "inference_tokens", 4]], metrics + ensure + client&.close + end + + # { "ct_1": {} } is the symbol :ct_1, and a caller who writes it means the + # credit type, so both spellings have to reach the same slot, knobs included. + def test_an_override_written_either_way_resolves_the_same_knobs + [{ ct_1: { "default_lease_size" => 42 } }, { "ct_1" => { default_lease_size: 42 } }].each do |overrides| + client = build_client(credit_leases: { mode: :server, overrides: overrides }) + + resolved = Leases.resolve_config(client.instance_variable_get(:@credit_lease_config), "ct_1") + + assert_equal 42, resolved.lease_size + client.close + end + end + + # A malformed usage is the caller asking to gate on a number that gates + # nothing, so with gating configured it is refused by the failure mode rather + # than quietly dropped, which would allow the call with no hold. + def test_an_invalid_usage_is_denied_under_fail_closed + client = build_client(credit_leases: { mode: :server }) + + [-5, Float::NAN].each do |usage| + result = client.check("inference", company: { "id" => "co_1" }, usage: usage, + on_acquire_failure: :fail_closed) + + refute_predicate result, :allowed? + assert_equal "invalid_usage", result.reason + assert_nil result.reservation + end + ensure + client&.close + end + + # With no gating there is nothing to refuse: the value would only reach the + # preflight, so the plain question still gets asked. + def test_an_invalid_usage_without_credit_leases_asks_the_plain_question + client = Schematic::SchematicClient.new(api_key: "sch_test", base_url: "https://api.schematichq.test", + logger: silent_logger) + body = nil + stub_request(:post, "https://api.schematichq.test/flags/inference/check") + .with { |req| body = JSON.parse(req.body) } + .to_return(status: 200, + body: JSON.generate({ "data" => { "flag" => "inference", "value" => true, "reason" => "plan" } }), + headers: { "Content-Type" => "application/json" }) + stub_request(:post, "https://c.schematichq.com/batch").to_return(status: 200, body: "") + + result = client.check("inference", company: { "id" => "co_1" }, usage: -5) + + assert_predicate result, :allowed? + refute body.key?("preflight") + ensure + client&.close + end + + # An account is free to define an ordinary entity key called "id", so the + # keys are looked up first and the cache's answer wins over the literal value. + def test_a_plain_id_key_resolves_through_the_cache + client = build_client(credit_leases: { mode: :client, prewarm_resolve_timeout_ms: 0 }) + datastream = Object.new + def datastream.get_cached_company(_company) = { "id" => "co_42" } + def datastream.close = nil + client.instance_variable_set(:@datastream_client, datastream) + + assert_equal "co_42", client.send(:resolve_company_id_with_wait, { "id" => "acme" }) + ensure + client&.close + end + + # Only once the key lookup comes up empty is a value read as the company's own + # id, by its prefix rather than by the name of the key it sits under. + def test_a_cache_miss_falls_back_to_a_prefixed_value_under_any_key + client = build_client(credit_leases: { mode: :client, prewarm_resolve_timeout_ms: 0 }) + datastream = Object.new + def datastream.get_cached_company(_company) = nil + def datastream.close = nil + client.instance_variable_set(:@datastream_client, datastream) + + assert_equal "comp_9", client.send(:resolve_company_id_with_wait, { "account_id" => "comp_9" }) + ensure + client&.close + end + + def test_a_cache_miss_without_a_prefixed_value_resolves_to_nothing + client = build_client(credit_leases: { mode: :client, prewarm_resolve_timeout_ms: 0 }) + datastream = Object.new + def datastream.get_cached_company(_company) = nil + def datastream.close = nil + client.instance_variable_set(:@datastream_client, datastream) + + assert_nil client.send(:resolve_company_id_with_wait, { "id" => "acme" }) + ensure + client&.close + end + + # Zero means cache-only, not "never resolve": a company the DataStream already + # holds still prewarms, it is only the active fetch that is skipped. + def test_a_zero_prewarm_timeout_still_resolves_from_the_cache + client = build_client(credit_leases: { mode: :client, prewarm_resolve_timeout_ms: 0 }) + datastream = Object.new + def datastream.get_cached_company(_company) = { "id" => "co_1" } + def datastream.get_company(_company) = raise("prewarm must not fetch at a zero timeout") + def datastream.close = nil + client.instance_variable_set(:@datastream_client, datastream) + + assert_equal "co_1", client.send(:resolve_company_id_with_wait, { "org_id" => "acme" }) + ensure + client&.close + end + + def test_a_zero_prewarm_timeout_never_fetches + client = build_client(credit_leases: { mode: :client, prewarm_resolve_timeout_ms: 0 }) + datastream = Object.new + def datastream.get_cached_company(_company) = nil + def datastream.get_company(_company) = raise("prewarm must not fetch at a zero timeout") + def datastream.close = nil + client.instance_variable_set(:@datastream_client, datastream) + + assert_nil client.send(:resolve_company_id_with_wait, { "org_id" => "acme" }) + ensure + client&.close + end + + # The generated transport takes its timeout from the client, not from a + # request, so a per-check timeout is carried and not yet applied. This pins + # the half this SDK owns: the request does carry it, in seconds, so it starts + # working the moment the transport honors it. + def test_a_per_check_timeout_is_carried_on_the_request_in_seconds + client = Schematic::SchematicClient.new(api_key: "sch_test", base_url: "https://api.schematichq.test", + logger: silent_logger) + + assert_empty client.send(:api_request_options, nil) + assert_in_delta 1.5, client.send(:api_request_options, 1500)[:timeout_in_seconds] + # And pins the half it does not: the generated transport never reads the + # key. When a regeneration makes it read one, this fails, and the README + # note saying the timeout is inert has to go with it. + transport = File.read(File.expand_path("../lib/schematic/internal/http/raw_client.rb", __dir__)) + + refute_includes transport.gsub(/^\s*#.*$/, ""), "timeout_in_seconds" + ensure + client&.close + end + + # identify is a buffer push. Asking it to prewarm must not turn it into an + # HTTP round trip on the caller's thread, which is what the flush is. The + # ordering the prewarm poll depends on still holds, because the flush and the + # poll run on the same background thread in that order. + def test_identify_with_a_prewarm_does_not_flush_on_the_calling_thread + client = build_client(credit_leases: { mode: :client }) + flushing = Queue.new + release = Queue.new + buffer = Object.new + buffer.define_singleton_method(:push) { |_event| nil } + buffer.define_singleton_method(:stop) { nil } + buffer.define_singleton_method(:flush) do + flushing << Thread.current + release.pop + end + client.instance_variable_set(:@event_buffer, buffer) + + identifying = Thread.new do + client.identify({ keys: { "user_id" => "u_1" }, company: { keys: { "id" => "co_1" } } }, + prewarm: ["ct_1"]) + end + + # Joined with a bound rather than waited on: a flush back on the calling + # thread would park here forever instead of failing. + assert identifying.join(5), "identify blocked on the event buffer flush" + refute_same identifying, flushing.pop + ensure + release << true + client&.close + end + + # A plain check over the API gates on the balance as it stands. A check with a + # usage has to gate on the balance the action will leave behind, so the + # preflight goes on the REST request too. + def test_a_fallback_check_sends_the_preflight_on_the_rest_request + client = Schematic::SchematicClient.new(api_key: "sch_test", base_url: "https://api.schematichq.test", + logger: silent_logger) + stub_request(:post, "https://api.schematichq.test/flags/inference/check") + .to_return(status: 200, body: JSON.generate({ "data" => { "flag" => "inference", "value" => true, + "reason" => "ok" } }), + headers: { "Content-Type" => "application/json" }) + + client.check("inference", company: { "id" => "co_1" }, usage: 1.5, event_subtype: "inference_tokens") + + assert_requested(:post, "https://api.schematichq.test/flags/inference/check") do |req| + preflight = JSON.parse(req.body)["preflight"] + preflight["event_usage"]["event_subtype"] == "inference_tokens" && + preflight["event_usage"]["quantity"] == 2 + end + end + + # A zero usage is documented as having no effect, so such a check is a plain + # one. Sending an empty preflight would cost it the flag cache for nothing. + def test_a_zero_usage_check_sends_no_preflight_and_stays_cacheable + client = Schematic::SchematicClient.new(api_key: "sch_test", base_url: "https://api.schematichq.test", + logger: silent_logger) + plain = stub_request(:post, "https://api.schematichq.test/flags/inference/check") + .to_return(status: 200, body: JSON.generate({ "data" => { "flag" => "inference", "value" => true, + "reason" => "ok" } }), + headers: { "Content-Type" => "application/json" }) + + client.check("inference", company: { "id" => "co_1" }, usage: 0) + client.check("inference", company: { "id" => "co_1" }, usage: 0) + + # One request, so the second check was served from the cache the first + # populated. + assert_requested plain, times: 1 + assert_requested(:post, "https://api.schematichq.test/flags/inference/check") do |req| + !JSON.parse(req.body).key?("preflight") + end + ensure + client&.close + end + + # The flag cache is keyed by flag, company and user, so a preflighted verdict + # and a plain one would share an entry while answering different questions. + def test_a_preflighted_check_neither_reads_nor_writes_the_flag_cache + # The client builds its own local flag cache, which is what a plain check + # populates. + client = Schematic::SchematicClient.new(api_key: "sch_test", base_url: "https://api.schematichq.test", + logger: silent_logger) + plain = stub_request(:post, "https://api.schematichq.test/flags/inference/check") + .with { |req| !JSON.parse(req.body).key?("preflight") } + .to_return(status: 200, body: JSON.generate({ "data" => { "flag" => "inference", "value" => true, + "reason" => "plain" } }), + headers: { "Content-Type" => "application/json" }) + preflighted = stub_request(:post, "https://api.schematichq.test/flags/inference/check") + .with { |req| JSON.parse(req.body).key?("preflight") } + .to_return(status: 200, body: JSON.generate({ "data" => { "flag" => "inference", + "value" => false, + "reason" => "preflighted" } }), + headers: { "Content-Type" => "application/json" }) + + # A plain check caches its verdict. + client.check_flag("inference", company: { "id" => "co_1" }) + client.check_flag("inference", company: { "id" => "co_1" }) + + assert_requested plain, times: 1 + + # The preflighted check goes to the API rather than reading that entry. + result = client.check("inference", company: { "id" => "co_1" }, usage: 10) + + refute_predicate result, :allowed? + assert_requested preflighted, times: 1 + + # And leaves the cached plain verdict as it found it. + assert client.check_flag("inference", company: { "id" => "co_1" }) + assert_requested plain, times: 1 + ensure + client&.close + end + + # Everything that is not :client or :server falls through to the auto + # behaviour, so a typo would silently pick a mode the caller did not ask for. + def test_a_misspelled_mode_falls_back_to_auto_and_says_so + logger = RecordingLogger.new + client = Schematic::SchematicClient.new( + api_key: "sch_test", base_url: "https://api.schematichq.test", logger: logger, + credit_leases: { mode: :serverr, default_reservation_ttl: 60_000 } + ) + + assert_equal :auto, client.instance_variable_get(:@credit_lease_mode) + assert(logger.warnings.any? { |w| w.include?("serverr") }) + ensure + client&.close + end + + # Releases are synchronous round trips, so without a bound a slow API would + # stretch close by one timeout per leased slot, right after the drain was + # carefully bounded. + def test_releasing_leases_on_close_stops_when_the_budget_runs_out + logger = RecordingLogger.new + store = Leases::LeaseStore.new(clock: clock.to_proc) + 10.times do |i| + store.replace(Leases::LeaseEntry.new(lease_id: "lse_#{i}", company_id: "co_#{i}", + credit_type_id: "ct_1", granted_amount: 1000, + expires_at: Time.now + 600)) + end + wire = Object.new + released = [] + wire.define_singleton_method(:release) do |lease_id:, **| + released << lease_id + sleep 0.05 + end + manager = Leases::LeaseManager.new(wire_client: wire, lease_store: store, logger: logger, + clock: -> { Time.now }) + + started = Process.clock_gettime(Process::CLOCK_MONOTONIC) + manager.release_all_local_leases(150) + elapsed_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started) * 1000 + + assert_operator elapsed_ms, :<, 500 + assert_operator released.size, :<, 10 + assert(logger.warnings.any? { |w| w.include?("left to server-side expiry") }) + end + + # The Java port had this as a blocker: a check with a usage but no leases + # falls through to the REST check, and when that call cannot be made at all + # the caller's default has to answer. + def test_an_unreachable_api_answers_with_the_caller_default + client = Schematic::SchematicClient.new(api_key: "sch_test", base_url: "https://api.schematichq.test", + logger: silent_logger, flag_defaults: { "inference" => false }) + stub_request(:post, "https://api.schematichq.test/flags/inference/check").to_timeout + + result = client.check("inference", company: { "id" => "co_1" }, usage: 10, default_value: true) + + assert_predicate result, :allowed? + ensure + client&.close + end + + # A request sent over a closed socket is dropped without an error, so a fetch + # would sit out the DataStream's own 30 second timeout, many times the + # prewarm budget. + def test_a_disconnected_datastream_resolves_within_the_budget + client = build_client(credit_leases: { mode: :client, prewarm_resolve_timeout_ms: 200 }) + datastream = Object.new + def datastream.connected? = false + def datastream.get_cached_company(_company) = nil + def datastream.get_company(_company) = raise("a disconnected datastream must not be fetched from") + def datastream.close = nil + client.instance_variable_set(:@datastream_client, datastream) + + started = Process.clock_gettime(Process::CLOCK_MONOTONIC) + result = client.send(:resolve_company_id_with_wait, { "org_id" => "acme" }) + elapsed_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started) * 1000 + + assert_nil result + assert_operator elapsed_ms, :<, 600 + ensure + client&.close + end + + # And a connected but slow fetch is capped at the budget rather than the + # DataStream's own timeout. + def test_a_slow_connected_fetch_returns_at_the_deadline + client = build_client(credit_leases: { mode: :client, prewarm_resolve_timeout_ms: 200 }) + datastream = Object.new + def datastream.connected? = true + def datastream.get_cached_company(_company) = nil + + def datastream.get_company(_company) + sleep 30 + nil + end + + def datastream.close = nil + client.instance_variable_set(:@datastream_client, datastream) + + started = Process.clock_gettime(Process::CLOCK_MONOTONIC) + result = client.send(:resolve_company_id_with_wait, { "org_id" => "acme" }) + elapsed_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started) * 1000 + + assert_nil result + assert_operator elapsed_ms, :<, 1000 + ensure + client&.close + end + + # A thread that would only log "no-op in server mode" is still a thread that + # close has to wait out. + def test_identify_with_a_prewarm_starts_no_thread_in_server_mode + # close flushes the buffer this identify fills. + stub_request(:post, %r{https://c\.schematichq\.com/.*}).to_return(status: 200, body: "{}") + client = build_client(credit_leases: { mode: :server, default_reservation_ttl: 60_000 }) + + client.identify({ keys: { "user_id" => "u_1" }, company: { keys: { "id" => "co_1" } } }, + prewarm: ["ct_1"]) + + assert_empty client.instance_variable_get(:@pending_prewarms) + ensure + client&.close + end + + # Documented as never raising, and a caller reading ids out of config can hand + # over nil without meaning to. + def test_prewarm_tolerates_a_nil_credit_type_id_list + client = build_client(credit_leases: { mode: :client }) + + assert_nil client.prewarm(nil, company: { "id" => "co_1" }) + assert_nil client.prewarm([], company: { "id" => "co_1" }) + ensure + client&.close + end + + # A knob that cannot mean anything is a configuration bug, and the place to + # say so is where the stack still points at the caller. + def test_unusable_numeric_knobs_are_rejected_at_construction + { + { sweep_interval_ms: 0 } => "sweep_interval_ms", + { default_reservation_ttl: -1 } => "default_reservation_ttl", + { default_lease_duration: "5m" } => "default_lease_duration", + { default_lease_size: 0 } => "default_lease_size", + { low_water_mark: 1 } => "low_water_mark", + { low_water_mark: 0 } => "low_water_mark", + { low_water_mark: Float::NAN } => "low_water_mark", + { low_water_mark: Float::INFINITY } => "low_water_mark", + { sweep_interval_ms: Float::INFINITY } => "sweep_interval_ms", + { default_reservation_ttl: Float::NAN } => "default_reservation_ttl", + { overrides: { "ct_1" => { default_lease_size: -5 } } } => "default_lease_size" + }.each do |knob, name| + error = assert_raises(ArgumentError) { build_client(credit_leases: { mode: :client }.merge(knob)) } + + assert_includes error.message, name + end + end + + # The raise escapes the constructor, so the caller never gets a client to + # close. Anything started before the check would run for the life of the + # process with nothing holding a reference to stop it. + def test_a_rejected_knob_leaves_no_thread_and_no_socket_behind + before = Thread.list.select(&:alive?) + sockets = [] + Schematic::DataStream::Client.define_singleton_method(:new) do |**| + sockets << :opened + raise "the socket must not be opened for a config that cannot be used" + end + + assert_raises(ArgumentError) do + build_client(credit_leases: { mode: :client, default_lease_size: -1 }, use_data_stream: true) + end + + assert_empty sockets + assert_empty(Thread.list.select(&:alive?) - before) + ensure + Schematic::DataStream::Client.singleton_class.send(:remove_method, :new) + end + + # The lease paths guard usage; the plain one has to as well, because Ruby has + # no type to stop a string reaching the preflight builder. + def test_an_unusable_usage_is_warned_about_and_ignored + logger = RecordingLogger.new + client = Schematic::SchematicClient.new(api_key: "sch_test", base_url: "https://api.schematichq.test", + logger: logger) + plain = stub_request(:post, "https://api.schematichq.test/flags/inference/check") + .to_return(status: 200, body: JSON.generate({ "data" => { "flag" => "inference", "value" => true, + "reason" => "ok" } }), + headers: { "Content-Type" => "application/json" }) + + ["3", -1, Float::NAN, Float::INFINITY].each do |usage| + result = client.check("inference", company: { "id" => "co_1" }, usage: usage) + + assert_predicate result, :allowed? + end + + assert_equal(4, logger.warnings.count { |w| w.include?("invalid usage") }) + # Each is an ordinary check: no preflight went out, so the first one's + # verdict was cacheable and served the other three. + assert_requested plain, times: 1 + assert_requested(:post, "https://api.schematichq.test/flags/inference/check") do |req| + !JSON.parse(req.body).key?("preflight") + end + ensure + client&.close + end + + # identify never raises into its caller, and a caller can hand it any shape. + def test_identify_with_a_prewarm_tolerates_a_malformed_company + stub_request(:post, %r{https://c\.schematichq\.com/.*}).to_return(status: 200, body: "{}") + client = build_client(credit_leases: { mode: :client }) + + [{ keys: { "user_id" => "u_1" }, company: "acme" }, + { keys: { "user_id" => "u_1" } }, + { company: { keys: nil } }].each do |body| + assert_nil client.identify(body, prewarm: ["ct_1"]) + end + ensure + client&.close + end + + # A DataStream request over a closed socket is dropped without an error, so + # waiting on it burns the 30 second resource timeout on a check that was + # always going to fall back. + def test_a_lease_check_against_a_disconnected_datastream_falls_back_at_once + client = build_client(credit_leases: { mode: :client }, flag_defaults: { "inference" => false }) + datastream = Schematic::DataStream::Client.new( + api_key: "sch_test", base_url: "wss://api.schematichq.test", logger: silent_logger, rules_engine: nil + ) + # Never started, so the socket is down. The flag is cached; the company is + # not, which is the pair that used to wait. + refute_predicate datastream, :connected? + datastream.instance_variable_get(:@flag_cache).set( + datastream.send(:flag_cache_key, "inference"), { id: "flag_1", key: "inference" } + ) + client.instance_variable_set(:@datastream_client, datastream) + stub_request(:post, "https://c.schematichq.com/batch").to_return(status: 200, body: "") + flag_check = stub_request(:post, "https://api.schematichq.test/flags/inference/check") + .to_return(status: 200, + body: JSON.generate({ "data" => { "flag" => "inference", "value" => true, "reason" => "plan" } }), + headers: { "Content-Type" => "application/json" }) + + started = Process.clock_gettime(Process::CLOCK_MONOTONIC) + result = client.check("inference", company: { "org_id" => "acme" }, usage: 10, default_value: false) + elapsed_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started) * 1000 + + # Resolved as the plain fallback, not the 30 second resource timeout, and + # nothing was held, so there is nothing to refund. + assert_requested flag_check + assert_predicate result, :allowed? + assert_nil result.reservation + assert_operator elapsed_ms, :<, 1000 + ensure + client&.close + end + + # A 200 whose value is nil is the API declining to answer, not a false. + def test_a_nil_value_from_the_api_falls_back_to_the_default + client = Schematic::SchematicClient.new(api_key: "sch_test", base_url: "https://api.schematichq.test", + logger: silent_logger, flag_defaults: { "inference" => false }) + stub_request(:post, "https://api.schematichq.test/flags/inference/check") + .to_return(status: 200, + body: JSON.generate({ "data" => { "flag" => "inference", "value" => nil, "reason" => "ok" } }), + headers: { "Content-Type" => "application/json" }) + stub_request(:post, "https://c.schematichq.com/batch").to_return(status: 200, body: "") + + result = client.check("inference", company: { "id" => "co_1" }, usage: 10, default_value: true) + + assert_predicate result, :allowed? + assert_equal "flag default", result.reason + ensure + client&.close + end + + # Same on the DataStream path, where the registered flag default stands in. + def test_a_nil_value_from_the_datastream_falls_back_to_the_flag_default + client = Schematic::SchematicClient.new(api_key: "sch_test", base_url: "https://api.schematichq.test", + logger: silent_logger, flag_defaults: { "inference" => true }) + datastream = Object.new + def datastream.connected? = true + def datastream.check_flag(_eval_ctx, flag_key) = { value: nil, reason: "no rules", flag_key: flag_key } + def datastream.close = nil + client.instance_variable_set(:@datastream_client, datastream) + stub_request(:post, "https://c.schematichq.com/batch").to_return(status: 200, body: "") + + response = client.check_flag_with_entitlement("inference", company: { "id" => "co_1" }) + + assert response.value + ensure + client&.close + end + + def test_prewarm_is_a_no_op_without_credit_leases + client = Schematic::SchematicClient.new(api_key: "sch_test", base_url: "https://api.schematichq.test", + logger: silent_logger) + + assert_nil client.prewarm(["ct_1"], company: { "id" => "co_1" }) + ensure + client&.close + end + + # Without credit_leases, check is a plain flag check with no gating. + def test_check_without_credit_leases_falls_through_to_a_plain_check + client = Schematic::SchematicClient.new(api_key: "sch_test", base_url: "https://api.schematichq.test", + logger: silent_logger, flag_defaults: { "inference" => true }) + stub_request(:post, "https://api.schematichq.test/flags/inference/check") + .to_return(status: 200, body: JSON.generate({ "data" => { "flag" => "inference", "value" => true, + "reason" => "ok" } }), + headers: { "Content-Type" => "application/json" }) + + result = client.check("inference", company: { "id" => "co_1" }, usage: 10) + + assert_predicate result, :allowed? + assert_nil result.reservation + ensure + client&.close + end +end diff --git a/test/custom.test.rb b/test/custom.test.rb index f330b72..1998452 100644 --- a/test/custom.test.rb +++ b/test/custom.test.rb @@ -2065,6 +2065,32 @@ def capture_event(&block) client.close WebMock.reset! end + + # The engine declining to answer is the one case default_value exists for, so + # the DataStream branch has to resolve it the way the offline and API branches + # do rather than reaching straight for the registered default. + it "honours the caller's default_value when DataStream returns no value" do + stub_request(:post, CAPTURE_URL).to_return(status: 200) + + client = Schematic::SchematicClient.new(api_key: "api_test_key_123") + client.set_flag_default("ds-no-value", false) + + mock_ds = Object.new + mock_ds.define_singleton_method(:connected?) { true } + mock_ds.define_singleton_method(:close) { nil } + mock_ds.define_singleton_method(:check_flag) do |_eval_ctx, flag_key| + { value: nil, flag_key: flag_key, reason: "no verdict" } + end + client.instance_variable_set(:@datastream_client, mock_ds) + + assert client.check_flag_with_entitlement("ds-no-value", default_value: true).value + assert client.check_flag_with_entitlement("ds-no-value", default_value: -> { true }).value + # With no caller default the registered one still stands in. + refute client.check_flag_with_entitlement("ds-no-value").value + + client.close + WebMock.reset! + end end # ============================================================================= diff --git a/test/lease_support.rb b/test/lease_support.rb new file mode 100644 index 0000000..4a64fcb --- /dev/null +++ b/test/lease_support.rb @@ -0,0 +1,516 @@ +# frozen_string_literal: true + +require "digest" +require "json" + +# Shared test doubles for the credit lease suite: a virtual clock, an +# in-process Redis stand-in, and a scripted wire client. +module LeaseSupport + # The fixed virtual instant every conformance vector and lease test starts + # from. Vectors express every *_at_ms field as an offset from it. + T0 = Time.utc(2026, 1, 1) + + # Only moves when a test advances it, so nothing here depends on wall time. + class VirtualClock + def initialize(start = T0) + @now = start + @mutex = Mutex.new + end + + def now + @mutex.synchronize { @now } + end + + def to_proc + -> { now } + end + + def advance_ms(millis) + @mutex.synchronize { @now += millis.to_f / 1000.0 } + end + + # An absolute position on the virtual timeline. + def at_ms(offset_ms) + T0 + (offset_ms.to_f / 1000.0) + end + end + + # An in-process stand-in for Redis, covering the command subset the lease and + # reservation stores use. + # + # No Lua runs here, as in the other SDKs' fakes: the registry is keyed by the + # exact script source the stores ship, and each entry runs a Ruby + # transliteration of it. Keying on the source is what makes a Lua change that + # is not mirrored here fail loudly with "unknown script" instead of silently + # diverging. The shipped Lua itself is what reaches a real Redis, and + # test/credits_test.rb pins its SHA1 against the reference implementation's. + class FakeRedis + class NoScriptError < StandardError + end + + # One entry per EXEC, naming the commands that transaction queued, so a + # test can tell one round trip from two. + attr_reader :transactions + + def initialize(clock) + @clock = clock + @transactions = [] + @hashes = {} + @zsets = {} + @expiries = {} + @scripts = {} + @loaded = {} + register_scripts + end + + # MULTI/EXEC: queue the commands, apply them on exec, and record what the + # transaction carried. + def multi + tx = Transaction.new(self) + yield tx + @transactions << tx.commands.map(&:first) + tx.apply + end + + # The queued half of a MULTI: it holds the commands the block issues and + # replays them against the client on exec. + class Transaction + attr_reader :commands + + def initialize(client) + @client = client + @commands = [] + end + + def hset(key, *pairs) + @commands << [:hset, [key, *pairs]] + self + end + + def pexpireat(key, millis) + @commands << [:pexpireat, [key, millis]] + self + end + + def apply + @commands.map { |name, args| @client.public_send(name, *args) } + end + end + + # --- string/hash commands ------------------------------------------------- + + def hgetall(key) + evict_expired + (@hashes[key] || {}).dup + end + + def hget(key, field) + evict_expired + (@hashes[key] || {})[field] + end + + def hset(key, *pairs) + evict_expired + pairs = pairs.first.to_a.flatten(1) if pairs.size == 1 && pairs.first.is_a?(Hash) + hash = (@hashes[key] ||= {}) + pairs.each_slice(2) { |field, value| hash[field.to_s] = value.to_s } + pairs.size / 2 + end + + def hdel(key, field) + evict_expired + (@hashes[key] || {}).delete(field.to_s) ? 1 : 0 + end + + def del(*keys) + evict_expired + keys.count do |key| + @expiries.delete(key) + existed = @hashes.key?(key) || @zsets.key?(key) + @hashes.delete(key) + @zsets.delete(key) + existed + end + end + + def exists?(key) + evict_expired + @hashes.key?(key) || @zsets.key?(key) + end + + def pexpireat(key, millis) + @expiries[key] = millis.to_i + evict_expired + 1 + end + + # --- sorted set commands -------------------------------------------------- + + def zadd(key, score, member) + evict_expired + zset = (@zsets[key] ||= {}) + added = zset.key?(member) ? 0 : 1 + zset[member] = score.to_f + added + end + + def zrem(key, member) + evict_expired + (@zsets[key] || {}).delete(member) ? 1 : 0 + end + + def zrangebyscore(key, min, max, limit: nil) + evict_expired + members = (@zsets[key] || {}).select { |_, score| score.between?(min.to_f, max.to_f) } + .sort_by { |member, score| [score, member] } + .map(&:first) + return members if limit.nil? + + offset, count = limit + members.drop(offset.to_i).first(count.to_i) + end + + def zcard(key) + evict_expired + (@zsets[key] || {}).size + end + + # --- scripting ------------------------------------------------------------ + + # This is the Redis EVAL command, not Kernel#eval: it looks the script up in + # a fixed registry of the store's own Lua and runs the matching Ruby + # transliteration. Nothing is ever compiled from the argument. + def eval(script, keys:, argv:) + body = @scripts[script] + raise "unknown script: #{Digest::SHA1.hexdigest(script)}" if body.nil? + + @loaded[Digest::SHA1.hexdigest(script)] = script + evict_expired + body.call(keys, argv.map(&:to_s)) + end + + def evalsha(sha, keys:, argv:) + script = @loaded[sha] + raise NoScriptError, "NOSCRIPT No matching script" if script.nil? + + eval(script, keys: keys, argv: argv) # rubocop:disable Security/Eval + end + + # What the Lua reads from redis.call('TIME'), in integer milliseconds. + def now_ms + (@clock.now.to_f * 1000).round + end + + private + + def evict_expired + now = now_ms + due = @expiries.select { |_, at| at <= now }.keys + due.each do |key| + @expiries.delete(key) + @hashes.delete(key) + @zsets.delete(key) + end + end + + def register_scripts + lease = Schematic::Credits::Leases::RedisLeaseStore + reservations = Schematic::Credits::Leases::RedisReservationStore + @scripts[lease::REPLACE_SCRIPT.source] = method(:run_replace) + @scripts[lease::TRY_RESERVE_SCRIPT.source] = method(:run_try_reserve) + @scripts[lease::REFUND_SCRIPT.source] = method(:run_refund) + @scripts[lease::EXTEND_SCRIPT.source] = method(:run_extend) + @scripts[reservations::CLAIM_SCRIPT.source] = method(:run_claim) + end + + def run_replace(keys, argv) + key = keys.first + existing_id = hget(key, "leaseId") + existing_expiry = (hget(key, "expiresAt") || "0").to_f + new_id, new_granted, new_expiry_raw, grace_raw, company_id, credit_type_id = argv + new_expiry = new_expiry_raw.to_f + grace = grace_raw.to_f + + return 0 if existing_id && existing_expiry > now_ms + + if existing_id == new_id + granted = (hget(key, "grantedAmount") || "0").to_f + add = new_granted.to_f - granted + if add.positive? + remaining = (hget(key, "localRemainingCredits") || "0").to_f + hset(key, "grantedAmount", new_granted, "localRemainingCredits", (remaining + add).to_s) + end + if new_expiry > existing_expiry + hset(key, "expiresAt", new_expiry_raw) + pexpireat(key, new_expiry + grace) + end + return 0 + end + + del(key) + hset(key, "leaseId", new_id, "companyId", company_id, "creditTypeId", credit_type_id, + "grantedAmount", new_granted, "localRemainingCredits", new_granted, "expiresAt", new_expiry_raw) + pexpireat(key, new_expiry + grace) + 1 + end + + def run_try_reserve(keys, argv) + key = keys.first + raw = hget(key, "localRemainingCredits") + return nil if raw.nil? + + lease_id = hget(key, "leaseId") + return nil if lease_id.nil? + return nil if (hget(key, "expiresAt") || "0").to_f <= now_ms + + remaining = raw.to_f + requested = argv[0].to_f + return nil if remaining < requested + + new_remaining = remaining - requested + hset(key, "localRemainingCredits", new_remaining.to_s) + [new_remaining.to_s, lease_id] + end + + def run_refund(keys, argv) + key = keys.first + raw_remaining = hget(key, "localRemainingCredits") + return 0 if raw_remaining.nil? + + required_lease = argv[1] + return 0 if required_lease && required_lease != "" && hget(key, "leaseId") != required_lease + + granted = (hget(key, "grantedAmount") || "0").to_f + new_balance = raw_remaining.to_f + argv[0].to_f + new_balance = granted if new_balance > granted + hset(key, "localRemainingCredits", new_balance.to_s) + 1 + end + + def run_extend(keys, argv) + key = keys.first + raw_granted = hget(key, "grantedAmount") + return 0 if raw_granted.nil? + + required_lease = argv[3] + return 0 if required_lease && required_lease != "" && hget(key, "leaseId") != required_lease + + target = argv[0].to_f + add = target - raw_granted.to_f + if add.positive? + remaining = (hget(key, "localRemainingCredits") || "0").to_f + hset(key, "grantedAmount", target.to_s, "localRemainingCredits", (remaining + add).to_s) + end + new_expiry = argv[1].to_f + grace = argv[2].to_f + if new_expiry > (hget(key, "expiresAt") || "0").to_f + hset(key, "expiresAt", argv[1]) + pexpireat(key, new_expiry + grace) + end + 1 + end + + def run_claim(keys, _argv) + key = keys.first + raw = hgetall(key) + return nil if raw.empty? + + del(key) + raw.to_a.flatten(1) + end + end + + # Stands in for the lease API: queued responses in, recorded calls out. + class ScriptedWireClient + AcquireCall = Struct.new(:company_id, :credit_type_id, :requested_amount, :expires_at) + ExtendCall = Struct.new(:lease_id, :additional_amount, :expires_at, :idempotency_key) + + attr_reader :acquire_calls, :extend_calls, :release_calls + # Run while a call is in flight, for emulating a sibling process winning the + # race or a close landing mid-call. + attr_accessor :during_acquire, :during_extend + + def initialize(clock) + @clock = clock + @mutex = Mutex.new + @acquire_responses = [] + @extend_responses = [] + @acquire_calls = [] + @extend_calls = [] + @release_calls = [] + end + + # Scripted responses nobody asked for, so a test can tell a run that took + # the path it described from one that merely passed its own assertions. + def pending_scripts + @mutex.synchronize { @acquire_responses + @extend_responses } + end + + def queue_acquire(script) + @mutex.synchronize { @acquire_responses << script } + end + + def queue_extend(script) + @mutex.synchronize { @extend_responses << script } + end + + def acquire(company_id:, credit_type_id:, requested_amount:, expires_at:, **) + during = nil + script = @mutex.synchronize do + @acquire_calls << AcquireCall.new(company_id, credit_type_id, requested_amount, expires_at) + during = @during_acquire + @during_acquire = nil + @acquire_responses.shift + end + during&.call + lease = scripted_lease(script, "unscripted acquire wire call") + Schematic::Credits::Leases::LeaseGrant.new( + lease_id: lease["lease_id"] || "lse_unnamed", + company_id: company_id, + credit_type_id: credit_type_id, + granted_amount: lease["granted_amount"].to_f, + expires_at: @clock.at_ms(lease["expires_at_ms"]) + ) + end + + def extend(lease_id:, additional_amount:, expires_at:, idempotency_key: nil, **) + during = nil + script = @mutex.synchronize do + @extend_calls << ExtendCall.new(lease_id, additional_amount, expires_at, idempotency_key) + during = @during_extend + @during_extend = nil + @extend_responses.shift + end + during&.call + lease = scripted_lease(script, "unscripted extend wire call") + granted = lease["granted_total"] || lease["granted_amount"] + Schematic::Credits::Leases::LeaseGrant.new( + lease_id: lease_id, + company_id: "co_wire", + credit_type_id: "ct_wire", + granted_amount: granted.to_f, + expires_at: @clock.at_ms(lease["expires_at_ms"]) + ) + end + + def release(lease_id:, **) + @mutex.synchronize { @release_calls << lease_id } + nil + end + + private + + def scripted_lease(script, missing) + raise missing if script.nil? + raise(script["error"] || missing) if script["error"] || script["lease"].nil? + + script["lease"] + end + end + + # Serves one flag and one company from "cache" and answers each evaluation + # with the next scripted result, in call order. + # + # The engine is an oracle here, as conformance/SPEC.md says: the vectors pin + # the orchestration around the rules engine, not the engine itself, which is + # shared WASM across the SDKs and has its own tests. + class ScriptedDataStream + EngineCall = Struct.new(:credit_balances, :options) + + attr_reader :calls + + def initialize(flag_key:, company:, results:, user: nil) + @flag_key = flag_key + @company = company + @user = user + @results = results.dup + @calls = [] + end + + # Scripted results nobody asked for. An extra call raises instead, so the + # pair pins the call count in both directions. + def pending_results + @results.dup + end + + def get_flag(_key) + { id: "flag_1", key: @flag_key } + end + + def get_company(_keys) + @company + end + + def get_user(_keys) + @user + end + + def check_flag_with_options(_flag, company, _user, options) + @calls << EngineCall.new((company || {})[:credit_balances] || {}, options) + raise "unscripted engine call for flag #{@flag_key}" if @results.empty? + + scripted = @results.shift + { + value: scripted["value"], + reason: scripted["reason"], + flag_key: @flag_key, + flag_id: "flag_1", + entitlement: entitlement_from(scripted["entitlement"]) + } + end + + private + + def entitlement_from(spec) + return nil if spec.nil? + + { + featureId: spec["feature_id"] || "feat_1", + featureKey: spec["feature_key"] || @flag_key, + valueType: spec["value_type"], + creditId: spec["credit_id"], + consumptionRate: spec["consumption_rate"], + eventSubtype: spec["event_subtype"], + metricResetAt: spec["metric_reset_at"] + }.compact + end + end + + # A lease-store refunder that fails once while armed, reproducing a process + # death between a reservation's claim and its refund. The reservation store + # refunds through whatever it is handed, so wrapping the store leaves the rest + # of the flow reading the real one. + class CrashingRefund + class SimulatedCrash < StandardError + end + + def initialize(target) + @target = target + @armed = false + end + + def arm + @armed = true + end + + def refund(company_id, credit_type_id, credits, pin_lease_id = nil) + if @armed + @armed = false + raise SimulatedCrash, "simulated crash before refund" + end + @target.refund(company_id, credit_type_id, credits, pin_lease_id) + end + + # Everything else the reservation store might reach for goes straight + # through, so the wrapper is only a seam for the refund. + def method_missing(name, ...) + return super unless @target.respond_to?(name) + + @target.public_send(name, ...) + end + + def respond_to_missing?(name, include_private = false) + @target.respond_to?(name, include_private) || super + end + end +end