Axum 0.8 Graceful WebSocket Drain-Fencing Standard

--- description: Axum 0.8 Graceful WebSocket Drain-Fencing Standard - Enforces deterministic drain-fencing and cooperative task teardown for bi-directional WebSocket streams in Axum 0.8 during cluster rollouts. Prevents dropped frames, dangling Tokio tasks, and broker desynchronization by decoupling socket upgrades from unmonitored runtime tasks using tokio_util::task::TaskTracker and broadcast cancellation tokens. globs: src/**/*.rs alwaysApply: true verified_by: "agentrules.org" license: "CC-BY-4.0" source: "https://agentrules.org/rules/axum-0-8-graceful-websocket-drain-fencing" --- # Axum 0.8 Graceful WebSocket Drain-Fencing Standard > **Severity**: CRITICAL | **Category**: backend > **Anti-Pattern Warning**: Spawning Unfenced Tokio Tasks for WebSocket Upgrades Without Cooperative Drain Propagation ## System Constraints You are a Principal Rust Systems Architect enforcing the Axum 0.8 Graceful WebSocket Drain-Fencing Standard. You mandate that no WebSocket upgrade handler spawns unmonitored Tokio tasks. All bi-directional stream loops must bind to a shared `tokio_util::sync::CancellationToken` and register against an orchestrated `tokio_util::task::TaskTracker`. When SIGTERM/SIGINT is received, the server must stop accepting incoming handshakes, notify existing clients via WebSocket Close frames (Code 1001 Going Away), flush pending outbound buffers with a strict deadline, and cleanly await task termination before the Tokio runtime exits. ## Strict Guidelines - Engineers MUST ALWAYS track every spawned WebSocket connection task using `tokio_util::task::TaskTracker`. - Engineers MUST ALWAYS pass a child or cloned `tokio_util::sync::CancellationToken` to every bi-directional streaming split loop (`Stream` / `Sink`). - Engineers MUST NOT invoke `tokio::spawn` inside Axum `WebSocketUpgrade::on_upgrade` without registering the future in the application's drain fence. - Code SHOULD explicitly transmit a WebSocket `Message::Close` frame with code `1001` (Going Away) upon receiving a cancellation signal prior to dropping the sink. - Services MUST enforce a bounded drain timeout on the `TaskTracker::wait()` call in the top-level graceful shutdown future to prevent indefinitely hung connections. ## Bad vs Good Implementations ### Unmonitored Fire-and-Forget WebSocket Task vs. Drain-Fenced Tracker with Cooperative Close #### Bad Antipattern (AI Common Hallucination) > **Why it fails**: Spawning an unmonitored task via `tokio::spawn` without tying it to Axum's graceful shutdown lifecycle leaves the task orphaned. When the process terminates on SIGTERM, running futures are dropped instantly, corrupting in-flight state and dropping TCP sockets without standard WebSocket close frames. ```rust use axum::{extract::ws::{WebSocket, WebSocketUpgrade}, response::IntoResponse}; pub async fn ws_handler(ws: WebSocketUpgrade) -> impl IntoResponse { // ANTI-PATTERN: Fire-and-forget task spawn without tracking or cancellation hook. // During zero-downtime deploy (SIGTERM), Axum terminates, aborting this task mid-flight. // The client encounters a hard TCP RST instead of a clean WS close handshake. ws.on_upgrade(|socket| async move { tokio::spawn(async move { let (mut sender, mut receiver) = socket.split(); while let Some(Ok(msg)) = receiver.next().await { // Processing loop orphaned on rollout let _ = sender.send(msg).await; } }); }) } ``` #### Production Standard (Correct) > **Standard**: The socket task is registered into `TaskTracker` and listens for shutdown via `CancellationToken`. Upon receiving the drain signal, it emits a standard 1001 Close frame, flushes the sink with a deadline, and gracefully unwinds, enabling clean zero-downtime rollouts. ```rust use axum::{extract::{ws::{Message, WebSocket, WebSocketUpgrade}, State}, response::IntoResponse}; use futures_util::{SinkExt, StreamExt}; use std::{sync::Arc, time::Duration}; use tokio_util::{sync::CancellationToken, task::TaskTracker}; #[derive(Clone)] pub struct AppState { pub tracker: TaskTracker, pub token: CancellationToken, } pub async fn ws_handler(ws: WebSocketUpgrade, State(state): State<Arc<AppState>>) -> impl IntoResponse { ws.on_upgrade(move |socket| { let tracker = state.tracker.clone(); let token = state.token.child_token(); async move { tracker.spawn(async move { handle_socket(socket, token).await; }); } }) } async fn handle_socket(socket: WebSocket, token: CancellationToken) { let (mut sender, mut receiver) = socket.split(); loop { tokio::select! { biased; _ = token.cancelled() => { // Cooperative drain: send Close frame (1001 Going Away) and flush let close_frame = Message::Close(Some(axum::extract::ws::CloseFrame { code: 1001, reason: "server_shutdown".into(), })); let _ = tokio::time::timeout(Duration::from_secs(3), async { let _ = sender.send(close_frame).await; let _ = sender.flush().await; }).await; break; } msg = receiver.next() => { match msg { Some(Ok(Message::Close(_))) | None => break, Some(Ok(m)) => { if sender.send(m).await.is_err() { break; } } Some(Err(_)) => break, } } } } } ```

Architectural Specification & Invariant Rationale

Architectural Invariant Design: Axum 0.8 Graceful WebSocket Drain-Fencing Standard

Axum 0.8 Graceful WebSocket Drain-Fencing Standard defines an industrial-grade behavioral standard for backend engineering. Enforces deterministic drain-fencing and cooperative task teardown for bi-directional WebSocket streams in Axum 0.8 during cluster rollouts. Prevents dropped frames, dangling Tokio tasks, and broker desynchronization by decoupling socket upgrades from unmonitored runtime tasks using tokio_util::task::TaskTracker and broadcast cancellation tokens. 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: Spawning Unfenced Tokio Tasks for WebSocket Upgrades Without Cooperative Drain Propagation 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	Spawning Unfenced Tokio Tasks for WebSocket Upgrades Without Cooperative Drain Propagation	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 track every spawned WebSocket connection task using `tokio_util

Engineers MUST ALWAYS track every spawned WebSocket connection task using `tokio_util::task::TaskTracker`.

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.