Axum 0.8 Cooperative SSE and WebSocket Stream Draining via Tokio CancellationToken

--- description: Axum 0.8 Cooperative SSE and WebSocket Stream Draining via Tokio CancellationToken - Mandates cooperative stream draining using tokio_util::sync::CancellationToken and bounded timeouts for Axum 0.8 long-lived SSE and WebSocket connections. Ensures graceful close frame dispatching, buffer flushes, and state finalization without abrupt socket severance or infinite shutdown hangs. globs: src/**/*.rs alwaysApply: true verified_by: "agentrules.org" license: "CC-BY-4.0" source: "https://agentrules.org/rules/axum-tokio-cancellation-token-stream-draining" --- # Axum 0.8 Cooperative SSE and WebSocket Stream Draining via Tokio CancellationToken > **Severity**: CRITICAL | **Category**: backend > **Anti-Pattern Warning**: Abrupt Socket Severance and Unbounded Process Leaks via Uncoordinated Stream Drops During Shutdown ## System Constraints You are a Principal Rust Concurrency Architect enforcing Axum 0.8 and Tokio cooperative shutdown standards. You mandate that all persistent streaming protocols (Server-Sent Events, WebSockets) track lifecycle state via child `CancellationToken`s and bounded draining procedures. Never abruptly drop TCP connections or leak uncoordinated background tasks. ## Strict Guidelines - Engineers MUST ALWAYS propagate a child `tokio_util::sync::CancellationToken` from application state to all active WebSocket handlers and SSE stream generators. - Engineers MUST NOT abruptly terminate active connection tasks using `tokio::select!` branches that drop the stream without emitting clean protocol termination frames (e.g., WS Close code 1001 'Going Away' or custom SSE shutdown events). - WebSocket connection loops MUST implement a dual-phase shutdown: phase one ceases accepting new ingress messages and sends a close frame; phase two drains in-flight egress buffers bounded by a strict `tokio::time::timeout`. - Server-Sent Events (`axum::response::sse::Sse`) MUST construct streams wrapped with `tokio_stream::StreamExt::take_until` or custom polling logic that emits a terminal sentinel event before completing. - Components NEVER await raw client responses indefinitely during shutdown; all draining logic MUST enforce non-configurable or bounded grace periods (recommended: 5–15 seconds) before forced socket abortion. ## Bad vs Good Implementations ### Cooperative WebSocket Stream Draining with Close Frame Protocol #### Bad Antipattern (AI Common Hallucination) > **Why it fails**: Dropping the socket directly inside the select branch instantly severs the TCP connection. The client receives an unhandled socket drop or connection reset instead of a clean WebSocket closure code, causing client-side reconnect storms and lost buffered outbound messages. ```rust use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; use axum::response::IntoResponse; use tokio_util::sync::CancellationToken; pub async fn ws_handler( ws: WebSocketUpgrade, token: CancellationToken, ) -> impl IntoResponse { ws.on_upgrade(move |mut socket| async move { loop { tokio::select! { // ANTI-PATTERN: Hard drop on cancellation causes TCP RST / broken pipe on client, // leaving client without clean status code 1001 (Going Away) and discarding buffered writes. _ = token.cancelled() => { tracing::info!("Server shutting down, dropping socket abruptly"); break; } Some(msg) = socket.recv() => { if let Ok(Message::Text(txt)) = msg { let _ = socket.send(Message::Text(format!("echo: {txt}"))).await; } } } } }) } ``` #### Production Standard (Correct) > **Standard**: When cancellation triggers, the handler enters a cooperative draining phase: it transmits a standard 1001 'Going Away' Close frame and executes a bounded drain using `tokio::time::timeout`, preventing socket starvation while guaranteeing the server can exit cleanly without hanging. ```rust use axum::extract::ws::{close_code, CloseFrame, Message, WebSocket, WebSocketUpgrade}; use axum::response::IntoResponse; use std::time::Duration; use tokio_util::sync::CancellationToken; pub async fn ws_handler( ws: WebSocketUpgrade, token: CancellationToken, ) -> impl IntoResponse { ws.on_upgrade(move |socket| handle_socket_cooperative(socket, token)) } async fn handle_socket_cooperative(mut socket: WebSocket, token: CancellationToken) { const DRAIN_TIMEOUT: Duration = Duration::from_secs(5); loop { tokio::select! { biased; _ = token.cancelled() => { tracing::info!("Shutdown signal received: initiating cooperative stream draining"); // Phase 1: Notify client with clean 1001 Going Away frame let close_frame = CloseFrame { code: close_code::AWAY, reason: "Server shutting down".into(), }; let _ = socket.send(Message::Close(Some(close_frame))).await; // Phase 2: Bounded drain of pending frames until client ACK or timeout let _ = tokio::time::timeout(DRAIN_TIMEOUT, async { while let Some(Ok(msg)) = socket.recv().await { if let Message::Close(_) = msg { break; } } }).await; break; } msg = socket.recv() => { match msg { Some(Ok(Message::Text(txt))) => { if socket.send(Message::Text(format!("echo: {txt}"))).await.is_err() { break; } } Some(Ok(Message::Close(_))) | None => break, _ => {} } } } } } ```

Architectural Specification & Invariant Rationale

Architectural Invariant Design: Axum 0.8 Cooperative SSE and WebSocket Stream Draining via Tokio CancellationToken

Axum 0.8 Cooperative SSE and WebSocket Stream Draining via Tokio CancellationToken defines an industrial-grade behavioral standard for backend engineering. Mandates cooperative stream draining using tokio_util::sync::CancellationToken and bounded timeouts for Axum 0.8 long-lived SSE and WebSocket connections. Ensures graceful close frame dispatching, buffer flushes, and state finalization without abrupt socket severance or infinite shutdown hangs. 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: Abrupt Socket Severance and Unbounded Process Leaks via Uncoordinated Stream Drops During Shutdown 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	Abrupt Socket Severance and Unbounded Process Leaks via Uncoordinated Stream Drops During Shutdown	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 propagate a child `tokio_util

Engineers MUST ALWAYS propagate a child `tokio_util::sync::CancellationToken` from application state to all active WebSocket handlers and SSE stream generators.

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.