Axum 0.8 Tokio Tower Bounded Graceful Drain & WebSocket Cancellation Protocol

--- description: Axum 0.8 Tokio Tower Bounded Graceful Drain & WebSocket Cancellation Protocol - Enforces deterministic, bounded connection draining during SIGTERM/SIGINT teardown in Axum 0.8 and Tokio runtimes. Requires tokio_util CancellationToken propagation across nested Tower service layers and coordinated Close frame broadcasts to stateful WebSockets before runtime drop to prevent dropped in-flight requests and orphaned socket descriptors. globs: src/**/*.rs, crates/**/*.rs alwaysApply: true verified_by: "agentrules.org" license: "CC-BY-4.0" source: "https://agentrules.org/rules/axum-08-tokio-tower-graceful-drain-websocket-cancellation" --- # Axum 0.8 Tokio Tower Bounded Graceful Drain & WebSocket Cancellation Protocol > **Severity**: CRITICAL | **Category**: backend > **Anti-Pattern Warning**: Unbounded Axum serve_with_graceful_shutdown causing abruptly aborted TCP sockets and unnotified WebSocket client disconnections ## System Constraints You are a Principal Rust Distributed Systems Engineer enforcing Axum 0.8 and Tokio graceful shutdown standards. You must mandate bounded shutdown timeouts, CancellationToken coordination across nested Tower middleware, and deterministic WebSocket Close frame handshakes before Tokio runtime shutdown. ## Strict Guidelines - Engineers MUST ALWAYS bound axum::serve graceful shutdown using tokio::time::timeout to prevent hung connections or non-terminating background tasks from holding the process indefinitely. - Engineers MUST instantiate a top-level tokio_util::sync::CancellationToken and clone it into Axum AppState or request extensions to signal shutdown to stateful long-running connections. - WebSocket handler loops MUST ALWAYS tokio::select! on the CancellationToken alongside socket frames, actively sending a CloseCode::Restart frame (or 1001 Going Away) and breaking the loop when triggered. - Engineers MUST NOT call std::process::exit directly inside signal handlers or application code, as it bypasses RAII drop handlers, pending Tower flush phases, and tracing flushes. - Tower services MUST implement graceful draining by coupling axum::serve(listener, router).with_graceful_shutdown(shutdown_signal) with a final tokio::time::sleep or join barrier for spawned connection tasks. ## Bad vs Good Implementations ### Graceful Draining and Coordinated WebSocket Teardown #### Bad Antipattern (AI Common Hallucination) > **Why it fails**: axum::serve is executed without with_graceful_shutdown, meaning incoming SIGTERM terminates the process abruptly. Stateful WebSockets receive no WebSocket close frame, client TCP streams are reset (RST), and in-flight requests running through Tower layers are truncated without completing their execution or dropping clean-up handles. ```rust use axum::{routing::get, Router}; use tokio::net::TcpListener; #[tokio::main] async fn main() { let app = Router::new().route("/ws", get(ws_handler)); let listener = TcpListener::bind("0.0.0.0:3000").await.unwrap(); // ANTI-PATTERN: Serve without graceful shutdown listener or bounded timeout. // Incoming SIGTERM drops runtime immediately; WebSocket connections get abruptly // terminated with TCP RST, and in-flight Tower middleware state is lost. axum::serve(listener, app).await.unwrap(); } async fn ws_handler(ws: axum::extract::ws::WebSocketUpgrade) -> axum::response::Response { ws.on_upgrade(|mut socket| async move { while let Some(Ok(msg)) = socket.recv().await { let _ = socket.send(msg).await; } }) } ``` #### Production Standard (Correct) > **Standard**: Implements dual OS signal traps (SIGINT/SIGTERM), triggers a top-level CancellationToken that alerts all active WebSocket select loops to issue 1001 Going Away close frames, and wraps axum::serve with a strict 15-second timeout to prevent deadlocks. ```rust use axum::{extract::{ws::{Message, WebSocket, WebSocketUpgrade}, State}, routing::get, Router}; use std::{sync::Arc, time::Duration}; use tokio::{net::TcpListener, signal, sync::broadcast, time::timeout}; use tokio_util::sync::CancellationToken; #[derive(Clone)] struct AppState { cancel_token: CancellationToken, } #[tokio::main] async fn main() -> Result<(), Box<dyn std::error::Error>> { let cancel_token = CancellationToken::new(); let state = AppState { cancel_token: cancel_token.clone() }; let app = Router::new().route("/ws", get(ws_handler)).with_state(state); let listener = TcpListener::bind("0.0.0.0:3000").await?; let drain_token = cancel_token.clone(); let server = axum::serve(listener, app) .with_graceful_shutdown(async move { wait_for_shutdown_signal().await; drain_token.cancel(); // Broadcast cancellation to all active WebSocket tasks }); // Bound draining phase to 15 seconds max if let Err(_) = timeout(Duration::from_secs(15), server).await { eprintln!("Server graceful drain timed out; forcing teardown"); } Ok(()) } async fn ws_handler(State(state): State<AppState>, ws: WebSocketUpgrade) -> axum::response::Response { ws.on_upgrade(move |socket| handle_ws(socket, state.cancel_token)) } async fn handle_ws(mut socket: WebSocket, cancel: CancellationToken) { loop { tokio::select! { _ = cancel.cancelled() => { // Protocol-compliant close frame before runtime teardown let _ = socket.send(Message::Close(Some(axum::extract::ws::CloseFrame { code: axum::extract::ws::close_code::AWAY, reason: "Server shutting down".into(), }))).await; break; } msg = socket.recv() => { match msg { Some(Ok(Message::Close(_))) | None => break, Some(Ok(msg)) => { let _ = socket.send(msg).await; } Some(Err(_)) => break, } } } } } async fn wait_for_shutdown_signal() { let ctrl_c = async { signal::ctrl_c().await.expect("failed to install Ctrl+C handler"); }; #[cfg(unix)] let terminate = async { signal::unix::signal(signal::unix::SignalKind::terminate()) .expect("failed to install SIGTERM handler") .recv().await; }; #[cfg(not(unix))] let terminate = std::future::pending::<()>(); tokio::select! { _ = ctrl_c => {}, _ = terminate => {} } } ```

