use super::lifecycle;
use super::{
active_open_count, compatible_idle_count, compatible_waiting_demand, complete_open_effect,
drain_drop_returns, pending_open_count, reject, request_position, reserved_open_count,
schedule_open_effects, wake_waiters, EngineInner, PoolBackend, PoolError, PoolState,
PooledConn, Request,
};
use asupersync::sync::Notify;
use asupersync::{time, Cx};
use std::sync::atomic::Ordering;
use std::sync::{Arc, Weak};
use std::time::{Duration, Instant};
pub(super) fn drop_conn<B: PoolBackend>(
state: &mut PoolState<B::Conn>,
inner: &EngineInner<B>,
conn: PooledConn<B::Conn>,
) {
state.to_drop.push_back(conn);
ensure_min_connections(state, inner);
inner.bg.notify_one();
}
fn ensure_min_connections<B: PoolBackend>(state: &mut PoolState<B::Conn>, inner: &EngineInner<B>) {
if state.open {
let reserved = reserved_open_count(state);
if reserved < state.config.min {
schedule_open_effects(state, state.config.min - reserved);
}
inner.bg.notify_one();
}
}
fn check_connection<B: PoolBackend>(
state: &mut PoolState<B::Conn>,
inner: &EngineInner<B>,
request_id: u64,
conn: PooledConn<B::Conn>,
) {
if !inner.backend.connection_is_open(&conn.conn) {
drop_conn(state, inner, conn);
return;
}
let max_lifetime = state.config.max_lifetime_session_secs;
if max_lifetime > 0
&& conn.time_created.elapsed() > Duration::from_secs(u64::from(max_lifetime))
{
drop_conn(state, inner, conn);
return;
}
let ping_interval = state.config.ping_interval_secs;
let requires_ping = if ping_interval == 0 {
true
} else if ping_interval > 0 {
conn.time_returned.elapsed() > Duration::from_secs(ping_interval.unsigned_abs())
} else {
false
};
let Some(position) = request_position(state, request_id) else {
let mut orphan = Request {
id: 0,
cclass: None,
cclass_matches: true,
wants_new: false,
requires_ping: false,
bg_processing: false,
is_extra: false,
is_replacing: false,
in_progress: false,
completed: false,
waiting: false,
conn: Some(conn),
error: None,
};
reject(state, inner, &mut orphan);
return;
};
let request = &mut state.requests[position];
request.conn = Some(conn);
if requires_ping {
request.requires_ping = true;
add_request_for_bg(state, inner, request_id);
} else {
request.completed = true;
}
}
fn add_request_for_bg<B: PoolBackend>(
state: &mut PoolState<B::Conn>,
inner: &EngineInner<B>,
request_id: u64,
) {
if let Some(position) = request_position(state, request_id) {
let request = &mut state.requests[position];
request.bg_processing = true;
request.completed = false;
inner.bg.notify_one();
}
}
pub(super) fn fulfill<B: PoolBackend>(
state: &mut PoolState<B::Conn>,
inner: &EngineInner<B>,
request_id: u64,
) -> Result<bool, PoolError> {
if !state.open {
return Err(PoolError::Closed);
}
let Some(position) = request_position(state, request_id) else {
return Err(PoolError::Internal("request lost".to_string()));
};
{
let request = &mut state.requests[position];
if let Some(error) = request.error.take() {
return Err(PoolError::Backend(error));
}
if request.completed {
return Ok(true);
}
if request.bg_processing {
return Ok(false);
}
}
let wants_new = state.requests[position].wants_new;
let request_cclass = state.requests[position].cclass.clone();
let cclass_matches = state.requests[position].cclass_matches;
if !wants_new {
let mut ix = state.free_used.len();
while ix > 0 {
ix -= 1;
let matches = request_cclass.is_none()
|| state.free_used[ix].cclass.as_deref() == request_cclass.as_deref();
if !matches {
continue;
}
let conn = state.free_used.remove(ix);
check_connection(state, inner, request_id, conn);
let Some(position) = request_position(state, request_id) else {
return Err(PoolError::Internal("request lost".to_string()));
};
let request = &state.requests[position];
if request.completed || request.requires_ping {
return Ok(request.completed);
}
ix = ix.min(state.free_used.len());
}
}
if cclass_matches {
while let Some(conn) = state.free_new.pop() {
check_connection(state, inner, request_id, conn);
let Some(position) = request_position(state, request_id) else {
return Err(PoolError::Internal("request lost".to_string()));
};
let request = &state.requests[position];
if request.completed || request.requires_ping {
return Ok(request.completed);
}
}
}
if let Some(position) = request_position(state, request_id) {
state.requests[position].requires_ping = false;
}
if reserved_open_count(state) >= state.config.max {
if let Some(victim) = state.free_new.pop() {
if let Some(position) = request_position(state, request_id) {
state.requests[position].is_replacing = true;
}
state.to_drop.push_back(victim);
add_request_for_bg(state, inner, request_id);
return Ok(false);
} else if let Some(victim) = state.free_used.pop() {
if let Some(position) = request_position(state, request_id) {
state.requests[position].is_replacing = true;
}
state.to_drop.push_back(victim);
add_request_for_bg(state, inner, request_id);
return Ok(false);
} else if state.force_get {
if let Some(position) = request_position(state, request_id) {
state.requests[position].is_extra = true;
}
add_request_for_bg(state, inner, request_id);
return Ok(false);
} else if state.config.getmode == super::POOL_GETMODE_NOWAIT {
return Err(PoolError::NoConnectionAvailable);
}
} else if cclass_matches {
let remaining_capacity = state.config.max.saturating_sub(reserved_open_count(state));
let pending = pending_open_count(state);
let idle =
compatible_idle_count(state, wants_new, request_cclass.as_deref(), cclass_matches);
let demand = compatible_waiting_demand(state);
let supply = pending.saturating_add(idle);
let shortfall = demand.saturating_sub(supply);
let desired = if pending == 0 {
state.config.increment.max(shortfall)
} else {
shortfall
};
if desired > 0 {
schedule_open_effects(state, desired.min(remaining_capacity));
}
}
add_request_for_bg(state, inner, request_id);
Ok(false)
}
fn get_next_request<C>(state: &mut PoolState<C>) -> Option<u64> {
let id = peek_next_request(state)?;
if let Some(position) = request_position(state, id) {
let request = &mut state.requests[position];
request.in_progress = request.waiting;
}
Some(id)
}
fn peek_next_request<C>(state: &PoolState<C>) -> Option<u64> {
for request in &state.requests {
if request.completed || request.error.is_some() || request.in_progress {
continue;
}
if !request.bg_processing {
continue;
}
if !request.waiting
|| request.requires_ping
|| request.is_replacing
|| request.is_extra
|| (!request.cclass_matches && reserved_open_count(state) < state.config.max)
{
return Some(request.id);
}
break;
}
None
}
fn post_process_request<B: PoolBackend>(
state: &mut PoolState<B::Conn>,
inner: &EngineInner<B>,
request_id: u64,
) {
let Some(position) = request_position(state, request_id) else {
return;
};
let request = &mut state.requests[position];
request.in_progress = false;
request.bg_processing = false;
if request.conn.is_some() {
request.completed = true;
let request = &mut state.requests[position];
if !request.waiting {
let mut request = state.requests.remove(position);
reject(state, inner, &mut request);
}
} else {
if request.requires_ping {
ensure_min_connections(state, inner);
}
let request = &mut state.requests[position];
if !request.waiting {
state.requests.remove(position);
}
}
wake_waiters(inner);
}
fn post_create_conn<B: PoolBackend>(
state: &mut PoolState<B::Conn>,
inner: &EngineInner<B>,
conn_id: u64,
created: Result<PooledConn<B::Conn>, String>,
) {
complete_open_effect(state, conn_id);
let conn = match created {
Ok(conn) => conn,
Err(error) => {
if state.open {
if let Some(request) = state.requests.iter_mut().find(|request| {
request.bg_processing
&& request.waiting
&& !request.requires_ping
&& !request.is_extra
&& !request.is_replacing
&& request.cclass_matches
&& request.conn.is_none()
}) {
request.bg_processing = false;
request.error = Some(error);
wake_waiters(inner);
}
}
return;
}
};
debug_assert_eq!(conn.id, conn_id);
if !state.open {
state.to_drop.push_back(conn);
inner.bg.notify_one();
return;
}
let mut conn = Some(conn);
let max = state.config.max;
let open_count = active_open_count(state).saturating_add(1);
for request in &mut state.requests {
if request.in_progress
|| request.conn.is_some()
|| !request.waiting
|| request.completed
|| request.error.is_some()
{
continue;
}
let candidate = conn.as_ref().expect("connection still available");
if request.cclass.is_none() || request.cclass.as_deref() == candidate.cclass.as_deref() {
request.conn = conn.take();
request.completed = true;
request.bg_processing = false;
wake_waiters(inner);
break;
} else if !request.cclass_matches && open_count >= max {
request.conn = conn.take();
request.is_replacing = true;
break;
}
}
if let Some(conn) = conn {
state.free_new.push(conn);
wake_waiters(inner);
}
}
pub(super) fn return_connection_helper<B: PoolBackend>(
state: &mut PoolState<B::Conn>,
inner: &EngineInner<B>,
mut conn: PooledConn<B::Conn>,
mut is_open: bool,
) {
if !is_open {
drop_conn(state, inner, conn);
return;
}
if conn.is_pool_extra {
conn.is_pool_extra = false;
let count_with_returned = active_open_count(state).saturating_add(u32::from(is_open));
if is_open && count_with_returned > state.config.max {
drop_conn(state, inner, conn);
return;
}
}
let mut returned = Some(conn);
if is_open {
let conn = returned.as_mut().expect("connection still available");
conn.time_returned = Instant::now();
let max_lifetime = state.config.max_lifetime_session_secs;
if max_lifetime != 0
&& conn.time_created.elapsed() > Duration::from_secs(u64::from(max_lifetime))
{
let conn = returned.take().expect("connection still available");
drop_conn(state, inner, conn);
is_open = false;
}
}
if is_open {
let mut conn = returned;
for request in &mut state.requests {
if request.in_progress
|| request.wants_new
|| request.conn.is_some()
|| !request.waiting
|| request.completed
|| request.error.is_some()
{
continue;
}
let candidate = conn.as_ref().expect("connection still available");
let matches = request.cclass.is_none()
|| request.cclass.as_deref() == candidate.cclass.as_deref();
if matches {
request.conn = conn.take();
request.completed = true;
request.bg_processing = false;
wake_waiters(inner);
return;
}
}
if let Some(conn) = conn {
state.free_used.push(conn);
}
}
}
fn sweep_idle_timeout<B: PoolBackend>(state: &mut PoolState<B::Conn>, inner: &EngineInner<B>) {
let timeout_secs = state.config.timeout_secs;
if timeout_secs == 0 {
return;
}
let limit = Duration::from_secs(u64::from(timeout_secs));
for list in ["new", "used"] {
loop {
if active_open_count(state) <= state.config.min {
return;
}
let conns = if list == "new" {
&mut state.free_new
} else {
&mut state.free_used
};
let Some(first) = conns.first() else {
break;
};
if first.time_returned.elapsed() < limit {
break;
}
let conn = conns.remove(0);
drop_conn(state, inner, conn);
}
}
}
pub(super) async fn reaper_main<B: PoolBackend>(weak: Weak<EngineInner<B>>, bg: Arc<Notify>) {
let cx = Cx::current();
let mut current_request: Option<u64> = None;
loop {
let Some(inner) = weak.upgrade() else {
return;
};
if let Some(cx) = cx.as_ref() {
if cx.checkpoint().is_err() {
return;
}
}
let stopping = inner.reaper_stop.load(Ordering::SeqCst);
let mut open;
{
let Ok(mut state) = inner.state.lock() else {
return;
};
if drain_drop_returns(&mut state, &inner).is_err() {
return;
}
open = state.open;
if current_request.is_none() && open {
current_request = get_next_request(&mut state);
}
}
if let Some(request_id) = current_request.take() {
if open {
process_request(&inner, request_id);
let Ok(mut state) = inner.state.lock() else {
return;
};
post_process_request(&mut state, &inner, request_id);
current_request = get_next_request(&mut state);
continue;
}
}
let (open_effect, cclass) = {
let Ok(mut state) = inner.state.lock() else {
return;
};
open = state.open;
let effect = if open {
if let Some(effect) = state.open_effects.pop_front() {
state.in_flight_open_effects.push_back(effect);
}
state.in_flight_open_effects.back().copied()
} else {
None
};
(effect, state.config.creation_cclass.clone())
};
if let Some(lifecycle::PoolEffect::Open { slot: conn_id, .. }) = open_effect {
let created = inner
.backend
.create_connection(conn_id, cclass.as_deref())
.map(|conn| PooledConn {
id: conn_id,
conn,
cclass,
time_created: Instant::now(),
time_returned: Instant::now(),
is_pool_extra: false,
ever_acquired: false,
});
let Ok(mut state) = inner.state.lock() else {
return;
};
post_create_conn(&mut state, &inner, conn_id, created);
continue;
}
let next_drop = {
let Ok(mut state) = inner.state.lock() else {
return;
};
state.to_drop.pop_back()
};
if let Some(conn) = next_drop {
inner.backend.close_connection(conn.id, conn.conn);
continue;
}
let timeout_armed;
{
let Ok(mut state) = inner.state.lock() else {
return;
};
sweep_idle_timeout(&mut state, &inner);
if (!state.open || stopping) && state.to_drop.is_empty() {
return;
}
let has_work = !state.open_effects.is_empty()
|| !state.to_drop.is_empty()
|| peek_next_request(&state).is_some();
if has_work {
continue;
}
timeout_armed = state.config.timeout_secs > 0
&& active_open_count(&state) > state.config.min
&& (!state.free_new.is_empty() || !state.free_used.is_empty());
}
drop(inner);
let notified = bg.notified();
if timeout_armed {
let _ = time::timeout(time::wall_now(), Duration::from_secs(1), notified).await;
} else {
notified.await;
}
}
}
fn process_request<B: PoolBackend>(inner: &Arc<EngineInner<B>>, request_id: u64) {
enum Work<C> {
Ping {
conn: PooledConn<C>,
ping_timeout_ms: u32,
},
Create {
conn_id: u64,
cclass: Option<String>,
is_extra: bool,
replaced: Option<PooledConn<C>>,
},
Nothing,
}
let work = {
let Ok(mut state) = inner.state.lock() else {
return;
};
let Some(position) = request_position(&state, request_id) else {
return;
};
let ping_timeout_ms = state.config.ping_timeout_ms;
let request = &mut state.requests[position];
if request.requires_ping {
match request.conn.take() {
Some(conn) => Work::Ping {
conn,
ping_timeout_ms,
},
None => Work::Nothing,
}
} else if request.is_replacing || request.is_extra || !request.cclass_matches {
let cclass = request.cclass.clone();
let is_extra = request.is_extra;
let replaced = request.conn.take();
let conn_id = state.next_conn_id;
state.next_conn_id += 1;
Work::Create {
conn_id,
cclass,
is_extra,
replaced,
}
} else {
Work::Nothing
}
};
match work {
Work::Ping {
conn,
ping_timeout_ms,
} => {
let healthy = inner.backend.ping_connection(&conn.conn, ping_timeout_ms);
if healthy {
let Ok(mut state) = inner.state.lock() else {
return;
};
if let Some(position) = request_position(&state, request_id) {
state.requests[position].conn = Some(conn);
}
} else {
inner.backend.close_connection(conn.id, conn.conn);
}
}
Work::Create {
conn_id,
cclass,
is_extra,
replaced,
} => {
let result = inner.backend.create_connection(conn_id, cclass.as_deref());
let Ok(mut state) = inner.state.lock() else {
return;
};
if let Some(old) = replaced {
drop_conn(&mut state, inner, old);
}
let Some(position) = request_position(&state, request_id) else {
if let Ok(conn) = result {
state.to_drop.push_back(PooledConn {
id: conn_id,
conn,
cclass,
time_created: Instant::now(),
time_returned: Instant::now(),
is_pool_extra: false,
ever_acquired: false,
});
inner.bg.notify_one();
}
return;
};
match result {
Ok(conn) => {
state.requests[position].conn = Some(PooledConn {
id: conn_id,
conn,
cclass,
time_created: Instant::now(),
time_returned: Instant::now(),
is_pool_extra: is_extra,
ever_acquired: false,
});
}
Err(error) => {
state.requests[position].error = Some(error);
}
}
}
Work::Nothing => {}
}
}