Axum 0.8 Tower Graceful Drain & CancellationToken Propagation Guardrail

--- description: Axum 0.8 Tower Graceful Drain & CancellationToken Propagation Guardrail - Mandates explicit CancellationToken propagation, TaskTracker integration, and bounded connection draining across all Axum 0.8 Tower services and streaming handlers (SSE, long-polling, WebSockets). Guarantees zero dropped in-flight requests, prevents infinite drain deadlocks, and eliminates memory leaks during zero-downtime rolling deployments under heavy concurrent asynchronous traffic. globs: src/**/*.rs alwaysApply: true verified_by: "agentrules.org" license: "CC-BY-4.0" source: "https://agentrules.org/rules/axum-08-tower-graceful-drain-cancellation-propagation" --- # Axum 0.8 Tower Graceful Drain & CancellationToken Propagation Guardrail > **Severity**: CRITICAL | **Category**: backend > **Anti-Pattern Warning**: Unbounded Streaming Handlers and Detached Tasks Blocking Axum Graceful Server Drain ## System Constraints You are a Principal Rust Architect enforcing deterministic graceful draining and async cancellation propagation in Axum 0.8 and Tower ecosystems. Axum 0.8's `axum::serve(...).with_graceful_shutdown(...)` halts acceptance of new TCP connections and waits for all active in-flight requests to complete before exiting. Long-polling loops, Server-Sent Events (SSE), WebSocket streams, and background worker loops will hang shutdown indefinitely unless every streaming handler and asynchronous worker explicitly hooks into a unified `tokio_util::sync::CancellationToken` and completes its stream lifecycle cleanly. You must strictly mandate that all persistent HTTP streams select over cancellation tokens, spawn with `TaskTracker` bounds, and emit terminal frames before terminating. ## Strict Guidelines - Engineers MUST ALWAYS pass a cloned `tokio_util::sync::CancellationToken` into Axum application state and select against `token.cancelled()` in all long-polling, SSE, and WebSocket handler loops. - Engineers MUST NOT use unbounded `loop { ... }` streams or unconstrained `ReceiverStream` instances inside Axum response handlers without cancellation branches. - Engineers MUST trigger `CancellationToken::cancel()` inside the `with_graceful_shutdown` closure to notify in-flight streams before the server begins TCP listener drainage. - Services MUST bound background tasks using `tokio_util::task::TaskTracker` instead of raw `tokio::spawn`, ensuring background pipelines complete or abort within orchestrated shutdown deadlines. - Components NEVER rely on OS-level `SIGKILL` or Kubernetes `terminationGracePeriodSeconds` timeouts to force-abort hanging connections, as this drops client frames mid-stream and triggers uncollected resource leaks. - SSE handlers SHOULD yield a terminal event (e.g., `event: shutdown`) and exit cleanly when cancellation is triggered, allowing downstream clients (browsers, proxies) to reconnect to alternate pods. ## Bad vs Good Implementations ### Unbounded SSE Stream Deadlocking Axum 0.8 Graceful Drain vs. Cancellation-Aware Draining #### Bad Antipattern (AI Common Hallucination) > **Why it fails**: The SSE handler generates an infinite stream without observing shutdown signals. Axum 0.8 server drain tracks in-flight request body completions; because the stream never yields None/terminates, the server blocks on shutdown forever until forcefully terminated, severing client connections and corrupting telemetry. ```rust use axum::{routing::get, response::sse::{Event, Sse}, Router}; use futures::stream::Stream; use std::{convert::Infallible, time::Duration}; // ANTI-PATTERN: Stream runs forever. In Axum 0.8, with_graceful_shutdown // waits for all in-flight request futures to complete. This handler never completes, // deadlocking rolling deployments until Kubernetes issues SIGKILL, dropping active clients. async fn sse_handler() -> Sse<impl Stream<Item = Result<Event, Infallible>>> { let stream = async_stream::stream! { loop { tokio::time::sleep(Duration::from_secs(1)).await; yield Ok(Event::default().data("heartbeat")); } }; Sse::new(stream) } pub async fn run_server() { let app = Router::new().route("/events", get(sse_handler)); let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap(); // Hangs indefinitely when SIGTERM is received! axum::serve(listener, app) .with_graceful_shutdown(async { tokio::signal::ctrl_c().await.ok(); }) .await .unwrap(); } ``` #### Production Standard (Correct) > **Standard**: The CancellationToken is held in AppState and cloned into the streaming handler. Upon receiving SIGINT/SIGTERM, the shutdown signal cancels the token. The handler breaks out of its infinite loop, emits a final 'server_shutdown' frame, and closes the stream, allowing Axum 0.8's server drain to finish immediately with zero dropped connections. ```rust use axum::{extract::State, routing::get, response::sse::{Event, Sse, KeepAlive}, Router}; use futures::stream::Stream; use std::{convert::Infallible, time::Duration}; use tokio_util::sync::CancellationToken; #[derive(Clone)] pub struct AppState { pub shutdown_token: CancellationToken, } async fn sse_handler( State(state): State<AppState>, ) -> Sse<impl Stream<Item = Result<Event, Infallible>>> { let token = state.shutdown_token.clone(); let stream = async_stream::stream! { loop { tokio::select! { biased; _ = token.cancelled() => { // Gracefully notify client of planned pod retirement and terminate yield Ok(Event::default().event("server_shutdown").data("draining")); break; } _ = tokio::time::sleep(Duration::from_secs(1)) => { yield Ok(Event::default().data("heartbeat")); } } } }; Sse::new(stream).keep_alive(KeepAlive::default()) } pub async fn run_server() { let token = CancellationToken::new(); let state = AppState { shutdown_token: token.clone() }; let app = Router::new().route("/events", get(sse_handler)).with_state(state); let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap(); let shutdown_signal = { let token = token.clone(); async move { tokio::signal::ctrl_c().await.ok(); // Propagate cancellation token to all active stream handlers token.cancel(); } }; axum::serve(listener, app) .with_graceful_shutdown(shutdown_signal) .await .unwrap(); } ```

Architectural Specification & Invariant Rationale

Architectural Invariant Design: Axum 0.8 Tower Graceful Drain & CancellationToken Propagation Guardrail

Axum 0.8 Tower Graceful Drain & CancellationToken Propagation Guardrail defines an industrial-grade behavioral standard for backend engineering. Mandates explicit CancellationToken propagation, TaskTracker integration, and bounded connection draining across all Axum 0.8 Tower services and streaming handlers (SSE, long-polling, WebSockets). Guarantees zero dropped in-flight requests, prevents infinite drain deadlocks, and eliminates memory leaks during zero-downtime rolling deployments under heavy concurrent asynchronous traffic. 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 Streaming Handlers and Detached Tasks Blocking Axum Graceful Server Drain 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 Streaming Handlers and Detached Tasks Blocking Axum Graceful Server Drain	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 cloned `tokio_util

Engineers MUST ALWAYS pass a cloned `tokio_util::sync::CancellationToken` into Axum application state and select against `token.cancelled()` in all long-polling, SSE, and WebSocket handler 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.