diff --git a/apps/gateway/src/cache.rs b/apps/gateway/src/cache.rs index c95edbab..99a925df 100644 --- a/apps/gateway/src/cache.rs +++ b/apps/gateway/src/cache.rs @@ -41,6 +41,18 @@ pub(crate) trait CacheStore: Send + Sync { /// Sets TTL only on first increment (new key / expired key). /// Returns the new count, or `None` on error (graceful fallback). async fn incr(&self, key: &str, ttl_secs: u64) -> Option; + + /// Add `delta` to a signed counter, returning the new total. + /// + /// Distinct from [`incr`], which is a +1 rate-limit counter over `u64`. + /// Spend counters are signed nano-dollar totals and move by a per-request + /// amount, so they need their own primitive rather than a loop over `incr`. + /// + /// Atomic per key: the read-modify-write happens under the entry lock, so + /// concurrent charges against one budget accumulate instead of racing on a + /// read-then-set. That is the whole point of having it — a + /// `get_raw`/`set_raw` pair would lose every charge but the last. + async fn incr_by(&self, key: &str, delta: i64, ttl_secs: u64) -> Option; } /// Extension methods for typed get/set on any `CacheStore`. @@ -151,6 +163,32 @@ impl CacheStore for InMemoryCacheStore { entry.data = count.to_string(); Some(count) } + + async fn incr_by(&self, key: &str, delta: i64, ttl_secs: u64) -> Option { + let now = Instant::now(); + let ttl = Duration::from_secs(ttl_secs); + + // `entry` holds the shard lock for the whole read-modify-write, so two + // concurrent charges on one budget cannot both read the same total. + let mut entry = self.map.entry(key.to_string()).or_insert(CachedEntry { + data: "0".to_string(), + expires_at: now + ttl, + }); + + if entry.expires_at <= now { + entry.data = "0".to_string(); + entry.expires_at = now + ttl; + } + + // An unparseable value is treated as 0 rather than propagating: this is + // a cache, and the durable BudgetSpend row is the floor that corrects it. + let total = entry.data.parse::().unwrap_or(0).saturating_add(delta); + entry.data = total.to_string(); + // Refresh the window on write so an actively-charged period never + // expires mid-period. + entry.expires_at = now + ttl; + Some(total) + } } // ── Tests ──────────────────────────────────────────────────────────────── @@ -189,6 +227,49 @@ mod tests { assert!(result.is_none()); } + #[tokio::test] + async fn incr_by_accumulates_concurrent_charges() { + // THE POINT OF incr_by. A get_raw/set_raw pair would let two concurrent + // charges both read 0 and both write their own delta, keeping only the + // last — which is exactly how spend used to slip past a cap. + let store = new_store(); + let key = "budget:spent:sec:org:m:2026-08"; + + let tasks: Vec<_> = (0..64) + .map(|_| { + let store = Arc::clone(&store); + let key = key.to_string(); + tokio::spawn(async move { store.incr_by(&key, 1_000, 60).await }) + }) + .collect(); + for t in tasks { + t.await.unwrap(); + } + + let total: i64 = store.get_raw(key).await.unwrap().parse().unwrap(); + assert_eq!(total, 64_000, "every concurrent charge must be counted"); + } + + #[tokio::test] + async fn incr_by_starts_from_zero_and_survives_a_set_floor() { + let store = new_store(); + let key = "budget:spent:sec:org:total"; + + assert_eq!(store.incr_by(key, 250, 60).await, Some(250)); + // A durable floor arriving from PostgreSQL, then more charges on top. + store.set_raw(key, "1000", 60).await; + assert_eq!(store.incr_by(key, 25, 60).await, Some(1025)); + } + + #[tokio::test] + async fn incr_by_treats_an_unparseable_value_as_zero() { + // The counter shares a namespace with JSON-valued caches; a stray value + // must not poison spend accounting. The durable row is the floor. + let store = new_store(); + store.set_raw("budget:spent:x", "not-a-number", 60).await; + assert_eq!(store.incr_by("budget:spent:x", 7, 60).await, Some(7)); + } + #[tokio::test] async fn del_removes_entry() { let store = new_store(); diff --git a/apps/gateway/src/telemetry.rs b/apps/gateway/src/telemetry.rs index 9b70d373..d5820747 100644 --- a/apps/gateway/src/telemetry.rs +++ b/apps/gateway/src/telemetry.rs @@ -42,10 +42,24 @@ async fn record_spend(pool: &PgPool, cache: &dyn CacheStore, key: &SpendKey, del match crate::db::upsert_budget_spend(pool, secret_id, organization_id, period_key, delta).await { Ok(total) => { + // RECONCILE, don't overwrite. The cache was already incremented for + // this charge (`charge_cache` below), and may have been incremented + // again for a charge that lands in the NEXT batch. A blind + // `set_raw(total)` would roll those newer charges back and let spend + // through twice. + // + // The durable row is a FLOOR: raise the counter to it when the cache + // is behind (cold start, eviction, a lost increment), never lower it. let counter = crate::budget::counter_key(secret_id, organization_id, period_key); - cache - .set_raw(&counter, &total.to_string(), crate::budget::PERIOD_TTL) - .await; + let cached = cache + .get_raw(&counter) + .await + .and_then(|raw| raw.parse::().ok()); + if cached.is_none_or(|c| c < total) { + cache + .set_raw(&counter, &total.to_string(), crate::budget::PERIOD_TTL) + .await; + } } Err(e) => { warn!(error = %e, secret_id = %secret_id, "budget: failed to record spend"); @@ -53,6 +67,26 @@ async fn record_spend(pool: &PgPool, cache: &dyn CacheStore, key: &SpendKey, del } } +/// Apply a charge to the hot counter the moment it is drained from the channel, +/// rather than waiting for the batch to reach PostgreSQL. +/// +/// `pre_forward` reads this counter to decide whether to deny with 402. Before +/// this, the counter only moved when the flush loop completed its upsert, so +/// requests arriving between a charge and its flush all read the same stale +/// total and were all admitted. The window was small — `collect_batch` returns +/// as soon as one event arrives, so the 5s interval only applies when idle — +/// but it was wide enough for concurrent requests against one budget. +/// +/// Fail-open, like every other budget read: a cache miss returns None and the +/// request proceeds. The durable floor in `record_spend` repairs the counter. +async fn charge_cache(cache: &dyn CacheStore, key: &SpendKey, delta: i64) { + let (secret_id, organization_id, period_key) = key; + let counter = crate::budget::counter_key(secret_id, organization_id, period_key); + cache + .incr_by(&counter, delta, crate::budget::PERIOD_TTL) + .await; +} + async fn insert_batch(pool: &PgPool, events: &[RequestEvent]) -> Result<(), sqlx::Error> { let filtered: Vec<&RequestEvent> = events .iter() @@ -181,7 +215,14 @@ async fn flush_loop( } } - // Persist the aggregated spend deltas (few keys per flush in practice). + // Move the hot counter FIRST, before the (slower) database round-trip: + // this is what `pre_forward` reads, so every microsecond it lags is a + // window in which a concurrent request sees a stale total. + for (key, delta) in &charges { + charge_cache(cache.as_ref(), key, *delta).await; + } + + // Then persist, and reconcile the counter against the durable floor. for (key, delta) in &charges { record_spend(&pool, cache.as_ref(), key, *delta).await; }