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        mut self,
153        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
154    ) -> Self {
155        let tid = TypeId::of::<T>();
156        if self
157            .descriptors
158            .iter()
159            .any(|d| d.type_id == tid && d.key.is_none())
160        {
161            return self;
162        }
163        self.push(ServiceLifetime::Singleton, None, f);
164        self
165    }
166
167    pub fn add<T: ?Sized + Send + Sync + 'static>(
168        mut self,
169        lt: ServiceLifetime,
170        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
171    ) -> Self {
172        self.push(lt, None, f);
173        self
174    }
175
176    pub fn instance<T: Send + Sync + 'static>(mut self, v: Arc<T>) -> Self {
177        let ff: ServiceFactory = Arc::new(move |_| Arc::new(v.clone()));
178        self.descriptors.push(ServiceDescriptor {
179            type_id: TypeId::of::<T>(),
180            type_name: std::any::type_name::<T>(),
181            key: None,
182            factory: ff,
183            async_factory: None,
184            lifetime: ServiceLifetime::Singleton,
185        });
186        self
187    }
188
189    pub fn singleton_value<T: Send + Sync + 'static>(self, v: T) -> Self {
190        self.instance(Arc::new(v))
191    }
192
193    pub fn build(self) -> Result<Arc<crate::provider::ServiceProvider>, crate::error::RdiError> {
194        let mut s = crate::entry::ServiceStore::new();
195        for (n, d) in self.descriptors.into_iter().enumerate() {
196            let e = crate::entry::ServiceEntry {
197                cache_key: n,
198                key: d.key,
199                type_name: d.type_name,
200                factory: d.factory,
201                async_factory: d.async_factory,
202                lifetime: d.lifetime,
203            };
204            s.entry(d.type_id).or_default().push(e);
205        }
206        crate::provider::ServiceProvider::new(s)
207    }
208
209    /// Build the provider, executing async factories for singleton services.
210    ///
211    /// Async singleton factories are run during build and their results are
212    /// cached as singleton instances. Async transient/scoped factories are
213    /// stored for later resolution.
214    ///
215    /// Returns `Arc<ServiceProvider>` so that `provider_arc()` works for
216    /// async resolution via `get_async` / `get_keyed_async`.
217    ///
218    /// Use this instead of [`build`](Self::build) when any service requires
219    /// async initialization (e.g., database connections, remote config).
220    pub async fn build_async(
221        self,
222    ) -> Result<Arc<crate::provider::ServiceProvider>, crate::error::RdiError> {
223        let mut s = crate::entry::ServiceStore::new();
224        for (n, d) in self.descriptors.into_iter().enumerate() {
225            let e = crate::entry::ServiceEntry {
226                cache_key: n,
227                key: d.key,
228                type_name: d.type_name,
229                factory: d.factory,
230                async_factory: d.async_factory,
231                lifetime: d.lifetime,
232            };
233            s.entry(d.type_id).or_default().push(e);
234        }
235        crate::provider::ServiceProvider::new_async(s).await
236    }
237
238    /// Build a `ServiceCollection` from all `#[rust_dix::inject]` annotations
239    /// in the current binary (across all crates).
240    #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
241    pub fn from_injected() -> Self {
242        let mut descriptors = Vec::new();
243        for reg in inventory::iter::<ServiceRegistration> {
244            let factory: ServiceFactory = Arc::new(move |r| (reg.factory)(r));
245            descriptors.push(ServiceDescriptor {
246                type_id: reg.type_id,
247                type_name: (reg.type_name_fn)(),
248                key: None,
249                factory,
250                async_factory: None,
251                lifetime: reg.lifetime,
252            });
253        }
254        Self { descriptors }
255    }
256
257    /// Register a new singleton service.
258    fn push<T: ?Sized + Send + Sync + 'static>(
259        &mut self,
260        lt: ServiceLifetime,
261        key: Option<String>,
262        f: impl Fn(&dyn crate::entry::IServiceResolver) -> Arc<T> + Send + Sync + 'static,
263    ) {
264        let sf: ServiceFactory = Arc::new(move |r| {
265            let val: Arc<T> = (f)(r);
266            Arc::new(val) as Arc<dyn Any + Send + Sync>
267        });
268        self.descriptors.push(ServiceDescriptor {
269            type_id: TypeId::of::<T>(),
270            type_name: std::any::type_name::<T>(),
271            key,
272            factory: sf,
273            async_factory: None,
274            lifetime: lt,
275        });
276    }
277
278    /// Register an async service.
279    fn push_async<T: ?Sized + Send + Sync + 'static>(
280        &mut self,
281        lt: ServiceLifetime,
282        key: Option<String>,
283        f: impl Fn(Arc<crate::provider::ServiceProvider>) -> Pin<Box<dyn Future<Output = Arc<T>> + Send>>
284            + Send
285            + Sync
286            + 'static,
287    ) {
288        let af: AsyncServiceFactory = Arc::new(move |sp: Arc<crate::provider::ServiceProvider>| {
289            let fut = (f)(sp);
290            Box::pin(async move {
291                let val: Arc<T> = fut.await;
292                Arc::new(val) as Arc<dyn Any + Send + Sync>
293            }) as Pin<Box<dyn Future<Output = Arc<dyn Any + Send + Sync>> + Send>>
294        });
295        // Sync factory: panics if called directly (user must use build_async).
296        let type_name = std::any::type_name::<T>();
297        let sf: ServiceFactory = Arc::new(move |_| {
298            panic!(
299                "async service '{}' resolved via sync path; use build_async() instead",
300                type_name
301            )
302        });
303        self.descriptors.push(ServiceDescriptor {
304            type_id: TypeId::of::<T>(),
305            type_name: std::any::type_name::<T>(),
306            key,
307            factory: sf,
308            async_factory: Some(af),
309            lifetime: lt,
310        });
311    }
312}
313impl Default for ServiceCollection {
314    fn default() -> Self {
315        Self::new()
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322    #[derive(Debug, PartialEq)]
323    struct G {
324        n: String,
325    }
326    #[derive(Debug, PartialEq)]
327    struct C {
328        v: i32,
329    }
330    #[test]
331    fn empty() {
332        let p = ServiceCollection::new().build().unwrap();
333        assert!(p.get_optional::<G>().is_none());
334    }
335    #[test]
336    fn singleton() {
337        let p = ServiceCollection::new()
338            .singleton(|_| Arc::new(G { n: "Hi".into() }))
339            .build()
340            .unwrap();
341        assert_eq!(p.get::<G>().unwrap().n, "Hi");
342    }
343    #[test]
344    fn singleton_caches() {
345        use std::sync::atomic::{AtomicUsize, Ordering};
346        static CNT: AtomicUsize = AtomicUsize::new(0);
347        let p = ServiceCollection::new()
348            .singleton(|_| {
349                CNT.fetch_add(1, Ordering::SeqCst);
350                Arc::new(C { v: 42 })
351            })
352            .build()
353            .unwrap();
354        let _ = p.get::<C>();
355        let _ = p.get::<C>();
356        assert_eq!(CNT.load(Ordering::SeqCst), 1);
357    }
358    #[test]
359    fn transient_not_cached() {
360        use std::sync::atomic::{AtomicUsize, Ordering};
361        static CNT: AtomicUsize = AtomicUsize::new(0);
362        let p = ServiceCollection::new()
363            .transient(|_| {
364                CNT.fetch_add(1, Ordering::SeqCst);
365                Arc::new(C { v: 1 })
366            })
367            .build()
368            .unwrap();
369        let _ = p.get::<C>();
370        let _ = p.get::<C>();
371        assert_eq!(CNT.load(Ordering::SeqCst), 2);
372    }
373    #[test]
374    fn instance() {
375        let g = Arc::new(G { n: "Inst".into() });
376        let p = ServiceCollection::new()
377            .instance(g.clone())
378            .build()
379            .unwrap();
380        assert!(Arc::ptr_eq(&g, &p.get::<G>().unwrap()));
381    }
382    #[test]
383    fn keyed_svc() {
384        let p = ServiceCollection::new()
385            .keyed_singleton("a", |_| Arc::new(G { n: "A".into() }))
386            .keyed_singleton("b", |_| Arc::new(G { n: "B".into() }))
387            .build()
388            .unwrap();
389        assert_eq!(p.get_keyed::<G>("a").unwrap().n, "A");
390        assert_eq!(p.get_keyed::<G>("b").unwrap().n, "B");
391    }
392    #[test]
393    fn get_all() {
394        let p = ServiceCollection::new()
395            .keyed_singleton("x", |_| Arc::new(C { v: 1 }))
396            .keyed_singleton("y", |_| Arc::new(C { v: 2 }))
397            .build()
398            .unwrap();
399        assert_eq!(p.get_all::<C>().len(), 2);
400    }
401    #[test]
402    fn try_add_skips() {
403        let p = ServiceCollection::new()
404            .singleton(|_| Arc::new(G { n: "First".into() }))
405            .try_add(|_| Arc::new(G { n: "Second".into() }))
406            .build()
407            .unwrap();
408        assert_eq!(p.get::<G>().unwrap().n, "First");
409    }
410}