Tokio / Chapter 4 of 6

Graceful shutdown and testing

Shut down a Tokio TCP server gracefully and test async code with controlled time. Examples cover cancellation and task cleanup.

2 topicsUpdated 22 Oct 2026

14 / 26

Graceful TCP shutdown#

Stop accepting, request cancellation, and wait for active connections before exiting.

Ctrl+C cancels a shared token. The listener checks cancellation while waiting for a connection slot or socket, and the task tracker waits for active connections. Add tokio-util as shown in Cancellation and task lifetimes.

Rustsrc/main.rs — bounded graceful echo server

use std::{sync::Arc, time::Duration};
use tokio::{io::{self, AsyncWriteExt}, net::TcpListener, sync::Semaphore};
use tokio_util::{sync::CancellationToken, task::TaskTracker};

#[tokio::main]
async fn main() -> io::Result<()> {
    let listener = TcpListener::bind("127.0.0.1:7003").await?;
    let stop = CancellationToken::new();
    let tracker = TaskTracker::new();
    let slots = Arc::new(Semaphore::new(100));
    let on_signal = stop.clone();
    tokio::spawn(async move {
        if let Err(error) = tokio::signal::ctrl_c().await {
            eprintln!("signal: {error}");
        }
        on_signal.cancel();
    });

    loop {
        let permit = tokio::select! {
            _ = stop.cancelled() => break,
            result = Arc::clone(&slots).acquire_owned() =>
                result.expect("semaphore is open"),
        };
        let stream = tokio::select! {
            _ = stop.cancelled() => break,
            accepted = listener.accept() => accepted?.0,
        };
        let child_stop = stop.clone();
        tracker.spawn(async move {
            let _permit = permit;
            let (mut read, mut write) = stream.into_split();
            tokio::select! {
                _ = child_stop.cancelled() => {
                    let _ = write.shutdown().await;
                }
                result = io::copy(&mut read, &mut write) => {
                    if let Err(error) = result { eprintln!("client: {error}"); }
                }
            }
        });
    }
    stop.cancel();
    tracker.close();
    if tokio::time::timeout(Duration::from_secs(5), tracker.wait())
        .await.is_err()
    {
        eprintln!("connections did not close before the deadline");
    }
    Ok(())
}

Shutdown deadline

On timeout, returning from main drops the runtime and cancels remaining tasks. If clients need to finish critical writes, track that state and choose a longer deadline.

15 / 26

Async tests and controlled time#

Use #[tokio::test] to create a runtime per test. Pause virtual time to test timers quickly and deterministically.

Add test-util to the Tokio features for start_paused and advance. Tests can use temporary loopback ports (127.0.0.1:0) to avoid collisions; assert protocol behavior rather than relying on sleep timing.

TOMLCargo.toml test feature

[dev-dependencies]
tokio = { version = "1", features = ["full", "test-util"] }

RustVirtual time: no real 10-second wait

use std::time::Duration;

#[tokio::test(start_paused = true)]
async fn retries_wait_before_repeating() {
    let task = tokio::spawn(async {
        tokio::time::sleep(Duration::from_secs(10)).await;
        42
    });
    tokio::task::yield_now().await; // register the sleep
    tokio::time::advance(Duration::from_secs(10)).await;
    assert_eq!(task.await.unwrap(), 42);
}

RustIsolated TCP client/server test

use tokio::{io::{AsyncReadExt, AsyncWriteExt}, net::{TcpListener, TcpStream}};

#[tokio::test]
async fn echo_round_trip() -> std::io::Result<()> {
    let server = TcpListener::bind("127.0.0.1:0").await?;
    let addr = server.local_addr()?;
    let task = tokio::spawn(async move {
        let (mut socket, _) = server.accept().await.unwrap();
        let mut bytes = [0u8; 4];
        socket.read_exact(&mut bytes).await.unwrap();
        socket.write_all(&bytes).await.unwrap();
    });
    let mut client = TcpStream::connect(addr).await?;
    client.write_all(b"ping").await?;
    let mut reply = [0u8; 4];
    client.read_exact(&mut reply).await?;
    assert_eq!(&reply, b"ping");
    task.await.unwrap();
    Ok(())
}

ShellRun tests

cargo test
cargo test echo_round_trip -- --nocapture

Paused time caveat

A paused clock auto-advances when the runtime has no other ready work. advance() jumps the clock; let spawned tasks register their timers first, as the yield above does.

Graceful TCP shutdown#

Async tests and controlled time#

Help improve CodeRef