refactor(katana-node): distinguish between launched node handle

bin/katana/src/cli/node.rs

@@ -225,23 +225,32 @@
         let sequencer_config = self.sequencer_config();
         let starknet_config = self.starknet_config()?;
 
-        // build the node and start it
-        let node = katana_node::start(server_config, sequencer_config, starknet_config).await?;
+        // Build the node
+        let node = katana_node::build(server_config, sequencer_config, starknet_config)
+            .await
+            .context("failed to build node")?;
 
         if !self.silent {
             #[allow(deprecated)]
             let genesis = &node.backend.config.genesis;
-            print_intro(&self, genesis, node.rpc.addr);
+            let server_address = node.server_config.addr();
+            print_intro(&self, genesis, &server_address);
         }
 
-        // Wait until an OS signal is received or TaskManager shutdown
+        // Launch the node
+        let handle = node.launch().await.context("failed to launch node")?;
+
+        // Wait until an OS signal (ie SIGINT, SIGTERM) is received or the node is shutdown.
         tokio::select! {
-            _ = dojo_utils::signal::wait_signals() => {},
-            _ = node.task_manager.wait_for_shutdown() => {}
+            _ = dojo_utils::signal::wait_signals() => {
+                // Gracefully shutdown the node before exiting
+                handle.stop().await?;
+            },
+
+            _ = handle.stopped() => { }
         }
 
-        info!("Shutting down...");
-        node.stop().await?;
+        info!("Shutting down.");
 
         Ok(())
     }
@@ -339,7 +348,7 @@
     }
 }
 
-fn print_intro(args: &NodeArgs, genesis: &Genesis, address: SocketAddr) {
+fn print_intro(args: &NodeArgs, genesis: &Genesis, address: &str) {
     let mut accounts = genesis.accounts().peekable();
     let account_class_hash = accounts.peek().map(|e| e.1.class_hash());
     let seed = &args.starknet.seed;

crates/dojo-test-utils/src/sequencer.rs

@@ -7,7 +7,7 @@
 #[allow(deprecated)]
 pub use katana_core::sequencer::SequencerConfig;
 use katana_executor::implementation::blockifier::BlockifierFactory;
-use katana_node::Handle;
+use katana_node::LaunchedNode;
 use katana_primitives::chain::ChainId;
 use katana_rpc::config::ServerConfig;
 use katana_rpc_api::ApiKind;
@@ -29,7 +29,7 @@
 #[allow(missing_debug_implementations)]
 pub struct TestSequencer {
     url: Url,
-    handle: Handle,
+    handle: LaunchedNode,
     account: TestAccount,
 }
 
@@ -45,19 +45,20 @@
             apis: vec![ApiKind::Starknet, ApiKind::Dev, ApiKind::Saya, ApiKind::Torii],
         };
 
-        let node = katana_node::start(server_config, config, starknet_config)
+        let node = katana_node::build(server_config, config, starknet_config)
             .await
             .expect("Failed to build node components");
+        let handle = node.launch().await.expect("Failed to launch node");
 
-        let url = Url::parse(&format!("http://{}", node.rpc.addr)).expect("Failed to parse URL");
+        let url = Url::parse(&format!("http://{}", handle.rpc.addr)).expect("Failed to parse URL");
 
-        let account = node.backend.config.genesis.accounts().next().unwrap();
+        let account = handle.node.backend.config.genesis.accounts().next().unwrap();
         let account = TestAccount {
             private_key: Felt::from_bytes_be(&account.1.private_key().unwrap().to_bytes_be()),
             account_address: Felt::from_bytes_be(&account.0.to_bytes_be()),
         };
 
-        TestSequencer { handle: node, account, url }
+        TestSequencer { handle, account, url }
     }
 
     pub fn account(&self) -> SingleOwnerAccount<JsonRpcClient<HttpTransport>, LocalWallet> {
@@ -79,15 +80,15 @@
     }
 
     pub fn backend(&self) -> &Arc<Backend<BlockifierFactory>> {
-        &self.handle.backend
+        &self.handle.node.backend
     }
 
     pub fn account_at_index(
         &self,
         index: usize,
     ) -> SingleOwnerAccount<JsonRpcClient<HttpTransport>, LocalWallet> {
         #[allow(deprecated)]
-        let accounts: Vec<_> = self.handle.backend.config.genesis.accounts().collect::<_>();
+        let accounts: Vec<_> = self.handle.node.backend.config.genesis.accounts().collect::<_>();
 
         let account = accounts[index];
         let private_key = Felt::from_bytes_be(&account.1.private_key().unwrap().to_bytes_be());

crates/katana/node/Cargo.toml

@@ -19,6 +19,7 @@ katana-tasks.workspace = true
 
 anyhow.workspace = true
 dojo-metrics.workspace = true
+futures.workspace = true
 hyper.workspace = true
 jsonrpsee.workspace = true
 num-traits.workspace = true

crates/katana/node/src/exit.rs

@@ -0,0 +1,41 @@
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use anyhow::Result;
+use futures::future::BoxFuture;
+use futures::FutureExt;
+
+use crate::LaunchedNode;
+
+/// A Future that is resolved once the node has been stopped including all of its running tasks.
+#[must_use = "futures do nothing unless polled"]
+pub struct NodeStoppedFuture<'a> {
+    fut: BoxFuture<'a, Result<()>>,
+}
+
+impl<'a> NodeStoppedFuture<'a> {
+    pub(crate) fn new(handle: &'a LaunchedNode) -> Self {
+        let fut = Box::pin(async {
+            handle.node.task_manager.wait_for_shutdown().await;
+            handle.stop().await?;
+            Ok(())
+        });
+        Self { fut }
+    }
+}
+
+impl<'a> Future for NodeStoppedFuture<'a> {
+    type Output = Result<()>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.get_mut();
+        this.fut.poll_unpin(cx)
+    }
+}
+
+impl<'a> core::fmt::Debug for NodeStoppedFuture<'a> {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.debug_struct("NodeStoppedFuture").field("fut", &"...").finish()
+    }
+}

