Handle ctrl+c to gracefully shutdown the server(s)

Cargo.toml

@@ -55,6 +55,7 @@ arrow2 = "0.16"
 arrow2_convert = "0.4.2"
 clap = "4.0"
 comfy-table = { version = "6.1", default-features = false }
+ctrlc = { version = "3.0", features = ["termination"] }
 ecolor = "0.21.0"
 eframe = { version = "0.21.3", default-features = false }
 egui = "0.21.0"

crates/re_viewer/Cargo.toml

@@ -111,7 +111,7 @@ wgpu.workspace = true
 arboard = { version = "3.2", default-features = false, features = [
   "image-data",
 ] }
-ctrlc = { version = "3.0", features = ["termination"] }
+ctrlc.workspace = true
 puffin_http = "0.11"
 puffin.workspace = true
 

crates/re_web_viewer_server/Cargo.toml

@@ -39,6 +39,7 @@ analytics = ["dep:re_analytics"]
 re_log.workspace = true
 
 anyhow.workspace = true
+ctrlc.workspace = true
 document-features = "0.2"
 futures-util = "0.3"
 hyper = { version = "0.14", features = ["full"] }

crates/re_web_viewer_server/src/lib.rs

@@ -157,8 +157,15 @@ impl WebViewerServer {
         Self { server }
     }
 
-    pub async fn serve(self) -> anyhow::Result<()> {
-        self.server.await?;
+    pub async fn serve(
+        self,
+        mut shutdown_rx: tokio::sync::broadcast::Receiver<()>,
+    ) -> anyhow::Result<()> {
+        self.server
+            .with_graceful_shutdown(async {
+                shutdown_rx.recv().await.ok();
+            })
+            .await?;
         Ok(())
     }
 }

crates/re_web_viewer_server/src/main.rs

@@ -6,8 +6,17 @@ async fn main() {
     re_log::setup_native_logging();
     let port = 9090;
     eprintln!("Hosting web-viewer on http://127.0.0.1:{port}");
+
+    // Shutdown server via Ctrl+C
+    let (shutdown_tx, shutdown_rx) = tokio::sync::broadcast::channel(1);
+    ctrlc::set_handler(move || {
+        re_log::debug!("Ctrl-C detected - Closing server.");
+        shutdown_tx.send(()).unwrap();
+    })
+    .expect("Error setting Ctrl-C handler");
+
     re_web_viewer_server::WebViewerServer::new(port)
-        .serve()
+        .serve(shutdown_rx)
         .await
         .unwrap();
 }

crates/re_ws_comms/src/server.rs

@@ -41,14 +41,25 @@ impl Server {
     }
 
     /// Accept new connections forever
-    pub async fn listen(self, rx: Receiver<LogMsg>) -> anyhow::Result<()> {
+    pub async fn listen(
+        self,
+        rx: Receiver<LogMsg>,
+        mut shutdown_rx: tokio::sync::broadcast::Receiver<()>,
+    ) -> anyhow::Result<()> {
         use anyhow::Context as _;
 
         let history = Arc::new(Mutex::new(Vec::new()));
 
         let log_stream = to_broadcast_stream(rx, history.clone());
 
-        while let Ok((tcp_stream, _)) = self.listener.accept().await {
+        loop {
+            let (tcp_stream, _) = tokio::select! {
+                res = self.listener.accept() => res?,
+                _ = shutdown_rx.recv() => {
+                    return Ok(());
+                }
+            };
+
             let peer = tcp_stream
                 .peer_addr()
                 .context("connected streams should have a peer address")?;
@@ -59,8 +70,6 @@ impl Server {
                 history.clone(),
             ));
         }
-
-        Ok(())
     }
 }
 

crates/rerun/Cargo.toml

@@ -97,6 +97,7 @@ backtrace = "0.3"
 clap = { workspace = true, features = ["derive"] }
 mimalloc.workspace = true
 puffin_http = "0.11"
+ctrlc.workspace = true
 tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
 
 # Native unix dependencies:

crates/rerun/src/run.rs

