Skip to main content

baukit_runtime/
lib.rs

1//! Process-wide lifecycle and listener composition for Baukit services.
2//!
3//! This crate owns mechanics shared by an entire process: shutdown signaling,
4//! bounded draining, background-task supervision, build identity, and the thin
5//! composition of public and operations listeners. Products continue to own
6//! configuration, routes, and application state.
7//!
8//! # Composition example
9//!
10//! A composition root can wire the crate to `baukit-config` and
11//! `baukit-telemetry` without making runtime own either product configuration or
12//! telemetry policy:
13//!
14//! ```no_run
15//! use std::sync::Arc;
16//! use axum::Router;
17//! use baukit_runtime::{
18//!     ProcessKind, RestartPolicy, ServiceInfo, ShutdownToken, TaskSupervisor,
19//!     ShutdownOrder, build_info, serve_listener_pair_with_shutdown_order,
20//! };
21//! use baukit_telemetry::{DeploymentEnvironment, TelemetryBuilder};
22//! use tokio::net::TcpListener;
23//! # use std::time::Duration;
24//! # async fn run_outbox_until(_: ShutdownToken) {}
25//! # fn product_routes() -> Router { Router::new() }
26//! # fn operations_routes() -> Router { Router::new() }
27//!
28//! #[tokio::main]
29//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
30//!     let service = ServiceInfo::new(
31//!         "orders",
32//!         ProcessKind::Api,
33//!         build_info!(),
34//!         DeploymentEnvironment::Local,
35//!     );
36//!     let telemetry = Arc::new(
37//!         TelemetryBuilder::new(service.telemetry_identity().clone()).init()?,
38//!     );
39//!     let shutdown = ShutdownToken::new(Duration::from_secs(30));
40//!     let signal_task = shutdown.spawn_signal_listener();
41//!
42//!     let mut workers = TaskSupervisor::new(shutdown.clone());
43//!     let worker_shutdown = shutdown.child_token();
44//!     workers.spawn("outbox", RestartPolicy::Restart { max_restarts: 3 }, move || {
45//!         let worker_shutdown = worker_shutdown.clone();
46//!         async move { run_outbox_until(worker_shutdown).await }
47//!     });
48//!
49//!     let api = TcpListener::bind("127.0.0.1:0").await?;
50//!     let ops = TcpListener::bind("127.0.0.1:0").await?;
51//!     let listener_result = serve_listener_pair_with_shutdown_order(
52//!         api,
53//!         Router::new().merge(product_routes()),
54//!         ops,
55//!         Router::new().merge(operations_routes()),
56//!         shutdown.clone(),
57//!         ShutdownOrder::OpsOutlivesApi,
58//!     ).await;
59//!     shutdown.trigger();
60//!     let _ = workers.join().await;
61//!
62//!     // Telemetry shutdown is synchronous, so keep it off the async executor.
63//!     // `run_during_drain` still bounds it by the process-wide drain deadline.
64//!     let telemetry_for_flush = Arc::clone(&telemetry);
65//!     shutdown.run_during_drain(async move {
66//!         tokio::task::spawn_blocking(move || telemetry_for_flush.shutdown()).await
67//!     }).await???;
68//!
69//!     signal_task.abort();
70//!     listener_result?;
71//!     Ok(())
72//! }
73//! ```
74
75#![deny(missing_docs)]
76
77use std::{
78    fmt,
79    future::Future,
80    io,
81    panic::{AssertUnwindSafe, catch_unwind},
82    sync::{Arc, Mutex, OnceLock},
83    time::Duration,
84};
85
86use axum::Router;
87pub use baukit_core::{BuildInfo, DeploymentEnvironment, ProcessKind, ServiceIdentity};
88use thiserror::Error;
89use tokio::{
90    net::TcpListener,
91    task::{JoinError, JoinHandle, JoinSet},
92    time::{Instant, timeout},
93};
94use tokio_util::sync::CancellationToken;
95
96/// The Rust compiler version used to build this crate.
97///
98/// [`build_info!`] combines this with package metadata captured at its call site.
99pub const RUST_VERSION: &str = env!("BAUKIT_RUST_VERSION");
100
101/// Captures build metadata from the binary crate in which the macro is invoked.
102///
103/// `GIT_COMMIT` should be set by the build or release pipeline. It falls back to
104/// `"unknown"` for local builds. The Rust compiler version is recorded by
105/// `baukit-runtime`'s build script.
106#[macro_export]
107macro_rules! build_info {
108    () => {
109        $crate::BuildInfo::new(
110            env!("CARGO_PKG_VERSION"),
111            option_env!("GIT_COMMIT").unwrap_or("unknown"),
112            $crate::RUST_VERSION,
113        )
114    };
115}
116
117/// Canonical service identity plus its complete process build metadata.
118///
119/// The service name follows the telemetry convention `<product>-<process>`.
120/// Keeping the canonical [`ServiceIdentity`] inside this value ensures runtime,
121/// logging, metrics, and traces cannot disagree about the process identity.
122#[derive(Clone, Debug, Eq, PartialEq)]
123pub struct ServiceInfo {
124    name: String,
125    identity: ServiceIdentity,
126    build: BuildInfo,
127}
128
129impl ServiceInfo {
130    /// Creates service information using Baukit's canonical service-name format.
131    pub fn new(
132        product: impl Into<String>,
133        process: ProcessKind,
134        build: BuildInfo,
135        environment: DeploymentEnvironment,
136    ) -> Self {
137        let product = product.into();
138        let identity = ServiceIdentity::new(
139            product,
140            process,
141            build.version().to_owned(),
142            build.commit().to_owned(),
143            environment,
144        );
145        let name = identity.service_name();
146        Self {
147            name,
148            identity,
149            build,
150        }
151    }
152
153    /// Returns the canonical `<product>-<process>` service name.
154    pub fn name(&self) -> &str {
155        &self.name
156    }
157
158    /// Returns the stable product identifier.
159    pub fn product(&self) -> &str {
160        self.identity.product()
161    }
162
163    /// Returns this process's kind.
164    pub const fn process(&self) -> ProcessKind {
165        self.identity.process()
166    }
167
168    /// Returns the service's build metadata.
169    pub const fn build(&self) -> &BuildInfo {
170        &self.build
171    }
172
173    /// Returns the identity accepted directly by `baukit-telemetry`.
174    pub const fn telemetry_identity(&self) -> &ServiceIdentity {
175        &self.identity
176    }
177}
178
179struct ShutdownState {
180    root: CancellationToken,
181    started_at: OnceLock<Instant>,
182    drain_hooks: Mutex<Vec<DrainHook>>,
183}
184
185type DrainHook = Box<dyn FnOnce() + Send + 'static>;
186
187impl fmt::Debug for ShutdownState {
188    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
189        let hook_count = self
190            .drain_hooks
191            .lock()
192            .unwrap_or_else(std::sync::PoisonError::into_inner)
193            .len();
194        formatter
195            .debug_struct("ShutdownState")
196            .field("root", &self.root)
197            .field("started_at", &self.started_at)
198            .field("drain_hook_count", &hook_count)
199            .finish()
200    }
201}
202
203/// A cloneable, awaitable process-shutdown handle with one shared drain deadline.
204///
205/// Child tokens observe parent shutdown, while calling [`ShutdownToken::trigger`]
206/// on any child still initiates process-wide shutdown. The deadline begins with
207/// the first trigger and is shared by all clones and children.
208#[derive(Clone, Debug)]
209pub struct ShutdownToken {
210    observed: CancellationToken,
211    state: Arc<ShutdownState>,
212    drain_timeout: Duration,
213}
214
215impl ShutdownToken {
216    /// Creates a shutdown handle with the maximum duration allowed for draining.
217    pub fn new(drain_timeout: Duration) -> Self {
218        let root = CancellationToken::new();
219        Self {
220            observed: root.clone(),
221            state: Arc::new(ShutdownState {
222                root,
223                started_at: OnceLock::new(),
224                drain_hooks: Mutex::new(Vec::new()),
225            }),
226            drain_timeout,
227        }
228    }
229
230    /// Returns a child which is cancelled by this token.
231    ///
232    /// Triggering the child deliberately triggers the shared process token too.
233    #[must_use]
234    pub fn child_token(&self) -> Self {
235        Self {
236            observed: self.observed.child_token(),
237            state: Arc::clone(&self.state),
238            drain_timeout: self.drain_timeout,
239        }
240    }
241
242    /// Initiates process-wide shutdown and returns whether this was the first trigger.
243    pub fn trigger(&self) -> bool {
244        let first = self.state.started_at.set(Instant::now()).is_ok();
245        if first {
246            let hooks = std::mem::take(
247                &mut *self
248                    .state
249                    .drain_hooks
250                    .lock()
251                    .unwrap_or_else(std::sync::PoisonError::into_inner),
252            );
253            for hook in hooks {
254                if catch_unwind(AssertUnwindSafe(hook)).is_err() {
255                    tracing::error!("a shutdown drain hook panicked");
256                }
257            }
258            self.state.root.cancel();
259        }
260        first
261    }
262
263    /// Registers a fast synchronous action to run when drain begins.
264    ///
265    /// Hooks run exactly once on the first [`ShutdownToken::trigger`], before
266    /// cancellation wakes asynchronous drain work. A hook registered after
267    /// drain has started runs immediately. This is the dependency-neutral
268    /// composition seam for closing `baukit_ops::TrafficGate` without a watcher
269    /// task:
270    ///
271    /// ```no_run
272    /// # use std::time::Duration;
273    /// # use baukit_runtime::ShutdownToken;
274    /// # #[derive(Clone)] struct TrafficGate;
275    /// # impl TrafficGate { fn stop_accepting(&self) {} }
276    /// # let traffic_gate = TrafficGate;
277    /// let shutdown = ShutdownToken::new(Duration::from_secs(30));
278    /// shutdown.on_drain({
279    ///     let traffic_gate = traffic_gate.clone();
280    ///     move || traffic_gate.stop_accepting()
281    /// });
282    /// ```
283    ///
284    /// Hooks should only flip in-memory state or perform similarly bounded work;
285    /// asynchronous cleanup belongs in [`ShutdownToken::run_during_drain`].
286    pub fn on_drain<F>(&self, hook: F)
287    where
288        F: FnOnce() + Send + 'static,
289    {
290        let hook: DrainHook = Box::new(hook);
291        let late_hook = {
292            let mut hooks = self
293                .state
294                .drain_hooks
295                .lock()
296                .unwrap_or_else(std::sync::PoisonError::into_inner);
297            if self.state.started_at.get().is_none() {
298                hooks.push(hook);
299                None
300            } else {
301                Some(hook)
302            }
303        };
304        if let Some(hook) = late_hook
305            && catch_unwind(AssertUnwindSafe(hook)).is_err()
306        {
307            tracing::error!("a late shutdown drain hook panicked");
308        }
309    }
310
311    /// Returns whether this token has observed shutdown.
312    pub fn is_cancelled(&self) -> bool {
313        self.observed.is_cancelled()
314    }
315
316    /// Waits until process-wide shutdown reaches this token.
317    pub async fn cancelled(&self) {
318        self.observed.cancelled().await;
319    }
320
321    /// Returns the configured drain timeout.
322    pub const fn drain_timeout(&self) -> Duration {
323        self.drain_timeout
324    }
325
326    /// Returns the shared absolute drain deadline after shutdown has begun.
327    pub fn deadline(&self) -> Option<Instant> {
328        self.state
329            .started_at
330            .get()
331            .map(|started_at| *started_at + self.drain_timeout)
332    }
333
334    /// Waits for shutdown, then runs `future` within the remaining drain budget.
335    ///
336    /// Every call uses the same absolute deadline, so sequential cleanup steps
337    /// cannot accidentally receive a fresh timeout. On expiry, the future is
338    /// dropped and [`DrainTimeout`] is returned.
339    ///
340    /// # Errors
341    ///
342    /// Returns [`DrainTimeout`] when `future` does not finish before the deadline.
343    pub async fn run_during_drain<F>(&self, future: F) -> Result<F::Output, DrainTimeout>
344    where
345        F: Future,
346    {
347        self.cancelled().await;
348        let remaining = self.remaining();
349        timeout(remaining, future)
350            .await
351            .map_err(|_| DrainTimeout::new(self.drain_timeout))
352    }
353
354    /// Waits for SIGINT/SIGTERM (or the platform Ctrl-C event), then triggers shutdown.
355    ///
356    /// # Errors
357    ///
358    /// Fails like [`wait_for_shutdown_signal`].
359    pub async fn shutdown_on_signal(&self) -> io::Result<ShutdownSignal> {
360        let signal = wait_for_shutdown_signal().await?;
361        self.trigger();
362        Ok(signal)
363    }
364
365    /// Spawns a Tokio task which triggers shutdown on the first process signal.
366    ///
367    /// The returned task reports signal-handler setup failures when joined.
368    #[must_use = "retain or join the signal listener so setup failures are observed"]
369    pub fn spawn_signal_listener(&self) -> tokio::task::JoinHandle<io::Result<ShutdownSignal>> {
370        let shutdown = self.clone();
371        tokio::spawn(async move { shutdown.shutdown_on_signal().await })
372    }
373
374    fn remaining(&self) -> Duration {
375        self.deadline()
376            .map(|deadline| deadline.saturating_duration_since(Instant::now()))
377            .unwrap_or(self.drain_timeout)
378    }
379}
380
381/// The operating-system event which initiated shutdown.
382#[derive(Clone, Copy, Debug, Eq, PartialEq)]
383pub enum ShutdownSignal {
384    /// Ctrl-C or SIGINT was received.
385    Interrupt,
386    /// SIGTERM was received on Unix.
387    Terminate,
388}
389
390/// Waits for the first supported process shutdown signal without triggering a token.
391///
392/// Factoring signal observation from [`ShutdownToken::trigger`] keeps lifecycle
393/// behavior directly testable without delivering real process signals.
394///
395/// # Errors
396///
397/// Returns an [`io::Error`] when the SIGTERM or Ctrl-C handler cannot be
398/// installed, or when the SIGTERM stream ends without a signal.
399#[cfg(unix)]
400pub async fn wait_for_shutdown_signal() -> io::Result<ShutdownSignal> {
401    use tokio::signal::unix::{SignalKind, signal};
402
403    let mut terminate = signal(SignalKind::terminate())?;
404    tokio::select! {
405        ctrl_c = tokio::signal::ctrl_c() => {
406            ctrl_c?;
407            Ok(ShutdownSignal::Interrupt)
408        }
409        received = terminate.recv() => received.map_or_else(
410            || Err(io::Error::other("SIGTERM signal stream ended")),
411            |_| Ok(ShutdownSignal::Terminate),
412        ),
413    }
414}
415
416/// Waits for the platform Ctrl-C event without triggering a token.
417///
418/// # Errors
419///
420/// Returns an [`io::Error`] when the Ctrl-C handler cannot be installed.
421#[cfg(not(unix))]
422pub async fn wait_for_shutdown_signal() -> io::Result<ShutdownSignal> {
423    tokio::signal::ctrl_c().await?;
424    Ok(ShutdownSignal::Interrupt)
425}
426
427/// Indicates that process draining exceeded its shared deadline.
428#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]
429#[error("shutdown drain deadline expired after {timeout:?}")]
430pub struct DrainTimeout {
431    timeout: Duration,
432}
433
434impl DrainTimeout {
435    const fn new(timeout: Duration) -> Self {
436        Self { timeout }
437    }
438
439    /// Returns the configured total drain timeout.
440    pub const fn timeout(self) -> Duration {
441        self.timeout
442    }
443}
444
445/// The response to a supervised task's unexpected termination.
446#[derive(Clone, Copy, Debug, Eq, PartialEq)]
447pub enum RestartPolicy {
448    /// Trigger process shutdown immediately.
449    FailProcess,
450    /// Restart up to `max_restarts` times, then trigger process shutdown.
451    Restart {
452        /// Maximum starts after the initial attempt.
453        max_restarts: usize,
454    },
455}
456
457/// Owns named background tasks tied to one [`ShutdownToken`].
458pub struct TaskSupervisor {
459    shutdown: ShutdownToken,
460    tasks: JoinSet<()>,
461}
462
463impl TaskSupervisor {
464    /// Creates an empty task supervisor.
465    pub fn new(shutdown: ShutdownToken) -> Self {
466        Self {
467            shutdown,
468            tasks: JoinSet::new(),
469        }
470    }
471
472    /// Spawns a named task from a restartable future factory.
473    ///
474    /// Both a normal return and a panic are unexpected before shutdown. The
475    /// selected `policy` either restarts the factory or triggers process shutdown.
476    pub fn spawn<F, Fut>(&mut self, name: impl Into<String>, policy: RestartPolicy, factory: F)
477    where
478        F: Fn() -> Fut + Send + Sync + 'static,
479        Fut: Future<Output = ()> + Send + 'static,
480    {
481        let shutdown = self.shutdown.clone();
482        let name = name.into();
483        let factory = Arc::new(factory);
484        self.tasks.spawn(async move {
485            supervise_task(name, policy, factory, shutdown).await;
486        });
487    }
488
489    /// Waits for shutdown and joins all supervised tasks within its drain deadline.
490    ///
491    /// Any tasks still alive when the deadline expires are aborted before this
492    /// method returns [`DrainTimeout`].
493    ///
494    /// # Errors
495    ///
496    /// Returns [`DrainTimeout`] when tasks are still running at the drain deadline.
497    pub async fn join(mut self) -> Result<(), DrainTimeout> {
498        self.shutdown.cancelled().await;
499        let remaining = self.shutdown.remaining();
500        let joined = timeout(remaining, async {
501            while let Some(result) = self.tasks.join_next().await {
502                if let Err(error) = result {
503                    tracing::error!(%error, "task supervisor failed while joining a worker");
504                }
505            }
506        })
507        .await;
508
509        if joined.is_err() {
510            self.tasks.abort_all();
511            while self.tasks.join_next().await.is_some() {}
512            return Err(DrainTimeout::new(self.shutdown.drain_timeout));
513        }
514
515        Ok(())
516    }
517
518    /// Returns the number of supervised task loops currently owned.
519    pub fn len(&self) -> usize {
520        self.tasks.len()
521    }
522
523    /// Returns whether no supervised tasks are currently owned.
524    pub fn is_empty(&self) -> bool {
525        self.tasks.is_empty()
526    }
527}
528
529async fn supervise_task<F, Fut>(
530    name: String,
531    policy: RestartPolicy,
532    factory: Arc<F>,
533    shutdown: ShutdownToken,
534) where
535    F: Fn() -> Fut + Send + Sync + 'static,
536    Fut: Future<Output = ()> + Send + 'static,
537{
538    let mut restarts = 0_usize;
539    loop {
540        let mut attempt = AbortOnDropTask(tokio::spawn(factory()));
541        let outcome = tokio::select! {
542            _ = shutdown.cancelled() => {
543                if let Err(error) = (&mut attempt.0).await
544                    && error.is_panic()
545                {
546                    tracing::error!(task = %name, %error, "supervised task panicked while draining");
547                }
548                return;
549            }
550            result = &mut attempt.0 => result,
551        };
552
553        if shutdown.is_cancelled() {
554            return;
555        }
556
557        match &outcome {
558            Ok(()) => tracing::error!(task = %name, "supervised task exited unexpectedly"),
559            Err(error) if error.is_panic() => {
560                tracing::error!(task = %name, %error, "supervised task panicked")
561            }
562            Err(error) => tracing::error!(task = %name, %error, "supervised task was cancelled"),
563        }
564
565        match policy {
566            RestartPolicy::FailProcess => {
567                shutdown.trigger();
568                return;
569            }
570            RestartPolicy::Restart { max_restarts } if restarts < max_restarts => {
571                restarts += 1;
572                tracing::warn!(task = %name, restarts, max_restarts, "restarting supervised task");
573            }
574            RestartPolicy::Restart { max_restarts } => {
575                tracing::error!(task = %name, restarts, max_restarts, "task restart limit exhausted");
576                shutdown.trigger();
577                return;
578            }
579        }
580    }
581}
582
583struct AbortOnDropTask<T>(JoinHandle<T>);
584
585impl<T> Drop for AbortOnDropTask<T> {
586    fn drop(&mut self) {
587        self.0.abort();
588    }
589}
590
591/// Identifies one member of an API/operations listener pair.
592#[derive(Clone, Copy, Debug, Eq, PartialEq)]
593pub enum ListenerKind {
594    /// The product API listener.
595    Api,
596    /// The operations listener.
597    Operations,
598}
599
600impl fmt::Display for ListenerKind {
601    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
602        formatter.write_str(match self {
603            Self::Api => "API",
604            Self::Operations => "operations",
605        })
606    }
607}
608
609/// An error returned while running the API and operations listeners.
610#[derive(Debug, Error)]
611pub enum ListenerPairError {
612    /// A listener returned an I/O error.
613    #[error("{listener} listener failed: {source}")]
614    Listener {
615        /// The listener which failed.
616        listener: ListenerKind,
617        /// The underlying server error.
618        #[source]
619        source: io::Error,
620    },
621    /// A listener task panicked or was externally cancelled.
622    #[error("listener task failed: {0}")]
623    Task(#[source] JoinError),
624    /// A listener stopped successfully before shutdown was requested.
625    #[error("{0} listener exited before shutdown")]
626    UnexpectedExit(ListenerKind),
627    /// One or both listeners did not drain before the process deadline.
628    #[error(transparent)]
629    Drain(#[from] DrainTimeout),
630}
631
632type ListenerResult = Result<ListenerKind, (ListenerKind, io::Error)>;
633
634/// Controls when the operations listener stops relative to the public API.
635#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
636pub enum ShutdownOrder {
637    /// Stops both listeners together when process shutdown begins.
638    #[default]
639    Together,
640    /// Keeps operations serving until the API has drained or the deadline expires.
641    OpsOutlivesApi,
642}
643
644/// Runs public API and operations routers concurrently with graceful shutdown.
645///
646/// Both TCP listeners are supplied already bound, keeping addresses and socket
647/// policy in the composition root. If either listener fails, process shutdown is
648/// triggered and the peer is drained. Once shutdown begins, both Axum servers
649/// receive the token through `with_graceful_shutdown`; remaining work is bounded
650/// by the token's shared deadline and forcefully aborted on expiry.
651///
652/// # Errors
653///
654/// Fails like [`serve_listener_pair_with_shutdown_order`].
655pub async fn serve_listener_pair(
656    api_listener: TcpListener,
657    api_router: Router,
658    operations_listener: TcpListener,
659    operations_router: Router,
660    shutdown: ShutdownToken,
661) -> Result<(), ListenerPairError> {
662    serve_listener_pair_with_shutdown_order(
663        api_listener,
664        api_router,
665        operations_listener,
666        operations_router,
667        shutdown,
668        ShutdownOrder::Together,
669    )
670    .await
671}
672
673/// Runs an API/operations listener pair with an explicit shutdown order.
674///
675/// [`ShutdownOrder::Together`] is identical to [`serve_listener_pair`]. With
676/// [`ShutdownOrder::OpsOutlivesApi`], process shutdown first stops the API from
677/// accepting new work and lets its in-flight requests drain. The operations
678/// listener continues serving readiness and metrics until the API server exits;
679/// it is then gracefully stopped within the same process-wide drain deadline.
680/// If that deadline expires first, both listener tasks are aborted.
681///
682/// # Errors
683///
684/// Returns [`ListenerPairError::Listener`] when a server fails with an I/O error,
685/// [`ListenerPairError::UnexpectedExit`] when one stops before shutdown,
686/// [`ListenerPairError::Task`] when a listener task panics, and
687/// [`ListenerPairError::Drain`] when the drain deadline expires.
688pub async fn serve_listener_pair_with_shutdown_order(
689    api_listener: TcpListener,
690    api_router: Router,
691    operations_listener: TcpListener,
692    operations_router: Router,
693    shutdown: ShutdownToken,
694    shutdown_order: ShutdownOrder,
695) -> Result<(), ListenerPairError> {
696    let mut listeners = JoinSet::new();
697    let operations_shutdown = match shutdown_order {
698        ShutdownOrder::Together => shutdown.observed.clone(),
699        ShutdownOrder::OpsOutlivesApi => CancellationToken::new(),
700    };
701    spawn_listener(
702        &mut listeners,
703        ListenerKind::Api,
704        api_listener,
705        api_router,
706        shutdown.observed.clone(),
707    );
708    spawn_listener(
709        &mut listeners,
710        ListenerKind::Operations,
711        operations_listener,
712        operations_router,
713        operations_shutdown.clone(),
714    );
715
716    let mut first_error = None;
717    while !listeners.is_empty() && !shutdown.is_cancelled() {
718        tokio::select! {
719            _ = shutdown.cancelled() => {}
720            result = listeners.join_next() => {
721                stop_operations_after_api(
722                    result.as_ref(),
723                    shutdown_order,
724                    &operations_shutdown,
725                );
726                record_listener_result(result, &shutdown, &mut first_error);
727            }
728        }
729    }
730
731    if listeners.is_empty() {
732        return first_error.map_or(Ok(()), Err);
733    }
734
735    let remaining = shutdown.remaining();
736    let drained = timeout(remaining, async {
737        while let Some(result) = listeners.join_next().await {
738            stop_operations_after_api(Some(&result), shutdown_order, &operations_shutdown);
739            record_listener_result(Some(result), &shutdown, &mut first_error);
740        }
741    })
742    .await;
743
744    if drained.is_err() {
745        listeners.abort_all();
746        while listeners.join_next().await.is_some() {}
747        return Err(DrainTimeout::new(shutdown.drain_timeout).into());
748    }
749
750    first_error.map_or(Ok(()), Err)
751}
752
753fn spawn_listener(
754    listeners: &mut JoinSet<ListenerResult>,
755    kind: ListenerKind,
756    listener: TcpListener,
757    router: Router,
758    shutdown: CancellationToken,
759) {
760    listeners.spawn(async move {
761        axum::serve(listener, router)
762            .with_graceful_shutdown(async move { shutdown.cancelled().await })
763            .await
764            .map(|()| kind)
765            .map_err(|error| (kind, error))
766    });
767}
768
769fn stop_operations_after_api(
770    result: Option<&Result<ListenerResult, JoinError>>,
771    shutdown_order: ShutdownOrder,
772    operations_shutdown: &CancellationToken,
773) {
774    if shutdown_order != ShutdownOrder::OpsOutlivesApi {
775        return;
776    }
777
778    let listener = match result {
779        Some(Ok(Ok(kind) | Err((kind, _)))) => Some(*kind),
780        Some(Err(_)) | None => None,
781    };
782    if listener == Some(ListenerKind::Api) {
783        operations_shutdown.cancel();
784    }
785}
786
787fn record_listener_result(
788    result: Option<Result<ListenerResult, JoinError>>,
789    shutdown: &ShutdownToken,
790    first_error: &mut Option<ListenerPairError>,
791) {
792    let error = match result {
793        Some(Ok(Ok(kind))) if !shutdown.is_cancelled() => {
794            Some(ListenerPairError::UnexpectedExit(kind))
795        }
796        Some(Ok(Ok(_))) | None => None,
797        Some(Ok(Err((listener, source)))) => Some(ListenerPairError::Listener { listener, source }),
798        Some(Err(error)) => Some(ListenerPairError::Task(error)),
799    };
800
801    if let Some(error) = error {
802        tracing::error!(%error, "listener pair initiated process shutdown");
803        if first_error.is_none() {
804            *first_error = Some(error);
805        }
806        shutdown.trigger();
807    }
808}
809
810#[cfg(test)]
811mod tests;
812
813// Compiles the README's examples so they cannot drift from the API.
814#[doc = include_str!("../README.md")]
815#[cfg(doctest)]
816struct ReadmeDoctests;