use std::sync::Arc;
use futures::future::join_all;
use serde::de::DeserializeOwned;
use tokio::{
sync::{Mutex, Semaphore},
task::JoinError,
};
use crate::{Error, Resolve, ResolveConfig, Script, script::ArgData};
#[derive(Debug, Clone)]
pub struct PooledResolve {
inner: Arc<InternalPool>,
}
#[derive(Debug)]
struct InternalPool {
instances: Mutex<Vec<Resolve>>,
permits: Semaphore,
}
impl PooledResolve {
pub async fn new(amount: usize) -> Result<Self, Error> {
Self::new_with_config(amount, ResolveConfig::default()).await
}
pub async fn new_with_config(amount: usize, config: ResolveConfig) -> Result<Self, Error> {
let mut handles = Vec::with_capacity(amount);
let conf = Arc::new(config);
for _ in 0..amount {
let config = conf.clone();
handles.push(tokio::spawn(async move {
Resolve::new_with_config(config.as_ref()).await
}));
}
let instances: Result<Result<Vec<Resolve>, Error>, JoinError> =
join_all(handles).await.into_iter().collect();
let instances = instances??;
let pool = InternalPool {
instances: Mutex::new(instances),
permits: Semaphore::new(amount),
};
Ok(Self {
inner: Arc::new(pool),
})
}
pub(crate) async fn on_lock<T: DeserializeOwned>(
script: Script<'_>,
instance: &Resolve,
) -> Result<T, Error> {
instance.execute::<T>(script).await
}
pub async fn execute<T>(&self, script: impl Into<Script<'_>>) -> Result<T, Error>
where
T: DeserializeOwned,
{
let script = script.into();
for arg in &script.args {
match arg {
ArgData::ArgRef(_) | ArgData::NamedArgRef { key: _, value: _ } => {
return Err(Error::CantHoldReferenceInPool);
}
_ => (),
}
}
let permit = self.inner.permits.acquire().await?;
let instance = {
let mut inst = self.inner.instances.lock().await;
inst.pop().ok_or(Error::OutOfSyncSemaphore)?
};
let value_result = Self::on_lock(script, &instance).await;
{
self.inner.instances.lock().await.push(instance);
}
drop(permit);
value_result
}
}