exum 0.1.10

A lightweight Axum syntax sugar library
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

use std::{any::{Any, TypeId}, collections::HashMap, pin::Pin, sync::Arc};
use inventory;

use tokio::sync::OnceCell;

use crate::ext::ArcAnyExt;

static GLOBAL_STATE_CONTAINER: OnceCell<Arc<LazyDependencyContainer>> = OnceCell::const_new();

pub async fn init_global_state() -> Arc<LazyDependencyContainer> {
    let container = LazyDependencyContainer::new();
    let _ = GLOBAL_STATE_CONTAINER.set(container.clone());
    container
}

pub fn global_container() -> Arc<LazyDependencyContainer> {
    GLOBAL_STATE_CONTAINER
        .get()
        .expect("State container not initialized")
        .clone()
}


pub struct StateDef {
    pub type_id: TypeId,
    pub prewarm: bool,
    pub init_fn: fn() -> Pin<Box<dyn Future<Output = Arc<dyn Any + Send + Sync>> + Send>>,
}
pub struct StateDefFn(pub fn() -> StateDef);

inventory::collect!(StateDefFn);
pub fn collect_states() -> Vec<StateDef> {
    inventory::iter::<StateDefFn>
        .into_iter()
        .map(|f| f.0())
        .collect()
}

type StateFuture =
    Pin<Box<dyn Future<Output = Arc<dyn Any + Send + Sync>> + Send>>;
pub struct LazyDependencyContainer {
  registry: HashMap<TypeId, fn() -> StateFuture>,
  prewarm_flags: HashMap<TypeId, bool>,
  instances: HashMap<TypeId, Arc<OnceCell<Arc<dyn Any + Send + Sync>>>>,
}
impl LazyDependencyContainer {
  pub fn new() -> Arc<Self> {
    let mut registry = HashMap::new();
    let mut prewarm_flags = HashMap::new();
    let mut instances = HashMap::new();

    for def in collect_states() {
        registry.insert(def.type_id, def.init_fn);
        prewarm_flags.insert(def.type_id, def.prewarm);
        instances.insert(def.type_id, Arc::new(OnceCell::new()));
    }

    Arc::new(Self {
      registry,
      prewarm_flags,
      instances,
    })
  }

  pub fn register(
      &mut self,
      type_id: TypeId,
      prewarm: bool,
      init: fn() -> StateFuture,
  ) {
      self.registry.insert(type_id, init);
      self.prewarm_flags.insert(type_id, prewarm);
      self.instances
          .insert(type_id, Arc::new(OnceCell::new()));
  }

  pub async fn get<T:'static + Send + Sync>(&self) -> Arc<T> {
    let type_id = TypeId::of::<T>();
    let init_fn = self.registry.get(&type_id).expect("No init function found for type");
    let cell = self.instances.get(&type_id).expect("No instance cell found for type");
    let instance = cell.get_or_init(|| init_fn()).await.clone();
    instance.clone().downcast_arc::<T>().unwrap()
  }

  pub async fn prewarm_all(&self) {
      for (type_id, prewarm) in &self.prewarm_flags {
          if *prewarm {
              let init_fn = self.registry.get(type_id).unwrap();
              let cell = self.instances.get(type_id).unwrap();
              let _ = cell.get_or_init(|| init_fn()).await;
          }
      }
  }
}