fix(storage): shut down dedicated runtime on engine drop to break self-cycle leak

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<Runtime>);
+/// 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<Option<Runtime>>,
+    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<Runtime> 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(_) => {}
+        }
+    }
 }

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

foyer-storage/src/engine/block/engine.rs

@@ -530,6 +530,25 @@ where
     load_holder: Holder,
 }
 
+impl<K, V, P> Drop for BlockEngineInner<K, V, P>
+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<BackgroundShutdownRuntime>`). 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<K, V, P> Clone for BlockEngine<K, V, P>
 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<BackgroundShutdownRuntime>`). 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<BackgroundShutdownRuntime>` 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<BackgroundShutdownRuntime> = 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<u64, Vec<u8>, 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::<FifoPicker>::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<BSR>` 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<BSR>` 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<BackgroundShutdownRuntime> was not released after drop(engine) (dedicated runtime leaked)"
+        );
+    }
 }