Skip to main content

tork_core/
lifespan.rs

1//! Application lifespan: typed startup and shutdown tied to a resource container.
2
3use std::future::Future;
4use std::net::SocketAddr;
5
6use crate::error::{Error, Result};
7use crate::resources::Resources;
8use crate::router::BoxFuture;
9use crate::state::StateMap;
10
11/// The startup and shutdown lifecycle of a resource container.
12///
13/// `startup` builds the container (acquiring pools, loading config, spawning
14/// workers); the value it returns has its resources registered for injection.
15/// `shutdown` releases those resources and is optional (it defaults to a no-op).
16///
17/// Implemented by `#[tork::lifespan]`. A lifespan type must also be a
18/// [`Resources`] container.
19pub trait Lifespan: Resources + Sized {
20    /// Builds the resource container.
21    fn startup(ctx: LifespanContext) -> impl Future<Output = Result<Self>> + Send;
22
23    /// Releases the container's resources. Defaults to a no-op.
24    fn shutdown(self) -> impl Future<Output = Result<()>> + Send {
25        async move {
26            let _ = self;
27            Ok(())
28        }
29    }
30}
31
32/// Context passed to a lifespan's `startup`.
33///
34/// Provides access to the process environment and is constructed by the
35/// framework when the application boots.
36pub struct LifespanContext {
37    _private: (),
38}
39
40impl LifespanContext {
41    /// Creates a lifespan context.
42    pub fn new() -> Self {
43        Self { _private: () }
44    }
45
46    /// Reads a required environment variable.
47    ///
48    /// # Errors
49    ///
50    /// Returns an error (code `MISSING_ENV`) if the variable is not set.
51    pub fn env(&self, key: &str) -> Result<String> {
52        std::env::var(key).map_err(|_| {
53            Error::internal(format!("required environment variable `{key}` is not set"))
54                .with_code("MISSING_ENV")
55        })
56    }
57}
58
59impl Default for LifespanContext {
60    fn default() -> Self {
61        Self::new()
62    }
63}
64
65/// Context passed to `on_ready` hooks, after the listener has bound.
66#[derive(Clone, Debug)]
67pub struct ReadyContext {
68    addr: SocketAddr,
69}
70
71impl ReadyContext {
72    /// Creates a ready context for a bound address.
73    pub fn new(addr: SocketAddr) -> Self {
74        Self { addr }
75    }
76
77    /// Returns the bound local address.
78    pub fn addr(&self) -> SocketAddr {
79        self.addr
80    }
81}
82
83/// Object-safe form of a [`Lifespan`], stored type-erased on the application.
84///
85/// `startup` runs the concrete startup, registers the produced container's
86/// resources, and stashes the container so `shutdown` can consume it by value.
87pub(crate) trait ErasedLifespan: Send + Sync {
88    fn startup<'a>(
89        &'a mut self,
90        ctx: LifespanContext,
91        registry: &'a mut StateMap,
92    ) -> BoxFuture<'a, Result<()>>;
93
94    fn shutdown(&mut self) -> BoxFuture<'_, Result<()>>;
95}
96
97/// Holds a lifespan and, after startup, the produced container.
98pub(crate) struct LifespanCell<L: Lifespan> {
99    container: Option<L>,
100}
101
102impl<L: Lifespan> LifespanCell<L> {
103    pub(crate) fn new() -> Self {
104        Self { container: None }
105    }
106}
107
108impl<L: Lifespan> ErasedLifespan for LifespanCell<L> {
109    fn startup<'a>(
110        &'a mut self,
111        ctx: LifespanContext,
112        registry: &'a mut StateMap,
113    ) -> BoxFuture<'a, Result<()>> {
114        Box::pin(async move {
115            let container = L::startup(ctx).await?;
116            container.register(registry);
117            self.container = Some(container);
118            Ok(())
119        })
120    }
121
122    fn shutdown(&mut self) -> BoxFuture<'_, Result<()>> {
123        Box::pin(async move {
124            match self.container.take() {
125                Some(container) => container.shutdown().await,
126                None => Ok(()),
127            }
128        })
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn missing_env_is_an_error() {
138        let ctx = LifespanContext::new();
139        let error = ctx.env("TORK_DEFINITELY_MISSING_VARIABLE_XYZ").unwrap_err();
140        assert_eq!(error.code(), "MISSING_ENV");
141    }
142
143    #[derive(Clone)]
144    struct Probe {
145        value: i64,
146    }
147
148    impl Resources for Probe {
149        fn register(&self, registry: &mut StateMap) {
150            registry.insert(self.value);
151        }
152    }
153
154    impl Lifespan for Probe {
155        async fn startup(_ctx: LifespanContext) -> Result<Self> {
156            Ok(Probe { value: 7 })
157        }
158    }
159
160    #[tokio::test]
161    async fn cell_registers_resources_and_shuts_down() {
162        let mut cell = LifespanCell::<Probe>::new();
163        let mut registry = StateMap::new();
164
165        cell.startup(LifespanContext::new(), &mut registry)
166            .await
167            .unwrap();
168        assert_eq!(registry.get::<i64>().map(|value| *value), Some(7));
169
170        cell.shutdown().await.unwrap();
171        // A second shutdown is a no-op (the container was already taken).
172        cell.shutdown().await.unwrap();
173    }
174}