1use std::any::{Any, TypeId};
8use std::collections::HashMap;
9use std::sync::Arc;
10
11use crate::{
12 Injectable, InjectableError, InjectableResult, PreDestruct, Provider, ProviderRegistry,
13 SingletonStore,
14};
15
16pub type SingletonCache = Arc<
17 tokio::sync::Mutex<HashMap<TypeId, Arc<tokio::sync::OnceCell<Arc<dyn Any + Send + Sync>>>>>,
18>;
19
20struct DestructorEntry {
25 instance: Arc<dyn PreDestruct>,
26 type_name: &'static str,
27}
28
29pub struct ResolveContext {
49 store: Arc<dyn SingletonStore>,
50 registry: Arc<ProviderRegistry>,
51 destructors: Arc<tokio::sync::Mutex<Vec<DestructorEntry>>>,
52 singleton_cache: SingletonCache,
56}
57
58impl ResolveContext {
59 pub fn new(store: Arc<dyn SingletonStore>, registry: Arc<ProviderRegistry>) -> Self {
61 Self {
62 store,
63 registry,
64 destructors: Arc::new(tokio::sync::Mutex::new(Vec::new())),
65 singleton_cache: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
66 }
67 }
68
69 pub fn from_store(store: Arc<dyn SingletonStore>) -> Self {
71 Self {
72 store,
73 registry: Arc::new(ProviderRegistry::new()),
74 destructors: Arc::new(tokio::sync::Mutex::new(Vec::new())),
75 singleton_cache: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
76 }
77 }
78
79 pub fn store(&self) -> &Arc<dyn SingletonStore> {
81 &self.store
82 }
83
84 pub fn registry(&self) -> &ProviderRegistry {
86 &self.registry
87 }
88
89 pub async fn extract<T>(&self) -> crate::InjectableResult<T>
105 where
106 T: crate::Extract + Send + Sync + 'static,
107 {
108 T::extract(self).await
109 }
110
111 pub async fn clone_from_singleton<T: Injectable + Clone>(&self) -> InjectableResult<T> {
117 Ok(Arc::unwrap_or_clone(
118 self.resolve_singleton_arc::<T>().await?,
119 ))
120 }
121
122 #[allow(clippy::expect_used)]
136 pub(crate) async fn resolve_singleton_arc<T: Injectable>(&self) -> InjectableResult<Arc<T>> {
137 let type_id = TypeId::of::<T>();
138
139 let cell = {
141 let mut cache = self.singleton_cache.lock().await;
142 Arc::clone(
143 cache
144 .entry(type_id)
145 .or_insert_with(|| Arc::new(tokio::sync::OnceCell::new())),
146 )
147 };
148
149 let arc_any = cell
151 .get_or_try_init(|| async {
152 let value = T::Provider::provide(self).await?;
153 let arc_t: Arc<T> = Arc::new(value);
155 Ok(Arc::new(arc_t) as Arc<dyn Any + Send + Sync>)
156 })
157 .await?;
158
159 let inner: &Arc<T> = (**arc_any)
161 .downcast_ref::<Arc<T>>()
162 .expect("TypeId guarantees type correctness; downcast cannot fail");
163 Ok(Arc::clone(inner))
164 }
165
166 pub async fn resolve_external<T: Send + Sync + 'static>(&self) -> InjectableResult<T> {
179 self.resolve_external_with_token::<T>(crate::registry::DEFAULT_TOKEN)
180 .await
181 }
182
183 pub async fn resolve_external_with_token<T: Send + Sync + 'static>(
195 &self,
196 token: &str,
197 ) -> InjectableResult<T> {
198 match self
199 .registry
200 .resolve_with_token::<T>(token, Arc::new(self.clone()))
201 .await
202 {
203 Some(result) => result,
204 None => Err(InjectableError::MissingDependency {
205 type_name: std::any::type_name::<T>(),
206 }),
207 }
208 }
209
210 pub async fn try_resolve_external<T: Send + Sync + 'static>(
214 &self,
215 ) -> Option<InjectableResult<T>> {
216 self.registry
217 .resolve_with_token::<T>(crate::registry::DEFAULT_TOKEN, Arc::new(self.clone()))
218 .await
219 }
220
221 pub async fn try_resolve_external_with_token<T: Send + Sync + 'static>(
225 &self,
226 token: &str,
227 ) -> Option<InjectableResult<T>> {
228 self.registry
229 .resolve_with_token::<T>(token, Arc::new(self.clone()))
230 .await
231 }
232
233 pub fn register_destructor(&self, instance: Arc<dyn PreDestruct>) {
239 if let Ok(mut destructors) = self.destructors.try_lock() {
243 destructors.push(DestructorEntry {
244 type_name: "",
245 instance,
246 });
247 }
248 }
249
250 pub fn register_destructor_with_name(
252 &self,
253 type_name: &'static str,
254 instance: Arc<dyn PreDestruct>,
255 ) {
256 if let Ok(mut destructors) = self.destructors.try_lock() {
257 destructors.push(DestructorEntry {
258 type_name,
259 instance,
260 });
261 }
262 }
263
264 pub async fn run_destructors(&self) -> Result<(), Vec<crate::InjectableError>> {
274 let mut destructors = self.destructors.lock().await;
275 let mut errors = Vec::new();
276
277 while let Some(entry) = destructors.pop() {
279 match entry.instance.pre_destruct().await {
280 Ok(()) => {}
281 Err(e) => {
282 errors.push(crate::InjectableError::LifecycleHookFailed {
283 type_name: entry.type_name,
284 hook: "pre_destruct",
285 reason: e.to_string(),
286 });
287 }
288 }
289 }
290
291 if errors.is_empty() {
292 Ok(())
293 } else {
294 Err(errors)
295 }
296 }
297
298 pub async fn destructor_count(&self) -> usize {
300 self.destructors.lock().await.len()
301 }
302}
303
304impl Clone for ResolveContext {
305 fn clone(&self) -> Self {
306 Self {
307 store: Arc::clone(&self.store),
308 registry: Arc::clone(&self.registry),
309 destructors: Arc::clone(&self.destructors),
310 singleton_cache: Arc::clone(&self.singleton_cache),
311 }
312 }
313}
314
315impl std::fmt::Debug for ResolveContext {
316 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
317 f.debug_struct("ResolveContext")
318 .field("store", &"Arc<dyn SingletonStore>")
319 .field("registry", &self.registry)
320 .finish()
321 }
322}
323
324#[cfg(test)]
325mod tests {
326 use super::*;
327 use crate::{DynProvider, EmptySingletonStore, HookResult, PreDestruct, ProviderRegistry};
328 use std::sync::Arc;
329
330 fn make_ctx() -> ResolveContext {
331 ResolveContext::new(
332 Arc::new(EmptySingletonStore),
333 Arc::new(ProviderRegistry::new()),
334 )
335 }
336
337 #[test]
338 fn from_store_creates_context() {
339 let ctx = ResolveContext::from_store(Arc::new(EmptySingletonStore));
340 assert!(ctx.registry().is_empty());
341 }
342
343 #[test]
344 fn store_and_registry_accessors() {
345 let ctx = make_ctx();
346 assert_eq!(ctx.store().len(), 0);
347 assert!(ctx.registry().is_empty());
348 }
349
350 #[test]
351 fn clone_shares_destructors() {
352 let ctx = make_ctx();
353 let ctx2 = ctx.clone();
354 assert!(Arc::ptr_eq(&ctx.destructors, &ctx2.destructors));
356 }
357
358 #[test]
359 fn debug_impl() {
360 let ctx = make_ctx();
361 let s = format!("{ctx:?}");
362 assert!(s.contains("ResolveContext"));
363 }
364
365 #[tokio::test]
366 async fn destructor_count_starts_zero() {
367 let ctx = make_ctx();
368 assert_eq!(ctx.destructor_count().await, 0);
369 }
370
371 #[tokio::test]
372 async fn run_destructors_empty_ok() {
373 let ctx = make_ctx();
374 assert!(ctx.run_destructors().await.is_ok());
375 }
376
377 #[tokio::test]
378 async fn register_destructor_increments_count() {
379 struct NoopDestructor;
380 #[async_trait::async_trait]
381 impl PreDestruct for NoopDestructor {
382 async fn pre_destruct(&self) -> HookResult {
383 Ok(())
384 }
385 }
386
387 let ctx = make_ctx();
388 ctx.register_destructor(Arc::new(NoopDestructor));
389 assert_eq!(ctx.destructor_count().await, 1);
390 }
391
392 #[tokio::test]
393 async fn register_destructor_with_name_increments_count() {
394 struct NoopDestructor;
395 #[async_trait::async_trait]
396 impl PreDestruct for NoopDestructor {
397 async fn pre_destruct(&self) -> HookResult {
398 Ok(())
399 }
400 }
401
402 let ctx = make_ctx();
403 ctx.register_destructor_with_name("TestType", Arc::new(NoopDestructor));
404 assert_eq!(ctx.destructor_count().await, 1);
405 }
406
407 #[tokio::test]
408 async fn run_destructors_calls_hooks() {
409 use std::sync::atomic::{AtomicBool, Ordering};
410
411 static CALLED: AtomicBool = AtomicBool::new(false);
412
413 struct FlagDestructor;
414 #[async_trait::async_trait]
415 impl PreDestruct for FlagDestructor {
416 async fn pre_destruct(&self) -> HookResult {
417 CALLED.store(true, Ordering::SeqCst);
418 Ok(())
419 }
420 }
421
422 CALLED.store(false, Ordering::SeqCst);
423 let ctx = make_ctx();
424 ctx.register_destructor(Arc::new(FlagDestructor));
425 ctx.run_destructors().await.unwrap();
426 assert!(CALLED.load(Ordering::SeqCst));
427 assert_eq!(ctx.destructor_count().await, 0);
428 }
429
430 #[tokio::test]
431 async fn run_destructors_collects_errors() {
432 struct FailingDestructor;
433 #[async_trait::async_trait]
434 impl PreDestruct for FailingDestructor {
435 async fn pre_destruct(&self) -> HookResult {
436 Err(Box::new(std::io::Error::new(
437 std::io::ErrorKind::Other,
438 "fail",
439 )))
440 }
441 }
442
443 let ctx = make_ctx();
444 ctx.register_destructor(Arc::new(FailingDestructor));
445 let result = ctx.run_destructors().await;
446 assert!(result.is_err());
447 assert_eq!(result.unwrap_err().len(), 1);
448 }
449
450 #[tokio::test]
451 async fn resolve_external_missing_returns_error() {
452 let ctx = make_ctx();
453 let result = ctx.resolve_external::<String>().await;
454 assert!(result.is_err());
455 }
456
457 #[tokio::test]
458 async fn resolve_external_registered_returns_value() {
459 let mut registry = ProviderRegistry::new();
460 registry.register("", DynProvider::from_value(42u32));
461 let ctx = ResolveContext::new(Arc::new(EmptySingletonStore), Arc::new(registry));
462 let v: u32 = ctx.resolve_external().await.unwrap();
463 assert_eq!(v, 42);
464 }
465
466 #[tokio::test]
467 async fn try_resolve_external_missing_returns_none() {
468 let ctx = make_ctx();
469 let result = ctx.try_resolve_external::<String>().await;
470 assert!(result.is_none());
471 }
472
473 #[tokio::test]
474 async fn try_resolve_external_registered_returns_some() {
475 let mut registry = ProviderRegistry::new();
476 registry.register("", DynProvider::from_value(99u32));
477 let ctx = ResolveContext::new(Arc::new(EmptySingletonStore), Arc::new(registry));
478 let result = ctx.try_resolve_external::<u32>().await.unwrap();
479 assert_eq!(result.unwrap(), 99u32);
480 }
481
482 #[tokio::test]
483 async fn resolve_external_with_token_resolves_named_provider() {
484 let mut registry = ProviderRegistry::new();
485 registry.register("primary", DynProvider::from_value(1u32));
486 registry.register("replica", DynProvider::from_value(2u32));
487 let ctx = ResolveContext::new(Arc::new(EmptySingletonStore), Arc::new(registry));
488 let primary: u32 = ctx.resolve_external_with_token("primary").await.unwrap();
489 let replica: u32 = ctx.resolve_external_with_token("replica").await.unwrap();
490 assert_eq!(primary, 1);
491 assert_eq!(replica, 2);
492 }
493
494 #[tokio::test]
495 async fn resolve_external_with_default_token_doesnt_find_named() {
496 let mut registry = ProviderRegistry::new();
497 registry.register("primary", DynProvider::from_value(1u32));
498 let ctx = ResolveContext::new(Arc::new(EmptySingletonStore), Arc::new(registry));
499 let result = ctx.resolve_external::<u32>().await;
501 assert!(
502 result.is_err(),
503 "default token should not resolve named provider"
504 );
505 }
506
507 #[tokio::test]
508 async fn try_resolve_external_with_token_returns_none_for_wrong_token() {
509 let mut registry = ProviderRegistry::new();
510 registry.register("primary", DynProvider::from_value(1u32));
511 let ctx = ResolveContext::new(Arc::new(EmptySingletonStore), Arc::new(registry));
512 let result = ctx.try_resolve_external_with_token::<u32>("replica").await;
513 assert!(result.is_none());
514 }
515}