rusk: implement RuskHttp builder and shutdown

Author: Neotamandua

rusk/src/lib/builder/http_only.rs

@@ -4,31 +4,52 @@
 //
 // Copyright (c) DUSK NETWORK. All rights reserved.
 
+use std::time::Duration;
+
 use tokio::sync::broadcast;
 use tracing::info;
 
-use crate::http::{DataSources, HttpServer, HttpServerConfig};
+use crate::http::{DataSources, HandleRequest, HttpServer, HttpServerConfig};
 
-#[derive(Default)]
 pub struct RuskHttpBuilder {
     http: Option<HttpServerConfig>,
+    data_sources: DataSources,
+    shutdown_timeout: Duration,
 }
 
 impl RuskHttpBuilder {
+    pub fn new() -> Self {
+        Self {
+            http: None,
+            data_sources: DataSources::default(),
+            shutdown_timeout: Duration::from_secs(30),
+        }
+    }
+
     pub fn with_http(mut self, http: HttpServerConfig) -> Self {
         self.http = Some(http);
         self
     }
 
-    pub async fn build_and_run(self) -> anyhow::Result<()> {
+    pub fn with_data_source(mut self, source: Box<dyn HandleRequest>) -> Self {
+        self.data_sources.sources.push(source);
+        self
+    }
+
+    pub fn with_shutdown_timeout(mut self, timeout: Duration) -> Self {
+        self.shutdown_timeout = timeout;
+        self
+    }
+
+    pub async fn build(self) -> anyhow::Result<RuskHttp> {
         let (_rues_sender, rues_receiver) = broadcast::channel(1);
 
-        let mut _ws_server = None;
+        let mut server = None;
         if let Some(http) = self.http {
             info!("Configuring HTTP");
 
             #[allow(unused_mut)]
-            let mut handler = DataSources::default();
+            let mut handler = self.data_sources;
 
             #[cfg(feature = "prover")]
             handler.sources.push(Box::new(rusk_prover::LocalProver));
@@ -38,7 +59,7 @@ impl RuskHttpBuilder {
                 _ => None,
             };
 
-            let (server, _) = HttpServer::bind(
+            let (http_server, _) = HttpServer::bind(
                 handler,
                 rues_receiver,
                 http.ws_event_channel_cap,
@@ -48,13 +69,51 @@ impl RuskHttpBuilder {
             )
             .await?;
 
-            _ws_server = Some(server);
+            server = Some(http_server);
         }
 
-        if let Some(s) = _ws_server {
-            s.wait().await?;
+        Ok(RuskHttp {
+            server,
+            shutdown_timeout: self.shutdown_timeout,
+        })
+    }
+
+    pub async fn build_and_run(self) -> anyhow::Result<()> {
+        let mut http = self.build().await?;
+        http.run().await
+    }
+}
+
+impl Default for RuskHttpBuilder {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+pub struct RuskHttp {
+    server: Option<HttpServer>,
+    shutdown_timeout: Duration,
+}
+
+impl RuskHttp {
+    pub async fn run(&mut self) -> anyhow::Result<()> {
+        if let Some(server) = &mut self.server {
+            server.wait().await?;
         }
+        Ok(())
+    }
 
+    pub async fn shutdown(&mut self) -> anyhow::Result<()> {
+        if let Some(server) = &mut self.server {
+            tokio::time::timeout(self.shutdown_timeout, server.shutdown())
+                .await
+                .map_err(|_| {
+                    anyhow::anyhow!(
+                        "HTTP server failed to shut down within {} seconds",
+                        self.shutdown_timeout.as_secs()
+                    )
+                })??;
+        }
         Ok(())
     }
 }

rusk/src/lib/http.rs

@@ -89,8 +89,9 @@ const RUSK_VERSION_HEADER: &str = "Rusk-Version";
 const RUSK_VERSION_STRICT_HEADER: &str = "Rusk-Version-Strict";
 
 pub struct HttpServer {
-    handle: task::JoinHandle<()>,
-    _shutdown: broadcast::Sender<Infallible>,
+    handle: Option<task::JoinHandle<()>>,
+    #[cfg_attr(feature = "chain", allow(dead_code))]
+    shutdown: Option<broadcast::Sender<Infallible>>,
 }
 
 pub struct HttpServerConfig {
@@ -102,8 +103,17 @@ pub struct HttpServerConfig {
 }
 
 impl HttpServer {
-    pub async fn wait(self) -> Result<(), JoinError> {
-        self.handle.await
+    pub async fn wait(&mut self) -> Result<(), JoinError> {
+        match self.handle.take() {
+            Some(handle) => handle.await,
+            None => Ok(()),
+        }
+    }
+
+    #[cfg(not(feature = "chain"))]
+    pub(crate) async fn shutdown(&mut self) -> Result<(), JoinError> {
+        self.shutdown.take(); // this closes the channel to signal shutdown
+        self.wait().await
     }
 
     pub async fn bind<A, H, P1, P2>(
@@ -141,8 +151,8 @@ impl HttpServer {
         ));
 
         let server = Self {
-            handle,
-            _shutdown: shutdown_sender,
+            handle: Some(handle),
+            shutdown: Some(shutdown_sender),
         };
         Ok((server, local_addr))
     }
@@ -256,16 +266,9 @@ async fn listening_loop<H>(
         ws_event_channel_cap,
     };
 
-    let runtime = tokio::runtime::Builder::new_multi_thread()
-        .worker_threads(4)
-        .thread_name("http")
-        .enable_all()
-        .build()
-        .expect("http runtime to be created");
     loop {
         tokio::select! {
             _ = shutdown.recv() => {
-                runtime.shutdown_background();
                 break;
             }
             r = listener.accept() => {
@@ -279,8 +282,9 @@ async fn listening_loop<H>(
                 let stream = TokioIo::new(stream);
                 let service = service.clone();
 
-                runtime.spawn(async move {
-                    let conn = http.serve_connection_with_upgrades(stream, service);
+                task::spawn(async move {
+                    let conn =
+                        http.serve_connection_with_upgrades(stream, service);
                     conn.await
                 });
             }