@@ -262,7 +262,9 @@ async fn run_impl(
                 // We are connecting to a server at a websocket address:
 
                 if args.web_viewer {
-                    return host_web_viewer(rerun_server_ws_url).await;
+                    let shutdown_rx = setup_ctrl_c_handler();
+                    let web_viewer = host_web_viewer(true, rerun_server_ws_url, shutdown_rx);
+                    return web_viewer.await;
                 } else {
                     #[cfg(feature = "native_viewer")]
                     return native_viewer_connect_to_ws_url(
@@ -311,13 +313,21 @@ async fn run_impl(
                 );
             }
 
+            // Make it possible to gracefully shutdown the servers on ctrl-c.
+            let shutdown_ws_server = setup_ctrl_c_handler();
+            let shutdown_web_viewer = shutdown_ws_server.resubscribe();
+
             // This is the server which the web viewer will talk to:
             let ws_server = re_ws_comms::Server::new(re_ws_comms::DEFAULT_WS_SERVER_PORT).await?;
-            let server_handle = tokio::spawn(ws_server.listen(rx));
+            let server_handle = tokio::spawn(ws_server.listen(rx, shutdown_ws_server));
 
-            let rerun_ws_server_url = re_ws_comms::default_server_url();
-            host_web_viewer(rerun_ws_server_url).await?;
+            // This is the server that serves the Wasm+HTML:
+            let ws_server_url = re_ws_comms::default_server_url();
+            let ws_server_handle =
+                tokio::spawn(host_web_viewer(true, ws_server_url, shutdown_web_viewer));
 
+            // Wait for both servers to shutdown.
+            ws_server_handle.await?.ok();
             return server_handle.await?;
         }
 
@@ -355,6 +365,16 @@ async fn run_impl(
     }
 }
 
+fn setup_ctrl_c_handler() -> tokio::sync::broadcast::Receiver<()> {
+    let (sender, receiver) = tokio::sync::broadcast::channel(1);
+    ctrlc::set_handler(move || {
+        re_log::debug!("Ctrl-C detected - Closing server.");
+        sender.send(()).unwrap();
+    })
+    .expect("Error setting Ctrl-C handler");
+    receiver
+}
+
 enum ArgumentCategory {
     /// A remote RRD file, served over http.
     RrdHttpUrl(String),
@@ -440,25 +460,14 @@ fn load_file_to_channel(path: &std::path::Path) -> anyhow::Result<Receiver<LogMs
 }
 
 #[cfg(feature = "web_viewer")]
-async fn host_web_viewer(rerun_ws_server_url: String) -> anyhow::Result<()> {
-    let web_port = 9090;
-    let viewer_url = format!("http://127.0.0.1:{web_port}?url={rerun_ws_server_url}");
-
-    let web_server = re_web_viewer_server::WebViewerServer::new(web_port);
-    let web_server_handle = tokio::spawn(web_server.serve());
-
-    let open = true;
-    if open {
-        webbrowser::open(&viewer_url).ok();
-    } else {
-        println!("Hosting Rerun Web Viewer at {viewer_url}.");
-    }
-
-    web_server_handle.await?
-}
+use crate::web_viewer::host_web_viewer;
 
 #[cfg(not(feature = "web_viewer"))]
