Axum Tokio In-Flight SSE & WebSocket Deterministic Channel Draining on SIGTERM

--- description: Axum Tokio In-Flight SSE & WebSocket Deterministic Channel Draining on SIGTERM - Long-lived HTTP streaming (SSE) and WebSocket connections in Axum must not rely solely on default hyper/axum graceful shutdown hooks, which stall indefinitely on open streams. Production services must coordinate Kubernetes SIGTERM teardown using CancellationToken hierarchies and TaskTracker primitives to stop ingress, flush buffered channel frames, dispatch WS 1001 (Going Away) close handshakes or SSE terminal events, and deterministically bound task draining before pod eviction. globs: src/**/*.rs alwaysApply: true verified_by: "agentrules.org" license: "CC-BY-4.0" source: "https://agentrules.org/rules/axum-tokio-sse-ws-graceful-draining-sigterm" --- # Axum Tokio In-Flight SSE & WebSocket Deterministic Channel Draining on SIGTERM > **Severity**: CRITICAL | **Category**: backend > **Anti-Pattern Warning**: Leaking Open Streams Across Container Teardown Causing Hard SIGKILLs and Truncated In-Flight Payloads ## System Constraints You are a Principal Rust Infrastructure Engineer enforcing bulletproof connection draining protocols for Axum and Tokio services. Under Kubernetes rolling updates, unmanaged SSE and WebSocket sessions block rolling deployments or get abruptly terminated by SIGKILL after terminationGracePeriodSeconds. Mandate the use of tokio_util::sync::CancellationToken, tokio_util::task::TaskTracker, explicit backpressure flushes, and protocol-compliant close frames to guarantee zero message drop. ## Strict Guidelines - Engineers MUST ALWAYS pair axum::serve::with_graceful_shutdown with an explicit tokio_util::sync::CancellationToken propagated to all active SSE and WebSocket stream loops. - Engineers MUST ALWAYS track all streaming tasks using tokio_util::task::TaskTracker or explicit join handles, ensuring the main process waits for task drainage before exiting. - Engineers MUST NOT spawn detached Tokio tasks (`tokio::spawn`) for WebSocket reader/writer loops without registering them to a TaskTracker or structured lifecycle supervisor. - WebSocket handlers MUST send a formal CloseFrame with code 1001 (Going Away) and allow a bounded window for the client TCP ACK before terminating upon cancellation. - SSE streams MUST evaluate cancellation within `tokio_stream` or `futures_util::stream::Stream` adapters, flushing pending channel items before emitting a terminal SSE comment or closing the underlying body chunk stream. - Services MUST implement a Kubernetes pre-drain grace period (readiness probe unregistration) prior to triggering channel cancellation to prevent race conditions with kube-proxy routing tables. - Channel pipelines MUST drain all currently buffered messages in bounded mpsc queues using non-blocking `try_recv()` loops before emitting close protocol payloads. ## Bad vs Good Implementations ### Detached WS/SSE Loops Blocking or Truncating on Pod SIGTERM #### Bad Antipattern (AI Common Hallucination) > **Why it fails**: Spawning an untracked, detached task inside `on_upgrade` that loops indefinitely on a channel prevents Axum's server from completing graceful shutdown cleanly. The container is forcibly murdered by Kubernetes via SIGKILL after `terminationGracePeriodSeconds`, dropping any in-flight frames and failing to send WebSocket close frames (RFC 6455 1001). ```rust use axum::{extract::ws::{WebSocket, WebSocketUpgrade}, response::Response, routing::get, Router}; use tokio::sync::broadcast; pub fn router(tx: broadcast::Sender<String>) -> Router { Router::new().route("/ws", get(move |ws: WebSocketUpgrade| { let tx = tx.clone(); async move { ws.on_upgrade(|socket| handle_socket(socket, tx)) } })) } async fn handle_socket(mut socket: WebSocket, tx: broadcast::Sender<String>) { let mut rx = tx.subscribe(); // ANTI-PATTERN: Detached task ignores SIGTERM. Axum's graceful_shutdown // hangs forever waiting for this connection to finish, ultimately // causing Kubernetes to send SIGKILL and drop unread buffer frames. tokio::spawn(async move { while let Ok(msg) = rx.recv().await { if socket.send(axum::extract::ws::Message::Text(msg)).await.is_err() { break; } } }); } ``` #### Production Standard (Correct) > **Standard**: The session is managed via `TaskTracker` and listens for `CancellationToken` cancellation. Upon SIGTERM, the loop exits `biased` selection, exhausts remaining queued mpsc messages via `try_recv()`, transmits a WebSocket 1001 'Away' close frame to inform the client to reconnect elsewhere, and cleanly closes the socket within a bounded timeout. ```rust use axum::{extract::ws::{close_code, CloseFrame, Message, WebSocket, WebSocketUpgrade}, response::Response, routing::get, Router}; use std::{borrow::Cow, sync::Arc, time::Duration}; use tokio::sync::mpsc; use tokio_util::{sync::CancellationToken, task::TaskTracker}; #[derive(Clone)] pub struct AppState { pub tracker: TaskTracker, pub token: CancellationToken, } pub fn router(state: AppState) -> Router { Router::new().route("/ws", get(move |ws: WebSocketUpgrade| { let s = state.clone(); async move { ws.on_upgrade(move |socket| { let tracker = s.tracker.clone(); let token = s.token.clone(); async move { tracker.spawn(async move { handle_draining_ws(socket, token).await; }); } }) } })) } async fn handle_draining_ws(mut socket: WebSocket, token: CancellationToken) { let (tx, mut rx) = mpsc::channel::<String>(128); // Protocol pump with deterministic SIGTERM cancellation interception loop { tokio::select! { biased; _ = token.cancelled() => { // 1. Drain residual queued messages with zero-drop guarantee while let Ok(msg) = rx.try_recv() { let _ = socket.send(Message::Text(msg)).await; } // 2. Dispatch protocol-compliant CloseFrame let close_frame = CloseFrame { code: close_code::AWAY, reason: Cow::Borrowed("Server pod initiating graceful rolling update"), }; let _ = socket.send(Message::Close(Some(close_frame))).await; // 3. Allow bounded flush window for TCP write buffer let _ = tokio::time::timeout(Duration::from_millis(500), socket.close()).await; break; } Some(msg) = rx.recv() => { if socket.send(Message::Text(msg)).await.is_err() { break; } } else => break, } } } ```

Architectural Specification & Invariant Rationale

Architectural Invariant Design: Axum Tokio In-Flight SSE & WebSocket Deterministic Channel Draining on SIGTERM

Axum Tokio In-Flight SSE & WebSocket Deterministic Channel Draining on SIGTERM defines an industrial-grade behavioral standard for backend engineering. Long-lived HTTP streaming (SSE) and WebSocket connections in Axum must not rely solely on default hyper/axum graceful shutdown hooks, which stall indefinitely on open streams. Production services must coordinate Kubernetes SIGTERM teardown using CancellationToken hierarchies and TaskTracker primitives to stop ingress, flush buffered channel frames, dispatch WS 1001 (Going Away) close handshakes or SSE terminal events, and deterministically bound task draining before pod eviction. 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: Leaking Open Streams Across Container Teardown Causing Hard SIGKILLs and Truncated In-Flight Payloads 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	Leaking Open Streams Across Container Teardown Causing Hard SIGKILLs and Truncated In-Flight Payloads	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 pair axum

Engineers MUST ALWAYS pair axum::serve::with_graceful_shutdown with an explicit tokio_util::sync::CancellationToken propagated to all active SSE and WebSocket stream loops.

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.