elemental-object 0.1.0

use scc::HashMap;
#[cfg(feature = "notify")]
use std::sync::atomic::AtomicBool;
use std::{
    any::{Any, TypeId},
    future::Future,
    pin::Pin,
    sync::{Arc, LazyLock, atomic::Ordering},
};
#[cfg(feature = "notify")]
use tokio::sync::Notify;

pub static POOL: LazyLock<ObjectPool> = LazyLock::new(ObjectPool::new);
pub type ShutdownFn<T> =
    Box<dyn Fn(Arc<T>) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>;
pub struct PoolEntry<T: Any + Send + Sync + ?Sized> {
    value: Option<Arc<T>>,
    #[cfg(feature = "notify")]
    notify: Arc<Notify>,
    #[cfg(feature = "notify")]
    active: AtomicBool,
    shutdown: Option<ShutdownFn<T>>,
}

impl<T: Any + Send + Sync + ?Sized> PoolEntry<T> {
    async fn shutdown(&mut self) {
        let old_value = self.value.take();
        let old_shutdown = self.shutdown.take();
        #[cfg(feature = "notify")]
        {
            // maybe one comsumer is going to waiting for value, but the value is shutdown, we should set active to false to avoid waiting forever.
            self.active.store(false, Ordering::Release);
            // maybe an acquire is waiting for this entry, we should notify it after shutdown, otherwise it may wait forever.
            self.notify.notify_waiters();
        }

        if let (Some(v), Some(s)) = (old_value, old_shutdown) {
            (s)(v).await;
        }
    }

    async fn reload(&mut self, value: Arc<T>, shutdown: Option<ShutdownFn<T>>) {
        let old_value = self.value.replace(value);
        let old_shutdown = std::mem::replace(&mut self.shutdown, shutdown);
        #[cfg(feature = "notify")]
        {
            self.active.store(true, Ordering::Release);
            self.notify.notify_waiters();
        }

        if let (Some(v), Some(s)) = (old_value, old_shutdown) {
            (s)(v).await;
        }
    }

    fn new() -> Self {
        Self {
            value: None,
            #[cfg(feature = "notify")]
            notify: Arc::new(Notify::new()),
            #[cfg(feature = "notify")]
            active: AtomicBool::new(true),
            shutdown: None,
        }
    }
}

pub struct ObjectPool {
    inner: HashMap<TypeId, PoolEntry<dyn Any + Send + Sync>>,
}

impl Default for ObjectPool {
    fn default() -> Self {
        Self::new()
    }
}

impl ObjectPool {
    /// Create a new Object Pool
    pub fn new() -> Self {
        Self {
            inner: HashMap::new(),
        }
    }

    pub async fn get_async<T: Any + Send + Sync>(&self) -> Option<Arc<T>> {
        let typeid = TypeId::of::<T>();
        let entry = self.inner.read_async(&typeid, |_, v| v.value.clone()).await;
        if let Some(entry) = entry {
            return Arc::downcast::<T>(entry?).ok();
        }

        None
    }

    // Overwrite and call the notifier if exists, otherwise just insert the value.
    pub async fn set_value<T: Any + Send + Sync>(
        &self,
        value: Arc<T>,
        shutdown: Option<ShutdownFn<T>>,
    ) {
        let type_id = TypeId::of::<T>();
        let shutdown = shutdown.map(|shutdown_fn| {
            Box::new(move |value_any: Arc<dyn Any + Send + Sync>| {
                let value = value_any
                    .downcast::<T>()
                    .expect("Can't downcast value in shutdown fn");
                shutdown_fn(value)
            }) as ShutdownFn<dyn Any + Send + Sync>
        });
        if self.inner.contains_async(&type_id).await {
            // Call shutdown of old value if exists, and update the value and shutdown fn.
            // It means the old value will be shutdown before the new value is set, so the shutdown fn can still access the old value.
            let mut entry = self.inner.get_async(&type_id).await.unwrap();
            entry.reload(value, shutdown).await;
            return;
        }

        let mut entry = PoolEntry::new();
        entry.value = Some(value.clone() as Arc<dyn Any + Send + Sync>);
        entry.shutdown = shutdown;
        // The entry is new, just insert it!
        self.inner.upsert_async(type_id, entry).await;
    }

    pub async fn remove_value<T: Any + Send + Sync>(&self) {
        let type_id = TypeId::of::<T>();
        if let Some(mut entry) = self.inner.get_async(&type_id).await {
            entry.shutdown().await;
        }
    }

    pub async fn remove_entry<T: Any + Send + Sync>(&self) {
        let type_id = TypeId::of::<T>();
        if let Some((_, mut entry)) = self.inner.remove_async(&type_id).await {
            entry.shutdown().await;
        }
    }

    #[cfg(feature = "notify")]
    pub async fn wait_value<T: Any + Send + Sync>(&self) {
        let type_id = TypeId::of::<T>();

        // If the value is shutdown, return immediately to avoid waiting forever.
        if let Some(active) = self
            .inner
            .read_async(&type_id, |_, v| v.active.load(Ordering::Acquire))
            .await
            && !active
        {
            return;
        }

        if self
            .inner
            .read_async(&type_id, |_, v| v.value.clone())
            .await
            .is_some()
        {
            // If the value exists, return immediately.
            return;
        }

        let notified = self
            .inner
            .entry_async(type_id)
            .await
            .or_insert(PoolEntry::new())
            .get()
            .notify
            .clone();

        notified.notified().await;
    }

    pub async fn fulfill_value<T: Any + Send + Sync>(
        &self,
        value: Arc<T>,
        shutdown: Option<ShutdownFn<T>>,
    ) {
        let type_id = TypeId::of::<T>();
        if let Some(fulfilled) = self
            .inner
            .read_async(&type_id, |_, v| v.value.is_some())
            .await
            && fulfilled
        {
            // Noops if already fulfilled
            return;
        }

        if let Some(mut entry) = self.inner.get_async(&type_id).await {
            entry.value = Some(value.clone() as Arc<dyn Any + Send + Sync>);
            let shutdown = shutdown.map(|shutdown_fn| {
                Box::new(move |value_any: Arc<dyn Any + Send + Sync>| {
                    let value = value_any
                        .downcast::<T>()
                        .expect("Can't downcast value in shutdown fn");
                    shutdown_fn(value)
                }) as ShutdownFn<dyn Any + Send + Sync>
            });
            entry.shutdown = shutdown;
            #[cfg(feature = "notify")]
            {
                entry.active.store(true, Ordering::Release);
                entry.notify.notify_waiters();
            }
        }
    }

    pub async fn shutdown(&self) {
        let mut iter = self.inner.begin_async().await;
        while let Some(mut entry) = iter {
            // `OccupiedEntry` can be sent across awaits and threads.
            // clean up the old value if it's exists
            entry.shutdown().await;

            iter = entry.next_async().await;
        }
    }
}