mirror of
https://github.com/tokio-rs/tokio.git
synced 2026-08-07 00:00:09 +02:00
121 lines
3.6 KiB
Rust
121 lines
3.6 KiB
Rust
//! 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<dyn Error>> {
|
|
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;
|
|
}
|