Axum 0.8 Tokio Cooperative Cancellation Token Propagation and Draining

--- description: Axum 0.8 Tokio Cooperative Cancellation Token Propagation and Draining - Enforces cooperative cancellation token propagation and connection tracking across long-lived Axum 0.8 streaming connections (Server-Sent Events and WebSockets). Prevents uncoordinated termination and leaked Tokio tasks during Kubernetes SIGTERM rolling updates by tracking in-flight streams via CancellationToken trees and graceful shutdown drain sequences. globs: src/**/*.rs alwaysApply: true verified_by: "agentrules.org" license: "CC-BY-4.0" source: "https://agentrules.org/rules/axum-tokio-graceful-cancellation-draining" --- # Axum 0.8 Tokio Cooperative Cancellation Token Propagation and Draining > **Severity**: CRITICAL | **Category**: backend > **Anti-Pattern Warning**: Unbounded Tokio Task Spawning and Detached Loops in Streaming Handlers Causing Hung Deployments and Leaked Connections ## System Constraints You are a Principal Rust Distributed Systems Engineer enforcing Axum 0.8 and Tokio graceful shutdown standards. Mandate cooperative cancellation across all persistent streams and WebSocket sessions using tokio_util::sync::CancellationToken and explicit connection draining guards. ## Strict Guidelines - Engineers MUST ALWAYS pass a child `tokio_util::sync::CancellationToken` or connection guard to every persistent connection handler (SSE, WebSocket, long-polling). - Engineers MUST NOT spawn detached `tokio::spawn` loops inside route handlers without linking them to a `CancellationToken` or holding an explicit tracker (e.g., `tokio_util::task::TaskTracker`). - Streaming endpoints MUST select over the `CancellationToken::cancelled()` future alongside incoming data streams to terminate loops deterministically during zero-downtime rolling deploys. - WebSocket handlers MUST send a formal WebSocket Close frame (`ws::Message::Close`) with a draining status code before unmounting during server shutdown. - The server root MUST configure `axum::serve(...).with_graceful_shutdown(drain_signal)` and ensure the graceful signal initiates tree cancellation on the root `CancellationToken` followed by a drain timeout. ## Bad vs Good Implementations ### WebSocket Session Handling with Cooperative Draining Guard #### Bad Antipattern (AI Common Hallucination) > **Why it fails**: Spawns a detached Tokio task without a reference to a cancellation token or task tracker. When Axum receives a shutdown signal, the server cannot coordinate draining with this connection, causing either stalled rolling updates or sudden TCP resets via SIGKILL. ```rust use axum::{extract::ws::{WebSocket, WebSocketUpgrade}, response::IntoResponse}; use tokio::time::{sleep, Duration}; pub async fn ws_handler(ws: WebSocketUpgrade) -> impl IntoResponse { ws.on_upgrade(|mut socket| async move { // BAD: Detached unbounded loop without cancellation awareness. // During SIGTERM / pod eviction, this task keeps running indefinitely // or gets abruptly dropped when SIGKILL occurs after terminationGracePeriodSeconds. tokio::spawn(async move { loop { sleep(Duration::from_secs(1)).await; if socket.send(axum::extract::ws::Message::Text("ping".into())).await.is_err() { break; } } }); }) } ``` #### Production Standard (Correct) > **Standard**: Constructs a child `CancellationToken` derived from the application-level shutdown tree. Uses `tokio::select!` to cooperatively react to graceful shutdown signals, sends an RFC 6455 1001 'Going Away' close frame, and closes the socket cleanly, allowing Kubernetes pods to terminate gracefully within the termination window. ```rust use axum::{extract::{ws::{CloseFrame, Message, WebSocket}, State, WebSocketUpgrade}, response::IntoResponse}; use std::sync::Arc; use tokio_util::sync::CancellationToken; #[derive(Clone)] pub struct AppState { pub cancel_token: CancellationToken, } pub async fn ws_handler(ws: WebSocketUpgrade, State(state): State<Arc<AppState>>) -> impl IntoResponse { let child_token = state.cancel_token.child_token(); ws.on_upgrade(move |socket| handle_socket(socket, child_token)) } async fn handle_socket(mut socket: WebSocket, cancel_token: CancellationToken) { let mut ticker = tokio::time::interval(std::time::Duration::from_secs(1)); loop { tokio::select! { _ = cancel_token.cancelled() => { let close_frame = CloseFrame { code: 1001, // Going Away reason: "Server shutting down".into(), }; let _ = socket.send(Message::Close(Some(close_frame))).await; let _ = socket.close().await; break; } _ = ticker.tick() => { if socket.send(Message::Text("ping".into())).await.is_err() { break; } } msg = socket.recv() => { match msg { Some(Ok(Message::Close(_))) | None => break, Some(Ok(_)) => {}, Some(Err(_)) => break, } } } } } ```

Architectural Specification & Invariant Rationale

Architectural Invariant Design: Axum 0.8 Tokio Cooperative Cancellation Token Propagation and Draining

Axum 0.8 Tokio Cooperative Cancellation Token Propagation and Draining defines an industrial-grade behavioral standard for backend engineering. Enforces cooperative cancellation token propagation and connection tracking across long-lived Axum 0.8 streaming connections (Server-Sent Events and WebSockets). Prevents uncoordinated termination and leaked Tokio tasks during Kubernetes SIGTERM rolling updates by tracking in-flight streams via CancellationToken trees and graceful shutdown drain sequences. 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 Tokio Task Spawning and Detached Loops in Streaming Handlers Causing Hung Deployments and Leaked Connections 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 Tokio Task Spawning and Detached Loops in Streaming Handlers Causing Hung Deployments and Leaked Connections	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 pass a child `tokio_util

Engineers MUST ALWAYS pass a child `tokio_util::sync::CancellationToken` or connection guard to every persistent connection handler (SSE, WebSocket, long-polling).

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.