1use 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
11pub trait Lifespan: Resources + Sized {
20 fn startup(ctx: LifespanContext) -> impl Future<Output = Result<Self>> + Send;
22
23 fn shutdown(self) -> impl Future<Output = Result<()>> + Send {
25 async move {
26 let _ = self;
27 Ok(())
28 }
29 }
30}
31
32pub struct LifespanContext {
37 _private: (),
38}
39
40impl LifespanContext {
41 pub fn new() -> Self {
43 Self { _private: () }
44 }
45
46 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#[derive(Clone, Debug)]
67pub struct ReadyContext {
68 addr: SocketAddr,
69}
70
71impl ReadyContext {
72 pub fn new(addr: SocketAddr) -> Self {
74 Self { addr }
75 }
76
77 pub fn addr(&self) -> SocketAddr {
79 self.addr
80 }
81}
82
83pub(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
97pub(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 cell.shutdown().await.unwrap();
173 }
174}