Skip to main content

rust_dix/
collection.rs

1use crate::entry::{AsyncServiceFactory, ServiceDescriptor, ServiceFactory, ServiceLifetime};
2use crate::registration::ServiceRegistration;
3use std::any::{Any, TypeId};
4use std::future::Future;
5use std::pin::Pin;
6use std::sync::Arc;
7
8pub struct ServiceCollection {
9    descriptors: Vec<ServiceDescriptor>,
10}
11
12impl ServiceCollection {
13    pub fn new() -> Self {
14        Self {
15            descriptors: Vec::new(),
16        }
17    }
18
19    pub fn singleton<T: ?Sized + Send + Sync + 'static>(
20        mut self,
21        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
22    ) -> Self {
23        self.push(ServiceLifetime::Singleton, None, f);
24        self
25    }
26
27    pub fn transient<T: ?Sized + Send + Sync + 'static>(
28        mut self,
29        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
30    ) -> Self {
31        self.push(ServiceLifetime::Transient, None, f);
32        self
33    }
34
35    pub fn scoped<T: ?Sized + Send + Sync + 'static>(
36        mut self,
37        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
38    ) -> Self {
39        self.push(ServiceLifetime::Scoped, None, f);
40        self
41    }
42
43    /// Register a keyed singleton service.
44    ///
45    /// The key distinguishes multiple implementations of the same trait.
46    /// Use `get_keyed::<T>(key)` or `try_get_keyed::<T>(key)` to resolve.
47    pub fn keyed_singleton<T: ?Sized + Send + Sync + 'static>(
48        mut self,
49        k: impl Into<String>,
50        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
51    ) -> Self {
52        self.push(ServiceLifetime::Singleton, Some(k.into()), f);
53        self
54    }
55
56    /// Register a keyed transient service (new instance each resolution).
57    pub fn keyed_transient<T: ?Sized + Send + Sync + 'static>(
58        mut self,
59        k: impl Into<String>,
60        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
61    ) -> Self {
62        self.push(ServiceLifetime::Transient, Some(k.into()), f);
63        self
64    }
65
66    /// Register a keyed scoped service (shared within a scope).
67    pub fn keyed_scoped<T: ?Sized + Send + Sync + 'static>(
68        mut self,
69        k: impl Into<String>,
70        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
71    ) -> Self {
72        self.push(ServiceLifetime::Scoped, Some(k.into()), f);
73        self
74    }
75
76    /// Register an async singleton service (initialized once during `build_async`).
77    pub fn async_singleton<T: ?Sized + Send + Sync + 'static>(
78        mut self,
79        f: impl Fn(Arc<crate::provider::ServiceProvider>) -> Pin<Box<dyn Future<Output = Arc<T>> + Send>>
80            + Send
81            + Sync
82            + 'static,
83    ) -> Self {
84        self.push_async(ServiceLifetime::Singleton, None, f);
85        self
86    }
87
88    /// Register an async transient service (new async instance each resolution).
89    pub fn async_transient<T: ?Sized + Send + Sync + 'static>(
90        mut self,
91        f: impl Fn(Arc<crate::provider::ServiceProvider>) -> Pin<Box<dyn Future<Output = Arc<T>> + Send>>
92            + Send
93            + Sync
94            + 'static,
95    ) -> Self {
96        self.push_async(ServiceLifetime::Transient, None, f);
97        self
98    }
99
100    /// Register an async scoped service (shared within a scope).
101    pub fn async_scoped<T: ?Sized + Send + Sync + 'static>(
102        mut self,
103        f: impl Fn(Arc<crate::provider::ServiceProvider>) -> Pin<Box<dyn Future<Output = Arc<T>> + Send>>
104            + Send
105            + Sync
106            + 'static,
107    ) -> Self {
108        self.push_async(ServiceLifetime::Scoped, None, f);
109        self
110    }
111
112    /// Register an async keyed singleton service.
113    pub fn async_keyed_singleton<T: ?Sized + Send + Sync + 'static>(
114        mut self,
115        k: impl Into<String>,
116        f: impl Fn(Arc<crate::provider::ServiceProvider>) -> Pin<Box<dyn Future<Output = Arc<T>> + Send>>
117            + Send
118            + Sync
119            + 'static,
120    ) -> Self {
121        self.push_async(ServiceLifetime::Singleton, Some(k.into()), f);
122        self
123    }
124
125    /// Register an async keyed transient service.
126    pub fn async_keyed_transient<T: ?Sized + Send + Sync + 'static>(
127        mut self,
128        k: impl Into<String>,
129        f: impl Fn(Arc<crate::provider::ServiceProvider>) -> Pin<Box<dyn Future<Output = Arc<T>> + Send>>
130            + Send
131            + Sync
132            + 'static,
133    ) -> Self {
134        self.push_async(ServiceLifetime::Transient, Some(k.into()), f);
135        self
136    }
137
138    /// Register an async keyed scoped service.
139    pub fn async_keyed_scoped<T: ?Sized + Send + Sync + 'static>(
140        mut self,
141        k: impl Into<String>,
142        f: impl Fn(Arc<crate::provider::ServiceProvider>) -> Pin<Box<dyn Future<Output = Arc<T>> + Send>>
143            + Send
144            + Sync
145            + 'static,
146    ) -> Self {
147        self.push_async(ServiceLifetime::Scoped, Some(k.into()), f);
148        self
149    }
150
151    pub fn try_add<T: ?Sized + Send + Sync + 'static>(
152        self,
153        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
154    ) -> Self {
155        self.try_add_with_lifetime(ServiceLifetime::Singleton, f)
156    }
157
158    /// Register a singleton only if no unkeyed registration for `T` exists yet.
159    pub fn try_add_singleton<T: ?Sized + Send + Sync + 'static>(
160        self,
161        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
162    ) -> Self {
163        self.try_add_with_lifetime(ServiceLifetime::Singleton, f)
164    }
165
166    /// Register a scoped service only if no unkeyed registration for `T` exists yet.
167    pub fn try_add_scoped<T: ?Sized + Send + Sync + 'static>(
168        self,
169        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
170    ) -> Self {
171        self.try_add_with_lifetime(ServiceLifetime::Scoped, f)
172    }
173
174    /// Register a transient service only if no unkeyed registration for `T` exists yet.
175    pub fn try_add_transient<T: ?Sized + Send + Sync + 'static>(
176        self,
177        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
178    ) -> Self {
179        self.try_add_with_lifetime(ServiceLifetime::Transient, f)
180    }
181
182    fn try_add_with_lifetime<T: ?Sized + Send + Sync + 'static>(
183        mut self,
184        lt: ServiceLifetime,
185        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
186    ) -> Self {
187        let tid = TypeId::of::<T>();
188        if self
189            .descriptors
190            .iter()
191            .any(|d| d.type_id == tid && d.key.is_none())
192        {
193            return self;
194        }
195        self.push(lt, None, f);
196        self
197    }
198
199    pub fn add<T: ?Sized + Send + Sync + 'static>(
200        mut self,
201        lt: ServiceLifetime,
202        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
203    ) -> Self {
204        self.push(lt, None, f);
205        self
206    }
207
208    pub fn instance<T: Send + Sync + 'static>(mut self, v: Arc<T>) -> Self {
209        let ff: ServiceFactory = Arc::new(move |_| Arc::new(v.clone()));
210        self.descriptors.push(ServiceDescriptor {
211            type_id: TypeId::of::<T>(),
212            type_name: std::any::type_name::<T>(),
213            key: None,
214            factory: ff,
215            async_factory: None,
216            lifetime: ServiceLifetime::Singleton,
217        });
218        self
219    }
220
221    pub fn singleton_value<T: Send + Sync + 'static>(self, v: T) -> Self {
222        self.instance(Arc::new(v))
223    }
224
225    pub fn build(self) -> Result<Arc<crate::provider::ServiceProvider>, crate::error::RdiError> {
226        let mut s = crate::entry::ServiceStore::new();
227        for (n, d) in self.descriptors.into_iter().enumerate() {
228            let e = crate::entry::ServiceEntry {
229                cache_key: n,
230                key: d.key,
231                type_name: d.type_name,
232                factory: d.factory,
233                async_factory: d.async_factory,
234                lifetime: d.lifetime,
235            };
236            s.entry(d.type_id).or_default().push(e);
237        }
238        crate::provider::ServiceProvider::new(s)
239    }
240
241    /// Build the provider, executing async factories for singleton services.
242    ///
243    /// Async singleton factories are run during build and their results are
244    /// cached as singleton instances. Async transient/scoped factories are
245    /// stored for later resolution.
246    ///
247    /// Returns `Arc<ServiceProvider>` so that `provider_arc()` works for
248    /// async resolution via `get_async` / `get_keyed_async`.
249    ///
250    /// Use this instead of [`build`](Self::build) when any service requires
251    /// async initialization (e.g., database connections, remote config).
252    pub async fn build_async(
253        self,
254    ) -> Result<Arc<crate::provider::ServiceProvider>, crate::error::RdiError> {
255        let mut s = crate::entry::ServiceStore::new();
256        for (n, d) in self.descriptors.into_iter().enumerate() {
257            let e = crate::entry::ServiceEntry {
258                cache_key: n,
259                key: d.key,
260                type_name: d.type_name,
261                factory: d.factory,
262                async_factory: d.async_factory,
263                lifetime: d.lifetime,
264            };
265            s.entry(d.type_id).or_default().push(e);
266        }
267        crate::provider::ServiceProvider::new_async(s).await
268    }
269
270    /// Build a `ServiceCollection` from all `#[rust_dix::inject]` annotations
271    /// in the current binary (across all crates).
272    #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
273    pub fn from_injected() -> Self {
274        let mut descriptors = Vec::new();
275        for reg in inventory::iter::<ServiceRegistration> {
276            let factory: ServiceFactory = Arc::new(move |r| (reg.factory)(r));
277            descriptors.push(ServiceDescriptor {
278                type_id: reg.type_id,
279                type_name: (reg.type_name_fn)(),
280                key: None,
281                factory,
282                async_factory: None,
283                lifetime: reg.lifetime,
284            });
285        }
286        Self { descriptors }
287    }
288
289    /// Register a new singleton service.
290    fn push<T: ?Sized + Send + Sync + 'static>(
291        &mut self,
292        lt: ServiceLifetime,
293        key: Option<String>,
294        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
295    ) {
296        let sf: ServiceFactory = Arc::new(move |r| {
297            let val: Arc<T> = (f)(r);
298            Arc::new(val) as Arc<dyn Any + Send + Sync>
299        });
300        self.descriptors.push(ServiceDescriptor {
301            type_id: TypeId::of::<T>(),
302            type_name: std::any::type_name::<T>(),
303            key,
304            factory: sf,
305            async_factory: None,
306            lifetime: lt,
307        });
308    }
309
310    /// Register an async service.
311    fn push_async<T: ?Sized + Send + Sync + 'static>(
312        &mut self,
313        lt: ServiceLifetime,
314        key: Option<String>,
315        f: impl Fn(Arc<crate::provider::ServiceProvider>) -> Pin<Box<dyn Future<Output = Arc<T>> + Send>>
316            + Send
317            + Sync
318            + 'static,
319    ) {
320        let af: AsyncServiceFactory = Arc::new(move |sp: Arc<crate::provider::ServiceProvider>| {
321            let fut = (f)(sp);
322            Box::pin(async move {
323                let val: Arc<T> = fut.await;
324                Arc::new(val) as Arc<dyn Any + Send + Sync>
325            }) as Pin<Box<dyn Future<Output = Arc<dyn Any + Send + Sync>> + Send>>
326        });
327        // Sync factory: typed error if somehow invoked; `build()` rejects async
328        // registrations before this path is reached.
329        let type_name = std::any::type_name::<T>();
330        let sf: ServiceFactory = Arc::new(move |_| {
331            std::panic::panic_any(crate::error::RdiError::AsyncServiceRequiresBuildAsync {
332                service: type_name.to_string(),
333            })
334        });
335        self.descriptors.push(ServiceDescriptor {
336            type_id: TypeId::of::<T>(),
337            type_name: std::any::type_name::<T>(),
338            key,
339            factory: sf,
340            async_factory: Some(af),
341            lifetime: lt,
342        });
343    }
344}
345impl Default for ServiceCollection {
346    fn default() -> Self {
347        Self::new()
348    }
349}
350
351#[cfg(test)]
352mod tests {
353    use super::*;
354    #[derive(Debug, PartialEq)]
355    struct G {
356        n: String,
357    }
358    #[derive(Debug, PartialEq)]
359    struct C {
360        v: i32,
361    }
362    #[test]
363    fn empty() {
364        let p = ServiceCollection::new().build().unwrap();
365        assert!(p.get_optional::<G>().is_none());
366    }
367    #[test]
368    fn singleton() {
369        let p = ServiceCollection::new()
370            .singleton(|_| Arc::new(G { n: "Hi".into() }))
371            .build()
372            .unwrap();
373        assert_eq!(p.get::<G>().unwrap().n, "Hi");
374    }
375    #[test]
376    fn singleton_caches() {
377        use std::sync::atomic::{AtomicUsize, Ordering};
378        static CNT: AtomicUsize = AtomicUsize::new(0);
379        let p = ServiceCollection::new()
380            .singleton(|_| {
381                CNT.fetch_add(1, Ordering::SeqCst);
382                Arc::new(C { v: 42 })
383            })
384            .build()
385            .unwrap();
386        let _ = p.get::<C>();
387        let _ = p.get::<C>();
388        assert_eq!(CNT.load(Ordering::SeqCst), 1);
389    }
390    #[test]
391    fn transient_not_cached() {
392        use std::sync::atomic::{AtomicUsize, Ordering};
393        static CNT: AtomicUsize = AtomicUsize::new(0);
394        let p = ServiceCollection::new()
395            .transient(|_| {
396                CNT.fetch_add(1, Ordering::SeqCst);
397                Arc::new(C { v: 1 })
398            })
399            .build()
400            .unwrap();
401        let _ = p.get::<C>();
402        let _ = p.get::<C>();
403        assert_eq!(CNT.load(Ordering::SeqCst), 2);
404    }
405    #[test]
406    fn instance() {
407        let g = Arc::new(G { n: "Inst".into() });
408        let p = ServiceCollection::new()
409            .instance(g.clone())
410            .build()
411            .unwrap();
412        assert!(Arc::ptr_eq(&g, &p.get::<G>().unwrap()));
413    }
414    #[test]
415    fn keyed_svc() {
416        let p = ServiceCollection::new()
417            .keyed_singleton("a", |_| Arc::new(G { n: "A".into() }))
418            .keyed_singleton("b", |_| Arc::new(G { n: "B".into() }))
419            .build()
420            .unwrap();
421        assert_eq!(p.get_keyed::<G>("a").unwrap().n, "A");
422        assert_eq!(p.get_keyed::<G>("b").unwrap().n, "B");
423    }
424    #[test]
425    fn get_all() {
426        let p = ServiceCollection::new()
427            .keyed_singleton("x", |_| Arc::new(C { v: 1 }))
428            .keyed_singleton("y", |_| Arc::new(C { v: 2 }))
429            .build()
430            .unwrap();
431        assert_eq!(p.get_all::<C>().len(), 2);
432    }
433    #[test]
434    fn try_add_skips() {
435        let p = ServiceCollection::new()
436            .singleton(|_| Arc::new(G { n: "First".into() }))
437            .try_add(|_| Arc::new(G { n: "Second".into() }))
438            .build()
439            .unwrap();
440        assert_eq!(p.get::<G>().unwrap().n, "First");
441    }
442}