Skip to Content
Course content

158: Graceful Shutdown in Tokio Applications

Click on the "Edit" button in the top corner of the screen to edit your slide content.

A few years ago, I was reviewing a production incident for a colleague who had built a high-throughput payment processing gateway. Every time they deployed a new version to their Kubernetes cluster, they noticed a spike in "incomplete transaction" errors in their database. The pods were being terminated, and while Kubernetes sends a SIGTERM to give the app time to wrap up, the app was just... vanishing. It was essentially pulling the plug on the server while it was in the middle of writing a critical record to the DB. The developer had assumed that since the process ended, everything just "stopped," but they weren't actually managing the lifecycle of their asynchronous tasks.

In a simple CLI tool, letting the process die instantly is usually fine. But in a production service, you want a graceful shutdown. This means stopping the acceptance of new requests, allowing current requests to finish their work, flushing buffers, and closing database connections cleanly. In Tokio, this isn't automatic because main is the boss; once main returns, the entire runtime shuts down and every other spawned task is dropped immediately, regardless of whether it's at a .await point or in the middle of a critical operation.

Capturing the Termination Signal

The first thing we need is a way to know that the outside world wants us to stop. On Unix systems, this is usually a SIGINT (Ctrl+C) or a SIGTERM (sent by orchestrators like Kubernetes). Tokio provides a convenient way to handle this without diving into raw libc calls. You can use tokio::signal::ctrl_c() for local development, but for production, you'll want to handle SIGTERM as well.

async fn shutdown_signal() {
    let ctrl_c = async {
        tokio::signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
            .expect("failed to install signal handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }
}

I usually wrap this in a helper function like the one above. By using tokio::select!, the first signal to arrive wins, and our program moves forward into the shutdown logic.

Broadcasting the Stop Order

Once main receives that signal, it needs to tell every background task—your telemetry loops, your request handlers, your queue consumers—to stop what they're doing. I strongly recommend using CancellationToken from the tokio-util crate. It's much cleaner than trying to pass around mpsc channels or atomic booleans.

The pattern is simple: you create a token, clone it for every task you spawn, and then call token.cancel() when the shutdown signal hits. Inside your tasks, you use tokio::select! to race your actual work against the token's cancelled() future.

use tokio_util::sync::CancellationToken;

async fn worker(id: u32, token: CancellationToken) {
    loop {
        tokio::select! {
            _ = token.cancelled() => {
                println!("Worker {} shutting down gracefully", id);
                return;
            }
            _ = tokio::time::sleep(std::time::Duration::from_secs(1)) => {
                println!("Worker {} doing some critical work...", id);
            }
        }
    }
}

The beauty of this approach is that it's non-blocking. Your worker doesn't have to periodically "poll" a boolean; it's woken up immediately by the runtime the moment the token is cancelled.

Waiting for the Last Task to Exit

Here is where most people trip up: calling token.cancel() isn't enough. If main finishes immediately after cancelling the token, the process exits and your workers are killed before they can even run their cleanup code. You need a way to wait for the tasks to actually finish.

I prefer using a JoinSet for this. A JoinSet allows you to spawn multiple tasks and then wait for all of them to resolve. If you're using an older version of Tokio or a different architecture, you can use a mpsc channel where the workers hold a sender; when all senders are dropped, the receiver in main will return None, signalling that everyone has left the building.

When you combine these three pieces—the signal listener, the cancellation token, and the task tracker—you get a robust shutdown sequence. You stop accepting new work, signal the existing work to wrap up, and block the main thread just long enough for the cleanup to complete. It's the difference between a crash and a professional exit.




📋 Practical Task

Exercise: Implementing a Graceful File-Writer Service

You are tasked with building a small background service that simulates writing logs to a file. The service must ensure that no data is lost when the application is shut down.

Requirements:

  • Create a Writer task that loops every 500ms, simulating a "write" operation by printing "Writing chunk [N] to disk..." to the console.
  • Use tokio_util::sync::CancellationToken to signal the Writer task to stop.
  • The Writer task must perform a "cleanup" operation when it receives the cancellation signal (print "Flushing buffers and closing file handle...") before returning.
  • In your main function, implement a shutdown signal handler using tokio::signal::ctrl_c().
  • Use a tokio::task::JoinSet to track the Writer task and ensure main does not exit until the Writer has completed its cleanup.

Success Criteria: When you press Ctrl+C, you should see the "Flushing buffers" message appear before the program exits. If the program exits immediately without that message, your shutdown is not graceful.

Rating
0 0

There are no comments for now.

to be the first to leave a comment.