diff --git a/src/credits/lease-store.ts b/src/credits/lease-store.ts index e970d71e..7285242b 100644 --- a/src/credits/lease-store.ts +++ b/src/credits/lease-store.ts @@ -32,11 +32,18 @@ export interface ILeaseStore { * sharing the backend: if a live (unexpired) lease is already present — * even one with a different `leaseId`, e.g. acquired by a sibling pod that * raced this one — leave it untouched so its already-debited - * `localRemainingCredits` wins. Returns `true` if it wrote a fresh row, - * `false` if it kept an existing live lease (the caller should release the - * lease it acquired ONLY if its ID differs from the installed one — the - * server hands a racing acquire the same active lease back, and releasing - * that would pull the shared lease out from under every sibling). + * `localRemainingCredits` wins. If the slot holds an *expired* row with + * the SAME `leaseId`, implementations must NOT rewrite it either: + * rewriting resets `localRemainingCredits` to the full grant, erasing + * debits whose reservations are still open (the idempotent server hands a + * racing acquire the same active lease back, and a stale response can + * land after the local row's expiry — especially when a concurrent extend + * pushed the server-side expiry forward). Instead, reconcile it like an + * extend: granted-to-total, expiry only forward, balance untouched. + * Returns `true` if it wrote a fresh row, `false` if it kept (or + * reconciled) an existing lease (the caller should release the lease it + * acquired ONLY if its ID differs from the installed one — releasing the + * shared lease would pull it out from under every sibling). */ replace(entry: Omit): Promise; /** @@ -115,9 +122,11 @@ export class LeaseStore implements ILeaseStore { * Replace (or insert) the lease for the given (company, creditType). * Resets `localRemainingCredits` to `grantedAmount` for a new lease. * If a live lease already occupies the slot — regardless of `leaseId` — it - * is left alone so any already-debited `localRemainingCredits` wins. - * Returns `true` if a fresh row was written, `false` if an existing live - * lease was kept. + * is left alone so any already-debited `localRemainingCredits` wins. If + * the slot holds an expired row with the SAME `leaseId`, the entry is + * reconciled (granted-to-total, expiry forward) instead of rewritten — + * see `ILeaseStore.replace`. Returns `true` if a fresh row was written, + * `false` if an existing lease was kept. */ async replace(entry: Omit): Promise { const key = leaseKey(entry.companyId, entry.creditTypeId); @@ -128,6 +137,24 @@ export class LeaseStore implements ILeaseStore { // already-debited localRemaining instead of clobbering it. return false; } + if (existing && existing.leaseId === entry.leaseId) { + // The SAME lease coming back over its own expired row: a stale + // acquire response for a lease the idempotent server handed to + // a racing sibling too (possibly since extended, so still + // active server-side). Rewriting would reset + // localRemainingCredits to the full grant, erasing debits + // whose reservations are still open — reconcile like an + // extend instead (see ILeaseStore.replace). + const add = entry.grantedAmount - existing.grantedAmount; + if (add > 0) { + existing.grantedAmount = entry.grantedAmount; + existing.localRemainingCredits += add; + } + if (entry.expiresAt.getTime() > existing.expiresAt.getTime()) { + existing.expiresAt = entry.expiresAt; + } + return false; + } this.leases.set(key, { ...entry, localRemainingCredits: entry.grantedAmount, diff --git a/src/credits/redis-lease-store.ts b/src/credits/redis-lease-store.ts index 98f19e97..3c3e345f 100644 --- a/src/credits/redis-lease-store.ts +++ b/src/credits/redis-lease-store.ts @@ -34,8 +34,13 @@ local now = (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000) * installed by a sibling instance that raced this acquire. Keeping the existing * live lease preserves its already-debited `localRemainingCredits`; the caller * should release the lease it acquired only if its ID differs from the - * installed one (see `CreditLeaseManager.acquire`). `now` is the Redis server - * clock (see `LEASE_NOW_MS`). + * installed one (see `CreditLeaseManager.acquire`). An expired row with the + * SAME `leaseId` (a stale acquire response for a lease the idempotent server + * handed to a racing sibling, possibly since extended) is reconciled like an + * extend — granted-to-total, expiry only forward, balance untouched — instead + * of rewritten, which would reset the balance and erase debits whose + * reservations are still open. `now` is the Redis server clock (see + * `LEASE_NOW_MS`). */ const REPLACE_SCRIPT = LEASE_NOW_MS + @@ -51,6 +56,22 @@ if existing_id and existing_expiry > now then return 0 end +if existing_id == new_id then + local granted = tonumber(redis.call('HGET', KEYS[1], 'grantedAmount') or '0') + local add = tonumber(new_granted) - granted + if add > 0 then + local remaining = tonumber(redis.call('HGET', KEYS[1], 'localRemainingCredits') or '0') + redis.call('HSET', KEYS[1], + 'grantedAmount', new_granted, + 'localRemainingCredits', tostring(remaining + add)) + end + if new_expiry > existing_expiry then + redis.call('HSET', KEYS[1], 'expiresAt', ARGV[3]) + redis.call('PEXPIREAT', KEYS[1], new_expiry + grace) + end + return 0 +end + redis.call('DEL', KEYS[1]) redis.call('HSET', KEYS[1], 'leaseId', new_id, @@ -189,7 +210,7 @@ export class RedisLeaseStore implements ILeaseStore { async get(companyId: string, creditTypeId: string): Promise { const raw = await this.client.hGetAll(this.hashKey(companyId, creditTypeId)); - if (!raw || !raw.leaseId) return undefined; + if (!raw?.leaseId) return undefined; return decodeEntry(raw); } diff --git a/tests/integration/redis-stores.test.ts b/tests/integration/redis-stores.test.ts index a7cde93c..b489ca70 100644 --- a/tests/integration/redis-stores.test.ts +++ b/tests/integration/redis-stores.test.ts @@ -125,6 +125,35 @@ describeIf("Redis stores against a real redis-server", () => { expect(entry?.localRemainingCredits).toBe(200); }); + it("replace reconciles an expired SAME-id lease instead of resetting its debited balance", async () => { + // A stale acquire response can hand back the lease already installed + // (the server is idempotent for an active slot) after the shared row + // expired — e.g. a concurrent extend pushed the server-side expiry + // forward. Rewriting would reset localRemainingCredits to the full + // grant, erasing debits whose reservations are still open. + const leaseStore = makeLeaseStore(); + const creditTypeId = await installLease(leaseStore, { leaseId: "lse_a", grantedAmount: 100 }); + expect(await leaseStore.tryReserve("co_1", creditTypeId, 40)).toBe(60); + + await client.hSet(leaseStore.hashKey("co_1", creditTypeId), "expiresAt", String(Date.now() - 1_000)); + const laterExpiry = new Date(Date.now() + 120_000); + const wrote = await leaseStore.replace({ + leaseId: "lse_a", + companyId: "co_1", + creditTypeId, + grantedAmount: 150, + expiresAt: laterExpiry, + }); + expect(wrote).toBe(false); + const entry = await leaseStore.get("co_1", creditTypeId); + expect(entry?.leaseId).toBe("lse_a"); + // Granted reconciled to the server total; the 40-credit debit survives. + expect(entry?.grantedAmount).toBe(150); + expect(entry?.localRemainingCredits).toBe(110); + // Expiry moves forward with the response. + expect(entry?.expiresAt.getTime()).toBe(laterExpiry.getTime()); + }); + it("tryReserve under contention admits exactly the available balance", async () => { const leaseStore = makeLeaseStore(); const creditTypeId = await installLease(leaseStore, { grantedAmount: 100 }); diff --git a/tests/races/harness.ts b/tests/races/harness.ts new file mode 100644 index 00000000..14b3d14b --- /dev/null +++ b/tests/races/harness.ts @@ -0,0 +1,676 @@ +/** + * Seeded randomized race harness for the lease/reservation stores. + * + * The conformance suite pins sequential semantics with deterministic vectors; + * the truly concurrent interleavings (single-flight, extend-vs-expiry, + * release-vs-sweeper, replace-vs-reserve) are prose invariants there. This + * harness exercises exactly those: per seed it spawns N async workers sharing + * a small set of (company, creditType) slots, each running a randomized op + * sequence against the store API with seeded micro-yields injected between + * store calls so the event-loop interleaving varies by seed. At quiescence it + * asserts the prose invariants (conservation, bounded leak, expiry + * monotonicity, pinned refunds/extends, exactly-once claims). + * + * Determinism: everything (op choice, amounts, yields, virtual-clock + * advances) draws from PRNGs derived from the seed, ids are counters, and the + * only awaited primitives are microtasks and `setImmediate` — whose relative + * ordering Node schedules deterministically — so a seed replays the exact + * schedule. `RACE_SEED=` re-runs one seed; `RACE_ITERATIONS=` scales the + * batch. + */ +import type { ILeaseStore } from "../../src/credits/lease-store"; +import { LeaseStore } from "../../src/credits/lease-store"; +import { RedisLeaseStore } from "../../src/credits/redis-lease-store"; +import { RedisReservationStore } from "../../src/credits/redis-reservation-store"; +import type { IReservationStore } from "../../src/credits/reservation-store"; +import { ReservationStore } from "../../src/credits/reservation-store"; +import type { Reservation } from "../../src/credits/types"; +import { makeFakeRedis } from "../unit/credits/fake-redis"; +import { deriveSeed, Rng } from "./prng"; + +const EPS = 1e-9; + +export interface Slot { + companyId: string; + creditTypeId: string; +} + +export function slotKeyOf(slot: Slot): string { + return `${slot.companyId}:${slot.creditTypeId}`; +} + +// --------------------------------------------------------------------------- +// Virtual clock +// +// Both stores (and fake-redis's TTL eviction) decide expiry via `Date.now()`, +// so overriding it gives the harness full, deterministic control of the +// timeline — clock advances are ordinary synchronous ops workers race against +// reserves/extends/sweeps. (The real Redis store reads the *server* clock via +// TIME; fake-redis emulates that with `Date.now()`, which this override also +// controls. The opt-in real-Redis integration suite is where server-clock +// authority itself is exercised.) +// --------------------------------------------------------------------------- + +export interface VirtualClock { + now(): number; + advance(ms: number): void; + uninstall(): void; +} + +const CLOCK_EPOCH_MS = 1_700_000_000_000; + +export function installVirtualClock(): VirtualClock { + const realNow = Date.now; + let current = CLOCK_EPOCH_MS; + Date.now = () => current; + return { + now: () => current, + advance: (ms: number) => { + current += ms; + }, + uninstall: () => { + Date.now = realNow; + }, + }; +} + +// --------------------------------------------------------------------------- +// Interleaving control +// --------------------------------------------------------------------------- + +/** + * Seeded micro-yield: 0-3 awaits, each either a microtask or a macrotask + * (`setImmediate`). Workers call this between store operations so the + * interleaving of their op sequences varies by seed. `setTimeout(0)` is + * deliberately excluded — its ordering relative to `setImmediate` is not + * deterministic in Node, which would break seed replay. + */ +export async function jitter(rng: Rng): Promise { + const yields = rng.int(4); + for (let i = 0; i < yields; i++) { + if (rng.chance(0.5)) { + await Promise.resolve(); + } else { + await new Promise((resolve) => setImmediate(resolve)); + } + } +} + +/** Drain pending microtasks/immediates (e.g. fire-and-forget releases) before verifying. */ +export async function settle(rounds = 25): Promise { + for (let i = 0; i < rounds; i++) { + await new Promise((resolve) => setImmediate(resolve)); + } +} + +/** + * Credit amounts are multiples of 0.25: fractional (exercising the Redis + * string-balance encoding) but exactly representable in binary, so the + * conservation sums below can be compared exactly. + */ +export function quarters(rng: Rng, maxQuarters: number, minQuarters = 1): number { + return rng.intBetween(minQuarters, maxQuarters) * 0.25; +} + +// --------------------------------------------------------------------------- +// Backends +// --------------------------------------------------------------------------- + +export interface StorePair { + leases: ILeaseStore; + reservations: IReservationStore; +} + +export interface Backend { + name: string; + make(): StorePair; +} + +// The sweep interval is irrelevant here (workers call `sweepExpired` +// explicitly, racing it against consumes); set it long so a stray +// `startSweep` could never fire mid-run. +export const backends: Backend[] = [ + { + name: "in-memory", + make: () => { + const leases = new LeaseStore(); + return { leases, reservations: new ReservationStore(leases, 60_000) }; + }, + }, + { + name: "redis (fake)", + make: () => { + const client = makeFakeRedis(); + const leases = new RedisLeaseStore({ client }); + return { + leases, + reservations: new RedisReservationStore({ client, leaseStore: leases, sweepIntervalMs: 60_000 }), + }; + }, + }, +]; + +// --------------------------------------------------------------------------- +// Model: what the harness knows it asked the stores to do +// --------------------------------------------------------------------------- + +export interface OpenReservation { + slotKey: string; + leaseId: string; + credits: number; +} + +export class RaceModel { + /** Leases whose `replace` returned written=true, keyed by lease id (ids are never reused). */ + installs = new Map(); + /** Every store `extend` issued, with its pin. Pins are only ever ids observed live in the slot. */ + extendsIssued: Array<{ pinLeaseId: string; total: number; expiresAtMs: number }> = []; + /** Reservations recorded via `add` and not yet claimed by a non-null `consume` this run. */ + open = new Map(); + /** Per-slot credits actually consumed (sum of non-null `consume` returns). */ + consumed = new Map(); + /** Per-slot credits deliberately leaked via simulated window-1 crashes (debit without record). */ + leaked = new Map(); + /** Per-reservation count of non-null `consume` returns (must be <= 1). */ + claims = new Map(); + /** Op-time violations (non-racy assertions checked inline). */ + violations: string[] = []; + + addTo(map: Map, key: string, credits: number): void { + map.set(key, (map.get(key) ?? 0) + credits); + } +} + +export interface RaceCtx { + seed: number; + clock: VirtualClock; + leases: ILeaseStore; + reservations: IReservationStore; + slots: Slot[]; + model: RaceModel; + nextId(prefix: string): string; +} + +export function makeReservation( + ctx: RaceCtx, + slot: Slot, + leaseId: string, + credits: number, + ttlMs: number, +): Reservation { + return { + id: ctx.nextId("res"), + leaseId, + companyId: slot.companyId, + creditTypeId: slot.creditTypeId, + eventSubtype: "inference_tokens", + quantityReserved: credits, + creditsReserved: credits, + consumptionRate: 1, + expiresAt: new Date(ctx.clock.now() + ttlMs), + evalCtx: { company: { id: slot.companyId } }, + }; +} + +// --------------------------------------------------------------------------- +// Shared ops +// --------------------------------------------------------------------------- + +export interface WeightedOp { + weight: number; + run(rng: Rng): Promise; +} + +export interface ReserveOpts { + /** Probability of simulating a window-1 crash (debit lands, record never written). */ + crashChance?: number; + /** Reservation TTL range [min, max] ms. */ + ttlMs: [number, number]; + /** Max reserve size in quarter-credits. */ + maxQuarters: number; +} + +/** + * The check-flow shape: observe the lease (skipping expired ones, like the + * manager does), debit via `try_reserve`, then record the reservation pinned + * to the OBSERVED lease id — deliberately leaving seeded yield points where + * expiry, a successor `replace`, or a sweep can interleave. + */ +export function reserveOp(ctx: RaceCtx, opts: ReserveOpts): WeightedOp { + return { + weight: 4, + run: async (rng) => { + const slot = rng.pick(ctx.slots); + const entry = await ctx.leases.get(slot.companyId, slot.creditTypeId); + if (!entry || entry.expiresAt.getTime() <= ctx.clock.now()) return; + const credits = quarters(rng, opts.maxQuarters); + await jitter(rng); + const post = await ctx.leases.tryReserve(slot.companyId, slot.creditTypeId, credits); + if (post === null) return; + if (post < -EPS) { + ctx.model.violations.push( + `tryReserve(${credits}) on ${slotKeyOf(slot)} returned a negative balance ${post}`, + ); + } + await jitter(rng); + if (opts.crashChance && rng.chance(opts.crashChance)) { + // Simulated window-1 crash: the debit landed but the record is + // never written. The sweeper can never refund it; conservation + // accounts for it via the leak budget. + ctx.model.addTo(ctx.model.leaked, slotKeyOf(slot), credits); + return; + } + const reservation = makeReservation(ctx, slot, entry.leaseId, credits, rng.intBetween(...opts.ttlMs)); + await ctx.reservations.add(reservation); + ctx.model.open.set(reservation.id, { slotKey: slotKeyOf(slot), leaseId: entry.leaseId, credits }); + }, + }; +} + +/** Settle a random open reservation — racing other consumers and the sweeper. */ +export function consumeOp(ctx: RaceCtx): WeightedOp { + return { + weight: 3, + run: async (rng) => { + const ids = Array.from(ctx.model.open.keys()); + if (ids.length === 0) return; + const id = rng.pick(ids); + const open = ctx.model.open.get(id); + if (!open) return; + const roll = rng.next(); + let requested: number; + if (roll < 0.15) { + requested = open.credits + quarters(rng, 20); // over-consume: clamps to reserved + } else if (roll < 0.25) { + requested = -quarters(rng, 8); // negative: clamps to 0, full refund + } else { + requested = quarters(rng, Math.max(1, Math.round(open.credits * 4)), 0); + } + await jitter(rng); + const consumed = await ctx.reservations.consume(id, requested); + if (consumed === null) return; // lost the claim race (sweeper or another consumer) + if (consumed < -EPS || consumed > open.credits + EPS) { + ctx.model.violations.push( + `consume(${id}, ${requested}) returned ${consumed}, outside [0, ${open.credits}]`, + ); + } + ctx.model.claims.set(id, (ctx.model.claims.get(id) ?? 0) + 1); + ctx.model.addTo(ctx.model.consumed, open.slotKey, consumed); + ctx.model.open.delete(id); + }, + }; +} + +/** Run the expired-reservation sweeper — racing settles on the same reservations. */ +export function sweepOp(ctx: RaceCtx): WeightedOp { + return { + weight: 2, + run: async (rng) => { + await jitter(rng); + await ctx.reservations.sweepExpired(new Date(ctx.clock.now())); + }, + }; +} + +/** + * Advance the virtual clock. Bounded to 300ms per op so a run's total advance + * stays well under the stores' TTL grace windows (30s reservation / 60s + * lease): rows must never TTL-evict un-swept mid-run, which would look like a + * lost refund to the conservation check. + */ +export function advanceOp(ctx: RaceCtx, weight = 2): WeightedOp { + return { + weight, + run: async (rng) => { + ctx.clock.advance(rng.intBetween(50, 300)); + }, + }; +} + +/** Non-finite/negative debits must be rejected without touching the balance. */ +export function invalidReserveOp(ctx: RaceCtx): WeightedOp { + return { + weight: 1, + run: async (rng) => { + const slot = rng.pick(ctx.slots); + const bad = rng.pick([Number.NaN, -3, Number.POSITIVE_INFINITY, Number.NEGATIVE_INFINITY]); + const result = await ctx.leases.tryReserve(slot.companyId, slot.creditTypeId, bad); + if (result !== null) { + ctx.model.violations.push(`tryReserve(${bad}) on ${slotKeyOf(slot)} returned ${result}, not null`); + } + }, + }; +} + +/** + * The manager's extend shape: observe the lease (skipping expired ones), + * then reconcile the store to a server-authoritative total pinned to the + * observed lease id — racing expiry and successor installs. + */ +export function extendObservedOp(ctx: RaceCtx, durationMs: [number, number]): WeightedOp { + return { + weight: 3, + run: async (rng) => { + const slot = rng.pick(ctx.slots); + const entry = await ctx.leases.get(slot.companyId, slot.creditTypeId); + if (!entry || entry.expiresAt.getTime() <= ctx.clock.now()) return; + // Mix of higher totals (real extends) and stale/lower totals + // (out-of-order applies that must converge as no-ops). + const total = entry.grantedAmount + rng.pick([-50, -0.25, 25, 100, 250.5]); + const expiresAtMs = ctx.clock.now() + rng.intBetween(...durationMs); + await jitter(rng); + await ctx.leases.extend(slot.companyId, slot.creditTypeId, total, new Date(expiresAtMs), entry.leaseId); + ctx.model.extendsIssued.push({ pinLeaseId: entry.leaseId, total, expiresAtMs }); + }, + }; +} + +export interface ReplaceOpts { + /** Lease duration range, or "immortal" for a lease that can never expire mid-run. */ + durationMs: [number, number] | "immortal"; + grantedQuarters: [number, number]; + weight?: number; + /** Only attempt the replace when the observed slot is empty/expired (the manager's shape). */ + onlyIfDead?: boolean; +} + +/** Install a fresh lease — racing sibling installs (keep-first) and in-flight reserves. */ +export function replaceFreshOp(ctx: RaceCtx, opts: ReplaceOpts): WeightedOp { + return { + weight: opts.weight ?? 2, + run: async (rng) => { + const slot = rng.pick(ctx.slots); + if (opts.onlyIfDead) { + const current = await ctx.leases.get(slot.companyId, slot.creditTypeId); + if (current && current.expiresAt.getTime() > ctx.clock.now()) return; + await jitter(rng); + } + const granted = quarters(rng, opts.grantedQuarters[1], opts.grantedQuarters[0]); + const expiresAtMs = + opts.durationMs === "immortal" + ? ctx.clock.now() + 10 * 365 * 24 * 60 * 60 * 1000 + : ctx.clock.now() + rng.intBetween(...opts.durationMs); + const leaseId = ctx.nextId("lse"); + const wrote = await ctx.leases.replace({ + leaseId, + companyId: slot.companyId, + creditTypeId: slot.creditTypeId, + grantedAmount: granted, + expiresAt: new Date(expiresAtMs), + }); + if (wrote) { + ctx.model.installs.set(leaseId, { slotKey: slotKeyOf(slot), granted, expiresAtMs }); + } + }, + }; +} + +/** + * Prove the expiry guard: observe the lease, synchronously advance the clock + * just past its expiry, and immediately probe `try_reserve`. The advance and + * the call have no yield between them, so any debit happens at + * `clock >= advancedTo`. If the probe succeeds and the slot still holds the + * SAME lease with an expiry `<= advancedTo`, the debit provably landed on an + * expired lease (ids are never reused and expiry only moves forward, so a + * live-at-debit lease would read a later expiry afterwards) — a + * server-side double-spend the local ledger alone cannot see. A success where + * the slot moved on to a live successor is legitimate (the debit hit the + * successor); it is accounted as an unrecorded hold, like a window-1 crash. + */ +export function expiredProbeOp(ctx: RaceCtx, weight = 2): WeightedOp { + return { + weight, + run: async (rng) => { + const slot = rng.pick(ctx.slots); + const before = await ctx.leases.get(slot.companyId, slot.creditTypeId); + if (!before) return; + const advancedTo = Math.max(ctx.clock.now(), before.expiresAt.getTime() + 1); + ctx.clock.advance(advancedTo - ctx.clock.now()); + const credits = 0.25; + const post = await ctx.leases.tryReserve(slot.companyId, slot.creditTypeId, credits); + if (post === null) return; + const after = await ctx.leases.get(slot.companyId, slot.creditTypeId); + if (after && after.leaseId === before.leaseId && after.expiresAt.getTime() <= advancedTo) { + ctx.model.violations.push( + `${slotKeyOf(slot)}: tryReserve succeeded against expired lease ${before.leaseId} ` + + `(expiry ${after.expiresAt.getTime()} <= clock ${advancedTo}; ` + + `the server already refunded this balance — serving it double-spends)`, + ); + return; + } + // Legitimate: the debit landed on a live successor. Nothing records + // this hold, so account it like a window-1 crash leak. + ctx.model.addTo(ctx.model.leaked, slotKeyOf(slot), credits); + }, + }; +} + +/** Drop the slot (remote release) — racing everything else. */ +export function dropOp(ctx: RaceCtx, weight = 1): WeightedOp { + return { + weight, + run: async (rng) => { + const slot = rng.pick(ctx.slots); + await jitter(rng); + await ctx.leases.drop(slot.companyId, slot.creditTypeId); + }, + }; +} + +// --------------------------------------------------------------------------- +// Invariant verification (at quiescence) +// --------------------------------------------------------------------------- + +export interface VerifyOpts { + /** + * Exact conservation: granted == remaining + open holds + consumed + + * crash-leaked, per slot. Only sound when the scenario guarantees a single + * never-expiring lease per slot (attribution of debits is then exact); it + * catches both double-refunds and lost refunds. + */ + conservation?: boolean; + /** + * Assert the exact final granted/expiry of the surviving lease from the + * model (install + pinned extends). Sound whenever the harness only pins + * extends to observed-live lease ids (ids are never reused, so a lease + * still installed at quiescence was installed continuously since its + * `replace` — every extend pinned to it applied). Off for manager-driven + * scenarios where installs happen inside the manager. + */ + exactLeaseState?: boolean; + /** Require every slot to still hold a known lease (single-immortal-lease scenarios). */ + requireLease?: boolean; +} + +export async function verifyInvariants(ctx: RaceCtx, opts: VerifyOpts): Promise { + const violations: string[] = [...ctx.model.violations]; + + // Reconcile the model's open reservations against the store: ones the + // sweeper claimed are gone (their hold refunded to the lease); ones still + // present are open holds. + const presentBySlot = new Map(); + for (const [id, open] of ctx.model.open) { + const inStore = await ctx.reservations.get(id); + if (!inStore) continue; // swept: full hold refunded (pinned), counts as consumed 0 + const list = presentBySlot.get(open.slotKey) ?? []; + list.push(open); + presentBySlot.set(open.slotKey, list); + } + + for (const slot of ctx.slots) { + const key = slotKeyOf(slot); + const present = presentBySlot.get(key) ?? []; + const presentSum = present.reduce((sum, r) => sum + r.credits, 0); + + // The reserved-credits index (a separate structure in the Redis store) + // must agree exactly with the reservations actually present. + const reservedTotal = await ctx.reservations.reservedCredits(slot.companyId, slot.creditTypeId); + if (Math.abs(reservedTotal - presentSum) > EPS) { + violations.push( + `${key}: reservedCredits reports ${reservedTotal} but open reservations sum to ${presentSum}`, + ); + } + + const entry = await ctx.leases.get(slot.companyId, slot.creditTypeId); + if (!entry) { + if (opts.requireLease) violations.push(`${key}: expected the immortal lease to survive, slot is empty`); + continue; + } + + const remaining = entry.localRemainingCredits; + const granted = entry.grantedAmount; + if (remaining < -EPS) violations.push(`${key}: negative balance ${remaining} on lease ${entry.leaseId}`); + if (remaining > granted + EPS) { + violations.push(`${key}: balance ${remaining} exceeds granted ${granted} on lease ${entry.leaseId}`); + } + + // No double-spend: credits held by open reservations pinned to the + // live lease were debited from it, so remaining + holds can never + // exceed granted. A double-refund or a refund crossing lease + // generations (a broken pin) pushes this over. + const pinnedOpen = present.filter((r) => r.leaseId === entry.leaseId).reduce((sum, r) => sum + r.credits, 0); + if (remaining + pinnedOpen > granted + EPS) { + violations.push( + `${key}: remaining ${remaining} + pinned open holds ${pinnedOpen} exceeds granted ${granted} ` + + `on lease ${entry.leaseId} (credits were minted)`, + ); + } + + const install = ctx.model.installs.get(entry.leaseId); + if (opts.exactLeaseState) { + if (!install) { + violations.push(`${key}: slot holds lease ${entry.leaseId} the harness never installed`); + } else { + const pinnedExtends = ctx.model.extendsIssued.filter((e) => e.pinLeaseId === entry.leaseId); + const expectedGranted = Math.max(install.granted, ...pinnedExtends.map((e) => e.total)); + const expectedExpiry = Math.max(install.expiresAtMs, ...pinnedExtends.map((e) => e.expiresAtMs)); + if (Math.abs(granted - expectedGranted) > EPS) { + violations.push( + `${key}: lease ${entry.leaseId} granted ${granted}, expected ${expectedGranted} ` + + `(install ${install.granted} + ${pinnedExtends.length} pinned extends; ` + + `an unpinned/stale extend leaked through)`, + ); + } + if (entry.expiresAt.getTime() !== expectedExpiry) { + violations.push( + `${key}: lease ${entry.leaseId} expiry ${entry.expiresAt.getTime()}, expected ${expectedExpiry} ` + + `(expiry must equal the max of install + pinned extend expiries — it only moves forward)`, + ); + } + } + } + + if (opts.conservation) { + const consumed = ctx.model.consumed.get(key) ?? 0; + const leaked = ctx.model.leaked.get(key) ?? 0; + const accounted = remaining + presentSum + consumed + leaked; + if (Math.abs(accounted - granted) > EPS) { + violations.push( + `${key}: conservation broken on lease ${entry.leaseId}: remaining ${remaining} + ` + + `open holds ${presentSum} + consumed ${consumed} + crash-leaked ${leaked} ` + + `= ${accounted}, expected granted ${granted} ` + + `(${accounted > granted ? "double refund" : "lost refund"})`, + ); + } + } + } + + for (const [id, count] of ctx.model.claims) { + if (count > 1) violations.push(`reservation ${id} was claimed ${count} times (consume must be exactly-once)`); + } + + return violations; +} + +// --------------------------------------------------------------------------- +// Runner +// --------------------------------------------------------------------------- + +export interface ScenarioRun { + ops: WeightedOp[]; + verify(): Promise; +} + +export interface RaceScenario { + name: string; + slots: Slot[]; + build(ctx: RaceCtx): Promise; +} + +function pickWeighted(ops: WeightedOp[], rng: Rng): WeightedOp { + const total = ops.reduce((sum, op) => sum + op.weight, 0); + let roll = rng.next() * total; + for (const op of ops) { + roll -= op.weight; + if (roll < 0) return op; + } + return ops[ops.length - 1]; +} + +export async function runScenario(scenario: RaceScenario, backend: Backend, seed: number): Promise { + const rng = new Rng(deriveSeed(seed, 0)); + const clock = installVirtualClock(); + try { + const { leases, reservations } = backend.make(); + let idCounter = 0; + const ctx: RaceCtx = { + seed, + clock, + leases, + reservations, + slots: scenario.slots, + model: new RaceModel(), + nextId: (prefix: string) => `${prefix}_${seed}_${++idCounter}`, + }; + const { ops, verify } = await scenario.build(ctx); + + const workerCount = rng.intBetween(3, 6); + const opsPerWorker = rng.intBetween(6, 12); + await Promise.all( + Array.from({ length: workerCount }, (_, workerIndex) => { + const workerRng = new Rng(deriveSeed(seed, workerIndex + 1)); + return (async () => { + for (let i = 0; i < opsPerWorker; i++) { + await jitter(workerRng); + await pickWeighted(ops, workerRng).run(workerRng); + } + })(); + }), + ); + + // Quiescence: all workers idle; drain anything fire-and-forget, then + // run one final sweep. The sweep reconciles reservation rows the + // backend TTL-evicted after large clock jumps (their per-slot index + // fields would otherwise read as reserved-credits drift) — the same + // reconciliation the production sweep loop performs. + await settle(); + await reservations.sweepExpired(new Date(clock.now())); + + const violations = await verify(); + if (violations.length > 0) { + throw new Error( + `Race invariant violated (scenario="${scenario.name}", backend="${backend.name}", seed=${seed}, ` + + `workers=${workerCount}, opsPerWorker=${opsPerWorker})\n` + + `Replay exactly this schedule with: RACE_SEED=${seed} yarn test tests/races\n` + + violations.map((v) => ` - ${v}`).join("\n"), + ); + } + } finally { + clock.uninstall(); + } +} + +/** Seeds for this run: `RACE_SEED` replays one; `RACE_ITERATIONS` scales the batch. */ +export function seedBatch(defaultIterations: number): number[] { + const single = process.env.RACE_SEED; + if (single !== undefined && single !== "") { + const seed = Number(single); + if (!Number.isInteger(seed)) throw new Error(`RACE_SEED must be an integer, got "${single}"`); + return [seed]; + } + const iterations = Number(process.env.RACE_ITERATIONS ?? defaultIterations); + if (!Number.isInteger(iterations) || iterations < 1) { + throw new Error(`RACE_ITERATIONS must be a positive integer, got "${process.env.RACE_ITERATIONS}"`); + } + return Array.from({ length: iterations }, (_, i) => i + 1); +} diff --git a/tests/races/prng.ts b/tests/races/prng.ts new file mode 100644 index 00000000..666904e2 --- /dev/null +++ b/tests/races/prng.ts @@ -0,0 +1,50 @@ +/** + * Seeded PRNG (mulberry32) for the race harness. Every random choice in a run + * — op selection, amounts, TTLs, and the micro-yields that vary event-loop + * interleaving — draws from an `Rng` derived from the run's seed, so a failing + * schedule replays exactly under `RACE_SEED=`. + */ +export class Rng { + private state: number; + + constructor(seed: number) { + this.state = seed >>> 0; + // mulberry32 degenerates on a zero state; nudge to a fixed constant so + // seed 0 is still a valid, distinct stream. + if (this.state === 0) this.state = 0x9e3779b9; + } + + /** Uniform float in [0, 1). */ + next(): number { + this.state = (this.state + 0x6d2b79f5) >>> 0; + let t = this.state; + t = Math.imul(t ^ (t >>> 15), t | 1); + t ^= t + Math.imul(t ^ (t >>> 7), t | 61); + return ((t ^ (t >>> 14)) >>> 0) / 4294967296; + } + + /** Uniform integer in [0, maxExclusive). */ + int(maxExclusive: number): number { + return Math.floor(this.next() * maxExclusive); + } + + /** Uniform integer in [min, max] (inclusive). */ + intBetween(min: number, max: number): number { + return min + this.int(max - min + 1); + } + + pick(items: readonly T[]): T { + return items[this.int(items.length)]; + } + + chance(probability: number): boolean { + return this.next() < probability; + } +} + +/** Deterministically derive a child seed (e.g. per worker) from a run seed. */ +export function deriveSeed(seed: number, salt: number): number { + let h = (Math.imul(seed, 0x9e3779b1) ^ Math.imul(salt + 1, 0x85ebca6b)) >>> 0; + h = Math.imul(h ^ (h >>> 13), 0xc2b2ae35) >>> 0; + return h; +} diff --git a/tests/races/races.test.ts b/tests/races/races.test.ts new file mode 100644 index 00000000..5a65d7bd --- /dev/null +++ b/tests/races/races.test.ts @@ -0,0 +1,311 @@ +/** + * Seeded concurrency race tests for the lease/reservation stores (and the + * lease manager on top of them). Each scenario targets one of the race pairs + * the conformance spec can only state as prose invariants: + * + * - concurrent acquires (replace) on one slot: keep-first, single live lease + * - extend racing lease expiry and successor installs: pinned extends never + * credit a successor; expiry only moves forward + * - settle (consume) racing the sweeper: exactly-once claims, no double or + * lost refunds + * - successor replace racing try_reserve: expired balances are never served + * - manager acquire/extend single-flight over a shared store: the installed + * lease is never released out from under siblings; extends never mint + * phantom credits + * - mixed chaos across slots: all of the above at once, plus drops and + * simulated window-1 crashes (bounded leak) + * + * Every run is deterministic per seed. A failure prints the seed; replay it + * with `RACE_SEED= yarn test tests/races`. `RACE_ITERATIONS=` scales + * the batch (default 40 per scenario/backend). + */ +import { CreditLeaseManager } from "../../src/credits/lease-manager"; +import type { Logger } from "../../src/logger"; + +import type { RaceCtx, RaceScenario, Slot, WeightedOp } from "./harness"; +import { + advanceOp, + backends, + consumeOp, + dropOp, + expiredProbeOp, + extendObservedOp, + invalidReserveOp, + quarters, + replaceFreshOp, + reserveOp, + runScenario, + seedBatch, + slotKeyOf, + sweepOp, + verifyInvariants, +} from "./harness"; +import { deriveSeed, Rng } from "./prng"; +import { ScriptedCreditsServer } from "./scripted-server"; + +const seeds = seedBatch(40); + +jest.setTimeout(600_000); + +const oneSlot: Slot[] = [{ companyId: "co_1", creditTypeId: "ct_1" }]; +const twoSlots: Slot[] = [ + { companyId: "co_1", creditTypeId: "ct_1" }, + { companyId: "co_1", creditTypeId: "ct_2" }, +]; + +const silentLogger: Logger = { error() {}, warn() {}, info() {}, debug() {} }; + +async function installLease(ctx: RaceCtx, slot: Slot, granted: number, expiresAtMs: number): Promise { + const leaseId = ctx.nextId("lse"); + const wrote = await ctx.leases.replace({ + leaseId, + companyId: slot.companyId, + creditTypeId: slot.creditTypeId, + grantedAmount: granted, + expiresAt: new Date(expiresAtMs), + }); + if (!wrote) throw new Error(`setup: failed to install lease for ${slotKeyOf(slot)}`); + ctx.model.installs.set(leaseId, { slotKey: slotKeyOf(slot), granted, expiresAtMs }); + return leaseId; +} + +const scenarios: RaceScenario[] = [ + { + // Documented gap: single-flight / concurrent acquires. All installs + // use a lease that can never expire, so exactly one concurrent + // `replace` may win the slot; the losers must keep the winner's + // already-debited balance intact (conservation stays exact). + name: "concurrent acquires (replace) on one slot", + slots: oneSlot, + async build(ctx) { + return { + ops: [ + replaceFreshOp(ctx, { durationMs: "immortal", grantedQuarters: [2000, 4000], weight: 3 }), + reserveOp(ctx, { ttlMs: [100, 800], maxQuarters: 100 }), + consumeOp(ctx), + sweepOp(ctx), + advanceOp(ctx), + invalidReserveOp(ctx), + ], + verify: async () => { + const violations = await verifyInvariants(ctx, { conservation: true, exactLeaseState: true }); + if (ctx.model.installs.size > 1) { + violations.push( + `${ctx.model.installs.size} concurrent replaces reported written=true for one live slot`, + ); + } + return violations; + }, + }; + }, + }, + { + // Documented gap: extend racing expiry. A short-lived lease is + // extended (pinned) while the clock races past its expiry and + // successors install over it. A pinned extend must never credit a + // successor, and a lease's final granted/expiry must equal exactly + // what its install plus its own pinned extends produced. + name: "extend racing expiry and successor installs", + slots: oneSlot, + async build(ctx) { + const setupRng = new Rng(deriveSeed(ctx.seed, 101)); + await installLease( + ctx, + ctx.slots[0], + quarters(setupRng, 2000, 800), + ctx.clock.now() + setupRng.intBetween(300, 1200), + ); + return { + ops: [ + advanceOp(ctx, 3), + extendObservedOp(ctx, [300, 1500]), + replaceFreshOp(ctx, { durationMs: [300, 1500], grantedQuarters: [800, 2000], onlyIfDead: true }), + reserveOp(ctx, { ttlMs: [100, 700], maxQuarters: 60 }), + consumeOp(ctx), + sweepOp(ctx), + expiredProbeOp(ctx), + ], + verify: () => verifyInvariants(ctx, { exactLeaseState: true }), + }; + }, + }, + { + // Documented gap: release (settle) racing the sweeper. One immortal + // lease; short-TTL reservations are settled by racing consumers while + // sweeps and clock advances fire concurrently, with occasional + // simulated window-1 crashes (debit without record). Conservation is + // exact: every credit is remaining, held, consumed, or accounted to + // the crash-leak budget — a double refund (sweeper + settle both + // winning a claim) or a lost refund breaks the equality. + name: "settle (consume) racing the sweeper", + slots: oneSlot, + async build(ctx) { + const setupRng = new Rng(deriveSeed(ctx.seed, 202)); + await installLease( + ctx, + ctx.slots[0], + quarters(setupRng, 1600, 1000), + ctx.clock.now() + 10 * 365 * 24 * 60 * 60 * 1000, + ); + return { + ops: [ + reserveOp(ctx, { ttlMs: [100, 700], maxQuarters: 400, crashChance: 0.1 }), + { ...consumeOp(ctx), weight: 4 }, + { ...sweepOp(ctx), weight: 3 }, + advanceOp(ctx, 3), + extendObservedOp(ctx, [300, 1500]), + invalidReserveOp(ctx), + ], + verify: () => verifyInvariants(ctx, { conservation: true, exactLeaseState: true, requireLease: true }), + }; + }, + }, + { + // Documented gap: replace racing try_reserve. Leases churn quickly + // (short durations, aggressive clock advances) while reserves are in + // flight, so debits race successor installs. A reserve must never be + // served by an expired balance and never mint credits on a successor. + name: "successor replace racing try_reserve", + slots: oneSlot, + async build(ctx) { + const setupRng = new Rng(deriveSeed(ctx.seed, 303)); + await installLease( + ctx, + ctx.slots[0], + quarters(setupRng, 1200, 400), + ctx.clock.now() + setupRng.intBetween(200, 1000), + ); + return { + ops: [ + replaceFreshOp(ctx, { + durationMs: [200, 1000], + grantedQuarters: [400, 1200], + onlyIfDead: true, + weight: 3, + }), + { ...reserveOp(ctx, { ttlMs: [100, 600], maxQuarters: 40 }), weight: 5 }, + consumeOp(ctx), + { ...sweepOp(ctx), weight: 1 }, + advanceOp(ctx, 3), + expiredProbeOp(ctx), + ], + verify: () => verifyInvariants(ctx, { exactLeaseState: true }), + }; + }, + }, + { + // Everything at once, across two slots: acquires, pinned extends, + // reserves, settles, sweeps, drops, clock advances, and window-1 + // crashes, all interleaved by seed. + name: "mixed chaos across slots", + slots: twoSlots, + async build(ctx) { + return { + ops: [ + replaceFreshOp(ctx, { durationMs: [200, 1200], grantedQuarters: [400, 1600], weight: 3 }), + extendObservedOp(ctx, [300, 1500]), + reserveOp(ctx, { ttlMs: [100, 700], maxQuarters: 60, crashChance: 0.1 }), + consumeOp(ctx), + sweepOp(ctx), + advanceOp(ctx, 3), + dropOp(ctx), + invalidReserveOp(ctx), + expiredProbeOp(ctx, 1), + ], + verify: () => verifyInvariants(ctx, { exactLeaseState: true }), + }; + }, + }, + { + // Documented gap: per-slot single-flight and the redundant-release + // rules, driven through the real CreditLeaseManager. Two managers + // ("pods") share one store and one scripted wire server whose calls + // carry seeded latency, so acquire responses race sibling installs + // and extend responses race expiry. The lease actually installed must + // never be released server-side, and the local view must never exceed + // the server-authoritative total. + name: "manager acquire/extend single-flight over a shared store", + slots: oneSlot, + async build(ctx) { + const buildRng = new Rng(deriveSeed(ctx.seed, 404)); + const server = new ScriptedCreditsServer(ctx.clock, new Rng(deriveSeed(ctx.seed, 405))); + const leaseDuration = buildRng.intBetween(800, 2500); + const managers = [0, 1].map( + () => + new CreditLeaseManager({ + creditsClient: server.client(), + leaseStore: ctx.leases, + logger: silentLogger, + config: { + defaultLeaseDuration: leaseDuration, + defaultReservationTTL: 60_000, + defaultLeaseSize: 500, + lowWaterMark: 0.25, + }, + }), + ); + const slot = ctx.slots[0]; + const ops: WeightedOp[] = [ + { + weight: 4, + run: async (rng) => { + await managers[rng.int(managers.length)].acquireIfNeeded(slot.companyId, slot.creditTypeId); + }, + }, + { + weight: 3, + run: async (rng) => { + const required = rng.chance(0.5) ? quarters(rng, 2400) : undefined; + await managers[rng.int(managers.length)].maybeExtendInBackground( + slot.companyId, + slot.creditTypeId, + required, + ); + }, + }, + reserveOp(ctx, { ttlMs: [100, 800], maxQuarters: 300 }), + consumeOp(ctx), + sweepOp(ctx), + advanceOp(ctx, 3), + expiredProbeOp(ctx, 1), + ]; + return { + ops, + verify: async () => { + const violations = await verifyInvariants(ctx, {}); + const entry = await ctx.leases.get(slot.companyId, slot.creditTypeId); + if (entry && entry.expiresAt.getTime() > ctx.clock.now()) { + if (server.released.has(entry.leaseId)) { + violations.push( + `live installed lease ${entry.leaseId} was released server-side ` + + `(a redundant-release pulled the shared lease out from under siblings)`, + ); + } + const srv = server.leasesById.get(entry.leaseId); + if (!srv) { + violations.push(`installed lease ${entry.leaseId} is unknown to the server`); + } else if (entry.grantedAmount > srv.granted + 1e-9) { + violations.push( + `local granted ${entry.grantedAmount} exceeds server-authoritative total ` + + `${srv.granted} for ${entry.leaseId} (phantom extend credits)`, + ); + } + } + return violations; + }, + }; + }, + }, +]; + +describe("lease/reservation concurrency race harness", () => { + for (const scenario of scenarios) { + for (const backend of backends) { + test(`${scenario.name} [${backend.name}] — ${seeds.length} seeded schedule(s)`, async () => { + for (const seed of seeds) { + await runScenario(scenario, backend, seed); + } + }); + } + } +}); diff --git a/tests/races/scripted-server.ts b/tests/races/scripted-server.ts new file mode 100644 index 00000000..a4ad2807 --- /dev/null +++ b/tests/races/scripted-server.ts @@ -0,0 +1,118 @@ +/** + * Deterministic in-process stand-in for the credits wire API, used by the + * manager-level race scenario. Mirrors the server contract the SDK relies on: + * + * - `acquire` is transactional and idempotent for an active slot: while a + * live, unreleased lease exists for (company, creditType) every acquire is + * handed that lease back; a fresh lease is minted only when the slot is + * free. The check-and-mint is atomic (no awaits inside), like the real + * server's transaction. + * - `extend` grows the lease's authoritative total and returns it; extending + * a released/expired lease errors. + * - `release` marks the lease released and frees the slot. + * + * Seeded jitter at the entry and exit of every call simulates wire latency, + * opening the windows (acquire response racing a sibling's install, extend + * response landing after expiry) the lease manager has to survive. + */ +import type { CreditsClient } from "../../src/api/resources/credits/client/Client"; + +import type { VirtualClock } from "./harness"; +import { jitter } from "./harness"; +import type { Rng } from "./prng"; + +interface ServerLease { + id: string; + companyId: string; + creditTypeId: string; + granted: number; + expiresAtMs: number; +} + +export class ScriptedCreditsServer { + readonly leasesById = new Map(); + readonly released = new Set(); + private readonly activeBySlot = new Map(); + private seq = 0; + + constructor( + private readonly clock: VirtualClock, + private readonly rng: Rng, + ) {} + + private activeLease(slotKey: string): ServerLease | undefined { + const id = this.activeBySlot.get(slotKey); + if (!id) return undefined; + const lease = this.leasesById.get(id); + if (!lease || this.released.has(id) || lease.expiresAtMs <= this.clock.now()) return undefined; + return lease; + } + + private toData(lease: ServerLease) { + // Snapshot inside the atomic section so the response carries the state + // as of the call, then travels through simulated latency. + return { + id: lease.id, + companyId: lease.companyId, + creditTypeId: lease.creditTypeId, + grantedAmount: lease.granted, + expiresAt: new Date(lease.expiresAtMs), + createdAt: new Date(0), + updatedAt: new Date(0), + }; + } + + client(): CreditsClient { + const acquireCreditLease = async (body: { + companyId: string; + creditTypeId: string; + requestedAmount: number; + expiresAt: Date; + }) => { + await jitter(this.rng); + const slotKey = `${body.companyId}:${body.creditTypeId}`; + let lease = this.activeLease(slotKey); + if (!lease) { + lease = { + id: `lse_srv_${++this.seq}`, + companyId: body.companyId, + creditTypeId: body.creditTypeId, + granted: body.requestedAmount, + expiresAtMs: body.expiresAt.getTime(), + }; + this.leasesById.set(lease.id, lease); + this.activeBySlot.set(slotKey, lease.id); + } + const data = this.toData(lease); + await jitter(this.rng); + return { data, params: {} }; + }; + + const extendCreditLease = async (leaseId: string, body: { additionalAmount: number; expiresAt: Date }) => { + await jitter(this.rng); + const lease = this.leasesById.get(leaseId); + if (!lease || this.released.has(leaseId) || lease.expiresAtMs <= this.clock.now()) { + throw new Error(`cannot extend inactive lease ${leaseId}`); + } + lease.granted += body.additionalAmount; + lease.expiresAtMs = Math.max(lease.expiresAtMs, body.expiresAt.getTime()); + const data = this.toData(lease); + await jitter(this.rng); + return { data, params: {} }; + }; + + const releaseCreditLease = async (leaseId: string) => { + await jitter(this.rng); + this.released.add(leaseId); + const lease = this.leasesById.get(leaseId); + if (lease) { + const slotKey = `${lease.companyId}:${lease.creditTypeId}`; + if (this.activeBySlot.get(slotKey) === leaseId) this.activeBySlot.delete(slotKey); + } + await jitter(this.rng); + return { data: {}, params: {} }; + }; + + return { acquireCreditLease, extendCreditLease, releaseCreditLease } as unknown as CreditsClient; + } +} diff --git a/tests/unit/credits/fake-redis.ts b/tests/unit/credits/fake-redis.ts index e9daf0ef..a3eb8693 100644 --- a/tests/unit/credits/fake-redis.ts +++ b/tests/unit/credits/fake-redis.ts @@ -38,7 +38,7 @@ export function makeFakeRedis(): RedisClient { const hset = (key: string, field: string, value: string): void => { if (!hashes.has(key)) hashes.set(key, new Map()); - hashes.get(key)!.set(field, value); + hashes.get(key)?.set(field, value); }; const hget = (key: string, field: string): string | null => { @@ -56,7 +56,7 @@ export function makeFakeRedis(): RedisClient { const zadd = (key: string, score: number, member: string): void => { if (!zsets.has(key)) zsets.set(key, new Map()); - zsets.get(key)!.set(member, score); + zsets.get(key)?.set(member, score); }; const zrem = (key: string, member: string): void => { @@ -85,6 +85,22 @@ export function makeFakeRedis(): RedisClient { if (existingId && existingExpiry > now) { return 0; } + if (existingId === newId) { + // Same-id reinstall over an expired row: reconcile like an + // extend (keep the debited balance) instead of rewriting. + const granted = Number(hget(leaseHashKey, "grantedAmount") ?? "0"); + const add = Number(newGranted) - granted; + if (add > 0) { + const remaining = Number(hget(leaseHashKey, "localRemainingCredits") ?? "0"); + hset(leaseHashKey, "grantedAmount", newGranted); + hset(leaseHashKey, "localRemainingCredits", String(remaining + add)); + } + if (newExpiry > existingExpiry) { + hset(leaseHashKey, "expiresAt", newExpiryStr); + expirations.set(leaseHashKey, newExpiry + Number(grace)); + } + return 0; + } del(leaseHashKey); hset(leaseHashKey, "leaseId", newId); hset(leaseHashKey, "companyId", companyId); diff --git a/tests/unit/credits/lease-store.test.ts b/tests/unit/credits/lease-store.test.ts index 83782332..9a4f9c76 100644 --- a/tests/unit/credits/lease-store.test.ts +++ b/tests/unit/credits/lease-store.test.ts @@ -249,4 +249,44 @@ describe("LeaseStore", () => { expect(store.get("co_1", "ct_1")?.leaseId).toBe("lse_2"); expect(store.get("co_1", "ct_1")?.localRemainingCredits).toBe(100); }); + + it("replace reconciles an expired SAME-id lease instead of resetting its debited balance", async () => { + // A stale acquire response can hand back the lease already installed + // (the server is idempotent for an active slot) after the local row + // expired — e.g. a concurrent extend pushed the server-side expiry + // forward. Rewriting would reset localRemainingCredits to the full + // grant, erasing debits whose reservations are still open. + const start = Date.now(); + const nowSpy = jest.spyOn(Date, "now").mockReturnValue(start); + try { + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(start + 60_000), + }); + await store.tryReserve("co_1", "ct_1", 40); + nowSpy.mockReturnValue(start + 61_000); // the local row expires + + const laterExpiry = new Date(start + 120_000); + const wrote = await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 150, + expiresAt: laterExpiry, + }); + expect(wrote).toBe(false); + const entry = store.get("co_1", "ct_1"); + expect(entry?.leaseId).toBe("lse_1"); + // Granted reconciled to the server total; the 40-credit debit survives. + expect(entry?.grantedAmount).toBe(150); + expect(entry?.localRemainingCredits).toBe(110); + // Expiry moves forward with the response. + expect(entry?.expiresAt.getTime()).toBe(laterExpiry.getTime()); + } finally { + nowSpy.mockRestore(); + } + }); }); diff --git a/tests/unit/credits/redis-lease-store.test.ts b/tests/unit/credits/redis-lease-store.test.ts index fbbf8dae..f07280cc 100644 --- a/tests/unit/credits/redis-lease-store.test.ts +++ b/tests/unit/credits/redis-lease-store.test.ts @@ -43,6 +43,43 @@ describe("RedisLeaseStore", () => { expect(entry?.localRemainingCredits).toBe(600); }); + it("replace reconciles an expired SAME-id lease instead of resetting its debited balance", async () => { + // A stale acquire response can hand back the lease already installed + // (the server is idempotent for an active slot) after the shared row + // expired — e.g. a concurrent extend pushed the server-side expiry + // forward. Rewriting would reset localRemainingCredits to the full + // grant, erasing debits whose reservations are still open. + const client = makeFakeRedis(); + const store = new RedisLeaseStore({ client }); + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 60_000), + }); + await store.tryReserve("co_1", "ct_1", 400); + // Expire the row in place (judged by the store clock). + await client.hSet(store.hashKey("co_1", "ct_1"), "expiresAt", String(Date.now() - 1_000)); + + const laterExpiry = new Date(Date.now() + 120_000); + const wrote = await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1500, + expiresAt: laterExpiry, + }); + expect(wrote).toBe(false); + const entry = await store.get("co_1", "ct_1"); + expect(entry?.leaseId).toBe("lse_1"); + // Granted reconciled to the server total; the 400-credit debit survives. + expect(entry?.grantedAmount).toBe(1500); + expect(entry?.localRemainingCredits).toBe(1100); + // Expiry moves forward with the response. + expect(entry?.expiresAt.getTime()).toBe(laterExpiry.getTime()); + }); + it("tryReserve atomically gates against the shared balance", async () => { const client = makeFakeRedis(); const store = new RedisLeaseStore({ client });