Graceful Shutdown

use std::sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
};
use std::time::Duration;

use axum::{Router, extract::State, http::StatusCode, routing::get};
use tokio::signal;
use tokio_util::{sync::CancellationToken, task::TaskTracker};
use tower_http::timeout::TimeoutLayer;

/// How long to keep serving after readiness starts failing.
const DRAIN_DELAY: Duration = Duration::from_secs(2);

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

/// Readiness: 200 while accepting work, 503 once shutdown has started.
async fn ready(State(state): State<AppState>) -> StatusCode {
    if state.ready.load(Ordering::Relaxed) {
        StatusCode::OK
    } else {
        StatusCode::SERVICE_UNAVAILABLE
    }
}

/// Liveness: the process is running; never fails during shutdown.
async fn live() -> StatusCode {
    StatusCode::OK
}

/// Resolves on Ctrl+C or SIGTERM, after marking the instance not ready and waiting `DRAIN_DELAY`.
async fn shutdown_signal(ready: Arc<AtomicBool>) {
    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 => {},
        () = terminate => {},
    }
    ready.store(false, Ordering::Relaxed); // GET /ready now answers 503
    println!("shutdown signal received: not ready, draining for {DRAIN_DELAY:?}");
    tokio::time::sleep(DRAIN_DELAY).await; // still serving while load balancers notice
}

async fn slow_report() -> &'static str {
    tokio::time::sleep(Duration::from_secs(3)).await; // an in-flight request during shutdown
    "report finished\n"
}

#[tokio::main]
async fn main() {
    let state = AppState { ready: Arc::new(AtomicBool::new(true)) };
    let tracker = TaskTracker::new();
    let cancel = CancellationToken::new();

    // a background worker that stops when told to
    let worker_cancel = cancel.clone();
    tracker.spawn(async move {
        let mut tick = tokio::time::interval(Duration::from_secs(1));
        loop {
            tokio::select! {
                () = worker_cancel.cancelled() => { println!("worker: finishing"); break; }
                _ = tick.tick() => { /* process a batch */ }
            }
        }
    });

    let app = Router::new()
        .route("/reports/slow", get(slow_report))
        // bounds how long any single request can delay shutdown
        .layer(TimeoutLayer::with_status_code(StatusCode::SERVICE_UNAVAILABLE, Duration::from_secs(10)))
        // probes are added after the timeout layer, so they are not wrapped by it
        .route("/ready", get(ready))
        .route("/live", get(live))
        .with_state(state.clone());

    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).with_graceful_shutdown(shutdown_signal(state.ready)).await.unwrap();
    println!("server stopped: no request in flight");

    cancel.cancel(); // 1. stop background work
    tracker.close();
    if tokio::time::timeout(Duration::from_secs(5), tracker.wait()).await.is_err() {
        println!("workers did not stop in time");
    }
    // 2. close pools: pool.close().await;  3. flush telemetry: tracer_provider.shutdown()
    println!("clean exit");
}

$ curl -s -o /dev/null -w '%{http_code}\n' localhost:3000/ready
200
$ curl -s localhost:3000/reports/slow & sleep 0.5; kill -TERM $(pgrep bookshelf-web); sleep 0.5
shutdown signal received: not ready, draining for 2s
$ curl -s -o /dev/null -w '%{http_code}\n' localhost:3000/ready   # during the drain delay
503
report finished                      # the in-flight request completed
server stopped: no request in flight
worker: finishing
clean exit

use std::{convert::Infallible, time::Duration};

use axum::{
    Router,
    extract::{
        State,
        ws::{CloseFrame, Message, WebSocket, WebSocketUpgrade, close_code},
    },
    response::{
        IntoResponse, Sse,
        sse::{Event, KeepAlive},
    },
    routing::get,
};
use futures::{Stream, StreamExt};
use tokio_stream::wrappers::IntervalStream;
use tokio_util::{sync::CancellationToken, task::TaskTracker};

#[derive(Clone)]
struct AppState {
    shutdown: CancellationToken,
    tracker: TaskTracker,
}