crates/katana/node/src/lib.rs

@@ -1,5 +1,7 @@
 #![cfg_attr(not(test), warn(unused_crate_dependencies))]
 
+mod exit;
+
 use std::future::IntoFuture;
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -19,6 +21,7 @@
 #[allow(deprecated)]
 use katana_core::sequencer::SequencerConfig;
 use katana_core::service::block_producer::BlockProducer;
+use katana_db::mdbx::DbEnv;
 use katana_executor::implementation::blockifier::BlockifierFactory;
 use katana_executor::{ExecutorFactory, SimulationFlag};
 use katana_pipeline::{stage, Pipeline};
@@ -49,27 +52,107 @@
 use tower_http::cors::{AllowOrigin, CorsLayer};
 use tracing::{info, trace};
 
-/// A handle to the instantiated Katana node.
+use crate::exit::NodeStoppedFuture;
+
+/// A handle to the launched node.
 #[allow(missing_debug_implementations)]
-pub struct Handle {
-    pub pool: TxPool,
+pub struct LaunchedNode {
+    pub node: Node,
+    /// Handle to the rpc server.
     pub rpc: RpcServer,
+}
+
+impl LaunchedNode {
+    /// Stops the node.
+    ///
+    /// This will instruct the node to stop and wait until it has actually stop.
+    pub async fn stop(&self) -> Result<()> {
+        // TODO: wait for the rpc server to stop instead of just stopping it.
+        self.rpc.handle.stop()?;
+        self.node.task_manager.shutdown().await;
+        Ok(())
+    }
+
+    /// Returns a future which resolves only when the node has stopped.
+    pub fn stopped(&self) -> NodeStoppedFuture<'_> {
+        NodeStoppedFuture::new(self)
+    }
+}
+
+/// A node instance.
+///
+/// The struct contains the handle to all the components of the node.
+#[must_use = "Node does nothing unless launched."]
+#[allow(missing_debug_implementations)]
+pub struct Node {
+    pub pool: TxPool,
+    pub db: Option<DbEnv>,
     pub task_manager: TaskManager,
     pub backend: Arc<Backend<BlockifierFactory>>,
     pub block_producer: BlockProducer<BlockifierFactory>,
+    pub server_config: ServerConfig,
+    #[allow(deprecated)]
+    pub sequencer_config: SequencerConfig,
 }
 
