Axum Tower Phase-Aligned Stream Draining & Graceful Backpressure Invariant

--- description: Axum Tower Phase-Aligned Stream Draining & Graceful Backpressure Invariant - Guarantees that Axum and Tower streaming endpoints (SSE, WebSockets, WebTransport) enforce phased, cooperative connection draining during pod lifecycle termination. Prevents connection resets, unbuffered drops, and indefinite rollout stalls by bounding drain phases with cancellation tokens and task tracking. globs: src/**/*.rs, crates/**/*.rs alwaysApply: true verified_by: "agentrules.org" license: "CC-BY-4.0" source: "https://agentrules.org/rules/axum-tower-graceful-drain-stream-invariants" --- # Axum Tower Phase-Aligned Stream Draining & Graceful Backpressure Invariant > **Severity**: CRITICAL | **Category**: backend > **Anti-Pattern Warning**: Uncoordinated Axum Graceful Shutdown Causing Abrupt SSE/WebSocket Dropping or Indefinite Pod Termination Deadlocks ## System Constraints You are a Principal Rust Architect enforcing Axum and Tower graceful drain and backpressure invariants. Mandate cooperative stream termination, task tracking via tokio_util::task::TaskTracker, dual-phased shutdown timeouts, and backpressure propagation during Kubernetes pod rollouts. ## Strict Guidelines - Engineers MUST ALWAYS couple Axum long-lived streaming endpoints (SSE, WebSocket, WebTransport) with a tokio_util::sync::CancellationToken to signal connection draining. - Engineers MUST NOT rely solely on axum::serve::with_graceful_shutdown without cooperative stream termination, as Hyper waits indefinitely for response streams to yield None. - Services MUST track background streaming tasks using tokio_util::task::TaskTracker rather than detached tokio::spawn calls. - Code MUST enforce a phased shutdown timeout: (Phase 1) signal cancellation and stop ingress, (Phase 2) bounded wait for in-flight streams, strictly less than kubernetes terminationGracePeriodSeconds. - Streams MUST emit protocol-compliant termination signals (e.g., SSE 'event: close', WebSocket 1001 Going Away) and flush pending frames before terminating. - Components NEVER drop channel receivers without draining buffer capacity, ensuring upstream publishers observe backpressure and avoid broken-pipe cascading errors. ## Bad vs Good Implementations ### Graceful SSE Stream Draining via CancellationToken and TaskTracker #### Bad Antipattern (AI Common Hallucination) > **Why it fails**: Hyper's HTTP/1.1 and HTTP/2 connection drain waits for active request/response bodies to terminate. Because this SSE stream loops indefinitely, the server never finishes shutting down naturally. Kubernetes sends SIGKILL after terminationGracePeriodSeconds, causing broken pipes and abrupt client-side reconnect storms. ```rust use axum::{response::sse::{Event, Sse}, routing::get, Router}; use futures::stream::Stream; use tokio::signal; async fn event_stream() -> Sse<impl Stream<Item = Result<Event, std::convert::Infallible>>> { // ANTI-PATTERN: Infinite stream has no awareness of shutdown signals. // Axum's with_graceful_shutdown will wait forever for this stream to finish, // blocking pod teardown until Kubernetes SIGKILL abruptly severs the TCP connection. let stream = async_stream::stream! { let mut ticker = tokio::time::interval(std::time::Duration::from_secs(1)); loop { ticker.tick().await; yield Ok(Event::default().data("tick")); } }; Sse::new(stream) } #[tokio::main] async fn main() { let app = Router::new().route("/events", get(event_stream)); 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.unwrap(); }) .await .unwrap(); } ``` #### Production Standard (Correct) > **Standard**: Uses a CancellationToken combined with tokio::select! to cooperatively terminate active SSE streams upon SIGTERM. The stream emits a final protocol-compliant event before closing, allowing clients to failover gracefully without connection resets. Bounded wait prevents deadlocks prior to pod eviction. ```rust use axum::{ extract::State, response::sse::{Event, KeepAlive, Sse}, routing::get, Router, }; use futures::stream::Stream; use std::time::Duration; use tokio::signal; use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; #[derive(Clone)] struct AppState { drain_token: CancellationToken, tracker: TaskTracker, } async fn sse_handler( State(state): State<AppState>, ) -> Sse<impl Stream<Item = Result<Event, std::convert::Infallible>>> { let drain_token = state.drain_token.clone(); let stream = async_stream::stream! { let mut ticker = tokio::time::interval(Duration::from_secs(1)); loop { tokio::select! { biased; _ = drain_token.cancelled() => { // Phase-aligned closure: inform client and terminate stream cleanly yield Ok(Event::default().event("shutdown").data("going_away")); break; } _ = ticker.tick() => { yield Ok(Event::default().data("heartbeat")); } } } }; Sse::new(stream).keep_alive(KeepAlive::default()) } #[tokio::main] async fn main() { let drain_token = CancellationToken::new(); let tracker = TaskTracker::new(); let state = AppState { drain_token: drain_token.clone(), tracker: tracker.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 server = axum::serve(listener, app); tokio::select! { res = server.with_graceful_shutdown(drain_token.clone().cancelled_owned()) => { res.unwrap(); } _ = signal::ctrl_c() => { // Phase 1: Signal all streams to cleanly exit drain_token.cancel(); tracker.close(); // Phase 2: Enforce bounded drain before Kubernetes SIGKILL cutoff (e.g., 20s of 30s period) if tokio::time::timeout(Duration::from_secs(20), tracker.wait()).await.is_err() { eprintln!("Drain timeout exceeded: forcing teardown"); } } } } ```

Architectural Specification & Invariant Rationale

Architectural Invariant Design: Axum Tower Phase-Aligned Stream Draining & Graceful Backpressure Invariant

Axum Tower Phase-Aligned Stream Draining & Graceful Backpressure Invariant defines an industrial-grade behavioral standard for backend engineering. Guarantees that Axum and Tower streaming endpoints (SSE, WebSockets, WebTransport) enforce phased, cooperative connection draining during pod lifecycle termination. Prevents connection resets, unbuffered drops, and indefinite rollout stalls by bounding drain phases with cancellation tokens and task tracking. 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: Uncoordinated Axum Graceful Shutdown Causing Abrupt SSE/WebSocket Dropping or Indefinite Pod Termination Deadlocks 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	Uncoordinated Axum Graceful Shutdown Causing Abrupt SSE/WebSocket Dropping or Indefinite Pod Termination Deadlocks	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 couple Axum long-lived streaming endpoints (SSE, WebSocket, WebTransport) with a tokio_util

Engineers MUST ALWAYS couple Axum long-lived streaming endpoints (SSE, WebSocket, WebTransport) with a tokio_util::sync::CancellationToken to signal connection draining.

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.