elemental_object/
pool.rs

1use scc::HashMap;
2#[cfg(feature = "notify")]
3use std::sync::atomic::AtomicBool;
4use std::{
5    any::{Any, TypeId},
6    future::Future,
7    pin::Pin,
8    sync::{Arc, LazyLock, atomic::Ordering},
9};
10#[cfg(feature = "notify")]
11use tokio::sync::Notify;
12
13pub static POOL: LazyLock<ObjectPool> = LazyLock::new(ObjectPool::new);
14pub type ShutdownFn<T> =
15    Box<dyn Fn(Arc<T>) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>;
16pub struct PoolEntry<T: Any + Send + Sync + ?Sized> {
17    value: Option<Arc<T>>,
18    #[cfg(feature = "notify")]
19    notify: Arc<Notify>,
20    #[cfg(feature = "notify")]
21    active: AtomicBool,
22    shutdown: Option<ShutdownFn<T>>,
23}
24
25impl<T: Any + Send + Sync + ?Sized> PoolEntry<T> {
26    async fn shutdown(&mut self) {
27        let old_value = self.value.take();
28        let old_shutdown = self.shutdown.take();
29        #[cfg(feature = "notify")]
30        {
31            // maybe one comsumer is going to waiting for value, but the value is shutdown, we should set active to false to avoid waiting forever.
32            self.active.store(false, Ordering::Release);
33            // maybe an acquire is waiting for this entry, we should notify it after shutdown, otherwise it may wait forever.
34            self.notify.notify_waiters();
35        }
36
37        if let (Some(v), Some(s)) = (old_value, old_shutdown) {
38            (s)(v).await;
39        }
40    }
41
42    async fn reload(&mut self, value: Arc<T>, shutdown: Option<ShutdownFn<T>>) {
43        let old_value = self.value.replace(value);
44        let old_shutdown = std::mem::replace(&mut self.shutdown, shutdown);
45        #[cfg(feature = "notify")]
46        {
47            self.active.store(true, Ordering::Release);
48            self.notify.notify_waiters();
49        }
50
51        if let (Some(v), Some(s)) = (old_value, old_shutdown) {
52            (s)(v).await;
53        }
54    }
55
56    fn new() -> Self {
57        Self {
58            value: None,
59            #[cfg(feature = "notify")]
60            notify: Arc::new(Notify::new()),
61            #[cfg(feature = "notify")]
62            active: AtomicBool::new(true),
63            shutdown: None,
64        }
65    }
66}
67
68pub struct ObjectPool {
69    inner: HashMap<TypeId, PoolEntry<dyn Any + Send + Sync>>,
70}
71
72impl Default for ObjectPool {
73    fn default() -> Self {
74        Self::new()
75    }
76}
77
78impl ObjectPool {
79    /// Create a new Object Pool
80    pub fn new() -> Self {
81        Self {
82            inner: HashMap::new(),
83        }
84    }
85
86    pub async fn get_async<T: Any + Send + Sync>(&self) -> Option<Arc<T>> {
87        let typeid = TypeId::of::<T>();
88        let entry = self.inner.read_async(&typeid, |_, v| v.value.clone()).await;
89        if let Some(entry) = entry {
90            return Arc::downcast::<T>(entry?).ok();
91        }
92
93        None
94    }
95
96    // Overwrite and call the notifier if exists, otherwise just insert the value.
97    pub async fn set_value<T: Any + Send + Sync>(
98        &self,
99        value: Arc<T>,
100        shutdown: Option<ShutdownFn<T>>,
101    ) {
102        let type_id = TypeId::of::<T>();
103        let shutdown = shutdown.map(|shutdown_fn| {
104            Box::new(move |value_any: Arc<dyn Any + Send + Sync>| {
105                let value = value_any
106                    .downcast::<T>()
107                    .expect("Can't downcast value in shutdown fn");
108                shutdown_fn(value)
109            }) as ShutdownFn<dyn Any + Send + Sync>
110        });
111        if self.inner.contains_async(&type_id).await {
112            // Call shutdown of old value if exists, and update the value and shutdown fn.
113            // It means the old value will be shutdown before the new value is set, so the shutdown fn can still access the old value.
114            let mut entry = self.inner.get_async(&type_id).await.unwrap();
115            entry.reload(value, shutdown).await;
116            return;
117        }
118
119        let mut entry = PoolEntry::new();
120        entry.value = Some(value.clone() as Arc<dyn Any + Send + Sync>);
121        entry.shutdown = shutdown;
122        // The entry is new, just insert it!
123        self.inner.upsert_async(type_id, entry).await;
124    }
125
126    pub async fn remove_value<T: Any + Send + Sync>(&self) {
127        let type_id = TypeId::of::<T>();
128        if let Some(mut entry) = self.inner.get_async(&type_id).await {
129            entry.shutdown().await;
130        }
131    }
132
133    pub async fn remove_entry<T: Any + Send + Sync>(&self) {
134        let type_id = TypeId::of::<T>();
135        if let Some((_, mut entry)) = self.inner.remove_async(&type_id).await {
136            entry.shutdown().await;
137        }
138    }
139
140    #[cfg(feature = "notify")]
141    pub async fn wait_value<T: Any + Send + Sync>(&self) {
142        let type_id = TypeId::of::<T>();
143
144        // If the value is shutdown, return immediately to avoid waiting forever.
145        if let Some(active) = self
146            .inner
147            .read_async(&type_id, |_, v| v.active.load(Ordering::Acquire))
148            .await
149            && !active
150        {
151            return;
152        }
153
154        if self
155            .inner
156            .read_async(&type_id, |_, v| v.value.clone())
157            .await
158            .is_some()
159        {
160            // If the value exists, return immediately.
161            return;
162        }
163
164        let notified = self
165            .inner
166            .entry_async(type_id)
167            .await
168            .or_insert(PoolEntry::new())
169            .get()
170            .notify
171            .clone();
172
173        notified.notified().await;
174    }
175
176    pub async fn fulfill_value<T: Any + Send + Sync>(
177        &self,
178        value: Arc<T>,
179        shutdown: Option<ShutdownFn<T>>,
180    ) {
181        let type_id = TypeId::of::<T>();
182        if let Some(fulfilled) = self
183            .inner
184            .read_async(&type_id, |_, v| v.value.is_some())
185            .await
186            && fulfilled
187        {
188            // Noops if already fulfilled
189            return;
190        }
191
192        if let Some(mut entry) = self.inner.get_async(&type_id).await {
193            entry.value = Some(value.clone() as Arc<dyn Any + Send + Sync>);
194            let shutdown = shutdown.map(|shutdown_fn| {
195                Box::new(move |value_any: Arc<dyn Any + Send + Sync>| {
196                    let value = value_any
197                        .downcast::<T>()
198                        .expect("Can't downcast value in shutdown fn");
199                    shutdown_fn(value)
200                }) as ShutdownFn<dyn Any + Send + Sync>
201            });
202            entry.shutdown = shutdown;
203            #[cfg(feature = "notify")]
204            {
205                entry.active.store(true, Ordering::Release);
206                entry.notify.notify_waiters();
207            }
208        }
209    }
210
211    pub async fn shutdown(&self) {
212        let mut iter = self.inner.begin_async().await;
213        while let Some(mut entry) = iter {
214            // `OccupiedEntry` can be sent across awaits and threads.
215            // clean up the old value if it's exists
216            entry.shutdown().await;
217
218            iter = entry.next_async().await;
219        }
220    }
221}
elemental_object/pool.rs

elemental_object/
pool.rs