Architectural Specification & Invariant Rationale

Architectural Invariant Design: Axum 0.8 Tokio Tower Bounded Graceful Drain & WebSocket Cancellation Protocol

Axum 0.8 Tokio Tower Bounded Graceful Drain & WebSocket Cancellation Protocol defines an industrial-grade behavioral standard for backend engineering. Enforces deterministic, bounded connection draining during SIGTERM/SIGINT teardown in Axum 0.8 and Tokio runtimes. Requires tokio_util CancellationToken propagation across nested Tower service layers and coordinated Close frame broadcasts to stateful WebSockets before runtime drop to prevent dropped in-flight requests and orphaned socket descriptors. Modern software architectures require strict adherence to deterministic patterns to guarantee high availability, maintainability, and zero security vulnerabilities.

Framework Paradigm Shift & Drift MechanicsLLMs are trained on millions of open-source repositories containing outdated code, legacy idioms, and unpatched antipatterns. When coding assistants generate implementations for this stack, they default to common statistical tokens: Unbounded Axum serve_with_graceful_shutdown causing abruptly aborted TCP sockets and unnotified WebSocket client disconnections Without an explicit system instruction file, developers spend hours debugging silent failures.

Failure Modes & Production Blast Radius

Hazard Analysis: What Breaks Without This Rule

The following failure modes have been cataloged from hundreds of hours of production code reviews where AI assistants produced non-compliant implementations:

Failure Mode	AI Symptom (The Anti-Pattern)	Severity	Production Impact & Remedy
AI Hallucination & Deprecation Drift	Unbounded Axum serve_with_graceful_shutdown causing abruptly aborted TCP sockets and unnotified WebSocket client disconnections	Critical	Violating this industrial standard introduces regression bugs, security risks, and unexpected runtime failures in production. Enforce explicit behavioral system directives across all AI coding assistants in your repository.
Silent Architecture Inconsistency	AI assistants produce mixed paradigms across files, creating fragmented codebase conventions.	Medium	Increases cognitive load during code reviews and leads to maintenance debt across engineering teams. Pin project conventions in .cursor/rules, CLAUDE.md, and AGENTS.md so all team agents follow identical guardrails.

Architectural Directives in Detail

Step-by-Step Production Directives

Each rule constraint below is injected into the AI agent's system prompt to establish deterministic boundaries:

Engineers MUST ALWAYS bound axum

Engineers MUST ALWAYS bound axum::serve graceful shutdown using tokio::time::timeout to prevent hung connections or non-terminating background tasks from holding the process indefinitely.

Rationale:AI models frequently generate unconstrained code here because statistical language distributions lack project-specific architectural context. Mandating this directive ensures that generated code adheres to verified industrial production requirements.