How to Handle Connection Close and Graceful Shutdown in iroh

Call router.shutdown().await to trigger a token-driven cascade that cancels all background tasks, closes the socket, and shuts down the runtime after a 100 ms grace period.

The iroh crate (from n0-computer/iroh) implements a fully-async, token-driven shutdown flow that guarantees every background task, connection, and runtime resource finishes cleanly before the process exits. Understanding this coordination mechanism is essential for building robust applications that avoid dangling connections or leaked tasks.

Understanding the Shutdown Architecture

iroh coordinates graceful shutdown through a hierarchy of cancellation tokens and stateful components. Rather than aborting tasks abruptly, the system propagates a shutdown signal that allows each component to perform its own cleanup.

Router

The Router is the top-level object that owns the network stack, protocol handlers, and the runtime. In iroh/src/protocol.rs, the Router::shutdown method (line 415) first checks Router::is_shutdown to prevent double-shutdown, then cancels the internal shutdown token (line 424). This single call initiates the entire shutdown cascade.

ShutdownState

The ShutdownState struct (defined in iroh/src/socket.rs at line 347) holds two distinct cancellation tokens: at_close_start (signals that the endpoint is beginning to close) and at_endpoint_closed (signals the endpoint has fully closed). When the router initiates shutdown, it cancels at_close_start (line 1140), propagating the signal to every task holding a child token.

CancellationToken

The implementation relies on CancellationToken from the tokio-util crate. This cheap, cloneable token allows any task to observe a graceful-shutdown request. Every long-living task—such as the relay actor, remote map, and transport actors—holds a child token (e.g., iroh/src/socket/transports/relay/actor.rs line 424). When the parent token is cancelled, the child’s cancelled() future resolves, causing the task to break out of its work loop cleanly.

The Graceful Shutdown Sequence

When you invoke graceful shutdown, the following sequence occurs:

1. Initiation: User code calls router.shutdown().await.
2. State Check: Router::shutdown verifies is_shutdown is false, then cancels the router-wide token.
3. Signal Propagation: ShutdownState::at_close_start is cancelled, notifying all listening tasks.
4. Protocol Cleanup: Each registered ProtocolHandler receives a shutdown request concurrently (protocol.rs lines 398-404), allowing protocols to close their own connections.
5. Socket Drain: The Socket task sees its shutdown_token cancelled (socket.rs line 1225), stops accepting new packets, and drains remaining work (observed in the select! branch at line 1507).
6. Grace Period: The system waits 100 ms (socket.rs lines 1173-1181) for tasks to finish before forcing runtime shutdown.
7. Runtime Termination: Finally, self.runtime.shutdown().await is called (socket.rs line 1192), and the closed flag on ShutdownState is set (line 1194).

Implementing Graceful Shutdown in Practice

Shutting Down the Router

The primary entry point for application shutdown is the Router itself. This pattern works for CLI tools, servers, and test harnesses:

use iroh::protocol::Router;
use iroh::runtime::Runtime;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // Build and start the router
    let router = Router::builder()
        .listen("tcp://0.0.0.0:12345")?
        .build()
        .await?;

    // ... application logic ...

    // Trigger graceful shutdown
    router.shutdown().await?;
    // All connections closed, background tasks cancelled, runtime shut down
    Ok(())
}

Observing Shutdown in Custom Tasks

When spawning custom background tasks, you must observe the shutdown token to participate in the graceful shutdown sequence:

use tokio::select;
use tokio_util::sync::CancellationToken;
use std::time::Duration;

async fn my_background_task(shutdown: CancellationToken) {
    loop {
        select! {
            // Normal work
            _ = tokio::time::sleep(Duration::from_millis(100)) => {
                // Perform work
            }
            
            // React to graceful shutdown
            _ = shutdown.cancelled() => {
                // Perform last-minute cleanup
                log::info!("my_background_task: shutting down gracefully");
                break;
            }
        }
    }
}

Spawning Tasks with Shutdown Tokens

Embed your tasks into the router's lifecycle by cloning the shutdown token:

// Clone the token from the router
let shutdown = router.shutdown_token.clone();

// Spawn task with shutdown awareness
tokio::spawn(my_background_task(shutdown));

When router.shutdown().await executes, the token is cancelled, causing your task to exit after completing its current iteration.

Key Source Files and Implementation Details

File	Purpose	Key Location
iroh/src/protocol.rs	Router definition and high-level shutdown	Router::shutdown at line 415; token cancellation at line 424
iroh/src/socket.rs	Socket abstraction and ShutdownState	ShutdownState struct at line 347; grace period at lines 1173-1181
iroh/src/socket/transports/relay/actor.rs	Transport actor implementation	Child token usage at line 424
iroh/src/runtime.rs	Tokio runtime wrapper	Runtime shutdown coordination

Summary

• Always call router.shutdown().await to initiate a clean shutdown sequence.
• Use CancellationToken from tokio-util to make custom tasks shutdown-aware.
• The 100 ms grace period prevents indefinite hangs while allowing tasks to finish.
• Protocol handlers receive concurrent shutdown notifications to close connections cleanly.
• Source files in iroh/src/protocol.rs and iroh/src/socket.rs contain the core implementation.

Frequently Asked Questions

What happens if I don't call router.shutdown()?

If the Router is dropped without explicit shutdown, the Drop implementation on Socket (socket.rs line 1225) cancels the shutdown token. However, relying on drop behavior skips the coordinated 100 ms grace period and may interrupt active protocol handlers mid-operation. Always call shutdown().await for production code.

How long does graceful shutdown wait for tasks to finish?

The implementation waits 100 milliseconds after signalling shutdown before forcing the runtime to stop (socket.rs lines 1173-1181). If tasks exceed this timeout, a warning is logged but shutdown proceeds to prevent process deadlock.

Can I run custom cleanup logic during shutdown?

Yes. Implement the ProtocolHandler trait and register your handler with the router. The shutdown method is called concurrently with other handlers (protocol.rs lines 398-404), giving you an opportunity to close connections, flush buffers, and cancel internal tokens before the runtime terminates.

Is the shutdown sequence cancellation-safe?

Absolutely. The system uses tokio_util::sync::CancellationToken throughout, which is designed for async cancellation safety. Tasks observe cancellation via the cancelled() future in a select! branch, allowing them to exit loops cleanly without losing data or leaking resources.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}