diff --git a/.fernignore b/.fernignore index 670bab1..6256ef6 100644 --- a/.fernignore +++ b/.fernignore @@ -16,6 +16,8 @@ src/main/java/com/schematic/api/HttpEventSender.java src/main/java/com/schematic/api/IdentifyOptions.java src/main/java/com/schematic/api/Schematic.java src/main/java/com/schematic/api/TrackOptions.java +src/main/java/com/schematic/api/credits/ +conformance/ src/main/java/com/schematic/api/cache/CacheProvider.java src/main/java/com/schematic/api/cache/CachedItem.java src/main/java/com/schematic/api/cache/LocalCache.java @@ -35,6 +37,7 @@ src/test/java/com/schematic/api/TestOfflineMode.java src/test/java/com/schematic/api/TestReadme.java src/test/java/com/schematic/api/TestSchematic.java src/test/java/com/schematic/api/cache/RedisCacheProviderTest.java +src/test/java/com/schematic/api/credits/ src/test/java/com/schematic/api/datastream/ src/test/java/com/schematic/webhook/ .fern/replay.lock diff --git a/README.md b/README.md index cc4b567..df431f9 100644 --- a/README.md +++ b/README.md @@ -165,6 +165,153 @@ user.put("user_id", "your-user-id"); boolean flagValue = schematic.checkFlag("some-flag-key", company, user); ``` +`checkFlagWithEntitlement` answers the same question and hands back the whole result: the value, the reason the rules engine gave, and the matched entitlement. + +## Credit Leases and Reservations + +For features metered by credit burndown (inference tokens, for example), `check` reserves credits for the work about to run and `trackWithReservation` settles the reservation with the 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. + +`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. + +### Setup + +```java +import com.schematic.api.Schematic; +import com.schematic.api.credits.CreditLeaseConfig; +import com.schematic.api.datastream.DatastreamOptions; +import java.time.Duration; +import redis.clients.jedis.JedisPooled; + +JedisPooled redisClient = new JedisPooled("localhost", 6379); + +Schematic schematic = Schematic.builder() + .apiKey("YOUR_API_KEY") + .datastreamOptions(DatastreamOptions.builder().build()) + .creditLeases(CreditLeaseConfig.builder() + .defaultLeaseSize(10000) // credits requested per lease + .defaultLeaseDuration(Duration.ofMinutes(5)) // lease lifetime + .defaultReservationTtl(Duration.ofSeconds(60)) // how long a reservation is held if no track settles it + .redisClient(redisClient) // lease and reservation state + .build()) + .build(); +``` + +Leases reuse the Redis client the DataStream cache is configured with, if there is one. The example above configures DataStream without a Redis cache, so it passes `redisClient` explicitly. Set it whenever the DataStream cache is local, or when lease state should live in a different Redis from the cache. With no Redis on either side the SDK falls back to per-process in-memory state, which gates one process only and warns at startup. + +Server mode needs only a TTL: + +```java +import com.schematic.api.Schematic; +import com.schematic.api.credits.CreditLeaseConfig; +import java.time.Duration; + +Schematic schematic = Schematic.builder() + .apiKey("YOUR_API_KEY") + .creditLeases(CreditLeaseConfig.builder() + .defaultReservationTtl(Duration.ofSeconds(60)) // just under an hour at most, which is as far out as the API will reserve credits + .build()) + .build(); +``` + +Only `mode` and `defaultReservationTtl` apply in server mode; the client warns at startup if a client-only option is set. + +### Checking and tracking + +```java +import com.schematic.api.credits.CheckOptions; +import com.schematic.api.credits.CheckResult; +import java.util.HashMap; +import java.util.Map; + +Map company = new HashMap<>(); +company.put("id", "your-company-id"); + +// Reserve up to maxTokens for this operation. +CheckResult result = schematic.check("inference", company, null, CheckOptions.builder() + .usage(maxTokens) // upper bound for this operation + .eventSubtype("inference_tokens") // the metered event + .build()); +if (!result.isAllowed()) { + throw new IllegalStateException("credit balance exceeded"); +} + +long tokensUsed = runInference(); + +// Report the actual usage; the unused slice of the reservation is refunded. +if (result.getReservation() != null) { + schematic.trackWithReservation(result.getReservation(), tokensUsed); +} else { + schematic.track("inference_tokens", company, null, null, tokensUsed); +} +``` + +A check can allow without reserving credits, when the feature is not credit-metered, when `usage` is 0, or when the check failed open, and that usage still has to be tracked. + +`usage` may be fractional, but credits are always sized in whole event units: a client-mode reservation records the fractional quantity, while the credits it reserves and the debit its settle makes are both `ceil(usage) x consumption rate`, so the local ledger moves by exactly what the track event bills. The integer fields on the wire round up for the same reason: the preflight quantity and the quantity a track event bills, so a partial unit is never billed as none. + +`usage` still gates a check that reserves nothing: it is 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. + +`CheckOptions.timeout` bounds every call a check waits on: the check-and-reserve call in server mode, the REST flag check a check can fall back to, and the client-mode lease acquire and extend. Lease calls are shared between concurrent checks, and a check that joins one somebody else opened waits no longer than its own timeout before giving up and taking its failure path, leaving that call running for the checks still on it. Background top-ups keep the client's own timeout. + +An unsettled reservation expires after `defaultReservationTtl` and its credits return to the lease. A late settle still bills the usage, since the track event carries a deterministic idempotency key that keeps it from double-billing, but it does not re-debit the local lease. Set `defaultReservationTtl` above the longest expected gap between the check and the settle. + +### Pre-warming + +Warm leases when the user is identified, so a session's first check does not wait on a lease acquire: + +```java +import com.schematic.api.IdentifyOptions; +import com.schematic.api.types.EventBodyIdentifyCompany; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +Map userKeys = new HashMap<>(); +userKeys.put("user_id", "your-user-id"); + +Map companyKeys = new HashMap<>(); +companyKeys.put("id", "your-company-id"); + +schematic.identify( + userKeys, + EventBodyIdentifyCompany.builder().keys(companyKeys).build(), + "Your User", + null, + IdentifyOptions.builder() + .prewarm(Collections.singletonList("credit-type-id")) + .build()); +``` + +Identifying with a prewarm flushes the event buffer first, so the server has the company before the warm-up asks for a lease against it. That makes it a session-start call, not one to put on every event. + +Or call `schematic.prewarm(companyKeys, creditTypeIds)` 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 + +In server mode, a check that times out after the server has already reserved leaves those credits reserved until the TTL expires, so keep `defaultReservationTtl` short there. + +A check that cannot gate, because the API is unreachable, Redis is down, or the lease is exhausted, fails closed by default. Override it per check: + +```java +import com.schematic.api.credits.CheckOptions; +import com.schematic.api.credits.OnAcquireFailure; + +CheckOptions options = CheckOptions.builder() + .usage(maxTokens) + .eventSubtype("inference_tokens") + .onAcquireFailure(OnAcquireFailure.FAIL_OPEN) + .build(); +``` + +In client mode `FAIL_OPEN` still evaluates the flag's rules with the credit balance assumed sufficient, so plan targeting and every non-credit condition apply and only the credit gate is bypassed. In server mode it returns the flag's default value, which is false unless the check passes `defaultValue` or the client configures a flag default. + +See [Credit Lease Options](#credit-lease-options) for the full set of options. + ## Webhook Verification Schematic can send webhooks to notify your application of events. To ensure the security of these webhooks, Schematic signs each request using HMAC-SHA256. The Java SDK provides utility functions to verify these signatures. @@ -277,6 +424,23 @@ Schematic schematic = Schematic.builder() .build(); ``` +### Credit Lease Options + +Set with `creditLeases(CreditLeaseConfig.builder()...build())`. Per-credit-type overrides take a `CreditLeaseOverride` under `override(creditTypeId, ...)`. + +| Option | Type | Default | Description | +|---|---|---|---| +| `mode` | `CreditLeaseMode` | `AUTO` | Where credits are reserved; `AUTO` picks client when DataStream is enabled, server otherwise | +| `defaultReservationTtl` | `Duration` | 60 seconds | How long an unsettled reservation is held | +| `defaultLeaseDuration` | `Duration` | 5 minutes | (client mode) Lease lifetime | +| `defaultLeaseSize` | `double` | 10000 | (client mode) Credits requested per lease acquire or extend | +| `lowWaterMark` | `double` | 0.25 | (client mode) Extend in the background when the lease balance dips below this fraction | +| `sweepInterval` | `Duration` | 1 second | (client mode) How often expired reservations are swept | +| `prewarmResolveTimeout` | `Duration` | 5 seconds | (client mode) How long `prewarm` waits for a freshly identified company to surface; zero resolves from the DataStream cache only | +| `redisClient` | `JedisPooled` | the DataStream cache's client | (client mode) Redis client for lease and reservation state | +| `redisKeyPrefix` | `String` | the DataStream cache's prefix | (client mode) Key prefix for lease and reservation keys | +| `overrides` | `Map` | none | (client mode) Per-credit-type overrides of the above, keyed by credit type id | + ### Offline Mode In development or testing environments, you may want to avoid making network requests when checking flags or submitting events. You can run Schematic in offline mode: diff --git a/build.gradle b/build.gradle index bb349a4..62227aa 100644 --- a/build.gradle +++ b/build.gradle @@ -36,6 +36,7 @@ dependencies { implementation 'com.dylibso.chicory:runtime:1.4.0' implementation 'com.dylibso.chicory:wasi:1.4.0' implementation 'redis.clients:jedis:5.2.0' + testImplementation 'com.github.codemonstur:embedded-redis:1.4.3' } diff --git a/conformance/README.md b/conformance/README.md new file mode 100644 index 0000000..3839580 --- /dev/null +++ b/conformance/README.md @@ -0,0 +1,10 @@ +# Credit lease conformance suite + +`SPEC.md` and `vectors/*.json` are copied verbatim from `conformance/` in +[schematic-node](https://github.com/SchematicHQ/schematic-node), the reference +implementation. Do not edit them here: change them there, then copy the new +versions across, so every SDK runs the same contract. + +The runner is the only language-specific piece. This SDK's lives in +`src/test/java/com/schematic/api/credits/conformance/`, and runs every vector +against both store backends: the in-memory stores and the Redis stores. 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/src/main/java/com/schematic/api/IdentifyOptions.java b/src/main/java/com/schematic/api/IdentifyOptions.java index 176654c..4c7d0d4 100644 --- a/src/main/java/com/schematic/api/IdentifyOptions.java +++ b/src/main/java/com/schematic/api/IdentifyOptions.java @@ -1,5 +1,9 @@ package com.schematic.api; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + /** * Optional metadata for an {@link Schematic#identify} event. * @@ -8,9 +12,14 @@ public final class IdentifyOptions { private final String idempotencyKey; + private final List prewarm; private IdentifyOptions(Builder builder) { this.idempotencyKey = builder.idempotencyKey; + // Copied, because the prewarm runs in the background and reads this list after identify + // has returned: a caller who reuses and mutates their own list would otherwise decide, + // after the fact, which credit types get warmed. + this.prewarm = builder.prewarm != null ? new ArrayList<>(builder.prewarm) : null; } public static Builder builder() { @@ -25,14 +34,32 @@ public String getIdempotencyKey() { return idempotencyKey; } + /** + * Credit type ids to warm a lease for once the identify is enqueued, so the first + * credit-gated check does not pay the acquire round trip. A no-op unless credit leases are + * configured on the client. + */ + public List getPrewarm() { + // Unmodifiable for the same reason the constructor copies: the prewarm this list drives + // runs after identify returns, and a caller editing it in between would move the work. + return prewarm == null ? null : Collections.unmodifiableList(prewarm); + } + public static final class Builder { private String idempotencyKey; + private List prewarm; public Builder idempotencyKey(String idempotencyKey) { this.idempotencyKey = idempotencyKey; return this; } + /** Credit type ids to warm a lease for after this identify. */ + public Builder prewarm(List prewarm) { + this.prewarm = prewarm; + return this; + } + public IdentifyOptions build() { return new IdentifyOptions(this); } diff --git a/src/main/java/com/schematic/api/Schematic.java b/src/main/java/com/schematic/api/Schematic.java index 2c426d5..79948da 100644 --- a/src/main/java/com/schematic/api/Schematic.java +++ b/src/main/java/com/schematic/api/Schematic.java @@ -7,6 +7,31 @@ import com.schematic.api.core.Environment; import com.schematic.api.core.NoOpHttpClient; import com.schematic.api.core.ObjectMappers; +import com.schematic.api.core.RequestOptions; +import com.schematic.api.credits.ApiLeaseWireClient; +import com.schematic.api.credits.CheckOptions; +import com.schematic.api.credits.CheckRequest; +import com.schematic.api.credits.CheckResult; +import com.schematic.api.credits.CreditAmounts; +import com.schematic.api.credits.CreditCheck; +import com.schematic.api.credits.CreditLeaseConfig; +import com.schematic.api.credits.CreditLeaseDefaults; +import com.schematic.api.credits.CreditLeaseManager; +import com.schematic.api.credits.CreditLeaseMode; +import com.schematic.api.credits.DataStreamCreditCheckSource; +import com.schematic.api.credits.InMemoryLeaseStore; +import com.schematic.api.credits.InMemoryReservationStore; +import com.schematic.api.credits.LeaseStore; +import com.schematic.api.credits.OnAcquireFailure; +import com.schematic.api.credits.PreflightOptions; +import com.schematic.api.credits.PrewarmCompanyResolver; +import com.schematic.api.credits.RedisLeaseStore; +import com.schematic.api.credits.RedisReservationStore; +import com.schematic.api.credits.Reservation; +import com.schematic.api.credits.ReservationSettlement; +import com.schematic.api.credits.ReservationStore; +import com.schematic.api.credits.ServerCreditCheck; +import com.schematic.api.datastream.CheckFlagOptions; import com.schematic.api.datastream.DataStreamClient; import com.schematic.api.datastream.DataStreamException; import com.schematic.api.datastream.DatastreamOptions; @@ -25,7 +50,9 @@ import com.schematic.api.types.EventBodyIdentifyCompany; import com.schematic.api.types.EventBodyTrack; import com.schematic.api.types.EventType; +import com.schematic.api.types.PreflightRequestBody; import com.schematic.api.types.RulesengineCheckFlagResult; +import java.time.Clock; import java.time.Duration; import java.time.OffsetDateTime; import java.util.ArrayList; @@ -33,9 +60,21 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; +import redis.clients.jedis.JedisPooled; public final class Schematic extends BaseSchematic implements AutoCloseable { + + // Namespaces the settle key so a reservation id can never collide with a caller's own + // idempotency key, and so a recovery emit and an accidental second settle collapse to one + // billed event across pods and restarts. + private static final String RESERVATION_TRACK_IDEMPOTENCY_PREFIX = "lease-reservation:"; + private final Duration eventBufferInterval; private final EventBuffer eventBuffer; private final List> flagCheckCacheProviders; @@ -45,8 +84,23 @@ public final class Schematic extends BaseSchematic implements AutoCloseable { private final Thread shutdownHook; private final boolean offline; private final HttpEventSender eventSender; - private final DataStreamClient dataStreamClient; + // Not final: a DataStream that fails to start leaves none, and auto mode reads this per check. + private volatile DataStreamClient dataStreamClient; private final DatastreamOptions datastreamOptions; + // Credit leases. Null throughout when the caller did not configure them, which is what every + // credit-aware path checks before doing anything. + private final CreditLeaseMode creditLeaseMode; + private final LeaseStore leaseStore; + private final ReservationStore reservations; + private final CreditLeaseManager creditLeaseManager; + private final CreditCheck creditCheck; + private final boolean leaseBackendShared; + private final Duration serverReservationTtl; + private final Duration prewarmResolveTimeout; + // Runs the prewarms identify kicks off, so the caller's identify does not wait on a lease + // acquire. Null when leases are not configured. + private final ExecutorService prewarms; + private volatile boolean closing; private Schematic(Builder builder) { super(buildClientOptions(builder.apiKey, builder)); @@ -87,13 +141,157 @@ private Schematic(Builder builder) { rulesEngine = null; } - this.dataStreamClient = new DataStreamClient( + DataStreamClient started = new DataStreamClient( this.datastreamOptions, this.apiKey, basePath, this.logger, rulesEngine, resolveSdkVersion()); - this.dataStreamClient.start(); + try { + started.start(); + } catch (RuntimeException e) { + // Nothing here can serve a local evaluation, so the client is dropped rather than + // kept as a handle every later path has to re-test. Checks fall to the API, and + // auto mode resolves to server-side gating on the strength of this field. + this.logger.error("DataStream failed to start, falling back to API checks: " + e.getMessage()); + try { + started.close(); + } catch (Exception closing) { + this.logger.debug("DataStream close after a failed start: " + closing); + } + started = null; + } + this.dataStreamClient = started; } else { this.dataStreamClient = null; } + // Credit leases and reservations, if the caller opted in. + CreditLeaseConfig creditLeases = builder.creditLeases; + if (creditLeases != null && this.offline) { + this.logger.warn("creditLeases is configured but the client is offline, so checks return flag defaults " + + "with no credit gating"); + } + CreditLeaseMode mode = null; + Duration serverTtl = CreditLeaseDefaults.RESERVATION_TTL; + Duration prewarmTimeout = CreditLeaseDefaults.PREWARM_RESOLVE_TIMEOUT; + LeaseStore leases = null; + ReservationStore holds = null; + CreditLeaseManager manager = null; + CreditCheck check = null; + boolean sharedBackend = false; + // A DataStream alone is not enough to gate locally: without a loaded engine every + // evaluation throws and a client-mode check falls through to a plain, ungated one. + boolean localGatingReady = this.dataStreamClient != null && this.dataStreamClient.hasRulesEngine(); + if (creditLeases != null && !this.offline) { + mode = creditLeases.getMode() != null ? creditLeases.getMode() : CreditLeaseMode.AUTO; + Duration configuredTtl = creditLeases.getDefaultReservationTtl() != null + ? creditLeases.getDefaultReservationTtl() + : CreditLeaseDefaults.RESERVATION_TTL; + // 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 it there would shorten holds for no reason. + Duration maxTtl = + CreditLeaseDefaults.MAX_RESERVATION_TTL.minus(CreditLeaseDefaults.RESERVATION_TTL_SKEW_ALLOWANCE); + serverTtl = mode == CreditLeaseMode.CLIENT || configuredTtl.compareTo(maxTtl) <= 0 ? configuredTtl : maxTtl; + if (mode != CreditLeaseMode.CLIENT && configuredTtl.compareTo(maxTtl) > 0) { + this.logger.warn("creditLeases.defaultReservationTtl of " + configuredTtl.toMillis() + + "ms is longer than the API will hold credits for; server-mode holds are clamped to " + + maxTtl.toMillis() + "ms"); + } + if (creditLeases.getPrewarmResolveTimeout() != null) { + prewarmTimeout = creditLeases.getPrewarmResolveTimeout(); + } + + // Server mode holds credits over the API, so none of the local plumbing is built and + // the options that only steer it would silently do nothing. Say so once, at startup. + if (mode == CreditLeaseMode.SERVER || (mode == CreditLeaseMode.AUTO && !localGatingReady)) { + String clientOnly = clientOnlyOptions(creditLeases); + if (!clientOnly.isEmpty()) { + this.logger.warn("creditLeases resolves to server mode, so " + clientOnly + + " will be ignored: those options only apply to client mode"); + } + } + // Auto with no DataStream is the server-mode default, not a misconfiguration. + // Client without DataStream is the degraded path, where every check falls back to a + // plain flag check with the usage ignored, so it still warns. + if (mode == CreditLeaseMode.AUTO && this.dataStreamClient == null) { + this.logger.info("creditLeases is configured and DataStream is not enabled, so credit reservations " + + "run in server mode, one check-and-reserve call per check"); + } + if (mode == CreditLeaseMode.CLIENT && this.dataStreamClient == null) { + this.logger.warn("creditLeases is configured but DataStream is not enabled, so check() falls back to " + + "plain flag checks with no credit gating"); + } else if (mode == CreditLeaseMode.CLIENT && !this.dataStreamClient.hasRulesEngine()) { + // Auto resolves this away by gating server-side. An explicit client mode is the + // caller's choice to keep, so say what it costs rather than overriding it. + this.logger.warn("creditLeases is set to client mode but the rules engine did not load, so every " + + "check() falls back to a plain flag check with no credit gating; use server mode until " + + "the engine is available"); + } + } + // The same readiness auto resolves on, so the stores, the manager and its sweeper are + // built exactly when a check will gate against them. Building them for an auto client + // that resolves to server mode leaves a sweeper polling an index nothing writes to. + boolean usesLeases = mode == CreditLeaseMode.CLIENT || (mode == CreditLeaseMode.AUTO && localGatingReady); + if (creditLeases != null && !this.offline && usesLeases) { + // Lease and hold state belongs in a shared cache so gating holds across horizontally + // scaled pods. An explicit client wins; otherwise reuse the one the DataStream caches + // are already configured with, so an existing Redis setup backs leases with no second + // client to wire up. + JedisPooled redisClient = inheritFromDataStream( + creditLeases.getRedisClient(), + this.dataStreamClient == null ? null : this.dataStreamClient.getRedisClient()); + String keyPrefix = inheritFromDataStream( + creditLeases.getRedisKeyPrefix(), + this.dataStreamClient == null ? null : this.dataStreamClient.getRedisKeyPrefix()); + if (redisClient != null) { + sharedBackend = true; + leases = new RedisLeaseStore(redisClient, keyPrefix, creditLeases.getDefaultLeaseDuration(), null); + holds = new RedisReservationStore(redisClient, leases, keyPrefix, null); + } else { + // Without a shared backend each pod gates against its own leases, which defeats + // the cross-pod protection that is the point of leasing, so warn rather than + // degrade silently. + this.logger.warn("creditLeases is enabled without a shared Redis backend, so lease and reservation " + + "state is per-process; configure a Redis client so leases gate across SDK instances"); + leases = new InMemoryLeaseStore(null); + holds = new InMemoryReservationStore(leases, null); + } + manager = new CreditLeaseManager( + new ApiLeaseWireClient(credits()), leases, holds, creditLeases, this.logger, Clock.systemUTC()); + manager.startSweep(); + check = new CreditCheck( + // Null rather than a source wrapping nothing: CreditCheck degrades to a plain + // check on a null source, and a wrapper would sail past that guard and fail + // on the first cached-flag read instead. + this.dataStreamClient == null ? null : new DataStreamCreditCheckSource(this.dataStreamClient), + leases, + holds, + manager, + this.logger, + Clock.systemUTC(), + body -> eventBuffer.push(CreateEventRequestBody.builder() + .eventType(EventType.FLAG_CHECK) + .body(EventBody.of(body)) + .sentAt(OffsetDateTime.now()) + .build()), + null); + } + this.creditLeaseMode = mode; + this.leaseStore = leases; + this.reservations = holds; + this.creditLeaseManager = manager; + this.creditCheck = check; + this.leaseBackendShared = sharedBackend; + this.serverReservationTtl = serverTtl; + this.prewarmResolveTimeout = prewarmTimeout; + this.prewarms = manager == null + ? null + : Executors.newSingleThreadExecutor(runnable -> { + Thread thread = new Thread(runnable, "SchematicCreditLeasePrewarm"); + // Daemon, so a prewarm in flight can never hold a shutting-down process open. + thread.setDaemon(true); + return thread; + }); + this.shutdownHook = new Thread( () -> { try { @@ -163,6 +361,7 @@ public static class Builder { private Map headers; private DatastreamOptions datastreamOptions; private String eventCaptureBaseUrl; + private CreditLeaseConfig creditLeases; public Builder apiKey(String apiKey) { this.apiKey = apiKey; @@ -224,6 +423,15 @@ public Builder datastreamOptions(DatastreamOptions datastreamOptions) { return this; } + /** + * Enables credit holds on {@link Schematic#check} and + * {@link Schematic#trackWithReservation}. Omit it to leave the client credit-unaware. + */ + public Builder creditLeases(CreditLeaseConfig creditLeases) { + this.creditLeases = creditLeases; + return this; + } + public Builder eventCaptureBaseUrl(String eventCaptureBaseUrl) { this.eventCaptureBaseUrl = eventCaptureBaseUrl; return this; @@ -333,10 +541,20 @@ public RulesengineCheckFlagResult checkFlagWithEntitlement( } private RulesengineCheckFlagResult defaultFlagResult(String flagKey, String reason, String err) { + return defaultFlagResult(flagKey, reason, err, null); + } + + /** + * The result a check falls back to when it cannot get an answer. {@code perCheckDefault} is + * the caller's own default for this one check, which outranks the client-wide one; null means + * the caller did not name one. + */ + private RulesengineCheckFlagResult defaultFlagResult( + String flagKey, String reason, String err, Boolean perCheckDefault) { return RulesengineCheckFlagResult.builder() .flagKey(flagKey) .reason(reason) - .value(getFlagDefault(flagKey)) + .value(perCheckDefault != null ? perCheckDefault : getFlagDefault(flagKey)) .err(err) .build(); } @@ -349,11 +567,16 @@ private RulesengineCheckFlagResult defaultFlagResult(String flagKey, String reas */ private RulesengineCheckFlagResult tryDatastreamCheckFlag( String flagKey, Map company, Map user) { + return tryDatastreamCheckFlag(flagKey, company, user, null); + } + + private RulesengineCheckFlagResult tryDatastreamCheckFlag( + String flagKey, Map company, Map user, CheckFlagOptions preflight) { if (dataStreamClient == null || !dataStreamClient.isConnected()) { return null; } try { - return dataStreamClient.checkFlag(flagKey, company, user); + return dataStreamClient.checkFlag(flagKey, company, user, preflight); } catch (Exception e) { logger.debug("Datastream flag check failed for " + flagKey + ", falling back to API: " + e.getMessage()); return null; @@ -545,27 +768,377 @@ private RulesengineCheckFlagResult toRulesengineResult(CheckFlagResponseData dat } /** - * Checks a flag via the Schematic API, using the flag check result cache. + * Checks a flag via the Schematic API, using the flag check result cache. A preflighted check + * skips that cache in both directions, since it asks a different question than the plain + * check the cache is keyed for. */ private RulesengineCheckFlagResult checkFlagViaApi( String flagKey, Map company, Map user) { + return checkFlagViaApi(flagKey, company, user, null, null, null); + } + + /** + * The REST flag check. {@code perCheckDefault} is what a failure resolves to, so a caller that + * named a default on this one check gets it rather than the client-wide one. + */ + private RulesengineCheckFlagResult checkFlagViaApi( + String flagKey, + Map company, + Map user, + Duration timeout, + PreflightOptions preflight, + Boolean perCheckDefault) { try { - RulesengineCheckFlagResult cached = getCachedFlag(flagKey, company, user); - if (cached != null) { - return cached; + // Null once a preflight that the API would ignore, such as a zero usage, has been + // dropped: such a check is a plain one and keeps the cache. + PreflightRequestBody preflightBody = preflight != null ? preflight.toRequestBody() : null; + // The cache is keyed by flag, company and user, and a preflighted check asks a + // different question than a plain one: whether the action about to run would be + // allowed. So a preflighted verdict is neither answered from the cache nor written + // back to it. + if (preflightBody == null) { + RulesengineCheckFlagResult cached = getCachedFlag(flagKey, company, user); + if (cached != null) { + return cached; + } } - CheckFlagRequestBody request = - CheckFlagRequestBody.builder().company(company).user(user).build(); - CheckFlagResponse response = features().checkFlag(flagKey, request); + CheckFlagRequestBody.Builder request = + CheckFlagRequestBody.builder().company(company).user(user); + if (preflightBody != null) { + request.preflight(preflightBody); + } + CheckFlagResponse response = timeout == null + ? features().checkFlag(flagKey, request.build()) + : features() + .checkFlag( + flagKey, + request.build(), + RequestOptions.builder() + .timeout(CreditAmounts.millisAsInt(timeout), TimeUnit.MILLISECONDS) + .build()); RulesengineCheckFlagResult result = toRulesengineResult(response.getData()); - cacheFlag(flagKey, result, company, user); + if (preflightBody == null) { + cacheFlag(flagKey, result, company, user); + } return result; } catch (Exception e) { logger.error("Error checking flag via API: " + e.getMessage()); - return defaultFlagResult(flagKey, "flag default", e.getMessage()); + return defaultFlagResult(flagKey, "flag default", e.getMessage(), perCheckDefault); + } + } + + /** + * Which reservation mode a {@code check()} with usage resolves to right now. Null means no + * credit gating at all: leases are not configured, or the client is offline. + * + *

Auto resolves per check rather than once at startup: a DataStream that failed to start + * leaves no client behind, and the checks that follow gate server-side instead of silently + * dropping to a plain, ungated flag check. A DataStream that is merely disconnected stays in + * client mode and degrades through the plain check, which has its own story for that. + * + *

A loaded rules engine is part of that readiness. Without one, every local evaluation + * throws, so a client-mode check would fall through to a plain flag check that takes no hold + * and debits nothing: credits handed out ungated for as long as the engine is missing. + */ + private CreditLeaseMode effectiveLeaseMode() { + if (creditLeaseMode == null || offline) { + return null; + } + if (creditLeaseMode != CreditLeaseMode.AUTO) { + return creditLeaseMode; + } + boolean clientPlumbingReady = creditCheck != null && leaseStore != null && reservations != null; + boolean localEvaluationReady = dataStreamClient != null && dataStreamClient.hasRulesEngine(); + return localEvaluationReady && clientPlumbingReady ? CreditLeaseMode.CLIENT : CreditLeaseMode.SERVER; + } + + /** + * Credit-aware feature check. With credit leases configured and a usage on the options, this + * gates the check against the company's credit balance and hands back a hold on success: pass + * it to {@link #trackWithReservation} when the work completes. + * + *

In client mode the hold is carved out of a local lease and the flag is evaluated locally; + * in server mode it is one check-and-reserve call that evaluates the flag and takes the hold + * server-side. + * + *

Without credit leases, or without a usage, this is a plain flag check that issues no + * hold. The caller's preflight still reaches whichever path answers it, local or the API, so + * the check gates on the post-call balance, just without a hold. + */ + public CheckResult check( + String flagKey, Map company, Map user, CheckOptions options) { + CheckOptions opts = options != null ? options : CheckOptions.builder().build(); + Callable fallback = () -> plainCheck(flagKey, company, user, opts); + CreditLeaseMode mode = effectiveLeaseMode(); + if (opts.getUsage() == null || mode == null) { + return plainCheck(flagKey, company, user, opts); + } + + CheckRequest request = new CheckRequest( + flagKey, + company, + user, + opts.getUsage(), + opts.getEventSubtype(), + opts.getOnAcquireFailure() == OnAcquireFailure.FAIL_OPEN, + opts.getTimeout()); + if (mode == CreditLeaseMode.SERVER) { + ServerCreditCheck serverCheck = + new ServerCreditCheck(features(), credits(), logger, serverReservationTtl, Clock.systemUTC()); + return serverCheck.check(request, opts.getTimeout(), () -> checkDefault(flagKey, opts), fallback); + } + // Client mode without the local plumbing keeps the old behavior: a plain, ungated check. + if (creditCheck == null) { + return plainCheck(flagKey, company, user, opts); + } + return creditCheck.check(request, fallback); + } + + /** + * Settles a hold issued by {@link #check}. In client mode it refunds the unspent slice to the + * lease and emits a track event carrying the lease id; in server mode the event carries the + * reservation id and the server settles the hold when it processes the event. + * + *

When the work outlived the hold's TTL and the sweeper already returned it, the usage + * still has to be billed, so the event goes out anyway. A deterministic idempotency key + * derived from the reservation id keeps that recovery emit, and an accidental second settle, + * from billing twice. + */ + public void trackWithReservation(Reservation reservation, double actualQuantity) { + trackWithReservation(reservation, actualQuantity, null); + } + + /** Settles a hold, attaching traits to the event it emits. */ + public void trackWithReservation(Reservation reservation, double actualQuantity, Map traits) { + if (offline) { + return; + } + // check() allows without a hold in several ordinary cases: the feature is not + // credit-metered, the check failed open, the usage was zero, or leases are not configured. + // Callers pass the result's reservation straight through, so take the null and say how to + // bill the usage instead of throwing on a settle with nothing to settle. + if (reservation == null) { + logger.error("trackWithReservation was called without a reservation: the check allowed without taking a " + + "hold, so there is nothing to settle. Report the usage with track() instead."); + return; + } + // A non-finite quantity must reach neither the store, where the clamp would claim the hold + // with no refund of the unspent slice, nor the billing event. Skipping the settle leaves + // the hold to its TTL, where the sweeper refunds all of it, so no credits are lost and + // nothing bogus is billed. + if (!CreditAmounts.isValidQuantity(actualQuantity)) { + logger.error("trackWithReservation: invalid actualQuantity " + actualQuantity + " for reservation " + + reservation.getId() + "; skipping the settle, the hold is refunded at its TTL"); + return; + } + + EventBodyTrack track; + boolean updateMetrics; + if (reservation.getMode() == CreditLeaseMode.SERVER || reservations == null) { + // Server mode holds the credits server-side, so there is nothing local to consume. A + // client-mode handle with no store still has to carry its lease id and its key, since + // dropping either would double-debit the grant or double-bill the usage. + track = ReservationSettlement.buildTrackEvent(reservation, actualQuantity); + updateMetrics = true; + } else { + try { + ReservationSettlement.SettleOutcome outcome = + ReservationSettlement.settle(reservations, reservation, actualQuantity); + track = outcome.getTrack(); + updateMetrics = outcome.isSettledLocally(); + if (!updateMetrics) { + logger.debug("trackWithReservation: reservation " + reservation.getId() + " was not settled " + + "locally (expired, already settled, or the store was unreachable); emitting the track " + + "keyed for server-side dedupe"); + } + } catch (RuntimeException e) { + // The local settle failed, likely an unreachable Redis. The usage still has to be + // billed, and the un-settled hold is reclaimed by the sweeper or at lease expiry. + logger.warn("trackWithReservation: failed to settle reservation " + reservation.getId() + " locally (" + + e + "); emitting the track anyway"); + track = ReservationSettlement.buildTrackEvent(reservation, actualQuantity); + updateMetrics = false; + } + } + + try { + eventBuffer.push(buildReservationSettleEvent(track, objectMapToJsonNode(traits), reservation.getId())); + // The cached metric moves only when this call moved local state with it. The event is + // keyed off the reservation id, so the server drops a retried settle as a duplicate; + // bumping the metric for one would have the caller's next local evaluation gate on + // usage that was counted twice. + if (updateMetrics) { + updateCompanyMetrics(track); + } + } catch (Exception e) { + logger.error("Error sending track event: " + e.getMessage()); + } + } + + /** + * Folds a track event into the cached company's metrics, so a local evaluation right after it + * gates on the usage just recorded instead of waiting for the stream to push the new figure + * back. A settle is a track, and reads the same way. + */ + private void updateCompanyMetrics(EventBodyTrack body) { + Map company = body.getCompany().orElse(null); + if (company == null || company.isEmpty() || dataStreamClient == null || !dataStreamClient.isConnected()) { + return; + } + try { + dataStreamClient.updateCompanyMetrics(body); + } catch (Exception e) { + logger.error("Failed to update company metrics: " + e.getMessage()); + } + } + + /** + * Warms a credit lease for each named credit type, so the first {@link #check} against it does + * not pay the acquire round trip. Failures are logged, never thrown. + * + *

The keys are looked up over the DataStream, which both resolves the id and warms the cache + * so the first check hits the lease path. Only a lookup that comes up empty falls back to + * reading a {@code comp_}-prefixed value as the id. + */ + public void prewarm(Map company, List creditTypeIds) { + if (creditLeaseManager == null || leaseStore == null) { + logger.debug( + effectiveLeaseMode() == CreditLeaseMode.SERVER + ? "prewarm is a no-op in server mode, since there is no local lease to warm" + : "prewarm was called but credit leases are not configured"); + return; + } + if (company == null || company.isEmpty()) { + logger.debug("prewarm needs company keys"); + return; + } + if (creditTypeIds == null || creditTypeIds.isEmpty()) { + logger.debug("prewarm was given no credit types"); + return; + } + 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: the client is closing, skipping the acquire"); + return; + } + String companyId = resolveCompanyIdWithWait(company); + if (companyId == null) { + logger.debug("prewarm: the company did not resolve within " + prewarmResolveTimeout.toMillis() + + "ms (the first check will acquire)"); + return; + } + for (String creditTypeId : creditTypeIds) { + try { + creditLeaseManager.acquireIfNeeded(companyId, creditTypeId); + } catch (RuntimeException e) { + logger.warn("prewarm: failed to acquire a lease for " + creditTypeId + ": " + e); + } + } + } + + /** + * Resolves the company id the way the server does: the keys are looked up first, whatever they + * are named, actively fetching over the DataStream so the lookup warms the cache as a side + * effect. Only when nothing matches is a value read as the company's own id, by its + * {@code comp_} prefix. Null when the company never surfaced within the prewarm resolve + * timeout and the keys carry no Schematic id. + */ + private String resolveCompanyIdWithWait(Map company) { + if (dataStreamClient == null) { + return PrewarmCompanyResolver.schematicId(company, PrewarmCompanyResolver.COMPANY_ID_PREFIX); + } + return PrewarmCompanyResolver.resolve( + company, + dataStreamClient::getCachedCompany, + dataStreamClient::getCompany, + prewarmResolveTimeout, + CreditLeaseDefaults.PREWARM_POLL_INTERVAL, + () -> closing, + error -> { + logger.debug("prewarm: the DataStream company fetch failed (" + error + ")"); + return null; + }); + } + + /** + * Whether the local engine declined to answer, leaving the DataStream client's stand-in + * verdict in place of a real one. The flag's own default stands in there, which is the right + * answer for a caller that named none and the wrong one for a caller that did. + */ + private static boolean declinedByEngine(RulesengineCheckFlagResult result) { + String reason = result.getReason(); + return "RULES_ENGINE_UNAVAILABLE".equals(reason) || "RULES_ENGINE_ERROR".equals(reason); + } + + /** The plain flag check a credit-aware check defers to, with the caller's preflight threaded through. */ + private CheckResult plainCheck( + String flagKey, Map company, Map user, CheckOptions options) { + PreflightOptions preflight = PreflightOptions.fromUsage(options.getUsage(), options.getEventSubtype()); + RulesengineCheckFlagResult result; + if (offline) { + boolean value = checkDefault(flagKey, options); + result = RulesengineCheckFlagResult.builder() + .flagKey(flagKey) + .reason("flag default") + .value(value) + .build(); + } else { + RulesengineCheckFlagResult dsResult = tryDatastreamCheckFlag( + flagKey, company, user, DataStreamCreditCheckSource.toEngineOptions(preflight)); + if (dsResult != null) { + // Reported before the substitution below, so the event records what the engine + // said rather than the default the caller resolved its own verdict with. + enqueueFlagCheckEvent(flagKey, dsResult, company, user); + // The engine declining to answer is the case defaultValue exists for, so resolve + // it the way the offline and API branches do rather than passing on the stand-in + // the DataStream client substituted. + result = declinedByEngine(dsResult) + ? RulesengineCheckFlagResult.builder() + .from(dsResult) + .value(checkDefault(flagKey, options)) + .build() + : dsResult; + } else { + // The API answers a preflight too, so the caller's usage gates the REST path the + // same way it gates a local evaluation. The caller's timeout applies, since this + // is the call it is waiting on. + result = checkFlagViaApi( + flagKey, company, user, options.getTimeout(), preflight, options.getDefaultValue()); + } } + return new CheckResult( + result.getValue(), + result.getValue(), + result.getReason(), + result.getFlagKey(), + result.getFlagId().orElse(null), + result.getEntitlement().orElse(null), + null, + result.getErr().orElse(null)); + } + + /** + * Builds the event that settles a hold. Package-private for unit-testing the mapping. The key + * is derived from the reservation id, so a recovery emit and an accidental second settle + * collapse to one billed event server-side. + */ + static CreateEventRequestBody buildReservationSettleEvent( + EventBodyTrack track, Map traits, String reservationId) { + EventBodyTrack body = + EventBodyTrack.builder().from(track).traits(traits).build(); + return buildTrackEvent( + EventBody.of(body), + TrackOptions.builder() + .idempotencyKey(RESERVATION_TRACK_IDEMPOTENCY_PREFIX + reservationId) + .build()); + } + + /** The caller's default for this flag: the per-check one when set, otherwise the client's. */ + private boolean checkDefault(String flagKey, CheckOptions options) { + return options.getDefaultValue() != null ? options.getDefaultValue() : getFlagDefault(flagKey); } public void identify( @@ -593,6 +1166,32 @@ public void identify( } catch (Exception e) { logger.error("Error sending identify event: " + e.getMessage()); } + + List creditTypeIds = options != null ? options.getPrewarm() : null; + if (creditTypeIds != null && !creditTypeIds.isEmpty() && prewarms != null) { + Map companyKeys = company != null ? company.getKeys() : null; + // Flush first so the server processes the identify promptly: without it the company + // can sit in the buffer for a full flush interval while the prewarm waits on us. + try { + prewarms.execute(() -> { + try { + eventBuffer.flush(); + } catch (RuntimeException e) { + logger.debug("identify: the flush before the prewarm failed: " + e); + } + try { + prewarm(companyKeys, creditTypeIds); + } catch (RuntimeException e) { + logger.warn("identify: the prewarm failed: " + e); + } + }); + } catch (RejectedExecutionException e) { + // identify still recorded the company; only the warm-up is dropped, and a + // caller identifying after close() should not be handed a shutdown race to + // catch. + logger.debug("identify: the client is closed, skipping the prewarm"); + } + } } public void track( @@ -637,15 +1236,7 @@ public void track( .build(); eventBuffer.push(buildTrackEvent(EventBody.of(body), options)); - - // Update cached company metrics if datastream is active - if (company != null && !company.isEmpty() && dataStreamClient != null && dataStreamClient.isConnected()) { - try { - dataStreamClient.updateCompanyMetrics(body); - } catch (Exception e2) { - logger.error("Failed to update company metrics: " + e2.getMessage()); - } - } + updateCompanyMetrics(body); } catch (Exception e) { logger.error("Error sending track event: " + e.getMessage()); } @@ -688,6 +1279,7 @@ static CreateEventRequestBody buildTrackEvent(EventBody body, TrackOptions optio @Override public void close() { + closing = true; try { // Remove shutdown hook if we're closing explicitly try { @@ -696,6 +1288,34 @@ public void close() { // Shutdown is already in progress, hook will run automatically } + if (creditLeaseManager != null) { + // Refuse new lease work first, so the waits below are waiting on work that is + // already unwinding rather than on 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. + creditLeaseManager.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. + long deadline = System.nanoTime() + CreditLeaseDefaults.SHUTDOWN_DRAIN_TIMEOUT.toNanos(); + if (prewarms != null) { + prewarms.shutdown(); + try { + if (!prewarms.awaitTermination(remaining(deadline).toMillis(), TimeUnit.MILLISECONDS)) { + logger.warn("Timed out waiting for in-flight prewarms on close"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + creditLeaseManager.drain(remaining(deadline)); + // Never release leases held in a shared backend: a sibling process is still + // drawing on them. + if (!leaseBackendShared) { + creditLeaseManager.releaseAllLocalLeases(remaining(deadline)); + } + creditLeaseManager.close(remaining(deadline)); + } + if (dataStreamClient != null) { dataStreamClient.close(); } @@ -706,6 +1326,53 @@ public void close() { } } + /** + * Resolves one lease Redis setting against the DataStream cache's. The client and the key + * prefix resolve independently: a lease client of its own does not cost a caller the + * DataStream prefix, which would split the key layout of a mixed fleet sharing those leases. + */ + static T inheritFromDataStream(T configured, T fromDataStream) { + return configured != null ? configured : fromDataStream; + } + + /** + * Names the configured options that only steer client-mode plumbing, so server mode can say + * once that it is ignoring them. + */ + private static String clientOnlyOptions(CreditLeaseConfig config) { + List names = new ArrayList<>(); + if (config.getDefaultLeaseDuration() != null) { + names.add("defaultLeaseDuration"); + } + if (config.getDefaultLeaseSize() != null) { + names.add("defaultLeaseSize"); + } + if (config.getLowWaterMark() != null) { + names.add("lowWaterMark"); + } + if (config.getSweepInterval() != null) { + names.add("sweepInterval"); + } + if (config.getRedisClient() != null) { + names.add("redisClient"); + } + if (config.getRedisKeyPrefix() != null) { + names.add("redisKeyPrefix"); + } + if (config.getPrewarmResolveTimeout() != null) { + names.add("prewarmResolveTimeout"); + } + if (config.getOverrides() != null && !config.getOverrides().isEmpty()) { + names.add("overrides"); + } + return String.join(", ", names); + } + + private static Duration remaining(long deadlineNanos) { + long left = deadlineNanos - System.nanoTime(); + return left <= 0 ? Duration.ZERO : Duration.ofNanos(left); + } + private boolean getFlagDefault(String flagKey) { return flagDefaults.getOrDefault(flagKey, false); } @@ -724,7 +1391,7 @@ private String buildCacheKey(String flagKey, Map company, Map objectMapToJsonNode(Map map) { + private static Map objectMapToJsonNode(Map map) { if (map == null) { return null; } diff --git a/src/main/java/com/schematic/api/credits/ApiLeaseWireClient.java b/src/main/java/com/schematic/api/credits/ApiLeaseWireClient.java new file mode 100644 index 0000000..3c1c503 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/ApiLeaseWireClient.java @@ -0,0 +1,100 @@ +package com.schematic.api.credits; + +import com.schematic.api.core.RequestOptions; +import com.schematic.api.resources.credits.CreditsClient; +import com.schematic.api.resources.credits.requests.AcquireCreditLeaseRequestBody; +import com.schematic.api.resources.credits.requests.ExtendCreditLeaseRequestBody; +import com.schematic.api.types.CreditLeaseResponseData; +import java.time.Duration; +import java.time.Instant; +import java.time.OffsetDateTime; +import java.time.ZoneOffset; +import java.util.UUID; +import java.util.concurrent.TimeUnit; + +/** Adapts the generated credits client to {@link LeaseWireClient}. */ +public final class ApiLeaseWireClient implements LeaseWireClient { + + private final CreditsClient credits; + + public ApiLeaseWireClient(CreditsClient credits) { + this.credits = credits; + } + + /** + * Acquire takes the client's default retry policy: the server hands back the slot's existing + * active lease rather than opening a second one, so a retry after a lost response returns the + * lease the first attempt created. + */ + @Override + public LeaseGrant acquire(String companyId, String creditTypeId, double requestedAmount, Instant expiresAt) { + return acquire(companyId, creditTypeId, requestedAmount, expiresAt, null); + } + + @Override + public LeaseGrant acquire( + String companyId, String creditTypeId, double requestedAmount, Instant expiresAt, Duration timeout) { + AcquireCreditLeaseRequestBody body = AcquireCreditLeaseRequestBody.builder() + .companyId(companyId) + .creditTypeId(creditTypeId) + .requestedAmount(requestedAmount) + .expiresAt(toOffsetDateTime(expiresAt)) + .build(); + return grantFrom( + timeout == null + ? credits.acquireCreditLease(body).getData() + : credits.acquireCreditLease(body, requestOptions(timeout)) + .getData()); + } + + @Override + public LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt) { + return extend(leaseId, additionalAmount, expiresAt, null); + } + + @Override + public LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt, Duration timeout) { + // An extend is an increment, so a retry without a key would grant the tranche twice. The + // key is minted once per extend, outside the call, so the transport's retries resend the + // same one and every attempt of this extend collapses to one grow, while a later extend + // gets its own key. + String idempotencyKey = UUID.randomUUID().toString(); + ExtendCreditLeaseRequestBody body = ExtendCreditLeaseRequestBody.builder() + .additionalAmount(additionalAmount) + .expiresAt(toOffsetDateTime(expiresAt)) + .idempotencyKey(idempotencyKey) + .build(); + return grantFrom( + timeout == null + ? credits.extendCreditLease(leaseId, body).getData() + : credits.extendCreditLease(leaseId, body, requestOptions(timeout)) + .getData()); + } + + private static RequestOptions requestOptions(Duration timeout) { + return RequestOptions.builder() + .timeout(CreditAmounts.millisAsInt(timeout), TimeUnit.MILLISECONDS) + .build(); + } + + @Override + public void release(String leaseId) { + credits.releaseCreditLease(leaseId); + } + + private static LeaseGrant grantFrom(CreditLeaseResponseData data) { + if (data == null) { + throw new IllegalStateException("credit lease response carried no data"); + } + return new LeaseGrant( + data.getId(), + data.getCompanyId(), + data.getCreditTypeId(), + data.getGrantedAmount(), + data.getExpiresAt().toInstant()); + } + + private static OffsetDateTime toOffsetDateTime(Instant instant) { + return OffsetDateTime.ofInstant(instant, ZoneOffset.UTC); + } +} diff --git a/src/main/java/com/schematic/api/credits/CheckOptions.java b/src/main/java/com/schematic/api/credits/CheckOptions.java new file mode 100644 index 0000000..6eba31e --- /dev/null +++ b/src/main/java/com/schematic/api/credits/CheckOptions.java @@ -0,0 +1,109 @@ +package com.schematic.api.credits; + +import java.time.Duration; + +/** The per-call knobs a credit-aware check takes. */ +public final class CheckOptions { + + private final Double usage; + private final String eventSubtype; + private final OnAcquireFailure onAcquireFailure; + private final Boolean defaultValue; + private final Duration timeout; + + private CheckOptions(Builder builder) { + this.usage = builder.usage; + this.eventSubtype = builder.eventSubtype; + this.onAcquireFailure = builder.onAcquireFailure; + this.defaultValue = builder.defaultValue; + this.timeout = builder.timeout; + } + + public static Builder builder() { + return new Builder(); + } + + public Double getUsage() { + return usage; + } + + public String getEventSubtype() { + return eventSubtype; + } + + public OnAcquireFailure getOnAcquireFailure() { + return onAcquireFailure == null ? OnAcquireFailure.FAIL_CLOSED : onAcquireFailure; + } + + public Boolean getDefaultValue() { + return defaultValue; + } + + public Duration getTimeout() { + return timeout; + } + + public static final class Builder { + private Double usage; + private String eventSubtype; + private OnAcquireFailure onAcquireFailure; + private Boolean defaultValue; + private Duration timeout; + + /** + * The units of the metered event the operation is about to record. The check holds + * {@code usage} times the entitlement's consumption rate. Omit it for a plain flag check. + * + *

A usage also rides on every check that takes no hold, as the preflight a local + * evaluation and the API both answer, so the verdict accounts for what this call is about + * to spend. A preflighted verdict is never cached, since the cache answers the plain + * question. + */ + public Builder usage(double usage) { + this.usage = usage; + return this; + } + + /** + * The event the usage applies to, which picks the credit condition to gate on when a flag + * meters more than one event. Defaults to the entitlement's own subtype. + */ + public Builder eventSubtype(String eventSubtype) { + this.eventSubtype = eventSubtype; + return this; + } + + /** + * What to do when the check cannot gate at all. Defaults to + * {@link OnAcquireFailure#FAIL_CLOSED}. + */ + public Builder onAcquireFailure(OnAcquireFailure onAcquireFailure) { + this.onAcquireFailure = onAcquireFailure; + return this; + } + + /** The value to fall back to, in place of the client's configured flag default. */ + public Builder defaultValue(boolean defaultValue) { + this.defaultValue = defaultValue; + return this; + } + + /** + * A per-call timeout for every call this check waits on: the check-and-reserve call in + * server mode, the REST flag check when the check falls back to one, and the client-mode + * lease acquire and extend. + * + *

Lease calls are single-flighted per company and credit type, so the timeout of + * whichever caller opened the flight governs everyone who joins it. Background top-ups, + * which no caller is waiting on, keep the client's own timeout. + */ + public Builder timeout(Duration timeout) { + this.timeout = timeout; + return this; + } + + public CheckOptions build() { + return new CheckOptions(this); + } + } +} diff --git a/src/main/java/com/schematic/api/credits/CheckRequest.java b/src/main/java/com/schematic/api/credits/CheckRequest.java new file mode 100644 index 0000000..4cb288b --- /dev/null +++ b/src/main/java/com/schematic/api/credits/CheckRequest.java @@ -0,0 +1,84 @@ +package com.schematic.api.credits; + +import java.time.Duration; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; + +/** One caller's ask, as the credit-gated flows take it. */ +public final class CheckRequest { + + private final String flagKey; + private final Map company; + private final Map user; + private final double usage; + private final String eventSubtype; + private final boolean failOpen; + private final Duration timeout; + + public CheckRequest( + String flagKey, + Map company, + Map user, + double usage, + String eventSubtype, + boolean failOpen) { + this(flagKey, company, user, usage, eventSubtype, failOpen, null); + } + + public CheckRequest( + String flagKey, + Map company, + Map user, + double usage, + String eventSubtype, + boolean failOpen, + Duration timeout) { + this.flagKey = flagKey; + this.company = copy(company); + this.user = copy(user); + this.usage = usage; + this.eventSubtype = eventSubtype; + this.failOpen = failOpen; + this.timeout = timeout; + } + + /** The caller's per-check timeout, or null for the client's own. */ + public Duration getTimeout() { + return timeout; + } + + private static Map copy(Map keys) { + if (keys == null || keys.isEmpty()) { + return Collections.emptyMap(); + } + return Collections.unmodifiableMap(new LinkedHashMap<>(keys)); + } + + public String getFlagKey() { + return flagKey; + } + + /** The caller's evaluation keys, threaded onto the hold so the settling event attributes usage the same way. */ + public Map getCompany() { + return company; + } + + public Map getUser() { + return user; + } + + public double getUsage() { + return usage; + } + + /** Names the event the usage applies to. Empty defers to the entitlement's own subtype. */ + public String getEventSubtype() { + return eventSubtype; + } + + /** Errs on the side of assuming the credits are there when the flow cannot gate. */ + public boolean isFailOpen() { + return failOpen; + } +} diff --git a/src/main/java/com/schematic/api/credits/CheckResult.java b/src/main/java/com/schematic/api/credits/CheckResult.java new file mode 100644 index 0000000..2ae7390 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/CheckResult.java @@ -0,0 +1,74 @@ +package com.schematic.api.credits; + +import com.schematic.api.types.RulesengineFeatureEntitlement; + +/** What a credit-aware check resolved, and the hold it took if it took one. */ +public final class CheckResult { + + private final boolean allowed; + private final boolean value; + private final String reason; + private final String flagKey; + private final String flagId; + private final RulesengineFeatureEntitlement entitlement; + private final Reservation reservation; + private final String err; + + public CheckResult( + boolean allowed, + boolean value, + String reason, + String flagKey, + String flagId, + RulesengineFeatureEntitlement entitlement, + Reservation reservation, + String err) { + this.allowed = allowed; + this.value = value; + this.reason = reason; + this.flagKey = flagKey; + this.flagId = flagId; + this.entitlement = entitlement; + this.reservation = reservation; + this.err = err; + } + + /** Whether the caller may proceed. */ + public boolean isAllowed() { + return allowed; + } + + /** The flag's boolean value, which the non-credit paths mirror onto {@link #isAllowed()}. */ + public boolean getValue() { + return value; + } + + public String getReason() { + return reason; + } + + public String getFlagKey() { + return flagKey; + } + + public String getFlagId() { + return flagId; + } + + public RulesengineFeatureEntitlement getEntitlement() { + return entitlement; + } + + /** + * The hold this check carved out, or null when it allowed without taking one. Pass it to + * {@code trackWithReservation} when the work completes. + */ + public Reservation getReservation() { + return reservation; + } + + /** The failure behind a result the SDK resolved itself, or null. */ + public String getErr() { + return err; + } +} diff --git a/src/main/java/com/schematic/api/credits/CreditAmounts.java b/src/main/java/com/schematic/api/credits/CreditAmounts.java new file mode 100644 index 0000000..90edf5b --- /dev/null +++ b/src/main/java/com/schematic/api/credits/CreditAmounts.java @@ -0,0 +1,65 @@ +package com.schematic.api.credits; + +import java.math.BigDecimal; +import java.time.Duration; + +/** Credit-amount helpers shared by the stores and the check flows. */ +public final class CreditAmounts { + + /** + * 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. + */ + public static boolean isValidQuantity(double value) { + return !Double.isNaN(value) && !Double.isInfinite(value) && value >= 0; + } + + /** + * Formats an amount for a Redis hash field. Plain decimal, shortest exact form, so a fleet of + * SDKs reading the same hash sees the same figures and no fractional rate is truncated. + */ + public static String format(double value) { + if (value == Math.rint(value) && !Double.isInfinite(value) && Math.abs(value) < 1e15) { + return Long.toString((long) value); + } + return BigDecimal.valueOf(value).stripTrailingZeros().toPlainString(); + } + + /** Parses an amount written by {@link #format}, or by another SDK's equivalent. */ + public static double parse(String raw, double fallback) { + if (raw == null || raw.isEmpty()) { + return fallback; + } + try { + return Double.parseDouble(raw); + } catch (NumberFormatException e) { + return fallback; + } + } + + /** Keeps local bookkeeping from ever debiting a lease past the hold it took. */ + public static double clampConsumption(double creditsConsumed, double creditsReserved) { + if (Double.isNaN(creditsConsumed) || creditsConsumed < 0) { + return 0; + } + return Math.min(creditsConsumed, creditsReserved); + } + + /** + * A duration as the milliseconds an int-typed request option takes, clamped at both ends. + * Casting alone wraps anything past the int range into a negative, which the transport reads + * as an instant timeout: the opposite of the long wait the caller asked for. A caller passing + * a negative or sub-millisecond duration lands on the same value, so the floor keeps it a + * timeout the transport can express rather than one that fails every call before it is sent. + */ + public static int millisAsInt(Duration timeout) { + long millis = timeout.toMillis(); + if (millis > Integer.MAX_VALUE) { + return Integer.MAX_VALUE; + } + return millis < 1 ? 1 : (int) millis; + } + + private CreditAmounts() {} +} diff --git a/src/main/java/com/schematic/api/credits/CreditCheck.java b/src/main/java/com/schematic/api/credits/CreditCheck.java new file mode 100644 index 0000000..c33a9d7 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/CreditCheck.java @@ -0,0 +1,503 @@ +package com.schematic.api.credits; + +import com.schematic.api.logger.SchematicLogger; +import com.schematic.api.types.EventBodyFlagCheck; +import com.schematic.api.types.RulesengineCheckFlagResult; +import com.schematic.api.types.RulesengineCompany; +import com.schematic.api.types.RulesengineEntitlementValueType; +import com.schematic.api.types.RulesengineFeatureEntitlement; +import com.schematic.api.types.RulesengineFlag; +import com.schematic.api.types.RulesengineUser; +import java.time.Clock; +import java.time.Instant; +import java.util.HashMap; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.Callable; +import java.util.function.Supplier; + +/** + * Gates one check against a local lease, returning a hold when it allows. + * + *

One check runs the rules engine twice. The first run is a probe against the company's real + * balance that names the credit being metered; the second gates the call against the lease's local + * balance, after the credits have already been debited. conformance/SPEC.md explains why each step + * is ordered the way it is, and the vectors pin it. + */ +public final class CreditCheck { + + private final CreditCheckDataStream dataStream; + private final LeaseStore leases; + private final ReservationStore reservations; + private final CreditLeaseManager manager; + private final SchematicLogger logger; + private final Clock clock; + private final FlagCheckReporter flagChecks; + private final Supplier reservationIds; + + public CreditCheck( + CreditCheckDataStream dataStream, + LeaseStore leases, + ReservationStore reservations, + CreditLeaseManager manager, + SchematicLogger logger, + Clock clock, + FlagCheckReporter flagChecks, + Supplier reservationIds) { + this.dataStream = dataStream; + this.leases = leases; + this.reservations = reservations; + this.manager = manager; + this.logger = logger; + this.clock = clock != null ? clock : Clock.systemUTC(); + this.flagChecks = flagChecks; + this.reservationIds = reservationIds != null + ? reservationIds + : () -> UUID.randomUUID().toString(); + } + + /** + * Runs the credit-gated check. + * + *

{@code fallback} is the plain flag check. Every step that cannot resolve a credit to meter + * defers to it, since the plain check has its own degradation story and issues no hold. A step + * that can resolve the credit but cannot gate on it goes through the caller's fail-open or + * fail-closed contract instead. + */ + public CheckResult check(CheckRequest request, Callable fallback) { + // A malformed usage must never reach the stores: NaN slips through every numeric + // comparison, and a NaN balance on a possibly shared lease would approve every later + // reserve. The caller asked for a contract covering exactly this, so resolve it through + // that rather than letting it surface as an opaque reserve failure. + if (!CreditAmounts.isValidQuantity(request.getUsage())) { + error("Lease check: invalid usage " + request.getUsage() + " for flag " + request.getFlagKey() + + "; must be a finite, non-negative number"); + return emit(request, staticFailure(request, "invalid_usage", null), null, null, null); + } + + // Nothing to reserve. The plain check still carries the preflight, so every rule evaluates + // normally; a zero-credit hold would only be a no-op. + if (request.getUsage() == 0) { + debug("Lease check: usage is 0 for flag " + request.getFlagKey() + + ", nothing to reserve, using a plain check"); + return fallBack(fallback); + } + + if (dataStream == null) { + debug("Lease check: no DataStream, using a plain check"); + return fallBack(fallback); + } + + RulesengineFlag flag; + try { + flag = dataStream.getFlag(request.getFlagKey()); + } catch (RuntimeException e) { + warn("Lease check: failed to load flag " + request.getFlagKey() + ": " + e); + flag = null; + } + if (flag == null) { + debug("Lease check: no cached flag for " + request.getFlagKey() + ", using a plain check"); + return fallBack(fallback); + } + + if (request.getCompany().isEmpty()) { + debug("Lease check: no company keys, using a plain check"); + return fallBack(fallback); + } + + // Resolve company and user the way a plain DataStream check does: cache first, then a live + // fetch. Evaluating without an entity the caller named would silently skip its targeted + // rules and overrides, so a miss defers to the plain check instead. + RulesengineCompany company; + try { + company = dataStream.getCompany(request.getCompany()); + } catch (RuntimeException e) { + debug("Lease check: company fetch failed (" + e + "), using a plain check"); + company = null; + } + if (company == null) { + return fallBack(fallback); + } + + RulesengineUser user = null; + if (!request.getUser().isEmpty()) { + try { + user = dataStream.getUser(request.getUser()); + } catch (RuntimeException e) { + debug("Lease check: user fetch failed (" + e + "), using a plain check"); + } + if (user == null) { + return fallBack(fallback); + } + } + + // Entitlement-first resolution. The probe runs against the real balance with no preflight: + // applying a credit cost to a lease-depleted server balance could fail the credit + // condition, drop the engine to a lower-priority rule, and hide the entitlement being + // looked for. + RulesengineCheckFlagResult probe; + try { + probe = dataStream.evaluateFlag(flag, company, user, null); + } catch (Exception e) { + // A probe failure is a resolution miss, not the gate, and no hold exists yet to cancel. + warn("Lease check: entitlement probe failed for flag " + request.getFlagKey() + " (" + e + + "), using a plain check"); + return fallBack(fallback); + } + if (probe == null) { + return fallBack(fallback); + } + + RulesengineFeatureEntitlement entitlement = probe.getEntitlement().orElse(null); + if (entitlement == null || !RulesengineEntitlementValueType.CREDIT.equals(entitlement.getValueType())) { + // A boolean or override grant, a numeric allocation, unlimited, or simply not + // entitled. The feature resolves without drawing a credit, so skip the lease and the + // reserve round trip entirely. + debug("Lease check: flag " + request.getFlagKey() + " matched a non-credit entitlement, using a plain " + + "check, no reservation"); + return fallBack(fallback); + } + + String creditId = entitlement.getCreditId().orElse(null); + double consumptionRate = entitlement.getConsumptionRate().orElse(0.0); + // The caller's subtype wins; otherwise the entitlement names the metered event. A credit + // entitlement with neither a resolvable subtype nor a positive rate can never be billed, + // so it is not gateable. + String eventSubtype = request.getEventSubtype(); + if (eventSubtype == null || eventSubtype.isEmpty()) { + eventSubtype = entitlement.getEventSubtype().orElse(null); + } + if (creditId == null + || creditId.isEmpty() + || consumptionRate <= 0 + || eventSubtype == null + || eventSubtype.isEmpty()) { + debug("Lease check: flag " + request.getFlagKey() + " has an incomplete credit entitlement, using a " + + "plain check"); + return fallBack(fallback); + } + + // 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. + double creditCost = Math.ceil(request.getUsage()) * consumptionRate; + String companyId = company.getId(); + String userId = user == null ? null : user.getId(); + + LeaseState lease = manager.acquireIfNeeded(companyId, creditId, request.getTimeout()); + if (lease == null) { + return failure(request, "lease_acquire_failed", flag, company, user, creditId, companyId, userId); + } + + // Resolved before the debit, not after it. Each of these can throw, and between the debit + // and the record that pins it there is nothing to refund a stranded slice: it would sit on + // the lease until expiry with no hold naming it. + String reservationId; + Instant expiresAt; + try { + ResolvedLeaseConfig resolved = manager.resolveConfig(creditId); + reservationId = reservationIds.get(); + expiresAt = clock.instant().plus(resolved.getReservationTtl()); + } catch (RuntimeException e) { + error("Lease check: could not prepare a reservation for " + companyId + "/" + creditId + ": " + e); + return failure(request, "lease_store_error", flag, company, user, creditId, companyId, userId); + } + + // tryReserve is the atomic gate: check and debit in one step, returning the post-debit + // balance so the pre-debit figure needs no second read, and the lease it debited. + ReserveResult reserve; + try { + reserve = leases.tryReserve(companyId, creditId, creditCost); + if (reserve == null) { + // Pass the cost as required credits so a single large request extends even while + // the ratio still sits above the water mark. + manager.maybeExtend(companyId, creditId, creditCost, request.getTimeout()); + reserve = leases.tryReserve(companyId, creditId, creditCost); + } + } catch (RuntimeException e) { + error("Lease check: reserve against " + companyId + "/" + creditId + " failed: " + e); + return failure(request, "lease_store_error", flag, company, user, creditId, companyId, userId); + } + if (reserve == null) { + return failure(request, "insufficient_lease_balance", flag, company, user, creditId, companyId, userId); + } + String debitedLeaseId = reserve.getLeaseId(); + if (debitedLeaseId == null || debitedLeaseId.isEmpty()) { + // A store that debited without naming the lease leaves the hold nothing to pin its + // refunds to, and pinning the acquired lease instead would refund and bill a lease + // that never held these credits. Hand the debit straight back, unpinned since there is + // no id to pin it to, and resolve through the caller's contract. + error("Lease check: reserve against " + companyId + "/" + creditId + " named no lease"); + try { + leases.refund(companyId, creditId, creditCost, null); + } catch (RuntimeException e) { + warn("Lease check: could not return an unattributed debit for " + companyId + "/" + creditId + " (" + e + + "); the slice is reclaimed at lease expiry"); + } + return failure(request, "lease_store_error", flag, company, user, creditId, companyId, userId); + } + + // Record the hold after the debit and before the gate. A crash between the debit and this + // add leaks at most this one hold, reclaimed when the lease expires server-side; recording + // first would instead leave a record with no debit, which a later consume would refund + // into a double-spend. + Reservation reservation = new Reservation( + reservationId, + // The lease the debit came out of, which the slot may have taken on since the + // acquire above: the window between them spans the extend's network call. A hold + // pinned to the lease the acquire returned would have its refunds dropped and + // would bill the wrong lease, so this is never the acquired id. + debitedLeaseId, + CreditLeaseMode.CLIENT, + companyId, + creditId, + eventSubtype, + request.getUsage(), + creditCost, + consumptionRate, + expiresAt, + request.getCompany(), + request.getUser()); + try { + reservations.add(reservation); + } catch (RuntimeException e) { + error("Lease check: failed to persist reservation " + reservation.getId() + ": " + e); + undoDebit(reservation); + return failure(request, "lease_store_error", flag, company, user, creditId, companyId, userId); + } + + // Gate against the lease's local view rather than the server's balance. The substituted + // figure is the pre-reservation balance (what tryReserve returned plus what it debited, + // exact as of the debit), and the credit cost tells the engine what this call costs, so it + // evaluates the same arithmetic tryReserve just enforced, plus every non-credit rule. + RulesengineCompany substituted = substituteCreditBalance(company, creditId, reserve.getBalance() + creditCost); + RulesengineCheckFlagResult result; + try { + result = dataStream.evaluateFlag( + flag, substituted, user, PreflightOptions.forCreditCost(creditId, creditCost)); + } catch (Exception e) { + error("Lease check: rules evaluation failed for flag " + request.getFlagKey() + ": " + e); + // The engine itself is down, so there is no fail-open re-evaluation to run: resolve + // the mode statically. + cancelReservation(reservation); + return emit(request, staticFailure(request, "wasm_error: " + e, flag), companyId, userId, null); + } + if (result == null) { + cancelReservation(reservation); + return emit(request, staticFailure(request, "wasm_error: no result", flag), companyId, userId, null); + } + + // Engine-evaluated exits report the engine's resolved ids, mirroring the plain DataStream + // path's flag_check event. + String resolvedCompanyId = result.getCompanyId().orElse(companyId); + String resolvedUserId = result.getUserId().orElse(userId); + String ruleId = result.getRuleId().orElse(null); + + if (!result.getValue()) { + cancelReservation(reservation); + return emit( + request, + new CheckResult( + false, + false, + orElse(result.getReason(), "denied_by_engine"), + orElse(result.getFlagKey(), request.getFlagKey()), + result.getFlagId().orElse(null), + result.getEntitlement().orElse(null), + null, + null), + resolvedCompanyId, + resolvedUserId, + ruleId); + } + + // Allowed against the substituted balance, so the hold stands. Top the lease up in the + // background now that it has been drawn down: the check that drew it down should not pay + // for the top-up. + manager.extendInBackground(companyId, creditId); + return emit( + request, + new CheckResult( + true, + true, + orElse(result.getReason(), "lease_reserved"), + orElse(result.getFlagKey(), request.getFlagKey()), + result.getFlagId().orElse(null), + result.getEntitlement().orElse(null), + reservation, + null), + resolvedCompanyId, + resolvedUserId, + ruleId); + } + + /** + * Resolves a check that could not gate: acquire failed, store unreachable, or the lease is + * exhausted. + * + *

Fail-closed denies. Fail-open means assume the credits are there, not skip the + * evaluation: the rules still run with the balance substituted to an effectively unlimited + * value, so plan targeting, overrides, and every non-credit condition still apply, and a + * company that is not entitled stays denied with the lease backend down. Only an error in that + * evaluation drops to a blanket allow. + */ + private CheckResult failure( + CheckRequest request, + String reason, + RulesengineFlag flag, + RulesengineCompany company, + RulesengineUser user, + String creditId, + String companyId, + String userId) { + if (!request.isFailOpen()) { + return emit(request, staticFailure(request, reason, flag), companyId, userId, null); + } + + RulesengineCheckFlagResult result; + try { + RulesengineCompany substituted = + substituteCreditBalance(company, creditId, CreditLeaseDefaults.FAIL_OPEN_BALANCE); + result = dataStream.evaluateFlag( + flag, substituted, user, PreflightOptions.fromUsage(request.getUsage(), request.getEventSubtype())); + } catch (Exception e) { + warn("Lease check: the fail-open evaluation failed (" + e + "); allowing"); + return emit(request, staticFailure(request, reason, flag), companyId, userId, null); + } + if (result == null) { + return emit(request, staticFailure(request, reason, flag), companyId, userId, null); + } + return emit( + request, + new CheckResult( + result.getValue(), + result.getValue(), + orElse(result.getReason(), "evaluated") + " (" + reason + "_fail_open)", + orElse(result.getFlagKey(), request.getFlagKey()), + result.getFlagId().orElse(flag.getId()), + result.getEntitlement().orElse(null), + null, + reason), + companyId, + userId, + null); + } + + /** + * Resolves a mode with no evaluation behind it: deny for fail-closed, blanket allow for + * fail-open. Used when the engine is the thing that failed, and when the fail-open evaluation + * itself errors. + */ + private static CheckResult staticFailure(CheckRequest request, String reason, RulesengineFlag flag) { + boolean allowed = request.isFailOpen(); + return new CheckResult( + allowed, + allowed, + allowed ? reason + "_fail_open" : reason, + request.getFlagKey(), + flag == null ? null : flag.getId(), + null, + null, + reason); + } + + /** + * Returns a debit whose reservation record never landed, rather than stranding it until lease + * expiry. Consume claims whatever slice of the add made it to the store and refunds it; + * nothing claimed means nothing landed, so the debit is refunded directly. Both are pinned to + * the lease the debit came out of. If the undo itself fails, accept the bounded leak: the + * slice comes back at lease expiry, which beats risking a double refund. + */ + private void undoDebit(Reservation reservation) { + try { + Double claimed = reservations.consume(reservation.getId(), 0); + if (claimed == null) { + leases.refund( + reservation.getCompanyId(), + reservation.getCreditTypeId(), + reservation.getCreditsReserved(), + reservation.getLeaseId()); + } + } catch (RuntimeException e) { + warn("Lease check: could not undo the local debit for " + reservation.getId() + " (" + e + + "); the slice is reclaimed at lease expiry"); + } + } + + /** Claims the hold and refunds all of it. Best effort: a failure leaves it for the sweeper. */ + private void cancelReservation(Reservation reservation) { + try { + reservations.consume(reservation.getId(), 0); + } catch (RuntimeException e) { + warn("Lease check: failed to cancel reservation " + reservation.getId() + " (" + e + + "); its hold is reclaimed by the sweeper or at lease expiry"); + } + } + + /** Copies the company with one credit balance replaced, leaving the cached entity untouched. */ + static RulesengineCompany substituteCreditBalance(RulesengineCompany company, String creditId, double balance) { + Map balances = new HashMap<>(company.getCreditBalances()); + balances.put(creditId, balance); + return RulesengineCompany.builder() + .from(company) + .creditBalances(balances) + .build(); + } + + /** + * Reports a credit-path resolution and passes the result straight through. Analytics must + * never change a verdict the caller is already acting on, so this only ever adds an event. + */ + private CheckResult emit(CheckRequest request, CheckResult result, String companyId, String userId, String ruleId) { + if (flagChecks == null) { + return result; + } + try { + flagChecks.report(EventBodyFlagCheck.builder() + .flagKey(result.getFlagKey()) + .reason(result.getReason()) + .value(result.getValue()) + .companyId(companyId) + .error(result.getErr()) + .flagId(result.getFlagId()) + .reqCompany(request.getCompany().isEmpty() ? null : request.getCompany()) + .reqUser(request.getUser().isEmpty() ? null : request.getUser()) + .ruleId(ruleId) + .userId(userId) + .build()); + } catch (RuntimeException e) { + error("Lease check: failed to report the flag check: " + e); + } + return result; + } + + private CheckResult fallBack(Callable fallback) { + try { + return fallback.call(); + } catch (Exception e) { + throw new IllegalStateException("plain flag check failed", e); + } + } + + private static String orElse(String value, String fallback) { + return value == null || value.isEmpty() ? fallback : value; + } + + private void debug(String message) { + if (logger != null) { + logger.debug(message); + } + } + + private void warn(String message) { + if (logger != null) { + logger.warn(message); + } + } + + private void error(String message) { + if (logger != null) { + logger.error(message); + } + } +} diff --git a/src/main/java/com/schematic/api/credits/CreditCheckDataStream.java b/src/main/java/com/schematic/api/credits/CreditCheckDataStream.java new file mode 100644 index 0000000..c886c38 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/CreditCheckDataStream.java @@ -0,0 +1,29 @@ +package com.schematic.api.credits; + +import com.schematic.api.types.RulesengineCheckFlagResult; +import com.schematic.api.types.RulesengineCompany; +import com.schematic.api.types.RulesengineFlag; +import com.schematic.api.types.RulesengineUser; +import java.util.Map; + +/** + * The slice of the DataStream client a credit-gated check touches. Narrow on purpose: it keeps + * this package off the wider DataStream surface and lets the conformance runner drive the flow + * without a socket or a WASM runtime. + */ +public interface CreditCheckDataStream { + + /** Reads a flag from the local cache, or returns null when it is not there. */ + RulesengineFlag getFlag(String flagKey); + + /** Resolves company keys, cache first, then over the wire. Null when it cannot be resolved. */ + RulesengineCompany getCompany(Map keys); + + /** Resolves user keys, cache first, then over the wire. Null when it cannot be resolved. */ + RulesengineUser getUser(Map keys); + + /** Runs the rules engine. A null {@code preflight} means no preflight. */ + RulesengineCheckFlagResult evaluateFlag( + RulesengineFlag flag, RulesengineCompany company, RulesengineUser user, PreflightOptions preflight) + throws Exception; +} diff --git a/src/main/java/com/schematic/api/credits/CreditLeaseConfig.java b/src/main/java/com/schematic/api/credits/CreditLeaseConfig.java new file mode 100644 index 0000000..8bf12a8 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/CreditLeaseConfig.java @@ -0,0 +1,266 @@ +package com.schematic.api.credits; + +import java.time.Duration; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; +import redis.clients.jedis.JedisPooled; + +/** + * Enables credit reservations on {@code check()} and {@code trackWithReservation()}. Leave it off + * the builder to keep the SDK credit-unaware, where {@code check()} is a plain flag check. + * + *

Client mode (local leases) needs datastream, so the SDK has the cached flag and company + * state it gates against. Without datastream the SDK gates in server mode instead: one + * check-and-reserve API call per check. + */ +public final class CreditLeaseConfig { + + private final CreditLeaseMode mode; + private final Duration defaultLeaseDuration; + private final Duration defaultReservationTtl; + private final Double defaultLeaseSize; + private final Double lowWaterMark; + private final Duration sweepInterval; + private final Duration prewarmResolveTimeout; + private final JedisPooled redisClient; + private final String redisKeyPrefix; + private final Map overrides; + + private CreditLeaseConfig(Builder builder) { + this.mode = builder.mode != null ? builder.mode : CreditLeaseMode.AUTO; + this.defaultLeaseDuration = builder.defaultLeaseDuration; + this.defaultReservationTtl = builder.defaultReservationTtl; + this.defaultLeaseSize = builder.defaultLeaseSize; + this.lowWaterMark = builder.lowWaterMark; + this.sweepInterval = builder.sweepInterval; + this.prewarmResolveTimeout = builder.prewarmResolveTimeout; + this.redisClient = builder.redisClient; + this.redisKeyPrefix = builder.redisKeyPrefix; + this.overrides = Collections.unmodifiableMap(new LinkedHashMap<>(builder.overrides)); + } + + public static Builder builder() { + return new Builder(); + } + + public CreditLeaseMode getMode() { + return mode; + } + + public Duration getDefaultLeaseDuration() { + return defaultLeaseDuration; + } + + public Duration getDefaultReservationTtl() { + return defaultReservationTtl; + } + + public Double getDefaultLeaseSize() { + return defaultLeaseSize; + } + + public Double getLowWaterMark() { + return lowWaterMark; + } + + public Duration getSweepInterval() { + return sweepInterval; + } + + public Duration getPrewarmResolveTimeout() { + return prewarmResolveTimeout; + } + + public JedisPooled getRedisClient() { + return redisClient; + } + + public String getRedisKeyPrefix() { + return redisKeyPrefix; + } + + public Map getOverrides() { + return overrides; + } + + /** The knobs for one credit type: its override wins, then this config, then the default. */ + public ResolvedLeaseConfig resolve(String creditTypeId) { + CreditLeaseOverride override = overrides.get(creditTypeId); + Duration leaseDuration = firstNonNull( + override != null ? override.getDefaultLeaseDuration() : null, + defaultLeaseDuration, + CreditLeaseDefaults.LEASE_DURATION); + Duration reservationTtl = firstNonNull( + override != null ? override.getDefaultReservationTtl() : null, + defaultReservationTtl, + CreditLeaseDefaults.RESERVATION_TTL); + Double leaseSize = firstNonNull( + override != null ? override.getDefaultLeaseSize() : null, + defaultLeaseSize, + CreditLeaseDefaults.LEASE_SIZE); + Double resolvedWaterMark = firstNonNull( + override != null ? override.getLowWaterMark() : null, lowWaterMark, CreditLeaseDefaults.LOW_WATER_MARK); + return new ResolvedLeaseConfig(leaseDuration, reservationTtl, leaseSize, resolvedWaterMark); + } + + private static T firstNonNull(T override, T configured, T fallback) { + if (override != null) { + return override; + } + return configured != null ? configured : fallback; + } + + public static final class Builder { + + private CreditLeaseMode mode; + private Duration defaultLeaseDuration; + private Duration defaultReservationTtl; + private Double defaultLeaseSize; + private Double lowWaterMark; + private Duration sweepInterval; + private Duration prewarmResolveTimeout; + private JedisPooled redisClient; + private String redisKeyPrefix; + private final Map overrides = new LinkedHashMap<>(); + + /** Where the credit hold lives. Defaults to {@link CreditLeaseMode#AUTO}. */ + public Builder mode(CreditLeaseMode mode) { + this.mode = mode; + return this; + } + + /** Lease lifetime requested at acquire and extend. Defaults to 5 minutes. */ + public Builder defaultLeaseDuration(Duration defaultLeaseDuration) { + this.defaultLeaseDuration = defaultLeaseDuration; + return this; + } + + /** + * Reservation lifetime. Defaults to 60 seconds. In server mode it is capped at an hour + * less a minute of room for clock skew, since an hour out is the furthest the API will + * hold credits and it measures that against its own clock. Size it above the longest + * expected gap between a check and its settle: a settle arriving after the TTL still + * bills the server but no longer re-debits the local lease, so the local balance reads + * high until the lease rolls over. + */ + public Builder defaultReservationTtl(Duration defaultReservationTtl) { + this.defaultReservationTtl = defaultReservationTtl; + return this; + } + + /** Credits requested per acquire, and the minimum extend tranche. Defaults to 10000. */ + public Builder defaultLeaseSize(double defaultLeaseSize) { + this.defaultLeaseSize = defaultLeaseSize; + return this; + } + + /** + * Fraction of the lease below which the SDK kicks off a background extend. Defaults to + * 0.25. + */ + public Builder lowWaterMark(double lowWaterMark) { + this.lowWaterMark = lowWaterMark; + return this; + } + + /** Expired-reservation sweep cadence. Defaults to 1 second. */ + public Builder sweepInterval(Duration sweepInterval) { + this.sweepInterval = sweepInterval; + return this; + } + + /** + * How long a prewarm waits for a freshly identified company to surface in the datastream + * cache. Zero skips the wait, so a prewarm gives up unless the company is already cached. + * Defaults to 5 seconds. + */ + public Builder prewarmResolveTimeout(Duration prewarmResolveTimeout) { + this.prewarmResolveTimeout = prewarmResolveTimeout; + return this; + } + + /** + * A pre-connected Redis client for lease and reservation state. Optional: without one the + * SDK reuses the datastream cache's Redis client, so an existing Redis setup backs leases + * automatically. Set this only to point lease state at a different Redis. + * + *

With no Redis at all the SDK falls back to per-process in-memory stores, which gate + * one process only, and says so in a warning. + */ + public Builder redisClient(JedisPooled redisClient) { + this.redisClient = redisClient; + return this; + } + + /** Key prefix for lease state. Falls back to the datastream cache's, then {@code schematic:}. */ + public Builder redisKeyPrefix(String redisKeyPrefix) { + this.redisKeyPrefix = redisKeyPrefix; + return this; + } + + /** Overrides the resolvable knobs for one credit type. */ + public Builder override(String creditTypeId, CreditLeaseOverride override) { + this.overrides.put(creditTypeId, override); + return this; + } + + public Builder overrides(Map overrides) { + if (overrides != null) { + this.overrides.putAll(overrides); + } + return this; + } + + public CreditLeaseConfig build() { + // Caught here rather than at the first check: a lease sized NaN or a water mark above + // one turns every later comparison into a silent no-op, and the symptom surfaces as + // checks that never gate rather than as the misconfiguration it is. + positiveAmount(defaultLeaseSize, "defaultLeaseSize"); + fraction(lowWaterMark, "lowWaterMark"); + positiveDuration(defaultLeaseDuration, "defaultLeaseDuration"); + positiveDuration(defaultReservationTtl, "defaultReservationTtl"); + positiveDuration(sweepInterval, "sweepInterval"); + for (Map.Entry entry : overrides.entrySet()) { + CreditLeaseOverride override = entry.getValue(); + if (override == null) { + continue; + } + String where = " for credit type " + entry.getKey(); + positiveAmount(override.getDefaultLeaseSize(), "defaultLeaseSize" + where); + fraction(override.getLowWaterMark(), "lowWaterMark" + where); + positiveDuration(override.getDefaultLeaseDuration(), "defaultLeaseDuration" + where); + positiveDuration(override.getDefaultReservationTtl(), "defaultReservationTtl" + where); + } + return new CreditLeaseConfig(this); + } + + private static void positiveAmount(Double value, String name) { + if (value == null) { + return; + } + if (Double.isNaN(value) || Double.isInfinite(value) || value <= 0) { + throw new IllegalArgumentException(name + " must be a positive finite number, got " + value); + } + } + + private static void fraction(Double value, String name) { + if (value == null) { + return; + } + if (Double.isNaN(value) || value <= 0 || value >= 1) { + throw new IllegalArgumentException( + name + " must be a fraction between 0 and 1, exclusive, got " + value); + } + } + + private static void positiveDuration(Duration value, String name) { + if (value == null) { + return; + } + if (value.isZero() || value.isNegative()) { + throw new IllegalArgumentException(name + " must be a positive duration, got " + value); + } + } + } +} diff --git a/src/main/java/com/schematic/api/credits/CreditLeaseDefaults.java b/src/main/java/com/schematic/api/credits/CreditLeaseDefaults.java new file mode 100644 index 0000000..11e0b0b --- /dev/null +++ b/src/main/java/com/schematic/api/credits/CreditLeaseDefaults.java @@ -0,0 +1,74 @@ +package com.schematic.api.credits; + +import java.time.Duration; + +/** + * Defaults shared by {@link CreditLeaseConfig}, {@link CreditLeaseManager} and the check flows. + * + *

The figures match the other Schematic SDKs: the semantics are pinned by the + * language-agnostic vectors in {@code conformance/}, and a fleet mixing SDKs shares one Redis. + */ +public final class CreditLeaseDefaults { + + /** Lease lifetime requested at acquire and extend: {@code expiresAt = now + duration}. */ + public static final Duration LEASE_DURATION = Duration.ofMinutes(5); + + /** + * Reservation lifetime, and the sweep deadline. Size it above the longest expected gap + * between a check and its settle: a settle arriving after the TTL still bills the server but + * no longer re-debits the lease, so the local balance reads high until the lease rolls over. + */ + public static final Duration RESERVATION_TTL = Duration.ofSeconds(60); + + /** + * The furthest out the API will hold credits, so a larger server-mode reservation TTL would + * fail every check. + */ + public static final Duration MAX_RESERVATION_TTL = Duration.ofHours(1); + + /** + * Held back from {@link #MAX_RESERVATION_TTL} because the API measures that hour against its + * own clock while the SDK computes {@code expiresAt} against the caller's: a client running + * ahead would otherwise be rejected at exactly the cap. + */ + public static final Duration RESERVATION_TTL_SKEW_ALLOWANCE = Duration.ofSeconds(60); + + /** Credits requested per acquire, and the minimum extend tranche. */ + public static final double LEASE_SIZE = 10_000d; + + /** Remaining/granted ratio at or below which a background extend is kicked off. */ + public static final double LOW_WATER_MARK = 0.25d; + + /** Cadence of the expired-reservation sweep. */ + public static final Duration SWEEP_INTERVAL = Duration.ofSeconds(1); + + /** + * How long a prewarm waits for a freshly identified company to surface in the datastream + * cache before giving up. Long enough to cover the buffer flush, server ingest and datastream + * push for a brand-new company; short enough that a misconfigured caller does not hang. + */ + public static final Duration PREWARM_RESOLVE_TIMEOUT = Duration.ofSeconds(5); + + /** Gap between prewarm resolve attempts. */ + public static final Duration PREWARM_POLL_INTERVAL = Duration.ofMillis(100); + + /** + * How long {@code 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 + * {@link #LEASE_DURATION}. + */ + public static final Duration SHUTDOWN_DRAIN_TIMEOUT = Duration.ofSeconds(5); + + /** Namespace for every lease and reservation key. */ + public static final String KEY_PREFIX = "schematic:"; + + /** + * The balance a fail-open evaluation substitutes for the metered credit: large enough that + * the credit gate always passes, and still exact as a JSON number, so the engine reads back + * what the SDK sent. This is Node's {@code Number.MAX_SAFE_INTEGER}, which the conformance + * vectors name {@code max_safe_integer}. + */ + public static final double FAIL_OPEN_BALANCE = 9007199254740991d; + + private CreditLeaseDefaults() {} +} diff --git a/src/main/java/com/schematic/api/credits/CreditLeaseManager.java b/src/main/java/com/schematic/api/credits/CreditLeaseManager.java new file mode 100644 index 0000000..95df327 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/CreditLeaseManager.java @@ -0,0 +1,766 @@ +package com.schematic.api.credits; + +import com.schematic.api.logger.SchematicLogger; +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * Owns lease rows for one client: acquire on first use or after expiry, extend when the local + * view dips below the water mark, release on close. + * + *

Acquire and extend each get their own best-effort single-flight map keyed by slot. + * Best-effort because callers racing ahead of the registration can still issue duplicate wire + * calls, which is safe: the server is idempotent for an active slot, {@link LeaseStore#replace} + * keeps the first live lease, and {@link LeaseStore#extend} reconciles to a total. + * + *

Every path here resolves rather than throws: callers route a missing lease through their + * fail-open or fail-closed handling, and several calls are fire-and-forget, where an exception + * has nowhere to go. + */ +public final class CreditLeaseManager implements AutoCloseable { + + /** + * How many in-flight extends one caller waits 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. + */ + private static final int MAX_EXTEND_JOINS = 2; + + private final LeaseWireClient wire; + private final LeaseStore leases; + private final ReservationStore reservations; + private final CreditLeaseConfig config; + private final SchematicLogger logger; + private final Clock clock; + private final Duration sweepInterval; + + // Kept separate so an in-flight extend can never satisfy an acquire, or the other way round. + private final ConcurrentHashMap acquireFlights = new ConcurrentHashMap<>(); + private final ConcurrentHashMap extendFlights = new ConcurrentHashMap<>(); + // Lease work nobody waits on: the redundant release a lost acquire race issues, and the + // background extends checks fire and forget. drain() waits these out so a close releases what + // they installed. + private final Set> background = ConcurrentHashMap.newKeySet(); + + private final ExecutorService executor; + private final ScheduledExecutorService sweeper; + private volatile boolean stopped; + // Held across the flag write in stop() and the re-check an acquire makes once it owns the + // slot's flight, which is what stops a lease landing in a slot close() has already swept. + private final Object stopLock = new Object(); + // Compare-and-set rather than a read then a write: two threads starting the sweep together + // would both pass a plain check and schedule a second sweeper onto the same store. + private final AtomicBoolean sweeping = new AtomicBoolean(); + + public CreditLeaseManager( + LeaseWireClient wire, + LeaseStore leases, + ReservationStore reservations, + CreditLeaseConfig config, + SchematicLogger logger, + Clock clock) { + this.wire = wire; + this.leases = leases; + this.reservations = reservations; + this.config = config != null ? config : CreditLeaseConfig.builder().build(); + this.logger = logger; + this.clock = clock != null ? clock : Clock.systemUTC(); + this.sweepInterval = this.config.getSweepInterval() != null + ? this.config.getSweepInterval() + : CreditLeaseDefaults.SWEEP_INTERVAL; + this.executor = Executors.newCachedThreadPool(daemonThreads("SchematicCreditLease")); + this.sweeper = Executors.newSingleThreadScheduledExecutor(daemonThreads("SchematicCreditLeaseSweep")); + } + + /** The lease knobs for one credit type. */ + public ResolvedLeaseConfig resolveConfig(String creditTypeId) { + return config.resolve(creditTypeId); + } + + /** + * Returns the slot's live lease, acquiring one over the wire if none is live. Returns null + * rather than throwing when the wire or the store is down, so the caller routes the outcome + * through fail-open or fail-closed. + */ + public LeaseState acquireIfNeeded(String companyId, String creditTypeId) { + return acquireIfNeeded(companyId, creditTypeId, null); + } + + /** + * Acquires under the caller's per-check timeout. The flight is shared, so the first caller's + * timeout governs everyone who joins it; a background caller passes null and takes the + * client's own. + */ + public LeaseState acquireIfNeeded(String companyId, String creditTypeId, Duration timeout) { + if (stopped) { + // A lease installed after releaseAllLocalLeases has listed the slots would be held + // until it expires server-side, with nobody left to release it. + debug("Not acquiring a credit lease for " + companyId + "/" + creditTypeId + ": the manager is stopped"); + return null; + } + LeaseState existing; + try { + existing = leases.get(companyId, creditTypeId); + } catch (RuntimeException e) { + error("Failed to read lease store for " + companyId + "/" + creditTypeId + ": " + e); + return null; + } + // Liveness here is judged on this process's clock, while the Redis store re-reads expiry + // against the Redis server's clock inside tryReserve. The two can disagree, so a lease + // this call hands back can still be refused there, and the check routes that through its + // fail-open handling. The stores keep an expired row for a grace window precisely so + // clocks within it agree on what is live. + if (existing != null && existing.isLiveAt(now())) { + return existing; + } + // An expired or absent slot is left for replace to overwrite: it guards on expiry and + // writes atomically. Dropping the stale row first would be a separate, non-atomic op that + // can interleave between a sibling's read and its replace, clobbering a lease that + // sibling just installed. Reading a stale entry in the gap is harmless, since every path + // that acts on a lease re-guards on expiry. + if (stopped) { + debug("Not acquiring a credit lease for " + companyId + "/" + creditTypeId + ": the manager is stopped"); + return null; + } + + String key = LeaseStore.leaseKey(companyId, creditTypeId); + Flight joined = acquireFlights.get(key); + if (joined != null) { + return joined.await(timeout); + } + Flight flight = new Flight(0); + Flight raced = acquireFlights.putIfAbsent(key, flight); + if (raced != null) { + return raced.await(timeout); + } + try { + boolean stoppedInTheGap; + synchronized (stopLock) { + // Under the lock stop() takes, so this either sees the stop or provably ran + // before it. The check at the top of the method can go stale between there and + // here, and a lease installed past that point is one close() has already finished + // looking for. + stoppedInTheGap = stopped; + } + if (stoppedInTheGap) { + debug("Not acquiring a credit lease for " + companyId + "/" + creditTypeId + + ": the manager stopped while the flight was being registered"); + return null; + } + LeaseState result = acquire(companyId, creditTypeId, timeout); + flight.result.complete(result); + return result; + } catch (RuntimeException e) { + error("Failed to acquire credit lease for " + companyId + "/" + creditTypeId + ": " + e); + return null; + } finally { + // Completing here and not only on the two paths above: an Error unwinding past both + // would leave the future unfinished, and every joiner parks on it forever. A no-op + // once the success path has already completed it. + flight.result.complete(null); + acquireFlights.remove(key, flight); + } + } + + private LeaseState acquire(String companyId, String creditTypeId, Duration timeout) { + ResolvedLeaseConfig resolved = resolveConfig(creditTypeId); + LeaseGrant grant; + try { + grant = wire.acquire( + companyId, creditTypeId, resolved.getLeaseSize(), now().plus(resolved.getLeaseDuration()), timeout); + } catch (RuntimeException e) { + error("Failed to acquire credit lease for " + companyId + "/" + creditTypeId + ": " + e); + return null; + } + boolean wrote; + try { + wrote = leases.replace(new LeaseGrant( + grant.getLeaseId(), + orElse(grant.getCompanyId(), companyId), + orElse(grant.getCreditTypeId(), creditTypeId), + grant.getGrantedAmount(), + grant.getExpiresAt())); + } catch (RuntimeException e) { + error("Failed to install credit lease " + grant.getLeaseId() + ": " + e); + return null; + } + + LeaseState current; + try { + current = leases.get(companyId, creditTypeId); + } catch (RuntimeException e) { + error("Failed to read lease store for " + companyId + "/" + creditTypeId + ": " + e); + return null; + } + if (wrote) { + debug("Acquired credit lease " + grant.getLeaseId() + " for " + companyId + "/" + creditTypeId + + " (granted=" + grant.getGrantedAmount() + ", expires=" + grant.getExpiresAt() + ")"); + return current; + } + + // A sibling holds the slot with a live lease, or the slot's expired row was reconciled in + // place. The server is idempotent for an active slot, so a racing acquire is normally + // handed back the SAME lease the sibling installed, and releasing it would pull the + // shared lease out from under every process drawing on it. Only a different lease is a + // redundant hold nobody will draw on, so only that one is released. An empty slot + // (expired in the gap) releases nothing either: this lease is likely what the next + // acquire is handed. + if (current != null && !current.getLeaseId().equals(grant.getLeaseId())) { + debug("Lost acquire race for " + companyId + "/" + creditTypeId + "; releasing redundant lease " + + grant.getLeaseId()); + // Tracked even once the manager is stopping: this lease is already granted and nobody + // will draw on it, so refusing the release would hold its credits until the server + // expires them. The drain waits it out within its own bound. + spawn( + () -> { + try { + wire.release(grant.getLeaseId()); + } catch (RuntimeException e) { + warn("Failed to release redundant credit lease " + grant.getLeaseId() + ": " + e); + } + }, + true); + } else { + debug("Lost acquire race for " + companyId + "/" + creditTypeId + "; the server returned the installed " + + "lease " + grant.getLeaseId() + ", nothing to release"); + } + return current; + } + + /** + * Extends the slot's lease when the local view warrants it, triggered by either the + * low-water-mark ratio (steady-state refresh) or a {@code requiredCredits} hint above the + * local remaining (a check just failed a reserve of that size). Pass null for + * {@code requiredCredits} to ask for the steady-state check only. + * + *

A caller arriving while an extend is in flight joins it. If its own shortfall is larger + * than what that extend asked for, it waits the flight out and then issues exactly one + * follow-up extend for the remaining difference; otherwise it would inherit a tranche-sized + * ask and fail its post-extend retry with credits still sitting on the server. + */ + public LeaseState maybeExtend(String companyId, String creditTypeId, Double requiredCredits) { + return maybeExtend(companyId, creditTypeId, requiredCredits, null); + } + + /** Extends under the caller's per-check timeout, or the client's own when null. */ + public LeaseState maybeExtend(String companyId, String creditTypeId, Double requiredCredits, Duration timeout) { + return maybeExtend(companyId, creditTypeId, requiredCredits, true, timeout); + } + + private LeaseState maybeExtend( + String companyId, String creditTypeId, Double requiredCredits, boolean joinInFlight, Duration timeout) { + if (stopped) { + // Extending past stop re-holds credits on a lease the close is about to release, or + // has already released. + debug("Not extending a credit lease for " + companyId + "/" + creditTypeId + ": the manager is stopped"); + return null; + } + String key = LeaseStore.leaseKey(companyId, creditTypeId); + // Joins are budgeted, extends of this caller's own are not: it waits out flights that ask + // for too little, but once the budget is spent it sends one extend of its own rather than + // joining again. Without the budget a caller could queue behind an unbounded run of other + // callers' follow-ups; without the extend of its own it would hand back a balance it + // already knows is short and fail its retry with credits still sitting on the server. + for (int joinsLeft = MAX_EXTEND_JOINS; ; joinsLeft--) { + LeaseState entry; + try { + entry = leases.get(companyId, creditTypeId); + } catch (RuntimeException e) { + warn("Failed to read lease store for " + companyId + "/" + creditTypeId + ": " + e); + return null; + } + if (entry == null) { + return null; + } + // Never extend an expired lease: the server treats it as released and has already + // refunded its remainder, so the only correct move is a fresh acquire on the next + // check. + if (!entry.isLiveAt(now())) { + return null; + } + ResolvedLeaseConfig resolved = resolveConfig(creditTypeId); + boolean belowWatermark = atOrBelowWatermark(entry, resolved); + boolean belowRequired = requiredCredits != null && entry.getLocalRemainingCredits() < requiredCredits; + if (!belowWatermark && !belowRequired) { + return entry; + } + // Size the extend to cover the request that triggered it: a single check needing more + // than remaining plus one tranche would otherwise fail its post-extend retry forever, + // however much balance the server has. The steady-state path keeps asking for the + // configured tranche. Sized here, one level above the wire call, so the flight + // registered below and the request body provably carry the same number for a joiner + // to compare against. + double shortfall = requiredCredits != null ? requiredCredits - entry.getLocalRemainingCredits() : 0; + double additionalAmount = Math.max(resolved.getLeaseSize(), shortfall); + + Flight inFlight = extendFlights.get(key); + if (inFlight != null && joinsLeft > 0) { + if (!joinInFlight) { + // The slot is already being topped up and nobody is waiting on this call's + // result, so parking on that flight would hold a pool thread for a wire call + // whose outcome this caller does not read. + return null; + } + LeaseState joined = inFlight.await(timeout); + if (!inFlight.isDone()) { + // The wait, not the flight, ran out of time. The extend runs on for everybody + // still on it, and reporting no lease sends this caller down the fail-open or + // fail-closed path its own timeout asked for. + debug("An extend in flight for " + companyId + "/" + creditTypeId + " outlasted the caller's " + + "timeout; not waiting on it"); + return null; + } + // The flight asked for at least what we need, which covers every watermark-driven + // joiner and any check the tranche fits. One wire call serves all of them, which + // is the point of single-flight. A flight that sent nothing covers nobody, so its + // ask does not stand in for ours. + if (inFlight.sentExtend && additionalAmount <= inFlight.requestedAdditional) { + return joined; + } + // It asked for less than we need. Go round to re-read the slot it just moved, so + // the next ask is sized against the balance it left rather than the one this call + // started from. + continue; + } + return startExtend( + key, companyId, creditTypeId, entry, resolved, requiredCredits, additionalAmount, timeout); + } + } + + /** + * Registers this call as the slot's flight and sends its extend. The registration overwrites + * rather than yields: a caller arriving here has spent its joins on the flight it would be + * overwriting, so yielding to that flight again is the one thing it must not do. + * Deregistration is identity-guarded, so the overwritten flight cannot evict this one on its + * way out. + */ + private LeaseState startExtend( + String key, + String companyId, + String creditTypeId, + LeaseState entry, + ResolvedLeaseConfig resolved, + Double requiredCredits, + double additionalAmount, + Duration timeout) { + Flight flight = new Flight(additionalAmount); + extendFlights.put(key, flight); + try { + // Re-read now that the slot's flight is ours. The row above was read before the + // flight check, so a previous extend can have landed and deregistered in between: + // that read says "below the mark" about a lease that has since been topped up, and + // sending on it bills a second tranche nobody needs. + LeaseState fresh = stillNeedsExtending(companyId, creditTypeId, requiredCredits, resolved); + if (fresh == null) { + return leases.get(companyId, creditTypeId); + } + flight.sentExtend = true; + LeaseState result = extend(fresh, resolved, additionalAmount, timeout); + flight.result.complete(result); + return result; + } catch (RuntimeException e) { + warn("Failed to extend credit lease " + entry.getLeaseId() + ": " + e); + return null; + } finally { + // Completing here and not only on the two paths above: an Error unwinding past both + // would leave the future unfinished, and every joiner parks on it forever. A no-op + // once the success path has already completed it. + flight.result.complete(null); + // Identity-guarded rather than an unconditional remove: a caller that spent its joins + // registers a flight of its own for the same key, and this one must not evict it. + extendFlights.remove(key, flight); + } + } + + /** + * Kicks off a water-mark extend without waiting for it: a check that just drew the lease down + * should not pay for the top-up. + */ + public void extendInBackground(String companyId, String creditTypeId) { + // Tested here, on the caller's thread, rather than inside the spawned task: every allowed + // check calls this, and a lease sitting comfortably above its water mark is the common + // case. Spawning first would queue a task per check onto an unbounded pool only to + // discover there was nothing to do. + // + // The flight is tested first, and for the same reason. A slot stays below its water mark + // for as long as the top-up is on the wire, so every check allowed in that window would + // otherwise spawn a task that reads the store, finds the flight it must not join, and + // returns having done nothing. + if (extendFlights.containsKey(LeaseStore.leaseKey(companyId, creditTypeId))) { + return; + } + if (!extendIsDue(companyId, creditTypeId)) { + return; + } + spawn(() -> maybeExtend(companyId, creditTypeId, null, false, null)); + } + + /** Whether the slot's lease has drawn down far enough to warrant a steady-state top-up. */ + private boolean extendIsDue(String companyId, String creditTypeId) { + if (stopped) { + return false; + } + LeaseState entry; + try { + entry = leases.get(companyId, creditTypeId); + } catch (RuntimeException e) { + warn("Failed to read lease store for " + companyId + "/" + creditTypeId + ": " + e); + return false; + } + return entry != null && entry.isLiveAt(now()) && atOrBelowWatermark(entry, resolveConfig(creditTypeId)); + } + + /** + * The slot's row if it still warrants the extend the caller sized, null if it no longer does. + * Read after winning the flight, so it reflects any extend that landed while this caller was + * deciding. + */ + private LeaseState stillNeedsExtending( + String companyId, String creditTypeId, Double requiredCredits, ResolvedLeaseConfig resolved) { + LeaseState fresh = leases.get(companyId, creditTypeId); + if (fresh == null || !fresh.isLiveAt(now())) { + return null; + } + boolean belowRequired = requiredCredits != null && fresh.getLocalRemainingCredits() < requiredCredits; + return atOrBelowWatermark(fresh, resolved) || belowRequired ? fresh : null; + } + + private static boolean atOrBelowWatermark(LeaseState entry, ResolvedLeaseConfig resolved) { + double ratio = entry.getLocalRemainingCredits() / Math.max(entry.getGrantedAmount(), 1); + return ratio <= resolved.getLowWaterMark(); + } + + private LeaseState extend( + LeaseState entry, ResolvedLeaseConfig resolved, double additionalAmount, Duration timeout) { + LeaseGrant grant; + try { + grant = wire.extend(entry.getLeaseId(), additionalAmount, now().plus(resolved.getLeaseDuration()), timeout); + } catch (RuntimeException e) { + warn("Failed to extend credit lease " + entry.getLeaseId() + ": " + e); + return null; + } + try { + // Reconcile to the server's authoritative TOTAL, with the store computing the delta + // against its own current total: per-process single-flight does not cover sibling + // processes. Pinned to the lease the server extended, so an expiry mid-call cannot + // mint the delta onto a successor. + leases.extend( + entry.getCompanyId(), + entry.getCreditTypeId(), + grant.getGrantedAmount(), + grant.getExpiresAt(), + entry.getLeaseId()); + debug("Extended credit lease " + entry.getLeaseId() + " to " + grant.getGrantedAmount() + " (was " + + entry.getGrantedAmount() + " at last read, expires " + grant.getExpiresAt() + ")"); + return leases.get(entry.getCompanyId(), entry.getCreditTypeId()); + } catch (RuntimeException e) { + warn("Failed to reconcile extended credit lease " + entry.getLeaseId() + ": " + e); + return null; + } + } + + /** + * Releases every live lease this process exclusively holds, returning their unspent + * remainders to the company balance immediately instead of waiting out the lease expiry. + * + *

Only a per-process store implements {@link LeaseLister}; a shared backend is skipped, + * since sibling processes still draw on those leases. Expired leases are skipped too: the + * server already swept them. Best-effort, with failures falling back to server-side expiry. + */ + public void releaseAllLocalLeases() { + releaseAllLocalLeases(CreditLeaseDefaults.SHUTDOWN_DRAIN_TIMEOUT); + } + + /** + * Releases every live lease this process exclusively holds, within {@code budget}. + * + *

The releases go out together and the budget bounds the whole set, not each one in turn. + * Issued serially, a process holding many slots behind a slow server would stretch a shutdown + * by the sum of them, and a single hung release would spend the entire budget on its own and + * abandon every lease behind it. Whatever has not landed by the deadline is left to + * server-side expiry, which is where a failed release leaves it too. + */ + public void releaseAllLocalLeases(Duration budget) { + if (!(leases instanceof LeaseLister)) { + return; + } + List entries; + try { + entries = ((LeaseLister) leases).list(); + } catch (RuntimeException e) { + warn("Failed to enumerate leases on close: " + e); + return; + } + long deadline = System.nanoTime() + Math.max(0, budget.toNanos()); + Instant now = now(); + List> releases = new ArrayList<>(); + for (LeaseState entry : entries) { + if (!entry.isLiveAt(now)) { + continue; + } + CompletableFuture released = new CompletableFuture<>(); + try { + executor.execute(() -> { + try { + wire.release(entry.getLeaseId()); + leases.drop(entry.getCompanyId(), entry.getCreditTypeId()); + debug("Released credit lease " + entry.getLeaseId() + " on close"); + } catch (RuntimeException e) { + warn("Failed to release credit lease " + entry.getLeaseId() + " on close (it will expire " + + "server-side): " + e); + } finally { + released.complete(null); + } + }); + releases.add(released); + } catch (RejectedExecutionException e) { + debug("Credit lease executor is shut down; leaving " + entry.getLeaseId() + " to server-side expiry"); + } + } + if (!awaitAll(releases, deadline)) { + warn("Ran out of shutdown budget releasing credit leases; any still held will expire server-side"); + } + } + + /** + * Runs the expired-reservation sweep on the configured interval. Safe to call twice; a no-op + * without a reservation store or after {@link #stop()}. + */ + public void startSweep() { + if (reservations == null || stopped || !sweeping.compareAndSet(false, true)) { + return; + } + long interval = Math.max(1, sweepInterval.toMillis()); + sweeper.scheduleWithFixedDelay( + () -> { + try { + reservations.sweepExpired(); + } catch (RuntimeException e) { + // Keep the loop alive: a sweep failure is transient (a Redis blip), and + // the next tick retries. + debug("Reservation sweep failed: " + e); + } + }, + interval, + interval, + TimeUnit.MILLISECONDS); + } + + /** + * Refuses new lease work. Idempotent, and paired with {@link #drain}: stopping first is what + * makes the drain terminate, since nothing can queue behind it. + */ + public void stop() { + synchronized (stopLock) { + stopped = true; + } + sweeper.shutdownNow(); + } + + /** + * Waits 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 {@code timeout} is abandoned + * rather than stalling the caller's shutdown, and the credits it holds fall back to + * server-side expiry. + */ + public void drain(Duration timeout) { + long deadline = System.nanoTime() + Math.max(0, timeout.toNanos()); + while (true) { + List> pending = new ArrayList<>(background); + for (Flight flight : acquireFlights.values()) { + pending.add(flight.result); + } + for (Flight flight : extendFlights.values()) { + pending.add(flight.result); + } + if (pending.isEmpty()) { + return; + } + // Settling one round can queue another (an acquire that loses its race fires a + // release), so keep going until nothing is left. + if (!awaitAll(pending, deadline)) { + warn("Timed out after " + timeout.toMillis() + "ms draining in-flight credit lease work; any " + + "credits it holds will be released by server-side expiry"); + return; + } + } + } + + private boolean awaitAll(List> pending, long deadlineNanos) { + for (CompletableFuture future : pending) { + if (!await(future, deadlineNanos)) { + return false; + } + } + return true; + } + + private boolean await(CompletableFuture future, long deadlineNanos) { + long remaining = deadlineNanos - System.nanoTime(); + if (remaining <= 0) { + return false; + } + try { + future.get(remaining, TimeUnit.NANOSECONDS); + return true; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return false; + } catch (TimeoutException e) { + return false; + } catch (RuntimeException | ExecutionException e) { + // A failed background step has already logged; the drain only cares that it landed. + return true; + } + } + + /** Stops the manager and drains what it has in flight. Leases are released by the caller. */ + @Override + public void close() { + close(CreditLeaseDefaults.SHUTDOWN_DRAIN_TIMEOUT); + } + + /** + * Stops the manager and drains what it has in flight within {@code budget}. Leases are + * released by the caller. + * + *

A caller closing several components under one deadline passes what is left of it, so the + * bound it promised is not reset to a full drain timeout here. + */ + public void close(Duration budget) { + stop(); + drain(budget); + executor.shutdown(); + } + + /** Runs a fire-and-forget step, refused once the manager is stopped. */ + private void spawn(Runnable step) { + spawn(step, false); + } + + /** + * Runs a fire-and-forget step. {@code afterStop} keeps work that has to happen even once the + * manager is stopping, which is what the drain is there to wait out; everything else is + * refused after {@link #stop()}, where it would touch a manager being torn down. Nothing here + * lets an exception escape: these paths are unawaited, so there is nobody to catch for them. + */ + private void spawn(Runnable step, boolean afterStop) { + if (stopped && !afterStop) { + return; + } + CompletableFuture landed = new CompletableFuture<>(); + background.add(landed); + try { + executor.execute(() -> { + try { + step.run(); + } catch (RuntimeException e) { + error("Background credit lease work failed: " + e); + } finally { + landed.complete(null); + background.remove(landed); + } + }); + } catch (RejectedExecutionException e) { + background.remove(landed); + debug("Credit lease executor is shut down; skipping background work"); + } + } + + private Instant now() { + return clock.instant(); + } + + private static String orElse(String value, String fallback) { + return value != null && !value.isEmpty() ? value : fallback; + } + + private static ThreadFactory daemonThreads(String name) { + return new ThreadFactory() { + @Override + public Thread newThread(Runnable runnable) { + Thread thread = new Thread(runnable, name); + thread.setDaemon(true); + return thread; + } + }; + } + + private void debug(String message) { + if (logger != null) { + logger.debug(message); + } + } + + private void warn(String message) { + if (logger != null) { + logger.warn(message); + } + } + + private void error(String message) { + if (logger != null) { + logger.error(message); + } + } + + /** One in-flight wire call for a slot, and the amount its extend asked the server for. */ + private final class Flight { + + private final double requestedAdditional; + private final CompletableFuture result = new CompletableFuture<>(); + // What the flight asked for only bounds a joiner's shortfall if the flight went out at + // all. A flight that re-read the slot and found the extend unnecessary sends nothing, and + // a joiner holding that as its answer would deny a check whose credits are still on the + // server. + private volatile boolean sentExtend; + + Flight(double requestedAdditional) { + this.requestedAdditional = requestedAdditional; + } + + boolean isDone() { + return result.isDone(); + } + + /** + * The flight's answer, waited for no longer than the joining caller's own deadline. The + * flight is shared, so a check that joins one someone else started would otherwise inherit + * a stranger's wire call and blow its timeout by however long that call runs. Abandoning + * the wait leaves the flight running for whoever else is on it, and the caller takes the + * failure path its mode chooses, the same as any other unresolved lease. Null timeout + * means the caller brought no deadline of its own. + */ + LeaseState await(Duration timeout) { + try { + return timeout == null + ? result.get() + : result.get(Math.max(1, timeout.toMillis()), TimeUnit.MILLISECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return null; + } catch (ExecutionException e) { + return null; + } catch (TimeoutException e) { + debug("Gave up waiting " + timeout.toMillis() + "ms on an in-flight credit lease call; it " + + "continues for the callers still on it"); + return null; + } + } + } +} diff --git a/src/main/java/com/schematic/api/credits/CreditLeaseMode.java b/src/main/java/com/schematic/api/credits/CreditLeaseMode.java new file mode 100644 index 0000000..ff5e822 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/CreditLeaseMode.java @@ -0,0 +1,20 @@ +package com.schematic.api.credits; + +/** + * Where a credit hold lives for a {@code check()} that passes usage. + */ +public enum CreditLeaseMode { + /** + * Local leases over DataStream: the SDK draws a tranche of credits up front and carves + * reservations out of it locally. Requires DataStream, and a shared Redis backend to gate + * across processes. + */ + CLIENT, + /** + * One check-and-reserve API call per check: the server evaluates the flag and takes the hold + * in the same round trip. No DataStream, no Redis, no local stores. + */ + SERVER, + /** Client mode when DataStream is enabled, server mode otherwise. The default. */ + AUTO +} diff --git a/src/main/java/com/schematic/api/credits/CreditLeaseOverride.java b/src/main/java/com/schematic/api/credits/CreditLeaseOverride.java new file mode 100644 index 0000000..0ec4581 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/CreditLeaseOverride.java @@ -0,0 +1,74 @@ +package com.schematic.api.credits; + +import java.time.Duration; + +/** + * Overrides the four resolvable lease knobs for one credit type. An unset field leaves the + * client-wide value in place. + */ +public final class CreditLeaseOverride { + + private final Duration defaultLeaseDuration; + private final Duration defaultReservationTtl; + private final Double defaultLeaseSize; + private final Double lowWaterMark; + + private CreditLeaseOverride(Builder builder) { + this.defaultLeaseDuration = builder.defaultLeaseDuration; + this.defaultReservationTtl = builder.defaultReservationTtl; + this.defaultLeaseSize = builder.defaultLeaseSize; + this.lowWaterMark = builder.lowWaterMark; + } + + public static Builder builder() { + return new Builder(); + } + + public Duration getDefaultLeaseDuration() { + return defaultLeaseDuration; + } + + public Duration getDefaultReservationTtl() { + return defaultReservationTtl; + } + + public Double getDefaultLeaseSize() { + return defaultLeaseSize; + } + + public Double getLowWaterMark() { + return lowWaterMark; + } + + public static final class Builder { + + private Duration defaultLeaseDuration; + private Duration defaultReservationTtl; + private Double defaultLeaseSize; + private Double lowWaterMark; + + public Builder defaultLeaseDuration(Duration defaultLeaseDuration) { + this.defaultLeaseDuration = defaultLeaseDuration; + return this; + } + + public Builder defaultReservationTtl(Duration defaultReservationTtl) { + this.defaultReservationTtl = defaultReservationTtl; + return this; + } + + public Builder defaultLeaseSize(double defaultLeaseSize) { + this.defaultLeaseSize = defaultLeaseSize; + return this; + } + + public Builder lowWaterMark(double lowWaterMark) { + this.lowWaterMark = lowWaterMark; + return this; + } + + public CreditLeaseOverride build() { + return new CreditLeaseOverride(this); + } + } +} diff --git a/src/main/java/com/schematic/api/credits/DataStreamCreditCheckSource.java b/src/main/java/com/schematic/api/credits/DataStreamCreditCheckSource.java new file mode 100644 index 0000000..536e00c --- /dev/null +++ b/src/main/java/com/schematic/api/credits/DataStreamCreditCheckSource.java @@ -0,0 +1,76 @@ +package com.schematic.api.credits; + +import com.schematic.api.datastream.CheckFlagOptions; +import com.schematic.api.datastream.DataStreamClient; +import com.schematic.api.types.RulesengineCheckFlagResult; +import com.schematic.api.types.RulesengineCompany; +import com.schematic.api.types.RulesengineFlag; +import com.schematic.api.types.RulesengineUser; +import java.util.Map; + +/** + * Serves a credit-gated check from the DataStream client. The translation from the flow's + * preflight to the engine's own options lives here, so the flow stays free of the engine's + * surface. + */ +public final class DataStreamCreditCheckSource implements CreditCheckDataStream { + + private final DataStreamClient dataStream; + + public DataStreamCreditCheckSource(DataStreamClient dataStream) { + this.dataStream = dataStream; + } + + @Override + public RulesengineFlag getFlag(String flagKey) { + return dataStream.getCachedFlag(flagKey); + } + + @Override + public RulesengineCompany getCompany(Map keys) { + return liveFetchIsPossible() ? dataStream.getCompany(keys) : dataStream.getCachedCompany(keys); + } + + @Override + public RulesengineUser getUser(Map keys) { + return liveFetchIsPossible() ? dataStream.getUser(keys) : dataStream.getCachedUser(keys); + } + + /** + * Whether a cache miss can still be answered over the socket. + * + *

Replicator mode has no socket to ask, and a disconnected client has nothing to send the + * request on, so a live fetch in either state only waits out its own timeout before returning + * nothing. A check would pay that wait per call before falling back to the plain check, which + * bails on the same two states without waiting. + */ + private boolean liveFetchIsPossible() { + return !dataStream.isReplicatorMode() && dataStream.isConnected(); + } + + @Override + public RulesengineCheckFlagResult evaluateFlag( + RulesengineFlag flag, RulesengineCompany company, RulesengineUser user, PreflightOptions preflight) + throws Exception { + return dataStream.evaluateFlagWithOptions(flag, company, user, toEngineOptions(preflight)); + } + + /** Translates the flow's preflight into the engine's options. */ + public static CheckFlagOptions toEngineOptions(PreflightOptions preflight) { + if (preflight == null) { + return null; + } + if (preflight.getCreditCost() != null) { + return CheckFlagOptions.creditCost(preflight.getCreditCost()); + } + if (preflight.getEventUsage() != null) { + return CheckFlagOptions.eventUsage( + preflight.getEventUsage().getEventSubtype(), + preflight.getEventUsage().getQuantity()); + } + if (preflight.getUsage() != null) { + return CheckFlagOptions.usage(preflight.getUsage()); + } + return null; + } +} diff --git a/src/main/java/com/schematic/api/credits/FlagCheckReporter.java b/src/main/java/com/schematic/api/credits/FlagCheckReporter.java new file mode 100644 index 0000000..a9703a8 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/FlagCheckReporter.java @@ -0,0 +1,13 @@ +package com.schematic.api.credits; + +import com.schematic.api.types.EventBodyFlagCheck; + +/** + * Reports a flag_check event for a check the credit flow resolved itself. The plain check paths + * enqueue one per check, so without this a credit-gated check would be invisible to flag-check + * analytics and to company last-seen. Fallback exits do not call it: the plain check they defer to + * reports its own. + */ +public interface FlagCheckReporter { + void report(EventBodyFlagCheck body); +} diff --git a/src/main/java/com/schematic/api/credits/InMemoryLeaseStore.java b/src/main/java/com/schematic/api/credits/InMemoryLeaseStore.java new file mode 100644 index 0000000..ae3aa3c --- /dev/null +++ b/src/main/java/com/schematic/api/credits/InMemoryLeaseStore.java @@ -0,0 +1,240 @@ +package com.schematic.api.credits; + +import java.time.Clock; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Supplier; + +/** + * Keeps lease slots in this process only, so it gates a single process. Swap in + * {@link RedisLeaseStore} to gate across several; both implement {@link LeaseStore}. + */ +public final class InMemoryLeaseStore implements LeaseStore, LeaseLister { + + private final Clock clock; + // Concurrent so list() can snapshot it without taking every slot lock; the compound + // read-modify-write below still runs under the slot's lock. + private final ConcurrentHashMap leases = new ConcurrentHashMap<>(); + // One lock per slot, so a reserve on one company never waits on another's. + private final ConcurrentHashMap locks = new ConcurrentHashMap<>(); + + public InMemoryLeaseStore() { + this(Clock.systemUTC()); + } + + public InMemoryLeaseStore(Clock clock) { + this.clock = clock != null ? clock : Clock.systemUTC(); + } + + @Override + public LeaseState get(String companyId, String creditTypeId) { + String key = LeaseStore.leaseKey(companyId, creditTypeId); + return withLock(key, () -> leases.get(key)); + } + + @Override + public boolean replace(LeaseGrant grant) { + String key = LeaseStore.leaseKey(grant.getCompanyId(), grant.getCreditTypeId()); + return withLock(key, () -> { + LeaseState existing = leases.get(key); + if (existing != null && existing.isLiveAt(now())) { + // A live lease already holds this slot: preserve its already-debited balance + // rather than clobbering it. + return false; + } + if (existing != null && existing.getLeaseId().equals(grant.getLeaseId())) { + // The same lease coming back over its own expired row: a stale acquire response + // for a lease the idempotent server also handed a racing sibling, which may since + // have extended it. Rewriting would reset the balance to the full grant and erase + // debits whose reservations are still open, so reconcile like an extend instead. + leases.put(key, reconcile(existing, grant.getGrantedAmount(), grant.getExpiresAt())); + return false; + } + leases.put( + key, + new LeaseState( + grant.getLeaseId(), + grant.getCompanyId(), + grant.getCreditTypeId(), + grant.getGrantedAmount(), + grant.getGrantedAmount(), + grant.getExpiresAt())); + return true; + }); + } + + @Override + public ReserveResult tryReserve(String companyId, String creditTypeId, double credits) { + // Reject a non-finite or negative debit outright: NaN passes every comparison below, and + // a NaN balance would approve every later reserve. + if (!CreditAmounts.isValidQuantity(credits)) { + return null; + } + String key = LeaseStore.leaseKey(companyId, creditTypeId); + return withLock(key, () -> { + LeaseState entry = leases.get(key); + if (entry == null) { + return null; + } + // Never reserve against an expired lease: the server treats it as released and has + // refunded the grant, so the local balance is stale. + if (!entry.isLiveAt(now())) { + return null; + } + if (entry.getLocalRemainingCredits() < credits) { + return null; + } + double balance = entry.getLocalRemainingCredits() - credits; + leases.put(key, withBalance(entry, balance)); + // The lease id is read under the same lock as the debit: the caller pins its hold to + // it, so a read after the lock could name a lease that replaced this one in between. + return new ReserveResult(balance, entry.getLeaseId()); + }); + } + + @Override + public void refund(String companyId, String creditTypeId, double credits, String pinLeaseId) { + if (credits <= 0) { + return; + } + String key = LeaseStore.leaseKey(companyId, creditTypeId); + withLock(key, () -> { + LeaseState entry = leases.get(key); + if (entry == null) { + return null; + } + if (pinLeaseId != null + && !pinLeaseId.isEmpty() + && !entry.getLeaseId().equals(pinLeaseId)) { + return null; + } + double balance = Math.min(entry.getLocalRemainingCredits() + credits, entry.getGrantedAmount()); + leases.put(key, withBalance(entry, balance)); + return null; + }); + } + + @Override + public void extend( + String companyId, String creditTypeId, double grantedTotal, Instant newExpiresAt, String pinLeaseId) { + String key = LeaseStore.leaseKey(companyId, creditTypeId); + withLock(key, () -> { + LeaseState entry = leases.get(key); + if (entry == null) { + return null; + } + if (pinLeaseId != null + && !pinLeaseId.isEmpty() + && !entry.getLeaseId().equals(pinLeaseId)) { + return null; + } + leases.put(key, reconcile(entry, grantedTotal, newExpiresAt)); + return null; + }); + } + + @Override + public void drop(String companyId, String creditTypeId) { + String key = LeaseStore.leaseKey(companyId, creditTypeId); + while (true) { + ReentrantLock lock = lockFor(key); + lock.lock(); + try { + if (locks.get(key) != lock) { + continue; + } + leases.remove(key); + // Retire the lock with the state it guarded, so the slot leaves both maps + // together. Retiring it last, and while holding it, is what lets a waiter notice + // and retake the replacement. + // + // Only a release on close drops a slot, so both maps hold an entry for every + // (company, credit type) this process has leased against, expired ones included, + // until it exits. That is bounded by the tenants one process actually serves, and + // an expired entry has to stay readable anyway: a reserve must see it and refuse + // rather than see nothing and look unleased. + locks.remove(key, lock); + return; + } finally { + lock.unlock(); + } + } + } + + @Override + public List list() { + return new ArrayList<>(leases.values()); + } + + /** + * Reconciles an entry to a server-authoritative total: the delta is credited to the balance, + * a total already applied is a no-op, and the expiry only ever moves forward, so an + * out-of-order apply cannot shorten a lease a concurrent extend already pushed out. + */ + private static LeaseState reconcile(LeaseState entry, double grantedTotal, Instant newExpiresAt) { + double granted = entry.getGrantedAmount(); + double balance = entry.getLocalRemainingCredits(); + double add = grantedTotal - granted; + if (add > 0) { + granted = grantedTotal; + balance += add; + } + Instant expiresAt = entry.getExpiresAt(); + if (newExpiresAt != null && newExpiresAt.isAfter(expiresAt)) { + expiresAt = newExpiresAt; + } + return new LeaseState( + entry.getLeaseId(), entry.getCompanyId(), entry.getCreditTypeId(), granted, balance, expiresAt); + } + + private static LeaseState withBalance(LeaseState entry, double balance) { + return new LeaseState( + entry.getLeaseId(), + entry.getCompanyId(), + entry.getCreditTypeId(), + entry.getGrantedAmount(), + balance, + entry.getExpiresAt()); + } + + /** + * Runs {@code body} under the slot's lock, and only ever under the lock currently registered + * for that key. + * + *

{@link #drop} retires a lock along with the state it guarded. A thread that was already + * waiting on that lock wakes up holding a retired one, while a newcomer serialises on the + * replacement, so it re-checks and retakes the replacement instead of mutating the slot + * behind the newcomer's back. A lock can only be retired by a drop, so the retry terminates. + */ + private T withLock(String key, Supplier body) { + while (true) { + ReentrantLock lock = lockFor(key); + lock.lock(); + try { + if (locks.get(key) != lock) { + continue; + } + return body.get(); + } finally { + lock.unlock(); + } + } + } + + private ReentrantLock lockFor(String key) { + ReentrantLock lock = locks.get(key); + if (lock != null) { + return lock; + } + ReentrantLock created = new ReentrantLock(); + ReentrantLock raced = locks.putIfAbsent(key, created); + return raced != null ? raced : created; + } + + private Instant now() { + return clock.instant(); + } +} diff --git a/src/main/java/com/schematic/api/credits/InMemoryReservationStore.java b/src/main/java/com/schematic/api/credits/InMemoryReservationStore.java new file mode 100644 index 0000000..55102f3 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/InMemoryReservationStore.java @@ -0,0 +1,92 @@ +package com.schematic.api.credits; + +import java.time.Clock; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Keeps the reservation table in this process, and refunds into the lease store it is handed. + * Swap in {@link RedisReservationStore} to share holds across processes. + */ +public final class InMemoryReservationStore implements ReservationStore { + + private final ReservationRefunder leases; + private final Clock clock; + private final Map reservations = new ConcurrentHashMap<>(); + + public InMemoryReservationStore(ReservationRefunder leases) { + this(leases, Clock.systemUTC()); + } + + public InMemoryReservationStore(ReservationRefunder leases, Clock clock) { + this.leases = leases; + this.clock = clock != null ? clock : Clock.systemUTC(); + } + + @Override + public void add(Reservation reservation) { + reservations.put(reservation.getId(), reservation); + } + + @Override + public Reservation get(String id) { + return reservations.get(id); + } + + @Override + public Double consume(String id, double creditsConsumed) { + // remove() is the claim: of two racing callers exactly one comes away with the record. + Reservation reservation = reservations.remove(id); + if (reservation == null) { + return null; + } + double consumed = CreditAmounts.clampConsumption(creditsConsumed, reservation.getCreditsReserved()); + double refund = reservation.getCreditsReserved() - consumed; + if (refund > 0) { + // 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 already went + // back to the company balance server-side. + leases.refund(reservation.getCompanyId(), reservation.getCreditTypeId(), refund, reservation.getLeaseId()); + } + return consumed; + } + + @Override + public double reservedCredits(String companyId, String creditTypeId) { + double total = 0; + for (Reservation reservation : reservations.values()) { + if (reservation.getCompanyId().equals(companyId) + && reservation.getCreditTypeId().equals(creditTypeId)) { + total += reservation.getCreditsReserved(); + } + } + return total; + } + + @Override + public int sweepExpired() { + Instant cutoff = clock.instant(); + List expired = new ArrayList<>(); + for (Reservation reservation : reservations.values()) { + if (!reservation.getExpiresAt().isAfter(cutoff)) { + expired.add(reservation.getId()); + } + } + int swept = 0; + for (String id : expired) { + // Routed through consume so the sweep claims exactly once too. + if (consume(id, 0) != null) { + swept++; + } + } + return swept; + } + + @Override + public int count() { + return reservations.size(); + } +} diff --git a/src/main/java/com/schematic/api/credits/LeaseGrant.java b/src/main/java/com/schematic/api/credits/LeaseGrant.java new file mode 100644 index 0000000..d2a0c5a --- /dev/null +++ b/src/main/java/com/schematic/api/credits/LeaseGrant.java @@ -0,0 +1,45 @@ +package com.schematic.api.credits; + +import java.time.Instant; + +/** + * 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. + */ +public final class LeaseGrant { + + private final String leaseId; + private final String companyId; + private final String creditTypeId; + private final double grantedAmount; + private final Instant expiresAt; + + public LeaseGrant(String leaseId, String companyId, String creditTypeId, double grantedAmount, Instant expiresAt) { + this.leaseId = leaseId; + this.companyId = companyId; + this.creditTypeId = creditTypeId; + this.grantedAmount = grantedAmount; + this.expiresAt = expiresAt; + } + + public String getLeaseId() { + return leaseId; + } + + public String getCompanyId() { + return companyId; + } + + public String getCreditTypeId() { + return creditTypeId; + } + + /** The server-authoritative TOTAL, not the increment an extend asked for. */ + public double getGrantedAmount() { + return grantedAmount; + } + + public Instant getExpiresAt() { + return expiresAt; + } +} diff --git a/src/main/java/com/schematic/api/credits/LeaseLister.java b/src/main/java/com/schematic/api/credits/LeaseLister.java new file mode 100644 index 0000000..c010cae --- /dev/null +++ b/src/main/java/com/schematic/api/credits/LeaseLister.java @@ -0,0 +1,14 @@ +package com.schematic.api.credits; + +import java.util.List; + +/** + * Implemented only by a per-process store, whose leases are exclusively this process's, so + * releasing them on close is safe. A shared backend must never enumerate and release: sibling + * processes still draw on those leases. + */ +public interface LeaseLister { + + /** A snapshot of every slot this process holds. */ + List list(); +} diff --git a/src/main/java/com/schematic/api/credits/LeaseState.java b/src/main/java/com/schematic/api/credits/LeaseState.java new file mode 100644 index 0000000..a36a740 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/LeaseState.java @@ -0,0 +1,79 @@ +package com.schematic.api.credits; + +import java.time.Instant; + +/** The local view of the one lease a (company, credit type) slot holds. */ +public final class LeaseState { + + private final String leaseId; + private final String companyId; + private final String creditTypeId; + private final double grantedAmount; + private final double localRemainingCredits; + private final Instant expiresAt; + + public LeaseState( + String leaseId, + String companyId, + String creditTypeId, + double grantedAmount, + double localRemainingCredits, + Instant expiresAt) { + this.leaseId = leaseId; + this.companyId = companyId; + this.creditTypeId = creditTypeId; + this.grantedAmount = grantedAmount; + this.localRemainingCredits = localRemainingCredits; + this.expiresAt = expiresAt; + } + + /** Server-issued lease id. */ + public String getLeaseId() { + return leaseId; + } + + public String getCompanyId() { + return companyId; + } + + public String getCreditTypeId() { + return creditTypeId; + } + + /** Server-authoritative total granted to this lease. It grows on extend. */ + public double getGrantedAmount() { + return grantedAmount; + } + + /** + * Granted minus outstanding holds and consumption. It starts at the full grant when the lease + * is installed. + */ + public double getLocalRemainingCredits() { + return localRemainingCredits; + } + + /** + * The instant past which the lease is dead: the server has refunded the remainder to the + * company balance, so the local balance is stale and must never serve another reserve. + */ + public Instant getExpiresAt() { + return expiresAt; + } + + /** Whether the lease is still live at {@code now}. */ + public boolean isLiveAt(Instant now) { + return expiresAt.isAfter(now); + } + + @Override + public String toString() { + return "LeaseState{leaseId=" + leaseId + + ", companyId=" + companyId + + ", creditTypeId=" + creditTypeId + + ", grantedAmount=" + grantedAmount + + ", localRemainingCredits=" + localRemainingCredits + + ", expiresAt=" + expiresAt + + "}"; + } +} diff --git a/src/main/java/com/schematic/api/credits/LeaseStore.java b/src/main/java/com/schematic/api/credits/LeaseStore.java new file mode 100644 index 0000000..d5dc43c --- /dev/null +++ b/src/main/java/com/schematic/api/credits/LeaseStore.java @@ -0,0 +1,72 @@ +package com.schematic.api.credits; + +import java.time.Instant; + +/** + * Holds at most one lease per (company, credit type) slot. Every mutation is atomic per slot: + * {@link InMemoryLeaseStore} gets that from a per-slot lock, {@link RedisLeaseStore} from + * single-key Lua. + */ +public interface LeaseStore extends ReservationRefunder { + + /** The (company, credit type) slot key, shared by every backend. */ + static String leaseKey(String companyId, String creditTypeId) { + return companyId + ":" + creditTypeId; + } + + /** + * A snapshot of the slot, expired or not, or null when the slot is empty. Callers re-guard on + * expiry. + */ + LeaseState get(String companyId, String creditTypeId); + + /** + * Installs a fresh lease at its full grant, if the slot is free to take, and reports whether + * it wrote. + * + *

A live lease holds the slot even when it carries a different id (a sibling process won + * the acquire race): its already-debited balance wins and this reports false. An expired row + * carrying the SAME id is not rewritten either, since that would reset the balance and erase + * debits whose reservations are still open; it is reconciled like an extend (granted to the + * incoming total, expiry forward only, balance untouched) and also reports false. Only a + * fresh write reports true, which is what tells the manager whether the lease it just + * acquired is redundant. + */ + boolean replace(LeaseGrant grant); + + /** + * Atomically checks and debits, returning the post-debit balance and the lease the credits + * came out of, or null when there is no lease, the lease has expired, the balance is short, + * or {@code credits} is not a finite non-negative number. Returning the balance rather than a + * boolean lets the caller derive the pre-debit figure as {@code balance + credits} without a + * racy follow-up read. + * + *

The lease id is read in the same atomic step as the debit. The slot's lease can be + * replaced between a caller's acquire and its reserve, so a hold pinned to the lease the + * caller last saw would send its refunds to a lease that never held the credits, and bill + * that lease for the usage. + */ + ReserveResult tryReserve(String companyId, String creditTypeId, double credits); + + /** + * Returns credits to the slot's balance, clamped at the granted amount. With a non-null + * {@code pinLeaseId} the refund applies only while the slot still holds that lease: a hold + * carved out of an expired lease must never inflate its successor, whose grant the server + * already issued whole. + */ + @Override + void refund(String companyId, String creditTypeId, double credits, String pinLeaseId); + + /** + * Reconciles the slot to the server-authoritative total. The delta is computed inside the + * store against the currently stored total, never from a caller-held pre-wire-call read: two + * processes extending concurrently from the same stale read would each apply a delta and mint + * phantom credits. A total a sibling already applied is a no-op, so applies converge in any + * order. Expiry only ever moves forward. A non-null {@code pinLeaseId} drops the whole extend + * when the slot holds a different lease. + */ + void extend(String companyId, String creditTypeId, double grantedTotal, Instant newExpiresAt, String pinLeaseId); + + /** Removes the slot entry, after a remote release. */ + void drop(String companyId, String creditTypeId); +} diff --git a/src/main/java/com/schematic/api/credits/LeaseWireClient.java b/src/main/java/com/schematic/api/credits/LeaseWireClient.java new file mode 100644 index 0000000..6cd0fad --- /dev/null +++ b/src/main/java/com/schematic/api/credits/LeaseWireClient.java @@ -0,0 +1,34 @@ +package com.schematic.api.credits; + +import java.time.Duration; +import java.time.Instant; + +/** + * The three lease calls the manager makes. Narrow on purpose: it keeps the manager independent of + * the generated client's request and response models, and lets tests script the server. + * + *

Every method throws on a wire failure; the manager resolves that to "no lease" so callers + * route it through their fail-open or fail-closed handling. + */ +public interface LeaseWireClient { + + LeaseGrant acquire(String companyId, String creditTypeId, double requestedAmount, Instant expiresAt); + + /** + * Acquires under the caller's per-check timeout. A null timeout means the client's own, which + * is what a background top-up nobody is waiting on takes. + */ + default LeaseGrant acquire( + String companyId, String creditTypeId, double requestedAmount, Instant expiresAt, Duration timeout) { + return acquire(companyId, creditTypeId, requestedAmount, expiresAt); + } + + LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt); + + /** Extends under the caller's per-check timeout, or the client's own when null. */ + default LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt, Duration timeout) { + return extend(leaseId, additionalAmount, expiresAt); + } + + void release(String leaseId); +} diff --git a/src/main/java/com/schematic/api/credits/OnAcquireFailure.java b/src/main/java/com/schematic/api/credits/OnAcquireFailure.java new file mode 100644 index 0000000..7b595f1 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/OnAcquireFailure.java @@ -0,0 +1,18 @@ +package com.schematic.api.credits; + +/** + * What a check does when it cannot gate: the wire call failed, the store is unreachable, or the + * lease is exhausted. + */ +public enum OnAcquireFailure { + /** + * Err on the side of assuming the credits are there. The rules engine still evaluates the + * flag, with the credit balance substituted to an effectively unlimited value, so plan + * targeting, overrides and every non-credit condition still apply and only the credit gate is + * bypassed. No reservation is issued. In server mode there is no local engine to re-run, so + * the check returns the caller's default value instead. + */ + FAIL_OPEN, + /** Deny, so the caller blocks the action. The default. */ + FAIL_CLOSED +} diff --git a/src/main/java/com/schematic/api/credits/PreflightOptions.java b/src/main/java/com/schematic/api/credits/PreflightOptions.java new file mode 100644 index 0000000..b9efe02 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/PreflightOptions.java @@ -0,0 +1,125 @@ +package com.schematic.api.credits; + +import com.schematic.api.types.PreflightEventUsageRequestBody; +import com.schematic.api.types.PreflightRequestBody; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; + +/** + * The preflight a credit-gated check asks an evaluation to answer: what this call is about to + * cost, before it has been recorded. + * + *

Plain data rather than the engine's own option type, so the flow can be driven and asserted + * on without a WASM runtime, and so the same value can be sent to the API and to the local engine. + */ +public final class PreflightOptions { + + /** A simulated quantity scoped to one event subtype. */ + public static final class EventUsage { + private final String eventSubtype; + private final double quantity; + + public EventUsage(String eventSubtype, double quantity) { + this.eventSubtype = eventSubtype; + this.quantity = quantity; + } + + public String getEventSubtype() { + return eventSubtype; + } + + public double getQuantity() { + return quantity; + } + } + + private final Map creditCost; + private final Double usage; + private final EventUsage eventUsage; + + private PreflightOptions(Map creditCost, Double usage, EventUsage eventUsage) { + this.creditCost = creditCost == null || creditCost.isEmpty() + ? null + : Collections.unmodifiableMap(new LinkedHashMap<>(creditCost)); + this.usage = usage; + this.eventUsage = eventUsage; + } + + /** + * The preflight a caller's usage implies. With a subtype the quantity goes out scoped to it so + * the engine matches that subtype's condition; without one it goes out as the generic knob. + * Returns null when there is no usage to declare. + */ + public static PreflightOptions fromUsage(Double usage, String eventSubtype) { + if (usage == null || !CreditAmounts.isValidQuantity(usage)) { + return null; + } + // Held as the caller gave it, fraction and all, so a client-mode hold can be sized off + // the same figure the reservation records. Rounding belongs at each boundary that needs + // it, not here: both the request body and the engine envelope take integers. + if (eventSubtype != null && !eventSubtype.isEmpty()) { + return new PreflightOptions(null, null, new EventUsage(eventSubtype, usage)); + } + return new PreflightOptions(null, usage, null); + } + + /** Prices one credit type directly, bypassing the engine's own quantity times rate arithmetic. */ + public static PreflightOptions forCreditCost(String creditTypeId, double cost) { + return new PreflightOptions(Collections.singletonMap(creditTypeId, cost), null, null); + } + + /** + * Casts a usage onto the integer the request body carries. A preflight asks an upper-bound + * question, so a fraction rounds up: the check must not pass on less usage than the operation + * is about to record. The engine envelope rounds the same way, in + * {@code WasmRulesEngine.putQuantity}, so a local evaluation and the API answer one figure. + */ + public static long preflightQuantity(double usage) { + return (long) Math.ceil(usage); + } + + public Map getCreditCost() { + return creditCost; + } + + public Double getUsage() { + return usage; + } + + public EventUsage getEventUsage() { + return eventUsage; + } + + /** + * The same preflight as the API's request body, for the paths that gate server-side. Null + * when nothing in it would change the answer, so the caller can send a plain request. + * + *

A zero usage and a zero event quantity are dropped: the API documents them as having no + * effect, and a request that carries one is still a preflighted request, which costs it the + * flag check cache for nothing. A zero credit cost stays, because that one says something, + * namely that this call is free rather than unpriced. + */ + public PreflightRequestBody toRequestBody() { + boolean hasUsage = usage != null && usage != 0; + boolean hasEventUsage = eventUsage != null && eventUsage.getQuantity() != 0; + if (creditCost == null && !hasUsage && !hasEventUsage) { + return null; + } + PreflightRequestBody.Builder builder = PreflightRequestBody.builder(); + if (creditCost != null) { + builder.creditCost(creditCost); + } + // The wire fields are integers, so the rounding happens on the way out, not on the way in. + if (hasUsage) { + builder.usage(preflightQuantity(usage)); + } + if (hasEventUsage) { + builder.eventUsage(PreflightEventUsageRequestBody.builder() + .eventSubtype(eventUsage.getEventSubtype()) + .quantity(preflightQuantity(eventUsage.getQuantity())) + .build()); + } + return builder.build(); + } +} diff --git a/src/main/java/com/schematic/api/credits/PrewarmCompanyResolver.java b/src/main/java/com/schematic/api/credits/PrewarmCompanyResolver.java new file mode 100644 index 0000000..cd2829f --- /dev/null +++ b/src/main/java/com/schematic/api/credits/PrewarmCompanyResolver.java @@ -0,0 +1,145 @@ +package com.schematic.api.credits; + +import com.schematic.api.types.RulesengineCompany; +import java.time.Duration; +import java.util.Map; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.function.BooleanSupplier; +import java.util.function.Function; + +/** + * Resolves the company id a prewarm acquires against, from keys that may carry no id. + * + *

Split out of the client so the waiting rule can be driven without a socket: a zero timeout + * means cache-only rather than no resolution at all, which is the difference between a prewarm + * that warms an already-cached company and one that never runs. + */ +public final class PrewarmCompanyResolver { + + /** Prefix Schematic's company secure ids carry, whatever key name they are passed under. */ + public static final String COMPANY_ID_PREFIX = "comp_"; + + private PrewarmCompanyResolver() {} + + /** + * The Schematic id hiding among a set of entity keys, recognised by its secure-id prefix. The + * server reads keys this way once its own key lookup has come up empty, so + * {@code {account_id: "comp_1"}} resolves and {@code {id: "acme"}} does not: the prefix + * decides, not the name of the key the value arrived under. + */ + public static String schematicId(Map keys, String prefix) { + if (keys == null) { + return null; + } + for (String value : keys.values()) { + if (value != null && value.startsWith(prefix)) { + return value; + } + } + return null; + } + + /** + * Returns the company id, resolved in the server's order: every supplied key/value pair is an + * ordinary entity key and gets looked up first, and only when nothing matches is a value read + * as the company's own id, by its {@code comp_} prefix. An account is free to define a key + * called {@code id} holding its own identifier, so the name alone settles nothing. Null when + * the keys never resolve and carry no Schematic id. + * + * @param keys the caller's company keys + * @param cached reads the company from the local cache only + * @param fetch resolves the company over the wire, warming the cache as a side effect + * @param timeout how long to keep fetching; zero or less is cache-only + * @param pollInterval how long to wait between fetches + * @param abort answers true when the caller has given up, for instance a closing client + * @param onFetchError reports a failed fetch, which is retried until the timeout + */ + public static String resolve( + Map keys, + Function, RulesengineCompany> cached, + Function, RulesengineCompany> fetch, + Duration timeout, + Duration pollInterval, + BooleanSupplier abort, + Function onFetchError) { + if (keys == null || keys.isEmpty()) { + return null; + } + try { + RulesengineCompany hit = cached.apply(keys); + if (hit != null) { + return hit.getId(); + } + } catch (RuntimeException e) { + // A cache that throws is a miss, not a failed prewarm: the fetch below answers the + // same question over the wire. + onFetchError.apply(e); + } + // A zero timeout is cache-only, not a refusal: the caller asked not to wait on the wire, + // and the cache has already answered above. + if (timeout == null || timeout.toMillis() <= 0) { + return schematicId(keys, COMPANY_ID_PREFIX); + } + + // Retry across the brief connecting window at boot. A new company needs the preceding + // identify ingested before the server can stream it back. + long deadline = System.nanoTime() + timeout.toNanos(); + // The fetch runs on its own thread so the timeout bounds the fetch itself, not just the + // gaps between attempts: a single call that never returns would otherwise hold the + // prewarm past every deadline the caller set. Daemon, so a stuck one cannot keep the + // process alive. + ExecutorService fetcher = Executors.newSingleThreadExecutor(runnable -> { + Thread thread = new Thread(runnable, "SchematicPrewarmResolve"); + thread.setDaemon(true); + return thread; + }); + try { + while (true) { + if (abort.getAsBoolean()) { + return null; + } + long remaining = deadline - System.nanoTime(); + if (remaining <= 0) { + return schematicId(keys, COMPANY_ID_PREFIX); + } + Future pending = fetcher.submit(() -> fetch.apply(keys)); + try { + RulesengineCompany resolved = pending.get(remaining, TimeUnit.NANOSECONDS); + if (resolved != null) { + return resolved.getId(); + } + } catch (TimeoutException e) { + pending.cancel(true); + return schematicId(keys, COMPANY_ID_PREFIX); + } catch (ExecutionException e) { + onFetchError.apply(asRuntime(e.getCause())); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return null; + } + if (System.nanoTime() >= deadline) { + // 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 schematicId(keys, COMPANY_ID_PREFIX); + } + try { + Thread.sleep(pollInterval.toMillis()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return null; + } + } + } finally { + fetcher.shutdownNow(); + } + } + + private static RuntimeException asRuntime(Throwable cause) { + return cause instanceof RuntimeException ? (RuntimeException) cause : new RuntimeException(cause); + } +} diff --git a/src/main/java/com/schematic/api/credits/RedisLeaseStore.java b/src/main/java/com/schematic/api/credits/RedisLeaseStore.java new file mode 100644 index 0000000..c84aecb --- /dev/null +++ b/src/main/java/com/schematic/api/credits/RedisLeaseStore.java @@ -0,0 +1,285 @@ +package com.schematic.api.credits; + +import java.time.Clock; +import java.time.Instant; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import redis.clients.jedis.JedisPooled; + +/** + * Keeps lease slots in Redis, one hash per (company, credit type), so every process pointed at + * the same Redis gates against one balance. + * + *

The key layout, hash fields and Lua scripts are identical to the Node, Go and Python SDKs', + * which is what lets a mixed fleet share one lease. Every script touches exactly one key, keeping + * them safe under Redis Cluster, where a multi-key script spanning slots raises CROSSSLOT. + */ +public final class RedisLeaseStore implements LeaseStore { + + private static final String LEASE_KEY_NAMESPACE = "credit-lease:"; + // How long after the declared expiry the row survives before Redis evicts it. It gives the + // sweeper a window to refund expired reservations before the lease state underneath them + // disappears. + private static final long LEASE_TTL_GRACE_MS = 60_000L; + + // Expiry is decided against the Redis server's clock, not the calling process's: with many + // processes sharing one lease, local clock skew would let them disagree on whether the lease + // is live. The snippet converts TIME to integer milliseconds (matching the stored expiresAt); + // replicate_commands() comes first, so the non-deterministic TIME read is allowed alongside + // writes on Redis 5 and 6. + private static final String LEASE_NOW_MS = "\n" + + "redis.replicate_commands()\n" + + "local t = redis.call('TIME')\n" + + "local now = (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000)\n"; + + /** + * Atomic replace. Writes the lease hash only when the slot is empty or the existing lease has + * expired. Returns 1 on write, 0 if a live lease already occupies the slot, even one with a + * different leaseId, e.g. installed by a sibling that raced this acquire. An expired row with + * the SAME leaseId is reconciled like an extend instead of rewritten, which would reset the + * balance and erase debits whose reservations are still open. + */ + private static final String REPLACE_SCRIPT = LEASE_NOW_MS + + "\n" + + "local existing_id = redis.call('HGET', KEYS[1], 'leaseId')\n" + + "local existing_expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0')\n" + + "local new_id = ARGV[1]\n" + + "local new_granted = ARGV[2]\n" + + "local new_expiry = tonumber(ARGV[3])\n" + + "local grace = tonumber(ARGV[4])\n" + + "\n" + + "if existing_id and existing_expiry > now then\n" + + " return 0\n" + + "end\n" + + "\n" + + "if existing_id == new_id then\n" + + " local granted = tonumber(redis.call('HGET', KEYS[1], 'grantedAmount') or '0')\n" + + " local add = tonumber(new_granted) - granted\n" + + " if add > 0 then\n" + + " local remaining = tonumber(redis.call('HGET', KEYS[1], 'localRemainingCredits') or '0')\n" + + " redis.call('HSET', KEYS[1],\n" + + " 'grantedAmount', new_granted,\n" + + " 'localRemainingCredits', tostring(remaining + add))\n" + + " end\n" + + " if new_expiry > existing_expiry then\n" + + " redis.call('HSET', KEYS[1], 'expiresAt', ARGV[3])\n" + + " redis.call('PEXPIREAT', KEYS[1], new_expiry + grace)\n" + + " end\n" + + " return 0\n" + + "end\n" + + "\n" + + "redis.call('DEL', KEYS[1])\n" + + "redis.call('HSET', KEYS[1],\n" + + " 'leaseId', new_id,\n" + + " 'companyId', ARGV[5],\n" + + " 'creditTypeId', ARGV[6],\n" + + " 'grantedAmount', new_granted,\n" + + " 'localRemainingCredits', new_granted,\n" + + " 'expiresAt', ARGV[3])\n" + + "redis.call('PEXPIREAT', KEYS[1], new_expiry + grace)\n" + + "return 1\n"; + + /** + * Atomic check-and-decrement on localRemainingCredits. Returns the post-debit balance as a + * string (a Lua number reply truncates to integer, which would corrupt fractional credit + * costs) alongside the leaseId it came out of; nil if there is no lease, the lease has + * expired, or there is insufficient remaining. The expiry guard compares against the Redis + * server clock, so a reserve against an expired-but-not-yet-evicted row during the TTL grace + * window is rejected. + */ + private static final String TRY_RESERVE_SCRIPT = LEASE_NOW_MS + + "\n" + + "local raw = redis.call('HGET', KEYS[1], 'localRemainingCredits')\n" + + "if not raw then return false end\n" + + "local lease_id = redis.call('HGET', KEYS[1], 'leaseId')\n" + + "if not lease_id then return false end\n" + + "local expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0')\n" + + "if expiry <= now then return false end\n" + + "local remaining = tonumber(raw)\n" + + "local requested = tonumber(ARGV[1])\n" + + "if remaining < requested then return false end\n" + + "local new_remaining = remaining - requested\n" + + "redis.call('HSET', KEYS[1], 'localRemainingCredits', tostring(new_remaining))\n" + + "return { tostring(new_remaining), lease_id }\n"; + + /** + * Refund credits, clamped at grantedAmount. ARGV[2], when non-empty, pins the refund to a + * specific 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, so crediting the successor would mint phantom credits. + */ + private static final String REFUND_SCRIPT = "\n" + + "local raw_remaining = redis.call('HGET', KEYS[1], 'localRemainingCredits')\n" + + "if not raw_remaining then return 0 end\n" + + "local required_lease = ARGV[2]\n" + + "if required_lease and required_lease ~= '' then\n" + + " local current_lease = redis.call('HGET', KEYS[1], 'leaseId')\n" + + " if current_lease ~= required_lease then return 0 end\n" + + "end\n" + + "local remaining = tonumber(raw_remaining)\n" + + "local granted = tonumber(redis.call('HGET', KEYS[1], 'grantedAmount') or '0')\n" + + "local refund = tonumber(ARGV[1])\n" + + "local new_balance = remaining + refund\n" + + "if new_balance > granted then new_balance = granted end\n" + + "redis.call('HSET', KEYS[1], 'localRemainingCredits', tostring(new_balance))\n" + + "return 1\n"; + + /** + * Reconcile the lease to the server-authoritative grantedAmount total (ARGV[1]), crediting + * the difference to localRemainingCredits. The delta is computed here, atomically against the + * hash's current total, never by the caller from a pre-wire-call read: per-process + * single-flight does not cover sibling processes, so two extending the same shared lease + * concurrently would each apply a delta against the same stale read and mint phantom credits. + * Expiry only ever moves forward. ARGV[4], when non-empty, pins the extend to a specific + * leaseId, mirroring the pin on the refund script. + */ + private static final String EXTEND_SCRIPT = "\n" + + "local raw_granted = redis.call('HGET', KEYS[1], 'grantedAmount')\n" + + "if not raw_granted then return 0 end\n" + + "local required_lease = ARGV[4]\n" + + "if required_lease and required_lease ~= '' then\n" + + " local current_lease = redis.call('HGET', KEYS[1], 'leaseId')\n" + + " if current_lease ~= required_lease then return 0 end\n" + + "end\n" + + "local granted = tonumber(raw_granted)\n" + + "local target = tonumber(ARGV[1])\n" + + "local add = target - granted\n" + + "if add > 0 then\n" + + " local remaining = tonumber(redis.call('HGET', KEYS[1], 'localRemainingCredits') or '0')\n" + + " redis.call('HSET', KEYS[1],\n" + + " 'grantedAmount', tostring(target),\n" + + " 'localRemainingCredits', tostring(remaining + add))\n" + + "end\n" + + "local new_expiry = tonumber(ARGV[2])\n" + + "local grace = tonumber(ARGV[3])\n" + + "local current_expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0')\n" + + "if new_expiry > current_expiry then\n" + + " redis.call('HSET', KEYS[1], 'expiresAt', ARGV[2])\n" + + " redis.call('PEXPIREAT', KEYS[1], new_expiry + grace)\n" + + "end\n" + + "return 1\n"; + + private final JedisPooled jedis; + private final String keyPrefix; + private final long defaultLeaseDurationMs; + private final Clock clock; + + public RedisLeaseStore(JedisPooled jedis) { + this(jedis, null, null, null); + } + + /** + * @param jedis a pre-configured client, shared with the datastream cache when one is set up + * @param keyPrefix namespace for lease keys; defaults to {@code schematic:} + * @param defaultLeaseDuration fallback expiry for an {@link #extend} that supplies none. The + * lease manager always supplies one, so this only matters for a direct caller. + * @param clock reads the current time for the expiry a bare extend derives. Lease liveness + * itself is decided by the Redis server's clock inside the scripts. + */ + public RedisLeaseStore(JedisPooled jedis, String keyPrefix, java.time.Duration defaultLeaseDuration, Clock clock) { + this.jedis = jedis; + this.keyPrefix = keyPrefix != null ? keyPrefix : CreditLeaseDefaults.KEY_PREFIX; + this.defaultLeaseDurationMs = + (defaultLeaseDuration != null ? defaultLeaseDuration : CreditLeaseDefaults.LEASE_DURATION).toMillis(); + this.clock = clock != null ? clock : Clock.systemUTC(); + } + + /** Public so the reservation store can target the same lease hash for refunds. */ + public String hashKey(String companyId, String creditTypeId) { + return keyPrefix + LEASE_KEY_NAMESPACE + LeaseStore.leaseKey(companyId, creditTypeId); + } + + @Override + public LeaseState get(String companyId, String creditTypeId) { + Map raw = jedis.hgetAll(hashKey(companyId, creditTypeId)); + if (raw == null || raw.get("leaseId") == null) { + return null; + } + return new LeaseState( + raw.get("leaseId"), + raw.get("companyId"), + raw.get("creditTypeId"), + CreditAmounts.parse(raw.get("grantedAmount"), 0), + CreditAmounts.parse(raw.get("localRemainingCredits"), 0), + Instant.ofEpochMilli((long) CreditAmounts.parse(raw.get("expiresAt"), 0))); + } + + @Override + public boolean replace(LeaseGrant grant) { + // No client clock here: the script reads now from the Redis server via TIME, so every + // process agrees on expiry. + Object result = jedis.eval( + REPLACE_SCRIPT, + Collections.singletonList(hashKey(grant.getCompanyId(), grant.getCreditTypeId())), + Arrays.asList( + grant.getLeaseId(), + CreditAmounts.format(grant.getGrantedAmount()), + Long.toString(grant.getExpiresAt().toEpochMilli()), + Long.toString(LEASE_TTL_GRACE_MS), + grant.getCompanyId(), + grant.getCreditTypeId())); + return result instanceof Long && (Long) result == 1L; + } + + @Override + public ReserveResult tryReserve(String companyId, String creditTypeId, double credits) { + // Reject a non-finite or negative debit before it reaches the script: "NaN" parses back + // to nan in Lua, slips through the comparison, and would poison the shared balance. + if (!CreditAmounts.isValidQuantity(credits)) { + return null; + } + Object result = jedis.eval( + TRY_RESERVE_SCRIPT, + Collections.singletonList(hashKey(companyId, creditTypeId)), + // Only the requested amount: now comes from the Redis server clock. + Collections.singletonList(CreditAmounts.format(credits))); + // A nil reply (could not reserve) arrives as null; success is a two-element reply of + // [post-debit balance, charged leaseId], both strings. + if (!(result instanceof List)) { + return null; + } + List reply = (List) result; + if (reply.size() < 2) { + return null; + } + return new ReserveResult(CreditAmounts.parse(String.valueOf(reply.get(0)), 0), String.valueOf(reply.get(1))); + } + + @Override + public void refund(String companyId, String creditTypeId, double credits, String pinLeaseId) { + if (credits <= 0) { + return; + } + jedis.eval( + REFUND_SCRIPT, + Collections.singletonList(hashKey(companyId, creditTypeId)), + // An empty string disables the lease pin: Lua has no nil ARGV. + Arrays.asList(CreditAmounts.format(credits), pinLeaseId != null ? pinLeaseId : "")); + } + + @Override + public void extend( + String companyId, String creditTypeId, double grantedTotal, Instant newExpiresAt, String pinLeaseId) { + long expiry = newExpiresAt != null + ? newExpiresAt.toEpochMilli() + : clock.instant().toEpochMilli() + defaultLeaseDurationMs; + jedis.eval( + EXTEND_SCRIPT, + Collections.singletonList(hashKey(companyId, creditTypeId)), + // grantedTotal is the server-authoritative TOTAL; the script computes the credit + // delta atomically against the stored total. + Arrays.asList( + CreditAmounts.format(grantedTotal), + Long.toString(expiry), + Long.toString(LEASE_TTL_GRACE_MS), + pinLeaseId != null ? pinLeaseId : "")); + } + + @Override + public void drop(String companyId, String creditTypeId) { + // A plain single-key delete: there is no secondary index to keep in sync. + jedis.del(hashKey(companyId, creditTypeId)); + } +} diff --git a/src/main/java/com/schematic/api/credits/RedisReservationStore.java b/src/main/java/com/schematic/api/credits/RedisReservationStore.java new file mode 100644 index 0000000..deec8d9 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/RedisReservationStore.java @@ -0,0 +1,351 @@ +package com.schematic.api.credits; + +import com.fasterxml.jackson.databind.JsonNode; +import com.schematic.api.core.ObjectMappers; +import java.time.Clock; +import java.time.Instant; +import java.util.Collections; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import redis.clients.jedis.AbstractTransaction; +import redis.clients.jedis.JedisPooled; + +/** + * Keeps the reservation table in Redis and refunds through the lease store it is handed. + * + *

Each hold is a hash, indexed by expiry in a sorted set so the sweeper can pop expired + * entries in O(log n), and by (company, credit type) so a balance display can sum a tenant's open + * holds. Every mutation is a single-key operation, or single-key Lua, or a transaction over one + * key, 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. + */ +public final class RedisReservationStore implements ReservationStore { + + private static final String RES_KEY_NAMESPACE = "credit-reservation:"; + // Sorted set scoring open reservations by expiresAt so the sweeper can pop expired entries in + // O(log n). Members encode the full (company, credit, id) tuple. + private static final String RES_INDEX_KEY = "credit-reservations:byExpiry"; + // Per-(company, credit) index of open holds, one hash of reservationId -> creditsReserved, so + // reservedCredits reads a tenant's holds with one HGETALL. The hash is also the source of + // truth for that sum: a field exists exactly while its reservation is open and unrefunded. + private static final String RES_BYCREDIT_NAMESPACE = "credit-reservations:byCredit:"; + // Buffer past expiresAt before Redis evicts the row, so the sweeper has a window to refund. + private static final long RES_TTL_GRACE_MS = 30_000L; + // Page size for the sweeper's ZRANGEBYSCORE. Without a limit, a backlog of expired holds + // (after a Redis outage or a long pause) would come back as one giant reply on every + // process's next tick; paging bounds the reply while the per-member ZREM keeps offset 0 + // advancing through the backlog. + private static final int SWEEP_BATCH_SIZE = 256; + // Bounds one sweep's work, and guards against an endless loop if ZREM persistently fails. + // Anything left over is picked up on the next tick. + private static final int MAX_SWEEP_BATCHES = 16; + // Absent from Schematic ids and from the UUID reservation id. + private static final String MEMBER_DELIMITER = "|"; + + /** + * Atomic claim: read the reservation hash and delete it in one step, returning its fields (or + * nil if it was already gone). Touches a single key. The atomic read-then-delete is what + * makes consume exactly-once: of two racing callers (a normal track and a sweeper, say) only + * one gets the fields back and proceeds to refund. The refund to the lease hash 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. + */ + private static final String CLAIM_SCRIPT = "\n" + + "local raw = redis.call('HGETALL', KEYS[1])\n" + + "if #raw == 0 then return nil end\n" + + "redis.call('DEL', KEYS[1])\n" + + "return raw\n"; + + private final JedisPooled jedis; + private final ReservationRefunder leases; + private final String keyPrefix; + private final Clock clock; + + public RedisReservationStore(JedisPooled jedis, ReservationRefunder leases) { + this(jedis, leases, null, null); + } + + /** + * @param keyPrefix namespace for reservation keys; pass the same prefix as the lease store + * @param clock decides the sweep cutoff + */ + public RedisReservationStore(JedisPooled jedis, ReservationRefunder leases, String keyPrefix, Clock clock) { + this.jedis = jedis; + this.leases = leases; + this.keyPrefix = keyPrefix != null ? keyPrefix : CreditLeaseDefaults.KEY_PREFIX; + this.clock = clock != null ? clock : Clock.systemUTC(); + } + + @Override + public void add(Reservation reservation) { + long expiresMs = reservation.getExpiresAt().toEpochMilli(); + String hashKey = hashKey(reservation.getId()); + Map hash = new LinkedHashMap<>(); + hash.put("id", reservation.getId()); + hash.put("leaseId", reservation.getLeaseId()); + hash.put("companyId", reservation.getCompanyId()); + hash.put("creditTypeId", reservation.getCreditTypeId()); + hash.put("eventSubtype", reservation.getEventSubtype()); + hash.put("quantityReserved", CreditAmounts.format(reservation.getQuantityReserved())); + hash.put("creditsReserved", CreditAmounts.format(reservation.getCreditsReserved())); + hash.put("consumptionRate", CreditAmounts.format(reservation.getConsumptionRate())); + hash.put("expiresAt", Long.toString(expiresMs)); + hash.put("evalCtx", encodeEvalCtx(reservation)); + // 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. Both commands + // touch the one key, so this is Cluster-safe. + try (AbstractTransaction txn = jedis.multi()) { + txn.hset(hashKey, hash); + txn.pexpireAt(hashKey, expiresMs + RES_TTL_GRACE_MS); + txn.exec(); + } + // The two indexes (expiry zset for the sweeper, per-tenant hash for reservedCredits) only + // depend on the hash existing, so they stay outside the transaction, where their keys are + // free to hash to other Cluster slots. A partial failure here at worst leaves an + // un-indexed reservation that the TTL reaps (its slice reclaimed when the lease expires), + // never a double-spend. + jedis.zadd( + indexKey(), + (double) expiresMs, + encodeMember(reservation.getCompanyId(), reservation.getCreditTypeId(), reservation.getId())); + jedis.hset( + byCreditKey(reservation.getCompanyId(), reservation.getCreditTypeId()), + reservation.getId(), + CreditAmounts.format(reservation.getCreditsReserved())); + } + + @Override + public Reservation get(String id) { + Map raw = jedis.hgetAll(hashKey(id)); + if (raw == null || raw.get("id") == null) { + return null; + } + return decode(raw); + } + + @Override + public Double consume(String id, double creditsConsumed) { + // Atomically claim (read and delete) the reservation hash. Only one caller wins; a + // duplicate or racing consume gets nil and reports nothing claimed. + Object claimed = jedis.eval(CLAIM_SCRIPT, Collections.singletonList(hashKey(id)), Collections.emptyList()); + Map raw = decodeFlat(claimed); + if (raw == null || raw.get("id") == null) { + return null; + } + + String companyId = raw.get("companyId"); + String creditTypeId = raw.get("creditTypeId"); + double reserved = CreditAmounts.parse(raw.get("creditsReserved"), 0); + + // Index cleanup, single-key ops. The per-tenant hash loses the slice BEFORE the refund + // below, so the lease (local remaining plus this hash) never transiently double-counts + // it. Both are best-effort: a failed cleanup must not abort the settle. The per-tenant + // field goes first, since the expiry index is what the sweeper would reach a surviving + // field through: dropping that first and then failing here would inflate reservedCredits + // forever. + ignoringFailures(() -> jedis.hdel(byCreditKey(companyId, creditTypeId), id)); + ignoringFailures(() -> jedis.zrem(indexKey(), encodeMember(companyId, creditTypeId, id))); + + double consumed = CreditAmounts.clampConsumption(creditsConsumed, reserved); + double refund = reserved - consumed; + if (refund > 0) { + // Delegated to the lease store, which owns the lease hash, so this cross-key write + // stays out of a single Lua script. Pinned to the reservation's lease so a hold + // carved out of an expired lease cannot inflate a successor's balance. + leases.refund(companyId, creditTypeId, refund, raw.get("leaseId")); + } + return consumed; + } + + @Override + public double reservedCredits(String companyId, String creditTypeId) { + Map byCredit; + try { + byCredit = jedis.hgetAll(byCreditKey(companyId, creditTypeId)); + } catch (RuntimeException e) { + // A display-path read, not a gate. A Redis blip here reads as nothing reserved rather + // than as an exception thrown at a caller asking what the balance looks like. + return 0; + } + if (byCredit == null) { + return 0; + } + double total = 0; + for (String value : byCredit.values()) { + total += CreditAmounts.parse(value, 0); + } + return total; + } + + @Override + public int sweepExpired() { + long cutoff = clock.instant().toEpochMilli(); + int swept = 0; + // Page through expired members (encoded company|credit|id, scored by expiresAt) rather + // than fetching them all at once. Each processed member is removed below, so re-reading + // at offset 0 advances through the backlog. + for (int batch = 0; batch < MAX_SWEEP_BATCHES; batch++) { + List expired = jedis.zrangeByScore(indexKey(), 0, (double) cutoff, 0, SWEEP_BATCH_SIZE); + if (expired == null || expired.isEmpty()) { + return swept; + } + for (String member : expired) { + String[] parts = decodeMember(member); + if (parts == null) { + // Nothing but add writes members, so this is belt-and-braces: drop it rather + // than let it wedge the sweeper. + ignoringFailures(() -> jedis.zrem(indexKey(), member)); + continue; + } + Double consumed = consume(parts[2], 0); + // Always drop the member just read. On the success path consume already removed + // it, so this is idempotent; it also covers the hash-evicted path below. + ignoringFailures(() -> jedis.zrem(indexKey(), member)); + if (consumed != null) { + swept++; + continue; + } + // No reservation hash: either a racing track consumed it (and reconciled the + // byCredit field, making this a no-op) or the hash TTL-evicted before the sweeper + // reached it, orphaning the field. Reconcile so reservedCredits stops summing an + // evicted hold. Deliberately no refund: without the hash, exactly-once cannot be + // arbitrated across racing sweepers, so the slice waits for the lease to expire + // server-side. + ignoringFailures(() -> jedis.hdel(byCreditKey(parts[0], parts[1]), parts[2])); + } + if (expired.size() < SWEEP_BATCH_SIZE) { + return swept; + } + } + return swept; + } + + @Override + public int count() { + try { + return (int) jedis.zcard(indexKey()); + } catch (RuntimeException e) { + return 0; + } + } + + private String hashKey(String id) { + return keyPrefix + RES_KEY_NAMESPACE + id; + } + + private String indexKey() { + return keyPrefix + RES_INDEX_KEY; + } + + private String byCreditKey(String companyId, String creditTypeId) { + return keyPrefix + RES_BYCREDIT_NAMESPACE + companyId + ":" + creditTypeId; + } + + /** + * Packs the whole tuple into an expiry-index member. The sweeper needs company and credit to + * clean the per-tenant hash even after the reservation hash has TTL-evicted, at which point + * the claim returns nil and cannot report them; otherwise the orphaned field would inflate + * reservedCredits forever. + */ + private static String encodeMember(String companyId, String creditTypeId, String id) { + return companyId + MEMBER_DELIMITER + creditTypeId + MEMBER_DELIMITER + id; + } + + private static String[] decodeMember(String member) { + String[] parts = member.split("\\|", -1); + return parts.length == 3 ? parts : null; + } + + private static String encodeEvalCtx(Reservation reservation) { + Map> ctx = new LinkedHashMap<>(); + if (reservation.getCompany() != null) { + ctx.put("company", reservation.getCompany()); + } + if (reservation.getUser() != null) { + ctx.put("user", reservation.getUser()); + } + try { + return ObjectMappers.JSON_MAPPER.writeValueAsString(ctx); + } catch (Exception e) { + return "{}"; + } + } + + /** Decodes the flat [field, value, ...] reply the claim script returns. */ + private static Map decodeFlat(Object raw) { + if (!(raw instanceof List)) { + return null; + } + List flat = (List) raw; + if (flat.isEmpty()) { + return null; + } + Map out = new LinkedHashMap<>(); + for (int i = 0; i + 1 < flat.size(); i += 2) { + out.put(String.valueOf(flat.get(i)), String.valueOf(flat.get(i + 1))); + } + return out; + } + + private static Reservation decode(Map raw) { + // Read field by field rather than binding the whole object to a map of string maps. The + // SDKs write the check's request body here, which carries more than company and user, so + // one sibling adding a field would otherwise fail the whole bind and silently drop the + // entity keys a recovered hold needs to bill its usage. + JsonNode ctx = null; + String encoded = raw.get("evalCtx"); + if (encoded != null && !encoded.isEmpty()) { + try { + ctx = ObjectMappers.JSON_MAPPER.readTree(encoded); + } catch (Exception e) { + ctx = null; + } + } + return new Reservation( + raw.get("id"), + raw.get("leaseId"), + CreditLeaseMode.CLIENT, + raw.get("companyId"), + raw.get("creditTypeId"), + raw.get("eventSubtype"), + CreditAmounts.parse(raw.get("quantityReserved"), 0), + CreditAmounts.parse(raw.get("creditsReserved"), 0), + CreditAmounts.parse(raw.get("consumptionRate"), 0), + Instant.ofEpochMilli((long) CreditAmounts.parse(raw.get("expiresAt"), 0)), + entityKeys(ctx, "company"), + entityKeys(ctx, "user")); + } + + /** One entity's keys out of a decoded eval context, null when it carries none. */ + private static Map entityKeys(JsonNode ctx, String field) { + if (ctx == null) { + return null; + } + JsonNode node = ctx.get(field); + if (node == null || !node.isObject()) { + return null; + } + Map keys = new LinkedHashMap<>(); + Iterator> fields = node.fields(); + while (fields.hasNext()) { + Map.Entry entry = fields.next(); + if (entry.getValue().isTextual()) { + keys.put(entry.getKey(), entry.getValue().asText()); + } + } + return keys; + } + + private static void ignoringFailures(Runnable step) { + try { + step.run(); + } catch (RuntimeException e) { + // Index bookkeeping only: a failure here is reconciled by the next sweep, and must + // not abort the settle that is already claimed. + } + } +} diff --git a/src/main/java/com/schematic/api/credits/Reservation.java b/src/main/java/com/schematic/api/credits/Reservation.java new file mode 100644 index 0000000..db27bba --- /dev/null +++ b/src/main/java/com/schematic/api/credits/Reservation.java @@ -0,0 +1,144 @@ +package com.schematic.api.credits; + +import java.time.Instant; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; + +/** + * One credit hold, issued by {@code check()} and settled by {@code trackWithReservation()}. + * + *

In client mode the hold is a local carve-out of a lease; in server mode the API holds the + * credits and the settling track event routes by reservation id. + */ +public final class Reservation { + + private final String id; + private final String leaseId; + private final CreditLeaseMode mode; + private final String companyId; + private final String creditTypeId; + private final String eventSubtype; + private final double quantityReserved; + private final double creditsReserved; + private final double consumptionRate; + private final Instant expiresAt; + private final Map company; + private final Map user; + + public Reservation( + String id, + String leaseId, + CreditLeaseMode mode, + String companyId, + String creditTypeId, + String eventSubtype, + double quantityReserved, + double creditsReserved, + double consumptionRate, + Instant expiresAt, + Map company, + Map user) { + this.id = id; + this.leaseId = leaseId; + this.mode = mode; + this.companyId = companyId; + this.creditTypeId = creditTypeId; + this.eventSubtype = eventSubtype; + this.quantityReserved = quantityReserved; + this.creditsReserved = creditsReserved; + this.consumptionRate = consumptionRate; + this.expiresAt = expiresAt; + this.company = copy(company); + this.user = copy(user); + } + + private static Map copy(Map keys) { + if (keys == null || keys.isEmpty()) { + return null; + } + return Collections.unmodifiableMap(new LinkedHashMap<>(keys)); + } + + /** Opaque reservation id. */ + public String getId() { + return id; + } + + /** + * The lease this hold was carved from, which pins its refunds. In server mode there is no + * lease, so this mirrors {@link #getId()} and the field stays populated. + */ + public String getLeaseId() { + return leaseId; + } + + /** Where the hold lives. */ + public CreditLeaseMode getMode() { + return mode; + } + + public String getCompanyId() { + return companyId; + } + + public String getCreditTypeId() { + return creditTypeId; + } + + /** The event the settling track is billed as. */ + public String getEventSubtype() { + return eventSubtype; + } + + /** The caller-declared usage, in event units. */ + public double getQuantityReserved() { + return quantityReserved; + } + + /** + * {@code ceil(quantityReserved) * consumptionRate}. Whole event units, since a fraction of an + * event is not something the server bills, so this is what the settle will charge. + */ + public double getCreditsReserved() { + return creditsReserved; + } + + /** The consumption rate at the time the hold was issued. */ + public double getConsumptionRate() { + return consumptionRate; + } + + /** When the hold expires and is swept back to the lease. */ + public Instant getExpiresAt() { + return expiresAt; + } + + /** + * The company keys the hold was issued for, threaded onto the track event so the server + * attributes the usage to the same company. Null when the check named none. + */ + public Map getCompany() { + return company; + } + + /** The user keys the hold was issued for. Null when the check named none. */ + public Map getUser() { + return user; + } + + @Override + public String toString() { + return "Reservation{id=" + id + + ", leaseId=" + leaseId + + ", mode=" + mode + + ", companyId=" + companyId + + ", creditTypeId=" + creditTypeId + + ", eventSubtype=" + eventSubtype + + ", quantityReserved=" + quantityReserved + + ", creditsReserved=" + creditsReserved + + ", consumptionRate=" + consumptionRate + + ", expiresAt=" + expiresAt + + "}"; + } +} diff --git a/src/main/java/com/schematic/api/credits/ReservationRefunder.java b/src/main/java/com/schematic/api/credits/ReservationRefunder.java new file mode 100644 index 0000000..5cfafed --- /dev/null +++ b/src/main/java/com/schematic/api/credits/ReservationRefunder.java @@ -0,0 +1,11 @@ +package com.schematic.api.credits; + +/** + * The slice of {@link LeaseStore} a reservation store needs. Narrowing it here is what lets the + * unspent-slice refund stay an ordinary single-key step rather than a cross-key script, and lets + * a test interpose on the claim-then-refund window. + */ +public interface ReservationRefunder { + + void refund(String companyId, String creditTypeId, double credits, String pinLeaseId); +} diff --git a/src/main/java/com/schematic/api/credits/ReservationSettlement.java b/src/main/java/com/schematic/api/credits/ReservationSettlement.java new file mode 100644 index 0000000..cccb278 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/ReservationSettlement.java @@ -0,0 +1,88 @@ +package com.schematic.api.credits; + +import com.schematic.api.types.EventBodyTrack; + +/** Consumes a client-mode hold against its lease and builds the event that bills it. */ +public final class ReservationSettlement { + + /** What a settle did locally, and what it owes the server. */ + public static final class SettleOutcome { + private final EventBodyTrack track; + private final boolean settledLocally; + + SettleOutcome(EventBodyTrack track, boolean settledLocally) { + this.track = track; + this.settledLocally = settledLocally; + } + + /** The billing event to emit. */ + public EventBodyTrack getTrack() { + return track; + } + + /** + * 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. + */ + public boolean isSettledLocally() { + return settledLocally; + } + } + + private ReservationSettlement() {} + + /** + * Settles a hold and builds its billing event. + * + *

The event comes from the caller-held reservation 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. + * + *

The debit rounds the actual up to whole event units, the same way the hold that it settles + * was sized. A fraction of an event is not something the server bills, so a raw-quantity debit + * would move the lease by less than the Track event charges, and the two would drift apart over + * a session. + */ + public static SettleOutcome settle(ReservationStore reservations, Reservation reservation, double actualQuantity) { + Double claimed = + reservations.consume(reservation.getId(), Math.ceil(actualQuantity) * reservation.getConsumptionRate()); + return new SettleOutcome(buildTrackEvent(reservation, actualQuantity), claimed != null); + } + + /** + * Builds the event that settles a hold, from the reservation alone. + * + *

Kept free of store access so the client can still bill the usage when the local settle + * fails against an unreachable store. In client mode the lease id routes the server-side + * consumption through the lease's sub-ledger instead of decrementing a grant the acquire + * already pre-debited; in server mode the hold lives on the server and settles by its own id, + * and the lease id is never sent, since the server prefers it when both are set. + */ + public static EventBodyTrack buildTrackEvent(Reservation reservation, double actualQuantity) { + EventBodyTrack._FinalStage event = + EventBodyTrack.builder().event(reservation.getEventSubtype()).quantity(settleQuantity(actualQuantity)); + if (reservation.getMode() == CreditLeaseMode.SERVER) { + event.reservationId(reservation.getId()); + } else { + event.leaseId(reservation.getLeaseId()); + } + if (reservation.getCompany() != null) { + event.company(reservation.getCompany()); + } + if (reservation.getUser() != null) { + event.user(reservation.getUser()); + } + return event.build(); + } + + /** + * Casts a settled usage onto the integer a track event records. The wire field is an integer, + * so a partial unit is billed as a whole one, which is the same rounding the hold and the debit + * already apply. + */ + public static long settleQuantity(double actualQuantity) { + return (long) Math.ceil(actualQuantity); + } +} diff --git a/src/main/java/com/schematic/api/credits/ReservationStore.java b/src/main/java/com/schematic/api/credits/ReservationStore.java new file mode 100644 index 0000000..5e9cd46 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/ReservationStore.java @@ -0,0 +1,45 @@ +package com.schematic.api.credits; + +/** + * Holds the open credit holds carved out of leases. + * + *

{@link #add} does not debit: the debit already landed in {@link LeaseStore#tryReserve}, and + * that ordering is what bounds a crash to a leaked hold rather than a double-spend. + */ +public interface ReservationStore { + + /** Registers a reservation. Idempotent on id. */ + void add(Reservation reservation); + + /** Looks up a reservation, or returns null once it has been claimed or swept. */ + Reservation get(String id); + + /** + * Claims a reservation exactly once and refunds its unspent slice. + * + *

The claim is atomic and comes first: a racing settle or sweep finds nothing to claim and + * refunds nothing. On a successful claim {@code creditsConsumed} is clamped to + * {@code [0, creditsReserved]}, the remainder is refunded to the lease (pinned to the + * reservation's lease), and the clamped figure is returned. Returns null when there was + * nothing to claim. A crash between the claim and the refund loses the refund; it never + * double-refunds. + */ + Double consume(String id, double creditsConsumed); + + /** + * Sums {@code creditsReserved} across the slot's open reservations. A hold counts exactly + * while it is in the table, so {@code localRemainingCredits + reservedCredits} stays exact + * between operations. + */ + double reservedCredits(String companyId, String creditTypeId); + + /** + * Removes every reservation past its TTL, refunding each full hold, and returns how many it + * swept. Refunds are pinned to the originating lease, so a hold carved from a lease that has + * since expired is dropped rather than credited to its successor. + */ + int sweepExpired(); + + /** The open reservations across every slot. */ + int count(); +} diff --git a/src/main/java/com/schematic/api/credits/ReserveResult.java b/src/main/java/com/schematic/api/credits/ReserveResult.java new file mode 100644 index 0000000..bdbf671 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/ReserveResult.java @@ -0,0 +1,29 @@ +package com.schematic.api.credits; + +/** + * A successful {@link LeaseStore#tryReserve}: the post-debit balance plus the id of the lease the + * credits actually came out of. + */ +public final class ReserveResult { + + private final double balance; + private final String leaseId; + + public ReserveResult(double balance, String leaseId) { + this.balance = balance; + this.leaseId = leaseId; + } + + /** The balance left on the lease after the debit. */ + public double getBalance() { + return balance; + } + + /** + * The lease the debit landed on. A caller pins its reservation to this, never to the lease + * its acquire handed back: the slot's lease can be replaced in between. + */ + public String getLeaseId() { + return leaseId; + } +} diff --git a/src/main/java/com/schematic/api/credits/ResolvedLeaseConfig.java b/src/main/java/com/schematic/api/credits/ResolvedLeaseConfig.java new file mode 100644 index 0000000..0b8b100 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/ResolvedLeaseConfig.java @@ -0,0 +1,35 @@ +package com.schematic.api.credits; + +import java.time.Duration; + +/** The lease knobs for one credit type, after the credit type's override and the defaults. */ +public final class ResolvedLeaseConfig { + + private final Duration leaseDuration; + private final Duration reservationTtl; + private final double leaseSize; + private final double lowWaterMark; + + public ResolvedLeaseConfig(Duration leaseDuration, Duration reservationTtl, double leaseSize, double lowWaterMark) { + this.leaseDuration = leaseDuration; + this.reservationTtl = reservationTtl; + this.leaseSize = leaseSize; + this.lowWaterMark = lowWaterMark; + } + + public Duration getLeaseDuration() { + return leaseDuration; + } + + public Duration getReservationTtl() { + return reservationTtl; + } + + public double getLeaseSize() { + return leaseSize; + } + + public double getLowWaterMark() { + return lowWaterMark; + } +} diff --git a/src/main/java/com/schematic/api/credits/ServerCreditCheck.java b/src/main/java/com/schematic/api/credits/ServerCreditCheck.java new file mode 100644 index 0000000..7e62479 --- /dev/null +++ b/src/main/java/com/schematic/api/credits/ServerCreditCheck.java @@ -0,0 +1,283 @@ +package com.schematic.api.credits; + +import com.schematic.api.core.BaseSchematicApiException; +import com.schematic.api.core.ObjectMappers; +import com.schematic.api.core.RequestOptions; +import com.schematic.api.logger.SchematicLogger; +import com.schematic.api.resources.credits.CreditsClient; +import com.schematic.api.resources.features.FeaturesClient; +import com.schematic.api.resources.features.requests.CheckAndReserveFlagRequestBody; +import com.schematic.api.types.ApiError; +import com.schematic.api.types.CheckAndReserveFlagResponseData; +import com.schematic.api.types.FeatureEntitlement; +import com.schematic.api.types.FlagCheckReservationResponseData; +import com.schematic.api.types.PreflightRequestBody; +import com.schematic.api.types.RulesengineFeatureEntitlement; +import java.time.Clock; +import java.time.Duration; +import java.time.OffsetDateTime; +import java.time.ZoneOffset; +import java.util.UUID; +import java.util.concurrent.Callable; +import java.util.concurrent.TimeUnit; +import java.util.function.BooleanSupplier; + +/** + * Gates one check server-side. A single check-and-reserve call does everything the client flow + * 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. + * + *

No flag_check event is reported here: the server logs the flag check for check-and-reserve + * itself, the way the REST check path does. + */ +public final class ServerCreditCheck { + + // Mirrors the reason the API returns on a 200 with value false for the same denial, so a + // caller matching on the reason has one string to match either way. + private static final String INSUFFICIENT_CREDITS_REASON = "Insufficient credits"; + + private final FeaturesClient features; + private final CreditsClient credits; + private final SchematicLogger logger; + private final Duration reservationTtl; + private final Clock clock; + + public ServerCreditCheck( + FeaturesClient features, + CreditsClient credits, + SchematicLogger logger, + Duration reservationTtl, + Clock clock) { + this.features = features; + this.credits = credits; + this.logger = logger; + this.reservationTtl = reservationTtl != null ? reservationTtl : CreditLeaseDefaults.RESERVATION_TTL; + this.clock = clock != null ? clock : Clock.systemUTC(); + } + + /** + * Runs the server-gated check. {@code getDefault} answers the caller's default for this flag, + * which the fail-open branch returns. {@code fallback} is the plain flag check, for the asks + * that need no hold at all. + */ + public CheckResult check( + CheckRequest request, Duration timeout, BooleanSupplier getDefault, Callable fallback) { + // 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. + if (!CreditAmounts.isValidQuantity(request.getUsage())) { + error("Server reservation: invalid usage " + request.getUsage() + " for flag " + request.getFlagKey() + + "; must be a finite, non-negative number"); + return failureResult(request, getDefault, "invalid_usage"); + } + + // Nothing to hold. The plain check still carries the preflight. + if (request.getUsage() == 0) { + debug("Server reservation: usage is 0 for flag " + request.getFlagKey() + + ", nothing to reserve, using a plain check"); + return fallBack(fallback); + } + + CheckAndReserveFlagRequestBody.Builder body = CheckAndReserveFlagRequestBody.builder() + .quantity(request.getUsage()) + .expiresAt(OffsetDateTime.ofInstant(clock.instant().plus(reservationTtl), ZoneOffset.UTC)); + if (!request.getCompany().isEmpty()) { + body.company(request.getCompany()); + } + if (!request.getUser().isEmpty()) { + body.user(request.getUser()); + } + PreflightOptions preflight = PreflightOptions.fromUsage(request.getUsage(), request.getEventSubtype()); + PreflightRequestBody preflightBody = preflight != null ? preflight.toRequestBody() : null; + if (preflightBody != null) { + body.preflight(preflightBody); + } + // One key per check, minted before the call so the transport's retries resend the same + // one: a 502 from a load balancer after the API committed the hold then collapses onto + // that hold instead of taking a second one and parking the first until its TTL. + body.idempotencyKey(UUID.randomUUID().toString()); + RequestOptions.Builder options = RequestOptions.builder(); + if (timeout != null) { + options.timeout(CreditAmounts.millisAsInt(timeout), TimeUnit.MILLISECONDS); + } + + CheckAndReserveFlagResponseData data; + try { + data = features.checkAndReserveFlag(request.getFlagKey(), body.build(), options.build()) + .getData(); + } catch (BaseSchematicApiException 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. + if (e.statusCode() == 402) { + return new CheckResult( + false, false, INSUFFICIENT_CREDITS_REASON, request.getFlagKey(), null, null, null, message(e)); + } + error("Server reservation: check-and-reserve for flag " + request.getFlagKey() + " failed: " + e); + return failureResult(request, getDefault, "server_reservation_failed"); + } catch (RuntimeException e) { + error("Server reservation: check-and-reserve for flag " + request.getFlagKey() + " failed: " + e); + return failureResult(request, getDefault, "server_reservation_failed"); + } + if (data == null) { + // A 200 with an empty body answers nothing, so resolve it through the caller's + // contract rather than letting the read below surface as a null dereference. + error("Server reservation: check-and-reserve for flag " + request.getFlagKey() + " carried no data"); + return failureResult(request, getDefault, "server_reservation_failed"); + } + + RulesengineFeatureEntitlement entitlement = + toRulesengineEntitlement(data.getEntitlement().orElse(null)); + CheckResult base = new CheckResult( + data.getValue(), + data.getValue(), + data.getReason(), + orElse(data.getFlag(), request.getFlagKey()), + data.getFlagId().orElse(null), + entitlement, + null, + data.getError().orElse(null)); + + // No hold 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 nothing to release. + FlagCheckReservationResponseData held = data.getReservation().orElse(null); + if (!data.getValue() || held == null) { + return base; + } + + // The settling 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, so release it now rather than leaving credits parked until the TTL. + String eventSubtype = request.getEventSubtype(); + if (eventSubtype == null || eventSubtype.isEmpty()) { + eventSubtype = held.getEventSubtype().orElse(null); + } + if (eventSubtype == null || eventSubtype.isEmpty()) { + error("Server reservation: reservation " + held.getId() + " for flag " + request.getFlagKey() + + " has no event subtype, releasing it, since it could never be settled"); + try { + credits.releaseCreditReservation(held.getId()); + } catch (RuntimeException e) { + warn("Server reservation: failed to release " + held.getId() + " (" + e + + "); its hold is refunded when it expires"); + } + if (!request.isFailOpen()) { + return failureResult(request, getDefault, "missing_event_subtype"); + } + // 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. + return new CheckResult( + base.isAllowed(), + base.getValue(), + base.getReason(), + base.getFlagKey(), + base.getFlagId(), + base.getEntitlement(), + null, + "missing_event_subtype"); + } + + Reservation reservation = new Reservation( + held.getId(), + // No lease exists in server mode; mirror the id so the field stays populated and a + // handle round-trips through code that reads it. + held.getId(), + CreditLeaseMode.SERVER, + held.getCompanyId(), + held.getCreditTypeId(), + eventSubtype, + held.getQuantityReserved(), + held.getCreditsReserved(), + held.getConsumptionRate(), + held.getExpiresAt().toInstant(), + request.getCompany(), + request.getUser()); + return new CheckResult( + true, + true, + data.getReason(), + orElse(data.getFlag(), request.getFlagKey()), + data.getFlagId().orElse(null), + entitlement, + reservation, + null); + } + + /** + * Resolves 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. + */ + private static CheckResult failureResult(CheckRequest request, BooleanSupplier getDefault, String reason) { + if (!request.isFailOpen()) { + return new CheckResult(false, false, reason, request.getFlagKey(), null, null, null, reason); + } + boolean value = getDefault != null && getDefault.getAsBoolean(); + return new CheckResult(value, value, reason + "_fail_open", request.getFlagKey(), null, null, null, reason); + } + + /** + * The API and the rules engine describe an entitlement with the same fields under two + * generated types, so the server's answer is remapped through its JSON rather than dropped. A + * shape the mapper cannot bridge costs the caller the entitlement detail, never the verdict. + */ + private RulesengineFeatureEntitlement toRulesengineEntitlement(FeatureEntitlement entitlement) { + if (entitlement == null) { + return null; + } + try { + return ObjectMappers.JSON_MAPPER.convertValue(entitlement, RulesengineFeatureEntitlement.class); + } catch (RuntimeException e) { + debug("Server reservation: could not read the entitlement off the response: " + e); + return null; + } + } + + private static String message(BaseSchematicApiException e) { + Object body = e.body(); + if (body instanceof ApiError) { + String error = ((ApiError) body).getError(); + if (error != null && !error.isEmpty()) { + return error; + } + } + return e.getMessage(); + } + + private static CheckResult fallBack(Callable fallback) { + try { + return fallback.call(); + } catch (Exception e) { + throw new IllegalStateException("plain flag check failed", e); + } + } + + private static String orElse(String value, String fallback) { + return value == null || value.isEmpty() ? fallback : value; + } + + private void debug(String message) { + if (logger != null) { + logger.debug(message); + } + } + + private void warn(String message) { + if (logger != null) { + logger.warn(message); + } + } + + private void error(String message) { + if (logger != null) { + logger.error(message); + } + } +} diff --git a/src/main/java/com/schematic/api/datastream/CheckFlagOptions.java b/src/main/java/com/schematic/api/datastream/CheckFlagOptions.java new file mode 100644 index 0000000..2773fc9 --- /dev/null +++ b/src/main/java/com/schematic/api/datastream/CheckFlagOptions.java @@ -0,0 +1,57 @@ +package com.schematic.api.datastream; + +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; + +/** + * The preflight a local evaluation answers: what the call being gated is about to cost, before it + * has been recorded. Serialized into the {@code options} envelope the rules engine reads. + */ +public final class CheckFlagOptions { + + private final Map creditCost; + private final Double usage; + private final String eventSubtype; + private final Double eventQuantity; + + private CheckFlagOptions(Map creditCost, Double usage, String eventSubtype, Double eventQuantity) { + this.creditCost = creditCost == null || creditCost.isEmpty() + ? null + : Collections.unmodifiableMap(new LinkedHashMap<>(creditCost)); + this.usage = usage; + this.eventSubtype = eventSubtype; + this.eventQuantity = eventQuantity; + } + + /** Prices the action per credit type, bypassing the engine's own quantity times rate arithmetic. */ + public static CheckFlagOptions creditCost(Map creditCost) { + return new CheckFlagOptions(creditCost, null, null, null); + } + + /** A simulated quantity, unscoped. */ + public static CheckFlagOptions usage(double usage) { + return new CheckFlagOptions(null, usage, null, null); + } + + /** A simulated quantity scoped to the event subtype whose condition should answer it. */ + public static CheckFlagOptions eventUsage(String eventSubtype, double quantity) { + return new CheckFlagOptions(null, null, eventSubtype, quantity); + } + + public Map getCreditCost() { + return creditCost; + } + + public Double getUsage() { + return usage; + } + + public String getEventSubtype() { + return eventSubtype; + } + + public Double getEventQuantity() { + return eventQuantity; + } +} diff --git a/src/main/java/com/schematic/api/datastream/DataStreamClient.java b/src/main/java/com/schematic/api/datastream/DataStreamClient.java index 1390106..fb84878 100644 --- a/src/main/java/com/schematic/api/datastream/DataStreamClient.java +++ b/src/main/java/com/schematic/api/datastream/DataStreamClient.java @@ -33,6 +33,7 @@ import okhttp3.OkHttpClient; import okhttp3.Request; import okhttp3.Response; +import redis.clients.jedis.JedisPooled; /** * High-level DataStream client that manages WebSocket connections (or replicator mode), @@ -60,6 +61,8 @@ public class DataStreamClient implements Closeable { private final SchematicLogger logger; private final ObjectMapper objectMapper; private final RulesEngine rulesEngine; + private final JedisPooled redisClient; + private final String redisKeyPrefix; // Typed entity caches private final CacheProvider flagCache; @@ -113,11 +116,12 @@ public DataStreamClient( this.rulesEngine = rulesEngine; // Build cache providers via factory: custom > Redis > local - redis.clients.jedis.JedisPooled redisClient = - DataStreamCacheFactory.buildRedisClient(options.getRedisCacheConfig()); + JedisPooled redisClient = DataStreamCacheFactory.buildRedisClient(options.getRedisCacheConfig()); String keyPrefix = options.getRedisCacheConfig() != null ? options.getRedisCacheConfig().getKeyPrefix() : "schematic:"; + this.redisClient = redisClient; + this.redisKeyPrefix = keyPrefix; this.flagCache = DataStreamCacheFactory.buildFlagCache(options, redisClient, keyPrefix); this.companyCache = DataStreamCacheFactory.buildCompanyCache(options, redisClient, keyPrefix); this.userCache = DataStreamCacheFactory.buildUserCache(options, redisClient, keyPrefix); @@ -146,6 +150,19 @@ public void start() { } } + /** + * The Redis client the caches were configured with, or null when they are local. Credit leases + * reuse it so an existing Redis setup gates them across pods with no second client to wire up. + */ + public JedisPooled getRedisClient() { + return redisClient; + } + + /** The key prefix the caches were configured with, which credit lease keys are built under. */ + public String getRedisKeyPrefix() { + return redisKeyPrefix; + } + /** * Returns whether the datastream is connected and ready for flag checks. */ @@ -156,6 +173,14 @@ public boolean isConnected() { return wsClient != null && wsClient.isReady(); } + /** + * Returns whether a rules engine is loaded and able to evaluate locally. False when the WASM + * engine failed to load, where every local evaluation throws and checks fall back to the API. + */ + public boolean hasRulesEngine() { + return rulesEngine != null && rulesEngine.isInitialized(); + } + /** * Returns whether this client is running in replicator mode. */ @@ -167,6 +192,17 @@ public boolean isReplicatorMode() { * Checks a flag using cached datastream data and the rules engine. */ public RulesengineCheckFlagResult checkFlag(String flagKey, Map company, Map user) { + return checkFlag(flagKey, company, user, null); + } + + /** + * Checks a flag using cached datastream data and the rules engine, with a preflight: what the + * call being gated is about to cost, before it has been recorded. + * + * @param preflight the preflight, or null for none + */ + public RulesengineCheckFlagResult checkFlag( + String flagKey, Map company, Map user, CheckFlagOptions preflight) { // Step 1: Get flag from cache RulesengineFlag flag = flagCache.get(flagCacheKey(flagKey)); if (flag == null) { @@ -202,13 +238,13 @@ public RulesengineCheckFlagResult checkFlag(String flagKey, Map // Step 3: Replicator mode - evaluate with whatever we have if (options.isReplicatorMode()) { - return evaluateFlag(flag, cachedCompany, cachedUser); + return evaluateFlag(flag, cachedCompany, cachedUser, preflight); } // Step 4: Direct mode - if all needed data is cached, evaluate immediately if ((!needsCompany || cachedCompany != null) && (!needsUser || cachedUser != null)) { log("debug", "All required resources found in cache for flag " + flagKey); - return evaluateFlag(flag, cachedCompany, cachedUser); + return evaluateFlag(flag, cachedCompany, cachedUser, preflight); } // Step 5: Direct mode - fetch missing entities via datastream and wait for response @@ -223,14 +259,15 @@ public RulesengineCheckFlagResult checkFlag(String flagKey, Map cachedUser = getUser(user); } - return evaluateFlag(flag, cachedCompany, cachedUser); + return evaluateFlag(flag, cachedCompany, cachedUser, preflight); } /** * Fetches a company via the datastream WebSocket, waiting for the response with a timeout. - * Deduplicates concurrent requests for the same entity. + * Deduplicates concurrent requests for the same entity. Returns null when the entity never + * arrives. */ - private RulesengineCompany getCompany(Map keys) { + public RulesengineCompany getCompany(Map keys) { // Check cache first RulesengineCompany cached = getCachedCompany(keys); if (cached != null) { @@ -261,7 +298,7 @@ private RulesengineCompany getCompany(Map keys) { * Fetches a user via the datastream WebSocket, waiting for the response with a timeout. * Deduplicates concurrent requests for the same entity. */ - private RulesengineUser getUser(Map keys) { + public RulesengineUser getUser(Map keys) { // Check cache first RulesengineUser cached = getCachedUser(keys); if (cached != null) { @@ -400,11 +437,50 @@ private void cleanupPendingUserRequests(Map keys, CompletableFut } } + /** + * Evaluates a flag with a preflight, propagating a failure instead of substituting the flag's + * default value. A credit-gated check needs that difference: a default returned as a verdict + * would leave the hold it took standing on an evaluation that never ran. + * + * @param preflight the preflight, or null for none + */ + public RulesengineCheckFlagResult evaluateFlagWithOptions( + RulesengineFlag flag, RulesengineCompany company, RulesengineUser user, CheckFlagOptions preflight) + throws Exception { + if (rulesEngine == null || !rulesEngine.isInitialized()) { + throw new DataStreamException("Rules engine not available for flag " + flag.getKey()); + } + RulesengineCheckFlagResult result = evaluate(flag, company, user, preflight); + return RulesengineCheckFlagResult.builder() + .from(result) + .companyId(result.getCompanyId().orElse(company != null ? company.getId() : null)) + .userId(result.getUserId().orElse(user != null ? user.getId() : null)) + .build(); + } + + /** + * Calls the rules engine, taking the preflight-free entry point when there is no preflight, so + * an engine that implements only that one still answers. + */ + private RulesengineCheckFlagResult evaluate( + RulesengineFlag flag, RulesengineCompany company, RulesengineUser user, CheckFlagOptions preflight) + throws Exception { + if (preflight == null) { + return rulesEngine.checkFlag(flag, company, user); + } + return rulesEngine.checkFlag(flag, company, user, preflight); + } + /** * Evaluates a flag using the rules engine. Falls back to the flag's default value * if the rules engine is not available. */ RulesengineCheckFlagResult evaluateFlag(RulesengineFlag flag, RulesengineCompany company, RulesengineUser user) { + return evaluateFlag(flag, company, user, null); + } + + RulesengineCheckFlagResult evaluateFlag( + RulesengineFlag flag, RulesengineCompany company, RulesengineUser user, CheckFlagOptions preflight) { boolean defaultValue = flag.getDefaultValue(); String flagKey = flag.getKey(); String flagId = flag.getId(); @@ -413,7 +489,7 @@ RulesengineCheckFlagResult evaluateFlag(RulesengineFlag flag, RulesengineCompany if (rulesEngine != null && rulesEngine.isInitialized()) { try { - RulesengineCheckFlagResult result = rulesEngine.checkFlag(flag, company, user); + RulesengineCheckFlagResult result = evaluate(flag, company, user, preflight); // The WASM engine returns a complete result — use it directly, // enriching with IDs from context if the engine didn't set them return RulesengineCheckFlagResult.builder() diff --git a/src/main/java/com/schematic/api/datastream/RulesEngine.java b/src/main/java/com/schematic/api/datastream/RulesEngine.java index fe1c9ba..31f2fb0 100644 --- a/src/main/java/com/schematic/api/datastream/RulesEngine.java +++ b/src/main/java/com/schematic/api/datastream/RulesEngine.java @@ -35,4 +35,17 @@ public interface RulesEngine { */ RulesengineCheckFlagResult checkFlag(RulesengineFlag flag, RulesengineCompany company, RulesengineUser user) throws Exception; + + /** + * Evaluates a flag with a preflight: what the call being gated is about to cost, before it has + * been recorded. An engine that cannot answer a preflight evaluates without one, which is the + * same verdict every caller got before preflights existed. + * + * @param options the preflight, or null for none + */ + default RulesengineCheckFlagResult checkFlag( + RulesengineFlag flag, RulesengineCompany company, RulesengineUser user, CheckFlagOptions options) + throws Exception { + return checkFlag(flag, company, user); + } } diff --git a/src/main/java/com/schematic/api/datastream/WasmRulesEngine.java b/src/main/java/com/schematic/api/datastream/WasmRulesEngine.java index 8e3627c..8b6e861 100644 --- a/src/main/java/com/schematic/api/datastream/WasmRulesEngine.java +++ b/src/main/java/com/schematic/api/datastream/WasmRulesEngine.java @@ -161,6 +161,13 @@ public String getVersionKey() { @Override public RulesengineCheckFlagResult checkFlag(RulesengineFlag flag, RulesengineCompany company, RulesengineUser user) throws Exception { + return checkFlag(flag, company, user, null); + } + + @Override + public RulesengineCheckFlagResult checkFlag( + RulesengineFlag flag, RulesengineCompany company, RulesengineUser user, CheckFlagOptions options) + throws Exception { if (!initialized) { throw new IllegalStateException("WASM rules engine not initialized"); } @@ -180,6 +187,10 @@ public RulesengineCheckFlagResult checkFlag(RulesengineFlag flag, RulesengineCom if (user != null) { envelope.set("user", sanitize(mapper.valueToTree(user), "user")); } + ObjectNode preflight = preflightNode(mapper, options); + if (preflight != null) { + envelope.set("options", preflight); + } String inputJson = mapper.writeValueAsString(envelope); String resultJson = callWasm(inputJson); @@ -192,6 +203,41 @@ public RulesengineCheckFlagResult checkFlag(RulesengineFlag flag, RulesengineCom return mapper.treeToValue(snakeNode, RulesengineCheckFlagResult.class); } + /** + * Writes a preflight quantity as the integer the engine reads. The engine's quantity fields + * are integers, and anything with a decimal point, a trailing {@code .0} included, fails to + * deserialize there and turns the whole evaluation into an error. A fraction rounds up rather + * than down, the same direction the request body takes: a preflight asks an upper bound, and + * the check must not pass on less usage than the call is about to record. + */ + private static void putQuantity(ObjectNode node, String name, double quantity) { + node.put(name, (long) Math.ceil(quantity)); + } + + /** + * Builds the snake_case {@code options} envelope the engine reads, or null when there is no + * preflight to declare. + */ + private static ObjectNode preflightNode(ObjectMapper mapper, CheckFlagOptions options) { + if (options == null) { + return null; + } + ObjectNode node = mapper.createObjectNode(); + if (options.getCreditCost() != null) { + node.set("credit_cost", mapper.valueToTree(options.getCreditCost())); + } + if (options.getUsage() != null) { + putQuantity(node, "usage", options.getUsage()); + } + if (options.getEventSubtype() != null && options.getEventQuantity() != null) { + ObjectNode eventUsage = mapper.createObjectNode(); + eventUsage.put("event_subtype", options.getEventSubtype()); + putQuantity(eventUsage, "quantity", options.getEventQuantity()); + node.set("event_usage", eventUsage); + } + return node.size() == 0 ? null : node; + } + /** * Calls into the WASM runtime. This method is synchronized because the WASM instance * uses shared linear memory — concurrent calls would corrupt each other's data. diff --git a/src/test/java/com/schematic/api/TestReadme.java b/src/test/java/com/schematic/api/TestReadme.java index 3015413..17917e0 100644 --- a/src/test/java/com/schematic/api/TestReadme.java +++ b/src/test/java/com/schematic/api/TestReadme.java @@ -7,10 +7,15 @@ import com.fasterxml.jackson.databind.JsonNode; import com.schematic.api.cache.LocalCache; import com.schematic.api.core.ObjectMappers; +import com.schematic.api.credits.CheckOptions; +import com.schematic.api.credits.CheckResult; +import com.schematic.api.credits.CreditLeaseConfig; +import com.schematic.api.credits.OnAcquireFailure; import com.schematic.api.logger.SchematicLogger; import com.schematic.api.resources.companies.CompaniesClient; import com.schematic.api.resources.companies.types.UpsertCompanyResponse; import com.schematic.api.types.CompanyDetailResponseData; +import com.schematic.api.types.EventBodyIdentifyCompany; import com.schematic.api.types.UpsertCompanyRequestBody; import java.time.Duration; import java.time.OffsetDateTime; @@ -129,4 +134,48 @@ void testClientWithOfflineModeAndDefaults() { assertTrue(schematic.isOffline()); assertTrue(schematic.checkFlag("some-flag-key", null, null)); } + + @Test + void testCreditLeaseCheckAndTrack() { + // Test the credit lease examples from README. Offline keeps the example off the network; + // the surface it exercises is the one the README shows. + Schematic schematic = Schematic.builder() + .apiKey("test_api_key") + .offline(true) + .creditLeases(CreditLeaseConfig.builder() + .defaultReservationTtl(Duration.ofSeconds(60)) + .build()) + .build(); + + Map company = new HashMap<>(); + company.put("id", "your-company-id"); + + CheckResult result = schematic.check( + "inference", + company, + null, + CheckOptions.builder() + .usage(100) + .eventSubtype("inference_tokens") + .onAcquireFailure(OnAcquireFailure.FAIL_OPEN) + .build()); + + assertFalse(result.isAllowed()); + assertNull(result.getReservation()); + + // A check can allow without a hold, and that usage still has to be tracked. + schematic.track("inference_tokens", company, null, null, 42L); + schematic.prewarm(company, Collections.singletonList("credit-type-id")); + + schematic.identify( + Collections.singletonMap("user_id", "your-user-id"), + EventBodyIdentifyCompany.builder().keys(company).build(), + "Your User", + null, + IdentifyOptions.builder() + .prewarm(Collections.singletonList("credit-type-id")) + .build()); + + schematic.close(); + } } diff --git a/src/test/java/com/schematic/api/TestSchematic.java b/src/test/java/com/schematic/api/TestSchematic.java index dab557d..8a8fd8e 100644 --- a/src/test/java/com/schematic/api/TestSchematic.java +++ b/src/test/java/com/schematic/api/TestSchematic.java @@ -8,6 +8,16 @@ import com.schematic.api.cache.CacheProvider; import com.schematic.api.cache.LocalCache; +import com.schematic.api.core.RequestOptions; +import com.schematic.api.credits.CheckOptions; +import com.schematic.api.credits.CheckResult; +import com.schematic.api.credits.CreditLeaseConfig; +import com.schematic.api.credits.CreditLeaseDefaults; +import com.schematic.api.credits.CreditLeaseMode; +import com.schematic.api.credits.Reservation; +import com.schematic.api.credits.ReservationSettlement; +import com.schematic.api.datastream.DataStreamClient; +import com.schematic.api.datastream.DatastreamOptions; import com.schematic.api.logger.SchematicLogger; import com.schematic.api.resources.features.FeaturesClient; import com.schematic.api.resources.features.types.CheckFlagResponse; @@ -21,17 +31,22 @@ import com.schematic.api.types.EventBodyIdentifyCompany; import com.schematic.api.types.EventBodyTrack; import com.schematic.api.types.EventType; +import com.schematic.api.types.PreflightRequestBody; import com.schematic.api.types.RulesengineCheckFlagResult; +import java.lang.reflect.Field; import java.time.Duration; +import java.time.Instant; import java.time.OffsetDateTime; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @@ -228,6 +243,366 @@ void buildTrackEvent_defaultsSentAtToNowWhenOptionOmitsIt() { assertEquals("idem-1", event.getIdempotencyKey().get()); } + // --- Credit-aware check plumbing --- + + @Test + void check_ThreadsThePerCheckTimeoutToTheApiFallback() { + FeaturesClient featuresClient = mock(FeaturesClient.class); + Schematic spySchematic = spy(schematic); + when(spySchematic.features()).thenReturn(featuresClient); + + CheckFlagResponse response = CheckFlagResponse.builder() + .data(CheckFlagResponseData.builder() + .flag("test_flag") + .reason("test_reason") + .value(true) + .build()) + .build(); + when(featuresClient.checkFlag(eq("test_flag"), any(CheckFlagRequestBody.class), any(RequestOptions.class))) + .thenReturn(response); + + CheckResult result = spySchematic.check( + "test_flag", + null, + null, + CheckOptions.builder().timeout(Duration.ofMillis(250)).build()); + + assertTrue(result.isAllowed()); + ArgumentCaptor options = ArgumentCaptor.forClass(RequestOptions.class); + verify(featuresClient).checkFlag(eq("test_flag"), any(CheckFlagRequestBody.class), options.capture()); + // The caller is waiting on this call, so its timeout has to reach it. + assertEquals(250, options.getValue().getTimeout().get()); + assertEquals(TimeUnit.MILLISECONDS, options.getValue().getTimeoutTimeUnit()); + } + + @Test + void check_SendsThePreflightOnTheApiFallback() { + FeaturesClient featuresClient = mock(FeaturesClient.class); + Schematic spySchematic = spy(schematic); + when(spySchematic.features()).thenReturn(featuresClient); + when(featuresClient.checkFlag(eq("test_flag"), any(CheckFlagRequestBody.class))) + .thenReturn(apiResponse(true)); + + CheckResult result = spySchematic.check( + "test_flag", + null, + null, + CheckOptions.builder() + .usage(7.2) + .eventSubtype("inference_tokens") + .build()); + + assertTrue(result.isAllowed()); + ArgumentCaptor body = ArgumentCaptor.forClass(CheckFlagRequestBody.class); + verify(featuresClient).checkFlag(eq("test_flag"), body.capture()); + PreflightRequestBody preflight = body.getValue().getPreflight().get(); + assertEquals("inference_tokens", preflight.getEventUsage().get().getEventSubtype()); + // A preflight asks an upper-bound question, so a fractional usage rounds up. + assertEquals(8L, preflight.getEventUsage().get().getQuantity()); + } + + @Test + void check_WithAPreflightNeitherReadsNorWritesTheFlagCache() { + FeaturesClient featuresClient = mock(FeaturesClient.class); + Schematic spySchematic = spy(schematic); + when(spySchematic.features()).thenReturn(featuresClient); + // Cache a plain verdict for this flag, company and user. + when(featuresClient.checkFlag(eq("test_flag"), any(CheckFlagRequestBody.class))) + .thenReturn(apiResponse(true)); + spySchematic.checkFlag("test_flag", null, null); + for (CacheProvider provider : spySchematic.getFlagCheckCacheProviders()) { + assertNotNull(provider.get("test_flag")); + } + + when(featuresClient.checkFlag(eq("test_flag"), any(CheckFlagRequestBody.class))) + .thenReturn(apiResponse(false)); + CheckResult preflighted = spySchematic.check( + "test_flag", null, null, CheckOptions.builder().usage(5).build()); + + // The cached plain verdict answers a different question, so it is not served here. + assertFalse(preflighted.isAllowed()); + for (CacheProvider provider : spySchematic.getFlagCheckCacheProviders()) { + // And the preflighted verdict must not become the answer a plain check reads back. + assertTrue(provider.get("test_flag").getValue()); + } + } + + @Test + void check_WithoutAPreflightStillUsesTheFlagCache() { + FeaturesClient featuresClient = mock(FeaturesClient.class); + Schematic spySchematic = spy(schematic); + when(spySchematic.features()).thenReturn(featuresClient); + when(featuresClient.checkFlag(eq("test_flag"), any(CheckFlagRequestBody.class))) + .thenReturn(apiResponse(true)); + + CheckResult first = spySchematic.check("test_flag", null, null, null); + CheckResult second = spySchematic.check("test_flag", null, null, null); + + assertTrue(first.isAllowed()); + assertTrue(second.isAllowed()); + verify(featuresClient, times(1)).checkFlag(eq("test_flag"), any(CheckFlagRequestBody.class)); + for (CacheProvider provider : spySchematic.getFlagCheckCacheProviders()) { + assertNotNull(provider.get("test_flag")); + } + } + + @Test + void check_WithAZeroUsageStaysPlainAndKeepsTheFlagCache() { + FeaturesClient featuresClient = mock(FeaturesClient.class); + Schematic spySchematic = spy(schematic); + when(spySchematic.features()).thenReturn(featuresClient); + when(featuresClient.checkFlag(eq("test_flag"), any(CheckFlagRequestBody.class))) + .thenReturn(apiResponse(true)); + + CheckOptions zeroUsage = CheckOptions.builder().usage(0).build(); + CheckResult first = spySchematic.check("test_flag", null, null, zeroUsage); + CheckResult second = spySchematic.check("test_flag", null, null, zeroUsage); + + assertTrue(first.isAllowed()); + assertTrue(second.isAllowed()); + ArgumentCaptor body = ArgumentCaptor.forClass(CheckFlagRequestBody.class); + // The API treats a zero usage as no usage, so sending it would cost the check its cache + // and buy nothing: one call answers both. + verify(featuresClient, times(1)).checkFlag(eq("test_flag"), body.capture()); + assertFalse(body.getValue().getPreflight().isPresent()); + } + + @Test + void check_UsesThePerCheckDefaultWhenTheApiFallbackFails() { + FeaturesClient featuresClient = mock(FeaturesClient.class); + Schematic spySchematic = spy(schematic); + when(spySchematic.features()).thenReturn(featuresClient); + when(featuresClient.checkFlag(eq("test_flag"), any(CheckFlagRequestBody.class))) + .thenThrow(new RuntimeException("connection refused")); + + CheckResult result = spySchematic.check( + "test_flag", + null, + null, + CheckOptions.builder().usage(5).defaultValue(true).build()); + + // A caller who named a default for this check gets it wherever the check lands on one, + // not just in offline mode. + assertTrue(result.isAllowed()); + assertEquals("flag default", result.getReason()); + } + + @Test + void check_WithoutAPerCheckDefaultFallsBackToTheClientDefault() { + FeaturesClient featuresClient = mock(FeaturesClient.class); + Schematic spySchematic = spy(schematic); + when(spySchematic.features()).thenReturn(featuresClient); + when(featuresClient.checkFlag(eq("test_flag"), any(CheckFlagRequestBody.class))) + .thenThrow(new RuntimeException("connection refused")); + spySchematic.setFlagDefault("test_flag", true); + + CheckResult unnamed = spySchematic.check( + "test_flag", null, null, CheckOptions.builder().usage(5).build()); + + assertTrue(unnamed.isAllowed()); + // And a per-check default still outranks the client-wide one. + CheckResult named = spySchematic.check( + "test_flag", + null, + null, + CheckOptions.builder().usage(5).defaultValue(false).build()); + assertFalse(named.isAllowed()); + } + + private static CheckFlagResponse apiResponse(boolean value) { + return CheckFlagResponse.builder() + .data(CheckFlagResponseData.builder() + .flag("test_flag") + .reason("test_reason") + .value(value) + .build()) + .build(); + } + + @Test + void inheritFromDataStream_ResolvesEachSettingOnItsOwn() { + // An explicit lease client must not cost the caller the DataStream key prefix: a mixed + // fleet sharing those leases would then read two different key layouts. + assertEquals("lease-client", Schematic.inheritFromDataStream("lease-client", "datastream-client")); + assertEquals("datastream:", Schematic.inheritFromDataStream(null, "datastream:")); + assertNull(Schematic.inheritFromDataStream(null, null)); + } + + @Test + void trackWithReservation_FoldsTheSettleIntoTheCachedCompanyMetrics() throws Exception { + DataStreamClient dataStream = mock(DataStreamClient.class); + when(dataStream.isConnected()).thenReturn(true); + Schematic spySchematic = spy(schematic); + setDataStreamClient(spySchematic, dataStream); + + spySchematic.trackWithReservation(reservation(CreditLeaseMode.SERVER), 7, null); + + // A settle is a track, so the cached usage has to move with it. Without this a local + // evaluation right after the settle gates on a figure the stream has not pushed back yet. + ArgumentCaptor recorded = ArgumentCaptor.forClass(EventBodyTrack.class); + verify(dataStream).updateCompanyMetrics(recorded.capture()); + assertEquals("inference_tokens", recorded.getValue().getEvent()); + assertEquals(7L, recorded.getValue().getQuantity().get()); + assertEquals( + Collections.singletonMap("id", "co_1"), + recorded.getValue().getCompany().get()); + } + + @Test + void check_ResolvesTheCallersDefaultWhenTheEngineDeclinesOnTheDataStreamBranch() throws Exception { + DataStreamClient dataStream = mock(DataStreamClient.class); + when(dataStream.isConnected()).thenReturn(true); + // What the DataStream client hands back when the engine cannot answer: the flag's own + // default standing in for a verdict. + when(dataStream.checkFlag(eq("test_flag"), any(), any(), any())) + .thenReturn(RulesengineCheckFlagResult.builder() + .flagKey("test_flag") + .reason("RULES_ENGINE_UNAVAILABLE") + .value(true) + .build()); + Schematic spySchematic = spy(schematic); + setDataStreamClient(spySchematic, dataStream); + spySchematic.setFlagDefault("test_flag", false); + + CheckResult named = spySchematic.check( + "test_flag", + null, + null, + CheckOptions.builder().usage(5).defaultValue(true).build()); + CheckResult unnamed = spySchematic.check( + "test_flag", null, null, CheckOptions.builder().usage(5).build()); + + // The engine declining is the case defaultValue exists for, so it applies on the branch + // that answers most checks and not just offline and on the API fallback. + assertTrue(named.isAllowed()); + // With no caller default the registered one stands in, rather than the flag's own. + assertFalse(unnamed.isAllowed()); + } + + @Test + void check_AutoModeFallsToServerGatingWhenTheDataStreamIsGone() throws Exception { + Schematic leased = Schematic.builder() + .apiKey("test_api_key") + .logger(logger) + .datastreamOptions(DatastreamOptions.builder().build()) + .creditLeases( + CreditLeaseConfig.builder().mode(CreditLeaseMode.AUTO).build()) + .build(); + try { + FeaturesClient featuresClient = mock(FeaturesClient.class); + Schematic spySchematic = spy(leased); + when(spySchematic.features()).thenReturn(featuresClient); + // What a DataStream that failed to start leaves behind. Auto has to notice at the + // check, not hold the verdict it reached at construction. + setDataStreamClient(spySchematic, null); + + spySchematic.check( + "test_flag", + Collections.singletonMap("id", "co_1"), + null, + CheckOptions.builder().usage(5).build()); + + // Server gating, not a plain ungated check: the credits still have to be held. + verify(featuresClient).checkAndReserveFlag(eq("test_flag"), any(), any()); + } finally { + leased.close(); + } + } + + @Test + void serverReservationTtl_IsClampedBelowTheApiCapButOnlyForServerMode() throws Exception { + Duration twoHours = Duration.ofHours(2); + Duration clamped = + CreditLeaseDefaults.MAX_RESERVATION_TTL.minus(CreditLeaseDefaults.RESERVATION_TTL_SKEW_ALLOWANCE); + + try (Schematic server = leasedClient(CreditLeaseMode.SERVER, twoHours); + Schematic client = leasedClient(CreditLeaseMode.CLIENT, twoHours)) { + // The API refuses a hold expiring more than its cap ahead of its own clock, and this + // TTL is measured on ours, so the clamp leaves a step for the difference. + assertEquals(clamped, serverReservationTtl(server)); + assertTrue(clamped.compareTo(CreditLeaseDefaults.MAX_RESERVATION_TTL) < 0); + // Client mode never sends it, so clamping there would shorten holds for nothing. + assertEquals(twoHours, serverReservationTtl(client)); + } + } + + private Schematic leasedClient(CreditLeaseMode mode, Duration ttl) { + return Schematic.builder() + .apiKey("test_api_key") + .logger(logger) + .creditLeases(CreditLeaseConfig.builder() + .mode(mode) + .defaultReservationTtl(ttl) + .build()) + .build(); + } + + private static void setDataStreamClient(Schematic target, DataStreamClient value) throws Exception { + Field field = Schematic.class.getDeclaredField("dataStreamClient"); + field.setAccessible(true); + field.set(target, value); + } + + private static Duration serverReservationTtl(Schematic target) throws Exception { + Field field = Schematic.class.getDeclaredField("serverReservationTtl"); + field.setAccessible(true); + return (Duration) field.get(target); + } + + // --- Reservation settles --- + + private static Reservation reservation(CreditLeaseMode mode) { + return new Reservation( + "res_1", + mode == CreditLeaseMode.SERVER ? "res_1" : "lse_1", + mode, + "co_1", + "ct_1", + "inference_tokens", + 10, + 100, + 10, + Instant.parse("2026-01-01T00:01:00Z"), + Collections.singletonMap("id", "co_1"), + null); + } + + @Test + void buildReservationSettleEvent_clientModeRoutesThroughTheLeaseAndKeysOffTheHold() { + EventBodyTrack track = ReservationSettlement.buildTrackEvent(reservation(CreditLeaseMode.CLIENT), 4); + + CreateEventRequestBody event = Schematic.buildReservationSettleEvent(track, null, "res_1"); + + assertEquals(EventType.TRACK, event.getEventType()); + assertEquals("lease-reservation:res_1", event.getIdempotencyKey().get()); + EventBodyTrack body = (EventBodyTrack) event.getBody().get().get(); + assertEquals("inference_tokens", body.getEvent()); + assertEquals(4L, body.getQuantity().get()); + // The lease id routes the server-side consumption through the lease's sub-ledger instead + // of decrementing a grant the acquire already pre-debited. + assertEquals("lse_1", body.getLeaseId().get()); + assertFalse(body.getReservationId().isPresent()); + } + + @Test + void buildReservationSettleEvent_serverModeRoutesByHoldIdAndNeverNamesALease() { + EventBodyTrack track = ReservationSettlement.buildTrackEvent(reservation(CreditLeaseMode.SERVER), 4); + + CreateEventRequestBody event = Schematic.buildReservationSettleEvent(track, null, "res_1"); + + EventBodyTrack body = (EventBodyTrack) event.getBody().get().get(); + assertEquals("res_1", body.getReservationId().get()); + // The server prefers the lease id when both are set, and there is no lease here. + assertFalse(body.getLeaseId().isPresent()); + } + + @Test + void buildReservationSettleEvent_billsAWholeUnitForAFractionalSettle() { + EventBodyTrack track = ReservationSettlement.buildTrackEvent(reservation(CreditLeaseMode.CLIENT), 0.2); + + assertEquals(1L, track.getQuantity().get()); + } + @Test void buildIdentifyEvent_appliesIdempotencyKey() { EventBody body = EventBody.of(EventBodyIdentify.builder() diff --git a/src/test/java/com/schematic/api/credits/ApiLeaseWireClientTest.java b/src/test/java/com/schematic/api/credits/ApiLeaseWireClientTest.java new file mode 100644 index 0000000..14df74d --- /dev/null +++ b/src/test/java/com/schematic/api/credits/ApiLeaseWireClientTest.java @@ -0,0 +1,83 @@ +package com.schematic.api.credits; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.schematic.api.resources.credits.CreditsClient; +import com.schematic.api.resources.credits.requests.AcquireCreditLeaseRequestBody; +import com.schematic.api.resources.credits.requests.ExtendCreditLeaseRequestBody; +import com.schematic.api.resources.credits.types.AcquireCreditLeaseResponse; +import com.schematic.api.resources.credits.types.ExtendCreditLeaseResponse; +import com.schematic.api.types.CreditLeaseResponseData; +import java.time.Instant; +import java.time.OffsetDateTime; +import java.time.ZoneOffset; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +class ApiLeaseWireClientTest { + + private static final Instant EXPIRES_AT = Instant.parse("2026-01-01T00:05:00Z"); + + private CreditsClient credits; + private ApiLeaseWireClient wire; + + @BeforeEach + void setUp() { + credits = mock(CreditsClient.class); + wire = new ApiLeaseWireClient(credits); + } + + private static CreditLeaseResponseData lease(double grantedAmount) { + return CreditLeaseResponseData.builder() + .companyId("co_1") + .createdAt(OffsetDateTime.ofInstant(EXPIRES_AT.minusSeconds(300), ZoneOffset.UTC)) + .creditTypeId("ct_1") + .expiresAt(OffsetDateTime.ofInstant(EXPIRES_AT, ZoneOffset.UTC)) + .grantedAmount(grantedAmount) + .id("lse_1") + .trackedAmount(0) + .updatedAt(OffsetDateTime.ofInstant(EXPIRES_AT.minusSeconds(300), ZoneOffset.UTC)) + .build(); + } + + @Test + void acquireMapsTheGrantOntoALeaseState() { + when(credits.acquireCreditLease(any(AcquireCreditLeaseRequestBody.class))) + .thenReturn( + AcquireCreditLeaseResponse.builder().data(lease(1000)).build()); + + LeaseGrant grant = wire.acquire("co_1", "ct_1", 1000, EXPIRES_AT); + + assertEquals("lse_1", grant.getLeaseId()); + assertEquals("co_1", grant.getCompanyId()); + assertEquals(1000.0, grant.getGrantedAmount()); + assertEquals(EXPIRES_AT, grant.getExpiresAt()); + } + + @Test + void everyExtendCarriesItsOwnIdempotencyKey() { + when(credits.extendCreditLease(eq("lse_1"), any(ExtendCreditLeaseRequestBody.class))) + .thenReturn( + ExtendCreditLeaseResponse.builder().data(lease(2000)).build()); + + wire.extend("lse_1", 1000, EXPIRES_AT); + wire.extend("lse_1", 1000, EXPIRES_AT); + + ArgumentCaptor body = ArgumentCaptor.forClass(ExtendCreditLeaseRequestBody.class); + verify(credits, org.mockito.Mockito.times(2)).extendCreditLease(eq("lse_1"), body.capture()); + String first = body.getAllValues().get(0).getIdempotencyKey().orElse(null); + String second = body.getAllValues().get(1).getIdempotencyKey().orElse(null); + assertTrue(first != null && !first.isEmpty()); + // An extend is an increment, so one key per extend: retries of the same call collapse, + // while a later extend still grows the lease. + assertNotEquals(first, second); + } +} diff --git a/src/test/java/com/schematic/api/credits/CreditCheckReservePinningTest.java b/src/test/java/com/schematic/api/credits/CreditCheckReservePinningTest.java new file mode 100644 index 0000000..d967f4e --- /dev/null +++ b/src/test/java/com/schematic/api/credits/CreditCheckReservePinningTest.java @@ -0,0 +1,196 @@ +package com.schematic.api.credits; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.schematic.api.types.RulesengineCheckFlagResult; +import com.schematic.api.types.RulesengineCompany; +import com.schematic.api.types.RulesengineEntitlementValueType; +import com.schematic.api.types.RulesengineFeatureEntitlement; +import com.schematic.api.types.RulesengineFlag; +import com.schematic.api.types.RulesengineUser; +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.time.ZoneOffset; +import java.util.Collections; +import java.util.Map; +import java.util.concurrent.Callable; +import org.junit.jupiter.api.Test; + +/** A hold is pinned to the lease the debit landed on, never to the one the acquire handed back. */ +class CreditCheckReservePinningTest { + + private static final Instant NOW = Instant.parse("2026-01-01T00:00:00Z"); + private static final Clock CLOCK = Clock.fixed(NOW, ZoneOffset.UTC); + + /** Answers both evaluations allow, and meters the credit the entitlement names. */ + private static final class AllowingDataStream implements CreditCheckDataStream { + @Override + public RulesengineFlag getFlag(String flagKey) { + return RulesengineFlag.builder() + .accountId("acct") + .defaultValue(false) + .environmentId("env") + .id("flag_1") + .key(flagKey) + .build(); + } + + @Override + public RulesengineCompany getCompany(Map keys) { + return RulesengineCompany.builder() + .accountId("acct") + .environmentId("env") + .id("co_1") + .creditBalances(Collections.singletonMap("ct_1", 5000.0)) + .build(); + } + + @Override + public RulesengineUser getUser(Map keys) { + return null; + } + + @Override + public RulesengineCheckFlagResult evaluateFlag( + RulesengineFlag flag, RulesengineCompany company, RulesengineUser user, PreflightOptions preflight) { + return RulesengineCheckFlagResult.builder() + .flagKey("inference") + .reason("ok") + .value(true) + .flagId("flag_1") + .entitlement(RulesengineFeatureEntitlement.builder() + .featureId("feat_1") + .featureKey("inference") + .valueType(RulesengineEntitlementValueType.CREDIT) + .creditId("ct_1") + .consumptionRate(10.0) + .eventSubtype("inference_tokens") + .build()) + .build(); + } + } + + /** Rewrites what the debit reports, so the flow can be held to what it pins. */ + private static final class RewritingLeaseStore implements LeaseStore { + private final LeaseStore delegate; + private final String debitedLeaseId; + + RewritingLeaseStore(LeaseStore delegate, String debitedLeaseId) { + this.delegate = delegate; + this.debitedLeaseId = debitedLeaseId; + } + + @Override + public LeaseState get(String companyId, String creditTypeId) { + return delegate.get(companyId, creditTypeId); + } + + @Override + public boolean replace(LeaseGrant grant) { + return delegate.replace(grant); + } + + @Override + public ReserveResult tryReserve(String companyId, String creditTypeId, double credits) { + ReserveResult result = delegate.tryReserve(companyId, creditTypeId, credits); + return result == null ? null : new ReserveResult(result.getBalance(), debitedLeaseId); + } + + @Override + public void refund(String companyId, String creditTypeId, double credits, String pinLeaseId) { + delegate.refund(companyId, creditTypeId, credits, pinLeaseId); + } + + @Override + public void extend( + String companyId, String creditTypeId, double grantedTotal, Instant newExpiresAt, String pinLeaseId) { + delegate.extend(companyId, creditTypeId, grantedTotal, newExpiresAt, pinLeaseId); + } + + @Override + public void drop(String companyId, String creditTypeId) { + delegate.drop(companyId, creditTypeId); + } + } + + private static final class Fixture { + final InMemoryLeaseStore backing = new InMemoryLeaseStore(CLOCK); + final LeaseStore leases; + final InMemoryReservationStore holds; + final CreditLeaseManager manager; + final CreditCheck flow; + boolean fellBack; + + Fixture(String debitedLeaseId) { + leases = new RewritingLeaseStore(backing, debitedLeaseId); + holds = new InMemoryReservationStore(leases, CLOCK); + LeaseWireClient wire = new LeaseWireClient() { + @Override + public LeaseGrant acquire( + String companyId, String creditTypeId, double requestedAmount, Instant expiresAt) { + throw new UnsupportedOperationException("the slot already holds a lease"); + } + + @Override + public LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt) { + throw new UnsupportedOperationException("no extend expected"); + } + + @Override + public void release(String leaseId) {} + }; + manager = new CreditLeaseManager( + wire, leases, holds, CreditLeaseConfig.builder().build(), null, CLOCK); + flow = new CreditCheck(new AllowingDataStream(), leases, holds, manager, null, CLOCK, null, null); + backing.replace(new LeaseGrant("lse_acquired", "co_1", "ct_1", 1000, NOW.plusSeconds(300))); + } + + CheckResult run() { + Callable fallback = () -> { + fellBack = true; + return new CheckResult(true, true, "fallback", "inference", null, null, null, null); + }; + CheckResult result = flow.check( + new CheckRequest( + "inference", Collections.singletonMap("id", "co_1"), null, 10, "inference_tokens", false), + fallback); + manager.drain(Duration.ofSeconds(5)); + manager.close(); + return result; + } + } + + @Test + void theHoldNamesTheLeaseTheDebitLandedOn() { + // The slot took on a different lease between the acquire and the debit, which the atomic + // reserve reports. Pinning the acquired lease instead would send this hold's refunds, and + // its billing, to a lease that never held the credits. + Fixture fixture = new Fixture("lse_debited"); + + CheckResult result = fixture.run(); + + assertFalse(fixture.fellBack); + assertTrue(result.isAllowed()); + assertNotNull(result.getReservation()); + assertEquals("lse_debited", result.getReservation().getLeaseId()); + } + + @Test + void aDebitThatNamesNoLeaseIsHandedBackRatherThanPinnedToTheAcquiredOne() { + Fixture fixture = new Fixture(null); + + CheckResult result = fixture.run(); + + assertFalse(result.isAllowed()); + assertNull(result.getReservation()); + assertEquals("lease_store_error", result.getReason()); + assertEquals(0, fixture.holds.count()); + // The debit went back to the lease rather than sitting there until it expires. + assertEquals(1000.0, fixture.backing.get("co_1", "ct_1").getLocalRemainingCredits()); + } +} diff --git a/src/test/java/com/schematic/api/credits/CreditCheckReserveSizingTest.java b/src/test/java/com/schematic/api/credits/CreditCheckReserveSizingTest.java new file mode 100644 index 0000000..dd9b059 --- /dev/null +++ b/src/test/java/com/schematic/api/credits/CreditCheckReserveSizingTest.java @@ -0,0 +1,339 @@ +package com.schematic.api.credits; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.schematic.api.types.RulesengineCheckFlagResult; +import com.schematic.api.types.RulesengineCompany; +import com.schematic.api.types.RulesengineEntitlementValueType; +import com.schematic.api.types.RulesengineFeatureEntitlement; +import com.schematic.api.types.RulesengineFlag; +import com.schematic.api.types.RulesengineUser; +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.time.ZoneId; +import java.time.ZoneOffset; +import java.util.Collections; +import java.util.Map; +import java.util.concurrent.Callable; +import org.junit.jupiter.api.Test; + +/** How large a hold a check takes, and what it leaves behind when it cannot take one. */ +class CreditCheckReserveSizingTest { + + private static final Instant NOW = Instant.parse("2026-01-01T00:00:00Z"); + private static final Clock CLOCK = Clock.fixed(NOW, ZoneOffset.UTC); + + /** Allows, and meters the credit at ten credits a unit. */ + private static class AllowingDataStream implements CreditCheckDataStream { + @Override + public RulesengineFlag getFlag(String flagKey) { + return RulesengineFlag.builder() + .accountId("acct") + .defaultValue(false) + .environmentId("env") + .id("flag_1") + .key(flagKey) + .build(); + } + + @Override + public RulesengineCompany getCompany(Map keys) { + return RulesengineCompany.builder() + .accountId("acct") + .environmentId("env") + .id("co_1") + .creditBalances(Collections.singletonMap("ct_1", 5000.0)) + .build(); + } + + @Override + public RulesengineUser getUser(Map keys) { + return null; + } + + @Override + public RulesengineCheckFlagResult evaluateFlag( + RulesengineFlag flag, RulesengineCompany company, RulesengineUser user, PreflightOptions preflight) { + return RulesengineCheckFlagResult.builder() + .flagKey("inference") + .reason("ok") + .value(true) + .flagId("flag_1") + .entitlement(RulesengineFeatureEntitlement.builder() + .featureId("feat_1") + .featureKey("inference") + .valueType(RulesengineEntitlementValueType.CREDIT) + .creditId("ct_1") + .consumptionRate(10.0) + .eventSubtype("inference_tokens") + .build()) + .build(); + } + } + + /** A clock nothing can read, standing in for anything the reservation prep can throw on. */ + private static final class BrokenClock extends Clock { + @Override + public ZoneId getZone() { + return ZoneOffset.UTC; + } + + @Override + public Clock withZone(ZoneId zone) { + return this; + } + + @Override + public Instant instant() { + throw new IllegalStateException("no clock"); + } + } + + /** Allows the probe, then throws when the gate asks. */ + private static final class ThrowingGateDataStream extends AllowingDataStream { + private int evaluations; + + @Override + public RulesengineCheckFlagResult evaluateFlag( + RulesengineFlag flag, RulesengineCompany company, RulesengineUser user, PreflightOptions preflight) { + if (++evaluations > 1) { + throw new IllegalStateException("the engine blew up"); + } + return super.evaluateFlag(flag, company, user, preflight); + } + } + + /** Delegates everything but the gate, which throws the way an unreachable store would. */ + private static final class UnreachableOnReserve implements LeaseStore { + private final LeaseStore delegate; + + UnreachableOnReserve(LeaseStore delegate) { + this.delegate = delegate; + } + + @Override + public LeaseState get(String companyId, String creditTypeId) { + return delegate.get(companyId, creditTypeId); + } + + @Override + public boolean replace(LeaseGrant grant) { + return delegate.replace(grant); + } + + @Override + public ReserveResult tryReserve(String companyId, String creditTypeId, double credits) { + throw new IllegalStateException("the lease store is unreachable"); + } + + @Override + public void refund(String companyId, String creditTypeId, double credits, String pinLeaseId) { + delegate.refund(companyId, creditTypeId, credits, pinLeaseId); + } + + @Override + public void extend( + String companyId, String creditTypeId, double grantedTotal, Instant newExpiresAt, String pinLeaseId) { + delegate.extend(companyId, creditTypeId, grantedTotal, newExpiresAt, pinLeaseId); + } + + @Override + public void drop(String companyId, String creditTypeId) { + delegate.drop(companyId, creditTypeId); + } + } + + /** Takes the debit but refuses to record the hold, the window undoDebit exists for. */ + private static final class UnreachableOnAdd implements ReservationStore { + private final ReservationStore delegate; + + UnreachableOnAdd(ReservationStore delegate) { + this.delegate = delegate; + } + + @Override + public void add(Reservation reservation) { + throw new IllegalStateException("the reservation store is unreachable"); + } + + @Override + public Reservation get(String id) { + return delegate.get(id); + } + + @Override + public Double consume(String id, double creditsConsumed) { + return delegate.consume(id, creditsConsumed); + } + + @Override + public double reservedCredits(String companyId, String creditTypeId) { + return delegate.reservedCredits(companyId, creditTypeId); + } + + @Override + public int sweepExpired() { + return delegate.sweepExpired(); + } + + @Override + public int count() { + return delegate.count(); + } + } + + private static final class Fixture { + final InMemoryLeaseStore leases = new InMemoryLeaseStore(CLOCK); + final InMemoryReservationStore holds = new InMemoryReservationStore(leases, CLOCK); + final CreditLeaseManager manager; + final CreditCheck flow; + boolean fellBack; + + Fixture(Clock flowClock) { + this(flowClock, new AllowingDataStream()); + } + + Fixture(Clock flowClock, CreditCheckDataStream source) { + this(flowClock, source, false, false); + } + + /** + * The two stores the flow writes through can each be made unreachable, which is the only + * way to reach the paths that give a debit back. + */ + Fixture(Clock flowClock, CreditCheckDataStream source, boolean breakReserve, boolean breakAdd) { + LeaseWireClient wire = new LeaseWireClient() { + @Override + public LeaseGrant acquire(String companyId, String creditTypeId, double amount, Instant expiresAt) { + throw new UnsupportedOperationException("the slot already holds a lease"); + } + + @Override + public LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt) { + throw new UnsupportedOperationException("no extend expected"); + } + + @Override + public void release(String leaseId) {} + }; + manager = new CreditLeaseManager( + wire, leases, holds, CreditLeaseConfig.builder().build(), null, CLOCK); + flow = new CreditCheck( + source, + breakReserve ? new UnreachableOnReserve(leases) : leases, + breakAdd ? new UnreachableOnAdd(holds) : holds, + manager, + null, + flowClock, + null, + null); + leases.replace(new LeaseGrant("lse_1", "co_1", "ct_1", 1000, NOW.plusSeconds(300))); + } + + CheckResult run(double usage) { + Callable fallback = () -> { + fellBack = true; + return new CheckResult(true, true, "fallback", "inference", null, null, null, null); + }; + CheckResult result = flow.check( + new CheckRequest( + "inference", + Collections.singletonMap("id", "co_1"), + null, + usage, + "inference_tokens", + false), + fallback); + manager.drain(Duration.ofSeconds(5)); + manager.close(); + return result; + } + } + + @Test + void aFractionalUsageHoldsAWholeEventUnitAtTheRate() { + Fixture fixture = new Fixture(CLOCK); + + CheckResult result = fixture.run(0.5); + + assertFalse(fixture.fellBack); + assertTrue(result.isAllowed()); + assertNotNull(result.getReservation()); + // Half an event is not something the server bills, so the hold is sized at the whole unit + // the settle will charge for, while the reservation still records what the caller + // declared. + assertEquals(10.0, result.getReservation().getCreditsReserved()); + assertEquals(0.5, result.getReservation().getQuantityReserved()); + assertEquals(990.0, fixture.leases.get("co_1", "ct_1").getLocalRemainingCredits()); + } + + @Test + void anEngineThrowAtTheGateCancelsTheHoldItTookFirst() { + Fixture fixture = new Fixture(CLOCK, new ThrowingGateDataStream()); + + CheckResult result = fixture.run(10); + + // A throw is not a verdict, so the check resolves through its fail-closed contract, and + // the credits debited before the gate go back rather than sitting on the lease. + assertFalse(fixture.fellBack); + assertFalse(result.isAllowed()); + assertNull(result.getReservation()); + assertTrue(result.getReason().startsWith("wasm_error")); + assertEquals(0, fixture.holds.count()); + assertEquals(1000.0, fixture.leases.get("co_1", "ct_1").getLocalRemainingCredits()); + } + + @Test + void aFailureWhilePreparingTheHoldLeavesNoDebitBehind() { + Fixture fixture = new Fixture(new BrokenClock()); + + CheckResult result = fixture.run(10); + + assertFalse(result.isAllowed()); + assertNull(result.getReservation()); + assertEquals("lease_store_error", result.getReason()); + assertEquals(0, fixture.holds.count()); + // Everything the hold needs is resolved before the debit, so a throw here cannot strand + // credits on a lease with no reservation naming them. + assertEquals(1000.0, fixture.leases.get("co_1", "ct_1").getLocalRemainingCredits()); + } + + @Test + void anUnreachableLeaseStoreAtTheGateFailsClosedWithoutAHold() { + Fixture fixture = new Fixture(CLOCK, new AllowingDataStream(), true, false); + + CheckResult result = fixture.run(10); + + // The gate is the one step that cannot be guessed at: without the atomic check and debit + // there is no telling whether the credits are there, so the check resolves through its + // fail-closed contract rather than allowing on an unknown balance. + assertFalse(fixture.fellBack); + assertFalse(result.isAllowed()); + assertNull(result.getReservation()); + assertEquals("lease_store_error", result.getReason()); + assertEquals(0, fixture.holds.count()); + assertEquals(1000.0, fixture.leases.get("co_1", "ct_1").getLocalRemainingCredits()); + } + + @Test + void aHoldThatCannotBeRecordedGivesItsCreditsBack() { + Fixture fixture = new Fixture(CLOCK, new AllowingDataStream(), false, true); + + CheckResult result = fixture.run(10); + + // The debit landed and the hold that would have settled it did not, so the credits go + // back to the lease they came out of. Left alone they would sit debited with nothing to + // settle or sweep them, and the slot would leak a tranche at a time. + assertFalse(fixture.fellBack); + assertFalse(result.isAllowed()); + assertNull(result.getReservation()); + assertEquals("lease_store_error", result.getReason()); + assertEquals(0, fixture.holds.count()); + assertEquals(1000.0, fixture.leases.get("co_1", "ct_1").getLocalRemainingCredits()); + } +} diff --git a/src/test/java/com/schematic/api/credits/CreditLeaseConfigValidationTest.java b/src/test/java/com/schematic/api/credits/CreditLeaseConfigValidationTest.java new file mode 100644 index 0000000..4074c6d --- /dev/null +++ b/src/test/java/com/schematic/api/credits/CreditLeaseConfigValidationTest.java @@ -0,0 +1,270 @@ +package com.schematic.api.credits; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.time.Duration; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +class CreditLeaseConfigValidationTest { + + @Test + void aNaNLeaseSizeIsRejected() { + IllegalArgumentException thrown = assertThrows( + IllegalArgumentException.class, + () -> CreditLeaseConfig.builder().defaultLeaseSize(Double.NaN).build()); + assertTrue(thrown.getMessage().contains("defaultLeaseSize")); + } + + @Test + void anInfiniteLeaseSizeIsRejected() { + assertThrows(IllegalArgumentException.class, () -> CreditLeaseConfig.builder() + .defaultLeaseSize(Double.POSITIVE_INFINITY) + .build()); + } + + @Test + void aZeroLeaseSizeIsRejected() { + assertThrows( + IllegalArgumentException.class, + () -> CreditLeaseConfig.builder().defaultLeaseSize(0).build()); + } + + @Test + void aNegativeLeaseSizeIsRejected() { + assertThrows( + IllegalArgumentException.class, + () -> CreditLeaseConfig.builder().defaultLeaseSize(-1).build()); + } + + @Test + void aNaNLowWaterMarkIsRejected() { + IllegalArgumentException thrown = assertThrows( + IllegalArgumentException.class, + () -> CreditLeaseConfig.builder().lowWaterMark(Double.NaN).build()); + assertTrue(thrown.getMessage().contains("lowWaterMark")); + } + + @Test + void aLowWaterMarkOfZeroIsRejected() { + assertThrows( + IllegalArgumentException.class, + () -> CreditLeaseConfig.builder().lowWaterMark(0).build()); + } + + @Test + void aLowWaterMarkOfOneIsRejected() { + assertThrows( + IllegalArgumentException.class, + () -> CreditLeaseConfig.builder().lowWaterMark(1).build()); + } + + @Test + void aNegativeLowWaterMarkIsRejected() { + assertThrows( + IllegalArgumentException.class, + () -> CreditLeaseConfig.builder().lowWaterMark(-0.5).build()); + } + + @Test + void aWaterMarkInsideTheOpenUnitIntervalIsAccepted() { + CreditLeaseConfig config = + CreditLeaseConfig.builder().lowWaterMark(0.99).build(); + assertEquals(0.99, config.getLowWaterMark(), 0); + } + + @Test + void aZeroLeaseDurationIsRejected() { + IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, () -> CreditLeaseConfig.builder() + .defaultLeaseDuration(Duration.ZERO) + .build()); + assertTrue(thrown.getMessage().contains("defaultLeaseDuration")); + } + + @Test + void aNegativeLeaseDurationIsRejected() { + assertThrows(IllegalArgumentException.class, () -> CreditLeaseConfig.builder() + .defaultLeaseDuration(Duration.ofSeconds(-1)) + .build()); + } + + @Test + void aZeroReservationTtlIsRejected() { + IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, () -> CreditLeaseConfig.builder() + .defaultReservationTtl(Duration.ZERO) + .build()); + assertTrue(thrown.getMessage().contains("defaultReservationTtl")); + } + + @Test + void aNegativeReservationTtlIsRejected() { + assertThrows(IllegalArgumentException.class, () -> CreditLeaseConfig.builder() + .defaultReservationTtl(Duration.ofMinutes(-5)) + .build()); + } + + @Test + void aZeroSweepIntervalIsRejected() { + IllegalArgumentException thrown = assertThrows( + IllegalArgumentException.class, + () -> CreditLeaseConfig.builder().sweepInterval(Duration.ZERO).build()); + assertTrue(thrown.getMessage().contains("sweepInterval")); + } + + @Test + void aNegativeSweepIntervalIsRejected() { + assertThrows(IllegalArgumentException.class, () -> CreditLeaseConfig.builder() + .sweepInterval(Duration.ofMillis(-1)) + .build()); + } + + @Test + void aFullyValidConfigBuilds() { + CreditLeaseConfig config = CreditLeaseConfig.builder() + .defaultLeaseSize(500) + .lowWaterMark(0.4) + .defaultLeaseDuration(Duration.ofMinutes(2)) + .defaultReservationTtl(Duration.ofSeconds(30)) + .sweepInterval(Duration.ofSeconds(2)) + .override( + "tokens", + CreditLeaseOverride.builder() + .defaultLeaseSize(50) + .lowWaterMark(0.1) + .defaultLeaseDuration(Duration.ofMinutes(1)) + .defaultReservationTtl(Duration.ofSeconds(15)) + .build()) + .build(); + + assertEquals(500, config.getDefaultLeaseSize(), 0); + assertEquals(Duration.ofSeconds(2), config.getSweepInterval()); + assertEquals(1, config.getOverrides().size()); + } + + @Test + void aConfigThatSetsNothingBuilds() { + CreditLeaseConfig config = CreditLeaseConfig.builder().build(); + + assertNotNull(config); + assertEquals(CreditLeaseMode.AUTO, config.getMode()); + } + + @Test + void anOverrideWithANaNLeaseSizeNamesTheCreditType() { + IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, () -> CreditLeaseConfig.builder() + .override( + "tokens", + CreditLeaseOverride.builder() + .defaultLeaseSize(Double.NaN) + .build()) + .build()); + assertTrue(thrown.getMessage().contains("defaultLeaseSize")); + assertTrue(thrown.getMessage().contains("for credit type tokens")); + } + + @Test + void anOverrideWithANegativeLeaseSizeIsRejected() { + assertThrows(IllegalArgumentException.class, () -> CreditLeaseConfig.builder() + .override( + "tokens", + CreditLeaseOverride.builder().defaultLeaseSize(-5).build()) + .build()); + } + + @Test + void anOverrideWithAWaterMarkOfOneNamesTheCreditType() { + IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, () -> CreditLeaseConfig.builder() + .override( + "tokens", CreditLeaseOverride.builder().lowWaterMark(1).build()) + .build()); + assertTrue(thrown.getMessage().contains("lowWaterMark")); + assertTrue(thrown.getMessage().contains("for credit type tokens")); + } + + @Test + void anOverrideWithAZeroLeaseDurationNamesTheCreditType() { + IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, () -> CreditLeaseConfig.builder() + .override( + "tokens", + CreditLeaseOverride.builder() + .defaultLeaseDuration(Duration.ZERO) + .build()) + .build()); + assertTrue(thrown.getMessage().contains("defaultLeaseDuration")); + assertTrue(thrown.getMessage().contains("for credit type tokens")); + } + + @Test + void anOverrideWithANegativeReservationTtlNamesTheCreditType() { + IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, () -> CreditLeaseConfig.builder() + .override( + "tokens", + CreditLeaseOverride.builder() + .defaultReservationTtl(Duration.ofSeconds(-1)) + .build()) + .build()); + assertTrue(thrown.getMessage().contains("defaultReservationTtl")); + assertTrue(thrown.getMessage().contains("for credit type tokens")); + } + + @Test + void anOverrideThatSetsNothingBuilds() { + CreditLeaseConfig config = CreditLeaseConfig.builder() + .override("tokens", CreditLeaseOverride.builder().build()) + .build(); + + assertEquals(1, config.getOverrides().size()); + } + + @Nested + class Resolution { + + private CreditLeaseConfig config() { + return CreditLeaseConfig.builder() + .defaultLeaseSize(1000) + .lowWaterMark(0.5) + .defaultLeaseDuration(Duration.ofMinutes(10)) + .defaultReservationTtl(Duration.ofSeconds(90)) + .override( + "tokens", + CreditLeaseOverride.builder() + .defaultLeaseSize(25) + .defaultLeaseDuration(Duration.ofMinutes(1)) + .build()) + .build(); + } + + @Test + void anOverriddenCreditTypeTakesTheKnobsItSetsAndTheConfigDefaultsItDoesNot() { + ResolvedLeaseConfig resolved = config().resolve("tokens"); + + assertEquals(25, resolved.getLeaseSize(), 0); + assertEquals(Duration.ofMinutes(1), resolved.getLeaseDuration()); + assertEquals(0.5, resolved.getLowWaterMark(), 0); + assertEquals(Duration.ofSeconds(90), resolved.getReservationTtl()); + } + + @Test + void aCreditTypeWithoutAnOverrideTakesEveryConfigDefault() { + ResolvedLeaseConfig resolved = config().resolve("seats"); + + assertEquals(1000, resolved.getLeaseSize(), 0); + assertEquals(Duration.ofMinutes(10), resolved.getLeaseDuration()); + assertEquals(0.5, resolved.getLowWaterMark(), 0); + assertEquals(Duration.ofSeconds(90), resolved.getReservationTtl()); + } + + @Test + void aConfigThatSetsNothingResolvesToTheLibraryDefaults() { + ResolvedLeaseConfig resolved = CreditLeaseConfig.builder().build().resolve("tokens"); + + assertEquals(CreditLeaseDefaults.LEASE_SIZE, resolved.getLeaseSize(), 0); + assertEquals(CreditLeaseDefaults.LEASE_DURATION, resolved.getLeaseDuration()); + assertEquals(CreditLeaseDefaults.LOW_WATER_MARK, resolved.getLowWaterMark(), 0); + assertEquals(CreditLeaseDefaults.RESERVATION_TTL, resolved.getReservationTtl()); + } + } +} diff --git a/src/test/java/com/schematic/api/credits/CreditLeaseManagerExtendTest.java b/src/test/java/com/schematic/api/credits/CreditLeaseManagerExtendTest.java new file mode 100644 index 0000000..31cfd35 --- /dev/null +++ b/src/test/java/com/schematic/api/credits/CreditLeaseManagerExtendTest.java @@ -0,0 +1,633 @@ +package com.schematic.api.credits; + +import static com.schematic.api.credits.TestThreads.parked; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.lang.reflect.Field; +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import org.junit.jupiter.api.Test; + +/** What the manager does with background top-ups, and what it refuses to spend a thread on. */ +class CreditLeaseManagerExtendTest { + + private static final Instant NOW = Instant.parse("2026-01-01T00:00:00Z"); + private static final Clock CLOCK = Clock.fixed(NOW, ZoneOffset.UTC); + private static final Instant LIVE = NOW.plusSeconds(300); + + /** Counts extends, and can hold them open until released. */ + private static class CountingWire implements LeaseWireClient { + final CountDownLatch started = new CountDownLatch(1); + final CountDownLatch release = new CountDownLatch(1); + final AtomicInteger extends_ = new AtomicInteger(); + final List released = Collections.synchronizedList(new ArrayList<>()); + private final boolean block; + + CountingWire(boolean block) { + this.block = block; + } + + @Override + public LeaseGrant acquire(String companyId, String creditTypeId, double amount, Instant expiresAt) { + return new LeaseGrant("lse_acquired", companyId, creditTypeId, amount, expiresAt); + } + + @Override + public LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt) { + extends_.incrementAndGet(); + started.countDown(); + if (block) { + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + return new LeaseGrant(leaseId, "co_1", "ct_1", 3000, expiresAt); + } + + @Override + public void release(String leaseId) { + released.add(leaseId); + } + } + + /** + * Serves the caller's first read as it stands, and every later one as a lease that has since + * been topped up. That is the shape of a stale read: a concurrent extend lands between the + * row the caller decided on and the moment it owns the slot's flight. + */ + private static class ToppedUpAfterFirstRead implements LeaseStore { + private final InMemoryLeaseStore delegate; + private final AtomicInteger reads = new AtomicInteger(); + + ToppedUpAfterFirstRead(InMemoryLeaseStore delegate) { + this.delegate = delegate; + } + + @Override + public LeaseState get(String companyId, String creditTypeId) { + LeaseState entry = delegate.get(companyId, creditTypeId); + if (entry == null || reads.incrementAndGet() == 1) { + return entry; + } + return new LeaseState(entry.getLeaseId(), companyId, creditTypeId, 3000, 2900, entry.getExpiresAt()); + } + + @Override + public boolean replace(LeaseGrant grant) { + return delegate.replace(grant); + } + + @Override + public ReserveResult tryReserve(String companyId, String creditTypeId, double credits) { + return delegate.tryReserve(companyId, creditTypeId, credits); + } + + @Override + public void refund(String companyId, String creditTypeId, double credits, String pinLeaseId) { + delegate.refund(companyId, creditTypeId, credits, pinLeaseId); + } + + @Override + public void extend( + String companyId, String creditTypeId, double grantedTotal, Instant newExpiresAt, String pinLeaseId) { + delegate.extend(companyId, creditTypeId, grantedTotal, newExpiresAt, pinLeaseId); + } + + @Override + public void drop(String companyId, String creditTypeId) { + delegate.drop(companyId, creditTypeId); + } + } + + /** Counts reads, so a task that would have read again is visible. */ + private static class CountingReads implements LeaseStore { + private final LeaseStore delegate; + final AtomicInteger reads = new AtomicInteger(); + + CountingReads(LeaseStore delegate) { + this.delegate = delegate; + } + + @Override + public LeaseState get(String companyId, String creditTypeId) { + reads.incrementAndGet(); + return delegate.get(companyId, creditTypeId); + } + + @Override + public boolean replace(LeaseGrant grant) { + return delegate.replace(grant); + } + + @Override + public ReserveResult tryReserve(String companyId, String creditTypeId, double credits) { + return delegate.tryReserve(companyId, creditTypeId, credits); + } + + @Override + public void refund(String companyId, String creditTypeId, double credits, String pinLeaseId) { + delegate.refund(companyId, creditTypeId, credits, pinLeaseId); + } + + @Override + public void extend( + String companyId, String creditTypeId, double grantedTotal, Instant newExpiresAt, String pinLeaseId) { + delegate.extend(companyId, creditTypeId, grantedTotal, newExpiresAt, pinLeaseId); + } + + @Override + public void drop(String companyId, String creditTypeId) { + delegate.drop(companyId, creditTypeId); + } + } + + /** Counts sweeps and holds nothing. */ + private static final class CountingReservations implements ReservationStore { + final AtomicInteger sweeps = new AtomicInteger(); + final CountDownLatch swept = new CountDownLatch(1); + + @Override + public void add(Reservation reservation) {} + + @Override + public Reservation get(String id) { + return null; + } + + @Override + public Double consume(String id, double creditsConsumed) { + return null; + } + + @Override + public double reservedCredits(String companyId, String creditTypeId) { + return 0; + } + + @Override + public int sweepExpired() { + swept.countDown(); + return sweeps.incrementAndGet(); + } + + @Override + public int count() { + return 0; + } + } + + private static CreditLeaseManager manager(LeaseWireClient wire, LeaseStore leases, ReservationStore holds) { + return manager(wire, leases, holds, CreditLeaseConfig.builder().build()); + } + + private static CreditLeaseManager manager( + LeaseWireClient wire, LeaseStore leases, ReservationStore holds, CreditLeaseConfig config) { + return new CreditLeaseManager(wire, leases, holds, config, null, CLOCK); + } + + /** How many threads are parked inside a lease flight right now, whichever pool they came from. */ + private static int threadsParkedOnAFlight() { + int count = 0; + for (StackTraceElement[] stack : Thread.getAllStackTraces().values()) { + for (StackTraceElement frame : stack) { + if (frame.getClassName().startsWith(CreditLeaseManager.class.getName() + "$Flight") + && "await".equals(frame.getMethodName())) { + count++; + break; + } + } + } + return count; + } + + @Test + void manyChecksDuringOneSlowExtendSendOneExtendAndParkNoThreads() throws Exception { + InMemoryLeaseStore leases = new InMemoryLeaseStore(CLOCK); + InMemoryReservationStore holds = new InMemoryReservationStore(leases, CLOCK); + CountingWire wire = new CountingWire(true); + CreditLeaseManager manager = manager(wire, leases, holds); + leases.replace(new LeaseGrant("lse_1", "co_1", "ct_1", 1000, LIVE)); + // Draw the lease under its water mark, so every check that follows warrants a top-up. + leases.tryReserve("co_1", "ct_1", 900); + + manager.extendInBackground("co_1", "ct_1"); + assertTrue(wire.started.await(5, TimeUnit.SECONDS)); + for (int i = 0; i < 200; i++) { + manager.extendInBackground("co_1", "ct_1"); + } + Thread.sleep(200); + + // Each of those checks found a top-up already on the wire. Joining it would have parked a + // pool thread apiece for the length of one network call. + assertEquals(1, wire.extends_.get()); + assertEquals(0, threadsParkedOnAFlight()); + wire.release.countDown(); + manager.close(); + } + + @Test + void aHealthyLeaseSendsNoBackgroundWorkToThePool() { + InMemoryLeaseStore backing = new InMemoryLeaseStore(CLOCK); + InMemoryReservationStore holds = new InMemoryReservationStore(backing, CLOCK); + CountingReads leases = new CountingReads(backing); + CountingWire wire = new CountingWire(false); + CreditLeaseManager manager = manager(wire, leases, holds); + backing.replace(new LeaseGrant("lse_1", "co_1", "ct_1", 1000, LIVE)); + backing.tryReserve("co_1", "ct_1", 10); + + for (int i = 0; i < 200; i++) { + manager.extendInBackground("co_1", "ct_1"); + } + + // The water-mark test runs on the caller's thread, so a comfortable lease queues nothing: + // exactly one store read per call, and no task behind it to read again. + assertEquals(0, wire.extends_.get()); + assertEquals(200, leases.reads.get()); + manager.close(); + } + + @Test + void aStaleRowDoesNotSendASecondExtend() { + InMemoryLeaseStore backing = new InMemoryLeaseStore(CLOCK); + backing.replace(new LeaseGrant("lse_1", "co_1", "ct_1", 1000, LIVE)); + backing.tryReserve("co_1", "ct_1", 900); + ToppedUpAfterFirstRead leases = new ToppedUpAfterFirstRead(backing); + InMemoryReservationStore holds = new InMemoryReservationStore(backing, CLOCK); + CountingWire wire = new CountingWire(false); + CreditLeaseManager manager = manager(wire, leases, holds); + + assertNotNull(manager.maybeExtend("co_1", "ct_1", null)); + + // The re-read once the flight is ours is what catches it. Without one the stale row bills + // a second tranche onto a lease the previous extend already topped up. + assertEquals(0, wire.extends_.get()); + manager.close(); + } + + @Test + void aLeaseThatLandsAfterStopIsLeftAloneRatherThanReleased() throws Exception { + InMemoryLeaseStore leases = new InMemoryLeaseStore(CLOCK); + InMemoryReservationStore holds = new InMemoryReservationStore(leases, CLOCK); + CountDownLatch onTheWire = new CountDownLatch(1); + CountDownLatch swept = new CountDownLatch(1); + List released = Collections.synchronizedList(new ArrayList<>()); + LeaseWireClient wire = new LeaseWireClient() { + @Override + public LeaseGrant acquire(String companyId, String creditTypeId, double amount, Instant expiresAt) { + onTheWire.countDown(); + try { + swept.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return new LeaseGrant("lse_late", companyId, creditTypeId, amount, expiresAt); + } + + @Override + public LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt) { + return null; + } + + @Override + public void release(String leaseId) { + released.add(leaseId); + } + }; + CreditLeaseManager manager = manager(wire, leases, holds); + + Thread caller = new Thread(() -> manager.acquireIfNeeded("co_1", "ct_1")); + caller.start(); + assertTrue(onTheWire.await(5, TimeUnit.SECONDS)); + manager.stop(); + manager.releaseAllLocalLeases(); + swept.countDown(); + caller.join(5000); + + // A lease that lands this late is left to the drain or to server-side expiry. Releasing + // it here would refund, on a shared backend, a lease sibling pods are still reserving + // against. + assertTrue(released.isEmpty(), "released " + released); + manager.close(); + } + + @Test + void aJoinerWhoseFlightSentNothingIssuesItsOwnExtend() throws Exception { + InMemoryLeaseStore backing = new InMemoryLeaseStore(CLOCK); + backing.replace(new LeaseGrant("lse_1", "co_1", "ct_1", 1000, LIVE)); + backing.tryReserve("co_1", "ct_1", 900); + InMemoryReservationStore holds = new InMemoryReservationStore(backing, CLOCK); + CountDownLatch ownerReRead = new CountDownLatch(1); + CountDownLatch joinerRegistered = new CountDownLatch(1); + AtomicInteger readOrder = new AtomicInteger(); + // The owner's first read sees the drawn-down row and decides to extend; its re-read, once + // the flight is registered, sees a slot another extend already topped up, so it sends + // nothing. The re-read is held open long enough for a joiner to queue behind that flight. + LeaseStore stalling = new ToppedUpAfterFirstRead(backing) { + @Override + public LeaseState get(String companyId, String creditTypeId) { + if (readOrder.incrementAndGet() == 2) { + ownerReRead.countDown(); + try { + joinerRegistered.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + return super.get(companyId, creditTypeId); + } + }; + CountingWire wire = new CountingWire(false); + CreditLeaseManager manager = manager(wire, stalling, holds); + + Thread owner = new Thread(() -> manager.maybeExtend("co_1", "ct_1", null)); + owner.start(); + assertTrue(ownerReRead.await(5, TimeUnit.SECONDS)); + LeaseState[] joined = new LeaseState[1]; + Thread joiner = new Thread(() -> joined[0] = manager.maybeExtend("co_1", "ct_1", 4000.0)); + joiner.start(); + // The joiner has to reach the flight it is queuing behind before the owner finishes, and + // its thread state is the signal: Flight.await is the one place maybeExtend blocks. + assertTrue(parked(joiner), "the joiner never parked on the owner's flight"); + joinerRegistered.countDown(); + owner.join(5000); + joiner.join(5000); + + // The owner asked for a tranche but never sent it, so its ask stands in for nobody. A + // joiner that took it as covering its own shortfall would deny a check whose credits are + // still sitting on the server. + assertEquals(1, wire.extends_.get()); + assertNotNull(joined[0]); + manager.close(); + } + + @Test + void theCallersTimeoutReachesTheAcquireAndTheExtendButNotABackgroundTopUp() throws Exception { + InMemoryLeaseStore leases = new InMemoryLeaseStore(CLOCK); + InMemoryReservationStore holds = new InMemoryReservationStore(leases, CLOCK); + List acquireTimeouts = Collections.synchronizedList(new ArrayList<>()); + List extendTimeouts = Collections.synchronizedList(new ArrayList<>()); + LeaseWireClient wire = new LeaseWireClient() { + @Override + public LeaseGrant acquire(String companyId, String creditTypeId, double amount, Instant expiresAt) { + throw new UnsupportedOperationException("the timeout-carrying overload is the one under test"); + } + + @Override + public LeaseGrant acquire( + String companyId, String creditTypeId, double amount, Instant expiresAt, Duration timeout) { + acquireTimeouts.add(timeout); + return new LeaseGrant("lse_1", companyId, creditTypeId, amount, expiresAt); + } + + @Override + public LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt) { + throw new UnsupportedOperationException("the timeout-carrying overload is the one under test"); + } + + @Override + public LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt, Duration timeout) { + extendTimeouts.add(timeout); + return new LeaseGrant(leaseId, "co_1", "ct_1", 5000, expiresAt); + } + + @Override + public void release(String leaseId) {} + }; + CreditLeaseManager manager = manager(wire, leases, holds); + Duration perCheck = Duration.ofMillis(250); + + manager.acquireIfNeeded("co_1", "ct_1", perCheck); + leases.tryReserve("co_1", "ct_1", 9900); + manager.maybeExtend("co_1", "ct_1", null, perCheck); + manager.extendInBackground("co_1", "ct_1"); + manager.drain(Duration.ofSeconds(5)); + + // The caller is waiting on these two, so its deadline is the one that counts. + assertEquals(Collections.singletonList(perCheck), acquireTimeouts); + // And the background top-up, which nobody is waiting on, keeps the client's own. + assertEquals(Arrays.asList(perCheck, null), extendTimeouts); + manager.close(); + } + + @Test + void startSweepSchedulesOneSweeperHoweverOftenItIsCalled() throws Exception { + InMemoryLeaseStore leases = new InMemoryLeaseStore(CLOCK); + CountingReservations holds = new CountingReservations(); + // An interval no test run reaches, so the queue holds the scheduled sweeps themselves + // rather than whatever happens to be pending at the moment of the assertion. + CreditLeaseManager manager = manager( + new CountingWire(false), + leases, + holds, + CreditLeaseConfig.builder().sweepInterval(Duration.ofHours(1)).build()); + + manager.startSweep(); + manager.startSweep(); + manager.startSweep(); + + Field field = CreditLeaseManager.class.getDeclaredField("sweeper"); + field.setAccessible(true); + ScheduledExecutorService sweeper = (ScheduledExecutorService) field.get(manager); + // Nothing has run at an hourly interval, so every schedule is still queued and + // shutdownNow hands them all back. Three of them would sweep three times as often as the + // one interval asks for, and counting them says so outright rather than inferring it from + // a tick count over a sleep. + assertEquals(1, sweeper.shutdownNow().size(), "startSweep scheduled more than one sweeper"); + assertEquals(0, holds.sweeps.get()); + manager.close(); + } + + @Test + void theScheduledSweeperActuallySweeps() throws Exception { + InMemoryLeaseStore leases = new InMemoryLeaseStore(CLOCK); + CountingReservations holds = new CountingReservations(); + CreditLeaseManager manager = manager( + new CountingWire(false), + leases, + holds, + CreditLeaseConfig.builder().sweepInterval(Duration.ofMillis(20)).build()); + + manager.startSweep(); + + assertTrue(holds.swept.await(5, TimeUnit.SECONDS), "the sweeper never ran"); + manager.stop(); + manager.close(); + } + + @Test + void aJoinerWaitsNoLongerThanItsOwnDeadline() throws Exception { + InMemoryLeaseStore leases = new InMemoryLeaseStore(CLOCK); + InMemoryReservationStore holds = new InMemoryReservationStore(leases, CLOCK); + CountingWire wire = new CountingWire(true); + CreditLeaseManager manager = manager(wire, leases, holds); + leases.replace(new LeaseGrant("lse_1", "co_1", "ct_1", 1000, LIVE)); + leases.tryReserve("co_1", "ct_1", 900); + + // Somebody else's top-up, held open on the wire for far longer than the joiner has. + Thread owner = new Thread(() -> manager.maybeExtend("co_1", "ct_1", null), "owner"); + owner.start(); + assertTrue(wire.started.await(5, TimeUnit.SECONDS)); + + LeaseState[] joined = new LeaseState[1]; + Thread impatient = new Thread( + () -> joined[0] = manager.maybeExtend("co_1", "ct_1", null, Duration.ofMillis(50)), "impatient"); + impatient.start(); + impatient.join(5000); + + // The flight runs on whatever timeout started it, a background refresh included, so a + // check with 50ms to spend must not inherit it. + assertFalse(impatient.isAlive(), "the joiner sat behind the owner's wire call"); + assertNull(joined[0]); + // It walked away from the wait rather than racing a second extend onto the same lease. + assertEquals(1, wire.extends_.get()); + + wire.release.countDown(); + owner.join(5000); + manager.close(); + } + + @Test + void checksDuringAnInFlightExtendQueueNoBackgroundWork() throws Exception { + InMemoryLeaseStore backing = new InMemoryLeaseStore(CLOCK); + InMemoryReservationStore holds = new InMemoryReservationStore(backing, CLOCK); + CountingReads leases = new CountingReads(backing); + CountingWire wire = new CountingWire(true); + CreditLeaseManager manager = manager(wire, leases, holds); + backing.replace(new LeaseGrant("lse_1", "co_1", "ct_1", 1000, LIVE)); + backing.tryReserve("co_1", "ct_1", 900); + + manager.extendInBackground("co_1", "ct_1"); + assertTrue(wire.started.await(5, TimeUnit.SECONDS)); + int readsWhileTheExtendRuns = leases.reads.get(); + for (int i = 0; i < 200; i++) { + manager.extendInBackground("co_1", "ct_1"); + } + + // The slot stays under its water mark for as long as the top-up is on the wire, so the + // water-mark test alone would have queued a task per check, each of them to read the + // store and find the flight it must not join. + assertEquals(readsWhileTheExtendRuns, leases.reads.get()); + assertEquals(1, wire.extends_.get()); + wire.release.countDown(); + manager.close(); + } + + @Test + void aCallerWaitsOutFlightsTooSmallForItAndThenExtendsForItself() throws Exception { + InMemoryLeaseStore backing = new InMemoryLeaseStore(CLOCK); + InMemoryReservationStore holds = new InMemoryReservationStore(backing, CLOCK); + List asks = Collections.synchronizedList(new ArrayList<>()); + CountDownLatch firstOnTheWire = new CountDownLatch(1); + CountDownLatch releaseFirst = new CountDownLatch(1); + CountDownLatch secondOnTheWire = new CountDownLatch(1); + CountDownLatch releaseSecond = new CountDownLatch(1); + AtomicInteger wireCalls = new AtomicInteger(); + AtomicLong grantedTotal = new AtomicLong(20000); + LeaseWireClient wire = new LeaseWireClient() { + @Override + public LeaseGrant acquire(String companyId, String creditTypeId, double amount, Instant expiresAt) { + return new LeaseGrant("lse_1", companyId, creditTypeId, amount, expiresAt); + } + + @Override + public LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt) { + asks.add(additionalAmount); + int call = wireCalls.incrementAndGet(); + try { + if (call == 1) { + firstOnTheWire.countDown(); + releaseFirst.await(10, TimeUnit.SECONDS); + } else if (call == 2) { + secondOnTheWire.countDown(); + releaseSecond.await(10, TimeUnit.SECONDS); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return new LeaseGrant( + leaseId, "co_1", "ct_1", grantedTotal.addAndGet((long) additionalAmount), expiresAt); + } + + @Override + public void release(String leaseId) {} + }; + AtomicInteger hungryReads = new AtomicInteger(); + // The hungry caller is held at the re-read that follows its first join, until a second, + // smaller flight is registered. That is the race under test: the flight it finds on the + // way back was sized for somebody else's shortfall. + LeaseStore sequenced = new CountingReads(backing) { + @Override + public LeaseState get(String companyId, String creditTypeId) { + if ("hungry".equals(Thread.currentThread().getName()) && hungryReads.incrementAndGet() == 2) { + try { + secondOnTheWire.await(10, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + return super.get(companyId, creditTypeId); + } + }; + CreditLeaseManager manager = manager( + wire, + sequenced, + holds, + CreditLeaseConfig.builder() + .defaultLeaseSize(2000) + .lowWaterMark(0.5) + .build()); + backing.replace(new LeaseGrant("lse_1", "co_1", "ct_1", 20000, LIVE)); + backing.tryReserve("co_1", "ct_1", 19000); + + // A modest shortfall claims the slot and sits on the wire: 4000 wanted, 1000 left, and a + // 2000 tranche floor, so it asks for 3000. + Thread small = new Thread(() -> manager.maybeExtend("co_1", "ct_1", 4000.0), "small"); + small.start(); + assertTrue(firstOnTheWire.await(5, TimeUnit.SECONDS)); + + // The hungry caller needs 18000 more than it has, so that flight covers nothing of its + // shortfall, and it joins. + LeaseState[] hungryGot = new LeaseState[1]; + Thread hungry = new Thread(() -> hungryGot[0] = manager.maybeExtend("co_1", "ct_1", 19000.0), "hungry"); + hungry.start(); + assertTrue(parked(hungry), "the hungry caller never joined the first flight"); + releaseFirst.countDown(); + small.join(5000); + + // A second caller claims the slot while the hungry one is re-reading, and asks for as + // little as the tranche floor allows. + Thread second = new Thread(() -> manager.maybeExtend("co_1", "ct_1", 5000.0), "second"); + second.start(); + assertTrue(secondOnTheWire.await(5, TimeUnit.SECONDS)); + releaseSecond.countDown(); + second.join(5000); + hungry.join(10000); + + // Two joins, then an extend of its own, sized against the balance those two flights left + // rather than the one it started from: 6000 held, 19000 wanted, so it asks for 13000. + // Inheriting either ask would have sent it back to its caller still short. + assertEquals(Arrays.asList(3000.0, 2000.0, 13000.0), asks); + assertNotNull(hungryGot[0]); + assertTrue( + hungryGot[0].getLocalRemainingCredits() >= 19000, + "the hungry caller came back with " + hungryGot[0].getLocalRemainingCredits()); + manager.close(); + } +} diff --git a/src/test/java/com/schematic/api/credits/CreditLeaseManagerShutdownTest.java b/src/test/java/com/schematic/api/credits/CreditLeaseManagerShutdownTest.java new file mode 100644 index 0000000..0464027 --- /dev/null +++ b/src/test/java/com/schematic/api/credits/CreditLeaseManagerShutdownTest.java @@ -0,0 +1,238 @@ +package com.schematic.api.credits; + +import static com.schematic.api.credits.TestThreads.parked; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; + +/** What the manager still owes once it is stopping. */ +class CreditLeaseManagerShutdownTest { + + private static final Instant NOW = Instant.parse("2026-01-01T00:00:00Z"); + private static final Clock CLOCK = Clock.fixed(NOW, ZoneOffset.UTC); + + /** Grants a fixed lease, and runs a hook while the acquire is on the wire. */ + private static final class RacingWire implements LeaseWireClient { + private final Runnable duringAcquire; + final List released = Collections.synchronizedList(new ArrayList<>()); + final List extended = Collections.synchronizedList(new ArrayList<>()); + + RacingWire(Runnable duringAcquire) { + this.duringAcquire = duringAcquire; + } + + @Override + public LeaseGrant acquire(String companyId, String creditTypeId, double requestedAmount, Instant expiresAt) { + duringAcquire.run(); + return new LeaseGrant("lse_wire", companyId, creditTypeId, requestedAmount, expiresAt); + } + + @Override + public LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt) { + extended.add(leaseId); + return new LeaseGrant(leaseId, "co_1", "ct_1", additionalAmount, expiresAt); + } + + @Override + public void release(String leaseId) { + released.add(leaseId); + } + } + + /** Grants a fixed lease, and takes its time handing one back. */ + private static final class SlowWire implements LeaseWireClient { + private final long releaseMillis; + // The lease whose release never returns, standing in for a hung connection. + private final String hangsOn; + final List released = Collections.synchronizedList(new ArrayList<>()); + + SlowWire(long releaseMillis) { + this(releaseMillis, null); + } + + SlowWire(long releaseMillis, String hangsOn) { + this.releaseMillis = releaseMillis; + this.hangsOn = hangsOn; + } + + @Override + public LeaseGrant acquire(String companyId, String creditTypeId, double requestedAmount, Instant expiresAt) { + return new LeaseGrant("lse_wire", companyId, creditTypeId, requestedAmount, expiresAt); + } + + @Override + public LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt) { + return new LeaseGrant(leaseId, "co_1", "ct_1", additionalAmount, expiresAt); + } + + @Override + public void release(String leaseId) { + try { + Thread.sleep(leaseId.equals(hangsOn) ? Long.MAX_VALUE : releaseMillis); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + released.add(leaseId); + } + } + + @Test + void anErrorDuringAnAcquireDoesNotStrandTheJoinersWaitingOnIt() throws Exception { + InMemoryLeaseStore leases = new InMemoryLeaseStore(CLOCK); + InMemoryReservationStore holds = new InMemoryReservationStore(leases, CLOCK); + CountDownLatch acquiring = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + LeaseWireClient wire = new LeaseWireClient() { + @Override + public LeaseGrant acquire(String companyId, String creditTypeId, double amount, Instant expiresAt) { + acquiring.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + throw new StackOverflowError("the wire blew the stack"); + } + + @Override + public LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt) { + return null; + } + + @Override + public void release(String leaseId) {} + }; + CreditLeaseManager manager = new CreditLeaseManager( + wire, leases, holds, CreditLeaseConfig.builder().build(), null, CLOCK); + + Thread first = new Thread(() -> { + try { + manager.acquireIfNeeded("co_1", "ct_1"); + } catch (Error expected) { + // An Error is the caller's to deal with; what matters is who else it takes down. + } + }); + first.start(); + assertTrue(acquiring.await(5, TimeUnit.SECONDS)); + + AtomicReference joined = new AtomicReference<>(); + AtomicReference joinerThrew = new AtomicReference<>(); + CountDownLatch joinerDone = new CountDownLatch(1); + Thread joiner = new Thread(() -> { + try { + joined.set(manager.acquireIfNeeded("co_1", "ct_1")); + } catch (Throwable t) { + // Counted either way, so a regression reads as a failed assertion here rather + // than as a thread that dies quietly and a latch that never finishes. + joinerThrew.set(t); + } finally { + joinerDone.countDown(); + } + }); + joiner.start(); + // Only fail the wire once the joiner is parked on the first flight. Releasing any sooner + // lets that flight finish and deregister, after which the joiner runs an acquire of its + // own and the test no longer covers joining a failed flight. Its thread state is the + // signal, since Flight.await is the one place acquireIfNeeded blocks. + assertTrue(parked(joiner), "the joiner never parked on the flight"); + release.countDown(); + + // Completing only on the return and the RuntimeException paths would park this joiner on + // an unfinished future for the life of the process. + assertTrue(joinerDone.await(10, TimeUnit.SECONDS), "the joiner never came back"); + assertNull(joinerThrew.get()); + // Null because it joined the failed flight, rather than acquiring on its own. + assertNull(joined.get()); + first.join(5000); + joiner.join(5000); + manager.close(); + } + + @Test + void shuttingDownStaysInsideItsBudgetWithManySlotsAndASlowWire() throws Exception { + InMemoryLeaseStore leases = new InMemoryLeaseStore(CLOCK); + InMemoryReservationStore holds = new InMemoryReservationStore(leases, CLOCK); + // One lease whose release never comes back, and forty-nine that take 100ms each. + SlowWire wire = new SlowWire(100, "lse_0"); + CreditLeaseManager manager = new CreditLeaseManager( + wire, leases, holds, CreditLeaseConfig.builder().build(), null, CLOCK); + for (int i = 0; i < 50; i++) { + leases.replace(new LeaseGrant("lse_" + i, "co_" + i, "ct_1", 1000, NOW.plusSeconds(300))); + } + + long startedAt = System.nanoTime(); + manager.releaseAllLocalLeases(Duration.ofSeconds(2)); + manager.close(Duration.ofMillis(100)); + long tookMillis = (System.nanoTime() - startedAt) / 1_000_000; + + // Releasing all fifty in turn is five seconds of shutdown, and one hung release would + // spend the whole budget on its own. Issued together, the budget bounds the set: the + // forty-nine land and the hung one is left to server-side expiry. + assertTrue(tookMillis < 4000, "shutdown took " + tookMillis + "ms"); + assertFalse(wire.released.contains("lse_0"), "the hung release should not have landed"); + for (int i = 1; i < 50; i++) { + assertTrue(wire.released.contains("lse_" + i), "lse_" + i + " was never released"); + } + } + + @Test + void aRedundantLeaseIsStillReleasedWhenTheAcquireLandsMidShutdown() { + InMemoryLeaseStore leases = new InMemoryLeaseStore(CLOCK); + InMemoryReservationStore holds = new InMemoryReservationStore(leases, CLOCK); + CreditLeaseManager[] manager = new CreditLeaseManager[1]; + // A sibling installs its own lease and the client starts closing, both while this acquire + // is on the wire: the lease it is granted is redundant the moment it lands. + RacingWire wire = new RacingWire(() -> { + leases.replace(new LeaseGrant("lse_sibling", "co_1", "ct_1", 1000, NOW.plusSeconds(300))); + manager[0].stop(); + }); + manager[0] = new CreditLeaseManager( + wire, leases, holds, CreditLeaseConfig.builder().build(), null, CLOCK); + + LeaseState result = manager[0].acquireIfNeeded("co_1", "ct_1"); + manager[0].drain(Duration.ofSeconds(5)); + + assertEquals("lse_sibling", result.getLeaseId()); + // Refusing this release would hold the granted credits until the server expires them. + assertEquals(Collections.singletonList("lse_wire"), wire.released); + manager[0].close(); + } + + @Test + void ordinaryBackgroundWorkIsStillRefusedOnceStopped() { + InMemoryLeaseStore leases = new InMemoryLeaseStore(CLOCK); + InMemoryReservationStore holds = new InMemoryReservationStore(leases, CLOCK); + RacingWire wire = new RacingWire(() -> {}); + CreditLeaseManager manager = new CreditLeaseManager( + wire, leases, holds, CreditLeaseConfig.builder().build(), null, CLOCK); + leases.replace(new LeaseGrant("lse_1", "co_1", "ct_1", 1000, NOW.plusSeconds(300))); + // Draw the lease under its low water mark, so a top-up is warranted. + leases.tryReserve("co_1", "ct_1", 900); + + manager.extendInBackground("co_1", "ct_1"); + manager.drain(Duration.ofSeconds(5)); + assertEquals(Collections.singletonList("lse_1"), wire.extended); + + manager.stop(); + manager.extendInBackground("co_1", "ct_1"); + manager.drain(Duration.ofSeconds(5)); + + // A top-up is work the shutdown has no reason to finish: it would install credits the + // release is about to hand back. + assertEquals(1, wire.extended.size()); + manager.close(); + } +} diff --git a/src/test/java/com/schematic/api/credits/DataStreamCreditCheckSourceTest.java b/src/test/java/com/schematic/api/credits/DataStreamCreditCheckSourceTest.java new file mode 100644 index 0000000..b7b141c --- /dev/null +++ b/src/test/java/com/schematic/api/credits/DataStreamCreditCheckSourceTest.java @@ -0,0 +1,94 @@ +package com.schematic.api.credits; + +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.schematic.api.datastream.DataStreamClient; +import com.schematic.api.types.RulesengineCompany; +import com.schematic.api.types.RulesengineUser; +import java.util.Collections; +import java.util.Map; +import org.junit.jupiter.api.Test; + +/** When the lease path is allowed to wait on the socket for an entity, and when it is not. */ +class DataStreamCreditCheckSourceTest { + + private static final Map KEYS = Collections.singletonMap("id", "co_1"); + + private static RulesengineCompany company() { + return RulesengineCompany.builder() + .accountId("acct") + .environmentId("env") + .id("co_1") + .build(); + } + + private static RulesengineUser user() { + return RulesengineUser.builder() + .accountId("acct") + .environmentId("env") + .id("user_1") + .build(); + } + + @Test + void replicatorModeReadsTheCacheRatherThanWaitingOnASocketItDoesNotHave() { + DataStreamClient dataStream = mock(DataStreamClient.class); + when(dataStream.isReplicatorMode()).thenReturn(true); + when(dataStream.getCachedCompany(KEYS)).thenReturn(company()); + when(dataStream.getCachedUser(KEYS)).thenReturn(user()); + DataStreamCreditCheckSource source = new DataStreamCreditCheckSource(dataStream); + + assertSame("co_1", source.getCompany(KEYS).getId()); + assertSame("user_1", source.getUser(KEYS).getId()); + + // A live fetch here has nothing to send the request on, so it only waits out its own + // timeout before answering nothing. + verify(dataStream, never()).getCompany(any()); + verify(dataStream, never()).getUser(any()); + } + + @Test + void aDisconnectedClientReadsTheCacheAndReturnsAMissImmediately() { + DataStreamClient dataStream = mock(DataStreamClient.class); + when(dataStream.isReplicatorMode()).thenReturn(false); + when(dataStream.isConnected()).thenReturn(false); + when(dataStream.getCachedCompany(KEYS)).thenReturn(null); + when(dataStream.getCachedUser(KEYS)).thenReturn(null); + DataStreamCreditCheckSource source = new DataStreamCreditCheckSource(dataStream); + + long startedAt = System.nanoTime(); + assertNull(source.getCompany(KEYS)); + assertNull(source.getUser(KEYS)); + long tookMillis = (System.nanoTime() - startedAt) / 1_000_000; + + // The plain check bails on this state without waiting; the lease path must not pay a + // resource timeout per call before falling back to it. + assertTrue(tookMillis < 500, "the lookups took " + tookMillis + "ms"); + verify(dataStream, never()).getCompany(any()); + verify(dataStream, never()).getUser(any()); + } + + @Test + void aConnectedClientStillFetchesOverTheSocket() { + DataStreamClient dataStream = mock(DataStreamClient.class); + when(dataStream.isReplicatorMode()).thenReturn(false); + when(dataStream.isConnected()).thenReturn(true); + when(dataStream.getCompany(KEYS)).thenReturn(company()); + when(dataStream.getUser(KEYS)).thenReturn(user()); + DataStreamCreditCheckSource source = new DataStreamCreditCheckSource(dataStream); + + assertSame("co_1", source.getCompany(KEYS).getId()); + assertSame("user_1", source.getUser(KEYS).getId()); + + // A cache miss a connected socket can still answer is worth the wait. + verify(dataStream, never()).getCachedCompany(any()); + verify(dataStream, never()).getCachedUser(any()); + } +} diff --git a/src/test/java/com/schematic/api/credits/InMemoryLeaseStoreTest.java b/src/test/java/com/schematic/api/credits/InMemoryLeaseStoreTest.java new file mode 100644 index 0000000..67862a9 --- /dev/null +++ b/src/test/java/com/schematic/api/credits/InMemoryLeaseStoreTest.java @@ -0,0 +1,89 @@ +package com.schematic.api.credits; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.lang.reflect.Field; +import java.time.Clock; +import java.time.Instant; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.junit.jupiter.api.Test; + +/** What the per-process store keeps, and what it lets go of. */ +class InMemoryLeaseStoreTest { + + private static final Instant NOW = Instant.parse("2026-01-01T00:00:00Z"); + private static final Clock CLOCK = Clock.fixed(NOW, ZoneOffset.UTC); + + @Test + void dropRetiresTheSlotLockAlongWithItsState() throws Exception { + InMemoryLeaseStore store = new InMemoryLeaseStore(CLOCK); + store.replace(new LeaseGrant("lse_1", "co_1", "ct_1", 1000, NOW.plusSeconds(300))); + store.tryReserve("co_1", "ct_1", 100); + assertEquals(1, lockCount(store)); + + store.drop("co_1", "ct_1"); + + // A long-lived process leases against many companies, and a lock left behind for each one + // is a slot's worth of memory nothing will ever read again. + assertEquals(0, lockCount(store)); + assertTrue(store.list().isEmpty()); + } + + @Test + void oneSlotIsStillExclusiveWhileItsLockIsBeingRetired() throws Exception { + InMemoryLeaseStore store = new InMemoryLeaseStore(CLOCK); + int rounds = 2000; + // Each lease grants exactly what one reserve takes, so a lease id can be charged once and + // only once. Seeing the same id twice means two threads debited it under different locks. + List charged = Collections.synchronizedList(new ArrayList<>()); + List failures = Collections.synchronizedList(new ArrayList<>()); + + Thread writer = new Thread(() -> { + for (int i = 0; i < rounds; i++) { + store.drop("co_1", "ct_1"); + store.replace(new LeaseGrant("lse_" + i, "co_1", "ct_1", 100, NOW.plusSeconds(300))); + } + }); + List readers = new ArrayList<>(); + for (int t = 0; t < 3; t++) { + readers.add(new Thread(() -> { + for (int i = 0; i < rounds; i++) { + ReserveResult result = store.tryReserve("co_1", "ct_1", 100); + if (result != null) { + charged.add(result.getLeaseId()); + } + } + })); + } + + List all = new ArrayList<>(readers); + all.add(writer); + for (Thread thread : all) { + thread.setUncaughtExceptionHandler((ignored, error) -> failures.add(error)); + thread.start(); + } + for (Thread thread : all) { + thread.join(30_000); + } + + assertEquals(Collections.emptyList(), failures); + Set distinct = new HashSet<>(charged); + assertEquals( + charged.size(), + distinct.size(), + "a lease was charged twice: " + charged.size() + " debits over " + distinct.size() + " leases"); + } + + private static int lockCount(InMemoryLeaseStore store) throws Exception { + Field field = InMemoryLeaseStore.class.getDeclaredField("locks"); + field.setAccessible(true); + return ((Map) field.get(store)).size(); + } +} diff --git a/src/test/java/com/schematic/api/credits/PrewarmCompanyResolverTest.java b/src/test/java/com/schematic/api/credits/PrewarmCompanyResolverTest.java new file mode 100644 index 0000000..d90f8df --- /dev/null +++ b/src/test/java/com/schematic/api/credits/PrewarmCompanyResolverTest.java @@ -0,0 +1,176 @@ +package com.schematic.api.credits; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.schematic.api.types.RulesengineCompany; +import java.time.Duration; +import java.util.Collections; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import org.junit.jupiter.api.Test; + +class PrewarmCompanyResolverTest { + + private static final Map KEYS = Collections.singletonMap("company_id", "acme"); + + private static RulesengineCompany company() { + return RulesengineCompany.builder() + .accountId("acct") + .environmentId("env") + .id("co_1") + .build(); + } + + private static final Function, RulesengineCompany> NOT_CACHED = keys -> null; + private static final Function IGNORE_ERRORS = error -> null; + + private static String resolve( + Map keys, + Function, RulesengineCompany> cached, + Function, RulesengineCompany> fetch, + Duration timeout) { + return PrewarmCompanyResolver.resolve( + keys, cached, fetch, timeout, Duration.ofMillis(1), () -> false, IGNORE_ERRORS); + } + + @Test + void aZeroTimeoutStillAnswersFromTheCache() { + AtomicInteger fetches = new AtomicInteger(); + + String id = resolve( + KEYS, + keys -> company(), + keys -> { + fetches.incrementAndGet(); + return company(); + }, + Duration.ZERO); + + assertEquals("co_1", id); + // Zero means cache-only, so nothing goes to the wire. + assertEquals(0, fetches.get()); + } + + @Test + void aZeroTimeoutGivesUpOnACacheMissWithoutFetching() { + AtomicInteger fetches = new AtomicInteger(); + + String id = resolve( + KEYS, + NOT_CACHED, + keys -> { + fetches.incrementAndGet(); + return company(); + }, + Duration.ZERO); + + assertNull(id); + assertEquals(0, fetches.get()); + } + + @Test + void aKeyNamedIdIsLookedUpLikeAnyOtherKey() { + Map keys = Collections.singletonMap("id", "acme"); + AtomicReference> lookedUp = new AtomicReference<>(); + + String id = resolve( + keys, + cachedKeys -> { + lookedUp.set(cachedKeys); + return company(); + }, + NOT_CACHED, + Duration.ofSeconds(5)); + + // An account is free to define an entity key called `id` holding its own identifier, so + // the lookup settles this, not the name of the key. + assertEquals("co_1", id); + assertEquals(keys, lookedUp.get()); + } + + @Test + void aLookupThatFindsNothingFallsBackToAPrefixedValueUnderAnyKeyName() { + String id = resolve(Collections.singletonMap("account_id", "comp_1"), NOT_CACHED, NOT_CACHED, Duration.ZERO); + + assertEquals("comp_1", id); + } + + @Test + void aLookupThatFindsNothingAndCarriesNoSchematicIdResolvesNothing() { + String id = resolve(Collections.singletonMap("id", "acme"), NOT_CACHED, NOT_CACHED, Duration.ZERO); + + assertNull(id); + } + + @Test + void aFetchThatNeverAnswersStillFallsBackToAPrefixedValue() { + String id = resolve( + Collections.singletonMap("account_id", "comp_1"), NOT_CACHED, keys -> null, Duration.ofMillis(20)); + + assertEquals("comp_1", id); + } + + @Test + void aCacheMissFetchesUntilTheCompanySurfaces() { + AtomicInteger fetches = new AtomicInteger(); + + String id = resolve( + KEYS, NOT_CACHED, keys -> fetches.incrementAndGet() < 3 ? null : company(), Duration.ofSeconds(5)); + + assertEquals("co_1", id); + assertEquals(3, fetches.get()); + } + + @Test + void aFetchThatNeverAnswersGivesUpAtTheTimeout() { + String id = resolve(KEYS, NOT_CACHED, keys -> null, Duration.ofMillis(20)); + + assertNull(id); + } + + @Test + void aCacheThatThrowsIsAMissRatherThanAFailedPrewarm() { + AtomicInteger fetches = new AtomicInteger(); + + String id = resolve( + KEYS, + keys -> { + throw new IllegalStateException("the cache is down"); + }, + keys -> { + fetches.incrementAndGet(); + return company(); + }, + Duration.ofMillis(200)); + + assertEquals("co_1", id); + assertEquals(1, fetches.get()); + } + + @Test + void aFetchThatHangsIsBoundedByTheResolveTimeout() { + long startedAt = System.nanoTime(); + String id = resolve( + KEYS, + NOT_CACHED, + keys -> { + try { + Thread.sleep(30_000); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return company(); + }, + Duration.ofMillis(100)); + long tookMillis = (System.nanoTime() - startedAt) / 1_000_000; + + // The timeout bounds the fetch itself, not just the gaps between attempts: one call that + // never answers would otherwise hold the prewarm past every deadline the caller set. + assertNull(id); + assertTrue(tookMillis < 5000, "the resolve took " + tookMillis + "ms"); + } +} diff --git a/src/test/java/com/schematic/api/credits/RedisReservationStoreFailureTest.java b/src/test/java/com/schematic/api/credits/RedisReservationStoreFailureTest.java new file mode 100644 index 0000000..ea0395a --- /dev/null +++ b/src/test/java/com/schematic/api/credits/RedisReservationStoreFailureTest.java @@ -0,0 +1,76 @@ +package com.schematic.api.credits; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.when; + +import java.time.Clock; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import redis.clients.jedis.JedisPooled; + +@ExtendWith(MockitoExtension.class) +class RedisReservationStoreFailureTest { + + @Mock + private JedisPooled jedis; + + private RedisReservationStore store() { + return new RedisReservationStore(jedis, refuserRefunder(), "test:", Clock.systemUTC()); + } + + @Test + void aRedisFailureReadsAsNothingReserved() { + when(jedis.hgetAll(anyString())).thenThrow(new RuntimeException("redis is down")); + + assertEquals(0.0, store().reservedCredits("comp-1", "tokens"), 0); + } + + @Test + void aRedisFailureReadsAsNoOpenReservations() { + when(jedis.zcard(anyString())).thenThrow(new RuntimeException("redis is down")); + + assertEquals(0, store().count()); + } + + /** + * The SDKs write the check's request body as the eval context, so a sibling can put fields + * beside company and user. Losing the entity keys over one of them would leave a recovered + * hold with nothing to attribute its usage to. + */ + @Test + void anEvalContextWithExtraFieldsStillYieldsTheEntityKeys() { + Map raw = new HashMap<>(); + raw.put("id", "res_1"); + raw.put("leaseId", "lse_1"); + raw.put("companyId", "comp_1"); + raw.put("creditTypeId", "tokens"); + raw.put("eventSubtype", "inference_tokens"); + raw.put("quantityReserved", "2"); + raw.put("creditsReserved", "4"); + raw.put("consumptionRate", "2"); + raw.put("expiresAt", "1767225600000"); + raw.put( + "evalCtx", + "{\"company\":{\"id\":\"comp_1\"},\"user\":{\"user_id\":\"u_1\"}," + + "\"preflight\":{\"event_usage\":{\"event_subtype\":\"inference_tokens\",\"quantity\":2}}}"); + when(jedis.hgetAll(anyString())).thenReturn(raw); + + Reservation reservation = store().get("res_1"); + + assertEquals(Collections.singletonMap("id", "comp_1"), reservation.getCompany()); + assertEquals(Collections.singletonMap("user_id", "u_1"), reservation.getUser()); + } + + /** Neither read touches the lease store, so any call through this one is a test failure. */ + private static ReservationRefunder refuserRefunder() { + return (companyId, creditTypeId, credits, pinLeaseId) -> { + throw new AssertionError("a display-path read must not refund"); + }; + } +} diff --git a/src/test/java/com/schematic/api/credits/ReservationSettlementTest.java b/src/test/java/com/schematic/api/credits/ReservationSettlementTest.java new file mode 100644 index 0000000..7482037 --- /dev/null +++ b/src/test/java/com/schematic/api/credits/ReservationSettlementTest.java @@ -0,0 +1,79 @@ +package com.schematic.api.credits; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.time.Clock; +import java.time.Instant; +import java.time.ZoneOffset; +import org.junit.jupiter.api.Test; + +/** What a settle bills the server, and what it debits the lease. */ +class ReservationSettlementTest { + + private static final Instant NOW = Instant.parse("2026-01-01T00:00:00Z"); + private static final Clock CLOCK = Clock.fixed(NOW, ZoneOffset.UTC); + + @Test + void aFractionalSettleDebitsAndBillsTheSameWholeUnits() { + InMemoryLeaseStore leases = new InMemoryLeaseStore(CLOCK); + InMemoryReservationStore holds = new InMemoryReservationStore(leases, CLOCK); + leases.replace(new LeaseGrant("lse_1", "co_1", "ct_1", 1000, NOW.plusSeconds(300))); + // A hold sized from the caller's declared usage of ten, at two credits a unit. + leases.tryReserve("co_1", "ct_1", 20); + Reservation reservation = new Reservation( + "rsv_1", + "lse_1", + CreditLeaseMode.CLIENT, + "co_1", + "ct_1", + "inference_tokens", + 10, + 20, + 2, + NOW.plusSeconds(60), + null, + null); + holds.add(reservation); + + ReservationSettlement.SettleOutcome outcome = ReservationSettlement.settle(holds, reservation, 7.2); + + assertTrue(outcome.isSettledLocally()); + // The event's quantity is an integer, so 7.2 bills as 8. + assertEquals(8L, outcome.getTrack().getQuantity().get()); + // And the lease is debited for those same eight units at 2 credits apiece, with the 4 + // credits left of the 20-credit hold going back. Debiting the raw 7.2 instead would move + // the local ledger by less than the event bills, and the two would drift over a session. + assertEquals(984.0, leases.get("co_1", "ct_1").getLocalRemainingCredits(), 1e-9); + } + + @Test + void aSettleBelowOneWholeUnitStillBillsOne() { + InMemoryLeaseStore leases = new InMemoryLeaseStore(CLOCK); + InMemoryReservationStore holds = new InMemoryReservationStore(leases, CLOCK); + leases.replace(new LeaseGrant("lse_1", "co_1", "ct_1", 1000, NOW.plusSeconds(300))); + leases.tryReserve("co_1", "ct_1", 2); + Reservation reservation = new Reservation( + "rsv_1", + "lse_1", + CreditLeaseMode.CLIENT, + "co_1", + "ct_1", + "inference_tokens", + 1, + 2, + 2, + NOW.plusSeconds(60), + null, + null); + holds.add(reservation); + + ReservationSettlement.SettleOutcome outcome = ReservationSettlement.settle(holds, reservation, 0.5); + + // The API takes the quantity as a float only to deserialize it, and rejects a non-integer + // while processing the event, so a raw 0.5 would be dropped server-side and never billed + // while the lease had already been debited for it. + assertEquals(1L, outcome.getTrack().getQuantity().get()); + assertEquals(998.0, leases.get("co_1", "ct_1").getLocalRemainingCredits(), 1e-9); + } +} diff --git a/src/test/java/com/schematic/api/credits/SchematicCreditLeaseTest.java b/src/test/java/com/schematic/api/credits/SchematicCreditLeaseTest.java new file mode 100644 index 0000000..1ec9cfe --- /dev/null +++ b/src/test/java/com/schematic/api/credits/SchematicCreditLeaseTest.java @@ -0,0 +1,171 @@ +package com.schematic.api.credits; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.contains; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.schematic.api.IdentifyOptions; +import com.schematic.api.Schematic; +import com.schematic.api.logger.SchematicLogger; +import com.schematic.api.resources.features.FeaturesClient; +import com.schematic.api.types.CheckFlagRequestBody; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.junit.jupiter.api.Test; + +/** The client-side wiring of credit leases: what each configured mode builds, and what it says. */ +class SchematicCreditLeaseTest { + + @Test + void serverModeWarnsAboutTheOptionsItIgnores() { + SchematicLogger logger = mock(SchematicLogger.class); + + try (Schematic schematic = Schematic.builder() + .apiKey("test_api_key") + .logger(logger) + .creditLeases(CreditLeaseConfig.builder() + .mode(CreditLeaseMode.SERVER) + .defaultLeaseSize(500) + .lowWaterMark(0.5) + .build()) + .build()) { + verify(logger).warn(contains("will be ignored")); + // No local plumbing exists to warm. + schematic.prewarm(Collections.singletonMap("id", "co_1"), Collections.singletonList("ct_1")); + verify(logger, never()).error(anyString()); + } + } + + @Test + void clientModeWithoutDataStreamWarnsThatChecksAreNotGated() { + SchematicLogger logger = mock(SchematicLogger.class); + + try (Schematic schematic = Schematic.builder() + .apiKey("test_api_key") + .logger(logger) + .creditLeases(CreditLeaseConfig.builder() + .mode(CreditLeaseMode.CLIENT) + .defaultReservationTtl(Duration.ofSeconds(30)) + .build()) + .build()) { + verify(logger).warn(contains("no credit gating")); + // No shared backend either, which is its own warning. + verify(logger, atLeastOnce()).warn(contains("shared Redis backend")); + assertFalse(schematic.isOffline()); + } + } + + @Test + void clientModeWithoutDataStreamDegradesToAPlainCheckThatHonoursThePerCheckDefault() { + SchematicLogger logger = mock(SchematicLogger.class); + + try (Schematic schematic = Schematic.builder() + .apiKey("test_api_key") + .logger(logger) + .creditLeases( + CreditLeaseConfig.builder().mode(CreditLeaseMode.CLIENT).build()) + .build()) { + FeaturesClient features = mock(FeaturesClient.class); + Schematic spied = spy(schematic); + when(spied.features()).thenReturn(features); + when(features.checkFlag(anyString(), any(CheckFlagRequestBody.class))) + .thenThrow(new RuntimeException("connection refused")); + + CheckResult result = spied.check( + "test_flag", + Collections.singletonMap("id", "co_1"), + null, + CheckOptions.builder().usage(5).defaultValue(true).build()); + + // The intended degradation, rather than a null pointer from a source that wraps + // nothing and so sails past the guard meant to catch this. + verify(logger).debug(contains("no DataStream, using a plain check")); + verify(logger, never()).warn(contains("NullPointerException")); + assertTrue(result.isAllowed()); + } + } + + @Test + void prewarmWithoutCreditTypesIsANoOp() { + SchematicLogger logger = mock(SchematicLogger.class); + + try (Schematic schematic = Schematic.builder() + .apiKey("test_api_key") + .logger(logger) + .creditLeases( + CreditLeaseConfig.builder().mode(CreditLeaseMode.CLIENT).build()) + .build()) { + // Nothing named is nothing to warm, not something to throw at a caller who passed a + // list their own configuration left empty. + schematic.prewarm(Collections.singletonMap("id", "co_1"), null); + schematic.prewarm(Collections.singletonMap("id", "co_1"), Collections.emptyList()); + verify(logger, never()).error(anyString()); + } + } + + @Test + void identifyWithAPrewarmAfterCloseDropsItInsteadOfThrowing() { + SchematicLogger logger = mock(SchematicLogger.class); + Schematic schematic = Schematic.builder() + .apiKey("test_api_key") + .logger(logger) + .creditLeases( + CreditLeaseConfig.builder().mode(CreditLeaseMode.CLIENT).build()) + .build(); + schematic.close(); + + // The prewarm executor is shut down by now, and a caller identifying on a closed client + // should not have to catch the shutdown race that queuing onto it loses. + schematic.identify( + Collections.singletonMap("id", "user_1"), + null, + null, + null, + IdentifyOptions.builder() + .prewarm(Collections.singletonList("ct_1")) + .build()); + + verify(logger).debug(contains("skipping the prewarm")); + } + + @Test + void aPrewarmListIsCopiedOutOfTheCallersHands() { + List creditTypes = new ArrayList<>(); + creditTypes.add("ct_1"); + + IdentifyOptions options = IdentifyOptions.builder().prewarm(creditTypes).build(); + // The prewarm runs in the background and reads this list after identify has returned, so a + // caller reusing their own list must not get to change what gets warmed after the fact. + creditTypes.clear(); + creditTypes.add("ct_2"); + + assertEquals(Collections.singletonList("ct_1"), options.getPrewarm()); + } + + @Test + void aTtlLongerThanTheApiWillHoldIsClampedForServerMode() { + SchematicLogger logger = mock(SchematicLogger.class); + + try (Schematic schematic = Schematic.builder() + .apiKey("test_api_key") + .logger(logger) + .creditLeases(CreditLeaseConfig.builder() + .mode(CreditLeaseMode.SERVER) + .defaultReservationTtl(Duration.ofHours(2)) + .build()) + .build()) { + verify(logger).warn(contains("longer than the API will hold")); + } + } +} diff --git a/src/test/java/com/schematic/api/credits/ServerCreditCheckTest.java b/src/test/java/com/schematic/api/credits/ServerCreditCheckTest.java new file mode 100644 index 0000000..f512c6f --- /dev/null +++ b/src/test/java/com/schematic/api/credits/ServerCreditCheckTest.java @@ -0,0 +1,311 @@ +package com.schematic.api.credits; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.schematic.api.core.BaseSchematicApiException; +import com.schematic.api.core.RequestOptions; +import com.schematic.api.errors.PaymentRequiredError; +import com.schematic.api.resources.credits.CreditsClient; +import com.schematic.api.resources.features.FeaturesClient; +import com.schematic.api.resources.features.requests.CheckAndReserveFlagRequestBody; +import com.schematic.api.resources.features.types.CheckAndReserveFlagResponse; +import com.schematic.api.types.ApiError; +import com.schematic.api.types.CheckAndReserveFlagResponseData; +import com.schematic.api.types.FlagCheckReservationResponseData; +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.time.OffsetDateTime; +import java.time.ZoneOffset; +import java.util.Collections; +import java.util.Map; +import java.util.concurrent.Callable; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +class ServerCreditCheckTest { + + private static final Instant NOW = Instant.parse("2026-01-01T00:00:00Z"); + private static final Map COMPANY = Collections.singletonMap("id", "co_1"); + + private FeaturesClient features; + private CreditsClient credits; + private ServerCreditCheck check; + + @BeforeEach + void setUp() { + features = mock(FeaturesClient.class); + credits = mock(CreditsClient.class); + check = new ServerCreditCheck( + features, credits, null, Duration.ofSeconds(60), Clock.fixed(NOW, ZoneOffset.UTC)); + } + + private static CheckRequest request(double usage, String eventSubtype, boolean failOpen) { + return new CheckRequest("inference", COMPANY, null, usage, eventSubtype, failOpen); + } + + private static Callable fallback(boolean[] called) { + return () -> { + called[0] = true; + return new CheckResult(true, true, "fallback", "inference", null, null, null, null); + }; + } + + private static CheckAndReserveFlagResponse response( + boolean value, String reason, FlagCheckReservationResponseData reservation) { + CheckAndReserveFlagResponseData._FinalStage data = CheckAndReserveFlagResponseData.builder() + .flag("inference") + .reason(reason) + .value(value) + .flagId("flag_1"); + if (reservation != null) { + data.reservation(reservation); + } + return CheckAndReserveFlagResponse.builder().data(data.build()).build(); + } + + private static FlagCheckReservationResponseData hold(String eventSubtype) { + FlagCheckReservationResponseData._FinalStage held = FlagCheckReservationResponseData.builder() + .companyId("co_1") + .consumptionRate(10) + .creditTypeId("ct_1") + .creditsReserved(100) + .expiresAt(OffsetDateTime.ofInstant(NOW.plusSeconds(60), ZoneOffset.UTC)) + .id("res_1") + .quantityReserved(10); + if (eventSubtype != null) { + held.eventSubtype(eventSubtype); + } + return held.build(); + } + + @Test + void takesAHoldAndReturnsAServerModeHandle() { + when(features.checkAndReserveFlag(eq("inference"), any(CheckAndReserveFlagRequestBody.class), any())) + .thenReturn(response(true, "ok", hold("inference_tokens"))); + + CheckResult result = check.check(request(10, null, false), null, () -> false, fallback(new boolean[1])); + + assertTrue(result.isAllowed()); + Reservation reservation = result.getReservation(); + assertNotNull(reservation); + assertEquals("res_1", reservation.getId()); + // No lease exists server-side, so the handle mirrors the hold id and settles by it. + assertEquals("res_1", reservation.getLeaseId()); + assertEquals(CreditLeaseMode.SERVER, reservation.getMode()); + assertEquals("inference_tokens", reservation.getEventSubtype()); + assertEquals(100.0, reservation.getCreditsReserved()); + } + + @Test + void sendsTheHoldWindowThePreflightAndAnIdempotencyKey() { + when(features.checkAndReserveFlag(eq("inference"), any(CheckAndReserveFlagRequestBody.class), any())) + .thenReturn(response(true, "ok", hold("inference_tokens"))); + + check.check(request(10, "inference_tokens", false), null, () -> false, fallback(new boolean[1])); + + ArgumentCaptor body = + ArgumentCaptor.forClass(CheckAndReserveFlagRequestBody.class); + ArgumentCaptor options = ArgumentCaptor.forClass(RequestOptions.class); + verify(features).checkAndReserveFlag(eq("inference"), body.capture(), options.capture()); + assertEquals(10.0, body.getValue().getQuantity().orElse(null)); + assertEquals( + OffsetDateTime.ofInstant(NOW.plusSeconds(60), ZoneOffset.UTC), + body.getValue().getExpiresAt().orElse(null)); + assertEquals( + "inference_tokens", + body.getValue().getPreflight().get().getEventUsage().get().getEventSubtype()); + assertTrue(body.getValue().getIdempotencyKey().isPresent()); + // The key makes the transport's retries safe, so they stay on. + assertFalse(options.getValue().getMaxRetries().isPresent()); + } + + @Test + void sendsAFractionalUsageAsTheQuantityToHold() { + when(features.checkAndReserveFlag(eq("inference"), any(CheckAndReserveFlagRequestBody.class), any())) + .thenReturn(response(true, "ok", hold("inference_tokens"))); + + check.check(request(0.5, "inference_tokens", false), null, () -> false, fallback(new boolean[1])); + + ArgumentCaptor body = + ArgumentCaptor.forClass(CheckAndReserveFlagRequestBody.class); + verify(features).checkAndReserveFlag(eq("inference"), body.capture(), any()); + // The wire quantity is a decimal, so the raw usage goes out and the server holds the same + // amount every other SDK would. Only the preflight rounds up, and only because that field + // is an integer. + assertEquals(0.5, body.getValue().getQuantity().orElse(null)); + assertEquals( + 1L, body.getValue().getPreflight().get().getEventUsage().get().getQuantity()); + } + + @Test + void clampsATimeoutTooLargeForTheRequestOption() { + when(features.checkAndReserveFlag(eq("inference"), any(CheckAndReserveFlagRequestBody.class), any())) + .thenReturn(response(true, "ok", hold("inference_tokens"))); + + check.check(request(10, "inference_tokens", false), Duration.ofDays(30), () -> false, fallback(new boolean[1])); + + ArgumentCaptor options = ArgumentCaptor.forClass(RequestOptions.class); + verify(features) + .checkAndReserveFlag(eq("inference"), any(CheckAndReserveFlagRequestBody.class), options.capture()); + // Casting alone would wrap this into a negative, which the transport reads as no time at + // all: the opposite of the long wait the caller asked for. + assertEquals(Integer.MAX_VALUE, options.getValue().getTimeout().orElse(null)); + } + + @Test + void sendsTheGenericUsagePreflightWithoutASubtype() { + when(features.checkAndReserveFlag(eq("inference"), any(CheckAndReserveFlagRequestBody.class), any())) + .thenReturn(response(true, "ok", hold("inference_tokens"))); + + check.check(request(7.2, null, false), null, () -> false, fallback(new boolean[1])); + + ArgumentCaptor body = + ArgumentCaptor.forClass(CheckAndReserveFlagRequestBody.class); + verify(features).checkAndReserveFlag(eq("inference"), body.capture(), any()); + // A preflight asks an upper-bound question, so a fraction rounds up. + assertEquals(8L, body.getValue().getPreflight().get().getUsage().orElse(null)); + } + + @Test + void deniesWithoutAHoldWhenTheServerSaysTheCreditsAreShort() { + when(features.checkAndReserveFlag(eq("inference"), any(CheckAndReserveFlagRequestBody.class), any())) + .thenReturn(response(false, "Insufficient credits", null)); + + CheckResult result = check.check(request(10, null, false), null, () -> false, fallback(new boolean[1])); + + assertFalse(result.isAllowed()); + assertNull(result.getReservation()); + assertEquals("Insufficient credits", result.getReason()); + } + + @Test + void allowsWithoutAHoldWhenTheServerReturnsNoReservation() { + when(features.checkAndReserveFlag(eq("inference"), any(CheckAndReserveFlagRequestBody.class), any())) + .thenReturn(response(true, "ok", null)); + + boolean[] fellBack = new boolean[1]; + CheckResult result = check.check(request(10, "inference_tokens", false), null, () -> false, fallback(fellBack)); + + // The feature is not credit-metered, so the server allowed it and held nothing. The + // caller gets the verdict and no handle, and has nothing to settle. + assertTrue(result.isAllowed()); + assertNull(result.getReservation()); + assertFalse(fellBack[0]); + verify(credits, never()).releaseCreditReservation(anyString()); + } + + @Test + void treatsAnyApiExceptionCarrying402AsADefinitiveDenial() { + // The endpoint answers 402 through whichever exception the transport builds; the status + // is what makes it definitive, not the class. Failing open must not override it. + when(features.checkAndReserveFlag(eq("inference"), any(CheckAndReserveFlagRequestBody.class), any())) + .thenThrow(new BaseSchematicApiException("payment required", 402, "out of credits")); + + CheckResult result = check.check(request(10, null, true), null, () -> true, fallback(new boolean[1])); + + assertFalse(result.isAllowed()); + assertNull(result.getReservation()); + assertEquals("Insufficient credits", result.getReason()); + } + + @Test + void treatsA402AsADefinitiveDenialEvenWhenFailingOpen() { + when(features.checkAndReserveFlag(eq("inference"), any(CheckAndReserveFlagRequestBody.class), any())) + .thenThrow(new PaymentRequiredError( + ApiError.builder().error("out of credits").build())); + + CheckResult result = check.check(request(10, null, true), null, () -> true, fallback(new boolean[1])); + + assertFalse(result.isAllowed()); + assertEquals("Insufficient credits", result.getReason()); + assertEquals("out of credits", result.getErr()); + } + + @Test + void failsClosedWhenTheCallErrors() { + when(features.checkAndReserveFlag(eq("inference"), any(CheckAndReserveFlagRequestBody.class), any())) + .thenThrow(new RuntimeException("wire down")); + + CheckResult result = check.check(request(10, null, false), null, () -> true, fallback(new boolean[1])); + + assertFalse(result.isAllowed()); + assertEquals("server_reservation_failed", result.getReason()); + assertEquals("server_reservation_failed", result.getErr()); + } + + @Test + void failsOpenToTheCallersDefaultWhenTheCallErrors() { + when(features.checkAndReserveFlag(eq("inference"), any(CheckAndReserveFlagRequestBody.class), any())) + .thenThrow(new RuntimeException("wire down")); + + CheckResult result = check.check(request(10, null, true), null, () -> true, fallback(new boolean[1])); + + assertTrue(result.isAllowed()); + assertEquals("server_reservation_failed_fail_open", result.getReason()); + assertEquals("server_reservation_failed", result.getErr()); + } + + @Test + void releasesAHoldThatNamesNoEventSubtype() { + when(features.checkAndReserveFlag(eq("inference"), any(CheckAndReserveFlagRequestBody.class), any())) + .thenReturn(response(true, "ok", hold(null))); + + CheckResult result = check.check(request(10, null, false), null, () -> false, fallback(new boolean[1])); + + verify(credits).releaseCreditReservation("res_1"); + assertFalse(result.isAllowed()); + assertEquals("missing_event_subtype", result.getErr()); + } + + @Test + void keepsTheServersVerdictForAnUnsettleableHoldWhenFailingOpen() { + when(features.checkAndReserveFlag(eq("inference"), any(CheckAndReserveFlagRequestBody.class), any())) + .thenReturn(response(true, "ok", hold(null))); + + CheckResult result = check.check(request(10, null, true), null, () -> false, fallback(new boolean[1])); + + verify(credits).releaseCreditReservation("res_1"); + assertTrue(result.isAllowed()); + assertNull(result.getReservation()); + assertEquals("missing_event_subtype", result.getErr()); + } + + @Test + void fallsBackToAPlainCheckWhenTheUsageIsZero() { + boolean[] called = {false}; + + CheckResult result = check.check(request(0, null, false), null, () -> false, fallback(called)); + + assertTrue(called[0]); + assertEquals("fallback", result.getReason()); + verify(features, never()).checkAndReserveFlag(any(), any(CheckAndReserveFlagRequestBody.class), any()); + } + + @Test + void resolvesAnInvalidUsageThroughTheFailureContractWithoutCallingTheApi() { + boolean[] called = {false}; + + CheckResult closed = check.check(request(-5, null, false), null, () -> true, fallback(called)); + CheckResult open = check.check(request(Double.NaN, null, true), null, () -> true, fallback(called)); + + assertFalse(called[0]); + assertFalse(closed.isAllowed()); + assertEquals("invalid_usage", closed.getReason()); + assertTrue(open.isAllowed()); + assertEquals("invalid_usage_fail_open", open.getReason()); + verify(features, never()).checkAndReserveFlag(any(), any(CheckAndReserveFlagRequestBody.class), any()); + } +} diff --git a/src/test/java/com/schematic/api/credits/TestThreads.java b/src/test/java/com/schematic/api/credits/TestThreads.java new file mode 100644 index 0000000..ed6e45e --- /dev/null +++ b/src/test/java/com/schematic/api/credits/TestThreads.java @@ -0,0 +1,29 @@ +package com.schematic.api.credits; + +import java.util.concurrent.TimeUnit; + +/** Thread-state helpers shared by the lease tests. */ +final class TestThreads { + + private TestThreads() {} + + /** + * Waits, boundedly, for a thread to block. Its state is the signal a test has that the thread + * reached the wait under test, which is what lets these tests orchestrate a race without a + * sleep long enough to be slow and short enough to be flaky. + */ + static boolean parked(Thread thread) throws InterruptedException { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10); + while (System.nanoTime() - deadline < 0) { + Thread.State state = thread.getState(); + if (state == Thread.State.WAITING || state == Thread.State.TIMED_WAITING) { + return true; + } + if (state == Thread.State.TERMINATED) { + return false; + } + Thread.sleep(1); + } + return false; + } +} diff --git a/src/test/java/com/schematic/api/credits/TrackWithReservationMetricsTest.java b/src/test/java/com/schematic/api/credits/TrackWithReservationMetricsTest.java new file mode 100644 index 0000000..ee9fc14 --- /dev/null +++ b/src/test/java/com/schematic/api/credits/TrackWithReservationMetricsTest.java @@ -0,0 +1,109 @@ +package com.schematic.api.credits; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.schematic.api.Schematic; +import com.schematic.api.datastream.DataStreamClient; +import com.schematic.api.logger.SchematicLogger; +import com.schematic.api.types.EventBodyTrack; +import java.lang.reflect.Field; +import java.time.Instant; +import java.util.Collections; +import java.util.UUID; +import org.junit.jupiter.api.Test; + +/** + * A settle moves the cached company metrics only when it moved local state with it. The event is + * keyed off the reservation id, so the server drops a retried settle as a duplicate: bumping the + * cached metric for one would have the caller's next local evaluation gate on usage counted twice. + */ +class TrackWithReservationMetricsTest { + + @Test + void aSettleThatClaimedTheHoldMovesTheCachedMetrics() { + DataStreamClient dataStream = mock(DataStreamClient.class); + when(dataStream.isConnected()).thenReturn(true); + + try (Schematic schematic = client(dataStream)) { + Reservation reservation = reservation(); + reservations(schematic).add(reservation); + + schematic.trackWithReservation(reservation, 3); + + verify(dataStream).updateCompanyMetrics(any(EventBodyTrack.class)); + } + } + + @Test + void aSettleThatDidNotSettleLocallyLeavesTheCachedMetricsAlone() { + DataStreamClient dataStream = mock(DataStreamClient.class); + // Lenient because reaching the connection check at all is the regression this guards. + lenient().when(dataStream.isConnected()).thenReturn(true); + + try (Schematic schematic = client(dataStream)) { + // The hold was never added to the store, so the settle claims nothing: the state a + // retried settle, an expired hold or an unreachable store all leave behind. + schematic.trackWithReservation(reservation(), 3); + + verify(dataStream, never()).updateCompanyMetrics(any(EventBodyTrack.class)); + } + } + + private static Schematic client(DataStreamClient dataStream) { + Schematic schematic = Schematic.builder() + .apiKey("test_api_key") + .logger(mock(SchematicLogger.class)) + .creditLeases( + CreditLeaseConfig.builder().mode(CreditLeaseMode.CLIENT).build()) + .build(); + // The builder only wires a DataStream from a live socket, and the metrics update is the + // one thing on this path that needs one. + set(schematic, "dataStreamClient", dataStream); + return schematic; + } + + private static Reservation reservation() { + return new Reservation( + UUID.randomUUID().toString(), + "lse_1", + CreditLeaseMode.CLIENT, + "comp_1", + "bilcr_1", + "tokens", + 10, + 10, + 1, + Instant.now().plusSeconds(300), + Collections.singletonMap("company_id", "acme"), + null); + } + + private static ReservationStore reservations(Schematic schematic) { + return (ReservationStore) read(schematic, "reservations"); + } + + private static void set(Schematic schematic, String name, Object value) { + try { + Field field = Schematic.class.getDeclaredField(name); + field.setAccessible(true); + field.set(schematic, value); + } catch (ReflectiveOperationException e) { + throw new AssertionError(e); + } + } + + private static Object read(Schematic schematic, String name) { + try { + Field field = Schematic.class.getDeclaredField(name); + field.setAccessible(true); + return field.get(schematic); + } catch (ReflectiveOperationException e) { + throw new AssertionError(e); + } + } +} diff --git a/src/test/java/com/schematic/api/credits/WasmCreditGateTest.java b/src/test/java/com/schematic/api/credits/WasmCreditGateTest.java new file mode 100644 index 0000000..3c0cd6a --- /dev/null +++ b/src/test/java/com/schematic/api/credits/WasmCreditGateTest.java @@ -0,0 +1,294 @@ +package com.schematic.api.credits; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assumptions.assumeTrue; + +import com.schematic.api.datastream.WasmRulesEngine; +import com.schematic.api.types.ComparableOperator; +import com.schematic.api.types.RulesengineCheckFlagResult; +import com.schematic.api.types.RulesengineCompany; +import com.schematic.api.types.RulesengineCondition; +import com.schematic.api.types.RulesengineConditionType; +import com.schematic.api.types.RulesengineEntitlementValueType; +import com.schematic.api.types.RulesengineFeatureEntitlement; +import com.schematic.api.types.RulesengineFlag; +import com.schematic.api.types.RulesengineRule; +import com.schematic.api.types.RulesengineRuleType; +import com.schematic.api.types.RulesengineUser; +import java.time.Clock; +import java.time.Instant; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.Callable; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +/** + * The flow's other tests script the rules engine, so the contract that matters most is never + * exercised against the real thing: reading the matched credit entitlement off the probe, + * substituting the lease balance into the company, and letting the engine's credit cost gate + * decide. These run the bundled WebAssembly end to end, so a drift in the option envelope or the + * entity shape fails here rather than mis-gating in production. + */ +class WasmCreditGateTest { + + private static final String FLAG_KEY = "infer"; + private static final String CREDIT_ID = "ct_1"; + private static final String SUBTYPE = "inference_tokens"; + private static final Instant NOW = Instant.parse("2026-01-01T00:00:00Z"); + + private static WasmRulesEngine engine; + + @BeforeAll + static void loadEngine() { + WasmRulesEngine candidate = new WasmRulesEngine(null); + try { + candidate.initialize(); + } catch (RuntimeException e) { + // The binary is fetched by scripts/download-wasm.sh, which needs a token. + candidate = null; + } + engine = candidate; + assumeTrue(engine != null, "the rules engine WASM binary is not present"); + } + + private static RulesengineCondition creditCondition() { + return RulesengineCondition.builder() + .accountId("acct") + .conditionType(RulesengineConditionType.CREDIT) + .environmentId("env") + .id("cond_credit") + .operator(ComparableOperator.LT) + .traitValue("") + .creditId(CREDIT_ID) + .consumptionRate(1.0) + .eventSubtype(SUBTYPE) + .build(); + } + + /** A membership condition, so the engine can deny for a reason that is not the balance. */ + private static RulesengineCondition companyCondition(String companyId) { + return RulesengineCondition.builder() + .accountId("acct") + .conditionType(RulesengineConditionType.COMPANY) + .environmentId("env") + .id("cond_company") + .operator(ComparableOperator.EQ) + .traitValue("") + .resourceIds(Collections.singletonList(companyId)) + .build(); + } + + private static RulesengineFlag creditFlag(RulesengineCondition... extra) { + List conditions = new ArrayList<>(); + conditions.add(creditCondition()); + conditions.addAll(Arrays.asList(extra)); + return RulesengineFlag.builder() + .accountId("acct") + .defaultValue(false) + .environmentId("env") + .id("flag_infer") + .key(FLAG_KEY) + .rules(Collections.singletonList(RulesengineRule.builder() + .accountId("acct") + .environmentId("env") + .id("rule_credit") + .name("Credit") + .priority(100) + .ruleType(RulesengineRuleType.PLAN_ENTITLEMENT) + .value(true) + .conditions(conditions) + .build())) + .build(); + } + + private static RulesengineCompany company(double balance) { + return RulesengineCompany.builder() + .accountId("acct") + .environmentId("env") + .id("co_1") + .creditBalances(Collections.singletonMap(CREDIT_ID, balance)) + .entitlements(Collections.singletonList(RulesengineFeatureEntitlement.builder() + .featureId("feat_infer") + .featureKey(FLAG_KEY) + .valueType(RulesengineEntitlementValueType.CREDIT) + .creditId(CREDIT_ID) + .consumptionRate(1.0) + .eventSubtype(SUBTYPE) + .creditTotal(balance) + .creditRemaining(balance) + .build())) + .build(); + } + + /** Serves fixed fixtures and the real engine. */ + private static final class EngineDataStream implements CreditCheckDataStream { + private final RulesengineFlag flag; + private final RulesengineCompany company; + + EngineDataStream(RulesengineFlag flag, RulesengineCompany company) { + this.flag = flag; + this.company = company; + } + + @Override + public RulesengineFlag getFlag(String flagKey) { + return flag; + } + + @Override + public RulesengineCompany getCompany(Map keys) { + return company; + } + + @Override + public RulesengineUser getUser(Map keys) { + return null; + } + + @Override + public RulesengineCheckFlagResult evaluateFlag( + RulesengineFlag flag, RulesengineCompany company, RulesengineUser user, PreflightOptions preflight) + throws Exception { + return engine.checkFlag(flag, company, user, DataStreamCreditCheckSource.toEngineOptions(preflight)); + } + } + + private static final class Fixture { + final InMemoryLeaseStore leases = new InMemoryLeaseStore(Clock.fixed(NOW, ZoneOffset.UTC)); + final InMemoryReservationStore holds = new InMemoryReservationStore(leases, Clock.fixed(NOW, ZoneOffset.UTC)); + final CreditLeaseManager manager; + final CreditCheck flow; + boolean fellBack; + int acquires; + + Fixture(RulesengineFlag flag, RulesengineCompany company) { + LeaseWireClient wire = new LeaseWireClient() { + @Override + public LeaseGrant acquire( + String companyId, String creditTypeId, double requestedAmount, Instant expiresAt) { + acquires++; + return new LeaseGrant("lse_wire", companyId, creditTypeId, requestedAmount, expiresAt); + } + + @Override + public LeaseGrant extend(String leaseId, double additionalAmount, Instant expiresAt) { + throw new UnsupportedOperationException("no extend expected"); + } + + @Override + public void release(String leaseId) {} + }; + manager = new CreditLeaseManager( + wire, leases, holds, CreditLeaseConfig.builder().build(), null, Clock.fixed(NOW, ZoneOffset.UTC)); + flow = new CreditCheck( + new EngineDataStream(flag, company), + leases, + holds, + manager, + null, + Clock.fixed(NOW, ZoneOffset.UTC), + null, + null); + } + + void installLease(double granted) { + leases.replace(new LeaseGrant("lse_1", "co_1", CREDIT_ID, granted, NOW.plusSeconds(300))); + } + + CheckResult run(double usage) { + Callable fallback = () -> { + fellBack = true; + return new CheckResult(true, true, "fallback", FLAG_KEY, null, null, null, null); + }; + CheckResult result = flow.check( + new CheckRequest(FLAG_KEY, Collections.singletonMap("id", "co_1"), null, usage, SUBTYPE, false), + fallback); + manager.drain(java.time.Duration.ofSeconds(5)); + manager.close(); + return result; + } + } + + @Test + void substitutesTheLeaseBalanceAndIssuesAHold() { + Fixture fixture = new Fixture(creditFlag(), company(100)); + fixture.installLease(10000); + + CheckResult result = fixture.run(50); + + assertFalse(fixture.fellBack); + assertTrue(result.isAllowed()); + assertNotNull(result.getReservation()); + assertEquals(CREDIT_ID, result.getReservation().getCreditTypeId()); + assertEquals(50.0, result.getReservation().getCreditsReserved()); + // The hold stays debited from the lease's local view. + assertEquals(9950.0, fixture.leases.get("co_1", CREDIT_ID).getLocalRemainingCredits()); + assertEquals(1, fixture.holds.count()); + } + + @Test + void cancelsTheHoldWhenTheRuleDeniesForANonCreditReason() { + // The balance is plentiful, but the membership condition excludes this company, so the + // engine denies and the hold taken before the gate has to go back. + Fixture fixture = new Fixture(creditFlag(companyCondition("co_other")), company(10000)); + fixture.installLease(10000); + + CheckResult result = fixture.run(50); + + assertFalse(fixture.fellBack); + assertFalse(result.isAllowed()); + assertNull(result.getReservation()); + assertEquals(10000.0, fixture.leases.get("co_1", CREDIT_ID).getLocalRemainingCredits()); + assertEquals(0, fixture.holds.count()); + } + + @Test + void aFractionalUsageIsCarriedRawAndRoundedUpOnlyAtTheEngine() throws Exception { + // The preflight keeps what the caller gave, so the hold can be sized from it exactly. + PreflightOptions preflight = PreflightOptions.fromUsage(0.5, SUBTYPE); + assertEquals(0.5, preflight.getEventUsage().getQuantity()); + // The engine's quantity is an integer, so the boundary rounds up, which is the direction + // an upper-bound question has to round. A balance of one unit still admits half a unit. + RulesengineCheckFlagResult half = engine.checkFlag( + creditFlag(), company(1), null, DataStreamCreditCheckSource.toEngineOptions(preflight)); + RulesengineCheckFlagResult two = engine.checkFlag( + creditFlag(), + company(1), + null, + DataStreamCreditCheckSource.toEngineOptions(PreflightOptions.fromUsage(2.0, SUBTYPE))); + + assertTrue(half.getValue()); + assertFalse(two.getValue()); + } + + @Test + void gatesExactlyAtTheBalanceBoundary() throws Exception { + // The contract the flow leans on: the preflight the SDK builds reaches the engine as the + // event-scoped usage it gates the credit condition with. + RulesengineFlag flag = creditFlag(); + RulesengineCompany company = company(100); + + RulesengineCheckFlagResult under = engine.checkFlag( + flag, + company, + null, + DataStreamCreditCheckSource.toEngineOptions(PreflightOptions.fromUsage(50.0, SUBTYPE))); + RulesengineCheckFlagResult over = engine.checkFlag( + flag, + company, + null, + DataStreamCreditCheckSource.toEngineOptions(PreflightOptions.fromUsage(150.0, SUBTYPE))); + + assertTrue(under.getValue()); + assertFalse(over.getValue()); + } +} diff --git a/src/test/java/com/schematic/api/credits/conformance/Backend.java b/src/test/java/com/schematic/api/credits/conformance/Backend.java new file mode 100644 index 0000000..bf60580 --- /dev/null +++ b/src/test/java/com/schematic/api/credits/conformance/Backend.java @@ -0,0 +1,85 @@ +package com.schematic.api.credits.conformance; + +import com.schematic.api.credits.InMemoryLeaseStore; +import com.schematic.api.credits.InMemoryReservationStore; +import com.schematic.api.credits.LeaseStore; +import com.schematic.api.credits.RedisLeaseStore; +import com.schematic.api.credits.RedisReservationStore; +import com.schematic.api.credits.ReservationStore; +import java.time.Instant; +import java.time.ZoneOffset; +import java.time.ZonedDateTime; +import redis.clients.jedis.JedisPooled; + +/** One store pair plus the seams a vector needs to drive it. */ +final class Backend { + + /** The backend names the vectors use in their {@code backends} field. */ + static final String IN_MEMORY = "in_memory"; + + static final String REDIS = "redis"; + + /** The fixed instant the in-memory timeline starts from. */ + private static final Instant T0 = + ZonedDateTime.of(2026, 1, 1, 0, 0, 0, 0, ZoneOffset.UTC).toInstant(); + + private static final String TEST_KEY_PREFIX = "schematic:"; + + final String name; + final VectorClock clock; + final LeaseStore leases; + final ReservationStore reservations; + final CrashingRefundLeaseStore crash; + + private Backend( + String name, + VectorClock clock, + LeaseStore leases, + ReservationStore reservations, + CrashingRefundLeaseStore crash) { + this.name = name; + this.clock = clock; + this.leases = leases; + this.reservations = reservations; + this.crash = crash; + } + + static Backend create(String name) { + if (IN_MEMORY.equals(name)) { + MutableClock clock = new MutableClock(T0); + VectorClock vectorClock = new VectorClock() { + @Override + public java.time.Clock clock() { + return clock; + } + + @Override + public void advance(long millis) { + clock.advance(millis); + } + + @Override + public Instant at(double offsetMillis) { + return T0.plusMillis((long) offsetMillis); + } + }; + InMemoryLeaseStore leases = new InMemoryLeaseStore(clock); + CrashingRefundLeaseStore crash = new CrashingRefundLeaseStore(leases); + return new Backend(name, vectorClock, leases, new InMemoryReservationStore(crash, clock), crash); + } + if (REDIS.equals(name)) { + JedisPooled jedis = EmbeddedRedis.client(); + jedis.flushAll(); + RedisVectorClock clock = new RedisVectorClock(jedis, TEST_KEY_PREFIX); + RedisLeaseStore leases = new RedisLeaseStore(jedis, TEST_KEY_PREFIX, null, clock.clock()); + CrashingRefundLeaseStore crash = new CrashingRefundLeaseStore(leases); + return new Backend( + name, + clock, + leases, + new RedisReservationStore(jedis, crash, TEST_KEY_PREFIX, clock.clock()), + crash); + } + throw new IllegalArgumentException("unknown conformance backend: " + name); + } +} diff --git a/src/test/java/com/schematic/api/credits/conformance/ConformanceVectorsTest.java b/src/test/java/com/schematic/api/credits/conformance/ConformanceVectorsTest.java new file mode 100644 index 0000000..1f157d7 --- /dev/null +++ b/src/test/java/com/schematic/api/credits/conformance/ConformanceVectorsTest.java @@ -0,0 +1,901 @@ +package com.schematic.api.credits.conformance; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; +import static org.junit.jupiter.api.DynamicTest.dynamicTest; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.schematic.api.credits.CheckRequest; +import com.schematic.api.credits.CheckResult; +import com.schematic.api.credits.CreditCheck; +import com.schematic.api.credits.CreditLeaseConfig; +import com.schematic.api.credits.CreditLeaseDefaults; +import com.schematic.api.credits.CreditLeaseManager; +import com.schematic.api.credits.CreditLeaseMode; +import com.schematic.api.credits.LeaseGrant; +import com.schematic.api.credits.LeaseLister; +import com.schematic.api.credits.LeaseState; +import com.schematic.api.credits.PreflightOptions; +import com.schematic.api.credits.Reservation; +import com.schematic.api.credits.ReservationSettlement; +import com.schematic.api.credits.ReserveResult; +import com.schematic.api.types.EventBodyTrack; +import com.schematic.api.types.RulesengineFeatureEntitlement; +import java.io.File; +import java.io.IOException; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.Callable; +import org.junit.jupiter.api.DynamicTest; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestFactory; + +/** + * Runs the language-agnostic conformance vectors against this SDK. + * + *

The vectors and the semantics they pin live in {@code 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, on every + * store backend it ships. + * + *

The vectors are grouped by what they drive, so a failure names the layer: the stores, the + * lease manager against a scripted wire client, or the whole check and track flow against a + * scripted rules engine. + */ +class ConformanceVectorsTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final List BACKENDS = Arrays.asList(Backend.IN_MEMORY, Backend.REDIS); + + private static final Set STORE_CATEGORIES = + new HashSet<>(Arrays.asList("lease_lifecycle", "reservation_lifecycle", "expiry", "crash_window")); + + private static final Duration DRAIN_TIMEOUT = Duration.ofSeconds(5); + + private static final Set MANAGER_CATEGORIES = new HashSet<>(Collections.singletonList("lease_manager")); + + private static final Set FLOW_CATEGORIES = + new HashSet<>(Arrays.asList("check_flow", "track_settle", "fractional_usage")); + + private static final Set RESERVATION_EXPECT_KEYS = keys( + "lease_id", "credit_type_id", "event_subtype", "quantity_reserved", "credits_reserved", "consumption_rate"); + + private static final Set TRACK_EXPECT_KEYS = keys("event", "quantity", "lease_id", "reservation_id"); + + private static final Set ENGINE_CALL_EXPECT_KEYS = + keys("credit_balance", "credit_cost", "event_usage", "usage"); + + private static final Set EVENT_USAGE_EXPECT_KEYS = keys("event_subtype", "quantity"); + + /** + * What each op asserts. A vector re-synced from the reference implementation can carry an + * expectation this runner has never heard of; without this map the vector would pass while + * asserting nothing, so an unrecognised key fails the same way an unknown op does. + */ + private static final Map> EXPECT_KEYS = expectKeys(); + + private static Map> expectKeys() { + Map> byOp = new HashMap<>(); + byOp.put("advance_clock", keys()); + byOp.put("replace_lease", keys("written")); + byOp.put("drop_lease", keys()); + byOp.put("try_reserve", keys("balance", "lease_id")); + byOp.put("refund_lease", keys()); + byOp.put("extend_lease", keys()); + byOp.put("get_lease", keys("exists", "lease_id", "granted_amount", "local_remaining_credits")); + byOp.put("add_reservation", keys()); + byOp.put("consume_reservation", keys("consumed", "throws")); + byOp.put("get_reservation", keys("exists")); + byOp.put("reserved_credits", keys("total")); + byOp.put("reservation_count", keys("count")); + byOp.put( + "check", + keys( + "allowed", + "reason", + "err", + "has_reservation", + "fallback_called", + "reservation", + "engine_calls", + "wire_extends", + "last_extend_additional_amount")); + byOp.put("track", keys("settled_locally", "track")); + byOp.put( + "acquire_if_needed", + keys("lease_id", "wire_acquires", "last_acquire_requested_amount", "released_lease_ids")); + byOp.put("maybe_extend", keys("wire_extends", "last_extend_additional_amount", "last_extend_lease_id")); + byOp.put("release_all_local_leases", keys("released_lease_ids", "remaining_slots")); + byOp.put("sweep_expired", keys("swept")); + return byOp; + } + + private static Set keys(String... names) { + return new HashSet<>(Arrays.asList(names)); + } + + @TestFactory + Collection storeVectors() { + return vectorsFor(STORE_CATEGORIES); + } + + @TestFactory + Collection managerVectors() { + return vectorsFor(MANAGER_CATEGORIES); + } + + @TestFactory + Collection flowVectors() { + return vectorsFor(FLOW_CATEGORIES); + } + + /** Every vector file has to be claimed by a group, or it would silently never run. */ + @Test + void everyCategoryIsCovered() { + Set claimed = new HashSet<>(STORE_CATEGORIES); + claimed.addAll(MANAGER_CATEGORIES); + claimed.addAll(FLOW_CATEGORIES); + for (JsonNode document : loadDocuments()) { + String category = document.get("category").asText(); + assertTrue(claimed.contains(category), "conformance category not claimed by any group: " + category); + } + } + + private Collection vectorsFor(Set categories) { + List tests = new ArrayList<>(); + for (JsonNode document : loadDocuments()) { + String category = document.get("category").asText(); + if (!categories.contains(category)) { + continue; + } + for (JsonNode vector : document.get("vectors")) { + for (String backend : BACKENDS) { + if (!runsOn(vector, backend)) { + continue; + } + String name = backend + " - " + category + " - " + + vector.get("name").asText(); + tests.add(dynamicTest(name, () -> runVector(backend, vector))); + } + } + } + assertFalse(tests.isEmpty(), "no conformance vectors found for " + categories); + return tests; + } + + private static boolean runsOn(JsonNode vector, String backend) { + JsonNode backends = vector.get("backends"); + if (backends == null || !backends.isArray()) { + return true; + } + for (JsonNode allowed : backends) { + if (allowed.asText().equals(backend)) { + return true; + } + } + return false; + } + + private static List loadDocuments() { + File directory = new File("conformance/vectors"); + File[] files = directory.listFiles((dir, name) -> name.endsWith(".json")); + assertNotNull(files, "conformance vectors directory not found at " + directory.getAbsolutePath()); + Arrays.sort(files); + List documents = new ArrayList<>(); + for (File file : files) { + try { + documents.add(MAPPER.readTree(file)); + } catch (IOException e) { + throw new IllegalStateException("failed to parse " + file, e); + } + } + return documents; + } + + private void runVector(String backendName, JsonNode vector) { + Harness harness = new Harness(Backend.create(backendName), vector.get("given")); + String vectorName = vector.get("name").asText(); + try { + harness.installGivenLeases(); + for (JsonNode op : vector.get("operations")) { + runOperation(harness, op, vectorName); + } + } finally { + harness.close(); + } + } + + private void runOperation(Harness h, JsonNode op, String vectorName) { + String name = op.get("op").asText(); + JsonNode expect = op.has("expect") ? op.get("expect") : MAPPER.createObjectNode(); + switch (name) { + case "advance_clock": + h.backend.clock.advance(op.get("ms").asLong()); + break; + case "replace_lease": + opReplaceLease(h, op, expect); + break; + case "drop_lease": + h.backend.leases.drop(text(op, "company_id"), text(op, "credit_type_id")); + break; + case "try_reserve": + opTryReserve(h, op, expect); + break; + case "refund_lease": + h.backend.leases.refund( + text(op, "company_id"), + text(op, "credit_type_id"), + op.get("credits").asDouble(), + text(op, "pin_lease_id")); + break; + case "extend_lease": + h.backend.leases.extend( + text(op, "company_id"), + text(op, "credit_type_id"), + op.get("granted_total").asDouble(), + op.has("expires_at_ms") + ? h.backend.clock.at(op.get("expires_at_ms").asDouble()) + : null, + text(op, "pin_lease_id")); + break; + case "get_lease": + opGetLease(h, op, expect); + break; + case "add_reservation": + opAddReservation(h, op); + break; + case "consume_reservation": + opConsumeReservation(h, op, expect); + break; + case "get_reservation": + assertEquals( + expect.get("exists").asBoolean(), + h.backend.reservations.get(h.resolveReservationId(op)) != null, + "get_reservation exists"); + break; + case "reserved_credits": + assertEquals( + expect.get("total").asDouble(), + h.backend.reservations.reservedCredits(text(op, "company_id"), text(op, "credit_type_id")), + "reserved_credits total"); + break; + case "reservation_count": + assertEquals(expect.get("count").asInt(), h.backend.reservations.count(), "reservation_count"); + break; + case "check": + opCheck(h, op, expect); + break; + case "track": + opTrack(h, op, expect); + break; + case "acquire_if_needed": + opAcquireIfNeeded(h, op, expect); + break; + case "maybe_extend": + opMaybeExtend(h, op, expect); + break; + case "release_all_local_leases": + opReleaseAllLocalLeases(h, expect); + break; + case "sweep_expired": + int swept = h.backend.reservations.sweepExpired(); + if (expect.has("swept")) { + assertEquals(expect.get("swept").asInt(), swept, "sweep_expired swept"); + } + break; + default: + fail("unknown conformance op: " + name); + } + assertExpectHandled(name, expect, vectorName); + } + + /** Fails on an expectation key no assertion above consumed, naming the key and the vector. */ + private static void assertExpectHandled(String op, JsonNode expect, String vectorName) { + Set known = EXPECT_KEYS.get(op); + assertNotNull(known, "conformance op " + op + " declares no expectation keys"); + assertHandledKeys(expect, known, "expect", op, vectorName); + assertHandledKeys(expect.get("reservation"), RESERVATION_EXPECT_KEYS, "expect.reservation", op, vectorName); + assertHandledKeys(expect.get("track"), TRACK_EXPECT_KEYS, "expect.track", op, vectorName); + JsonNode calls = expect.get("engine_calls"); + if (calls == null || !calls.isArray()) { + return; + } + for (int i = 0; i < calls.size(); i++) { + String where = "expect.engine_calls[" + i + "]"; + assertHandledKeys(calls.get(i), ENGINE_CALL_EXPECT_KEYS, where, op, vectorName); + assertHandledKeys( + calls.get(i).get("event_usage"), EVENT_USAGE_EXPECT_KEYS, where + ".event_usage", op, vectorName); + } + } + + private static void assertHandledKeys( + JsonNode node, Set known, String where, String op, String vectorName) { + if (node == null || !node.isObject()) { + return; + } + List unhandled = new ArrayList<>(); + Iterator fields = node.fieldNames(); + while (fields.hasNext()) { + String field = fields.next(); + if (!known.contains(field)) { + unhandled.add(field); + } + } + if (!unhandled.isEmpty()) { + fail("unhandled conformance expectation " + unhandled + " in " + where + " of op " + op + " in vector " + + vectorName); + } + } + + private static void opReplaceLease(Harness h, JsonNode op, JsonNode expect) { + boolean wrote = h.backend.leases.replace(new LeaseGrant( + text(op, "lease_id"), + text(op, "company_id"), + text(op, "credit_type_id"), + op.get("granted_amount").asDouble(), + h.backend.clock.at(op.get("expires_at_ms").asDouble()))); + if (expect.has("written")) { + assertEquals(expect.get("written").asBoolean(), wrote, "replace_lease written"); + } + } + + private static void opTryReserve(Harness h, JsonNode op, JsonNode expect) { + ReserveResult result = h.backend.leases.tryReserve( + text(op, "company_id"), + text(op, "credit_type_id"), + op.get("credits").asDouble()); + if (expect.has("balance")) { + if (expect.get("balance").isNull()) { + assertNull(result, "try_reserve should have been refused"); + } else { + assertNotNull(result, "try_reserve should have succeeded"); + assertEquals(expect.get("balance").asDouble(), result.getBalance(), "try_reserve balance"); + } + } + // The charged lease is what a caller pins its reservation to, so a vector that names one + // is checking the pin, not just the arithmetic. + if (expect.has("lease_id")) { + if (expect.get("lease_id").isNull()) { + assertNull(result, "try_reserve should have been refused"); + } else { + assertNotNull(result, "try_reserve should have succeeded"); + assertEquals(expect.get("lease_id").asText(), result.getLeaseId(), "try_reserve lease_id"); + } + } + } + + private static void opGetLease(Harness h, JsonNode op, JsonNode expect) { + LeaseState entry = h.backend.leases.get(text(op, "company_id"), text(op, "credit_type_id")); + if (expect.has("exists")) { + assertEquals(expect.get("exists").asBoolean(), entry != null, "get_lease exists"); + } + if (expect.has("lease_id")) { + if (expect.get("lease_id").isNull()) { + assertNull(entry, "get_lease should have found no lease"); + } else { + assertNotNull(entry, "get_lease should have found a lease"); + assertEquals(expect.get("lease_id").asText(), entry.getLeaseId(), "get_lease lease_id"); + } + } + if (expect.has("granted_amount")) { + assertNotNull(entry, "get_lease should have found a lease"); + assertEquals(expect.get("granted_amount").asDouble(), entry.getGrantedAmount(), "get_lease granted"); + } + if (expect.has("local_remaining_credits")) { + assertNotNull(entry, "get_lease should have found a lease"); + assertEquals( + expect.get("local_remaining_credits").asDouble(), + entry.getLocalRemainingCredits(), + "get_lease local_remaining_credits"); + } + } + + private static void opAddReservation(Harness h, JsonNode op) { + Map company = new HashMap<>(); + company.put("id", text(op, "company_id")); + h.backend.reservations.add(new Reservation( + text(op, "id"), + text(op, "lease_id"), + CreditLeaseMode.CLIENT, + text(op, "company_id"), + text(op, "credit_type_id"), + text(op, "event_subtype"), + op.get("quantity_reserved").asDouble(), + op.get("credits_reserved").asDouble(), + op.get("consumption_rate").asDouble(), + h.backend.clock.at(op.get("expires_at_ms").asDouble()), + company, + null)); + } + + private static void opConsumeReservation(Harness h, JsonNode op, JsonNode expect) { + String id = h.resolveReservationId(op); + double credits = op.get("credits").asDouble(); + if (op.has("crash_before_refund") && op.get("crash_before_refund").asBoolean()) { + h.backend.crash.arm(); + // Driven by the vector rather than asserting the vector back at itself, which would + // pass whatever it said and break on a vector that omits the key. + if (expect.has("throws") && expect.get("throws").asBoolean()) { + assertThrows( + CrashingRefundLeaseStore.SimulatedCrash.class, + () -> h.backend.reservations.consume(id, credits), + "consume_reservation should have crashed before the refund"); + } else { + assertDoesNotThrow( + () -> h.backend.reservations.consume(id, credits), + "consume_reservation should not have crashed"); + } + return; + } + Double consumed = h.backend.reservations.consume(id, credits); + if (expect.has("consumed")) { + if (expect.get("consumed").isNull()) { + assertNull(consumed, "consume_reservation should have claimed nothing"); + } else { + assertNotNull(consumed, "consume_reservation should have claimed the hold"); + assertEquals(expect.get("consumed").asDouble(), consumed, "consume_reservation consumed"); + } + } + } + + private static void opCheck(Harness h, JsonNode op, JsonNode expect) { + String flagKey = op.has("flag_key") ? op.get("flag_key").asText() : "flag"; + JsonNode companySpec = op.get("company"); + String companyId = companySpec != null && companySpec.has("id") + ? companySpec.get("id").asText() + : "co_1"; + Map balances = new LinkedHashMap<>(); + if (companySpec != null && companySpec.has("credit_balances")) { + Iterator> fields = + companySpec.get("credit_balances").fields(); + while (fields.hasNext()) { + Map.Entry field = fields.next(); + balances.put(field.getKey(), field.getValue().asDouble()); + } + } + List results = new ArrayList<>(); + if (op.has("engine")) { + for (JsonNode scripted : op.get("engine")) { + results.add(new ScriptedCheckDataStream.Result( + scripted.get("value").asBoolean(), + scripted.has("reason") ? scripted.get("reason").asText() : null, + entitlementFrom(scripted.get("entitlement")))); + } + } + if (op.has("server")) { + JsonNode server = op.get("server"); + if (server.has("acquire")) { + h.wire.queueAcquire(acquireScript(h, server.get("acquire"))); + } + if (server.has("extend")) { + h.wire.queueExtend(extendScript(h, server.get("extend"))); + } + } + + ScriptedCheckDataStream dataStream = + new ScriptedCheckDataStream(flagKey, ScriptedCheckDataStream.company(companyId, balances), results); + boolean[] fellBack = {false}; + Callable fallback = () -> { + fellBack[0] = true; + return new CheckResult(true, true, "fallback", flagKey, null, null, null, null); + }; + + CreditCheck flow = new CreditCheck( + dataStream, + h.backend.leases, + h.backend.reservations, + h.manager(), + null, + h.backend.clock.clock(), + null, + null); + CheckResult result = flow.check( + new CheckRequest( + flagKey, + Collections.singletonMap("id", companyId), + null, + op.has("usage") ? op.get("usage").asDouble() : 0, + op.has("event_subtype") ? op.get("event_subtype").asText() : null, + "fail-open".equals(text(op, "on_acquire_failure"))), + fallback); + h.manager().drain(DRAIN_TIMEOUT); + + if (expect.has("allowed")) { + assertEquals(expect.get("allowed").asBoolean(), result.isAllowed(), "allowed"); + } + if (expect.has("reason")) { + assertEquals(expect.get("reason").asText(), result.getReason(), "reason"); + } + if (expect.has("err")) { + assertEquals(expect.get("err").asText(), result.getErr(), "err"); + } + if (expect.has("has_reservation")) { + assertEquals(expect.get("has_reservation").asBoolean(), result.getReservation() != null, "has_reservation"); + } + if (expect.has("fallback_called")) { + assertEquals(expect.get("fallback_called").asBoolean(), fellBack[0], "fallback_called"); + } + assertReservation(expect.get("reservation"), result.getReservation()); + assertEngineCalls(expect.get("engine_calls"), dataStream.calls, creditIdFor(op, balances)); + if (expect.has("wire_extends")) { + assertEquals(expect.get("wire_extends").asInt(), h.wire.extendCalls.size(), "wire_extends"); + } + if (expect.has("last_extend_additional_amount")) { + assertFalse(h.wire.extendCalls.isEmpty(), "no extend was made"); + assertEquals( + expect.get("last_extend_additional_amount").asDouble(), + h.wire.extendCalls.get(h.wire.extendCalls.size() - 1).additionalAmount, + "last_extend_additional_amount"); + } + if (op.has("save_reservation_as") && result.getReservation() != null) { + h.handles.put(op.get("save_reservation_as").asText(), result.getReservation()); + } + } + + private static void opTrack(Harness h, JsonNode op, JsonNode expect) { + Reservation reservation = h.handles.get(op.get("handle").asText()); + assertNotNull( + reservation, "unknown reservation handle: " + op.get("handle").asText()); + ReservationSettlement.SettleOutcome outcome = ReservationSettlement.settle( + h.backend.reservations, reservation, op.get("actual_quantity").asDouble()); + + if (expect.has("settled_locally")) { + assertEquals(expect.get("settled_locally").asBoolean(), outcome.isSettledLocally(), "settled_locally"); + } + JsonNode want = expect.get("track"); + if (want == null) { + return; + } + EventBodyTrack track = outcome.getTrack(); + if (want.has("event")) { + assertEquals(want.get("event").asText(), track.getEvent(), "track event"); + } + if (want.has("quantity")) { + // Asserted exactly as the vector states it. Rounding the expectation here would let a + // vector that states a fractional quantity pass in this SDK and fail in the others. + assertEquals(want.get("quantity").asLong(), track.getQuantity().orElse(null), "track quantity"); + } + if (want.has("lease_id")) { + assertEquals(want.get("lease_id").asText(), track.getLeaseId().orElse(null), "track lease_id"); + } + if (want.has("reservation_id")) { + assertEquals( + want.get("reservation_id").asText(), + track.getReservationId().orElse(null), + "track reservation_id"); + } + } + + private static void assertReservation(JsonNode want, Reservation reservation) { + if (want == null) { + return; + } + assertNotNull(reservation, "expected a reservation"); + if (want.has("lease_id")) { + assertEquals(want.get("lease_id").asText(), reservation.getLeaseId(), "reservation lease_id"); + } + if (want.has("credit_type_id")) { + assertEquals( + want.get("credit_type_id").asText(), reservation.getCreditTypeId(), "reservation credit_type_id"); + } + if (want.has("event_subtype")) { + assertEquals( + want.get("event_subtype").asText(), reservation.getEventSubtype(), "reservation event_subtype"); + } + if (want.has("quantity_reserved")) { + assertEquals( + want.get("quantity_reserved").asDouble(), + reservation.getQuantityReserved(), + "reservation quantity_reserved"); + } + if (want.has("credits_reserved")) { + assertEquals( + want.get("credits_reserved").asDouble(), + reservation.getCreditsReserved(), + "reservation credits_reserved"); + } + if (want.has("consumption_rate")) { + assertEquals( + want.get("consumption_rate").asDouble(), + reservation.getConsumptionRate(), + "reservation consumption_rate"); + } + } + + private static void assertEngineCalls( + JsonNode want, List calls, String creditId) { + if (want == null) { + return; + } + assertEquals(want.size(), calls.size(), "engine call count"); + for (int i = 0; i < want.size(); i++) { + JsonNode expected = want.get(i); + ScriptedCheckDataStream.EngineCall got = calls.get(i); + if (expected.has("credit_balance")) { + assertNotNull(creditId, "engine_calls needs a credit id to assert a balance against"); + assertNotNull(got.creditBalances, "the engine was given no company"); + assertEquals( + expectedCreditBalance(expected.get("credit_balance")), + got.creditBalances.get(creditId), + "engine call " + i + " credit_balance"); + } + if (expected.has("credit_cost")) { + assertNotNull(got.preflight, "engine call " + i + " carried no preflight"); + assertNotNull(got.preflight.getCreditCost(), "engine call " + i + " carried no credit cost"); + assertEquals( + expected.get("credit_cost").asDouble(), + got.preflight.getCreditCost().get(creditId), + "engine call " + i + " credit_cost"); + } + if (expected.has("event_usage")) { + assertNotNull(got.preflight, "engine call " + i + " carried no preflight"); + PreflightOptions.EventUsage eventUsage = got.preflight.getEventUsage(); + assertNotNull(eventUsage, "engine call " + i + " carried no event usage"); + assertEquals( + expected.get("event_usage").get("event_subtype").asText(), + eventUsage.getEventSubtype(), + "engine call " + i + " event_subtype"); + assertEquals( + expected.get("event_usage").get("quantity").asLong(), + eventUsage.getQuantity(), + "engine call " + i + " event quantity"); + } + if (expected.has("usage")) { + assertNotNull(got.preflight, "engine call " + i + " carried no preflight"); + assertEquals( + (Long) expected.get("usage").asLong(), got.preflight.getUsage(), "engine call " + i + " usage"); + } + } + } + + /** A balance expectation is a number, or the name of the fail-open substitution. */ + private static Double expectedCreditBalance(JsonNode raw) { + if (raw.isTextual() && "max_safe_integer".equals(raw.asText())) { + return CreditLeaseDefaults.FAIL_OPEN_BALANCE; + } + return raw.asDouble(); + } + + /** + * 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. + */ + private static String creditIdFor(JsonNode op, Map balances) { + if (op.has("engine")) { + for (JsonNode scripted : op.get("engine")) { + JsonNode entitlement = scripted.get("entitlement"); + if (entitlement != null && entitlement.has("credit_id")) { + return entitlement.get("credit_id").asText(); + } + } + } + for (String creditId : balances.keySet()) { + return creditId; + } + return null; + } + + private static RulesengineFeatureEntitlement entitlementFrom(JsonNode spec) { + if (spec == null) { + return null; + } + return ScriptedCheckDataStream.entitlement( + spec.get("value_type").asText(), + spec.has("credit_id") ? spec.get("credit_id").asText() : null, + spec.has("consumption_rate") ? spec.get("consumption_rate").asDouble() : null, + spec.has("event_subtype") ? spec.get("event_subtype").asText() : null); + } + + private static void opAcquireIfNeeded(Harness h, JsonNode op, JsonNode expect) { + if (op.has("server")) { + h.wire.queueAcquire(acquireScript(h, op.get("server"))); + } + if (op.has("install_during_wire")) { + JsonNode install = op.get("install_during_wire"); + h.wire.duringAcquire = () -> h.backend.leases.replace(new LeaseGrant( + install.get("lease_id").asText(), + install.get("company_id").asText(), + install.get("credit_type_id").asText(), + install.get("granted_amount").asDouble(), + h.backend.clock.at(install.get("expires_at_ms").asDouble()))); + } + LeaseState entry = h.manager().acquireIfNeeded(text(op, "company_id"), text(op, "credit_type_id")); + h.manager().drain(DRAIN_TIMEOUT); + + if (expect.has("lease_id")) { + if (expect.get("lease_id").isNull()) { + assertNull(entry, "acquire_if_needed should have found no lease"); + } else { + assertNotNull(entry, "acquire_if_needed should have returned a lease"); + assertEquals(expect.get("lease_id").asText(), entry.getLeaseId(), "acquire_if_needed lease_id"); + } + } + if (expect.has("wire_acquires")) { + assertEquals(expect.get("wire_acquires").asInt(), h.wire.acquireCalls.size(), "wire_acquires"); + } + if (expect.has("last_acquire_requested_amount")) { + assertFalse(h.wire.acquireCalls.isEmpty(), "no acquire was made"); + assertEquals( + expect.get("last_acquire_requested_amount").asDouble(), + h.wire.acquireCalls.get(h.wire.acquireCalls.size() - 1).requestedAmount, + "last_acquire_requested_amount"); + } + assertReleasedLeaseIds(h, expect); + } + + private static void opMaybeExtend(Harness h, JsonNode op, JsonNode expect) { + if (op.has("server")) { + h.wire.queueExtend(extendScript(h, op.get("server"))); + } + Double required = + op.has("required_credits") ? op.get("required_credits").asDouble() : null; + h.manager().maybeExtend(text(op, "company_id"), text(op, "credit_type_id"), required); + h.manager().drain(DRAIN_TIMEOUT); + + if (expect.has("wire_extends")) { + assertEquals(expect.get("wire_extends").asInt(), h.wire.extendCalls.size(), "wire_extends"); + } + if (expect.has("last_extend_additional_amount")) { + assertFalse(h.wire.extendCalls.isEmpty(), "no extend was made"); + assertEquals( + expect.get("last_extend_additional_amount").asDouble(), + h.wire.extendCalls.get(h.wire.extendCalls.size() - 1).additionalAmount, + "last_extend_additional_amount"); + } + if (expect.has("last_extend_lease_id")) { + assertFalse(h.wire.extendCalls.isEmpty(), "no extend was made"); + assertEquals( + expect.get("last_extend_lease_id").asText(), + h.wire.extendCalls.get(h.wire.extendCalls.size() - 1).leaseId, + "last_extend_lease_id"); + } + } + + private static void opReleaseAllLocalLeases(Harness h, JsonNode expect) { + h.manager().releaseAllLocalLeases(); + assertReleasedLeaseIds(h, expect); + if (expect.has("remaining_slots")) { + assertTrue(h.backend.leases instanceof LeaseLister, "remaining_slots needs an enumerable store"); + assertEquals( + expect.get("remaining_slots").asInt(), + ((LeaseLister) h.backend.leases).list().size(), + "remaining_slots"); + } + } + + private static void assertReleasedLeaseIds(Harness h, JsonNode expect) { + if (!expect.has("released_lease_ids")) { + return; + } + List released = new ArrayList<>(); + for (JsonNode id : expect.get("released_lease_ids")) { + released.add(id.asText()); + } + assertEquals(released, h.wire.releasedLeaseIds, "released_lease_ids"); + } + + private static ScriptedWireClient.Script acquireScript(Harness h, JsonNode server) { + if (server.has("error")) { + return ScriptedWireClient.Script.error(server.get("error").asText()); + } + JsonNode lease = server.get("lease"); + return ScriptedWireClient.Script.lease( + lease.get("lease_id").asText(), + lease.get("granted_amount").asDouble(), + h.backend.clock.at(lease.get("expires_at_ms").asDouble())); + } + + private static ScriptedWireClient.Script extendScript(Harness h, JsonNode server) { + if (server.has("error")) { + return ScriptedWireClient.Script.error(server.get("error").asText()); + } + JsonNode lease = server.get("lease"); + return ScriptedWireClient.Script.lease( + lease.has("lease_id") ? lease.get("lease_id").asText() : null, + lease.get("granted_total").asDouble(), + h.backend.clock.at(lease.get("expires_at_ms").asDouble())); + } + + private static String text(JsonNode op, String field) { + JsonNode value = op.get(field); + return value == null || value.isNull() ? null : value.asText(); + } + + /** One vector's stores, clock and reservation handles. */ + static final class Harness { + + final Backend backend; + final JsonNode given; + final Map handles = new HashMap<>(); + final ScriptedWireClient wire = new ScriptedWireClient(); + private CreditLeaseManager manager; + + Harness(Backend backend, JsonNode given) { + this.backend = backend; + this.given = given == null ? MAPPER.createObjectNode() : given; + } + + /** The manager under the vector's config, built on first use. */ + CreditLeaseManager manager() { + if (manager == null) { + JsonNode config = config(); + CreditLeaseConfig.Builder builder = CreditLeaseConfig.builder(); + if (config.has("lease_duration_ms")) { + builder.defaultLeaseDuration( + Duration.ofMillis(config.get("lease_duration_ms").asLong())); + } + if (config.has("reservation_ttl_ms")) { + builder.defaultReservationTtl( + Duration.ofMillis(config.get("reservation_ttl_ms").asLong())); + } + if (config.has("lease_size")) { + builder.defaultLeaseSize(config.get("lease_size").asDouble()); + } + if (config.has("low_water_mark")) { + builder.lowWaterMark(config.get("low_water_mark").asDouble()); + } + manager = new CreditLeaseManager( + wire, backend.leases, backend.reservations, builder.build(), null, backend.clock.clock()); + } + return manager; + } + + void installGivenLeases() { + JsonNode leases = given.get("leases"); + if (leases == null) { + return; + } + for (JsonNode lease : leases) { + boolean wrote = backend.leases.replace(new LeaseGrant( + lease.get("lease_id").asText(), + lease.get("company_id").asText(), + lease.get("credit_type_id").asText(), + lease.get("granted_amount").asDouble(), + backend.clock.at(lease.get("expires_at_ms").asDouble()))); + assertTrue(wrote, "given.leases must install cleanly"); + } + } + + JsonNode config() { + JsonNode config = given.get("config"); + return config == null ? MAPPER.createObjectNode() : config; + } + + /** Maps a vector's handle back to the id the check that issued it returned. */ + String resolveReservationId(JsonNode op) { + if (op.has("handle")) { + Reservation reservation = handles.get(op.get("handle").asText()); + assertNotNull( + reservation, + "unknown reservation handle: " + op.get("handle").asText()); + return reservation.getId(); + } + assertTrue(op.has("id"), "op " + op.get("op").asText() + " needs an id or handle"); + return op.get("id").asText(); + } + + void close() { + handles.clear(); + if (manager != null) { + manager.close(); + } + } + } +} diff --git a/src/test/java/com/schematic/api/credits/conformance/CrashingRefundLeaseStore.java b/src/test/java/com/schematic/api/credits/conformance/CrashingRefundLeaseStore.java new file mode 100644 index 0000000..6ce2236 --- /dev/null +++ b/src/test/java/com/schematic/api/credits/conformance/CrashingRefundLeaseStore.java @@ -0,0 +1,67 @@ +package com.schematic.api.credits.conformance; + +import com.schematic.api.credits.LeaseGrant; +import com.schematic.api.credits.LeaseState; +import com.schematic.api.credits.LeaseStore; +import com.schematic.api.credits.ReserveResult; +import java.time.Instant; + +/** + * A lease store whose refund fails once while armed, reproducing a process death between a + * reservation's claim and its refund. The reservation store refunds through whatever store it is + * handed, so wrapping that one leaves the rest of the vector reading the real store. + */ +final class CrashingRefundLeaseStore implements LeaseStore { + + static final class SimulatedCrash extends RuntimeException { + SimulatedCrash() { + super("simulated crash before refund"); + } + } + + private final LeaseStore target; + private volatile boolean armed; + + CrashingRefundLeaseStore(LeaseStore target) { + this.target = target; + } + + void arm() { + armed = true; + } + + @Override + public LeaseState get(String companyId, String creditTypeId) { + return target.get(companyId, creditTypeId); + } + + @Override + public boolean replace(LeaseGrant grant) { + return target.replace(grant); + } + + @Override + public ReserveResult tryReserve(String companyId, String creditTypeId, double credits) { + return target.tryReserve(companyId, creditTypeId, credits); + } + + @Override + public void refund(String companyId, String creditTypeId, double credits, String pinLeaseId) { + if (armed) { + armed = false; + throw new SimulatedCrash(); + } + target.refund(companyId, creditTypeId, credits, pinLeaseId); + } + + @Override + public void extend( + String companyId, String creditTypeId, double grantedTotal, Instant newExpiresAt, String pinLeaseId) { + target.extend(companyId, creditTypeId, grantedTotal, newExpiresAt, pinLeaseId); + } + + @Override + public void drop(String companyId, String creditTypeId) { + target.drop(companyId, creditTypeId); + } +} diff --git a/src/test/java/com/schematic/api/credits/conformance/EmbeddedRedis.java b/src/test/java/com/schematic/api/credits/conformance/EmbeddedRedis.java new file mode 100644 index 0000000..2324bbf --- /dev/null +++ b/src/test/java/com/schematic/api/credits/conformance/EmbeddedRedis.java @@ -0,0 +1,57 @@ +package com.schematic.api.credits.conformance; + +import java.io.IOException; +import java.net.ServerSocket; +import redis.clients.jedis.JedisPooled; +import redis.embedded.RedisServer; + +/** + * One embedded Redis for the whole test JVM. + * + *

A real server rather than a fake: the stores' Lua is byte-identical to the other SDKs', and + * only a real Redis runs it, TIME and all. + */ +final class EmbeddedRedis { + + private static RedisServer server; + private static JedisPooled client; + + static synchronized JedisPooled client() { + if (client != null) { + return client; + } + try { + int port = freePort(); + server = RedisServer.newRedisServer().port(port).build(); + server.start(); + client = new JedisPooled("localhost", port); + Runtime.getRuntime().addShutdownHook(new Thread(EmbeddedRedis::stop)); + return client; + } catch (Exception e) { + throw new IllegalStateException("Failed to start the embedded Redis the conformance suite runs on", e); + } + } + + private static synchronized void stop() { + if (client != null) { + client.close(); + client = null; + } + if (server != null) { + try { + server.stop(); + } catch (Exception e) { + // Shutting down anyway. + } + server = null; + } + } + + private static int freePort() throws IOException { + try (ServerSocket socket = new ServerSocket(0)) { + return socket.getLocalPort(); + } + } + + private EmbeddedRedis() {} +} diff --git a/src/test/java/com/schematic/api/credits/conformance/MutableClock.java b/src/test/java/com/schematic/api/credits/conformance/MutableClock.java new file mode 100644 index 0000000..b12dada --- /dev/null +++ b/src/test/java/com/schematic/api/credits/conformance/MutableClock.java @@ -0,0 +1,35 @@ +package com.schematic.api.credits.conformance; + +import java.time.Clock; +import java.time.Instant; +import java.time.ZoneId; +import java.time.ZoneOffset; + +/** A clock that only moves when a test moves it, so nothing depends on wall time. */ +final class MutableClock extends Clock { + + private volatile Instant now; + + MutableClock(Instant start) { + this.now = start; + } + + @Override + public ZoneId getZone() { + return ZoneOffset.UTC; + } + + @Override + public Clock withZone(ZoneId zone) { + return this; + } + + @Override + public Instant instant() { + return now; + } + + void advance(long millis) { + now = now.plusMillis(millis); + } +} diff --git a/src/test/java/com/schematic/api/credits/conformance/RedisVectorClock.java b/src/test/java/com/schematic/api/credits/conformance/RedisVectorClock.java new file mode 100644 index 0000000..d0fd9bc --- /dev/null +++ b/src/test/java/com/schematic/api/credits/conformance/RedisVectorClock.java @@ -0,0 +1,75 @@ +package com.schematic.api.credits.conformance; + +import java.time.Clock; +import java.time.Instant; +import java.util.List; +import java.util.Set; +import redis.clients.jedis.JedisPooled; +import redis.clients.jedis.resps.Tuple; + +/** + * The vector timeline against a real Redis. + * + *

The lease scripts decide expiry against the Redis server's own clock, and a real server's + * clock cannot be moved, so this moves the data instead: advancing the timeline by N milliseconds + * brings every stored deadline N milliseconds nearer, which is the same comparison from the other + * side. The SDK-side clock stays the system clock, so both sides read one frame. Vector instants + * are translated into that frame by {@link #at}. + */ +final class RedisVectorClock implements VectorClock { + + private final JedisPooled jedis; + private final String keyPrefix; + private final Instant start = Instant.now(); + private long shiftedMillis; + + RedisVectorClock(JedisPooled jedis, String keyPrefix) { + this.jedis = jedis; + this.keyPrefix = keyPrefix; + } + + @Override + public Clock clock() { + return Clock.systemUTC(); + } + + @Override + public void advance(long millis) { + shiftedMillis += millis; + shiftHashDeadlines(keyPrefix + "credit-lease:*", millis); + shiftHashDeadlines(keyPrefix + "credit-reservation:*", millis); + shiftIndexScores(keyPrefix + "credit-reservations:byExpiry", millis); + } + + @Override + public Instant at(double offsetMillis) { + return start.plusMillis((long) offsetMillis - shiftedMillis); + } + + private void shiftHashDeadlines(String pattern, long millis) { + Set keys = jedis.keys(pattern); + for (String key : keys) { + String raw = jedis.hget(key, "expiresAt"); + if (raw != null) { + jedis.hset(key, "expiresAt", Long.toString(Long.parseLong(raw) - millis)); + } + long pttl = jedis.pttl(key); + if (pttl <= 0) { + continue; + } + // Past its grace window the row is one Redis would already have evicted. + if (pttl - millis <= 0) { + jedis.del(key); + } else { + jedis.pexpire(key, pttl - millis); + } + } + } + + private void shiftIndexScores(String key, long millis) { + List members = jedis.zrangeWithScores(key, 0, -1); + for (Tuple member : members) { + jedis.zadd(key, member.getScore() - millis, member.getElement()); + } + } +} diff --git a/src/test/java/com/schematic/api/credits/conformance/ScriptedCheckDataStream.java b/src/test/java/com/schematic/api/credits/conformance/ScriptedCheckDataStream.java new file mode 100644 index 0000000..01dd4db --- /dev/null +++ b/src/test/java/com/schematic/api/credits/conformance/ScriptedCheckDataStream.java @@ -0,0 +1,128 @@ +package com.schematic.api.credits.conformance; + +import com.schematic.api.credits.CreditCheckDataStream; +import com.schematic.api.credits.PreflightOptions; +import com.schematic.api.types.RulesengineCheckFlagResult; +import com.schematic.api.types.RulesengineCompany; +import com.schematic.api.types.RulesengineEntitlementValueType; +import com.schematic.api.types.RulesengineFeatureEntitlement; +import com.schematic.api.types.RulesengineFlag; +import com.schematic.api.types.RulesengineUser; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; +import java.util.List; +import java.util.Map; + +/** + * 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 it, not the engine, which is shared across the SDKs and has its own + * tests. + */ +final class ScriptedCheckDataStream implements CreditCheckDataStream { + + /** One evaluation the flow asked for, kept so a vector can assert on what it was given. */ + static final class EngineCall { + final Map creditBalances; + final PreflightOptions preflight; + + EngineCall(Map creditBalances, PreflightOptions preflight) { + this.creditBalances = creditBalances; + this.preflight = preflight; + } + } + + /** One scripted engine answer. */ + static final class Result { + final boolean value; + final String reason; + final RulesengineFeatureEntitlement entitlement; + + Result(boolean value, String reason, RulesengineFeatureEntitlement entitlement) { + this.value = value; + this.reason = reason; + this.entitlement = entitlement; + } + } + + private final String flagKey; + private final RulesengineFlag flag; + private final RulesengineCompany company; + private final Deque results = new ArrayDeque<>(); + final List calls = new ArrayList<>(); + + ScriptedCheckDataStream(String flagKey, RulesengineCompany company, List results) { + this.flagKey = flagKey; + this.flag = RulesengineFlag.builder() + .accountId("acct_1") + .defaultValue(false) + .environmentId("env_1") + .id("flag_1") + .key(flagKey) + .build(); + this.company = company; + this.results.addAll(results); + } + + static RulesengineCompany company(String id, Map creditBalances) { + return RulesengineCompany.builder() + .accountId("acct_1") + .environmentId("env_1") + .id(id) + .creditBalances(creditBalances) + .build(); + } + + static RulesengineFeatureEntitlement entitlement( + String valueType, String creditId, Double consumptionRate, String eventSubtype) { + RulesengineFeatureEntitlement._FinalStage builder = RulesengineFeatureEntitlement.builder() + .featureId("feat_1") + .featureKey("feature") + .valueType(RulesengineEntitlementValueType.valueOf(valueType)); + if (creditId != null) { + builder.creditId(creditId); + } + if (consumptionRate != null) { + builder.consumptionRate(consumptionRate); + } + if (eventSubtype != null) { + builder.eventSubtype(eventSubtype); + } + return builder.build(); + } + + @Override + public RulesengineFlag getFlag(String key) { + return flag; + } + + @Override + public RulesengineCompany getCompany(Map keys) { + return company; + } + + @Override + public RulesengineUser getUser(Map keys) { + return null; + } + + @Override + public RulesengineCheckFlagResult evaluateFlag( + RulesengineFlag flag, RulesengineCompany company, RulesengineUser user, PreflightOptions preflight) { + calls.add(new EngineCall(company == null ? null : company.getCreditBalances(), preflight)); + Result scripted = results.poll(); + if (scripted == null) { + throw new IllegalStateException("unscripted engine call in a check op for flag " + flagKey); + } + RulesengineCheckFlagResult._FinalStage result = RulesengineCheckFlagResult.builder() + .flagKey(flagKey) + .reason(scripted.reason) + .value(scripted.value) + .flagId("flag_1"); + if (scripted.entitlement != null) { + result.entitlement(scripted.entitlement); + } + return result.build(); + } +} diff --git a/src/test/java/com/schematic/api/credits/conformance/ScriptedWireClient.java b/src/test/java/com/schematic/api/credits/conformance/ScriptedWireClient.java new file mode 100644 index 0000000..2b5fa3b --- /dev/null +++ b/src/test/java/com/schematic/api/credits/conformance/ScriptedWireClient.java @@ -0,0 +1,119 @@ +package com.schematic.api.credits.conformance; + +import com.schematic.api.credits.LeaseGrant; +import com.schematic.api.credits.LeaseWireClient; +import java.time.Instant; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; +import java.util.List; + +/** Stands in for the lease API: queued responses in, recorded calls out. */ +final class ScriptedWireClient implements LeaseWireClient { + + /** One acquire or extend the vector scripts the server to answer with. */ + static final class Script { + + final String leaseId; + final double grantedAmount; + final Instant expiresAt; + final String error; + + private Script(String leaseId, double grantedAmount, Instant expiresAt, String error) { + this.leaseId = leaseId; + this.grantedAmount = grantedAmount; + this.expiresAt = expiresAt; + this.error = error; + } + + static Script lease(String leaseId, double grantedAmount, Instant expiresAt) { + return new Script(leaseId, grantedAmount, expiresAt, null); + } + + static Script error(String message) { + return new Script(null, 0, null, message); + } + } + + static final class AcquireCall { + + final String companyId; + final String creditTypeId; + final double requestedAmount; + final Instant expiresAt; + + AcquireCall(String companyId, String creditTypeId, double requestedAmount, Instant expiresAt) { + this.companyId = companyId; + this.creditTypeId = creditTypeId; + this.requestedAmount = requestedAmount; + this.expiresAt = expiresAt; + } + } + + static final class ExtendCall { + + final String leaseId; + final double additionalAmount; + final Instant expiresAt; + + ExtendCall(String leaseId, double additionalAmount, Instant expiresAt) { + this.leaseId = leaseId; + this.additionalAmount = additionalAmount; + this.expiresAt = expiresAt; + } + } + + private final Deque