opt(katana): graceful shutdown on critical tasks termination

bin/katana/src/cli/node.rs

@@ -234,9 +234,14 @@
             print_intro(&self, genesis, node.rpc.addr);
         }
 
-        // Wait until Ctrl + C is pressed, then shutdown
-        ctrl_c().await?;
-        node.rpc.handle.stop()?;
+        // Wait until ctrl-c signal is received or TaskManager signals shutdown
+        tokio::select! {
+            _ = ctrl_c() => {},
+            _ = node.task_manager.wait_for_shutdown() => {}
+        }
+
+        info!("Shutting down...");
+        node.stop().await?;
 
         Ok(())
     }

crates/katana/node/Cargo.toml

@@ -14,6 +14,7 @@ katana-primitives.workspace = true
 katana-provider.workspace = true
 katana-rpc.workspace = true
 katana-rpc-api.workspace = true
+katana-tasks.workspace = true
 
 anyhow.workspace = true
 dojo-metrics.workspace = true
@@ -22,7 +23,6 @@ jsonrpsee.workspace = true
 num-traits.workspace = true
 serde_json.workspace = true
 starknet.workspace = true
-tokio.workspace = true
 tower = { workspace = true, features = [ "full" ] }
 tower-http = { workspace = true, features = [ "full" ] }
 tracing.workspace = true

crates/katana/node/src/lib.rs

@@ -41,6 +41,7 @@
 use katana_rpc_api::starknet::{StarknetApiServer, StarknetTraceApiServer, StarknetWriteApiServer};
 use katana_rpc_api::torii::ToriiApiServer;
 use katana_rpc_api::ApiKind;
+use katana_tasks::TaskManager;
 use num_traits::ToPrimitive;
 use starknet::core::types::{BlockId, BlockStatus, MaybePendingBlockWithTxHashes};
 use starknet::core::utils::parse_cairo_short_string;
@@ -54,10 +55,21 @@
 pub struct Handle {
     pub pool: TxPool,
     pub rpc: RpcServer,
+    pub task_manager: TaskManager,
     pub backend: Arc<Backend<BlockifierFactory>>,
     pub block_producer: Arc<BlockProducer<BlockifierFactory>>,
 }
 
+impl Handle {
+    /// Stops the Katana node.
+    pub async fn stop(self) -> Result<()> {
+        // TODO: wait for the rpc server to stop
+        self.rpc.handle.stop()?;
+        self.task_manager.shutdown().await;
+        Ok(())
+    }
+}
+
 /// Build the core Katana components from the given configurations and start running the node.
 // TODO: placeholder until we implement a dedicated class that encapsulate building the node
 // components
@@ -209,8 +221,11 @@
 
     let block_producer = Arc::new(block_producer);
 
-    // TODO: avoid dangling task, or at least store the handle to the NodeService
-    tokio::spawn(NodeService::new(
+    // Create a TaskManager using the ambient Tokio runtime
+    let task_manager = TaskManager::current();
+
+    // Spawn the NodeService as a critical task
+    task_manager.build_task().critical().name("NodeService").spawn(NodeService::new(
         pool.clone(),
         miner,
         block_producer.clone(),
@@ -223,7 +238,7 @@
     let node_components = (pool.clone(), backend.clone(), block_producer.clone(), validator);
     let rpc = spawn(node_components, server_config).await?;
 
-    Ok(Handle { backend, block_producer, pool, rpc })
+    Ok(Handle { backend, block_producer, pool, rpc, task_manager })
 }
 
 // Moved from `katana_rpc` crate

crates/katana/tasks/src/manager.rs

@@ -42,21 +42,19 @@ impl TaskManager {
         self.spawn_inner(fut)
     }
 
-    /// Wait until all spawned tasks are completed.
-    pub async fn wait(&self) {
-        // need to close the tracker first before waiting
-        let _ = self.tracker.close();
-        self.tracker.wait().await;
-        // reopen the tracker for spawning future tasks
-        let _ = self.tracker.reopen();
+    /// Wait for the shutdown signal to be received.
+    pub async fn wait_for_shutdown(&self) {
+        self.on_cancel.cancelled().await;
     }
 
-    /// Consumes the manager and wait until all tasks are finished, either due to completion or
+    /// Shutdowns the manger and wait until all tasks are finished, either due to completion or
     /// cancellation.
-    pub async fn wait_shutdown(self) {
+    ///
+    /// No task can be spawned on the manager after this method is called.
+    pub async fn shutdown(self) {
+        self.wait_for_shutdown().await;
         // need to close the tracker first before waiting
         let _ = self.tracker.close();
-        let _ = self.on_cancel.cancelled().await;
         self.tracker.wait().await;
     }
 
@@ -70,6 +68,16 @@ impl TaskManager {
         TaskBuilder::new(self)
     }
 
+    /// Wait until all spawned tasks are completed.
+    #[cfg(test)]
+    async fn wait(&self) {
+        // need to close the tracker first before waiting
+        let _ = self.tracker.close();
+        self.tracker.wait().await;
+        // reopen the tracker for spawning future tasks
+        let _ = self.tracker.reopen();
+    }
+
     fn spawn_inner<F>(&self, task: F) -> TaskHandle<F::Output>
     where
         F: Future + Send + 'static,
@@ -156,20 +164,20 @@ mod tests {
         manager.build_task().graceful_shutdown().spawn(future::ready(()));
 
         // wait until all task spawned to the manager have been completed
-        manager.wait_shutdown().await;
+        manager.shutdown().await;
     }
 
     #[tokio::test]
     async fn critical_task_implicit_graceful_shutdown() {
         let manager = TaskManager::current();
         manager.build_task().critical().spawn(future::ready(()));
-        manager.wait_shutdown().await;
+        manager.shutdown().await;
     }
 
     #[tokio::test]
     async fn critical_task_graceful_shudown_on_panicked() {
         let manager = TaskManager::current();
         manager.build_task().critical().spawn(async { panic!("panicking") });
-        manager.wait_shutdown().await;
+        manager.shutdown().await;
     }
 }

crates/katana/tasks/src/task.rs

@@ -79,15 +79,15 @@ impl<'a> TaskBuilder<'a> {
         let Self { manager, instrument, graceful_shutdown, .. } = self;
 
         // creates a future that will send a cancellation signal to the manager when the future is
-        // completed.
-        let fut = if graceful_shutdown {
+        // completed, regardless of success or error.
+        let fut = {
             let ct = manager.on_cancel.clone();
-            Either::Left(fut.map(move |a| {
-                ct.cancel();
-                a
-            }))
-        } else {
-            Either::Right(fut)
+            fut.map(move |res| {
+                if graceful_shutdown {
+                    ct.cancel();
+                }
+                res
+            })
         };
 
         let fut = if instrument {