Skeleton: axum, tracing, graceful shutdown

Dependencies

Drop this into infergw/Cargo.toml:

[package]
name = "infergw"
version = "0.1.0"
edition.workspace = true
rust-version.workspace = true

[dependencies]
axum               = "0.7"
tokio              = { version = "1", features = ["macros", "rt-multi-thread", "signal", "time", "sync"] }
tower              = { version = "0.4", features = ["timeout", "limit"] }
tower-http         = { version = "0.5", features = ["trace", "timeout"] }
hyper              = { version = "1", features = ["full"] }
tracing            = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
serde              = { version = "1", features = ["derive"] }
serde_json         = "1"
anyhow             = "1"
thiserror          = "1"

That's the floor — we'll add reqwest, futures, metrics, opentelemetry, and friends in later chapters.

A small Config struct

Read everything from env vars at startup. Twelve-factor by default. Create infergw/src/config.rs:

use std::env;

#[derive(Debug, Clone)]
pub struct Config {
    pub bind_addr: String,
    pub upstream_url: String,
    pub upstream_timeout_ms: u64,
    pub shutdown_grace_ms: u64,
}

impl Config {
    pub fn from_env() -> anyhow::Result<Self> {
        Ok(Self {
            bind_addr: env::var("INFERGW_BIND").unwrap_or_else(|_| "0.0.0.0:8080".into()),
            upstream_url: env::var("INFERGW_UPSTREAM")
                .unwrap_or_else(|_| "http://127.0.0.1:9000".into()),
            upstream_timeout_ms: env::var("INFERGW_UPSTREAM_TIMEOUT_MS")
                .ok()
                .and_then(|s| s.parse().ok())
                .unwrap_or(30_000),
            shutdown_grace_ms: env::var("INFERGW_SHUTDOWN_GRACE_MS")
                .ok()
                .and_then(|s| s.parse().ok())
                .unwrap_or(15_000),
        })
    }
}

Defaults match the mock from Chapter 00. Override at runtime when we Docker-ize in Chapter 06.

Structured tracing

One of the points of choosing Rust here is that tracing gives you compile-time-checked structured logging that's also the foundation for OpenTelemetry spans in Chapter 05. Set it up once and benefit forever.

Create infergw/src/telemetry.rs:

use tracing_subscriber::{fmt, prelude::*, EnvFilter};

pub fn init_tracing() {
    let filter = EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| EnvFilter::new("infergw=info,tower_http=info"));

    // JSON when INFERGW_LOG_JSON=1, pretty otherwise.
    let json_mode = std::env::var("INFERGW_LOG_JSON")
        .ok()
        .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
        .unwrap_or(false);

    let registry = tracing_subscriber::registry().with(filter);

    if json_mode {
        registry.with(fmt::layer().json().with_current_span(true)).init();
    } else {
        registry.with(fmt::layer().with_target(true).compact()).init();
    }
}

Router and health endpoints

Two endpoints to start: /healthz (liveness — am I running?) and /readyz (readiness — should traffic come to me?). The K8s difference is real: failing readiness drains traffic, failing liveness restarts the pod.

Create infergw/src/routes.rs:

use axum::{http::StatusCode, response::IntoResponse, routing::get, Router};
use std::sync::{
    atomic::{AtomicBool, Ordering},
    Arc,
};

#[derive(Clone)]
pub struct AppState {
    pub ready: Arc<AtomicBool>,
}

impl AppState {
    pub fn new() -> Self {
        Self {
            ready: Arc::new(AtomicBool::new(true)),
        }
    }
}

pub fn router(state: AppState) -> Router {
    Router::new()
        .route("/healthz", get(healthz))
        .route("/readyz", get(readyz))
        .with_state(state)
}

async fn healthz() -> impl IntoResponse {
    (StatusCode::OK, "ok")
}

async fn readyz(
    axum::extract::State(state): axum::extract::State<AppState>,
) -> impl IntoResponse {
    if state.ready.load(Ordering::Relaxed) {
        (StatusCode::OK, "ready")
    } else {
        (StatusCode::SERVICE_UNAVAILABLE, "draining")
    }
}

Graceful shutdown

This is the bit most tutorials skip. When K8s sends SIGTERM, you want to: flip readyz to SERVICE_UNAVAILABLE (so the load balancer stops sending new requests), wait for in-flight requests to finish, then exit. If you can't drain in 15 seconds, give up and exit anyway.

Create infergw/src/shutdown.rs:

use crate::routes::AppState;
use std::{sync::atomic::Ordering, time::Duration};
use tokio::signal;

/// Future that completes on SIGINT or SIGTERM, flips readiness to false,
/// then yields so axum can finish in-flight requests up to the grace period.
pub async fn shutdown_signal(state: AppState, grace: Duration) {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install Ctrl-C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install SIGTERM handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c    => tracing::info!("received Ctrl-C, draining..."),
        _ = terminate => tracing::info!("received SIGTERM, draining..."),
    }

    // Stop accepting new requests at the LB layer.
    state.ready.store(false, Ordering::Relaxed);

    // Give in-flight requests a chance to complete.
    tokio::time::sleep(grace).await;
    tracing::info!("grace period elapsed, exiting");
}

Putting it together

Now the entry point. infergw/src/main.rs:

mod config;
mod routes;
mod shutdown;
mod telemetry;

