feat(jstzd): implement shutdown signal in jstzd server

jstz-dev · Nov 29, 2024 · 3e03410 · 3e03410
1 parent f5fbf93
commit 3e03410
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 5 deletions.
diff --git a/crates/jstzd/src/task/jstzd.rs b/crates/jstzd/src/task/jstzd.rs
@@ -16,7 +16,11 @@ use octez::r#async::{
 };
 use serde::Serialize;
 use std::sync::Arc;
-use tokio::{net::TcpListener, sync::RwLock, task::JoinHandle};
+use tokio::{
+    net::TcpListener,
+    sync::{oneshot, RwLock},
+    task::JoinHandle,
+};
 
 struct Jstzd {
     octez_node: Arc<RwLock<OctezNode>>,
@@ -182,6 +186,7 @@ struct ServerState {
     jstzd_config_json: serde_json::Map<String, serde_json::Value>,
     jstzd: Option<Jstzd>,
     server_handle: Option<JoinHandle<()>>,
+    shutdown_tx: Option<oneshot::Sender<()>>,
 }
 
 #[async_trait]
@@ -195,10 +200,12 @@ impl AsyncDrop for JstzdServerInner {
 pub struct JstzdServer {
     port: u16,
     inner: Arc<AsyncDropper<JstzdServerInner>>,
+    shutdown_rx: Option<oneshot::Receiver<()>>,
 }
 
 impl JstzdServer {
     pub fn new(config: JstzdConfig, port: u16) -> Self {
+        let (tx, rx) = oneshot::channel();
         Self {
             port,
             inner: Arc::new(AsyncDropper::new(JstzdServerInner {
@@ -211,8 +218,16 @@ impl JstzdServer {
                     jstzd_config: Some(config),
                     jstzd: None,
                     server_handle: None,
+                    shutdown_tx: Some(tx),
                 })),
             })),
+            shutdown_rx: Some(rx),
+        }
+    }
+
+    pub async fn wait(&mut self) {
+        if let Some(rx) = self.shutdown_rx.take() {
+            let _ = rx.await;
         }
     }
 
@@ -294,6 +309,9 @@ async fn shutdown(state: &mut ServerState) -> Result<()> {
     }
     state.jstzd_config.take();
     state.jstzd_config_json.clear();
+    if let Some(v) = state.shutdown_tx.take() {
+        let _ = v.send(());
+    }
     Ok(())
 }
 

diff --git a/crates/jstzd/tests/jstzd_test.rs b/crates/jstzd/tests/jstzd_test.rs
@@ -27,11 +27,18 @@ async fn jstzd_test() {
 
     fetch_config_test(config, jstzd_port).await;
 
-    reqwest::Client::new()
-        .put(&format!("http://localhost:{}/shutdown", jstzd_port))
-        .send()
+    tokio::spawn(async move {
+        tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;
+        reqwest::Client::new()
+            .put(&format!("http://localhost:{}/shutdown", jstzd_port))
+            .send()
+            .await
+            .unwrap();
+    });
+
+    tokio::time::timeout(tokio::time::Duration::from_secs(30), jstzd.wait())
         .await
-        .unwrap();
+        .expect("should not wait too long for the server to be taken down");
 
     ensure_jstzd_components_are_down(&jstzd, &rpc_endpoint, jstzd_port).await;