use super::pool_object::PoolObject;
use crate::config::{AutoPoolConfig, PickStrategy};
use parking_lot::lock_api::{MutexGuard, RawMutex};
use parking_lot::{Condvar, Mutex};
use rand::Rng;
use std::time::Duration;
pub struct AutoPool<T: Send> {
config: AutoPoolConfig,
storage: Mutex<Vec<T>>,
condvar: Condvar,
}
impl<T: Send + 'static> AutoPool<T> {
pub fn new(items: impl IntoIterator<Item = T>) -> Self { Self::new_with_config(AutoPoolConfig::default(), items) }
pub fn new_with_config(config: AutoPoolConfig, items: impl IntoIterator<Item = T>) -> Self {
let objects = items.into_iter().collect();
Self {
config,
storage: Mutex::new(objects),
condvar: Condvar::new(),
}
}
pub fn get(&'_ self) -> Option<PoolObject<'_, T>> { self.get_with_timeout(self.config.wait_duration) }
#[cfg(feature = "async")]
pub async fn get_async(&'_ self) -> Option<PoolObject<'_, T>> {
if self.config.wait_duration.is_zero() {
return self.get();
}
let start_time = std::time::Instant::now();
while std::time::Instant::now() - start_time < self.config.wait_duration {
if let Some(obj) = self.get_with_timeout(self.config.lock_duration) {
return Some(obj);
}
smol::Timer::after(self.config.sleep_duration).await;
}
None
}
pub fn add(&self, item: T) {
self.storage.lock().push(item);
self.condvar.notify_one();
}
pub fn size(&self) -> usize { self.storage.lock().len() }
pub fn shrink_to_fit(&self) { self.storage.lock().shrink_to_fit(); }
fn get_with_timeout(&'_ self, timeout: Duration) -> Option<PoolObject<'_, T>> {
let mut locked_storage = self.storage.lock();
while locked_storage.is_empty() {
let wait_res = self.condvar.wait_for(&mut locked_storage, timeout);
if wait_res.timed_out() {
return None;
}
}
self.extract_object(locked_storage)
}
fn extract_object<R>(&'_ self, mut locked_storage: MutexGuard<R, Vec<T>>) -> Option<PoolObject<'_, T>>
where
R: RawMutex,
{
let inner = match self.config.pick_strategy {
PickStrategy::LIFO => locked_storage.pop(),
PickStrategy::RANDOM => match locked_storage.len() {
0 => None,
1 => locked_storage.pop(),
items_cnt => {
let index = rand::rng().next_u64() as usize % items_cnt;
locked_storage.swap(index, items_cnt - 1);
locked_storage.pop()
}
},
};
inner.map(|inner| PoolObject::new(inner, self))
}
}