Axum 0.8 & Tower In-Flight Bidirectional Stream Drain Boundary

--- description: Axum 0.8 & Tower In-Flight Bidirectional Stream Drain Boundary - Enforces deterministic graceful shutdown and connection drain fences for upgraded bidirectional connections (WebSockets, SSE, and gRPC) in Axum 0.8 and Tower services. Prevents abrupt Kubernetes SIGKILL resets by mandating protocol-level drain signaling, TaskTracker tracking, and bounded drain timeouts prior to pod eviction. globs: src/**/*.rs, crates/**/*.rs alwaysApply: true verified_by: "agentrules.org" license: "CC-BY-4.0" source: "https://agentrules.org/rules/axum-tower-bidirectional-stream-drain-boundary" --- # Axum 0.8 & Tower In-Flight Bidirectional Stream Drain Boundary > **Severity**: CRITICAL | **Category**: backend > **Anti-Pattern Warning**: Leaking In-Flight Upgraded Connections Past SIGTERM Leading to Abrupt SIGKILL Termination ## System Constraints You are a Principal Systems Architect enforcing strict stream lifecycle and shutdown boundaries for Axum 0.8 and Tower services running on Kubernetes. You must prevent catastrophic socket resets by enforcing TaskTracker task boundaries, protocol-level termination frames (e.g., WebSocket 1001 Going Away), and readiness probe flipping on SIGTERM. ## Strict Guidelines - Engineers MUST NOT rely solely on `axum::serve(listener, app).with_graceful_shutdown(...)` to terminate upgraded connections, because Hyper's default graceful shutdown does not close upgraded full-duplex sockets (such as WebSockets). - Services MUST instantiate a `tokio_util::sync::CancellationToken` and a `tokio_util::task::TaskTracker` at the process root to manage long-lived bidirectional streams. - WebSocket and streaming handlers MUST track their spawned execution futures within the global `TaskTracker` and select across the client socket stream and `CancellationToken::cancelled()`. - Upon receiving `SIGTERM`, services MUST immediately fail readiness checks (return HTTP 503 at `/readyz`) while preserving liveness checks, to ensure ingress controllers stop routing new traffic before draining commences. - When a shutdown signal is tripped, streaming workers MUST emit a protocol-compliant termination frame (e.g., `Message::Close` with `close_code::AWAY` for WebSockets) before closing the underlying transport. - The orchestrator MUST enforce a hard drain fence via `tokio::time::timeout` on `tracker.wait()` strictly shorter than the Kubernetes `terminationGracePeriodSeconds` (e.g., a 25-second drain fence for a 30-second grace period) to guarantee clean process exit before SIGKILL. ## Bad vs Good Implementations ### Graceful Upgraded Connection Drain in Axum 0.8 #### Bad Antipattern (AI Common Hallucination) > **Why it fails**: Hyper and Axum graceful shutdown logic only halts acceptance of new TCP connections and drains active HTTP request-response cycles. Upgraded protocols (like WebSockets) are severed from hyper's internal bookkeeping. The connection remains open indefinitely until Kubernetes hits the pod termination timeout and sends SIGKILL, resetting client sockets with unhandled TCP RST packets. ```rust use axum::{extract::ws::{WebSocket, WebSocketUpgrade}, response::IntoResponse, routing::get, Router}; use tokio::signal; async fn ws_handler(ws: WebSocketUpgrade) -> impl IntoResponse { ws.on_upgrade(|mut socket: WebSocket| async move { // BUG: Infinite read loop with no cancellation boundary. // axum::serve with_graceful_shutdown will wait forever on this connection, // causing Kubernetes to send SIGKILL after terminationGracePeriodSeconds. while let Some(Ok(msg)) = socket.recv().await { if socket.send(msg).await.is_err() { break; } } }) } #[tokio::main] async fn main() { let app = Router::new().route("/ws", get(ws_handler)); let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap(); axum::serve(listener, app) .with_graceful_shutdown(async { signal::ctrl_c().await.ok(); }) .await .unwrap(); } ``` #### Production Standard (Correct) > **Standard**: The production-grade pattern registers a TaskTracker and CancellationToken. When SIGTERM is intercepted, cancellation is broadcast to every active WebSocket stream. Each stream transmits a WebSocket 1001 Going Away frame and closes cleanly. The application then enforces a strict 25-second bounded fence via `tokio::time::timeout` on `tracker.wait()`, safely under the 30-second Kubernetes SIGKILL window. ```rust use axum::{extract::{ws::{close_code, CloseFrame, Message, WebSocket}, State, WebSocketUpgrade}, response::IntoResponse, routing::get, Router}; use std::{borrow::Cow, sync::Arc, time::Duration}; use tokio::{net::TcpListener, signal::unix::{signal, SignalKind}, time::timeout}; use tokio_util::{sync::CancellationToken, task::TaskTracker}; #[derive(Clone)] struct AppState { tracker: TaskTracker, token: CancellationToken, } async fn ws_handler(State(state): State<AppState>, ws: WebSocketUpgrade) -> impl IntoResponse { ws.on_upgrade(move |socket| { let token = state.token.clone(); // Track connection inside the bounded task tracker state.tracker.track_future(async move { handle_socket(socket, token).await; }) }) } async fn handle_socket(mut socket: WebSocket, token: CancellationToken) { loop { tokio::select! { _ = token.cancelled() => { // Protocol-level drain: Emit 1001 Going Away close frame let frame = CloseFrame { code: close_code::AWAY, reason: Cow::Borrowed("Server draining") }; let _ = socket.send(Message::Close(Some(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, } } } } } #[tokio::main] async fn main() { let tracker = TaskTracker::new(); let token = CancellationToken::new(); let state = AppState { tracker: tracker.clone(), token: token.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 mut sigterm = signal(SignalKind::terminate()).unwrap(); tokio::select! { _ = sigterm.recv() => println!("SIGTERM received: initiating drain"), _ = tokio::signal::ctrl_c() => println!("Ctrl+C received: initiating drain"), } // 1. Broadcast cancellation to active streaming workers token.cancel(); tracker.close(); // 2. Bound connection drain to 25s (Kubernetes terminationGracePeriodSeconds is 30s) const DRAIN_TIMEOUT: Duration = Duration::from_secs(25); let drain_result = timeout(DRAIN_TIMEOUT, async { tokio::join!(server.with_graceful_shutdown(async {}), tracker.wait()); }).await; if drain_result.is_err() { eprintln!("Drain fence timeout exceeded: forced connection shutdown"); } } ```

Architectural Specification & Invariant Rationale

Architectural Invariant Design: Axum 0.8 & Tower In-Flight Bidirectional Stream Drain Boundary

Axum 0.8 & Tower In-Flight Bidirectional Stream Drain Boundary defines an industrial-grade behavioral standard for backend engineering. Enforces deterministic graceful shutdown and connection drain fences for upgraded bidirectional connections (WebSockets, SSE, and gRPC) in Axum 0.8 and Tower services. Prevents abrupt Kubernetes SIGKILL resets by mandating protocol-level drain signaling, TaskTracker tracking, and bounded drain timeouts prior to 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 In-Flight Upgraded Connections Past SIGTERM Leading to Abrupt SIGKILL Termination 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 In-Flight Upgraded Connections Past SIGTERM Leading to Abrupt SIGKILL Termination	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 NOT rely solely on `axum

Engineers MUST NOT rely solely on `axum::serve(listener, app).with_graceful_shutdown(...)` to terminate upgraded connections, because Hyper's default graceful shutdown does not close upgraded full-duplex sockets (such as WebSockets).

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.