-async fn host_web_viewer(_rerun_ws_server_url: String) -> anyhow::Result<()> {
+pub async fn host_web_viewer(
+    _open_browser: bool,
+    _ws_server_url: String,
+    _shutdown_rx: tokio::sync::broadcast::Receiver<()>,
+) -> anyhow::Result<()> {
     panic!("Can't host web-viewer - rerun was not compiled with the 'web_viewer' feature");
 }
 

crates/rerun/src/web_viewer.rs

@@ -6,6 +6,8 @@ use re_log_types::LogMsg;
 struct RemoteViewerServer {
     web_server_join_handle: tokio::task::JoinHandle<()>,
     sender: re_smart_channel::Sender<LogMsg>,
+    #[allow(dead_code)] // Unused currently, but can be used later to cancel a serve.
+    shutdown_tx: tokio::sync::broadcast::Sender<()>,
 }
 
 impl Drop for RemoteViewerServer {
@@ -18,40 +20,58 @@ impl Drop for RemoteViewerServer {
 impl RemoteViewerServer {
     pub fn new(tokio_rt: &tokio::runtime::Runtime, open_browser: bool) -> Self {
         let (rerun_tx, rerun_rx) = re_smart_channel::smart_channel(re_smart_channel::Source::Sdk);
+        let (shutdown_tx, shutdown_rx_ws_server) = tokio::sync::broadcast::channel(1);
+        let shutdown_rx_web_server = shutdown_tx.subscribe();
 
         let web_server_join_handle = tokio_rt.spawn(async move {
             // This is the server which the web viewer will talk to:
             let ws_server = re_ws_comms::Server::new(re_ws_comms::DEFAULT_WS_SERVER_PORT)
                 .await
                 .unwrap();
-            let ws_server_handle = tokio::spawn(ws_server.listen(rerun_rx)); // TODO(emilk): use tokio_rt ?
+            let ws_server_handle = tokio::spawn(ws_server.listen(rerun_rx, shutdown_rx_ws_server)); // TODO(emilk): use tokio_rt ?
 
             // This is the server that serves the Wasm+HTML:
-            let web_port = 9090;
-            let web_server = re_web_viewer_server::WebViewerServer::new(web_port);
-            let web_server_handle = tokio::spawn(async move {
-                web_server.serve().await.unwrap();
-            });
-
             let ws_server_url = re_ws_comms::default_server_url();
-            let viewer_url = format!("http://127.0.0.1:{web_port}?url={ws_server_url}");
-            if open_browser {
-                webbrowser::open(&viewer_url).ok();
-            } else {
-                re_log::info!("Web server is running - view it at {viewer_url}");
-            }
+            let web_server_handle = tokio::spawn(host_web_viewer(
+                open_browser,
+                ws_server_url,
+                shutdown_rx_web_server,
+            ));
 
             ws_server_handle.await.unwrap().unwrap();
-            web_server_handle.await.unwrap();
+            web_server_handle.await.unwrap().unwrap();
         });
 
         Self {
             web_server_join_handle,
             sender: rerun_tx,
+            shutdown_tx,
         }
     }
 }
 
+/// Start a web server for localhost and optionally spawns a browser to view it.
+#[cfg(feature = "web_viewer")]
+pub async fn host_web_viewer(
+    open_browser: bool,
+    ws_server_url: String,
+    shutdown_rx: tokio::sync::broadcast::Receiver<()>,
+) -> anyhow::Result<()> {
+    let web_port = 9090;
+    let viewer_url = format!("http://127.0.0.1:{web_port}?url={ws_server_url}");
+
+    let web_server = re_web_viewer_server::WebViewerServer::new(web_port);
+    let web_server_handle = tokio::spawn(web_server.serve(shutdown_rx));
+
+    if open_browser {
+        webbrowser::open(&viewer_url).ok();
+    } else {
+        re_log::info!("Web server is running - view it at {viewer_url}");
+    }
+
+    web_server_handle.await?
+}
+
 impl crate::sink::LogSink for RemoteViewerServer {
     fn send(&self, msg: LogMsg) {
         if let Err(err) = self.sender.send(msg) {

rerun_py/rerun_sdk/rerun/script_helpers.py

@@ -18,9 +18,7 @@
 ```
 
 """
-import contextlib
 from argparse import ArgumentParser, Namespace
-from time import sleep
 
 import rerun as rr
 
@@ -93,6 +91,10 @@ def script_teardown(args: Namespace) -> None:
 
     """
     if args.serve:
+        import signal
+        from threading import Event
+
+        exit = Event()
+        signal.signal(signal.SIGINT, lambda sig, frame: exit.set())
         print("Sleeping while serving the web viewer. Abort with Ctrl-C")
-        with contextlib.suppress(Exception):
-            sleep(1_000_000_000)
+        exit.wait()