From f5b36376d999cb3322bcc8974b6774197f7b3688 Mon Sep 17 00:00:00 2001 From: epicvinny Date: Mon, 17 Aug 2026 21:21:55 -0300 Subject: [PATCH 1/8] feat(pool): decay warm-process refill target after idle TTL The warm pool ratchets its geometric refill target up after bursts of acquisitions and never decays it, so a single burst pins the excess warm processes (and their memory) for the lifetime of the server. Add PoolConfig::idle_ttl: when no acquisition happens within the TTL, the maintenance worker decays the fill target back to the low watermark and drains the excess entries. Acquisitions reset the idle clock. Wire it as [pool.firecracker] idle_ttl_secs (default 600, 0 disables decay), documented in config/default.toml and the configuration reference. The network-slot and ublk pools keep the historical ratchet behavior (idle_ttl: None): they hold no guest processes, so idle decay buys nothing there. Tests: idle_ttl_decays_fill_target_and_drains_to_low_watermark, maintenance_worker_wakes_on_idle_ttl_and_drains, acquisitions_reset_idle_ttl_clock (cargo test -p warm-pool: 16 ok). --- config/default.toml | 4 + crates/warm-pool/src/lib.rs | 193 +++++++++++++++++++++++++++- docs/src/configuration/reference.md | 1 + src/cfg.rs | 12 ++ src/sandbox/network/manager.rs | 1 + storage/ublk-daemon/src/main.rs | 2 + storage/ublk-daemon/src/server.rs | 1 + 7 files changed, 207 insertions(+), 7 deletions(-) diff --git a/config/default.toml b/config/default.toml index 00f94bf7..c5b3d686 100644 --- a/config/default.toml +++ b/config/default.toml @@ -298,6 +298,10 @@ enabled = true startup_prewarm = true # Maximum concurrent process creations in one refill batch. fill_concurrency = 4 +# Seconds without warm-process acquisitions before the pool decays its refill +# target back to the low watermark and drains the excess warm processes +# (0 = never decay). +idle_ttl_secs = 600 [memory_snapshot] # Generated OverlayBD config dedicated to memory snapshot devices. diff --git a/crates/warm-pool/src/lib.rs b/crates/warm-pool/src/lib.rs index c288b46e..e1f0f15d 100644 --- a/crates/warm-pool/src/lib.rs +++ b/crates/warm-pool/src/lib.rs @@ -3,6 +3,7 @@ //! This crate provides reusable pool mechanics for resources that are expensive //! to create but can be reset and reused. It handles: //! - Watermark-based refill/drain decisions +//! - Idle TTL decay of the geometric refill target //! - Background maintenance worker with condvar signaling //! - Shutdown coordination with safe resource cleanup //! - Process exit hooks for static singleton pools @@ -12,6 +13,7 @@ use std::collections::VecDeque; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Condvar, Mutex}; +use std::time::{Duration, Instant}; /// Action computed by watermark logic for the maintenance worker. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -50,6 +52,11 @@ pub struct PoolConfig { /// Advisory flag for callers that can prewarm once a reusable resource /// shape is known. `WarmPool` itself only owns generic pool mechanics. pub startup_prewarm: bool, + /// Maximum time without acquisitions before the geometric fill target + /// decays back to the low watermark and idle resources above it are + /// drained. `None` keeps the fill target ratcheted for the process + /// lifetime. + pub idle_ttl: Option, } impl PoolConfig { @@ -88,11 +95,15 @@ pub struct WarmPool { /// Watermark config. config: PoolConfig, /// Current refill target. Starts at the low watermark and grows toward the - /// high watermark under acquisition pressure. This intentionally ratchets - /// upward for the process lifetime: after a node observes bursty demand, it - /// keeps extra warm capacity instead of shrinking back to cold-start - /// behavior. + /// high watermark under acquisition pressure. With `idle_ttl` unset this + /// intentionally ratchets upward for the process lifetime: after a node + /// observes bursty demand, it keeps extra warm capacity instead of + /// shrinking back to cold-start behavior. When `idle_ttl` is set, a + /// sustained idle period decays the target back to the low watermark so + /// warm capacity (and the resources it holds) is released. fill_target: Mutex, + /// Last time an acquisition was attempted. Drives idle TTL decay. + last_acquisition: Mutex, /// Background maintenance worker state. maintenance_signal: Mutex, /// Wakes the maintenance worker. @@ -113,6 +124,7 @@ impl WarmPool { Self { pool: Mutex::new(VecDeque::new()), fill_target: Mutex::new(fill_target), + last_acquisition: Mutex::new(Instant::now()), config, maintenance_signal: Mutex::new(PoolMaintenanceSignal::default()), maintenance_cv: Condvar::new(), @@ -144,6 +156,7 @@ impl WarmPool { /// Compute the maintenance action based on current pool size. pub fn compute_maintenance_action(&self, pool_len: usize) -> PoolMaintenanceAction { + let idle_expired = self.decay_fill_target_if_idle(); let fill_target = self.current_fill_target(); if pool_len < fill_target { let to_fill = fill_target.saturating_sub(pool_len); @@ -151,11 +164,19 @@ impl WarmPool { return PoolMaintenanceAction::Fill(to_fill); } } - if pool_len > self.config.high_watermark { + // After an idle TTL decay the pool shrinks toward the decayed fill + // target (the low watermark); otherwise only the high watermark caps + // idle resources. + let drain_target = if idle_expired { + fill_target + } else { + self.config.high_watermark + }; + if pool_len > drain_target { // Drain the full excess in one maintenance cycle. Resource-specific // cleanup happens outside the pool lock, and shutdown paths already // have to tolerate draining the whole pool. - let to_drain = pool_len - self.config.high_watermark; + let to_drain = pool_len - drain_target; if to_drain > 0 { return PoolMaintenanceAction::Drain(to_drain); } @@ -167,6 +188,33 @@ impl WarmPool { (*self.fill_target.lock().unwrap()).min(self.config.high_watermark) } + /// Collapse the geometric fill target back to the low watermark when no + /// acquisition happened within `idle_ttl`. Returns true when the TTL had + /// expired, so callers drain excess idle resources toward the low + /// watermark instead of the high one. The idle clock restarts on each + /// expiry so a fully decayed pool does not retrigger every cycle. + fn decay_fill_target_if_idle(&self) -> bool { + let Some(ttl) = self.config.idle_ttl else { + return false; + }; + let mut last = self.last_acquisition.lock().unwrap(); + if last.elapsed() < ttl { + return false; + } + *last = Instant::now(); + let low = self.config.low_watermark.min(self.config.high_watermark); + let mut target = self.fill_target.lock().unwrap(); + *target = (*target).min(low); + true + } + + /// Time until the idle TTL expires, if decay is configured. + fn idle_ttl_remaining(&self) -> Option { + let ttl = self.config.idle_ttl?; + let elapsed = self.last_acquisition.lock().unwrap().elapsed(); + Some(ttl.saturating_sub(elapsed)) + } + fn grow_fill_target_after_pressure(&self, pool_len: usize) { if pool_len >= self.config.low_watermark || self.config.high_watermark == 0 { return; @@ -183,6 +231,7 @@ impl WarmPool { } fn record_acquisition_pressure(&self, pool_len: usize) { + *self.last_acquisition.lock().unwrap() = Instant::now(); self.grow_fill_target_after_pressure(pool_len); if matches!( self.compute_maintenance_action(pool_len), @@ -343,7 +392,22 @@ impl WarmPool { if !has_immediate_work { let mut signal = self.maintenance_signal.lock().unwrap(); while !signal.stop && !signal.pending { - signal = self.maintenance_cv.wait(signal).unwrap(); + match self.idle_ttl_remaining() { + Some(remaining) => { + // Wake when the idle TTL expires even if nothing + // requested maintenance, so the fill target can + // decay and excess idle resources drain. + let (new_signal, timeout) = + self.maintenance_cv.wait_timeout(signal, remaining).unwrap(); + signal = new_signal; + if timeout.timed_out() && !signal.stop && !signal.pending { + signal.pending = true; + } + } + None => { + signal = self.maintenance_cv.wait(signal).unwrap(); + } + } } if signal.stop { break; @@ -402,6 +466,7 @@ mod tests { high_watermark: 32, maintenance_enabled: true, startup_prewarm: false, + idle_ttl: None, } .validate(); assert_eq!(config.low_watermark, 32); @@ -415,6 +480,7 @@ mod tests { high_watermark: 10, maintenance_enabled: true, startup_prewarm: false, + idle_ttl: None, }); assert_eq!( pool.compute_maintenance_action(2), @@ -433,6 +499,7 @@ mod tests { high_watermark: 10, maintenance_enabled: true, startup_prewarm: false, + idle_ttl: None, }); assert_eq!( @@ -462,6 +529,7 @@ mod tests { high_watermark: 10, maintenance_enabled: false, startup_prewarm: false, + idle_ttl: None, }); assert_eq!(pool.try_acquire(), None); @@ -490,6 +558,7 @@ mod tests { high_watermark: 10, maintenance_enabled: true, startup_prewarm: false, + idle_ttl: None, }); assert!(!pool.maintenance_signal.lock().unwrap().pending); @@ -504,6 +573,7 @@ mod tests { high_watermark: 4, maintenance_enabled: true, startup_prewarm: false, + idle_ttl: None, }); assert_eq!( pool.compute_maintenance_action(8), @@ -522,6 +592,7 @@ mod tests { high_watermark: 10, maintenance_enabled: false, startup_prewarm: false, + idle_ttl: None, }); assert!(pool.try_acquire().is_none()); } @@ -533,6 +604,7 @@ mod tests { high_watermark: 10, maintenance_enabled: false, startup_prewarm: false, + idle_ttl: None, }); pool.release(42).unwrap(); assert_eq!(pool.try_acquire(), Some(42)); @@ -545,6 +617,7 @@ mod tests { high_watermark: 2, maintenance_enabled: false, startup_prewarm: false, + idle_ttl: None, }); assert!(pool.release(1).is_ok()); assert!(pool.release(2).is_ok()); @@ -558,6 +631,7 @@ mod tests { high_watermark: 2, maintenance_enabled: true, startup_prewarm: false, + idle_ttl: None, }); assert!(pool.release(1).is_ok()); assert!(pool.release(2).is_ok()); @@ -572,6 +646,7 @@ mod tests { high_watermark: 10, maintenance_enabled: false, startup_prewarm: false, + idle_ttl: None, }); pool.release(1).unwrap(); pool.release(2).unwrap(); @@ -588,6 +663,7 @@ mod tests { high_watermark: 10, maintenance_enabled: false, startup_prewarm: false, + idle_ttl: None, }); pool.release(42).unwrap(); pool.drain_all(); @@ -601,8 +677,111 @@ mod tests { high_watermark: 10, maintenance_enabled: false, startup_prewarm: false, + idle_ttl: None, }); pool.drain_all(); assert_eq!(pool.release(42), Err(42)); } + + #[test] + fn idle_ttl_decays_fill_target_and_drains_to_low_watermark() { + let pool = WarmPool::::new(PoolConfig { + low_watermark: 2, + high_watermark: 10, + maintenance_enabled: false, + startup_prewarm: false, + idle_ttl: Some(Duration::from_millis(50)), + }); + + // Ratchet the fill target up to 8 under acquisition pressure. + pool.release(1).unwrap(); + pool.release(2).unwrap(); + assert_eq!(pool.try_acquire(), Some(1)); + assert_eq!(pool.try_acquire(), Some(2)); + assert_eq!( + pool.compute_maintenance_action(0), + PoolMaintenanceAction::Fill(8) + ); + + for value in 1..=6 { + pool.release(value).unwrap(); + } + + std::thread::sleep(Duration::from_millis(80)); + + // Idle past the TTL: excess drains toward the decayed fill target + // (low watermark) instead of lingering below the high watermark. + assert_eq!( + pool.compute_maintenance_action(6), + PoolMaintenanceAction::Drain(4) + ); + assert_eq!( + pool.compute_maintenance_action(2), + PoolMaintenanceAction::Idle + ); + assert_eq!( + pool.compute_maintenance_action(1), + PoolMaintenanceAction::Fill(1) + ); + } + + #[test] + fn acquisitions_reset_idle_ttl_clock() { + let pool = WarmPool::::new(PoolConfig { + low_watermark: 2, + high_watermark: 10, + maintenance_enabled: false, + startup_prewarm: false, + idle_ttl: Some(Duration::from_millis(300)), + }); + + assert_eq!(pool.try_acquire(), None); + std::thread::sleep(Duration::from_millis(100)); + assert_eq!(pool.try_acquire(), None); + std::thread::sleep(Duration::from_millis(150)); + + // 150ms since the last acquisition: the 300ms TTL has not expired and + // the ratcheted fill target is preserved. + assert_eq!( + pool.compute_maintenance_action(0), + PoolMaintenanceAction::Fill(8) + ); + + std::thread::sleep(Duration::from_millis(200)); + assert_eq!( + pool.compute_maintenance_action(0), + PoolMaintenanceAction::Fill(2) + ); + } + + #[test] + fn maintenance_worker_wakes_on_idle_ttl_and_drains() { + let pool: &'static WarmPool = Box::leak(Box::new(WarmPool::new(PoolConfig { + low_watermark: 0, + high_watermark: 4, + maintenance_enabled: true, + startup_prewarm: false, + idle_ttl: Some(Duration::from_millis(50)), + }))); + pool.start_maintenance_worker(move || { + if let PoolMaintenanceAction::Drain(to_drain) = + pool.compute_maintenance_action(pool.len()) + { + for _ in 0..to_drain { + pool.try_drain_one(); + } + } + }); + + pool.release(1).unwrap(); + pool.release(2).unwrap(); + pool.release(3).unwrap(); + + // No acquisition happens, so the idle TTL wake-up drains everything + // (low watermark is 0) without any explicit maintenance request. + std::thread::sleep(Duration::from_millis(300)); + assert_eq!(pool.len(), 0); + + pool.drain_all(); + } } diff --git a/docs/src/configuration/reference.md b/docs/src/configuration/reference.md index d2ad8049..fdd505d5 100644 --- a/docs/src/configuration/reference.md +++ b/docs/src/configuration/reference.md @@ -294,6 +294,7 @@ Component sections: | `[pool.firecracker]` | `maintenance_enabled` | boolean | `true` | Enable the background Firecracker process maintenance worker | | `[pool.firecracker]` | `startup_prewarm` | boolean | `true` | Spawn warm Firecracker entries up to the low watermark during server startup | | `[pool.firecracker]` | `fill_concurrency` | integer | `4` | Maximum number of warm Firecracker processes created concurrently by one maintenance refill batch | +| `[pool.firecracker]` | `idle_ttl_secs` | integer | `600` | Seconds without warm-process acquisitions before the refill target decays back to the low watermark and excess warm processes are drained (`0` disables decay) | Validation rules: diff --git a/src/cfg.rs b/src/cfg.rs index 405543f2..59ad87be 100644 --- a/src/cfg.rs +++ b/src/cfg.rs @@ -1,5 +1,6 @@ use std::path::{Path, PathBuf}; use std::sync::OnceLock; +use std::time::Duration; pub(crate) mod image; pub(crate) mod network; @@ -215,6 +216,12 @@ pub struct FirecrackerProcessPoolConfig { pub startup_prewarm: bool, #[config(default = 4usize)] pub fill_concurrency: usize, + /// Seconds without pool acquisitions before the geometric fill target + /// decays back to the low watermark and excess warm processes are drained. + /// 0 disables decay and keeps the fill target ratcheted for the process + /// lifetime. + #[config(default = 600u64)] + pub idle_ttl_secs: u64, } #[derive(Debug, Clone)] @@ -764,6 +771,9 @@ impl AppConfig { high_watermark: self.pool.high_watermark, maintenance_enabled: pool.enabled && pool.maintenance_enabled, startup_prewarm: pool.startup_prewarm, + // The network slot pool holds no processes; keep its historical + // ratchet behavior without idle decay. + idle_ttl: None, } } @@ -781,6 +791,7 @@ impl AppConfig { // reusable device shape is image/size dependent. maintenance_enabled: false, startup_prewarm: pool.startup_prewarm, + idle_ttl: None, }) } @@ -797,6 +808,7 @@ impl AppConfig { high_watermark: self.pool.high_watermark, maintenance_enabled: pool.maintenance_enabled, startup_prewarm: pool.startup_prewarm, + idle_ttl: (pool.idle_ttl_secs > 0).then(|| Duration::from_secs(pool.idle_ttl_secs)), }, fill_concurrency: pool.fill_concurrency, }) diff --git a/src/sandbox/network/manager.rs b/src/sandbox/network/manager.rs index fce7136a..ffafd03f 100644 --- a/src/sandbox/network/manager.rs +++ b/src/sandbox/network/manager.rs @@ -115,6 +115,7 @@ impl NetworkManager { high_watermark, maintenance_enabled, startup_prewarm: true, + idle_ttl: None, }, address_plan: NetworkAddressPlan::default(), netns_dir: std::env::temp_dir().join("aenv-network-tests/netns"), diff --git a/storage/ublk-daemon/src/main.rs b/storage/ublk-daemon/src/main.rs index 4a945026..b150d01a 100644 --- a/storage/ublk-daemon/src/main.rs +++ b/storage/ublk-daemon/src/main.rs @@ -141,6 +141,7 @@ fn load_pool_config( .startup_prewarm .or_else(|| pool.and_then(|pool| pool.startup_prewarm)) .unwrap_or(true), + idle_ttl: None, })) } @@ -232,6 +233,7 @@ fn default_pool_config() -> warm_pool::PoolConfig { high_watermark: 64, maintenance_enabled: false, startup_prewarm: true, + idle_ttl: None, } } diff --git a/storage/ublk-daemon/src/server.rs b/storage/ublk-daemon/src/server.rs index 1803e150..d9542d49 100644 --- a/storage/ublk-daemon/src/server.rs +++ b/storage/ublk-daemon/src/server.rs @@ -1658,6 +1658,7 @@ mod tests { high_watermark: 1, maintenance_enabled: false, startup_prewarm: false, + idle_ttl: None, } } From 00784ce311ab2084120436fffbf3e923f3a3ed26 Mon Sep 17 00:00:00 2001 From: epicvinny Date: Mon, 17 Aug 2026 22:21:26 -0300 Subject: [PATCH 2/8] fix(pool): persist idle decay state and normalize zero idle TTL decay_fill_target_if_idle reset last_acquisition on expiry, so a single action computation consumed the decay event: if the drain cycle only partially succeeded (e.g. resource cleanup failed mid-cycle) or another caller computed the action first, the next computation fell back to the high watermark and retained the excess for another full TTL. Keep a persistent decaying flag instead: set on expiry, cleared by any acquisition, and cleared only once the pool reaches the decayed fill target, so the drain target stays pinned to the low watermark until the excess is actually drained. PoolConfig::validate now normalizes idle_ttl = Some(Duration::ZERO) to None, matching the documented external config semantics of 0 = never decay. A zero TTL made idle_ttl_remaining always return zero, so the maintenance worker scheduled back-to-back cycles (busy loop) even after reaching the target. Tests: widen TTL/sleep margins so loaded CI workers cannot cross the deadline (pre-expiry margin is now 800ms instead of 50ms), poll the maintenance-worker test with a 10s deadline instead of a fixed 300ms sleep, and cover the zero-TTL normalization and the non-consumable decay. --- crates/warm-pool/src/lib.rs | 102 ++++++++++++++++++++++++++++++------ 1 file changed, 87 insertions(+), 15 deletions(-) diff --git a/crates/warm-pool/src/lib.rs b/crates/warm-pool/src/lib.rs index e1f0f15d..61b59be2 100644 --- a/crates/warm-pool/src/lib.rs +++ b/crates/warm-pool/src/lib.rs @@ -70,6 +70,15 @@ impl PoolConfig { ); self.low_watermark = self.high_watermark; } + // A zero idle TTL would expire on every cycle: `idle_ttl_remaining` + // would always return zero and the maintenance worker would schedule + // back-to-back cycles (busy loop) even after the pool reached its + // target. The external config documents 0 as "never decay", so + // normalize to `None` here; internal uses can then assume `Some(d)` + // always carries d > 0. + if self.idle_ttl == Some(Duration::ZERO) { + self.idle_ttl = None; + } self } } @@ -104,6 +113,12 @@ pub struct WarmPool { fill_target: Mutex, /// Last time an acquisition was attempted. Drives idle TTL decay. last_acquisition: Mutex, + /// Persistent "draining after idle decay" state. Set when the idle TTL + /// expires, cleared by any acquisition, and cleared once the pool has + /// drained to the decayed fill target. Keeps the drain target pinned to + /// the low watermark across partially-failed drain cycles, so the decay + /// event cannot be consumed by a single action computation. + decaying: AtomicBool, /// Background maintenance worker state. maintenance_signal: Mutex, /// Wakes the maintenance worker. @@ -125,6 +140,7 @@ impl WarmPool { pool: Mutex::new(VecDeque::new()), fill_target: Mutex::new(fill_target), last_acquisition: Mutex::new(Instant::now()), + decaying: AtomicBool::new(false), config, maintenance_signal: Mutex::new(PoolMaintenanceSignal::default()), maintenance_cv: Condvar::new(), @@ -156,7 +172,7 @@ impl WarmPool { /// Compute the maintenance action based on current pool size. pub fn compute_maintenance_action(&self, pool_len: usize) -> PoolMaintenanceAction { - let idle_expired = self.decay_fill_target_if_idle(); + let decaying = self.decay_fill_target_if_idle(); let fill_target = self.current_fill_target(); if pool_len < fill_target { let to_fill = fill_target.saturating_sub(pool_len); @@ -167,7 +183,7 @@ impl WarmPool { // After an idle TTL decay the pool shrinks toward the decayed fill // target (the low watermark); otherwise only the high watermark caps // idle resources. - let drain_target = if idle_expired { + let drain_target = if decaying { fill_target } else { self.config.high_watermark @@ -181,6 +197,11 @@ impl WarmPool { return PoolMaintenanceAction::Drain(to_drain); } } + if decaying { + // The pool reached the decayed fill target: the decay cycle is + // complete and the high watermark caps idle resources again. + self.decaying.store(false, Ordering::Release); + } PoolMaintenanceAction::Idle } @@ -189,22 +210,27 @@ impl WarmPool { } /// Collapse the geometric fill target back to the low watermark when no - /// acquisition happened within `idle_ttl`. Returns true when the TTL had - /// expired, so callers drain excess idle resources toward the low - /// watermark instead of the high one. The idle clock restarts on each - /// expiry so a fully decayed pool does not retrigger every cycle. + /// acquisition happened within `idle_ttl`. Returns true while the pool is + /// draining after a decay, so callers drain excess idle resources toward + /// the low watermark instead of the high one. The idle clock restarts on + /// each expiry so a fully decayed pool does not retrigger every cycle, + /// while the `decaying` flag persists until the pool actually reaches the + /// decayed target: a partially-failed drain cycle or an interleaved + /// computation cannot consume the decay event and retain the excess for + /// another full TTL. fn decay_fill_target_if_idle(&self) -> bool { let Some(ttl) = self.config.idle_ttl else { return false; }; let mut last = self.last_acquisition.lock().unwrap(); if last.elapsed() < ttl { - return false; + return self.decaying.load(Ordering::Acquire); } *last = Instant::now(); let low = self.config.low_watermark.min(self.config.high_watermark); let mut target = self.fill_target.lock().unwrap(); *target = (*target).min(low); + self.decaying.store(true, Ordering::Release); true } @@ -232,6 +258,9 @@ impl WarmPool { fn record_acquisition_pressure(&self, pool_len: usize) { *self.last_acquisition.lock().unwrap() = Instant::now(); + // New demand cancels any in-progress idle decay: the fill target may + // grow again and the high watermark caps idle resources. + self.decaying.store(false, Ordering::Release); self.grow_fill_target_after_pressure(pool_len); if matches!( self.compute_maintenance_action(pool_len), @@ -683,6 +712,21 @@ mod tests { assert_eq!(pool.release(42), Err(42)); } + #[test] + fn zero_idle_ttl_is_normalized_to_none() { + // A zero TTL would expire every cycle and busy-loop the maintenance + // worker; the external config documents 0 as "never decay". + let config = PoolConfig { + low_watermark: 0, + high_watermark: 2, + maintenance_enabled: true, + startup_prewarm: false, + idle_ttl: Some(Duration::ZERO), + } + .validate(); + assert_eq!(config.idle_ttl, None); + } + #[test] fn idle_ttl_decays_fill_target_and_drains_to_low_watermark() { let pool = WarmPool::::new(PoolConfig { @@ -690,7 +734,7 @@ mod tests { high_watermark: 10, maintenance_enabled: false, startup_prewarm: false, - idle_ttl: Some(Duration::from_millis(50)), + idle_ttl: Some(Duration::from_millis(200)), }); // Ratchet the fill target up to 8 under acquisition pressure. @@ -707,7 +751,7 @@ mod tests { pool.release(value).unwrap(); } - std::thread::sleep(Duration::from_millis(80)); + std::thread::sleep(Duration::from_millis(400)); // Idle past the TTL: excess drains toward the decayed fill target // (low watermark) instead of lingering below the high watermark. @@ -715,10 +759,28 @@ mod tests { pool.compute_maintenance_action(6), PoolMaintenanceAction::Drain(4) ); + // The decay is not consumed by the first computation: if a drain + // cycle only partially succeeds, the next computation keeps draining + // toward the low watermark instead of retaining the excess for + // another full TTL. + assert_eq!( + pool.compute_maintenance_action(6), + PoolMaintenanceAction::Drain(4) + ); + assert_eq!( + pool.compute_maintenance_action(5), + PoolMaintenanceAction::Drain(3) + ); + // Once the pool reaches the decayed target the decay cycle ends and + // the high watermark caps idle resources again. assert_eq!( pool.compute_maintenance_action(2), PoolMaintenanceAction::Idle ); + assert_eq!( + pool.compute_maintenance_action(6), + PoolMaintenanceAction::Idle + ); assert_eq!( pool.compute_maintenance_action(1), PoolMaintenanceAction::Fill(1) @@ -732,22 +794,23 @@ mod tests { high_watermark: 10, maintenance_enabled: false, startup_prewarm: false, - idle_ttl: Some(Duration::from_millis(300)), + idle_ttl: Some(Duration::from_millis(1000)), }); assert_eq!(pool.try_acquire(), None); std::thread::sleep(Duration::from_millis(100)); assert_eq!(pool.try_acquire(), None); - std::thread::sleep(Duration::from_millis(150)); + std::thread::sleep(Duration::from_millis(200)); - // 150ms since the last acquisition: the 300ms TTL has not expired and - // the ratcheted fill target is preserved. + // 200ms since the last acquisition: the 1000ms TTL has not expired + // (wide margin for loaded CI workers) and the ratcheted fill target + // is preserved. assert_eq!( pool.compute_maintenance_action(0), PoolMaintenanceAction::Fill(8) ); - std::thread::sleep(Duration::from_millis(200)); + std::thread::sleep(Duration::from_millis(1000)); assert_eq!( pool.compute_maintenance_action(0), PoolMaintenanceAction::Fill(2) @@ -779,7 +842,16 @@ mod tests { // No acquisition happens, so the idle TTL wake-up drains everything // (low watermark is 0) without any explicit maintenance request. - std::thread::sleep(Duration::from_millis(300)); + // Poll with a generous deadline instead of asserting after a fixed + // sleep, so CI scheduling delays cannot fail the test. + let deadline = Instant::now() + Duration::from_secs(10); + while pool.len() > 0 { + assert!( + Instant::now() < deadline, + "maintenance worker did not drain the pool after the idle TTL" + ); + std::thread::sleep(Duration::from_millis(20)); + } assert_eq!(pool.len(), 0); pool.drain_all(); From 2d4412f252b433a14fda6f2c65b3353af9964df2 Mon Sep 17 00:00:00 2001 From: epicvinny Date: Mon, 17 Aug 2026 23:23:07 -0300 Subject: [PATCH 3/8] fix(pool): compute idle decay from one synchronized demand state The decay decision was not an atomic snapshot: last_acquisition, decaying, and fill_target lived under separate locks, so a computation concurrent with an acquisition could observe the new acquisition timestamp together with the stale decaying flag and old fill target and return a low-watermark Drain for demand that had just resumed; the unconditional decaying clear could also erase a newer decay transition started by another caller. Move the three fields into a DemandState struct under a single mutex. compute_maintenance_action and record_acquisition each hold the lock for their full transition, so the decay lifecycle, fill target, and idle clock always move together. Tests: acquisition_clears_decay_state_before_next_computation covers the resumed-demand-mid-drain transition deterministically, and concurrent_acquisition_and_decay_keep_state_consistent races four threads of release/acquire/compute against a 5ms TTL, asserting the fill target stays within the watermarks. --- crates/warm-pool/src/lib.rs | 290 +++++++++++++++++++++++------------- 1 file changed, 190 insertions(+), 100 deletions(-) diff --git a/crates/warm-pool/src/lib.rs b/crates/warm-pool/src/lib.rs index 61b59be2..6b5d361f 100644 --- a/crates/warm-pool/src/lib.rs +++ b/crates/warm-pool/src/lib.rs @@ -83,6 +83,107 @@ impl PoolConfig { } } +/// Demand-side pool state: last acquisition time, geometric fill target, and +/// the idle-decay lifecycle. Kept under one mutex so every maintenance action +/// is computed from a single atomic snapshot of demand: a concurrent +/// acquisition can never be observed halfway (new timestamp but stale decay +/// state or fill target). +#[derive(Debug)] +struct DemandState { + /// Last time an acquisition was attempted. Drives idle TTL decay. + last_acquisition: Instant, + /// Current refill target. Starts at the low watermark and grows toward the + /// high watermark under acquisition pressure. With `idle_ttl` unset this + /// intentionally ratchets upward for the process lifetime: after a node + /// observes bursty demand, it keeps extra warm capacity instead of + /// shrinking back to cold-start behavior. When `idle_ttl` is set, a + /// sustained idle period decays the target back to the low watermark so + /// warm capacity (and the resources it holds) is released. + fill_target: usize, + /// Persistent "draining after idle decay" state. Set when the idle TTL + /// expires, cleared by any acquisition, and cleared once the pool has + /// drained to the decayed fill target. Keeps the drain target pinned to + /// the low watermark across partially-failed drain cycles, so the decay + /// event cannot be consumed by a single action computation. + decaying: bool, +} + +impl DemandState { + /// Compute the maintenance action from one synchronized demand snapshot. + fn compute_maintenance_action( + &mut self, + config: &PoolConfig, + pool_len: usize, + ) -> PoolMaintenanceAction { + // Apply the idle TTL decay transition, if any. The idle clock + // restarts on each expiry so a fully decayed pool does not retrigger + // every cycle, while `decaying` persists until the pool actually + // reaches the decayed target: a partially-failed drain cycle or an + // interleaved computation cannot consume the decay event and retain + // the excess for another full TTL. + if let Some(ttl) = config.idle_ttl { + if self.last_acquisition.elapsed() >= ttl { + self.last_acquisition = Instant::now(); + let low = config.low_watermark.min(config.high_watermark); + self.fill_target = self.fill_target.min(low); + self.decaying = true; + } + } + + let fill_target = self.fill_target.min(config.high_watermark); + if pool_len < fill_target { + let to_fill = fill_target.saturating_sub(pool_len); + if to_fill > 0 { + return PoolMaintenanceAction::Fill(to_fill); + } + } + // After an idle TTL decay the pool shrinks toward the decayed fill + // target (the low watermark); otherwise only the high watermark caps + // idle resources. + let drain_target = if self.decaying { + fill_target + } else { + config.high_watermark + }; + if pool_len > drain_target { + // Drain the full excess in one maintenance cycle. Resource-specific + // cleanup happens outside the pool lock, and shutdown paths already + // have to tolerate draining the whole pool. + let to_drain = pool_len - drain_target; + if to_drain > 0 { + return PoolMaintenanceAction::Drain(to_drain); + } + } + if self.decaying { + // The pool reached the decayed fill target: the decay cycle is + // complete and the high watermark caps idle resources again. + self.decaying = false; + } + PoolMaintenanceAction::Idle + } + + /// Record an acquisition attempt: resets the idle clock, cancels any + /// in-progress decay, and grows the fill target geometrically when the + /// pool dipped below the low watermark. + fn record_acquisition(&mut self, config: &PoolConfig, pool_len: usize) { + self.last_acquisition = Instant::now(); + // New demand cancels any in-progress idle decay: the fill target may + // grow again and the high watermark caps idle resources. + self.decaying = false; + if pool_len >= config.low_watermark || config.high_watermark == 0 { + return; + } + + let low = config.low_watermark.min(config.high_watermark); + self.fill_target = self + .fill_target + .max(low) + .max(1) + .saturating_mul(2) + .min(config.high_watermark); + } +} + /// Generic warm pool for reusable resources. /// /// `T` is the pooled resource type. Resource-specific create/reset/delete @@ -103,22 +204,9 @@ pub struct WarmPool { pool: Mutex>, /// Watermark config. config: PoolConfig, - /// Current refill target. Starts at the low watermark and grows toward the - /// high watermark under acquisition pressure. With `idle_ttl` unset this - /// intentionally ratchets upward for the process lifetime: after a node - /// observes bursty demand, it keeps extra warm capacity instead of - /// shrinking back to cold-start behavior. When `idle_ttl` is set, a - /// sustained idle period decays the target back to the low watermark so - /// warm capacity (and the resources it holds) is released. - fill_target: Mutex, - /// Last time an acquisition was attempted. Drives idle TTL decay. - last_acquisition: Mutex, - /// Persistent "draining after idle decay" state. Set when the idle TTL - /// expires, cleared by any acquisition, and cleared once the pool has - /// drained to the decayed fill target. Keeps the drain target pinned to - /// the low watermark across partially-failed drain cycles, so the decay - /// event cannot be consumed by a single action computation. - decaying: AtomicBool, + /// Demand state under a single mutex so decay decisions are atomic + /// snapshots of the last acquisition, fill target, and decay lifecycle. + demand_state: Mutex, /// Background maintenance worker state. maintenance_signal: Mutex, /// Wakes the maintenance worker. @@ -138,9 +226,11 @@ impl WarmPool { let fill_target = config.low_watermark.min(config.high_watermark); Self { pool: Mutex::new(VecDeque::new()), - fill_target: Mutex::new(fill_target), - last_acquisition: Mutex::new(Instant::now()), - decaying: AtomicBool::new(false), + demand_state: Mutex::new(DemandState { + last_acquisition: Instant::now(), + fill_target, + decaying: false, + }), config, maintenance_signal: Mutex::new(PoolMaintenanceSignal::default()), maintenance_cv: Condvar::new(), @@ -172,96 +262,24 @@ impl WarmPool { /// Compute the maintenance action based on current pool size. pub fn compute_maintenance_action(&self, pool_len: usize) -> PoolMaintenanceAction { - let decaying = self.decay_fill_target_if_idle(); - let fill_target = self.current_fill_target(); - if pool_len < fill_target { - let to_fill = fill_target.saturating_sub(pool_len); - if to_fill > 0 { - return PoolMaintenanceAction::Fill(to_fill); - } - } - // After an idle TTL decay the pool shrinks toward the decayed fill - // target (the low watermark); otherwise only the high watermark caps - // idle resources. - let drain_target = if decaying { - fill_target - } else { - self.config.high_watermark - }; - if pool_len > drain_target { - // Drain the full excess in one maintenance cycle. Resource-specific - // cleanup happens outside the pool lock, and shutdown paths already - // have to tolerate draining the whole pool. - let to_drain = pool_len - drain_target; - if to_drain > 0 { - return PoolMaintenanceAction::Drain(to_drain); - } - } - if decaying { - // The pool reached the decayed fill target: the decay cycle is - // complete and the high watermark caps idle resources again. - self.decaying.store(false, Ordering::Release); - } - PoolMaintenanceAction::Idle - } - - fn current_fill_target(&self) -> usize { - (*self.fill_target.lock().unwrap()).min(self.config.high_watermark) - } - - /// Collapse the geometric fill target back to the low watermark when no - /// acquisition happened within `idle_ttl`. Returns true while the pool is - /// draining after a decay, so callers drain excess idle resources toward - /// the low watermark instead of the high one. The idle clock restarts on - /// each expiry so a fully decayed pool does not retrigger every cycle, - /// while the `decaying` flag persists until the pool actually reaches the - /// decayed target: a partially-failed drain cycle or an interleaved - /// computation cannot consume the decay event and retain the excess for - /// another full TTL. - fn decay_fill_target_if_idle(&self) -> bool { - let Some(ttl) = self.config.idle_ttl else { - return false; - }; - let mut last = self.last_acquisition.lock().unwrap(); - if last.elapsed() < ttl { - return self.decaying.load(Ordering::Acquire); - } - *last = Instant::now(); - let low = self.config.low_watermark.min(self.config.high_watermark); - let mut target = self.fill_target.lock().unwrap(); - *target = (*target).min(low); - self.decaying.store(true, Ordering::Release); - true + self.demand_state + .lock() + .unwrap() + .compute_maintenance_action(&self.config, pool_len) } /// Time until the idle TTL expires, if decay is configured. fn idle_ttl_remaining(&self) -> Option { let ttl = self.config.idle_ttl?; - let elapsed = self.last_acquisition.lock().unwrap().elapsed(); + let elapsed = self.demand_state.lock().unwrap().last_acquisition.elapsed(); Some(ttl.saturating_sub(elapsed)) } - fn grow_fill_target_after_pressure(&self, pool_len: usize) { - if pool_len >= self.config.low_watermark || self.config.high_watermark == 0 { - return; - } - - let low = self.config.low_watermark.min(self.config.high_watermark); - let mut target = self.fill_target.lock().unwrap(); - let next = (*target) - .max(low) - .max(1) - .saturating_mul(2) - .min(self.config.high_watermark); - *target = next; - } - fn record_acquisition_pressure(&self, pool_len: usize) { - *self.last_acquisition.lock().unwrap() = Instant::now(); - // New demand cancels any in-progress idle decay: the fill target may - // grow again and the high watermark caps idle resources. - self.decaying.store(false, Ordering::Release); - self.grow_fill_target_after_pressure(pool_len); + self.demand_state + .lock() + .unwrap() + .record_acquisition(&self.config, pool_len); if matches!( self.compute_maintenance_action(pool_len), PoolMaintenanceAction::Fill(_) @@ -856,4 +874,76 @@ mod tests { pool.drain_all(); } + + #[test] + fn acquisition_clears_decay_state_before_next_computation() { + // Regression: an acquisition landing after the idle TTL expiry must + // cancel the drain-to-low transition atomically; a later computation + // must not observe the new acquisition timestamp together with the + // stale decaying state and return a low-watermark Drain for resumed + // demand. + let pool = WarmPool::::new(PoolConfig { + low_watermark: 2, + high_watermark: 10, + maintenance_enabled: false, + startup_prewarm: false, + idle_ttl: Some(Duration::from_millis(200)), + }); + + for value in 1..=6 { + pool.release(value).unwrap(); + } + std::thread::sleep(Duration::from_millis(400)); + + // TTL expired: the decay starts and drains toward the low watermark. + assert_eq!( + pool.compute_maintenance_action(6), + PoolMaintenanceAction::Drain(4) + ); + + // Demand resumes mid-drain: the acquisition resets the idle clock + // and clears the decaying state in the same locked section, so the + // high watermark caps idle resources again. + assert!(pool.try_acquire().is_some()); + assert_eq!( + pool.compute_maintenance_action(5), + PoolMaintenanceAction::Idle + ); + } + + #[test] + fn concurrent_acquisition_and_decay_keep_state_consistent() { + // Regression: acquisitions and action computations race on the + // demand state from multiple threads; the fill target must stay + // within watermarks and no thread may observe a torn state. + let pool = std::sync::Arc::new(WarmPool::::new(PoolConfig { + low_watermark: 2, + high_watermark: 8, + maintenance_enabled: false, + startup_prewarm: false, + idle_ttl: Some(Duration::from_millis(5)), + })); + + let mut handles = Vec::new(); + for _ in 0..4 { + let pool = std::sync::Arc::clone(&pool); + handles.push(std::thread::spawn(move || { + for value in 0..200u32 { + pool.release(value).unwrap(); + let _ = pool.try_acquire(); + let action = pool.compute_maintenance_action(pool.len()); + if let PoolMaintenanceAction::Drain(n) | PoolMaintenanceAction::Fill(n) = action + { + assert!(n <= 8, "action size beyond high watermark: {n}"); + } + } + })); + } + for handle in handles { + handle.join().unwrap(); + } + + let state = pool.demand_state.lock().unwrap(); + assert!(state.fill_target <= 8); + } } From ffdedf54372732ecc869413e87168887070e1516 Mon Sep 17 00:00:00 2001 From: epicvinny Date: Tue, 18 Aug 2026 13:16:44 -0300 Subject: [PATCH 4/8] fix(pool): re-validate maintenance drains against live demand state Address remaining review threads on idle-TTL decay: - A computed Drain action could go stale before execution: an acquisition after the computation clears the decaying state via record_acquisition, but the maintenance cycle still drained the previously returned count toward the low watermark. Maintenance drains now claim each resource through WarmPool::try_drain_one_for_maintenance, which re-checks the drain decision under the demand-state and pool locks at execution time, so resumed demand actually cancels an in-progress decay. - Decay completion was decided from a pool length sampled separately from the demand state, so a concurrent release could strand resources above the decayed target for another full TTL. Completion is now marked only when try_drain_one_for_maintenance verifies, under lock, that the pool actually reached the decayed fill target. - TTL tests no longer rely on wall-clock sleeps: DemandState methods take an injected clock and tests rewind last_acquisition directly; only the maintenance-worker condvar test keeps a bounded polling deadline. --- crates/warm-pool/src/lib.rs | 238 +++++++++++++++++++++++++++----- src/sandbox/firecracker/pool.rs | 5 +- src/sandbox/network/manager.rs | 5 +- 3 files changed, 212 insertions(+), 36 deletions(-) diff --git a/crates/warm-pool/src/lib.rs b/crates/warm-pool/src/lib.rs index 6b5d361f..c0a2efae 100644 --- a/crates/warm-pool/src/lib.rs +++ b/crates/warm-pool/src/lib.rs @@ -102,28 +102,35 @@ struct DemandState { fill_target: usize, /// Persistent "draining after idle decay" state. Set when the idle TTL /// expires, cleared by any acquisition, and cleared once the pool has - /// drained to the decayed fill target. Keeps the drain target pinned to - /// the low watermark across partially-failed drain cycles, so the decay - /// event cannot be consumed by a single action computation. + /// verifiably drained to the decayed fill target (see + /// `WarmPool::try_drain_one_for_maintenance`). Keeps the drain target + /// pinned to the low watermark across partially-failed drain cycles, so + /// the decay event cannot be consumed by a single action computation. decaying: bool, } impl DemandState { /// Compute the maintenance action from one synchronized demand snapshot. + /// + /// The action is a point-in-time plan: callers executing a `Drain` must + /// re-validate each removal with `WarmPool::try_drain_one_for_maintenance` + /// because an interleaved acquisition cancels the decay and a concurrent + /// release/acquire changes the pool length. fn compute_maintenance_action( &mut self, config: &PoolConfig, pool_len: usize, + now: Instant, ) -> PoolMaintenanceAction { // Apply the idle TTL decay transition, if any. The idle clock // restarts on each expiry so a fully decayed pool does not retrigger - // every cycle, while `decaying` persists until the pool actually + // every cycle, while `decaying` persists until the pool verifiably // reaches the decayed target: a partially-failed drain cycle or an // interleaved computation cannot consume the decay event and retain // the excess for another full TTL. if let Some(ttl) = config.idle_ttl { - if self.last_acquisition.elapsed() >= ttl { - self.last_acquisition = Instant::now(); + if now.saturating_duration_since(self.last_acquisition) >= ttl { + self.last_acquisition = now; let low = config.low_watermark.min(config.high_watermark); self.fill_target = self.fill_target.min(low); self.decaying = true; @@ -154,19 +161,14 @@ impl DemandState { return PoolMaintenanceAction::Drain(to_drain); } } - if self.decaying { - // The pool reached the decayed fill target: the decay cycle is - // complete and the high watermark caps idle resources again. - self.decaying = false; - } PoolMaintenanceAction::Idle } /// Record an acquisition attempt: resets the idle clock, cancels any /// in-progress decay, and grows the fill target geometrically when the /// pool dipped below the low watermark. - fn record_acquisition(&mut self, config: &PoolConfig, pool_len: usize) { - self.last_acquisition = Instant::now(); + fn record_acquisition(&mut self, config: &PoolConfig, pool_len: usize, now: Instant) { + self.last_acquisition = now; // New demand cancels any in-progress idle decay: the fill target may // grow again and the high watermark caps idle resources. self.decaying = false; @@ -261,11 +263,19 @@ impl WarmPool { } /// Compute the maintenance action based on current pool size. + /// + /// The returned action is a point-in-time plan computed from `pool_len` + /// and the demand state. A `Drain` plan can go stale before the caller + /// finishes executing it (an interleaved acquisition cancels an + /// in-progress decay; concurrent releases/acquisitions change the pool + /// length), so maintenance drains must claim each resource through + /// `try_drain_one_for_maintenance`, which re-validates the decision at + /// execution time. pub fn compute_maintenance_action(&self, pool_len: usize) -> PoolMaintenanceAction { self.demand_state .lock() .unwrap() - .compute_maintenance_action(&self.config, pool_len) + .compute_maintenance_action(&self.config, pool_len, Instant::now()) } /// Time until the idle TTL expires, if decay is configured. @@ -276,10 +286,11 @@ impl WarmPool { } fn record_acquisition_pressure(&self, pool_len: usize) { + let now = Instant::now(); self.demand_state .lock() .unwrap() - .record_acquisition(&self.config, pool_len); + .record_acquisition(&self.config, pool_len, now); if matches!( self.compute_maintenance_action(pool_len), PoolMaintenanceAction::Fill(_) @@ -359,6 +370,44 @@ impl WarmPool { pool.pop_back() } + /// Drain one idle resource as part of a maintenance drain, re-validating + /// the drain decision against the live demand state and pool length at + /// execution time. + /// + /// A previously computed `Drain` action can be stale by the time the + /// caller executes it: an interleaved acquisition cancels an in-progress + /// decay, and a concurrent release/acquire changes the pool length. + /// Re-checking under the demand-state and pool locks while claiming each + /// resource keeps the drain consistent — decay-driven drains stop as soon + /// as an acquisition cancels the decay, and decay completion is marked + /// only once the pool is verifiably at/below the decayed target, so a + /// stale length snapshot cannot end the cycle early and strand resources. + pub fn try_drain_one_for_maintenance(&self) -> Option { + let mut state = self.demand_state.lock().unwrap(); + let mut pool = self.pool.lock().unwrap(); + let drain_target = if state.decaying { + state.fill_target.min(self.config.high_watermark) + } else { + self.config.high_watermark + }; + if pool.len() <= drain_target { + if state.decaying { + // The pool verifiably reached the decayed fill target: the + // decay cycle is complete and the high watermark caps idle + // resources again. + state.decaying = false; + } + return None; + } + let resource = pool.pop_back(); + if state.decaying && pool.len() <= drain_target { + // This removal brought the pool to the decayed target: the decay + // cycle is complete. + state.decaying = false; + } + resource + } + /// Return a resource to the pool. /// /// If maintenance is enabled, enqueues the resource even when the pool is @@ -745,6 +794,12 @@ mod tests { assert_eq!(config.idle_ttl, None); } + /// Rewind the idle clock so the idle TTL has expired, without sleeping. + fn expire_idle_ttl(pool: &WarmPool) { + let ttl = pool.config().idle_ttl.expect("idle TTL configured"); + pool.demand_state.lock().unwrap().last_acquisition = Instant::now() - ttl; + } + #[test] fn idle_ttl_decays_fill_target_and_drains_to_low_watermark() { let pool = WarmPool::::new(PoolConfig { @@ -752,7 +807,7 @@ mod tests { high_watermark: 10, maintenance_enabled: false, startup_prewarm: false, - idle_ttl: Some(Duration::from_millis(200)), + idle_ttl: Some(Duration::from_secs(60)), }); // Ratchet the fill target up to 8 under acquisition pressure. @@ -769,7 +824,7 @@ mod tests { pool.release(value).unwrap(); } - std::thread::sleep(Duration::from_millis(400)); + expire_idle_ttl(&pool); // Idle past the TTL: excess drains toward the decayed fill target // (low watermark) instead of lingering below the high watermark. @@ -789,12 +844,18 @@ mod tests { pool.compute_maintenance_action(5), PoolMaintenanceAction::Drain(3) ); - // Once the pool reaches the decayed target the decay cycle ends and - // the high watermark caps idle resources again. - assert_eq!( - pool.compute_maintenance_action(2), - PoolMaintenanceAction::Idle - ); + + // Execute the drain through the maintenance claim path: each removal + // is re-validated, and the decay completes only once the pool + // verifiably reaches the decayed target. + for _ in 0..4 { + assert!(pool.try_drain_one_for_maintenance().is_some()); + assert_eq!(pool.demand_state.lock().unwrap().decaying, pool.len() > 2); + } + assert_eq!(pool.len(), 2); + assert!(pool.try_drain_one_for_maintenance().is_none()); + + // The decay cycle ended: the high watermark caps idle resources again. assert_eq!( pool.compute_maintenance_action(6), PoolMaintenanceAction::Idle @@ -812,29 +873,136 @@ mod tests { high_watermark: 10, maintenance_enabled: false, startup_prewarm: false, - idle_ttl: Some(Duration::from_millis(1000)), + idle_ttl: Some(Duration::from_secs(60)), }); assert_eq!(pool.try_acquire(), None); - std::thread::sleep(Duration::from_millis(100)); assert_eq!(pool.try_acquire(), None); - std::thread::sleep(Duration::from_millis(200)); - // 200ms since the last acquisition: the 1000ms TTL has not expired - // (wide margin for loaded CI workers) and the ratcheted fill target - // is preserved. + // Half the TTL has elapsed since the last acquisition: no decay, the + // ratcheted fill target is preserved. + let ttl = pool.config().idle_ttl.unwrap(); + pool.demand_state.lock().unwrap().last_acquisition = Instant::now() - ttl / 2; assert_eq!( pool.compute_maintenance_action(0), PoolMaintenanceAction::Fill(8) ); - std::thread::sleep(Duration::from_millis(1000)); + expire_idle_ttl(&pool); assert_eq!( pool.compute_maintenance_action(0), PoolMaintenanceAction::Fill(2) ); } + #[test] + fn demand_state_idle_ttl_uses_injected_clock() { + let config = PoolConfig { + low_watermark: 2, + high_watermark: 10, + maintenance_enabled: false, + startup_prewarm: false, + idle_ttl: Some(Duration::from_secs(30)), + } + .validate(); + let t0 = Instant::now(); + let mut state = DemandState { + last_acquisition: t0, + fill_target: 8, + decaying: false, + }; + + // Just before expiry: no decay. + assert_eq!( + state.compute_maintenance_action(&config, 8, t0 + Duration::from_secs(29)), + PoolMaintenanceAction::Idle + ); + // At expiry: decay kicks in and drains toward the low watermark. + assert_eq!( + state.compute_maintenance_action(&config, 8, t0 + Duration::from_secs(30)), + PoolMaintenanceAction::Drain(6) + ); + + // An acquisition resets the idle clock and cancels the decay. + state.record_acquisition(&config, 7, t0 + Duration::from_secs(31)); + assert!(!state.decaying); + assert_eq!( + state.compute_maintenance_action(&config, 7, t0 + Duration::from_secs(40)), + PoolMaintenanceAction::Idle + ); + } + + #[test] + fn acquisition_cancels_stale_decay_drain_at_execution_time() { + // Regression: a Drain action computed during decay goes stale when an + // acquisition lands before the caller executes it; each removal must + // be re-validated so resumed demand actually cancels the in-progress + // decay instead of being drained toward the low watermark. + let pool = WarmPool::::new(PoolConfig { + low_watermark: 2, + high_watermark: 10, + maintenance_enabled: false, + startup_prewarm: false, + idle_ttl: Some(Duration::from_secs(60)), + }); + for value in 1..=6 { + pool.release(value).unwrap(); + } + expire_idle_ttl(&pool); + assert_eq!( + pool.compute_maintenance_action(6), + PoolMaintenanceAction::Drain(4) + ); + + // Demand resumes after the computation but before the drain executes. + assert!(pool.try_acquire().is_some()); + + // The stale drain must not remove resources held for resumed demand. + assert!(pool.try_drain_one_for_maintenance().is_none()); + assert_eq!(pool.len(), 5); + } + + #[test] + fn decay_completion_uses_verified_pool_length() { + // Regression: decay completion must be decided from a pool-length + // snapshot synchronized with the pool mutation, not from a length the + // caller sampled earlier. A release landing after the sample raises + // the pool above the decayed target; the decay must stay active until + // the pool verifiably reaches it. + let pool = WarmPool::::new(PoolConfig { + low_watermark: 2, + high_watermark: 10, + maintenance_enabled: false, + startup_prewarm: false, + idle_ttl: Some(Duration::from_secs(60)), + }); + for value in 1..=4 { + pool.release(value).unwrap(); + } + expire_idle_ttl(&pool); + assert_eq!( + pool.compute_maintenance_action(4), + PoolMaintenanceAction::Drain(2) + ); + + // Drain one (pool 3), then a concurrent release raises the pool back + // to 4 — below the high watermark, so nothing re-signals maintenance. + assert!(pool.try_drain_one_for_maintenance().is_some()); + pool.release(5).unwrap(); + + // The decay is still active and drains the extra resource too. + assert!(pool.demand_state.lock().unwrap().decaying); + assert_eq!( + pool.compute_maintenance_action(pool.len()), + PoolMaintenanceAction::Drain(2) + ); + assert!(pool.try_drain_one_for_maintenance().is_some()); + assert!(pool.try_drain_one_for_maintenance().is_some()); + assert!(!pool.demand_state.lock().unwrap().decaying); + assert!(pool.try_drain_one_for_maintenance().is_none()); + assert_eq!(pool.len(), 2); + } + #[test] fn maintenance_worker_wakes_on_idle_ttl_and_drains() { let pool: &'static WarmPool = Box::leak(Box::new(WarmPool::new(PoolConfig { @@ -849,7 +1017,9 @@ mod tests { pool.compute_maintenance_action(pool.len()) { for _ in 0..to_drain { - pool.try_drain_one(); + if pool.try_drain_one_for_maintenance().is_none() { + break; + } } } }); @@ -863,7 +1033,7 @@ mod tests { // Poll with a generous deadline instead of asserting after a fixed // sleep, so CI scheduling delays cannot fail the test. let deadline = Instant::now() + Duration::from_secs(10); - while pool.len() > 0 { + while !pool.is_empty() { assert!( Instant::now() < deadline, "maintenance worker did not drain the pool after the idle TTL" @@ -887,13 +1057,13 @@ mod tests { high_watermark: 10, maintenance_enabled: false, startup_prewarm: false, - idle_ttl: Some(Duration::from_millis(200)), + idle_ttl: Some(Duration::from_secs(60)), }); for value in 1..=6 { pool.release(value).unwrap(); } - std::thread::sleep(Duration::from_millis(400)); + expire_idle_ttl(&pool); // TTL expired: the decay starts and drains toward the low watermark. assert_eq!( diff --git a/src/sandbox/firecracker/pool.rs b/src/sandbox/firecracker/pool.rs index eb0c3f1a..afbc444b 100644 --- a/src/sandbox/firecracker/pool.rs +++ b/src/sandbox/firecracker/pool.rs @@ -234,7 +234,10 @@ impl FirecrackerPool { } PoolMaintenanceAction::Drain(to_drain) => { for _ in 0..to_drain { - let Some(warm) = self.pool.try_drain_one() else { + // Re-validate each removal at execution time: the computed + // drain count can be stale if an acquisition cancelled an + // in-progress decay or the pool length changed. + let Some(warm) = self.pool.try_drain_one_for_maintenance() else { break; }; self.cleanup_warm_blocking(warm, false)?; diff --git a/src/sandbox/network/manager.rs b/src/sandbox/network/manager.rs index ffafd03f..3bd8a6a2 100644 --- a/src/sandbox/network/manager.rs +++ b/src/sandbox/network/manager.rs @@ -327,7 +327,10 @@ impl NetworkManager { } PoolMaintenanceAction::Drain(to_drain) => { for _ in 0..to_drain { - let maybe_slot = self.pool.try_drain_one(); + // Re-validate each removal at execution time: the computed + // drain count can be stale if an acquisition cancelled an + // in-progress decay or the pool length changed. + let maybe_slot = self.pool.try_drain_one_for_maintenance(); let Some(slot) = maybe_slot else { break; }; From cebbd99876f9d7c23f874942de81f906fd9e7464 Mon Sep 17 00:00:00 2001 From: epicvinny Date: Tue, 18 Aug 2026 14:47:10 -0300 Subject: [PATCH 5/8] fix(warm-pool): make acquisition pop and demand record atomic under drain lock order --- crates/warm-pool/src/lib.rs | 33 +++++++++++++++++++++++++-------- 1 file changed, 25 insertions(+), 8 deletions(-) diff --git a/crates/warm-pool/src/lib.rs b/crates/warm-pool/src/lib.rs index c0a2efae..423fd70f 100644 --- a/crates/warm-pool/src/lib.rs +++ b/crates/warm-pool/src/lib.rs @@ -285,12 +285,11 @@ impl WarmPool { Some(ttl.saturating_sub(elapsed)) } - fn record_acquisition_pressure(&self, pool_len: usize) { - let now = Instant::now(); - self.demand_state - .lock() - .unwrap() - .record_acquisition(&self.config, pool_len, now); + /// Wake maintenance when the pool is below its fill target. Called after + /// the acquisition has been recorded under the demand-state lock and + /// both locks have been released (`compute_maintenance_action` re-locks + /// the demand state). + fn request_maintenance_if_fill(&self, pool_len: usize) { if matches!( self.compute_maintenance_action(pool_len), PoolMaintenanceAction::Fill(_) @@ -323,11 +322,19 @@ impl WarmPool { if self.is_shutting_down() { return None; } + // Lock order is demand_state → pool everywhere (see + // `try_drain_one_for_maintenance`): the pop and the acquisition + // record must be atomic, otherwise a maintenance drain can lock + // `demand_state` in the gap, observe a stale `decaying`, and remove + // a resource that resumed demand has already claimed. + let mut state = self.demand_state.lock().unwrap(); let mut pool = self.pool.lock().unwrap(); let resource = pool.pop_front(); let next_pool_len = pool.len(); + state.record_acquisition(&self.config, next_pool_len, Instant::now()); drop(pool); - self.record_acquisition_pressure(next_pool_len); + drop(state); + self.request_maintenance_if_fill(next_pool_len); resource } @@ -338,14 +345,19 @@ impl WarmPool { if self.is_shutting_down() { return None; } + // Same atomic pop-and-record under the demand_state → pool lock + // order as `try_acquire`. + let mut state = self.demand_state.lock().unwrap(); let mut pool = self.pool.lock().unwrap(); let resource = pool .iter() .position(&mut predicate) .and_then(|idx| pool.remove(idx)); let next_pool_len = pool.len(); + state.record_acquisition(&self.config, next_pool_len, Instant::now()); drop(pool); - self.record_acquisition_pressure(next_pool_len); + drop(state); + self.request_maintenance_if_fill(next_pool_len); resource } @@ -382,6 +394,11 @@ impl WarmPool { /// as an acquisition cancels the decay, and decay completion is marked /// only once the pool is verifiably at/below the decayed target, so a /// stale length snapshot cannot end the cycle early and strand resources. + /// + /// The lock order here (demand_state → pool) is the same one acquisition + /// paths use to pop a resource and record the acquisition atomically + /// (`try_acquire`/`try_acquire_where`), so a drain can never observe a + /// half-recorded acquisition (resource popped but `decaying` still set). pub fn try_drain_one_for_maintenance(&self) -> Option { let mut state = self.demand_state.lock().unwrap(); let mut pool = self.pool.lock().unwrap(); From 225c252d52627b4e8f73a7be9cf4dcde05ce9ab3 Mon Sep 17 00:00:00 2001 From: epicvinny Date: Tue, 18 Aug 2026 15:14:23 -0300 Subject: [PATCH 6/8] fix(warm-pool): complete idle decay when pool is already at the decayed target --- crates/warm-pool/src/lib.rs | 41 +++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/crates/warm-pool/src/lib.rs b/crates/warm-pool/src/lib.rs index 423fd70f..91656516 100644 --- a/crates/warm-pool/src/lib.rs +++ b/crates/warm-pool/src/lib.rs @@ -141,9 +141,22 @@ impl DemandState { if pool_len < fill_target { let to_fill = fill_target.saturating_sub(pool_len); if to_fill > 0 { + // Refilling to the decayed target completes the decay + // lifecycle: once the pool is back at the low watermark there + // is nothing left to drain. + self.decaying = false; return PoolMaintenanceAction::Fill(to_fill); } } + if self.decaying && pool_len <= fill_target { + // The pool is already at/below the decayed target (e.g. it was + // empty when the TTL expired): the decay lifecycle completes + // here, because `try_drain_one_for_maintenance` only runs for + // Drain actions and would never clear the flag otherwise — a + // later release would be drained immediately with no new idle + // period. + self.decaying = false; + } // After an idle TTL decay the pool shrinks toward the decayed fill // target (the low watermark); otherwise only the high watermark caps // idle resources. @@ -979,6 +992,34 @@ mod tests { assert_eq!(pool.len(), 5); } + #[test] + fn idle_ttl_decay_completes_without_drain_when_pool_at_decayed_target() { + // Regression: when the TTL expires while the pool is already at/below + // the decayed target, the decay lifecycle must complete without a + // Drain action — `try_drain_one_for_maintenance` only runs for Drain + // actions, so a lingering `decaying` flag would let a later release + // be drained immediately with no new idle period. + let pool = WarmPool::::new(PoolConfig { + low_watermark: 0, + high_watermark: 10, + maintenance_enabled: true, + startup_prewarm: false, + idle_ttl: Some(Duration::from_secs(60)), + }); + expire_idle_ttl(&pool); + assert_eq!( + pool.compute_maintenance_action(0), + PoolMaintenanceAction::Idle + ); + assert!(!pool.demand_state.lock().unwrap().decaying); + + // A fresh release after the completed decay is normal idle capacity, + // not decay excess: maintenance must not drain it. + pool.release(1).unwrap(); + assert!(pool.try_drain_one_for_maintenance().is_none()); + assert_eq!(pool.len(), 1); + } + #[test] fn decay_completion_uses_verified_pool_length() { // Regression: decay completion must be decided from a pool-length From 1fb7b5aa3ba4cba2a358c80604b5ff6d71c19581 Mon Sep 17 00:00:00 2001 From: epicvinny Date: Tue, 18 Aug 2026 15:41:32 -0300 Subject: [PATCH 7/8] fix(warm-pool): compute maintenance actions from the live pool length under both locks --- crates/warm-pool/src/lib.rs | 135 ++++++++++++++++++------------ src/sandbox/firecracker/pool.rs | 2 +- src/sandbox/network/manager.rs | 33 ++++++-- storage/ublk-daemon/src/server.rs | 6 +- 4 files changed, 109 insertions(+), 67 deletions(-) diff --git a/crates/warm-pool/src/lib.rs b/crates/warm-pool/src/lib.rs index 91656516..0a3040ae 100644 --- a/crates/warm-pool/src/lib.rs +++ b/crates/warm-pool/src/lib.rs @@ -275,20 +275,24 @@ impl WarmPool { self.len() == 0 } - /// Compute the maintenance action based on current pool size. + /// Compute the maintenance action from the live pool length. /// - /// The returned action is a point-in-time plan computed from `pool_len` - /// and the demand state. A `Drain` plan can go stale before the caller - /// finishes executing it (an interleaved acquisition cancels an - /// in-progress decay; concurrent releases/acquisitions change the pool - /// length), so maintenance drains must claim each resource through - /// `try_drain_one_for_maintenance`, which re-validates the decision at - /// execution time. - pub fn compute_maintenance_action(&self, pool_len: usize) -> PoolMaintenanceAction { - self.demand_state - .lock() - .unwrap() - .compute_maintenance_action(&self.config, pool_len, Instant::now()) + /// Locks in the crate-wide demand_state → pool order and reads the pool + /// length under the pool lock while holding the demand state, so decay + /// lifecycle transitions (which mutate `decaying`) and the returned plan + /// are always computed from a live, synchronized length — never from a + /// caller-supplied snapshot that may be stale. + /// + /// The returned action is still a point-in-time plan. A `Drain` plan can + /// go stale before the caller finishes executing it (an interleaved + /// acquisition cancels an in-progress decay; concurrent + /// releases/acquisitions change the pool length), so maintenance drains + /// must claim each resource through `try_drain_one_for_maintenance`, + /// which re-validates the decision at execution time. + pub fn compute_maintenance_action(&self) -> PoolMaintenanceAction { + let mut state = self.demand_state.lock().unwrap(); + let pool_len = self.pool.lock().unwrap().len(); + state.compute_maintenance_action(&self.config, pool_len, Instant::now()) } /// Time until the idle TTL expires, if decay is configured. @@ -301,10 +305,10 @@ impl WarmPool { /// Wake maintenance when the pool is below its fill target. Called after /// the acquisition has been recorded under the demand-state lock and /// both locks have been released (`compute_maintenance_action` re-locks - /// the demand state). - fn request_maintenance_if_fill(&self, pool_len: usize) { + /// demand_state → pool). + fn request_maintenance_if_fill(&self) { if matches!( - self.compute_maintenance_action(pool_len), + self.compute_maintenance_action(), PoolMaintenanceAction::Fill(_) ) { self.request_maintenance(); @@ -347,7 +351,7 @@ impl WarmPool { state.record_acquisition(&self.config, next_pool_len, Instant::now()); drop(pool); drop(state); - self.request_maintenance_if_fill(next_pool_len); + self.request_maintenance_if_fill(); resource } @@ -370,7 +374,7 @@ impl WarmPool { state.record_acquisition(&self.config, next_pool_len, Instant::now()); drop(pool); drop(state); - self.request_maintenance_if_fill(next_pool_len); + self.request_maintenance_if_fill(); resource } @@ -551,13 +555,10 @@ impl WarmPool { break; } - has_immediate_work = { - let pool_len = self.pool.lock().unwrap().len(); - !matches!( - self.compute_maintenance_action(pool_len), - PoolMaintenanceAction::Idle - ) - }; + has_immediate_work = !matches!( + self.compute_maintenance_action(), + PoolMaintenanceAction::Idle + ); } } @@ -608,12 +609,19 @@ mod tests { startup_prewarm: false, idle_ttl: None, }); + // Two idle resources below the low watermark: fill the missing two. + pool.release(1).unwrap(); + pool.release(2).unwrap(); assert_eq!( - pool.compute_maintenance_action(2), + pool.compute_maintenance_action(), PoolMaintenanceAction::Fill(2) ); + // Seven idle resources within the watermarks: no action. + for value in 3..=7 { + pool.release(value).unwrap(); + } assert_eq!( - pool.compute_maintenance_action(7), + pool.compute_maintenance_action(), PoolMaintenanceAction::Idle ); } @@ -629,7 +637,7 @@ mod tests { }); assert_eq!( - pool.compute_maintenance_action(0), + pool.compute_maintenance_action(), PoolMaintenanceAction::Fill(2) ); @@ -637,13 +645,13 @@ mod tests { pool.release(2).unwrap(); assert_eq!(pool.try_acquire(), Some(1)); assert_eq!( - pool.compute_maintenance_action(1), + pool.compute_maintenance_action(), PoolMaintenanceAction::Fill(3) ); assert_eq!(pool.try_acquire(), Some(2)); assert_eq!( - pool.compute_maintenance_action(0), + pool.compute_maintenance_action(), PoolMaintenanceAction::Fill(8) ); } @@ -660,19 +668,19 @@ mod tests { assert_eq!(pool.try_acquire(), None); assert_eq!( - pool.compute_maintenance_action(0), + pool.compute_maintenance_action(), PoolMaintenanceAction::Fill(4) ); assert_eq!(pool.try_acquire(), None); assert_eq!( - pool.compute_maintenance_action(0), + pool.compute_maintenance_action(), PoolMaintenanceAction::Fill(8) ); assert_eq!(pool.try_acquire_where(|_| true), None); assert_eq!( - pool.compute_maintenance_action(0), + pool.compute_maintenance_action(), PoolMaintenanceAction::Fill(10) ); } @@ -701,12 +709,20 @@ mod tests { startup_prewarm: false, idle_ttl: None, }); + // Eight idle resources above the high watermark: drain the excess. + for value in 1..=8 { + pool.release(value).unwrap(); + } assert_eq!( - pool.compute_maintenance_action(8), + pool.compute_maintenance_action(), PoolMaintenanceAction::Drain(4) ); + // At the high watermark: no action. + for _ in 0..4 { + assert!(pool.try_drain_one().is_some()); + } assert_eq!( - pool.compute_maintenance_action(4), + pool.compute_maintenance_action(), PoolMaintenanceAction::Idle ); } @@ -846,7 +862,7 @@ mod tests { assert_eq!(pool.try_acquire(), Some(1)); assert_eq!(pool.try_acquire(), Some(2)); assert_eq!( - pool.compute_maintenance_action(0), + pool.compute_maintenance_action(), PoolMaintenanceAction::Fill(8) ); @@ -859,7 +875,7 @@ mod tests { // Idle past the TTL: excess drains toward the decayed fill target // (low watermark) instead of lingering below the high watermark. assert_eq!( - pool.compute_maintenance_action(6), + pool.compute_maintenance_action(), PoolMaintenanceAction::Drain(4) ); // The decay is not consumed by the first computation: if a drain @@ -867,18 +883,21 @@ mod tests { // toward the low watermark instead of retaining the excess for // another full TTL. assert_eq!( - pool.compute_maintenance_action(6), + pool.compute_maintenance_action(), PoolMaintenanceAction::Drain(4) ); + // Simulate a partially successful drain cycle: after one verified + // removal, the computation keeps draining the remaining excess. + assert!(pool.try_drain_one_for_maintenance().is_some()); assert_eq!( - pool.compute_maintenance_action(5), + pool.compute_maintenance_action(), PoolMaintenanceAction::Drain(3) ); - // Execute the drain through the maintenance claim path: each removal - // is re-validated, and the decay completes only once the pool - // verifiably reaches the decayed target. - for _ in 0..4 { + // Execute the rest of the drain through the maintenance claim path: + // each removal is re-validated, and the decay completes only once + // the pool verifiably reaches the decayed target. + for _ in 0..3 { assert!(pool.try_drain_one_for_maintenance().is_some()); assert_eq!(pool.demand_state.lock().unwrap().decaying, pool.len() > 2); } @@ -886,12 +905,18 @@ mod tests { assert!(pool.try_drain_one_for_maintenance().is_none()); // The decay cycle ended: the high watermark caps idle resources again. + for value in 7..=10 { + pool.release(value).unwrap(); + } assert_eq!( - pool.compute_maintenance_action(6), + pool.compute_maintenance_action(), PoolMaintenanceAction::Idle ); + for _ in 0..5 { + assert!(pool.try_drain_one().is_some()); + } assert_eq!( - pool.compute_maintenance_action(1), + pool.compute_maintenance_action(), PoolMaintenanceAction::Fill(1) ); } @@ -914,13 +939,13 @@ mod tests { let ttl = pool.config().idle_ttl.unwrap(); pool.demand_state.lock().unwrap().last_acquisition = Instant::now() - ttl / 2; assert_eq!( - pool.compute_maintenance_action(0), + pool.compute_maintenance_action(), PoolMaintenanceAction::Fill(8) ); expire_idle_ttl(&pool); assert_eq!( - pool.compute_maintenance_action(0), + pool.compute_maintenance_action(), PoolMaintenanceAction::Fill(2) ); } @@ -980,7 +1005,7 @@ mod tests { } expire_idle_ttl(&pool); assert_eq!( - pool.compute_maintenance_action(6), + pool.compute_maintenance_action(), PoolMaintenanceAction::Drain(4) ); @@ -1008,7 +1033,7 @@ mod tests { }); expire_idle_ttl(&pool); assert_eq!( - pool.compute_maintenance_action(0), + pool.compute_maintenance_action(), PoolMaintenanceAction::Idle ); assert!(!pool.demand_state.lock().unwrap().decaying); @@ -1039,7 +1064,7 @@ mod tests { } expire_idle_ttl(&pool); assert_eq!( - pool.compute_maintenance_action(4), + pool.compute_maintenance_action(), PoolMaintenanceAction::Drain(2) ); @@ -1051,7 +1076,7 @@ mod tests { // The decay is still active and drains the extra resource too. assert!(pool.demand_state.lock().unwrap().decaying); assert_eq!( - pool.compute_maintenance_action(pool.len()), + pool.compute_maintenance_action(), PoolMaintenanceAction::Drain(2) ); assert!(pool.try_drain_one_for_maintenance().is_some()); @@ -1072,7 +1097,7 @@ mod tests { }))); pool.start_maintenance_worker(move || { if let PoolMaintenanceAction::Drain(to_drain) = - pool.compute_maintenance_action(pool.len()) + pool.compute_maintenance_action() { for _ in 0..to_drain { if pool.try_drain_one_for_maintenance().is_none() { @@ -1125,7 +1150,7 @@ mod tests { // TTL expired: the decay starts and drains toward the low watermark. assert_eq!( - pool.compute_maintenance_action(6), + pool.compute_maintenance_action(), PoolMaintenanceAction::Drain(4) ); @@ -1134,7 +1159,7 @@ mod tests { // high watermark caps idle resources again. assert!(pool.try_acquire().is_some()); assert_eq!( - pool.compute_maintenance_action(5), + pool.compute_maintenance_action(), PoolMaintenanceAction::Idle ); } @@ -1159,7 +1184,7 @@ mod tests { for value in 0..200u32 { pool.release(value).unwrap(); let _ = pool.try_acquire(); - let action = pool.compute_maintenance_action(pool.len()); + let action = pool.compute_maintenance_action(); if let PoolMaintenanceAction::Drain(n) | PoolMaintenanceAction::Fill(n) = action { assert!(n <= 8, "action size beyond high watermark: {n}"); diff --git a/src/sandbox/firecracker/pool.rs b/src/sandbox/firecracker/pool.rs index afbc444b..020a5fd1 100644 --- a/src/sandbox/firecracker/pool.rs +++ b/src/sandbox/firecracker/pool.rs @@ -228,7 +228,7 @@ impl FirecrackerPool { } fn run_maintenance_cycle(&self) -> Result<()> { - match self.pool.compute_maintenance_action(self.pool.len()) { + match self.pool.compute_maintenance_action() { PoolMaintenanceAction::Fill(to_fill) => { self.runtime.block_on(self.fill_warm_entries(to_fill))?; } diff --git a/src/sandbox/network/manager.rs b/src/sandbox/network/manager.rs index 3bd8a6a2..75fea8fb 100644 --- a/src/sandbox/network/manager.rs +++ b/src/sandbox/network/manager.rs @@ -288,12 +288,12 @@ impl NetworkManager { } #[cfg(test)] - fn compute_maintenance_action(&self, pool_len: usize) -> PoolMaintenanceAction { - self.pool.compute_maintenance_action(pool_len) + fn compute_maintenance_action(&self) -> PoolMaintenanceAction { + self.pool.compute_maintenance_action() } fn run_pool_maintenance_cycle(&self) -> Result<()> { - let action = self.pool.compute_maintenance_action(self.pool.len()); + let action = self.pool.compute_maintenance_action(); match action { PoolMaintenanceAction::Fill(to_fill) => { @@ -901,12 +901,22 @@ mod tests { #[test] fn maintenance_action_fills_to_current_target() { let manager = NetworkManager::new(true, 4, 10); + // Two warm slots below the low watermark: fill the missing two. + for idx in 1..=2u32 { + manager.allocated.insert(idx as usize); + manager.pool.release(test_slot(idx)).unwrap(); + } assert_eq!( - manager.compute_maintenance_action(2), + manager.compute_maintenance_action(), PoolMaintenanceAction::Fill(2) ); + // Seven warm slots within the watermarks: no action. + for idx in 3..=7u32 { + manager.allocated.insert(idx as usize); + manager.pool.release(test_slot(idx)).unwrap(); + } assert_eq!( - manager.compute_maintenance_action(7), + manager.compute_maintenance_action(), PoolMaintenanceAction::Idle ); } @@ -914,12 +924,21 @@ mod tests { #[test] fn maintenance_action_drains_above_high_watermark() { let manager = NetworkManager::new(true, 2, 4); + // Eight warm slots above the high watermark: drain the excess. + for idx in 1..=8u32 { + manager.allocated.insert(idx as usize); + manager.pool.release(test_slot(idx)).unwrap(); + } assert_eq!( - manager.compute_maintenance_action(8), + manager.compute_maintenance_action(), PoolMaintenanceAction::Drain(4) ); + // At the high watermark: no action. + for _ in 0..4 { + assert!(manager.pool.try_drain_one().is_some()); + } assert_eq!( - manager.compute_maintenance_action(4), + manager.compute_maintenance_action(), PoolMaintenanceAction::Idle ); } diff --git a/storage/ublk-daemon/src/server.rs b/storage/ublk-daemon/src/server.rs index d9542d49..97801eac 100644 --- a/storage/ublk-daemon/src/server.rs +++ b/storage/ublk-daemon/src/server.rs @@ -1442,7 +1442,7 @@ fn schedule_idle_pool_refill( pool.refill_inflight.store(false, Ordering::Release); if matches!( - pool.idle.compute_maintenance_action(pool.idle.len()), + pool.idle.compute_maintenance_action(), PoolMaintenanceAction::Fill(_) ) { schedule_idle_pool_refill(pool, ctrl_ring, virtual_size); @@ -1455,9 +1455,7 @@ async fn refill_idle_pool( ctrl_ring: IoRingHandle, virtual_size: u64, ) -> Result<()> { - let PoolMaintenanceAction::Fill(to_create) = - pool.idle.compute_maintenance_action(pool.idle.len()) - else { + let PoolMaintenanceAction::Fill(to_create) = pool.idle.compute_maintenance_action() else { return Ok(()); }; From 5acb28388e6933e76a66c35d7c330b9972d8f9d8 Mon Sep 17 00:00:00 2001 From: epicvinny Date: Tue, 18 Aug 2026 15:53:30 -0300 Subject: [PATCH 8/8] fix(warm-pool): complete idle decay only on a verified at-target pool length --- crates/warm-pool/src/lib.rs | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/crates/warm-pool/src/lib.rs b/crates/warm-pool/src/lib.rs index 0a3040ae..8dc86f54 100644 --- a/crates/warm-pool/src/lib.rs +++ b/crates/warm-pool/src/lib.rs @@ -138,25 +138,25 @@ impl DemandState { } let fill_target = self.fill_target.min(config.high_watermark); + if self.decaying && pool_len <= fill_target { + // The live pool (this length is always read under the pool lock + // by `WarmPool::compute_maintenance_action`) is verifiably at or + // below the decayed target — including when it was empty when + // the TTL expired and no Drain action ever ran. The decay + // lifecycle completes only on this synchronized observation, + // never on a merely PLANNED refill: if the fill then fails or a + // concurrent release raises the pool again, `decaying` is + // already false only because the pool genuinely reached the + // target, and resources between the low and high watermarks are + // legitimately retained by the high watermark. + self.decaying = false; + } if pool_len < fill_target { let to_fill = fill_target.saturating_sub(pool_len); if to_fill > 0 { - // Refilling to the decayed target completes the decay - // lifecycle: once the pool is back at the low watermark there - // is nothing left to drain. - self.decaying = false; return PoolMaintenanceAction::Fill(to_fill); } } - if self.decaying && pool_len <= fill_target { - // The pool is already at/below the decayed target (e.g. it was - // empty when the TTL expired): the decay lifecycle completes - // here, because `try_drain_one_for_maintenance` only runs for - // Drain actions and would never clear the flag otherwise — a - // later release would be drained immediately with no new idle - // period. - self.decaying = false; - } // After an idle TTL decay the pool shrinks toward the decayed fill // target (the low watermark); otherwise only the high watermark caps // idle resources.