-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(())
+impl Node {
+    /// Start the node.
+    ///
+    /// This method will start all the node process, running them until the node is stopped.
+    pub async fn launch(self) -> Result<LaunchedNode> {
+        // Metrics recorder must be initialized before calling any of the metrics macros, in order
+        // for it to be registered.
+
+        if let Some(addr) = self.server_config.metrics {
+            let prometheus_handle = prometheus_exporter::install_recorder("katana")?;
+            let mut reports = Vec::new();
+
+            if let Some(ref db) = self.db {
+                reports.push(Box::new(db.clone()) as Box<dyn Report>);
+            }
+
+            prometheus_exporter::serve(
+                addr,
+                prometheus_handle,
+                metrics_process::Collector::default(),
+                reports,
+            )
+            .await?;
+
+            info!(%addr, "Metrics endpoint started.");
+        }
+
+        let pool = self.pool.clone();
+        let backend = self.backend.clone();
+        let block_producer = self.block_producer.clone();
+        let validator = self.block_producer.validator().clone();
+
+        // --- build sequencing stage
+
+        #[allow(deprecated)]
+        let sequencing = stage::Sequencing::new(
+            pool.clone(),
+            backend.clone(),
+            self.task_manager.clone(),
+            block_producer.clone(),
+            self.sequencer_config.messaging.clone(),
+        );
+
+        // --- build and start the pipeline
+
+        let mut pipeline = Pipeline::new();
+        pipeline.add_stage(Box::new(sequencing));
+
+        self.task_manager.spawn(pipeline.into_future());
+
+        let node_components = (pool, backend, block_producer, validator);
+        let rpc = spawn(node_components, self.server_config.clone()).await?;
+
+        Ok(LaunchedNode { node: self, rpc })
     }
 }
 
-/// Build the core Katana components from the given configurations and start running the node.
+/// Build the core Katana components from the given configurations.
 // TODO: placeholder until we implement a dedicated class that encapsulate building the node
 // components
 //
@@ -79,11 +162,11 @@
 //
 // NOTE: Don't rely on this function as it is mainly used as a placeholder for now.
 #[allow(deprecated)]
-pub async fn start(
+pub async fn build(
     server_config: ServerConfig,
     sequencer_config: SequencerConfig,
     mut starknet_config: StarknetConfig,
-) -> Result<Handle> {
+) -> Result<Node> {
     // --- build executor factory
 
     let cfg_env = CfgEnv {
@@ -189,52 +272,17 @@
     let validator = block_producer.validator();
     let pool = TxPool::new(validator.clone(), FiFo::new());
 
-    // --- build metrics service
-
-    // Metrics recorder must be initialized before calling any of the metrics macros, in order for
-    // it to be registered.
-    if let Some(addr) = server_config.metrics {
-        let prometheus_handle = prometheus_exporter::install_recorder("katana")?;
-        let reports = db.map(|db| vec![Box::new(db) as Box<dyn Report>]).unwrap_or_default();
-
-        prometheus_exporter::serve(
-            addr,
-            prometheus_handle,
-            metrics_process::Collector::default(),
-            reports,
-        )
-        .await?;
-
-        info!(%addr, "Metrics endpoint started.");
-    }
-
-    // --- create a TaskManager using the ambient Tokio runtime
-
-    let task_manager = TaskManager::current();
-
-    // --- build sequencing stage
-
-    let sequencing = stage::Sequencing::new(
-        pool.clone(),
-        backend.clone(),
-        task_manager.clone(),
-        block_producer.clone(),
-        sequencer_config.messaging.clone(),
-    );
-
-    // --- build and start the pipeline
-
-    let mut pipeline = Pipeline::new();
-    pipeline.add_stage(Box::new(sequencing));
-
-    task_manager.spawn(pipeline.into_future());
-
-    // --- spawn rpc server
-
-    let node_components = (pool.clone(), backend.clone(), block_producer.clone(), validator);
-    let rpc = spawn(node_components, server_config).await?;
+    let node = Node {
+        db,
+        pool,
+        backend,
+        server_config,
+        block_producer,
+        sequencer_config,
+        task_manager: TaskManager::current(),
+    };
 
-    Ok(Handle { backend, block_producer, pool, rpc, task_manager })
+    Ok(node)
 }
 
 // Moved from `katana_rpc` crate