forked from jmug/stoatchat
feat: initial work on tenor gif searching
This commit is contained in:
committed by
Angelo Kontaxis
parent
bfe4018e43
commit
b0c977b324
23
crates/core/coalesced/src/config.rs
Normal file
23
crates/core/coalesced/src/config.rs
Normal file
@@ -0,0 +1,23 @@
|
||||
#[derive(Clone, PartialEq, Eq, Debug)]
|
||||
pub struct CoalescionServiceConfig {
|
||||
/// How many tasks are running at once
|
||||
pub max_concurrent: Option<usize>,
|
||||
/// Whether to queue tasks once `max_concurrent` is reached
|
||||
#[cfg(feature = "queue")]
|
||||
pub queue_requests: bool,
|
||||
/// Max amount of tasks in the buffer queue
|
||||
#[cfg(feature = "queue")]
|
||||
pub max_queue: Option<usize>,
|
||||
}
|
||||
|
||||
impl Default for CoalescionServiceConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_concurrent: Some(100),
|
||||
#[cfg(feature = "queue")]
|
||||
queue_requests: true,
|
||||
#[cfg(feature = "queue")]
|
||||
max_queue: Some(100)
|
||||
}
|
||||
}
|
||||
}
|
||||
20
crates/core/coalesced/src/error.rs
Normal file
20
crates/core/coalesced/src/error.rs
Normal file
@@ -0,0 +1,20 @@
|
||||
use std::fmt::Display;
|
||||
|
||||
#[derive(Clone, PartialEq, Eq, Debug)]
|
||||
pub enum Error {
|
||||
RecvError,
|
||||
MaxConcurrent,
|
||||
MaxQueue,
|
||||
}
|
||||
|
||||
impl Display for Error {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Error::RecvError => write!(f, "Unable to receive data from the channel"),
|
||||
Error::MaxConcurrent => write!(f, "Max number of tasks running at once"),
|
||||
Error::MaxQueue => write!(f, "Max number of tasks in queue"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for Error {}
|
||||
7
crates/core/coalesced/src/lib.rs
Normal file
7
crates/core/coalesced/src/lib.rs
Normal file
@@ -0,0 +1,7 @@
|
||||
mod config;
|
||||
mod error;
|
||||
mod service;
|
||||
|
||||
pub use config::CoalescionServiceConfig;
|
||||
pub use error::Error;
|
||||
pub use service::CoalescionService;
|
||||
168
crates/core/coalesced/src/service.rs
Normal file
168
crates/core/coalesced/src/service.rs
Normal file
@@ -0,0 +1,168 @@
|
||||
use std::{collections::{HashMap}, fmt::Debug, hash::Hash, sync::Arc, future::Future};
|
||||
|
||||
use tokio::{sync::{watch::{channel as watch_channel, Receiver}, RwLock, Mutex}};
|
||||
|
||||
#[cfg(feature = "cache")]
|
||||
use lru::LruCache;
|
||||
|
||||
#[cfg(feature = "queue")]
|
||||
use indexmap::IndexMap;
|
||||
|
||||
use crate::{CoalescionServiceConfig, Error};
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
#[allow(clippy::type_complexity)]
|
||||
pub struct CoalescionService<Id: Hash + Eq, Value> {
|
||||
config: Arc<CoalescionServiceConfig>,
|
||||
watchers: Arc<RwLock<HashMap<Id, Receiver<Option<Result<Arc<Value>, Error>>>>>>,
|
||||
#[cfg(feature = "queue")]
|
||||
queue: Arc<RwLock<IndexMap<Id, Receiver<Option<Result<Arc<Value>, Error>>>>>>,
|
||||
#[cfg(feature = "cache")]
|
||||
cache: Option<Arc<Mutex<LruCache<Id, Arc<Value>>>>>,
|
||||
}
|
||||
|
||||
impl<Id: Hash + PartialEq + Eq + Clone + Ord, Value> CoalescionService<Id, Value> {
|
||||
pub fn new() -> Self {
|
||||
Self::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<Id, Arc<Value>>) -> Self {
|
||||
Self {
|
||||
cache: Some(Arc::new(Mutex::new(cache))),
|
||||
..Self::from_config(config)
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_for(&self, mut receiver: Receiver<Option<Result<Arc<Value>, Error>>>) -> Result<Arc<Value>, Error> {
|
||||
receiver
|
||||
.wait_for(|v| v.is_some())
|
||||
.await
|
||||
.map_err(|_| Error::RecvError)
|
||||
.and_then(|r| r.clone().unwrap())
|
||||
}
|
||||
|
||||
async fn insert_and_execute<F: FnOnce() -> Fut, Fut: Future<Output = Value>>(&self, id: Id, func: F) -> Result<Arc<Value>, 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());
|
||||
});
|
||||
|
||||
#[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
|
||||
}
|
||||
|
||||
pub async fn execute<F: FnOnce() -> Fut, Fut: Future<Output = Value>>(&self, id: Id, func: F) -> Result<Arc<Value>, Error> {
|
||||
#[cfg(feature = "cache")]
|
||||
if let Some(cache) = self.cache.as_ref() {
|
||||
if let Some(value) = cache.lock().await.get(&id) {
|
||||
return Ok(value.clone())
|
||||
}
|
||||
};
|
||||
|
||||
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());
|
||||
});
|
||||
|
||||
return response
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
Err(Error::MaxConcurrent)
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "queue"))]
|
||||
Err(Error::MaxConcurrent)
|
||||
}
|
||||
_ => {
|
||||
self.insert_and_execute(id, func).await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn current_task_count(&self) -> usize {
|
||||
self.watchers.read().await.len()
|
||||
}
|
||||
|
||||
pub async fn current_queue_len(&self) -> usize {
|
||||
self.queue.read().await.len()
|
||||
}
|
||||
}
|
||||
|
||||
impl<Id: Hash + PartialEq + Eq + Clone + Ord, Value> Default for CoalescionService<Id, Value> {
|
||||
fn default() -> Self {
|
||||
Self::from_config(CoalescionServiceConfig::default())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user