siderite_core/
lifespan.rs

1//! Application lifespan: startup/shutdown hooks and scoped resources.
2//!
3//! Hooks run in registration order at startup and in **reverse** order at
4//! shutdown; hooks of mounted child apps run after the parent's at startup
5//! and before them at shutdown. A failing startup hook aborts startup with
6//! [`ServerError::Lifespan`] after reverse-order cleanup. At shutdown every
7//! hook runs even if one fails; the first error is reported.
8//!
9//! [`App::lifespan_resource`] is the equivalent of FastAPI's lifespan
10//! context manager: a value created at startup, injected with
11//! [`Resource<T>`], and cleaned up at shutdown.
12
13use crate::app::{App, BoxFuture, LifespanHook};
14use crate::error::{ApiError, ServerError};
15use crate::extract::FromRequestParts;
16use futures_util::FutureExt;
17use std::panic::AssertUnwindSafe;
18use std::time::Duration;
19use tokio::time::Instant;
20
21mod managed;
22pub use managed::ManagedLifespan;
23
24/// Default budget for request draining and lifespan teardown.
25pub const DEFAULT_SHUTDOWN_BUDGET: Duration = Duration::from_secs(30);
26
27use http::StatusCode;
28use http::request::Parts;
29use std::future::Future;
30use std::ops::Deref;
31use std::sync::{Arc, Mutex, PoisonError, RwLock};
32
33/// The startup and shutdown hooks collected from an app tree.
34pub struct Lifespan {
35    pub(crate) startup: Vec<LifespanHook>,
36    pub(crate) shutdown: Vec<LifespanHook>,
37    pub(crate) shutdown_budget: Duration,
38    pub(crate) background: Vec<crate::background::TaskManager>,
39    pub(crate) transports: Vec<crate::server::tasks::TaskOwner>,
40    pub(crate) server: Option<crate::server::ServerOwners>,
41}
42
43impl Default for Lifespan {
44    fn default() -> Self {
45        Self {
46            startup: Vec::new(),
47            shutdown: Vec::new(),
48            shutdown_budget: DEFAULT_SHUTDOWN_BUDGET,
49            background: Vec::new(),
50            transports: Vec::new(),
51            server: None,
52        }
53    }
54}
55
56impl std::fmt::Debug for Lifespan {
57    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58        f.debug_struct("Lifespan")
59            .field("startup_hooks", &self.startup.len())
60            .field("shutdown_hooks", &self.shutdown.len())
61            .finish()
62    }
63}
64
65impl Lifespan {
66    /// Start a cleanup supervisor without waiting for initialization.
67    ///
68    /// Args:
69    ///     self: The hooks whose ownership transfers to the supervisor.
70    ///
71    /// Returns:
72    ///     An owner used to await readiness and request shutdown.
73    ///
74    /// # Errors
75    /// Missing Tokio runtime. Use this API for cancellation-safe ownership;
76    /// initializers must clean their own work until they return successfully.
77    pub fn supervise(self) -> Result<ManagedLifespan, ServerError> {
78        ManagedLifespan::new(self)
79    }
80
81    /// Run startup and unwind successful initialization on error or panic.
82    ///
83    /// Args:
84    ///     self: Lifespan retained by the caller.
85    ///
86    /// Returns:
87    ///     Success, or the original startup error after cleanup.
88    ///
89    /// # Errors
90    /// Hook or cleanup failure. Direct callers must not cancel this method
91    /// without subsequently calling shutdown; prefer [`Self::supervise`].
92    pub async fn startup(&mut self) -> Result<(), ServerError> {
93        let result = self.start_hooks().await;
94        if result.is_err() {
95            let _ = self.shutdown().await;
96        }
97        result
98    }
99
100    async fn start_hooks(&mut self) -> Result<(), ServerError> {
101        for hook in std::mem::take(&mut self.startup) {
102            call_hook(hook, "startup")
103                .await
104                .map_err(ServerError::Lifespan)?;
105        }
106        if let Some(server) = &self.server {
107            server.readiness.ready();
108        }
109        Ok(())
110    }
111
112    /// Run every shutdown hook in reverse order within the configured budget.
113    ///
114    /// Args:
115    ///     self: Lifespan whose remaining teardown hooks are consumed.
116    ///
117    /// Returns:
118    ///     Success, or the first failure after all cleanup attempts.
119    ///
120    /// # Errors
121    /// Hook failure/panic or deadline expiry. Use a supervisor to shield this
122    /// operation from caller cancellation. Hooks must cooperate with polling.
123    pub async fn shutdown(&mut self) -> Result<(), ServerError> {
124        self.shutdown_at(Instant::now() + self.shutdown_budget)
125            .await
126    }
127
128    async fn shutdown_at(&mut self, deadline: Instant) -> Result<(), ServerError> {
129        let mut first_error = None;
130        let server = self.server.take();
131        if let Some(server) = &server
132            && !server.stop_at(deadline).await
133        {
134            first_error.get_or_insert(ServerError::ShutdownTimeout);
135        }
136        for owner in std::mem::take(&mut self.transports) {
137            if !owner.stop_at(deadline).await {
138                first_error.get_or_insert(ServerError::ShutdownTimeout);
139            }
140        }
141        for manager in std::mem::take(&mut self.background).into_iter().rev() {
142            if let Err(error) = manager.shutdown_at(deadline).await {
143                tracing::error!(%error, "background shutdown failed");
144                first_error.get_or_insert(error);
145            }
146        }
147        while let Some(hook) = self.shutdown.pop() {
148            let result = tokio::time::timeout_at(deadline, call_hook(hook, "shutdown")).await;
149            let error = match result {
150                Ok(Ok(())) => continue,
151                Ok(Err(error)) => ServerError::Lifespan(error),
152                Err(_) => ServerError::ShutdownTimeout,
153            };
154            tracing::error!(%error, "shutdown hook failed");
155            first_error.get_or_insert(error);
156        }
157        if let Some(server) = server {
158            server.readiness.stopped();
159        }
160        first_error.map_or(Ok(()), Err)
161    }
162}
163
164async fn call_hook(hook: LifespanHook, phase: &str) -> Result<(), ApiError> {
165    match AssertUnwindSafe(async move { hook().await })
166        .catch_unwind()
167        .await
168    {
169        Ok(result) => result,
170        Err(_) => {
171            tracing::error!(phase, "lifespan hook panicked");
172            Err(ApiError::internal(format!("{phase} hook panicked")))
173        }
174    }
175}
176
177/// Handler argument giving access to a value created by
178/// [`App::lifespan_resource`]. Responds `503` until startup has completed
179/// (and again after shutdown).
180#[derive(Debug)]
181pub struct Resource<T>(pub Arc<T>);
182
183impl<T> Clone for Resource<T> {
184    fn clone(&self) -> Self {
185        Self(Arc::clone(&self.0))
186    }
187}
188
189impl<T> Deref for Resource<T> {
190    type Target = T;
191
192    fn deref(&self) -> &T {
193        &self.0
194    }
195}
196
197/// Holds the teardown future between startup and shutdown.
198type TeardownCell = Mutex<Option<BoxFuture<Result<(), ApiError>>>>;
199
200/// Shared cell installed as a request extension; filled at startup.
201struct Slot<T>(Arc<RwLock<Option<Arc<T>>>>);
202
203impl<T> Clone for Slot<T> {
204    fn clone(&self) -> Self {
205        Self(Arc::clone(&self.0))
206    }
207}
208
209impl<T: Send + Sync + 'static> FromRequestParts for Resource<T> {
210    async fn from_request_parts(parts: &mut Parts) -> Result<Self, ApiError> {
211        let slot = parts.extensions.get::<Slot<T>>().ok_or_else(|| {
212            ApiError::internal(format!(
213                "no lifespan resource of type `{}` registered",
214                std::any::type_name::<T>()
215            ))
216        })?;
217        let value = slot
218            .0
219            .read()
220            .unwrap_or_else(PoisonError::into_inner)
221            .clone();
222        value.map(Resource).ok_or_else(|| {
223            ApiError::new(
224                StatusCode::SERVICE_UNAVAILABLE,
225                "The service is not ready yet.",
226            )
227        })
228    }
229}
230
231impl App {
232    /// Set the total graceful-shutdown budget, including resource teardown.
233    ///
234    /// Args:
235    ///     budget: Positive, representable duration (default: 30 seconds).
236    ///
237    /// Returns:
238    ///     The configured application; invalid budgets fail validation.
239    #[must_use]
240    pub fn shutdown_timeout(mut self, budget: Duration) -> Self {
241        if budget.is_zero() || std::time::Instant::now().checked_add(budget).is_none() {
242            self.config_errors.push("invalid shutdown timeout".into());
243        } else {
244            self.shutdown_budget = budget;
245        }
246        self
247    }
248
249    /// Run `hook` at startup (before the socket is bound).
250    #[must_use]
251    pub fn on_startup<F, Fut>(mut self, hook: F) -> Self
252    where
253        F: FnOnce() -> Fut + Send + 'static,
254        Fut: Future<Output = Result<(), ApiError>> + Send + 'static,
255    {
256        self.startup_hooks.push(boxed(hook));
257        self
258    }
259
260    /// Run `hook` at shutdown (hooks run in reverse registration order).
261    #[must_use]
262    pub fn on_shutdown<F, Fut>(mut self, hook: F) -> Self
263    where
264        F: FnOnce() -> Fut + Send + 'static,
265        Fut: Future<Output = Result<(), ApiError>> + Send + 'static,
266    {
267        self.shutdown_hooks.push(boxed(hook));
268        self
269    }
270
271    /// Create a resource at startup and tear it down at shutdown.
272    ///
273    /// `init` returns the value and a *teardown future* (which may capture
274    /// clones of whatever it needs to close). Handlers receive the value via
275    /// [`Resource<T>`]. Resources are torn down in reverse registration order.
276    #[must_use]
277    pub fn lifespan_resource<T, F, Fut, C>(mut self, init: F) -> Self
278    where
279        T: Send + Sync + 'static,
280        F: FnOnce() -> Fut + Send + 'static,
281        Fut: Future<Output = Result<(T, C), ApiError>> + Send + 'static,
282        C: Future<Output = Result<(), ApiError>> + Send + 'static,
283    {
284        let slot: Slot<T> = Slot(Arc::default());
285        let teardown: Arc<TeardownCell> = Arc::default();
286        let (start_slot, start_teardown) = (slot.clone(), Arc::clone(&teardown));
287        self.startup_hooks.push(boxed(move || async move {
288            let (value, close) = init().await?;
289            *start_slot.0.write().unwrap_or_else(PoisonError::into_inner) = Some(Arc::new(value));
290            let close: BoxFuture<Result<(), ApiError>> = Box::pin(close);
291            *start_teardown
292                .lock()
293                .unwrap_or_else(PoisonError::into_inner) = Some(close);
294            Ok(())
295        }));
296        let stop_slot = slot.clone();
297        self.shutdown_hooks.push(boxed(move || async move {
298            stop_slot
299                .0
300                .write()
301                .unwrap_or_else(PoisonError::into_inner)
302                .take();
303            let close = teardown
304                .lock()
305                .unwrap_or_else(PoisonError::into_inner)
306                .take();
307            match close {
308                Some(close) => close.await,
309                None => Ok(()),
310            }
311        }));
312        self.router_layers
313            .push(Box::new(move |router| router.layer(axum::Extension(slot))));
314        self
315    }
316}
317
318fn boxed<F, Fut>(hook: F) -> LifespanHook
319where
320    F: FnOnce() -> Fut + Send + 'static,
321    Fut: Future<Output = Result<(), ApiError>> + Send + 'static,
322{
323    Box::new(move || Box::pin(hook()))
324}
siderite_core/lifespan.rs

siderite_core/
lifespan.rs
