1use std::future::Future;
7use std::pin::Pin;
8use std::sync::Arc;
9
10use tokio::sync::OnceCell;
11
12use super::DependencyError;
13
14type AsyncFactory<T> =
16 dyn Fn() -> Pin<Box<dyn Future<Output = Result<T, DependencyError>> + Send>> + Send + Sync;
17
18pub struct LazyProvider<T: Clone + Send + Sync + 'static> {
23 cell: Arc<OnceCell<T>>,
24 factory: Arc<AsyncFactory<T>>,
25}
26
27impl<T: Clone + Send + Sync + 'static> Clone for LazyProvider<T> {
28 fn clone(&self) -> Self {
29 Self {
30 cell: Arc::clone(&self.cell),
31 factory: Arc::clone(&self.factory),
32 }
33 }
34}
35
36impl<T: Clone + Send + Sync + 'static> LazyProvider<T> {
37 pub fn new<F, Fut>(factory: F) -> Self
39 where
40 F: Fn() -> Fut + Send + Sync + 'static,
41 Fut: Future<Output = Result<T, DependencyError>> + Send + 'static,
42 {
43 Self {
44 cell: Arc::new(OnceCell::new()),
45 factory: Arc::new(move || Box::pin(factory())),
46 }
47 }
48
49 pub async fn get(&self) -> Result<T, DependencyError> {
51 let factory = &self.factory;
52 self.cell
53 .get_or_try_init(|| factory())
54 .await
55 .cloned()
56 }
57}
58
59#[async_trait::async_trait]
61trait LazyInit: Send + Sync {
62 async fn init(&self) -> Result<(), DependencyError>;
63}
64
65struct LazyEntry<T: Clone + Send + Sync + 'static> {
66 #[allow(dead_code)]
67 name: String,
68 provider: LazyProvider<T>,
69}
70
71#[async_trait::async_trait]
72impl<T: Clone + Send + Sync + 'static> LazyInit for LazyEntry<T> {
73 async fn init(&self) -> Result<(), DependencyError> {
74 self.provider.get().await?;
75 Ok(())
76 }
77}
78
79pub struct LazyContainer {
81 entries: Vec<Arc<dyn LazyInit>>,
82}
83
84impl Default for LazyContainer {
85 fn default() -> Self {
86 Self::new()
87 }
88}
89
90impl LazyContainer {
91 pub fn new() -> Self {
93 Self {
94 entries: Vec::new(),
95 }
96 }
97
98 pub fn register_lazy<T, F, Fut>(&mut self, name: &str, factory: F)
100 where
101 T: Clone + Send + Sync + 'static,
102 F: Fn() -> Fut + Send + Sync + 'static,
103 Fut: Future<Output = Result<T, DependencyError>> + Send + 'static,
104 {
105 let provider = LazyProvider::new(factory);
106 self.entries.push(Arc::new(LazyEntry {
107 name: name.to_string(),
108 provider,
109 }));
110 }
111
112 pub async fn warm_up(&self) -> Result<(), DependencyError> {
118 let mut handles = Vec::with_capacity(self.entries.len());
119 for entry in &self.entries {
120 let entry = Arc::clone(entry);
121 handles.push(tokio::spawn(async move { entry.init().await }));
122 }
123 for handle in handles {
124 handle
125 .await
126 .map_err(|e| DependencyError::InitializationFailed {
127 name: "warm_up".to_string(),
128 source: Box::new(e),
129 })??;
130 }
131 Ok(())
132 }
133}