async fn events(State(state): State<AppState>) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    let ticks = IntervalStream::new(tokio::time::interval(Duration::from_secs(1)))
        .map(|_| Ok(Event::default().event("loan").data("tick")));
    // the stream -- and with it the response body -- ends when shutdown starts
    Sse::new(ticks.take_until(state.shutdown.cancelled_owned())).keep_alive(KeepAlive::default())
}

async fn ws(ws: WebSocketUpgrade, State(state): State<AppState>) -> impl IntoResponse {
    let tracker = state.tracker.clone();
    // track the socket task so main can wait for its close frame
    ws.on_upgrade(move |socket| tracker.track_future(echo(socket, state.shutdown)))
}

async fn echo(mut socket: WebSocket, shutdown: CancellationToken) {
    loop {
        tokio::select! {
            () = shutdown.cancelled() => {
                let frame = CloseFrame { code: close_code::AWAY, reason: "server shutting down".into() };
                let _ = socket.send(Message::Close(Some(frame))).await;
                break;
            }
            msg = socket.recv() => match msg {
                Some(Ok(Message::Text(text))) => {
                    if socket.send(Message::Text(text)).await.is_err() {
                        break;
                    }
                }
                Some(Ok(_)) => {}
                Some(Err(_)) | None => break,
            },
        }
    }
}

async fn shutdown_signal() {
    let ctrl_c = async { tokio::signal::ctrl_c().await.expect("failed to install Ctrl+C handler") };
    #[cfg(unix)]
    let terminate = async {
        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
            .expect("failed to install SIGTERM handler")
            .recv()
            .await;
    };
    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();
    tokio::select! {
        () = ctrl_c => {},
        () = terminate => {},
    }
}

#[tokio::main]
async fn main() {
    let shutdown = CancellationToken::new();
    let tracker = TaskTracker::new();
    let state = AppState { shutdown: shutdown.clone(), tracker: tracker.clone() };
    let app = Router::new().route("/events", get(events)).route("/ws", get(ws)).with_state(state);

    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app)
        .with_graceful_shutdown(async move {
            shutdown_signal().await;
            shutdown.cancel(); // end SSE streams and close WebSockets *before* hyper waits for connections
        })
        .await
        .unwrap();

    tracker.close();
    tracker.wait().await; // every WebSocket has sent its close frame
    println!("clean exit");
}

use std::time::Duration;

use axum::{Router, routing::get};
use tokio_util::sync::CancellationToken;

const DRAIN_DEADLINE: Duration = Duration::from_secs(20);

async fn stream_forever() -> axum::body::Body {
    // a body that never ends and ignores shutdown -- the case the deadline is for
    let stream = futures::stream::pending::<Result<axum::body::Bytes, std::io::Error>>();
    axum::body::Body::from_stream(stream)
}

#[tokio::main]
async fn main() {
    let app = Router::new().route("/stuck", get(stream_forever));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();

    let started = CancellationToken::new();
    let signal_started = started.clone();
    let server = axum::serve(listener, app).with_graceful_shutdown(async move {
        tokio::signal::ctrl_c().await.expect("failed to install Ctrl+C handler");
        signal_started.cancel(); // starts the deadline clock
    });
    let deadline = async {
        started.cancelled().await;
        tokio::time::sleep(DRAIN_DEADLINE).await;
    };

    tokio::select! {
        result = server => result.expect("server error"),
        () = deadline => eprintln!("drain deadline exceeded: abandoning open connections"),
    }
    // cleanup runs in both cases: stop workers, close pools, flush telemetry
    println!("exiting");
}

spec:
  terminationGracePeriodSeconds: 30        # SIGKILL after this; must exceed preStop + drain + cleanup
  containers:
    - name: bookshelf-web
      lifecycle:
        preStop:
          sleep:
            seconds: 5                     # keep serving while endpoints are updated (sleep action, Kubernetes 1.30+)
      readinessProbe:
        httpGet: { path: /ready, port: 3000 }
        periodSeconds: 2
        failureThreshold: 1
      livenessProbe:
        httpGet: { path: /live, port: 3000 }
        periodSeconds: 10