feat(websockets): coordinate chat server shutdown

Author: robjtede

Cargo.lock

@@ -5,6 +5,9 @@ edition.workspace = true
 rust-version.workspace = true
 
 [dependencies]
+err-report = { workspace = true }
 rustls = { workspace = true }
+tokio = { workspace = true, features = ["rt", "signal", "macros"] }
+tokio-util = { workspace = true, features = ["rt"] }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }

examples-common/Cargo.toml

@@ -1,6 +1,51 @@
+use std::io;
+
+use err_report::Report;
+use tokio_util::sync::CancellationToken;
 use tracing::level_filters::LevelFilter;
 use tracing_subscriber::{EnvFilter, prelude::*};
 
+/// Spawns a task that cancels `shutdown` on SIGINT, SIGTERM, or SIGQUIT on Unix. On other
+/// platforms, listens for Ctrl-C.
+///
+/// Signal errors are logged with tracing and also cancel the token. The task is detached and stops
+/// when the Tokio runtime exits. Call this within a Tokio runtime.
+pub fn watch_shutdown_signal(shutdown: CancellationToken) {
+    tokio::spawn(async move {
+        let result: io::Result<()> = async {
+            #[cfg(unix)]
+            {
+                use tokio::signal::unix::{SignalKind, signal};
+
+                let mut interrupt = signal(SignalKind::interrupt())?;
+                let mut terminate = signal(SignalKind::terminate())?;
+                let mut quit = signal(SignalKind::quit())?;
+
+                tokio::select! {
+                    _ = interrupt.recv() => {},
+                    _ = terminate.recv() => {},
+                    _ = quit.recv() => {},
+                }
+
+                Ok(())
+            }
+
+            #[cfg(not(unix))]
+            tokio::signal::ctrl_c().await
+        }
+        .await;
+
+        if let Err(err) = result {
+            tracing::error!(
+                error = %Report::new(&err),
+                "failed to listen for shutdown signals"
+            );
+        }
+
+        shutdown.cancel();
+    });
+}
+
 /// Initializes standard tracing subscriber.
 pub fn init_standard_logger() {
     tracing_subscriber::registry()

examples-common/src/lib.rs

@@ -16,6 +16,10 @@ futures-util.workspace = true
 ractor = { version = "0.16", default-features = false, features = ["tokio_runtime"] }
 tokio = { workspace = true, features = ["rt-multi-thread", "time", "macros"] }
 
-env_logger.workspace = true
-log.workspace = true
+tokio-util = { workspace = true, features = ["rt", "time"] }
+color-eyre.workspace = true
+eyre.workspace = true
+examples-common.workspace = true
+
+tracing.workspace = true
 rand.workspace = true

websockets/chat/Cargo.toml

@@ -28,6 +28,10 @@ cd websockets/chat
 cargo run --bin websocket-chat-server
 ```
 
+Press Ctrl-C to stop the server. On Unix, SIGINT, SIGTERM, and SIGQUIT start shutdown. The server
+closes WebSocket sessions with close code 1001 (Going Away). After the HTTP server finishes graceful
+shutdown, the chat actor stops.
+
 ## Browser Client
 
 - Open in browser: <http://localhost:8080/>.

websockets/chat/README.md

@@ -11,7 +11,10 @@ use actix_files::{Files, NamedFile};
 use actix_web::{
     App, Error, HttpRequest, HttpResponse, HttpServer, Responder, middleware::Logger, web,
 };
+use examples_common::{init_standard_logger, watch_shutdown_signal};
+use eyre::WrapErr as _;
 use ractor::{Actor, ActorRef};
+use tokio_util::sync::CancellationToken;
 
 mod server;
 mod session;
@@ -24,11 +27,12 @@ async fn index() -> impl Responder {
 async fn chat_route(
     req: HttpRequest,
     stream: web::Payload,
-    srv: web::Data<ActorRef<server::Command>>,
+    web::ThinData(srv): web::ThinData<ActorRef<server::Command>>,
+    web::ThinData(shutdown): web::ThinData<CancellationToken>,
 ) -> Result<HttpResponse, Error> {
     let (response, session, messages) = actix_ws::handle(&req, stream)?;
 
-    tokio::task::spawn_local(session::chat_ws(srv.get_ref().clone(), session, messages));
+    tokio::task::spawn_local(session::chat_ws(srv, session, messages, shutdown));
 
     Ok(response)
 }
@@ -40,8 +44,9 @@ async fn get_count(count: web::Data<AtomicUsize>) -> impl Responder {
 }
 
 #[tokio::main]
-async fn main() -> std::io::Result<()> {
-    env_logger::init_from_env(env_logger::Env::new().default_filter_or("info"));
+async fn main() -> eyre::Result<()> {
+    color_eyre::install()?;
+    init_standard_logger();
 
     // set up applications state
     // keep a count of the number of visitors
@@ -51,27 +56,37 @@ async fn main() -> std::io::Result<()> {
     let (server, server_task) =
         server::ChatServer::spawn(None, server::ChatServer, app_state.clone())
             .await
-            .map_err(std::io::Error::other)?;
-    let shutdown = server.clone();
+            .wrap_err("failed to start chat actor")?;
+    let actor = server.clone();
+    let shutdown = CancellationToken::new();
+    watch_shutdown_signal(shutdown.clone());
+    let worker_shutdown = shutdown.clone();
 
-    log::info!("starting HTTP server at http://localhost:8080");
+    tracing::info!("starting HTTP server at http://localhost:8080");
 
-    let result = HttpServer::new(move || {
+    let http_server = HttpServer::new(move || {
         App::new()
             .app_data(web::Data::from(app_state.clone()))
-            .app_data(web::Data::new(server.clone()))
+            .app_data(web::ThinData(server.clone()))
+            .app_data(web::ThinData(worker_shutdown.clone()))
             .service(web::resource("/").to(index))
             .route("/count", web::get().to(get_count))
             .route("/ws", web::get().to(chat_route))
             .service(Files::new("/static", "./static"))
             .wrap(Logger::default())
     })
+    .shutdown_signal(shutdown.clone().cancelled_owned())
     .workers(2)
-    .bind(("127.0.0.1", 8080))?
-    .run()
-    .await;
+    .bind(("127.0.0.1", 8080));
+
+    let result = match http_server {
+        Ok(http_server) => http_server.run().await.wrap_err("HTTP server failed"),
+        Err(err) => Err(err).wrap_err("failed to bind HTTP server"),
+    };
+
+    shutdown.cancel();
+    actor.stop(None);
+    server_task.await.wrap_err("chat actor task failed")?;
 
-    shutdown.stop(None);
-    server_task.await.map_err(std::io::Error::other)?;
     result
 }

websockets/chat/src/main.rs

@@ -108,7 +108,7 @@ impl Actor for ChatServer {
     ) -> Result<(), ActorProcessingErr> {
         match message {
             Command::Connect { addr, reply } => {
-                println!("Someone joined");
+                tracing::info!("Someone joined");
 
                 // notify all users in same room
                 state.send_message("main", "Someone joined", 0);
@@ -127,7 +127,7 @@ impl Actor for ChatServer {
                 let _ = reply.send(id);
             }
             Command::Disconnect { id } => {
-                println!("Someone disconnected");
+                tracing::info!("Someone disconnected");
 
                 let mut rooms: Vec<String> = Vec::new();
 

websockets/chat/src/server.rs

@@ -1,14 +1,15 @@
 //! WebSocket connection handling and chat command parsing.
 //!
 //! Aggregates fragmented messages, checks heartbeats, and forwards chat messages through the actor.
-//! Removes the session from the actor when the connection ends.
+//! Closes connections on shutdown and removes each session from the actor when it ends.
 
 use std::time::{Duration, Instant};
 
 use actix_ws::AggregatedMessage;
 use futures_util::StreamExt as _;
 use ractor::{ActorProcessingErr, ActorRef};
 use tokio::sync::mpsc;
+use tokio_util::{future::FutureExt as _, sync::CancellationToken};
 
 use crate::server::Command;
 
@@ -20,6 +21,7 @@ pub async fn chat_ws(
     server: ActorRef<Command>,
     mut session: actix_ws::Session,
     messages: actix_ws::MessageStream,
+    shutdown: CancellationToken,
 ) {
     let (addr, mut inbox) = mpsc::unbounded_channel();
     let Ok(id) = ractor::call!(server, |reply| Command::Connect { addr, reply }) else {
@@ -39,6 +41,8 @@ pub async fn chat_ws(
 
     let reason = loop {
         tokio::select! {
+            _ = shutdown.cancelled() => break Some(actix_ws::CloseCode::Away.into()),
+
             incoming = messages.next() => {
                 match incoming {
                     Some(Ok(AggregatedMessage::Ping(bytes))) => {
@@ -52,16 +56,18 @@ pub async fn chat_ws(
                         }
                     }
                     Some(Ok(AggregatedMessage::Close(reason))) => break reason,
-                    Some(Ok(AggregatedMessage::Binary(_))) => log::warn!("unexpected binary message"),
+                    Some(Ok(AggregatedMessage::Binary(_))) => tracing::warn!("unexpected binary message"),
                     Some(Err(_)) | None => break None,
                 }
             }
+
             outgoing = inbox.recv() => {
                 match outgoing {
                     Some(text) => if session.text(text).await.is_err() { break None; },
                     None => break None,
                 }
             }
+
             _ = timer.tick() => {
                 if heartbeat.elapsed() > CLIENT_TIMEOUT || session.ping(b"").await.is_err() {
                     break None;
@@ -71,7 +77,9 @@ pub async fn chat_ws(
     };
 
     let _ = server.cast(Command::Disconnect { id });
-    let _ = session.close(reason).await;
+
+    // A slow peer must not hold up shutdown while the close frame is queued.
+    let _ = session.close(reason).timeout(Duration::from_secs(1)).await;
 }
 
 async fn process_text(