gracefully handle sigint

Author: lherman-cs

Cargo.lock

@@ -33,6 +33,7 @@ mimalloc = { version = "0.1.48", features = ["v3"] }
 hyper = "1.7.0"
 pulsebeam-runtime = { version = "^0.1.0", path = "../pulsebeam-runtime" }
 thread-priority = "3.0.0"
+tokio-util = "0.7.16"
 
 [dev-dependencies]
 turmoil = "0.6.6"

pulsebeam/Cargo.toml

@@ -1,4 +1,5 @@
 use mimalloc::MiMalloc;
+use tokio_util::sync::CancellationToken;
 
 #[global_allocator]
 static GLOBAL: MiMalloc = MiMalloc;
@@ -55,15 +56,17 @@ fn main() {
         .build()
         .unwrap();
 
+    let shutdown = CancellationToken::new();
     let io_rt = tokio::runtime::Builder::new_current_thread()
         .enable_all()
         .thread_name("io")
         .build()
         .unwrap();
-    io_rt.block_on(run(cpu_rt));
+    io_rt.block_on(run(shutdown.clone(), &cpu_rt));
+    shutdown.cancel();
 }
 
-pub async fn run(cpu_rt: rt::Runtime) {
+pub async fn run(shutdown: CancellationToken, cpu_rt: &rt::Runtime) {
     let external_ip = select_host_address();
     let external_addr: SocketAddr = format!("{}:3478", external_ip).parse().unwrap();
     let local_addr: SocketAddr = "0.0.0.0:3478".parse().unwrap();
@@ -73,17 +76,18 @@ pub async fn run(cpu_rt: rt::Runtime) {
             .expect("bind to udp socket");
     let http_addr: SocketAddr = "0.0.0.0:3000".parse().unwrap();
     tracing::info!(
-        "✅ Signaling server listening. Clients should connect to http://{}:3000 or http://localhost:3000",
+        "server listening at http://{}:3000 or http://localhost:3000",
         external_ip,
     );
 
     // Run the main logic and signal handler concurrently
     tokio::select! {
-        Err(err) = node::run(cpu_rt, external_addr, unified_socket, http_addr) => {
-            tracing::warn!("node exited with an error: {err}");
+        Err(err) = node::run(shutdown.clone(), cpu_rt, external_addr, unified_socket, http_addr) => {
+            tracing::warn!("node exited with error: {err}");
         }
         _ = tokio::signal::ctrl_c() => {
-            tracing::info!("received SIGINT, shutting down gracefully");
+            tracing::info!("received SIGINT, shutting down gracefully...");
+                shutdown.cancel();
         }
     }
 }

pulsebeam/src/main.rs

@@ -2,29 +2,30 @@ use crate::{api, controller, system};
 use futures::{FutureExt, StreamExt, stream::FuturesUnordered};
 use pulsebeam_runtime::{actor, net, rt};
 use std::{net::SocketAddr, sync::Arc, time::Duration};
+use tokio_util::sync::CancellationToken;
 use tower_http::cors::{AllowOrigin, CorsLayer};
 
 pub async fn run(
-    cpu_rt: rt::Runtime,
+    shutdown: CancellationToken,
+    cpu_rt: &rt::Runtime,
     external_addr: SocketAddr,
     unified_socket: net::UnifiedSocket<'static>,
     http_addr: SocketAddr,
 ) -> anyhow::Result<()> {
-    // Configure CORS
     let cors = CorsLayer::very_permissive()
         .allow_origin(AllowOrigin::mirror_request())
         .expose_headers([hyper::header::LOCATION])
         .max_age(Duration::from_secs(86400));
 
-    // Spawn system and controller actors
     let mut join_set = FuturesUnordered::new();
 
     let (system_ctx, system_join) = system::SystemContext::spawn(unified_socket);
     join_set.push(system_join.map(|_| ()).boxed());
 
     let (controller_ready_tx, controller_ready_rx) = tokio::sync::oneshot::channel();
 
-    // TODO: handle cleanup
+    // Spawn controller on CPU runtime
+    let shutdown_for_controller = shutdown.clone();
     cpu_rt.spawn(async move {
         let controller_actor = controller::ControllerActor::new(
             system_ctx,
@@ -33,28 +34,52 @@ pub async fn run(
         );
         let (controller_handle, controller_join) =
             actor::spawn(controller_actor, actor::RunnerConfig::default());
-        controller_ready_tx.send(controller_handle).unwrap();
+        let _ = controller_ready_tx.send(controller_handle);
         tracing::debug!("controller is ready");
-        controller_join.await;
+
+        tokio::select! {
+            _ = controller_join => {}
+            _ = shutdown_for_controller.cancelled() => {
+                tracing::debug!("controller received shutdown");
+            }
+        }
     });
 
     tracing::debug!("waiting on controller to be ready");
     let controller_handle = controller_ready_rx.await?;
-    // Set up signaling router
+
+    // HTTP API
     let api_cfg = api::ApiConfig {
         base_path: "/api/v1".to_string(),
         default_host: http_addr.to_string(),
     };
     let router = api::router(controller_handle, api_cfg).layer(cors);
-    tracing::debug!("listening on {http_addr}");
 
+    let shutdown_for_http = shutdown.clone();
     let signaling = async move {
         let listener = tokio::net::TcpListener::bind(http_addr).await.unwrap();
-        axum::serve(listener, router).await.unwrap();
+        tracing::debug!("listening on {http_addr}");
+        tokio::select! {
+            res = axum::serve(listener, router) => {
+                if let Err(e) = res {
+                    tracing::error!("http server error: {e}");
+                }
+            }
+            _ = shutdown_for_http.cancelled() => {
+                tracing::info!("http server shutting down");
+            }
+        }
     };
+
     join_set.push(tokio::spawn(signaling).map(|_| ()).boxed());
 
-    // Wait for all tasks to complete
-    while join_set.next().await.is_some() {}
+    // Wait for all tasks to complete OR shutdown
+    tokio::select! {
+        _ = async { while join_set.next().await.is_some() {} } => {}
+        _ = shutdown.cancelled() => {
+            tracing::info!("node received shutdown");
+        }
+    }
+
     Ok(())
 }