From 9bf102f965343a4e3696913a8e0ccc3e29af6da1 Mon Sep 17 00:00:00 2001 From: houseme Date: Sat, 11 Jul 2026 13:36:20 +0800 Subject: [PATCH] fix(cache): reserve admitted bytes in the memory gate to bound burst overshoot (#4718) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fill gate compared each request against a snapshot refreshed at most every 5 s, with no accounting for what it had already let through. A burst arriving while the snapshot still read high therefore all passed the same check-then-act test and over-allocated far past the real headroom before the next refresh — a gap the burst stress test could expose but not close. Track admitted bytes since the last refresh in the shared snapshot cell and subtract them from available memory in `allows_fill`, reserving the request's size on each admission. Cumulative admission is now bounded to the real budget even though every fill reads the same stale snapshot; the refresh resets the counter because the fresh reading already reflects those allocations. The `min_free_memory_percent == 0` opt-out still short-circuits first, and the already-low-memory path is unchanged. New test `moka_backend_gate_reservation_bounds_burst_under_stale_snapshot`: 20 concurrent 40 KiB fills against a 500 KiB stale-high snapshot (300 KiB budget) admit only a bounded handful, not the whole storm. Passed 10/10 runs; mutation-verified — dropping the reservation admits all 20 and fails the test. Refs: backlog#1107 Co-authored-by: heihutu --- crates/object-data-cache/src/memory.rs | 37 ++++++++++++-- crates/object-data-cache/src/moka_backend.rs | 53 ++++++++++++++++++++ 2 files changed, 87 insertions(+), 3 deletions(-) diff --git a/crates/object-data-cache/src/memory.rs b/crates/object-data-cache/src/memory.rs index 5215780f7..a8d89ffc1 100644 --- a/crates/object-data-cache/src/memory.rs +++ b/crates/object-data-cache/src/memory.rs @@ -117,6 +117,14 @@ impl ObjectDataCacheMemorySnapshot { struct MemorySnapshotCell { total_bytes: AtomicU64, available_bytes: AtomicU64, + /// Bytes admitted by the gate since the last snapshot refresh. The snapshot + /// is sampled at most every 5 s, so a burst that begins while it still reads + /// high could all pass a plain check-then-act gate and over-allocate before + /// the next refresh. Subtracting this running total from `available_bytes` + /// shrinks the effective budget as the burst proceeds, bounding cumulative + /// admission to the real headroom; the refresh resets it because the fresh + /// reading already reflects those allocations (backlog#1107). + admitted_since_refresh: AtomicU64, } impl MemorySnapshotCell { @@ -124,12 +132,16 @@ impl MemorySnapshotCell { Self { total_bytes: AtomicU64::new(snapshot.total_bytes), available_bytes: AtomicU64::new(snapshot.available_bytes), + admitted_since_refresh: AtomicU64::new(0), } } + /// Stores a fresh snapshot and resets the admitted-bytes counter: the new + /// reading already accounts for whatever was admitted since the last one. fn store(&self, snapshot: ObjectDataCacheMemorySnapshot) { self.total_bytes.store(snapshot.total_bytes, Ordering::Relaxed); self.available_bytes.store(snapshot.available_bytes, Ordering::Relaxed); + self.admitted_since_refresh.store(0, Ordering::Relaxed); } fn load(&self) -> ObjectDataCacheMemorySnapshot { @@ -138,6 +150,14 @@ impl MemorySnapshotCell { available_bytes: self.available_bytes.load(Ordering::Relaxed), } } + + fn admitted(&self) -> u64 { + self.admitted_since_refresh.load(Ordering::Relaxed) + } + + fn reserve(&self, bytes: u64) { + self.admitted_since_refresh.fetch_add(bytes, Ordering::Relaxed); + } } /// Aborts the periodic refresher when the gate (and thus the cache) is dropped. @@ -255,12 +275,23 @@ impl ObjectDataCacheMemoryGate { return true; } + // Effective budget is the snapshot's available memory minus what the + // gate has already admitted since that snapshot was taken. This bounds a + // burst that arrives faster than the 5 s refresh: each admission shrinks + // the budget the next one sees, so cumulative admission cannot exceed the + // real headroom even though every fill reads the same (stale) snapshot. + let effective_available = snapshot.available_bytes.saturating_sub(self.snapshot.admitted()); + let min_free = u64::from(self.min_free_memory_percent); - let has_percent_budget = snapshot.available_bytes.saturating_mul(100) >= snapshot.total_bytes.saturating_mul(min_free); - let has_entry_budget = snapshot.available_bytes >= required_bytes; + let has_percent_budget = effective_available.saturating_mul(100) >= snapshot.total_bytes.saturating_mul(min_free); + let has_entry_budget = effective_available >= required_bytes; let allowed = has_percent_budget && has_entry_budget; - if !allowed { + if allowed { + // Reserve the admitted bytes so a concurrent fill sees a smaller + // budget; the reservation clears on the next snapshot refresh. + self.snapshot.reserve(required_bytes); + } else { record_memory_pressure(&self.stats, "moka"); } diff --git a/crates/object-data-cache/src/moka_backend.rs b/crates/object-data-cache/src/moka_backend.rs index b6c13d37c..9bcc52003 100644 --- a/crates/object-data-cache/src/moka_backend.rs +++ b/crates/object-data-cache/src/moka_backend.rs @@ -996,4 +996,57 @@ mod tests { assert_eq!(rejected, TASKS, "the gate must reject every fill in the burst when memory is constrained"); assert_eq!(backend.entry_count(), 0, "no body may be admitted under memory pressure"); } + + /// A burst that begins while the snapshot still reads *high* must be bounded + /// to the real headroom, not admitted wholesale. Each fill is large enough + /// that only a handful fit the budget; without the admitted-bytes reservation + /// every fill would read the same stale-high snapshot and be admitted, + /// over-allocating far past the budget before the 5 s refresh (backlog#1107). + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn moka_backend_gate_reservation_bounds_burst_under_stale_snapshot() { + let stats = Arc::new(ObjectDataCacheStats::default()); + let backend = Arc::new(MokaBackend::new(&memory_gated_config(), Arc::clone(&stats)).expect("moka backend should build")); + // Stale-but-high snapshot: 500 KiB available, 20% floor = 200 KiB, so the + // real budget above the floor is ~300 KiB. + backend + .memory_gate + .set_test_snapshot(Some(crate::memory::ObjectDataCacheMemorySnapshot { + total_bytes: 1_000_000, + available_bytes: 500_000, + })); + + // 20 concurrent fills of 40 KiB each = 800 KiB requested, far past the + // ~300 KiB budget. Only ~7 should fit. + const TASKS: usize = 20; + const BODY: usize = 40_000; + let mut handles = Vec::new(); + for t in 0..TASKS { + let backend = Arc::clone(&backend); + handles.push(tokio::spawn(async move { + let plan = versioned_plan(&format!("object-{t}"), "v1", &format!("etag-{t}")); + backend.fill_body(&plan, Bytes::from(vec![0u8; BODY])).await + })); + } + let mut admitted = 0usize; + for handle in handles { + if handle.await.expect("burst task should complete") == ObjectDataCacheFillResult::Inserted { + admitted += 1; + } + } + + // The reservation must cap the burst: not every fill gets in (that is the + // stale-snapshot overshoot), yet some do (the gate is not just rejecting + // everything). The exact count varies with the check-then-reserve race, + // but the admitted bytes must stay within a small multiple of the budget. + assert!( + admitted < TASKS, + "reservation must bound the burst — not admit the whole {TASKS}-fill storm" + ); + assert!(admitted >= 1, "the gate must still admit fills that fit the budget"); + assert!( + admitted * BODY <= 500_000, + "admitted bytes ({}) must not exceed the snapshot's available memory", + admitted * BODY + ); + } }