use std::time::Duration;
use tower_http::trace::TraceLayer;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    telemetry::init_tracing();
    let cfg = config::Config::from_env()?;
    tracing::info!(?cfg, "starting infergw");

    let state = routes::AppState::new();
    let app = routes::router(state.clone()).layer(TraceLayer::new_for_http());

    let listener = tokio::net::TcpListener::bind(&cfg.bind_addr).await?;
    tracing::info!(addr = %cfg.bind_addr, "listening");

    let grace = Duration::from_millis(cfg.shutdown_grace_ms);
    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown::shutdown_signal(state, grace))
        .await?;

    tracing::info!("infergw shut down cleanly");
    Ok(())
}

That's the whole skeleton. ~150 lines of code, every line earning its keep.

Run it

From the workspace root:

cargo run -p infergw

You should see:

INFO infergw: starting infergw cfg=Config { bind_addr: "0.0.0.0:8080", ... }
INFO infergw: listening addr=0.0.0.0:8080

From another terminal:

curl -i http://127.0.0.1:8080/healthz
# HTTP/1.1 200 OK
# ...
# ok

curl -i http://127.0.0.1:8080/readyz
# HTTP/1.1 200 OK
# ...
# ready

Now press Ctrl-C in the server terminal. Watch what happens:

^C
INFO infergw::shutdown: received Ctrl-C, draining...
INFO infergw::shutdown: grace period elapsed, exiting
INFO infergw: infergw shut down cleanly

During the grace window, hit /readyz again from a third terminal — you'll get 503 Service Unavailable: draining. That's the load balancer's cue to stop routing to this pod.

// In routes.rs, add to the router:
.route("/slow", get(slow))

async fn slow() -> impl IntoResponse {
    tokio::time::sleep(std::time::Duration::from_secs(10)).await;
    "done"
}

Checkpoint

• cargo run -p infergw binds to :8080.
• GET /healthz returns 200 ok.
• GET /readyz returns 200 ready, switches to 503 draining on shutdown.
• Ctrl-C drains in-flight requests up to the grace period then exits cleanly.
• Logs are structured and respect RUST_LOG.

Keep both processes running into the next chapter — we're about to make the gateway actually do something.

→ Chapter 02: Proxy & streaming

✓ Check yourself

Answer from memory before opening each one.

1. Start the gateway and run curl -i against /readyz. Press Ctrl-C in the server terminal and run the same curl again within a few seconds. What do you see each time, and what does it mean if the second curl says connection refused?

   Show answer

   The first call returns 200 with ready. The second returns 503 with draining, because the signal handler flips the flag and then sleeps for the grace period while the server keeps answering. Connection refused means the process already exited, so the grace sleep is not running or the grace is set to zero. See graceful shutdown.

2. Why does the gateway expose both /healthz and /readyz, and which one must fail while it drains?

   Show answer

   Failing readiness takes the pod out of the Service so no new traffic arrives, and the process keeps running. Failing liveness makes Kubernetes restart the container. Draining must fail readiness only, otherwise the orchestrator kills the requests you are trying to finish. See router and health endpoints.

3. You press Ctrl-C on an idle gateway with default settings and it takes about 15 seconds to exit. Your orchestrator sends SIGKILL after 10 seconds. What do you change?

   Show answer

   The handler sleeps the full grace period before axum is told to stop, so an idle process still waits it out. Lower INFERGW_SHUTDOWN_GRACE_MS below the kill timeout, or raise the kill timeout above the grace plus your longest stream. Pick the grace from how long the load balancer needs to notice, not from habit. See graceful shutdown.

4. A teammate proposes replacing the sleep with a counter that waits until in-flight requests reach zero. What does that buy, and what must you keep?

   Show answer

   It exits as soon as the last request finishes instead of burning the full grace. Keep an upper bound, because one stuck stream would otherwise block exit forever. Keep a short minimum wait after flipping readiness too, or you stop accepting before the balancer has stopped sending. See graceful shutdown.

5. You export RUST_LOG=warn and the listening line vanishes from the logs. Why?

   Show answer

   The default filter is used only when RUST_LOG is unset. Setting it replaces the default, so the info-level lines from infergw are filtered out. Set RUST_LOG=infergw=info,warn to keep them. See structured tracing.

Takeaway

• Liveness answers whether to restart the process, and readiness answers whether to send it traffic. Draining fails readiness only.
• Graceful shutdown is an ordered sequence: flip readiness, give the balancer time to notice, stop accepting, let in-flight work finish, then exit.
• The grace sleep here is a fixed upper bound, so even an idle gateway takes the full grace to exit.
• The orchestrator's kill timeout has to exceed the grace plus the longest in-flight request, or you cut off the work you meant to drain.
• Config comes from environment variables with defaults that match the local mock, so one binary runs everywhere.
• One tracing setup gives structured logs now and is the base the OpenTelemetry spans attach to in chapter 05.

The gateway answers health checks but forwards nothing yet. → Proxy & streaming

Further reading

• axum crate documentation — docs.rs, accessed 2026. The router, extractors and state patterns used in every chapter.
• axum::serve::WithGracefulShutdown — docs.rs, accessed 2026. Exactly what the server does once the shutdown future completes.
• Graceful shutdown — Tokio project docs, accessed 2026. The pattern of detecting a signal, telling tasks to stop, and waiting for them.
• Configure liveness, readiness and startup probes — Kubernetes docs, accessed 2026. What each probe does to a pod when it fails.
• tracing-subscriber crate documentation — docs.rs, accessed 2026. Filters, formatting layers and the registry used in the telemetry module.