47. A production HTTP server with Axum: limits, timeouts, tracing, shutdown

Implements: HTTPS Ch 19 (the Go version) and Go §62.

Axum is a thin layer over Hyper (the HTTP implementation) and Tower (a middleware abstraction: Service<Request> -> Future<Response>). Its main idea is extractors: a handler's parameter types say what it needs from the request, and Axum parses and validates them before the handler runs.

// src/bin/api.rs
use axum::extract::{Path, State};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use axum::{Json, Router};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, RwLock};
use std::time::Duration;
use tower_http::{limit::RequestBodyLimitLayer, timeout::TimeoutLayer, trace::TraceLayer};

#[derive(Clone, Default)]
struct AppState {                                   // cloned per request: keep it cheap (Arcs inside)
    users: Arc<RwLock<HashMap<u64, User>>>,
    next_id: Arc<AtomicU64>,
}

#[derive(Clone, Serialize)]
struct User { id: u64, name: String, email: String }

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]                       // reject {"name":..,"is_admin":true} (mass assignment)
struct NewUser { name: String, email: String }

/// One error type for the whole API, mapped to HTTP in exactly one place.
enum ApiError { NotFound, Invalid(String) }

impl IntoResponse for ApiError {
    fn into_response(self) -> Response {
        let (status, msg) = match self {
            ApiError::NotFound => (StatusCode::NOT_FOUND, "not found".to_string()),
            ApiError::Invalid(m) => (StatusCode::UNPROCESSABLE_ENTITY, m),
        };
        (status, Json(serde_json::json!({ "error": msg }))).into_response()
    }
}

async fn healthz() -> &'static str { "ok" }

async fn create_user(State(st): State<AppState>, Json(body): Json<NewUser>) -> Result<(StatusCode, Json<User>), ApiError> {
    let name = body.name.trim();
    if name.is_empty() || name.len() > 100 { return Err(ApiError::Invalid("name must be 1-100 characters".into())); }
    if !body.email.contains('@') || body.email.len() > 254 { return Err(ApiError::Invalid("invalid email".into())); }
    let id = st.next_id.fetch_add(1, Ordering::Relaxed) + 1;
    let user = User { id, name: name.into(), email: body.email.to_lowercase() };
    st.users.write().unwrap().insert(id, user.clone());   // lock held briefly, never across .await
    Ok((StatusCode::CREATED, Json(user)))
}

async fn get_user(State(st): State<AppState>, Path(id): Path<u64>) -> Result<Json<User>, ApiError> {
    st.users.read().unwrap().get(&id).cloned().map(Json).ok_or(ApiError::NotFound)
}

fn app(state: AppState) -> Router {
    Router::new()
        .route("/healthz", get(healthz))
        .route("/users", post(create_user))
        .route("/users/{id}", get(get_user))               // axum 0.8 path syntax: {id}
        .layer(TraceLayer::new_for_http())                 // one span per request, with latency
        .layer(TimeoutLayer::with_status_code(StatusCode::SERVICE_UNAVAILABLE, Duration::from_secs(10)))
        .layer(RequestBodyLimitLayer::new(64 * 1024))      // 413 for bodies over 64 KiB
        .with_state(state)
}

async fn shutdown_signal() {
    let ctrl_c = async { tokio::signal::ctrl_c().await.expect("install Ctrl-C handler") };
    #[cfg(unix)]
    let term = async {
        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
            .expect("install SIGTERM handler")
            .recv()
            .await;
    };
    #[cfg(not(unix))]
    let term = std::future::pending::<()>();
    tokio::select! { _ = ctrl_c => {}, _ = term => {} }
    tracing::info!("shutdown signal received, draining connections");
}

#[tokio::main]
async fn main() -> std::io::Result<()> {
    tracing_subscriber::fmt().with_target(false).init();
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await?;
    tracing::info!("listening on {}", listener.local_addr()?);
    axum::serve(listener, app(AppState::default()))
        .with_graceful_shutdown(shutdown_signal())        // stop accepting, finish in-flight requests
        .await
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::{to_bytes, Body};
    use axum::http::Request;
    use tower::ServiceExt;                                // .oneshot(): call the Router without a socket

    async fn call(app: &Router, method: &str, uri: &str, body: &str) -> (StatusCode, String) {
        let req = Request::builder().method(method).uri(uri)
            .header("content-type", "application/json").body(Body::from(body.to_string())).unwrap();
        let resp = app.clone().oneshot(req).await.unwrap();
        let status = resp.status();
        let bytes = to_bytes(resp.into_body(), usize::MAX).await.unwrap();
        (status, String::from_utf8_lossy(&bytes).into_owned())
    }

    #[tokio::test]
    async fn user_lifecycle_and_limits() {
        let app = app(AppState::default());
        let (s, b) = call(&app, "POST", "/users", r#"{"name":"Ada","email":"ADA@x.io"}"#).await;
        assert_eq!(s, StatusCode::CREATED);
        assert!(b.contains(r#""email":"ada@x.io""#));
        assert_eq!(call(&app, "GET", "/users/1", "").await.0, StatusCode::OK);
        assert_eq!(call(&app, "GET", "/users/99", "").await.0, StatusCode::NOT_FOUND);
        assert_eq!(call(&app, "GET", "/users/abc", "").await.0, StatusCode::BAD_REQUEST);  // Path<u64> rejected
        assert_eq!(call(&app, "POST", "/users", r#"{"name":"x","email":"x@y","is_admin":true}"#).await.0,
                   StatusCode::UNPROCESSABLE_ENTITY);                                    // deny_unknown_fields
        assert_eq!(call(&app, "POST", "/users", r#"{"name":"","email":"a@b"}"#).await.0, StatusCode::UNPROCESSABLE_ENTITY);
        let huge = format!(r#"{{"name":"{}","email":"a@b"}}"#, "x".repeat(100_000));
        assert_eq!(call(&app, "POST", "/users", &huge).await.0, StatusCode::PAYLOAD_TOO_LARGE);
    }
}

cargo run --bin api &
curl -s localhost:3000/users -H 'content-type: application/json' -d '{"name":"Ada","email":"ada@example.com"}'
# {"id":1,"name":"Ada","email":"ada@example.com"}
curl -s -o /dev/null -w '%{http_code}\n' localhost:3000/users/42      # 404
kill -TERM %1                                                          # graceful drain, exit 0
cargo test --bin api

What each layer defends against, mapped to the HTTPS guide:

One gap to know about: axum::serve does not expose a header read timeout. A Slowloris client that sends one header byte every few seconds is limited only by the OS and by TimeoutLayer, and TimeoutLayer starts counting only after the headers have been parsed. Production deployments either sit behind a proxy that enforces it (nginx, Envoy, or a cloud load balancer), or serve the Router with hyper-util's connection builder, which has header_read_timeout. Ch 50 shows that builder.