use super::connection::{Floating, Idle, Live};
use crate::connection::ConnectOptions;
use crate::connection::Connection;
use crate::database::Database;
use crate::error::Error;
use crate::pool::{deadline_as_timeout, PoolOptions};
use crossbeam_queue::ArrayQueue;
use std::cmp;
use std::mem;
use std::ptr;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use mco::co;
use mco::std::sync::Blocker;
use crate::pool::semaphore::{BoxSemaphore, PermitGuard};
const WAKE_ALL_PERMITS: usize = usize::MAX / 2;
pub struct SharedPool<DB: Database> {
pub(super) connect_options: <DB::Connection as Connection>::Options,
pub(super) idle_conns: ArrayQueue<Idle<DB>>,
pub(super) semaphore: BoxSemaphore,
pub(super) size: AtomicU32,
is_closed: AtomicBool,
pub(super) options: PoolOptions<DB>,
}
impl<DB: Database> SharedPool<DB> {
pub(super) fn new_arc(
options: PoolOptions<DB>,
connect_options: <DB::Connection as Connection>::Options,
) -> Arc<Self> {
let capacity = options.max_connections as usize;
let _ = capacity
.checked_add(WAKE_ALL_PERMITS)
.expect("max_connections exceeds max capacity of the pool");
let pool = Self {
connect_options,
idle_conns: ArrayQueue::new(capacity),
semaphore: BoxSemaphore::new(capacity),
size: AtomicU32::new(0),
is_closed: AtomicBool::new(false),
options,
};
let pool = Arc::new(pool);
spawn_reaper(&pool);
pool
}
pub(super) fn size(&self) -> u32 {
self.size.load(Ordering::Acquire)
}
pub(super) fn num_idle(&self) -> usize {
self.idle_conns.len()
}
pub(super) fn is_closed(&self) -> bool {
self.is_closed.load(Ordering::Acquire)
}
pub(super) fn close(&self) {
let already_closed = self.is_closed.swap(true, Ordering::AcqRel);
if !already_closed {
self.semaphore.release_left(WAKE_ALL_PERMITS);
}
while let Some(idle) = self.idle_conns.pop() {
let _ = idle.live.float(self).close();
}
}
#[inline]
pub(super) fn try_acquire(&self) -> Option<Floating<'_, Idle<DB>>> {
if self.is_closed() {
return None;
}
let permit = self.semaphore.try_acquire();
if let Some(permit) = permit {
return self.pop_idle(permit).ok();
} else {
None
}
}
fn pop_idle<'a>(
&'a self,
permit: PermitGuard<'a>,
) -> Result<Floating<'a, Idle<DB>>, PermitGuard<'a>> {
if let Some(idle) = self.idle_conns.pop() {
Ok(Floating::from_idle(idle, self, permit))
} else {
Err(permit)
}
}
pub(super) fn release(&self, mut floating: Floating<'_, Live<DB>>) {
if let Some(test) = &self.options.after_release {
if !test(&mut floating.raw) {
return;
}
}
let Floating { inner: idle, guard } = floating.into_idle();
if !self.idle_conns.push(idle).is_ok() {
panic!("BUG: connection queue overflow in release()");
}
guard.release_permit();
}
pub(super) fn try_increment_size<'a>(
&'a self,
permit: PermitGuard<'a>,
) -> Result<DecrementSizeGuard<'a>, PermitGuard<'a>> {
match self
.size
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |size| {
size.checked_add(1)
.filter(|size| size <= &self.options.max_connections)
}) {
Ok(_) => Ok(DecrementSizeGuard::from_permit(self, permit)),
Err(_) => Err(permit),
}
}
#[allow(clippy::needless_lifetimes)]
pub(super) fn acquire<'s>(&'s self) -> Result<Floating<'s, Live<DB>>, Error> {
if self.is_closed() {
return Err(Error::PoolClosed);
}
let deadline = Instant::now() + self.options.connect_timeout;
loop {
let permit = self.semaphore.acquire();
if self.is_closed() {
return Err(Error::PoolClosed);
}
let guard = match self.pop_idle(permit) {
Ok(conn) => match check_conn(conn, &self.options) {
Ok(live) => return Ok(live),
Err(guard) => guard,
},
Err(permit) => if let Ok(guard) = self.try_increment_size(permit) {
guard
} else {
log::debug!("woke but was unable to acquire idle connection or open new one; retrying");
continue;
}
};
return self.connection(deadline, guard);
}
}
pub(super) fn connection<'s>(
&'s self,
deadline: Instant,
guard: DecrementSizeGuard<'s>,
) -> Result<Floating<'s, Live<DB>>, Error> {
if self.is_closed() {
return Err(Error::PoolClosed);
}
let mut backoff = Duration::from_millis(10);
let max_backoff = deadline_as_timeout::<DB>(deadline)? / 5;
loop {
let timeout = deadline_as_timeout::<DB>(deadline)?;
match self.connect_options.connect(timeout) {
Ok(mut raw) => {
if let Some(callback) = &self.options.after_connect {
callback(&mut raw)?;
}
return Ok(Floating::new_live(raw, guard));
}
Err(Error::Io(e)) if e.kind() == std::io::ErrorKind::ConnectionRefused => (),
Err(Error::Database(error)) if error.code().as_deref() == Some("57P03") => (),
Err(e) => return Err(e),
_ => {
return Err(Error::PoolTimedOut);
}
}
mco::coroutine::sleep(backoff);
backoff = cmp::min(backoff * 2, max_backoff);
}
}
}
fn is_beyond_lifetime<DB: Database>(live: &Live<DB>, options: &PoolOptions<DB>) -> bool {
options
.max_lifetime
.map_or(false, |max| live.created.elapsed() > max)
}
fn is_beyond_idle<DB: Database>(idle: &Idle<DB>, options: &PoolOptions<DB>) -> bool {
options
.idle_timeout
.map_or(false, |timeout| idle.since.elapsed() > timeout)
}
fn check_conn<'s: 'p, 'p, DB: Database>(
mut conn: Floating<'s, Idle<DB>>,
options: &'p PoolOptions<DB>,
) -> Result<Floating<'s, Live<DB>>, DecrementSizeGuard<'s>> {
if is_beyond_lifetime(&conn, options) {
return Err(conn.close());
} else if options.test_before_acquire {
if let Err(e) = conn.ping() {
log::info!("ping on idle connection returned error: {}", e);
return Err(conn.close());
}
} else if let Some(test) = &options.before_acquire {
match test(&mut conn.live.raw) {
Ok(false) => {
return Err(conn.close());
}
Err(error) => {
log::info!("in `before_acquire`: {}", error);
return Err(conn.close());
}
Ok(true) => {}
}
}
Ok(conn.into_live())
}
fn spawn_reaper<DB: Database>(pool: &Arc<SharedPool<DB>>) {
let period = match (pool.options.max_lifetime, pool.options.idle_timeout) {
(Some(it), None) | (None, Some(it)) => it,
(Some(a), Some(b)) => cmp::min(a, b),
(None, None) => return,
};
let pool = Arc::clone(&pool);
co!(move || {
while !pool.is_closed() {
if !pool.idle_conns.is_empty() {
do_reap(&pool);
}
mco::coroutine::sleep(period);
}
});
}
fn do_reap<DB: Database>(pool: &SharedPool<DB>) {
let max_reaped = pool.size().saturating_sub(pool.options.min_connections);
let (reap, keep) = (0..max_reaped)
.filter_map(|_| pool.try_acquire())
.partition::<Vec<_>, _>(|conn| {
is_beyond_idle(conn, &pool.options) || is_beyond_lifetime(conn, &pool.options)
});
for conn in keep {
pool.release(conn.into_live());
}
for conn in reap {
let _ = conn.close();
}
}
pub(in crate::pool) struct DecrementSizeGuard<'a> {
size: &'a AtomicU32,
semaphore: &'a BoxSemaphore,
dropped: bool,
}
impl<'a> DecrementSizeGuard<'a> {
pub fn new_permit<DB: Database>(pool: &'a SharedPool<DB>) -> Self {
Self {
size: &pool.size,
semaphore: &pool.semaphore,
dropped: false,
}
}
pub fn from_permit<DB: Database>(
pool: &'a SharedPool<DB>,
mut permit: PermitGuard<'a>,
) -> Self {
permit.release();
Self::new_permit(pool)
}
pub fn same_pool<DB: Database>(&self, pool: &'a SharedPool<DB>) -> bool {
ptr::eq(self.size, &pool.size)
}
fn release_permit(self) {
self.semaphore.release();
self.cancel();
}
pub fn cancel(self) {
mem::forget(self);
}
}
impl Drop for DecrementSizeGuard<'_> {
fn drop(&mut self) {
assert!(!self.dropped, "double-dropped!");
self.dropped = true;
self.size.fetch_sub(1, Ordering::SeqCst);
self.semaphore.release();
}
}