Axum 0.8 Tower Graceful Connection Drain & WebSocket Handshake Fencing

--- description: Axum 0.8 Tower Graceful Connection Drain & WebSocket Handshake Fencing - Enforces deterministic zero-downtime rolling deploys in Axum 0.8 by decoupling TCP connection draining from stateful HTTP upgrade handshakes. Mandates an atomic handshake fence to reject new WebSocket upgrades with 503 during shutdown, alongside bounded dual-phase horizons that dispatch RFC 6455 1001 Close frames to active sessions before hyper socket termination. globs: src/**/*.rs, Cargo.toml alwaysApply: true verified_by: "agentrules.org" license: "CC-BY-4.0" source: "https://agentrules.org/rules/axum-graceful-drain-websocket-fencing" --- # Axum 0.8 Tower Graceful Connection Drain & WebSocket Handshake Fencing > **Severity**: CRITICAL | **Category**: backend > **Anti-Pattern Warning**: Unbounded WebSocket Upgrades Stalling Axum Graceful Drains and Causing SIGKILL Dropped Frames ## System Constraints You are a Principal Systems Architect enforcing mission-critical network invariants in Axum 0.8 and Tower. You must prevent rolling deploy timeouts and abrupt TCP socket severance by fencing new WebSocket upgrades upon shutdown initiation, dispatching bounded RFC 6455 close frames, and isolating short-lived HTTP drain from long-lived upgraded streams. ## Strict Guidelines - Engineers MUST ALWAYS implement a shared `tokio_util::sync::CancellationToken` or watch channel to signal graceful shutdown across both the Axum server runner and upgrade handlers. - Engineers MUST ALWAYS fence incoming WebSocket handshakes by checking shutdown state prior to completing `WebSocketUpgrade::on_upgrade`, returning `StatusCode::SERVICE_UNAVAILABLE` with `Connection: close` once shutdown has begun. - Engineers MUST ALWAYS bound WebSocket session lifetimes with a dual-phase timeout horizon, dispatching WS 1001 ('Going Away') close frames before hyper forces TCP connection termination. - Engineers MUST NOT pass unfenced handlers directly to `axum::serve(...).with_graceful_shutdown(...)` when handling stateful upgrades, as hyper's internal graceful drain will wait indefinitely for upgrade streams unless externally terminated. - Components MUST track active WebSocket sessions using `tokio_util::task::TaskTracker` to guarantee deterministic drain completion before process exit. ## Bad vs Good Implementations ### Unfenced WebSocket Upgrade and Unbounded Graceful Shutdown #### Bad Antipattern (AI Common Hallucination) > **Why it fails**: When SIGTERM is received, axum::serve attempts to drain connections. However, WebSocket upgrades break out of hyper's standard HTTP request-response cycle. Without an explicit cancellation fence, in-flight handshakes still succeed and active WebSockets never exit their loops, causing orchestrators (e.g., Kubernetes) to forcefully issue SIGKILL and drop client packets. ```rust use axum::{routing::get, Router, extract::ws::{WebSocketUpgrade, WebSocket}}; use tokio::net::TcpListener; async fn ws_handler(ws: WebSocketUpgrade) -> axum::response::Response { // ANTI-PATTERN: Accepts upgrades during shutdown, loops infinitely without close horizon ws.on_upgrade(|mut socket: WebSocket| async move { while let Some(Ok(msg)) = socket.recv().await { let _ = socket.send(msg).await; } }) } pub async fn run_server() { let app = Router::new().route("/ws", get(ws_handler)); let listener = TcpListener::bind("0.0.0.0:3000").await.unwrap(); // Hyper waits forever for active upgraded TCP connections to terminate naturally axum::serve(listener, app) .with_graceful_shutdown(async { tokio::signal::ctrl_c().await.unwrap(); }) .await .unwrap(); } ``` #### Production Standard (Correct) > **Standard**: This pattern isolates the HTTP upgrade handshake from the server drain horizon. It uses a CancellationToken to return 503 Service Unavailable to any handshake received after shutdown begins. Active WebSocket connections receive an RFC 6455 1001 Close frame with a 5-second per-socket budget, and the TaskTracker enforces a global 15-second bounded horizon before hyper exits, guaranteeing zero-downtime rolling deploys. ```rust use axum::{ extract::{ws::{close_code, CloseFrame, Message, WebSocket, WebSocketUpgrade}, State}, http::{header, StatusCode}, response::{IntoResponse, Response}, routing::get, Router, }; use std::{borrow::Cow, sync::Arc, time::Duration}; use tokio::{net::TcpListener, time::timeout}; use tokio_util::{sync::CancellationToken, task::TaskTracker}; #[derive(Clone)] struct AppState { shutdown_token: CancellationToken, tracker: TaskTracker, } async fn ws_handler(ws: WebSocketUpgrade, State(state): State<AppState>) -> Response { // INVARIANT: Reject incoming upgrades immediately once shutdown starts if state.shutdown_token.is_cancelled() { return ( StatusCode::SERVICE_UNAVAILABLE, [(header::CONNECTION, "close")], "Server is shutting down", ).into_response(); } let token = state.shutdown_token.clone(); state.tracker.track_future(async move { ws.on_upgrade(move |socket| handle_socket(socket, token)) }).await } async fn handle_socket(mut socket: WebSocket, token: CancellationToken) { loop { tokio::select! { _ = token.cancelled() => { // INVARIANT: Send WS Close frame 1001 (Going Away) with bounded timeout let close_frame = CloseFrame { code: close_code::AWAY, reason: Cow::Borrowed("Server draining") }; let _ = timeout(Duration::from_secs(5), socket.send(Message::Close(Some(close_frame)))).await; let _ = socket.close().await; break; } msg = socket.recv() => { match msg { Some(Ok(Message::Close(_))) | None => break, Some(Ok(msg)) => { if socket.send(msg).await.is_err() { break; } } Some(Err(_)) => break, } } } } } pub async fn run_server() { let token = CancellationToken::new(); let tracker = TaskTracker::new(); let state = AppState { shutdown_token: token.clone(), tracker: tracker.clone() }; let app = Router::new().route("/ws", get(ws_handler)).with_state(state); let listener = TcpListener::bind("0.0.0.0:3000").await.unwrap(); let server = axum::serve(listener, app); let shutdown_signal = { let token = token.clone(); async move { tokio::signal::ctrl_c().await.ok(); token.cancel(); // Fence new upgrades } }; // Run server with graceful TCP drain let _ = server.with_graceful_shutdown(shutdown_signal).await; // Enforce hard bounded horizon for active WebSocket sessions tracker.close(); if timeout(Duration::from_secs(15), tracker.wait()).await.is_err() { eprintln!("WebSocket drain horizon exceeded; forcing socket termination"); } } ```

Architectural Specification & Invariant Rationale

Architectural Invariant Design: Axum 0.8 Tower Graceful Connection Drain & WebSocket Handshake Fencing

Axum 0.8 Tower Graceful Connection Drain & WebSocket Handshake Fencing defines an industrial-grade behavioral standard for backend engineering. Enforces deterministic zero-downtime rolling deploys in Axum 0.8 by decoupling TCP connection draining from stateful HTTP upgrade handshakes. Mandates an atomic handshake fence to reject new WebSocket upgrades with 503 during shutdown, alongside bounded dual-phase horizons that dispatch RFC 6455 1001 Close frames to active sessions before hyper socket termination. 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 WebSocket Upgrades Stalling Axum Graceful Drains and Causing SIGKILL Dropped Frames 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 WebSocket Upgrades Stalling Axum Graceful Drains and Causing SIGKILL Dropped Frames	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 implement a shared `tokio_util

Engineers MUST ALWAYS implement a shared `tokio_util::sync::CancellationToken` or watch channel to signal graceful shutdown across both the Axum server runner and upgrade handlers.

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.