fix: gracefullly shutdown runners (#47)

* fix: gracefullly shutdown runners

Use `Store::shutdown_runners()` to gracefully shutdown runners.`

Signed-off-by: MrCroxx <[email protected]>

* log warn if send failed instead of panic

Signed-off-by: MrCroxx <[email protected]>

---------

Signed-off-by: MrCroxx <[email protected]>

Author: MrCroxx

foyer-storage-bench/src/main.rs

@@ -260,6 +260,8 @@ async fn main() {
         &metrics_dump_end,
     );
     println!("\nTotal:\n{}", analysis);
+
+    store.shutdown_runners().await.unwrap();
 }
 
 async fn bench(args: Args, store: Arc<TStore>, metrics: Metrics, stop: oneshot::Receiver<()>) {

foyer-storage/src/flusher.rs

@@ -17,9 +17,9 @@ use std::sync::Arc;
 use foyer_common::queue::AsyncQueue;
 use foyer_intrusive::{core::adapter::Link, eviction::EvictionPolicy};
 use itertools::Itertools;
-use tokio::sync::{
-    mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender},
-    Mutex,
+use tokio::{
+    sync::{broadcast, mpsc, Mutex},
+    task::JoinHandle,
 };
 
 use crate::{
@@ -38,7 +38,7 @@ pub struct FlushTask {
 struct FlusherInner {
     sequence: usize,
 
-    task_txs: Vec<UnboundedSender<FlushTask>>,
+    task_txs: Vec<mpsc::UnboundedSender<FlushTask>>,
 }
 
 pub struct Flusher {
@@ -63,7 +63,9 @@ impl Flusher {
         &self,
         buffers: Arc<AsyncQueue<Vec<u8, A>>>,
         region_manager: Arc<RegionManager<A, D, E, EL>>,
-    ) where
+        stop_rxs: Vec<broadcast::Receiver<()>>,
+    ) -> Vec<JoinHandle<()>>
+    where
         A: BufferAllocator,
         D: Device<IoBufferAllocator = A>,
         E: EvictionPolicy<RegionEpItemAdapter<EL>, Link = EL>,
@@ -73,25 +75,30 @@ impl Flusher {
 
         #[allow(clippy::type_complexity)]
         let (mut txs, rxs): (
-            Vec<UnboundedSender<FlushTask>>,
-            Vec<UnboundedReceiver<FlushTask>>,
-        ) = (0..self.runners).map(|_| unbounded_channel()).unzip();
+            Vec<mpsc::UnboundedSender<FlushTask>>,
+            Vec<mpsc::UnboundedReceiver<FlushTask>>,
+        ) = (0..self.runners).map(|_| mpsc::unbounded_channel()).unzip();
         inner.task_txs.append(&mut txs);
 
         let runners = rxs
             .into_iter()
-            .map(|rx| Runner {
-                task_rx: rx,
+            .zip_eq(stop_rxs.into_iter())
+            .map(|(task_rx, stop_rx)| Runner {
+                task_rx,
                 buffers: buffers.clone(),
                 region_manager: region_manager.clone(),
+                stop_rx,
             })
             .collect_vec();
 
+        let mut handles = vec![];
         for runner in runners {
-            tokio::spawn(async move {
+            let handle = tokio::spawn(async move {
                 runner.run().await.unwrap();
             });
+            handles.push(handle);
         }
+        handles
     }
 
     pub fn runners(&self) -> usize {
@@ -113,10 +120,12 @@ where
     E: EvictionPolicy<RegionEpItemAdapter<EL>, Link = EL>,
     EL: Link,
 {
-    task_rx: UnboundedReceiver<FlushTask>,
+    task_rx: mpsc::UnboundedReceiver<FlushTask>,
     buffers: Arc<AsyncQueue<Vec<u8, A>>>,
 
     region_manager: Arc<RegionManager<A, D, E, EL>>,
+
+    stop_rx: broadcast::Receiver<()>,
 }
 
 impl<A, D, E, EL> Runner<A, D, E, EL>
@@ -128,71 +137,79 @@ where
 {
     async fn run(mut self) -> Result<()> {
         loop {
-            if let Some(task) = self.task_rx.recv().await {
-                // TODO(MrCroxx): seal buffer
+            tokio::select! {
+                Some(task) = self.task_rx.recv() => {
+                    self.handle(task).await?;
+                }
+                _ = self.stop_rx.recv() => {
+                    tracing::info!("[flusher] exit");
+                    return Ok(())
+                }
+            }
+        }
+    }
 
-                tracing::info!("[flusher] receive flush task, region: {}", task.region_id);
+    async fn handle(&self, task: FlushTask) -> Result<()> {
+        tracing::info!("[flusher] receive flush task, region: {}", task.region_id);
 
-                let region = self.region_manager.region(&task.region_id);
+        let region = self.region_manager.region(&task.region_id);
 
-                tracing::trace!("[flusher] step 1");
+        tracing::trace!("[flusher] step 1");
 
-                {
-                    // step 1: write buffer back to device
-                    let slice = region.load(.., 0).await?.unwrap();
+        {
+            // step 1: write buffer back to device
+            let slice = region.load(.., 0).await?.unwrap();
 
-                    // wait all physical readers (from previous version) and writers done
-                    let guard = region.exclusive(false, true, false).await;
+            // wait all physical readers (from previous version) and writers done
+            let guard = region.exclusive(false, true, false).await;
 
-                    tracing::trace!("[flusher] write region {} back to device", task.region_id);
+            tracing::trace!("[flusher] write region {} back to device", task.region_id);
 
-                    let mut offset = 0;
-                    let len = region.device().io_size();
-                    while offset < region.device().region_size() {
-                        let start = offset;
-                        let end = std::cmp::min(offset + len, region.device().region_size());
+            let mut offset = 0;
+            let len = region.device().io_size();
+            while offset < region.device().region_size() {
+                let start = offset;
+                let end = std::cmp::min(offset + len, region.device().region_size());
 
-                        let s = unsafe { Slice::new(&slice.as_ref()[start..end]) };
-                        region
-                            .device()
-                            .write(s, region.id(), offset as u64, len)
-                            .await?;
-                        offset += len;
-                    }
-                    drop(guard);
-                    slice.destroy().await;
-                }
+                let s = unsafe { Slice::new(&slice.as_ref()[start..end]) };
+                region
+                    .device()
+                    .write(s, region.id(), offset as u64, len)
+                    .await?;
+                offset += len;
+            }
+            drop(guard);
+            slice.destroy().await;
+        }
+
+        tracing::trace!("[flusher] step 2");
 
-                tracing::trace!("[flusher] step 2");
+        let buffer = {
+            // step 2: detach buffer
+            let mut guard = region.exclusive(false, false, true).await;
 
-                let buffer = {
-                    // step 2: detach buffer
-                    let mut guard = region.exclusive(false, false, true).await;
+            let buffer = guard.detach_buffer();
 
-                    let buffer = guard.detach_buffer();
+            tracing::trace!(
+                "[flusher] region {}, writers: {}, buffered readers: {}, physical readers: {}",
+                region.id(),
+                guard.writers(),
+                guard.buffered_readers(),
+                guard.physical_readers()
+            );
 
-                    tracing::trace!(
-                            "[flusher] region {}, writers: {}, buffered readers: {}, physical readers: {}",
-                            region.id(),
-                            guard.writers(),
-                            guard.buffered_readers(),
-                            guard.physical_readers()
-                        );
+            drop(guard);
+            buffer
+        };
 
-                    drop(guard);
-                    buffer
-                };
+        tracing::trace!("[flusher] step 3");
 
-                tracing::trace!("[flusher] step 3");
+        // step 3: release buffer
+        self.buffers.release(buffer);
+        self.region_manager.set_region_evictable(&region.id()).await;
 
-                // step 3: release buffer
-                self.buffers.release(buffer);
-                self.region_manager.set_region_evictable(&region.id()).await;
+        tracing::info!("[flusher] finish flush task, region: {}", task.region_id);
 
-                tracing::info!("[flusher] finish flush task, region: {}", task.region_id);
-            } else {
-                return Ok(());
-            }
-        }
+        Ok(())
     }
 }

foyer-storage/src/lib.rs

@@ -16,6 +16,7 @@
 #![feature(strict_provenance)]
 #![feature(trait_alias)]
 #![feature(get_mut_unchecked)]
+#![feature(let_chains)]
 #![allow(clippy::type_complexity)]
 
 use device::io_buffer::AlignedAllocator;