From 6752f50154025a1aa2c231b643cdb78bb4c3892f Mon Sep 17 00:00:00 2001 From: Sim-hu <93906619+Sim-hu@users.noreply.github.com> Date: Mon, 30 Mar 2026 18:05:53 +0900 Subject: [PATCH] examples: add graceful shutdown example (#7962) --- examples/Cargo.toml | 4 ++ examples/graceful-shutdown.rs | 120 ++++++++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+) create mode 100644 examples/graceful-shutdown.rs diff --git a/examples/Cargo.toml b/examples/Cargo.toml index 22622670f..056ab17f7 100644 --- a/examples/Cargo.toml +++ b/examples/Cargo.toml @@ -49,6 +49,10 @@ path = "connect-udp.rs" name = "echo-tcp" path = "echo-tcp.rs" +[[example]] +name = "graceful-shutdown" +path = "graceful-shutdown.rs" + [[example]] name = "echo-udp" path = "echo-udp.rs" diff --git a/examples/graceful-shutdown.rs b/examples/graceful-shutdown.rs new file mode 100644 index 000000000..8faf63b63 --- /dev/null +++ b/examples/graceful-shutdown.rs @@ -0,0 +1,120 @@ +//! Graceful shutdown example. +//! +//! This example follows the same approach described in the +//! [Graceful Shutdown tutorial](https://tokio.rs/tokio/topics/shutdown): +//! +//! - A [`CancellationToken`] tells tasks to stop accepting new work. +//! - A [`TaskTracker`] waits for in-flight work to complete. +//! +//! It runs a TCP echo server on `127.0.0.1:6142`. When Ctrl+C is +//! pressed, the server stops accepting connections and waits for all +//! active connections to finish before exiting. +//! +//! Start the server: +//! +//! cargo run --example graceful-shutdown +//! +//! Then connect with: +//! +//! nc 127.0.0.1 6142 +//! +//! Press Ctrl+C on the server to trigger a graceful shutdown. + +#![warn(rust_2018_idioms)] + +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::time::{self, Duration}; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; + +use std::error::Error; +use std::net::SocketAddr; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let listener = TcpListener::bind("127.0.0.1:6142").await?; + println!("listening on 127.0.0.1:6142"); + + let token = CancellationToken::new(); + let tracker = TaskTracker::new(); + + loop { + tokio::select! { + result = listener.accept() => { + let (socket, addr) = match result { + Ok(conn) => conn, + Err(e) => { + // Transient errors (e.g. fd exhaustion) are recoverable, + // so we log and continue. A production server might add a + // backoff or break on fatal errors to avoid a busy loop. + eprintln!("failed to accept: {e}"); + continue; + } + }; + println!("accepted connection from {addr}"); + + let token = token.clone(); + tracker.spawn(handle_connection(socket, addr, token)); + } + + _ = tokio::signal::ctrl_c() => { + println!("\nshutdown signal received, waiting for connections to finish"); + break; + } + } + } + + // Signal all tasks to stop and wait for them to complete. + token.cancel(); + tracker.close(); + tracker.wait().await; + + println!("shutdown complete"); + Ok(()) +} + +async fn handle_connection(mut socket: TcpStream, addr: SocketAddr, token: CancellationToken) { + tokio::select! { + _ = echo(&mut socket) => {} + _ = token.cancelled() => { + notify_shutdown(&mut socket).await; + } + } + + println!("connection from {addr} closed"); +} + +/// Reads lines from the client and writes them back. +/// +/// Called for every accepted connection. Runs until the client disconnects +/// or a read/write error occurs. +async fn echo(socket: &mut TcpStream) { + let (reader, mut writer) = socket.split(); + let mut reader = BufReader::new(reader); + let mut line = String::new(); + + loop { + match reader.read_line(&mut line).await { + Ok(0) | Err(_) => return, + Ok(_) => { + if writer.write_all(line.as_bytes()).await.is_err() { + return; + } + line.clear(); + } + } + } +} + +/// Sends a shutdown notice to the client before closing the connection. +/// +/// Called when the cancellation token fires. Uses a timeout so that a +/// slow or unresponsive client cannot hold up the server shutdown. +async fn notify_shutdown(socket: &mut TcpStream) { + let _ = time::timeout( + Duration::from_secs(1), + socket.write_all(b"server shutting down\n"), + ) + .await; +}