refactor: add simpler shutdown handling

Cargo.toml

@@ -20,6 +20,7 @@ members = [
     "peer-svc",
     "pipeline",
     "recon",
+    "shutdown",
     "sql",
     "validation",
     "beetle/iroh-bitswap",
@@ -182,6 +183,7 @@ serde_qs = "0.10.1"
 serde_with = "2.1"
 sha2 = { version = "0.10", default-features = false }
 sha3 = "0.10"
+shutdown = { path = "./shutdown/" }
 smallvec = "1.10"
 # pragma optimize hangs forver on 0.8, possibly due to libsqlite-sys upgrade
 sqlx = { version = "0.7", features = ["sqlite", "runtime-tokio", "chrono"] }

anchor-service/Cargo.toml

@@ -29,3 +29,6 @@ chrono.workspace = true
 
 [features]
 test-network = []
+
+[dev-dependencies]
+shutdown.workspace = true

anchor-service/src/anchor_batch.rs

@@ -72,7 +72,7 @@ impl AnchorService {
     /// - Store the TimeEvents using the AnchorClient
     ///
     /// This function will run indefinitely, or until the process is shutdown.
-    pub async fn run(&mut self, shutdown_signal: impl Future<Output = ()>) {
+    pub async fn run(mut self, shutdown_signal: impl Future<Output = ()>) {
         let shutdown_signal = shutdown_signal.fuse();
         pin_mut!(shutdown_signal);
 
@@ -235,7 +235,8 @@ mod tests {
     use ceramic_core::NodeKey;
     use ceramic_sql::sqlite::SqlitePool;
     use expect_test::expect_file;
-    use tokio::{sync::broadcast, time::sleep};
+    use shutdown::Shutdown;
+    use tokio::time::sleep;
 
     use super::AnchorService;
     use crate::{MockAnchorEventService, MockCas};
@@ -248,28 +249,22 @@ mod tests {
         let node_id = NodeKey::random().id();
         let anchor_interval = Duration::from_millis(5);
         let anchor_batch_size = 1000000;
-        let mut anchor_service = AnchorService::new(
+        let anchor_service = AnchorService::new(
             tx_manager,
             event_service.clone(),
             pool,
             node_id,
             anchor_interval,
             anchor_batch_size,
         );
-        let (shutdown_signal_tx, mut shutdown_signal) = broadcast::channel::<()>(1);
-        tokio::spawn(async move {
-            anchor_service
-                .run(async move {
-                    let _ = shutdown_signal.recv().await;
-                })
-                .await
-        });
+        let shutdown = Shutdown::new();
+        tokio::spawn(anchor_service.run(shutdown.wait_fut()));
         while event_service.events.lock().unwrap().is_empty() {
             sleep(Duration::from_millis(1)).await;
         }
         expect_file!["./test-data/test_anchor_service_run.txt"]
             .assert_debug_eq(&event_service.events.lock().unwrap());
-        shutdown_signal_tx.send(()).unwrap();
+        shutdown.shutdown();
     }
 
     #[tokio::test]
@@ -280,28 +275,21 @@ mod tests {
         let node_id = NodeKey::random().id();
         let anchor_interval = Duration::from_millis(5);
         let anchor_batch_size = 1000000;
-        let mut anchor_service = AnchorService::new(
+        let anchor_service = AnchorService::new(
             tx_manager,
             event_service.clone(),
             pool,
             node_id,
             anchor_interval,
             anchor_batch_size,
         );
-        let (shutdown_signal_tx, mut shutdown_signal) = broadcast::channel::<()>(1);
-        // let mut shutdown_signal = shutdown_signal_rx.resubscribe();
-        tokio::spawn(async move {
-            anchor_service
-                .run(async move {
-                    let _ = shutdown_signal.recv().await;
-                })
-                .await
-        });
+        let shutdown = Shutdown::new();
+        tokio::spawn(anchor_service.run(shutdown.wait_fut()));
         while event_service.events.lock().unwrap().is_empty() {
             sleep(Duration::from_millis(1)).await;
         }
         expect_file!["./test-data/test_anchor_service_run_1.txt"]
             .assert_debug_eq(&event_service.events.lock().unwrap());
-        shutdown_signal_tx.send(()).unwrap();
+        shutdown.shutdown();
     }
 }

api/Cargo.toml

@@ -26,6 +26,7 @@ recon.workspace = true
 serde.workspace = true
 serde_ipld_dagcbor.workspace = true
 serde_json.workspace = true
+shutdown.workspace = true
 swagger.workspace = true
 tokio.workspace = true
 tracing.workspace = true

api/src/server.rs

@@ -52,6 +52,7 @@ use datafusion::logical_expr::{col, lit, BuiltInWindowFunction, Expr, ExprFuncti
 use futures::TryFutureExt;
 use multiaddr::Protocol;
 use recon::Key;
+use shutdown::Shutdown;
 use swagger::{ApiError, ByteArray};
 #[cfg(not(target_env = "msvc"))]
 use tikv_jemalloc_ctl::epoch;
@@ -401,7 +402,7 @@ where
         model: Arc<M>,
         p2p: P,
         pipeline: Option<SessionContext>,
-        shutdown_signal: broadcast::Receiver<()>,
+        shutdown_signal: Shutdown,
     ) -> Self {
         let (tx, event_rx) = tokio::sync::mpsc::channel::<EventInsert>(1024);
         let event_store = model.clone();
@@ -433,7 +434,7 @@ where
         event_store: Arc<M>,
         mut event_rx: tokio::sync::mpsc::Receiver<EventInsert>,
         node_id: NodeId,
-        mut shutdown_signal: broadcast::Receiver<()>,
+        shutdown_signal: Shutdown,
     ) -> tokio::task::JoinHandle<()> {
         tokio::spawn(async move {
             let mut interval = tokio::time::interval(Duration::from_millis(FLUSH_INTERVAL_MS));
@@ -455,7 +456,7 @@ where
                             events.extend(buf);
                         }
                     }
-                    _ = shutdown_signal.recv() => {
+                    _ = shutdown_signal.wait_fut() => {
                         tracing::debug!("Insert many task got shutdown signal");
                         shutdown = true;
                     }

api/src/tests.rs

@@ -32,6 +32,7 @@ use mockall::{mock, predicate};
 use multiaddr::Multiaddr;
 use multibase::Base;
 use recon::Key;
+use shutdown::Shutdown;
 use test_log::test;
 use tokio::join;
 
@@ -202,8 +203,8 @@ where
     M: EventService + 'static,
     P: P2PService,
 {
-    let (_, rx) = tokio::sync::broadcast::channel(1);
-    Server::new(node_id, network, interest, model, p2p, pipeline, rx)
+    let shutdown = Shutdown::new();
+    Server::new(node_id, network, interest, model, p2p, pipeline, shutdown)
 }
 
 #[test(tokio::test)]