use crate::connection::Connection;
use crate::raw::RawConnection;
use async_std::sync::Arc;
use crossbeam_queue::{ArrayQueue, SegQueue};
use futures_channel::oneshot;
use std::sync::atomic::{AtomicU32, Ordering};
use std::time::Instant;
#[derive(Clone)]
pub struct Pool(Arc<SharedPool>);
impl Pool {
pub fn builder() -> Builder {
Builder::new()
}
pub async fn acquire(&self) -> crate::Result<Connection> {
let raw = self.0.acquire().await?;
Ok(Connection {
raw: Some(raw),
pool: Some(self.0.clone()),
})
}
pub async fn get(&self, key: &str) -> crate::Result<Option<String>> {
self.acquire().await?.get(key).await
}
}
pub struct Builder {
max_size: u32,
}
impl Builder {
pub fn new() -> Self {
Self { max_size: 1 }
}
pub fn max_size(mut self, size: u32) -> Self {
self.max_size = size;
self
}
pub fn build(self, url: &str) -> Pool {
Pool(Arc::new(SharedPool {
max_size: self.max_size,
url: url.to_owned(),
size: AtomicU32::new(0),
idle: ArrayQueue::new(self.max_size as usize),
waiters: SegQueue::new(),
}))
}
}
pub(crate) struct SharedPool {
max_size: u32,
size: AtomicU32,
url: String,
idle: ArrayQueue<Idle>,
waiters: SegQueue<oneshot::Sender<RawConnection>>,
}
impl SharedPool {
#[inline]
fn try_acquire(&self) -> Option<RawConnection> {
if let Ok(idle) = self.idle.pop() {
return Some(idle.raw);
}
None
}
async fn acquire(&self) -> crate::Result<RawConnection> {
if let Some(live) = self.try_acquire() {
return Ok(live);
}
loop {
let size = self.size.load(Ordering::Acquire);
if size >= self.max_size {
let (sender, receiver) = oneshot::channel();
self.waiters.push(sender);
return Ok(receiver
.await
.expect("waiter dropped without dropping pool"));
}
if self.size.compare_and_swap(size, size + 1, Ordering::AcqRel) == size {
return RawConnection::open(&self.url).await;
}
}
}
pub(crate) fn release(&self, mut raw: RawConnection) {
while let Ok(waiter) = self.waiters.pop() {
raw = match waiter.send(raw) {
Ok(()) => {
return;
}
Err(raw) => raw,
};
}
let _ = self.idle.push(Idle {
raw,
since: Instant::now(),
});
}
}
struct Idle {
raw: RawConnection,
#[allow(unused)]
since: Instant,
}