chio_http_serve/
drain.rs

1use std::future::Future;
2use std::time::Duration;
3
4use tokio::sync::watch;
5use tracing::warn;
6
7use crate::signal::wait_for_shutdown;
8
9/// How a bounded drain finished.
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub enum DrainOutcome {
12    /// Every in-flight request completed before the drain deadline.
13    Clean,
14    /// The drain deadline elapsed and the remaining connections were
15    /// force-closed.
16    Forced,
17}
18
19/// Terminal error from a bounded serve. Either variant is a non-zero process
20/// exit; neither is ever swallowed.
21#[derive(Debug, thiserror::Error)]
22pub enum ServeError {
23    /// The serve future itself failed.
24    #[error("serve I/O error: {0}")]
25    Io(#[from] std::io::Error),
26    /// The post-drain flush of a receipt store failed. This is loud on purpose:
27    /// it means a receipt may not be durable, so the operator must see a
28    /// non-zero exit rather than a silent success.
29    #[error("post-drain receipt flush failed: {0}")]
30    Flush(String),
31}
32
33/// Run a graceful serve future to completion, bounding only the post-signal
34/// drain, then run a flush hook.
35///
36/// `server` is the site's `axum::serve(..).with_graceful_shutdown(ctrl.signalled())`
37/// future, `shutdown` is a receiver from the same controller, and `on_drained`
38/// flushes the site's receipt store (or is a no-op).
39///
40/// The drain deadline arms only once the shutdown signal fires, never at server
41/// start: a naive timeout around the whole serve future would kill a healthy
42/// server `drain_timeout` after boot. So the serve runs unbounded until either
43/// it ends on its own or the signal arrives; only then does the remaining
44/// in-flight drain get a deadline.
45///
46/// Flush-hook contract: on a clean drain the flush runs after every accepted
47/// connection has finished, so it has the store to itself. On a forced drain a
48/// connection outlived the deadline; axum serves each connection on a detached
49/// task that a bounded drain can neither join nor cancel, so `on_drained` can run
50/// while such a straggler is still finishing and MUST be safe against a concurrent
51/// writer. Sites keep this safe by making every acknowledged receipt durable
52/// synchronously in the handler (or via a commit actor that acks only after the
53/// write reaches the WAL), never in the flush hook, so a straggler can never cost
54/// an acknowledged receipt.
55pub async fn run_until_drained<S, D>(
56    server: S,
57    shutdown: watch::Receiver<bool>,
58    drain_timeout: Duration,
59    on_drained: D,
60) -> Result<DrainOutcome, ServeError>
61where
62    // `axum::serve(..).with_graceful_shutdown(..)` is an `IntoFuture`, not a
63    // `Future`; accept it directly so a call site passes it as-is.
64    S: std::future::IntoFuture<Output = std::io::Result<()>>,
65    D: Future<Output = Result<(), String>>,
66{
67    // Box-pin rather than stack-pin so the forced-drain path can drop the serve
68    // future to force-close a connection still stuck at the deadline. A stack
69    // `Pin<&mut _>` only borrows and cannot be dropped early; an owned
70    // `Pin<Box<_>>` can.
71    let mut server = Box::pin(server.into_future());
72    let signalled = wait_for_shutdown(shutdown);
73
74    let outcome = tokio::select! {
75        // Phase 1: serve until the server exits on its own or the signal fires.
76        result = &mut server => match result {
77            Ok(()) => DrainOutcome::Clean,
78            Err(source) => return Err(ServeError::Io(source)),
79        },
80        () = signalled => {
81            // Phase 2: graceful shutdown has stopped the accept loop; bound only
82            // the remaining in-flight drain.
83            match tokio::time::timeout(drain_timeout, &mut server).await {
84                Ok(Ok(())) => DrainOutcome::Clean,
85                Ok(Err(source)) => return Err(ServeError::Io(source)),
86                Err(_elapsed) => {
87                    let drain_ms = u64::try_from(drain_timeout.as_millis()).unwrap_or(u64::MAX);
88                    warn!(
89                        drain_timeout_ms = drain_ms,
90                        "drain deadline exceeded; force-closing remaining connections"
91                    );
92                    DrainOutcome::Forced
93                }
94            }
95        }
96    };
97
98    // On a forced drain the serve future is still parked waiting on a connection
99    // that outlived the deadline. Drop it to release the accept loop and the
100    // listener before the flush. This does not guarantee the connection's handler
101    // has stopped: axum runs it on a detached task, so it may still be finishing
102    // (see the flush-hook contract above). A clean outcome already ran the future
103    // to completion, so the drop is a no-op there.
104    if matches!(outcome, DrainOutcome::Forced) {
105        drop(server);
106    }
107
108    on_drained.await.map_err(ServeError::Flush)?;
109    Ok(outcome)
110}
chio_http_serve/drain.rs

chio_http_serve/
drain.rs
