Skip to main content

allframe_core/di/
lazy.rs

1//! Lazy initialization for DI containers
2//!
3//! Provides `LazyProvider<T>` for deferred initialization and `LazyContainer`
4//! for bulk warm-up of lazy bindings.
5
6use std::future::Future;
7use std::pin::Pin;
8use std::sync::Arc;
9
10use tokio::sync::OnceCell;
11
12use super::DependencyError;
13
14/// Type alias for an async factory that produces a value of type T.
15type AsyncFactory<T> =
16    dyn Fn() -> Pin<Box<dyn Future<Output = Result<T, DependencyError>> + Send>> + Send + Sync;
17
18/// A provider that lazily initializes a value on first access.
19///
20/// Thread-safe: concurrent calls to `get()` will only trigger one initialization;
21/// other callers wait for it to complete.
22pub 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    /// Create a new lazy provider with the given async factory.
38    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    /// Get the value, initializing it on first call.
50    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/// Trait for type-erased lazy initialization (used by LazyContainer).
60#[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
79/// Container that holds multiple lazy providers and can warm them up concurrently.
80pub 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    /// Create a new empty container.
92    pub fn new() -> Self {
93        Self {
94            entries: Vec::new(),
95        }
96    }
97
98    /// Register a lazy provider for type T with a name and factory.
99    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    /// Initialize all lazy bindings concurrently.
113    ///
114    /// Spawns each initialization as a separate tokio task and waits for all
115    /// to complete. If any initialization fails, the error is returned after
116    /// all tasks finish.
117    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}