diff --git a/foyer-common/src/spawn.rs b/foyer-common/src/spawn.rs index 3b52e389..b7a1ed58 100644 --- a/foyer-common/src/spawn.rs +++ b/foyer-common/src/spawn.rs @@ -14,9 +14,8 @@ use std::{ fmt::Debug, - mem::ManuallyDrop, ops::{Deref, DerefMut}, - sync::Arc, + sync::{Arc, Mutex}, }; use tokio::{ @@ -26,46 +25,73 @@ use tokio::{ use crate::error::{Error, ErrorKind, Result}; -/// A wrapper around [`Runtime`] that shuts down the runtime in the background when dropped. +/// A wrapper around a dedicated tokio [`Runtime`] that can shut the runtime down in the background. /// -/// This is necessary because directly dropping a nested runtime is not allowed in a parent runtime. -pub struct BackgroundShutdownRuntime(ManuallyDrop); +/// The runtime is shared by [`Arc`], so dropping the *last* clone does **not** deterministically run +/// the shutdown: tasks spawned on the runtime may keep a clone alive themselves (e.g. via a captured +/// [`Spawner`]), forming a self-cycle that prevents the strong count from ever reaching zero. To +/// break such a cycle, call [`BackgroundShutdownRuntime::shutdown`] explicitly; the wrapper also +/// defers the shutdown to a background step when the last clone is eventually dropped. +/// +/// This is necessary because directly dropping a nested runtime is not allowed in a parent runtime; +/// [`BackgroundShutdownRuntime::shutdown`] (like [`Runtime::shutdown_background`]) may be called from +/// within the runtime itself or a parent runtime. +pub struct BackgroundShutdownRuntime { + runtime: Mutex>, + handle: Handle, +} impl Debug for BackgroundShutdownRuntime { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_tuple("BackgroundShutdownRuntime").finish() + f.debug_struct("BackgroundShutdownRuntime").finish() + } +} + +impl BackgroundShutdownRuntime { + /// Deterministically shut the dedicated runtime down in the background. + /// + /// This is idempotent: the first call initiates the shutdown and subsequent calls are no-ops. + /// + /// This may be called from within the runtime itself or from a parent runtime. After this + /// returns, the runtime is being torn down in the background: all spawned tasks are dropped + /// (releasing any [`Arc`] clones they captured) and the worker threads exit. + pub fn shutdown(&self) { + if let Some(runtime) = self.runtime.lock().unwrap().take() { + #[cfg(madsim)] + drop(runtime); + #[cfg(not(madsim))] + runtime.shutdown_background(); + } } } impl Drop for BackgroundShutdownRuntime { fn drop(&mut self) { - // Safety: The runtime is only dropped once here. - let runtime = unsafe { ManuallyDrop::take(&mut self.0) }; - - #[cfg(madsim)] - drop(runtime); - #[cfg(not(madsim))] - runtime.shutdown_background(); + self.shutdown(); } } impl Deref for BackgroundShutdownRuntime { - type Target = Runtime; + type Target = Handle; fn deref(&self) -> &Self::Target { - &self.0 + &self.handle } } impl DerefMut for BackgroundShutdownRuntime { fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.0 + &mut self.handle } } impl From for BackgroundShutdownRuntime { fn from(runtime: Runtime) -> Self { - Self(ManuallyDrop::new(runtime)) + let handle = runtime.handle().clone(); + Self { + runtime: Mutex::new(Some(runtime)), + handle, + } } } @@ -141,4 +167,22 @@ impl Spawner { pub fn current() -> Self { Spawner::Handle(Handle::current()) } + + /// Deterministically shut down the dedicated runtime, if this is a [`Spawner::Runtime`]. + /// + /// This breaks the self-cycle where tasks spawned on the dedicated runtime keep a clone of the + /// spawner (and thus the runtime) alive, preventing the runtime from ever being dropped. After + /// this call the runtime is torn down in the background; all spawned tasks are dropped, + /// releasing any spawner clones they captured. + /// + /// For [`Spawner::Handle`], there is no owned runtime to shut down (the handle references an + /// externally-owned runtime whose lifetime the caller controls), so this is a no-op. + /// + /// This is safe to call from within the runtime itself or a parent runtime, and is idempotent. + pub fn shutdown(&self) { + match self { + Spawner::Runtime(rt) => rt.shutdown(), + Spawner::Handle(_) => {} + } + } } diff --git a/foyer-storage/Cargo.toml b/foyer-storage/Cargo.toml index fbe2c081..ab60f233 100644 --- a/foyer-storage/Cargo.toml +++ b/foyer-storage/Cargo.toml @@ -61,7 +61,7 @@ bytesize = { workspace = true } foyer-memory = { workspace = true, features = ["test_utils"] } tempfile = { workspace = true } test-log = { workspace = true, features = ["trace", "color"] } -tokio = { workspace = true } +tokio = { workspace = true, features = ["rt-multi-thread", "macros"] } [lints] workspace = true diff --git a/foyer-storage/src/engine/block/engine.rs b/foyer-storage/src/engine/block/engine.rs index 1b9d19d8..ac440852 100644 --- a/foyer-storage/src/engine/block/engine.rs +++ b/foyer-storage/src/engine/block/engine.rs @@ -530,6 +530,25 @@ where load_holder: Holder, } +impl Drop for BlockEngineInner +where + K: StorageKey, + V: StorageValue, + P: Properties, +{ + fn drop(&mut self) { + // Tear down the dedicated runtime, if any. + // + // The flusher/reclaimer workers are spawned *on* the dedicated `Spawner::Runtime` and each + // captures a clone of the spawner (`Arc`). This forms a self-cycle + // — `Arc -> Runtime -> spawned task -> Spawner (Arc)` — that keeps the strong count at or + // above one forever, so the last-`Arc`-drop shutdown would never fire on its own. Forcing the + // shutdown here drops the spawned tasks (releasing their spawner clones) and lets the + // runtime be reclaimed when the store is dropped. For `Spawner::Handle` this is a no-op. + self._spawner.shutdown(); + } +} + impl Clone for BlockEngine where K: StorageKey, @@ -847,10 +866,10 @@ where #[cfg(test)] mod tests { - use std::{fs::File, path::Path}; + use std::{fs::File, path::Path, time::Duration}; use bytesize::ByteSize; - use foyer_common::hasher::ModHasher; + use foyer_common::{hasher::ModHasher, spawn::BackgroundShutdownRuntime}; use foyer_memory::{Cache, CacheBuilder, CacheEntry, FifoConfig, TestProperties}; use itertools::Itertools; @@ -1403,4 +1422,76 @@ mod tests { .unwrap(); assert_eq!(engine.inner.block_manager.blocks(), (1 + 2 + 4) * MB / (64 * KB)); } + + #[cfg(not(madsim))] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_dedicated_runtime_released_on_engine_drop() { + // The block engine spawns its flusher/reclaimer workers on the dedicated `Spawner::Runtime`, + // and each worker captures a clone of the spawner (`Arc`). This + // formed a self-cycle that prevented the runtime from ever being dropped. Dropping the engine + // must tear the dedicated runtime down so the `Arc` is released. + let dir = tempfile::tempdir().unwrap(); + let dedicated = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .thread_name("foyer-eng-rt") + .enable_all() + .build() + .unwrap(); + let bsr: Arc = Arc::new(dedicated.into()); + let weak = Arc::downgrade(&bsr); + let spawner = Spawner::Runtime(bsr); + let io_engine = io_engine_for_test(spawner.clone()).await; + let device = FsDeviceBuilder::new(dir.path()) + .with_capacity(ByteSize::kib(256).as_u64() as _) + .build() + .unwrap(); + let builder: BlockEngineConfig, TestProperties> = BlockEngineConfig { + device, + block_size: 16 * 1024, + compression: Compression::None, + indexer_shards: 4, + recover_concurrency: 2, + flushers: 2, + reclaimers: 1, + clean_block_threshold: 1, + admission_filter: StorageFilter::new(), + eviction_pickers: vec![Box::::default()], + reinsertion_filter: StorageFilter::new().with_condition(RejectAll), + enable_tombstone_log: false, + buffer_pool_size: 16 * 1024 * 1024, + blob_index_size: 4 * 1024, + submit_queue_size_threshold: 16 * 1024 * 1024 * 2, + flush_switch: Switch::default(), + load_holder: Holder::default(), + marker: PhantomData, + }; + let engine = Box::new(builder) + .build(EngineBuildContext { + io_engine, + metrics: Arc::new(Metrics::noop()), + spawner, + recover_mode: RecoverMode::Strict, + }) + .await + .unwrap(); + // The flusher/reclaimer workers are spawned during `build()` and hold `Arc` clones, but + // the only *external* strong ref is the one we moved into the engine. + assert!( + weak.upgrade().is_some(), + "runtime must be alive while the engine is alive" + ); + + drop(engine); + + // Poll until the background runtime shutdown releases the last `Arc` clone. Before the + // fix, this never happened and the runtime leaked permanently. + let deadline = Instant::now() + Duration::from_secs(5); + while Instant::now() < deadline && weak.upgrade().is_some() { + std::thread::sleep(Duration::from_millis(10)); + } + assert!( + weak.upgrade().is_none(), + "Arc was not released after drop(engine) (dedicated runtime leaked)" + ); + } } diff --git a/foyer-storage/src/store.rs b/foyer-storage/src/store.rs index 4254cada..4f047c04 100644 --- a/foyer-storage/src/store.rs +++ b/foyer-storage/src/store.rs @@ -509,8 +509,10 @@ where #[cfg(test)] mod tests { - use foyer_common::hasher::ModHasher; - use foyer_memory::CacheBuilder; + use std::time::{Duration, Instant}; + + use foyer_common::{code::DefaultHasher, hasher::ModHasher, spawn::BackgroundShutdownRuntime}; + use foyer_memory::{CacheBuilder, CacheProperties}; use super::*; use crate::{ @@ -716,4 +718,198 @@ mod tests { let l2 = store.load(&1).await.unwrap(); assert!(matches!(l2, Load::Miss)); } + + // Regression tests for the dedicated-runtime leak. Dropping a `Store` built with a dedicated + // `Spawner::Runtime` must release the `Arc` so the runtime is shut + // down in the background (worker threads exit). Before the fix, a self-cycle + // (`Arc -> Runtime -> spawned task -> Spawner(Arc)`) kept the strong count above zero + // forever, so `BackgroundShutdownRuntime::drop` / `shutdown_background()` never ran. + + /// Count OS threads whose `comm` starts with `prefix`. Linux-only probe; returns 0 elsewhere. + #[cfg(not(madsim))] + fn count_threads_named(prefix: &str) -> usize { + #[cfg(target_os = "linux")] + { + std::fs::read_dir("/proc/self/task") + .map(|rd| { + rd.filter_map(|e| e.ok()) + .filter_map(|e| std::fs::read_to_string(e.path().join("comm")).ok()) + .filter(|s| s.trim().starts_with(prefix)) + .count() + }) + .unwrap_or(0) + } + #[cfg(not(target_os = "linux"))] + { + let _ = prefix; + 0 + } + } + + /// Poll until the given weak ref can no longer be upgraded, or `timeout` elapses. + #[cfg(not(madsim))] + fn wait_until_weak_dropped(weak: &std::sync::Weak, timeout: Duration) -> bool { + let deadline = Instant::now() + timeout; + while Instant::now() < deadline { + if weak.upgrade().is_none() { + return true; + } + std::thread::sleep(Duration::from_millis(10)); + } + weak.upgrade().is_none() + } + + /// Poll until no OS thread whose `comm` starts with `prefix` remains, or `timeout` elapses. + #[cfg(not(madsim))] + fn wait_until_threads_gone(prefix: &str, timeout: Duration) -> bool { + let deadline = Instant::now() + timeout; + while Instant::now() < deadline { + if count_threads_named(prefix) == 0 { + return true; + } + std::thread::sleep(Duration::from_millis(10)); + } + count_threads_named(prefix) == 0 + } + + /// Build a dedicated-runtime `Store` for the leak tests, returning the store and a weak ref to the + /// runtime. The only strong ref to the `Arc` is moved into the store. + #[cfg(not(madsim))] + async fn build_dedicated_runtime_store( + temp: &tempfile::TempDir, + thread_name: &str, + memory: Cache>, + ) -> ( + Store, DefaultHasher, CacheProperties>, + std::sync::Weak, + ) { + let metrics = Arc::new(Metrics::noop()); + let dedicated = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .thread_name(thread_name) + .enable_all() + .build() + .unwrap(); + let bsr: Arc = Arc::new(dedicated.into()); + let weak = Arc::downgrade(&bsr); + let store = StoreBuilder::new("test", memory, metrics) + .with_io_engine_config(PsyncIoEngineConfig::new()) + .with_engine_config( + BlockEngineConfig::new( + FsDeviceBuilder::new(temp.path()) + .with_capacity(256 * 1024) + .build() + .unwrap(), + ) + .with_flushers(2) + .with_reclaimers(1) + .with_block_size(16 * 1024), + ) + .with_spawner(Spawner::Runtime(bsr)) + .build() + .await + .unwrap(); + (store, weak) + } + + #[cfg(not(madsim))] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_dedicated_runtime_released_on_drop() { + let dir = tempfile::tempdir().unwrap(); + let memory: Cache> = CacheBuilder::new(16).build(); + let (store, weak) = build_dedicated_runtime_store(&dir, "foyer-drop-rt", memory.clone()).await; + + let e = memory.insert(1, b"v1".to_vec()); + store.enqueue(e.piece(), false); + store.wait().await; + + // The dedicated runtime's eager worker threads must be alive while the store exists. + // Thread-name probing via /proc/self/task is Linux-only; skip on other platforms. + #[cfg(target_os = "linux")] + assert!( + count_threads_named("foyer-drop-rt") >= 2, + "dedicated runtime worker threads should exist while store is alive" + ); + assert!(weak.upgrade().is_some(), "runtime must be alive while store is alive"); + + drop(store); + drop(memory); + + // After dropping the store (with no external strong ref to the runtime), the runtime must be + // torn down: the `Arc` must reach zero and the worker threads must + // exit. Before the fix, `weak.upgrade()` stayed `Some` forever. + assert!( + wait_until_weak_dropped(&weak, Duration::from_secs(5)), + "Arc was not released after drop(store) (dedicated runtime leaked)" + ); + assert!( + wait_until_threads_gone("foyer-drop-rt", Duration::from_secs(5)), + "dedicated runtime worker threads should exit after drop(store)" + ); + } + + #[cfg(not(madsim))] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_dedicated_runtime_released_on_close_then_drop() { + let dir = tempfile::tempdir().unwrap(); + let memory: Cache> = CacheBuilder::new(16).build(); + let (store, weak) = build_dedicated_runtime_store(&dir, "foyer-close-rt", memory.clone()).await; + + let e = memory.insert(1, b"v1".to_vec()); + store.enqueue(e.piece(), false); + store.wait().await; + + store.close().await.unwrap(); + assert!( + weak.upgrade().is_some(), + "runtime must stay alive while the store is alive, even after close()" + ); + + drop(store); + drop(memory); + + assert!( + wait_until_weak_dropped(&weak, Duration::from_secs(5)), + "Arc was not released after close()+drop() (dedicated runtime leaked)" + ); + assert!( + wait_until_threads_gone("foyer-close-rt", Duration::from_secs(5)), + "dedicated runtime worker threads should exit after close()+drop()" + ); + } + + #[cfg(not(madsim))] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_dedicated_runtime_no_accumulation() { + const N: usize = 4; + let mut dirs = vec![]; + let mut weaks = vec![]; + let mut names = vec![]; + for i in 0..N { + dirs.push(tempfile::tempdir().unwrap()); + let name = format!("foyer-accum-rt-{i}"); + names.push(name.clone()); + let memory: Cache> = CacheBuilder::new(16).build(); + let (store, weak) = build_dedicated_runtime_store(&dirs[i], &name, memory.clone()).await; + let e = memory.insert(i as u64, vec![i as u8; 1024]); + store.enqueue(e.piece(), false); + store.wait().await; + drop(store); + drop(memory); + weaks.push(weak); + } + // Every dropped runtime must be fully torn down — no accumulation across caches. Before the + // fix, all `N` dropped runtimes stayed alive (one leaked runtime per cache), so the strong + // count and the per-runtime worker threads never went away. + for (i, weak) in weaks.iter().enumerate() { + assert!( + wait_until_weak_dropped(weak, Duration::from_secs(5)), + "Arc for iteration {i} was not released; runtimes accumulate" + ); + assert!( + wait_until_threads_gone(&names[i], Duration::from_secs(5)), + "dedicated runtime worker threads for iteration {i} should exit" + ); + } + } }