use std::{any::Any, collections::HashMap, fmt::Debug, future::Future, hash::Hash, sync::Arc}; use tokio::sync::{ watch::{channel as watch_channel, Receiver}, RwLock, }; #[cfg(feature = "cache")] use lru::LruCache; #[cfg(feature = "queue")] use indexmap::IndexMap; use crate::{CoalescionServiceConfig, Error}; #[derive(Debug, Clone)] #[allow(clippy::type_complexity)] /// # Coalescion service /// /// See module description for example usage. pub struct CoalescionService { config: Arc, watchers: Arc, Error>>>>>>, #[cfg(feature = "queue")] queue: Arc, Error>>>>>>, #[cfg(feature = "cache")] cache: Option>>>>, } impl CoalescionService { pub fn new() -> Self { Default::default() } pub fn from_config(config: CoalescionServiceConfig) -> Self { Self { config: Arc::new(config), watchers: Arc::new(RwLock::new(HashMap::new())), #[cfg(feature = "queue")] queue: Arc::new(RwLock::new(IndexMap::new())), #[cfg(feature = "cache")] cache: None, } } #[cfg(feature = "cache")] pub fn from_cache( config: CoalescionServiceConfig, cache: LruCache>, ) -> Self { Self { cache: Some(Arc::new(Mutex::new(cache))), ..Self::from_config(config) } } async fn wait_for( &self, mut receiver: Receiver, Error>>>, ) -> Result, Error> { receiver .wait_for(|v| v.is_some()) .await .map_err(|_| Error::RecvError) .and_then(|r| r.clone().unwrap()) .and_then(|arc| Arc::downcast(arc).map_err(|_| Error::DowncastError)) } async fn insert_and_execute< Value: Send + Sync + 'static, F: FnOnce() -> Fut, Fut: Future, >( &self, id: Id, func: F, ) -> Result, Error> { let (send, recv) = watch_channel(None); self.watchers.write().await.insert(id.clone(), recv); let value = Ok(Arc::new(func().await)); send.send_modify(|opt| { opt.replace(value.clone().map(|v| v as Arc)); }); #[cfg(feature = "cache")] if let Some(cache) = self.cache.as_ref() { if let Ok(value) = &value { cache.lock().await.push(id.clone(), value.clone()); } }; self.watchers.write().await.remove(&id); value } /// Coalesces an function, the actual function may not run if one with the same id is already running, /// queued to be ran, or cached, the id should be globally unique for this specific action. pub async fn execute< Value: Send + Sync + 'static, F: FnOnce() -> Fut, Fut: Future, >( &self, id: Id, func: F, ) -> Result, Error> { #[cfg(feature = "cache")] if let Some(cache) = self.cache.as_ref() { if let Some(value) = cache.lock().await.get(&id) { return Arc::downcast::(value.clone()).map_err(|_| Error::DowncastError); } }; let (receiver, length) = { let watchers = self.watchers.read().await; let length = watchers.len(); (watchers.get(&id).cloned(), length) }; if let Some(receiver) = receiver { self.wait_for(receiver).await } else { match self.config.max_concurrent { Some(max_concurrent) if length >= max_concurrent => { #[cfg(feature = "queue")] if self.config.queue_requests { let (receiver, length) = { let queue = self.queue.read().await; (queue.get(&id).cloned(), queue.len()) }; if let Some(receiver) = receiver { return self.wait_for(receiver).await; } else { if self .config .max_queue .is_some_and(|max_queue| max_queue >= length) { return Err(Error::MaxQueue); }; let (send, recv) = watch_channel(None); self.queue.write().await.insert(id.clone(), recv); loop { let length = self.watchers.read().await.len(); if length < max_concurrent { let first_key = { let queue = self.queue.read().await; queue.first().map(|v| v.0).cloned() }; if first_key == Some(id.clone()) { self.queue.write().await.shift_remove(&id); let response = self.insert_and_execute(id, func).await; send.send_modify(|opt| { opt.replace( response .clone() .map(|v| v as Arc), ); }); return response; } } } } } else { Err(Error::MaxConcurrent) } #[cfg(not(feature = "queue"))] Err(Error::MaxConcurrent) } _ => self.insert_and_execute(id, func).await, } } } /// Fetches the amount of currently running tasks pub async fn current_task_count(&self) -> usize { self.watchers.read().await.len() } #[cfg(feature = "queue")] /// Fetches the current length of the queue pub async fn current_queue_len(&self) -> usize { self.queue.read().await.len() } } impl Default for CoalescionService { fn default() -> Self { Self::from_config(CoalescionServiceConfig::default()) } }