pub mod replica;
pub use replica::{ReplicaPool, ReplicaStrategy};
pub mod sharding;
pub use sharding::{ModuloShardChooser, QueryHints, ShardChooser, ShardedPool, ShardedPoolStats};
use std::collections::VecDeque;
use std::future::Future;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, Weak};
use std::time::Duration;
use asupersync::{
Budget, CancelReason, Cx, Outcome, Time,
combinator::{Either, Select},
runtime::RuntimeBuilder,
sync::{Notify, OnceCell},
time::TimerDriverHandle,
};
use sqlmodel_core::error::{PoolError, PoolErrorKind};
use sqlmodel_core::{Connection, Error};
#[derive(Debug, Clone)]
pub struct PoolConfig {
pub min_connections: usize,
pub max_connections: usize,
pub idle_timeout_ms: u64,
pub acquire_timeout_ms: u64,
pub max_lifetime_ms: u64,
pub test_on_checkout: bool,
}
impl Default for PoolConfig {
fn default() -> Self {
Self {
min_connections: 1,
max_connections: 10,
idle_timeout_ms: 600_000, acquire_timeout_ms: 30_000, max_lifetime_ms: 1_800_000, test_on_checkout: true,
}
}
}
impl PoolConfig {
#[must_use]
pub fn new(max_connections: usize) -> Self {
Self {
max_connections,
..Default::default()
}
}
#[must_use]
pub fn min_connections(mut self, n: usize) -> Self {
self.min_connections = n;
self
}
#[must_use]
pub fn idle_timeout(mut self, ms: u64) -> Self {
self.idle_timeout_ms = ms;
self
}
#[must_use]
pub fn acquire_timeout(mut self, ms: u64) -> Self {
self.acquire_timeout_ms = ms;
self
}
#[must_use]
pub fn max_lifetime(mut self, ms: u64) -> Self {
self.max_lifetime_ms = ms;
self
}
#[must_use]
pub fn test_on_checkout(mut self, enabled: bool) -> Self {
self.test_on_checkout = enabled;
self
}
}
#[derive(Debug, Clone, Default)]
pub struct PoolStats {
pub total_connections: usize,
pub idle_connections: usize,
pub active_connections: usize,
pub pending_requests: usize,
pub connections_created: u64,
pub connections_closed: u64,
pub acquires: u64,
pub timeouts: u64,
}
struct ConnectionMeta<C> {
conn: C,
created_at: Time,
last_used: Time,
clock: TimerDriverHandle,
}
impl<C> ConnectionMeta<C> {
fn new(conn: C, clock: TimerDriverHandle) -> Self {
let now = clock.now();
Self {
conn,
created_at: now,
last_used: now,
clock,
}
}
fn touch(&mut self) {
self.last_used = self.clock.now();
}
fn age(&self) -> Duration {
Duration::from_nanos(self.clock.now().duration_since(self.created_at))
}
fn idle_time(&self) -> Duration {
Duration::from_nanos(self.clock.now().duration_since(self.last_used))
}
}
struct PoolInner<C: Connection> {
config: PoolConfig,
idle: VecDeque<ConnectionMeta<C>>,
active_count: usize,
total_count: usize,
waiter_count: usize,
closed: bool,
}
impl<C: Connection> PoolInner<C> {
fn new(config: PoolConfig) -> Self {
Self {
config,
idle: VecDeque::new(),
active_count: 0,
total_count: 0,
waiter_count: 0,
closed: false,
}
}
fn can_create_new(&self) -> bool {
!self.closed && self.total_count < self.config.max_connections
}
fn stats(&self) -> PoolStats {
PoolStats {
total_connections: self.total_count,
idle_connections: self.idle.len(),
active_connections: self.active_count,
pending_requests: self.waiter_count,
..Default::default()
}
}
}
struct PoolShared<C: Connection> {
inner: Mutex<PoolInner<C>>,
conn_available: Notify,
active_drained: OnceCell<()>,
retirement_failure: Mutex<Option<Arc<str>>>,
connections_created: AtomicU64,
connections_closed: AtomicU64,
acquires: AtomicU64,
timeouts: AtomicU64,
clock: TimerDriverHandle,
}
impl<C: Connection> PoolShared<C> {
fn new(config: PoolConfig, clock: TimerDriverHandle) -> Self {
Self {
inner: Mutex::new(PoolInner::new(config)),
conn_available: Notify::new(),
active_drained: OnceCell::new(),
retirement_failure: Mutex::new(None),
connections_created: AtomicU64::new(0),
connections_closed: AtomicU64::new(0),
acquires: AtomicU64::new(0),
timeouts: AtomicU64::new(0),
clock,
}
}
fn lock_or_recover(&self) -> std::sync::MutexGuard<'_, PoolInner<C>> {
self.inner.lock().unwrap_or_else(|poisoned| {
tracing::error!(
"Pool mutex poisoned; recovering for read-only access. \
A thread panicked while holding the lock."
);
poisoned.into_inner()
})
}
#[allow(clippy::result_large_err)] fn lock_or_error(
&self,
operation: &'static str,
) -> Result<std::sync::MutexGuard<'_, PoolInner<C>>, Error> {
self.inner
.lock()
.map_err(|_| Error::Pool(PoolError::poisoned(operation)))
}
fn release_active_slot(&self, operation: &'static str) {
let mut accounting_underflow = false;
let (drained, notify_open_waiter) = match self.inner.lock() {
Ok(mut inner) => {
if inner.active_count == 0 || inner.total_count == 0 {
tracing::error!(
operation,
active_count = inner.active_count,
total_count = inner.total_count,
"attempted to release an unaccounted pool slot"
);
accounting_underflow = true;
}
inner.active_count = inner.active_count.saturating_sub(1);
inner.total_count = inner.total_count.saturating_sub(1);
if accounting_underflow && inner.closed && inner.active_count == 0 {
inner.total_count = 0;
}
(inner.closed && inner.active_count == 0, !inner.closed)
}
Err(poisoned) => {
tracing::error!(
operation,
"Pool mutex poisoned while releasing an active slot; \
recovering to prevent stranded drain accounting"
);
let error = Error::Pool(PoolError::poisoned(operation));
self.record_retirement_failure(operation, &error);
let mut inner = poisoned.into_inner();
if inner.active_count == 0 || inner.total_count == 0 {
accounting_underflow = true;
}
inner.active_count = inner.active_count.saturating_sub(1);
inner.total_count = inner.total_count.saturating_sub(1);
if accounting_underflow && inner.closed && inner.active_count == 0 {
inner.total_count = 0;
}
(inner.closed && inner.active_count == 0, !inner.closed)
}
};
if accounting_underflow {
let error = Error::Custom(format!(
"pool accounting underflow while releasing active slot during {operation}"
));
self.record_retirement_failure(operation, &error);
}
if notify_open_waiter {
self.conn_available.notify_one();
}
if drained {
let _ = self.active_drained.set(());
}
}
fn record_retirement_failure(&self, context: &'static str, error: &Error) {
let message: Arc<str> = format!("{context}: {error}").into();
let mut failure = self.retirement_failure.lock().unwrap_or_else(|poisoned| {
tracing::error!(
context,
"Pool retirement-failure mutex poisoned; recovering"
);
poisoned.into_inner()
});
if failure.is_none() {
*failure = Some(message);
}
}
fn retirement_failure_error(&self) -> Option<Error> {
let failure = self
.retirement_failure
.lock()
.unwrap_or_else(|poisoned| {
tracing::error!("Pool retirement-failure mutex poisoned; recovering");
poisoned.into_inner()
})
.clone();
failure.map(|message| Error::Custom(format!("pool retirement failed: {message}")))
}
}
#[allow(clippy::result_large_err)] fn close_connection_blocking<C: Connection>(conn: C, context: &'static str) -> Result<(), Error> {
let runtime = match RuntimeBuilder::current_thread().build() {
Ok(runtime) => runtime,
Err(error) => {
tracing::warn!(
context,
error = %error,
"failed to build runtime while closing pooled connection"
);
drop(conn);
return Err(Error::Custom(format!(
"failed to build runtime while closing pooled connection: {error}"
)));
}
};
let cx = runtime.request_cx_with_budget(Budget::INFINITE);
let result = runtime.block_on(async { conn.close_for_pool(&cx).await });
if let Err(error) = &result {
tracing::warn!(
context,
error = %error,
"failed to close pooled connection explicitly"
);
}
result
}
struct ActiveSlotGuard<C: Connection> {
pool: Arc<PoolShared<C>>,
armed: bool,
}
impl<C: Connection> ActiveSlotGuard<C> {
fn new(pool: Arc<PoolShared<C>>) -> Self {
Self { pool, armed: true }
}
fn disarm(&mut self) {
self.armed = false;
}
fn release(&mut self, operation: &'static str) {
if self.armed {
self.armed = false;
self.pool.release_active_slot(operation);
}
}
}
impl<C: Connection> Drop for ActiveSlotGuard<C> {
fn drop(&mut self) {
self.release("connection factory future drop");
}
}
struct ActiveConnectionGuard<C: Connection> {
pool: Arc<PoolShared<C>>,
meta: Option<ConnectionMeta<C>>,
armed: bool,
}
enum RetirementOutcome {
Closed,
Cancelled(CancelReason),
Failed(Error),
}
impl<C: Connection> ActiveConnectionGuard<C> {
fn new(pool: Arc<PoolShared<C>>, meta: ConnectionMeta<C>) -> Self {
Self {
pool,
meta: Some(meta),
armed: true,
}
}
fn connection(&self) -> &C {
&self
.meta
.as_ref()
.expect("active connection guard already consumed")
.conn
}
fn into_meta(mut self) -> ConnectionMeta<C> {
self.armed = false;
self.meta
.take()
.expect("active connection guard already consumed")
}
fn release(&mut self, operation: &'static str) {
if self.armed {
self.armed = false;
self.pool.connections_closed.fetch_add(1, Ordering::Relaxed);
self.pool.release_active_slot(operation);
}
}
async fn close(mut self, cx: &Cx, context: &'static str) -> RetirementOutcome {
let meta = self
.meta
.take()
.expect("active connection guard already consumed");
let close_future = Box::pin(meta.conn.close_for_pool(cx));
let cancellation_latch = OnceCell::<()>::new();
let cancellation_future = Box::pin(cancellation_latch.wait(cx));
let selected = Select::new(close_future, cancellation_future).await;
let outcome = match selected {
Ok(Either::Left(result)) => match result {
Ok(()) => RetirementOutcome::Closed,
Err(error) => {
tracing::warn!(
context,
error = %error,
"failed to close pooled connection explicitly"
);
self.pool.record_retirement_failure(context, &error);
RetirementOutcome::Failed(error)
}
},
Ok(Either::Right(_)) => RetirementOutcome::Cancelled(
cx.cancel_reason()
.unwrap_or_else(|| CancelReason::user("pool retirement cancelled")),
),
Err(error) => {
tracing::error!(
context,
error = %error,
"fresh pool retirement select completed inconsistently"
);
let error = Error::Custom(format!(
"pool retirement select completed inconsistently: {error}"
));
self.pool.record_retirement_failure(context, &error);
RetirementOutcome::Failed(error)
}
};
self.release(context);
outcome
}
}
impl<C: Connection> Drop for ActiveConnectionGuard<C> {
fn drop(&mut self) {
if !self.armed {
return;
}
drop(self.meta.take());
self.release("checkout validation future drop");
}
}
pub struct Pool<C: Connection> {
shared: Arc<PoolShared<C>>,
}
impl<C: Connection> Pool<C> {
#[must_use]
pub fn new(config: PoolConfig) -> Self {
Self::with_timer_driver(config, TimerDriverHandle::with_wall_clock())
}
#[must_use]
pub fn with_timer_driver(config: PoolConfig, clock: TimerDriverHandle) -> Self {
Self {
shared: Arc::new(PoolShared::new(config, clock)),
}
}
#[must_use]
pub fn config(&self) -> PoolConfig {
let inner = self.shared.lock_or_recover();
inner.config.clone()
}
#[must_use]
pub fn stats(&self) -> PoolStats {
let inner = self.shared.lock_or_recover();
let mut stats = inner.stats();
stats.connections_created = self.shared.connections_created.load(Ordering::Relaxed);
stats.connections_closed = self.shared.connections_closed.load(Ordering::Relaxed);
stats.acquires = self.shared.acquires.load(Ordering::Relaxed);
stats.timeouts = self.shared.timeouts.load(Ordering::Relaxed);
stats
}
#[must_use]
pub fn at_capacity(&self) -> bool {
let inner = self.shared.lock_or_recover();
inner.total_count >= inner.config.max_connections
}
#[must_use]
pub fn is_closed(&self) -> bool {
let inner = self.shared.lock_or_recover();
inner.closed
}
pub async fn acquire<F, Fut>(&self, cx: &Cx, factory: F) -> Outcome<PooledConnection<C>, Error>
where
F: Fn() -> Fut,
Fut: Future<Output = Outcome<C, Error>>,
{
let clock = self.shared.clock.clone();
let mut deadline = clock.now() + Duration::from_millis(self.config().acquire_timeout_ms);
let mut budget_limited = false;
if let Some(budget_deadline) = cx.budget().deadline
&& budget_deadline < deadline
{
deadline = budget_deadline;
budget_limited = true;
}
let test_on_checkout = self.config().test_on_checkout;
let max_lifetime = Duration::from_millis(self.config().max_lifetime_ms);
let idle_timeout = Duration::from_millis(self.config().idle_timeout_ms);
loop {
if cx.is_cancel_requested() {
return Outcome::Cancelled(CancelReason::user("pool acquire cancelled"));
}
if clock.now() >= deadline {
self.shared.timeouts.fetch_add(1, Ordering::Relaxed);
let message = if budget_limited {
"acquire timeout: budget deadline reached before a connection was available"
} else {
"acquire timeout: no connections available"
};
return Outcome::Err(Error::Pool(PoolError {
kind: PoolErrorKind::Timeout,
message: message.to_string(),
source: None,
}));
}
let (action, retired) = {
let mut inner = match self.shared.lock_or_error("acquire") {
Ok(guard) => guard,
Err(e) => return Outcome::Err(e),
};
let mut retired = Vec::with_capacity(inner.idle.len());
let action = if inner.closed {
AcquireAction::PoolClosed
} else {
let mut found_conn = None;
while let Some(mut meta) = inner.idle.pop_front() {
if meta.age() > max_lifetime {
inner.active_count += 1;
retired
.push(ActiveConnectionGuard::new(Arc::clone(&self.shared), meta));
continue;
}
if meta.idle_time() > idle_timeout {
inner.active_count += 1;
retired
.push(ActiveConnectionGuard::new(Arc::clone(&self.shared), meta));
continue;
}
if !retired.is_empty() {
inner.idle.push_front(meta);
break;
}
meta.touch();
inner.active_count += 1;
found_conn = Some(meta);
break;
}
if !retired.is_empty() {
AcquireAction::RetireAndRetry
} else if let Some(meta) = found_conn {
AcquireAction::ValidateExisting(meta)
} else if inner.can_create_new() {
inner.total_count += 1;
inner.active_count += 1;
AcquireAction::CreateNew
} else {
inner.waiter_count += 1;
AcquireAction::Wait
}
};
(action, retired)
};
for guard in retired {
match guard
.close(cx, "expired pooled connection retirement")
.await
{
RetirementOutcome::Closed => {}
RetirementOutcome::Cancelled(reason) => {
return Outcome::Cancelled(reason);
}
RetirementOutcome::Failed(error) => return Outcome::Err(error),
}
}
match action {
AcquireAction::RetireAndRetry => {
continue;
}
AcquireAction::PoolClosed => {
return Outcome::Err(Error::Pool(PoolError {
kind: PoolErrorKind::Closed,
message: "pool has been closed".to_string(),
source: None,
}));
}
AcquireAction::ValidateExisting(meta) => {
match self.validate_and_wrap(cx, meta, test_on_checkout).await {
Outcome::Ok(Some(pooled)) => return Outcome::Ok(pooled),
Outcome::Ok(None) => {
tracing::warn!(
"pooled connection failed its checkout ping; closed and replaced"
);
continue;
}
Outcome::Err(e) => return Outcome::Err(e),
Outcome::Cancelled(r) => return Outcome::Cancelled(r),
Outcome::Panicked(p) => return Outcome::Panicked(p),
}
}
AcquireAction::CreateNew => {
let mut slot_guard = ActiveSlotGuard::new(Arc::clone(&self.shared));
match factory().await {
Outcome::Ok(conn) => {
self.shared
.connections_created
.fetch_add(1, Ordering::Relaxed);
let publish = match self.shared.lock_or_error("acquire_publish") {
Ok(inner) if inner.closed => {
let retirement = ActiveConnectionGuard::new(
Arc::clone(&self.shared),
ConnectionMeta::new(conn, clock.clone()),
);
slot_guard.disarm();
FactoryPublish::Retire {
guard: retirement,
error: Error::Pool(PoolError {
kind: PoolErrorKind::Closed,
message: "pool has been closed".to_string(),
source: None,
}),
context: "acquire publish after pool closure",
}
}
Ok(_inner) => {
self.shared.acquires.fetch_add(1, Ordering::Relaxed);
let meta = ConnectionMeta::new(conn, clock.clone());
slot_guard.disarm();
FactoryPublish::Published(PooledConnection::new(
meta,
Arc::downgrade(&self.shared),
))
}
Err(error) => {
let retirement = ActiveConnectionGuard::new(
Arc::clone(&self.shared),
ConnectionMeta::new(conn, clock.clone()),
);
slot_guard.disarm();
FactoryPublish::Retire {
guard: retirement,
error,
context: "acquire publish bookkeeping failure",
}
}
};
match publish {
FactoryPublish::Published(pooled) => {
return Outcome::Ok(pooled);
}
FactoryPublish::Retire {
guard,
error,
context,
} => match guard.close(cx, context).await {
RetirementOutcome::Closed => return Outcome::Err(error),
RetirementOutcome::Cancelled(reason) => {
return Outcome::Cancelled(reason);
}
RetirementOutcome::Failed(close_error) => {
return Outcome::Err(close_error);
}
},
}
}
Outcome::Err(e) => {
slot_guard.release("connection factory error");
return Outcome::Err(e);
}
Outcome::Cancelled(reason) => {
slot_guard.release("connection factory cancellation");
return Outcome::Cancelled(reason);
}
Outcome::Panicked(info) => {
slot_guard.release("connection factory panic");
return Outcome::Panicked(info);
}
}
}
AcquireAction::Wait => {
let remaining = Duration::from_nanos(
deadline.as_nanos().saturating_sub(clock.now().as_nanos()),
);
if remaining.is_zero() {
if let Ok(mut inner) = self.shared.lock_or_error("acquire_timeout") {
inner.waiter_count -= 1;
}
self.shared.timeouts.fetch_add(1, Ordering::Relaxed);
return Outcome::Err(Error::Pool(PoolError {
kind: PoolErrorKind::Timeout,
message: "acquire timeout: no connections available".to_string(),
source: None,
}));
}
let wait_time = remaining.min(Duration::from_millis(100));
{
let mut notified = std::pin::pin!(self.shared.conn_available.notified());
let mut slice =
std::pin::pin!(asupersync::time::sleep(cx.now(), wait_time));
let _ = Select::new(notified.as_mut(), slice.as_mut()).await;
}
{
if let Ok(mut inner) = self.shared.lock_or_error("acquire_wake") {
inner.waiter_count = inner.waiter_count.saturating_sub(1);
}
}
}
}
}
}
async fn validate_and_wrap(
&self,
cx: &Cx,
meta: ConnectionMeta<C>,
test_on_checkout: bool,
) -> Outcome<Option<PooledConnection<C>>, Error> {
let guard = ActiveConnectionGuard::new(Arc::clone(&self.shared), meta);
if test_on_checkout {
match guard.connection().ping(cx).await {
Outcome::Ok(()) => {
self.shared.acquires.fetch_add(1, Ordering::Relaxed);
Outcome::Ok(Some(PooledConnection::new(
guard.into_meta(),
Arc::downgrade(&self.shared),
)))
}
Outcome::Err(_) | Outcome::Cancelled(_) | Outcome::Panicked(_) => {
match guard
.close(cx, "pooled connection checkout validation failure")
.await
{
RetirementOutcome::Closed => {}
RetirementOutcome::Cancelled(reason) => {
return Outcome::Cancelled(reason);
}
RetirementOutcome::Failed(error) => return Outcome::Err(error),
}
Outcome::Ok(None)
}
}
} else {
self.shared.acquires.fetch_add(1, Ordering::Relaxed);
Outcome::Ok(Some(PooledConnection::new(
guard.into_meta(),
Arc::downgrade(&self.shared),
)))
}
}
fn begin_idle_retirement(&self) -> Vec<ActiveConnectionGuard<C>> {
match self.shared.inner.lock() {
Ok(mut inner) => {
let mut retired = Vec::with_capacity(inner.idle.len());
while let Some(meta) = inner.idle.pop_front() {
inner.active_count += 1;
retired.push(ActiveConnectionGuard::new(Arc::clone(&self.shared), meta));
}
retired
}
Err(_poisoned) => {
tracing::error!(
"Pool mutex poisoned during idle retirement; \
idle connections cannot be retired safely"
);
Vec::new()
}
}
}
fn begin_close(&self) -> Vec<ActiveConnectionGuard<C>> {
let retired = match self.shared.inner.lock() {
Ok(mut inner) => {
inner.closed = true;
let mut retired = Vec::with_capacity(inner.idle.len());
while let Some(meta) = inner.idle.pop_front() {
inner.active_count += 1;
retired.push(ActiveConnectionGuard::new(Arc::clone(&self.shared), meta));
}
retired
}
Err(poisoned) => {
tracing::error!(
"Pool mutex poisoned during close; attempting recovery. \
Pool state may be inconsistent."
);
let mut inner = poisoned.into_inner();
inner.closed = true;
let mut retired = Vec::with_capacity(inner.idle.len());
while let Some(meta) = inner.idle.pop_front() {
inner.active_count += 1;
retired.push(ActiveConnectionGuard::new(Arc::clone(&self.shared), meta));
}
retired
}
};
self.shared.conn_available.notify_waiters();
retired
}
fn close_retirements_blocking(
&self,
retired: Vec<ActiveConnectionGuard<C>>,
context: &'static str,
) {
for guard in retired {
let runtime = match RuntimeBuilder::current_thread().build() {
Ok(runtime) => runtime,
Err(error) => {
tracing::warn!(
context,
error = %error,
"failed to build runtime while retiring pooled connection"
);
let error = Error::Custom(format!(
"failed to build runtime while retiring pooled connection: {error}"
));
self.shared.record_retirement_failure(context, &error);
drop(guard);
continue;
}
};
let cx = runtime.request_cx_with_budget(Budget::INFINITE);
let _ = runtime.block_on(guard.close(&cx, context));
}
}
pub fn clear_idle(&self) {
let retired = self.begin_idle_retirement();
self.close_retirements_blocking(retired, "pool clear_idle");
}
pub fn close(&self) {
let retired = self.begin_close();
self.close_retirements_blocking(retired, "pool close");
if self.shared.lock_or_recover().active_count == 0 {
let _ = self.shared.active_drained.set(());
}
}
pub async fn close_and_drain(&self, cx: &Cx) -> Outcome<(), Error> {
let retired = self.begin_close();
let mut direct_failure = None;
for guard in retired {
match guard.close(cx, "pool close-and-drain").await {
RetirementOutcome::Closed => {}
RetirementOutcome::Cancelled(reason) => {
return Outcome::Cancelled(reason);
}
RetirementOutcome::Failed(error) => {
direct_failure.get_or_insert(error);
}
}
}
let poisoned_failure = match self.shared.inner.lock() {
Ok(inner) => {
drop(inner);
None
}
Err(poisoned) => {
let error = Error::Pool(PoolError::poisoned("close_and_drain"));
self.shared
.record_retirement_failure("close_and_drain", &error);
drop(poisoned.into_inner());
Some(error)
}
};
if self.shared.lock_or_recover().active_count == 0 {
let _ = self.shared.active_drained.set(());
}
if self.shared.active_drained.wait(cx).await.is_ok() {
if let Some(error) = direct_failure {
Outcome::Err(error)
} else if let Some(error) = poisoned_failure {
Outcome::Err(error)
} else if let Some(error) = self.shared.retirement_failure_error() {
Outcome::Err(error)
} else {
Outcome::Ok(())
}
} else {
let reason = cx
.cancel_reason()
.unwrap_or_else(|| CancelReason::user("pool drain cancelled"));
Outcome::Cancelled(reason)
}
}
#[must_use]
pub fn idle_count(&self) -> usize {
let inner = self.shared.lock_or_recover();
inner.idle.len()
}
#[must_use]
pub fn active_count(&self) -> usize {
let inner = self.shared.lock_or_recover();
inner.active_count
}
#[must_use]
pub fn total_count(&self) -> usize {
let inner = self.shared.lock_or_recover();
inner.total_count
}
}
impl<C: Connection> Drop for Pool<C> {
fn drop(&mut self) {
self.close();
}
}
enum AcquireAction<C> {
RetireAndRetry,
PoolClosed,
ValidateExisting(ConnectionMeta<C>),
CreateNew,
Wait,
}
enum FactoryPublish<C: Connection> {
Published(PooledConnection<C>),
Retire {
guard: ActiveConnectionGuard<C>,
error: Error,
context: &'static str,
},
}
pub struct PooledConnection<C: Connection> {
meta: Option<ConnectionMeta<C>>,
pool: Weak<PoolShared<C>>,
}
impl<C: Connection> PooledConnection<C> {
fn new(meta: ConnectionMeta<C>, pool: Weak<PoolShared<C>>) -> Self {
Self {
meta: Some(meta),
pool,
}
}
pub fn detach(mut self) -> C {
let conn = self.meta.take().expect("connection already detached").conn;
if let Some(pool) = self.pool.upgrade() {
pool.connections_closed.fetch_add(1, Ordering::Relaxed);
pool.release_active_slot("pooled connection detach");
}
conn
}
#[must_use]
pub fn age(&self) -> Duration {
self.meta.as_ref().map_or(Duration::ZERO, |m| m.age())
}
#[must_use]
pub fn idle_time(&self) -> Duration {
self.meta.as_ref().map_or(Duration::ZERO, |m| m.idle_time())
}
}
impl<C: Connection> std::ops::Deref for PooledConnection<C> {
type Target = C;
fn deref(&self) -> &Self::Target {
&self
.meta
.as_ref()
.expect("connection already returned to pool")
.conn
}
}
impl<C: Connection> std::ops::DerefMut for PooledConnection<C> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self
.meta
.as_mut()
.expect("connection already returned to pool")
.conn
}
}
impl<C: Connection> Drop for PooledConnection<C> {
fn drop(&mut self) {
if let Some(mut meta) = self.meta.take() {
meta.touch(); if let Some(pool) = self.pool.upgrade() {
let mut inner = match pool.inner.lock() {
Ok(guard) => guard,
Err(poisoned) => {
drop(poisoned);
tracing::error!(
"Pool mutex poisoned during connection return; \
connection will be closed instead of returned. A thread panicked while holding the lock."
);
let context = "pooled connection drop poisoned";
if let Err(error) = close_connection_blocking(meta.conn, context) {
pool.record_retirement_failure(context, &error);
}
pool.connections_closed.fetch_add(1, Ordering::Relaxed);
pool.release_active_slot(context);
return;
}
};
if inner.closed {
drop(inner);
let context = "pooled connection drop closed pool";
if let Err(error) = close_connection_blocking(meta.conn, context) {
pool.record_retirement_failure(context, &error);
}
pool.connections_closed.fetch_add(1, Ordering::Relaxed);
pool.release_active_slot("pooled connection drop closed pool");
return;
}
let max_lifetime = Duration::from_millis(inner.config.max_lifetime_ms);
if meta.age() > max_lifetime {
drop(inner);
let context = "pooled connection drop max lifetime";
if let Err(error) = close_connection_blocking(meta.conn, context) {
pool.record_retirement_failure(context, &error);
}
pool.connections_closed.fetch_add(1, Ordering::Relaxed);
pool.release_active_slot("pooled connection drop max lifetime");
return;
}
inner.active_count -= 1;
inner.idle.push_back(meta);
drop(inner);
pool.conn_available.notify_one();
} else {
let _ = close_connection_blocking(meta.conn, "pooled connection drop missing pool");
}
}
}
}
impl<C: Connection + std::fmt::Debug> std::fmt::Debug for PooledConnection<C> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PooledConnection")
.field("conn", &self.meta.as_ref().map(|m| &m.conn))
.field("age", &self.age())
.field("idle_time", &self.idle_time())
.finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
use asupersync::lab::explorer::{DporExplorer, ExplorerConfig};
use asupersync::lab::{LabConfig, LabRuntime};
use asupersync::time::VirtualClock;
use asupersync::types::RegionId;
use asupersync::{Budget, Time};
use sqlmodel_core::connection::{
Dialect, IsolationLevel, PreparedStatement, TransactionMode, TransactionOps,
};
use sqlmodel_core::error::{
ConnectionError, ConnectionErrorKind, QueryError, QueryErrorKind, TransactionErrorKind,
};
use sqlmodel_core::{RetryPolicy, TransactionOptions, retry_transaction};
use sqlmodel_core::{Row, Value};
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicUsize};
use std::task::{Context, Poll, Wake, Waker};
fn test_clock() -> TimerDriverHandle {
TimerDriverHandle::with_wall_clock()
}
fn virtual_clock() -> (Arc<VirtualClock>, TimerDriverHandle) {
let clock = Arc::new(VirtualClock::new());
let driver = TimerDriverHandle::with_virtual_clock(Arc::clone(&clock));
(clock, driver)
}
fn backdate<C>(meta: &mut ConnectionMeta<C>, ago: Duration) {
meta.created_at = meta
.clock
.now()
.saturating_sub_nanos(u64::try_from(ago.as_nanos()).expect("duration fits u64 nanos"));
}
#[derive(Debug)]
struct MockConnection {
id: u32,
ping_should_fail: Arc<AtomicBool>,
pool_close_calls: Arc<AtomicUsize>,
pool_close_pending: bool,
pool_close_should_fail: bool,
pool_shared: Option<Weak<PoolShared<MockConnection>>>,
pool_lock_was_free: Option<Arc<AtomicBool>>,
}
impl MockConnection {
fn new(id: u32) -> Self {
Self {
id,
ping_should_fail: Arc::new(AtomicBool::new(false)),
pool_close_calls: Arc::new(AtomicUsize::new(0)),
pool_close_pending: false,
pool_close_should_fail: false,
pool_shared: None,
pool_lock_was_free: None,
}
}
#[allow(dead_code)]
fn with_ping_behavior(id: u32, should_fail: Arc<AtomicBool>) -> Self {
Self {
id,
ping_should_fail: should_fail,
pool_close_calls: Arc::new(AtomicUsize::new(0)),
pool_close_pending: false,
pool_close_should_fail: false,
pool_shared: None,
pool_lock_was_free: None,
}
}
fn with_pool_close_counter(id: u32, pool_close_calls: Arc<AtomicUsize>) -> Self {
Self {
id,
ping_should_fail: Arc::new(AtomicBool::new(false)),
pool_close_calls,
pool_close_pending: false,
pool_close_should_fail: false,
pool_shared: None,
pool_lock_was_free: None,
}
}
fn with_pool_close_probe(
id: u32,
pool_close_calls: Arc<AtomicUsize>,
pool_shared: Weak<PoolShared<MockConnection>>,
pool_lock_was_free: Arc<AtomicBool>,
) -> Self {
Self {
id,
ping_should_fail: Arc::new(AtomicBool::new(false)),
pool_close_calls,
pool_close_pending: false,
pool_close_should_fail: false,
pool_shared: Some(pool_shared),
pool_lock_was_free: Some(pool_lock_was_free),
}
}
fn with_pending_pool_close(id: u32, pool_close_calls: Arc<AtomicUsize>) -> Self {
Self {
id,
ping_should_fail: Arc::new(AtomicBool::new(false)),
pool_close_calls,
pool_close_pending: true,
pool_close_should_fail: false,
pool_shared: None,
pool_lock_was_free: None,
}
}
fn with_failing_pool_close(id: u32) -> Self {
Self {
id,
ping_should_fail: Arc::new(AtomicBool::new(false)),
pool_close_calls: Arc::new(AtomicUsize::new(0)),
pool_close_pending: false,
pool_close_should_fail: true,
pool_shared: None,
pool_lock_was_free: None,
}
}
}
struct GatedFactory {
ready: Arc<AtomicBool>,
conn: Option<MockConnection>,
}
#[derive(Default)]
struct WakeCounter {
wakes: AtomicUsize,
}
impl Wake for WakeCounter {
fn wake(self: Arc<Self>) {
self.wakes.fetch_add(1, Ordering::Relaxed);
}
fn wake_by_ref(self: &Arc<Self>) {
self.wakes.fetch_add(1, Ordering::Relaxed);
}
}
impl Future for GatedFactory {
type Output = Outcome<MockConnection, Error>;
fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
if self.ready.load(Ordering::Acquire) {
Poll::Ready(Outcome::Ok(
self.conn
.take()
.expect("gated factory polled after completion"),
))
} else {
Poll::Pending
}
}
}
struct MockTx;
#[allow(clippy::unused_async_trait_impl)]
impl TransactionOps for MockTx {
async fn query(&self, _cx: &Cx, _sql: &str, _params: &[Value]) -> Outcome<Vec<Row>, Error> {
Outcome::Ok(vec![])
}
async fn query_one(
&self,
_cx: &Cx,
_sql: &str,
_params: &[Value],
) -> Outcome<Option<Row>, Error> {
Outcome::Ok(None)
}
async fn execute(&self, _cx: &Cx, _sql: &str, _params: &[Value]) -> Outcome<u64, Error> {
Outcome::Ok(0)
}
async fn savepoint(&self, _cx: &Cx, _name: &str) -> Outcome<(), Error> {
Outcome::Ok(())
}
async fn rollback_to(&self, _cx: &Cx, _name: &str) -> Outcome<(), Error> {
Outcome::Ok(())
}
async fn release(&self, _cx: &Cx, _name: &str) -> Outcome<(), Error> {
Outcome::Ok(())
}
async fn commit(self, _cx: &Cx) -> Outcome<(), Error> {
Outcome::Ok(())
}
async fn rollback(self, _cx: &Cx) -> Outcome<(), Error> {
Outcome::Ok(())
}
}
#[allow(clippy::unused_async_trait_impl)]
impl Connection for MockConnection {
type Tx<'conn> = MockTx;
async fn query(&self, _cx: &Cx, _sql: &str, _params: &[Value]) -> Outcome<Vec<Row>, Error> {
Outcome::Ok(vec![])
}
async fn query_one(
&self,
_cx: &Cx,
_sql: &str,
_params: &[Value],
) -> Outcome<Option<Row>, Error> {
Outcome::Ok(None)
}
async fn execute(&self, _cx: &Cx, _sql: &str, _params: &[Value]) -> Outcome<u64, Error> {
Outcome::Ok(0)
}
async fn insert(&self, _cx: &Cx, _sql: &str, _params: &[Value]) -> Outcome<i64, Error> {
Outcome::Ok(0)
}
async fn batch(
&self,
_cx: &Cx,
_statements: &[(String, Vec<Value>)],
) -> Outcome<Vec<u64>, Error> {
Outcome::Ok(vec![])
}
async fn begin(&self, _cx: &Cx) -> Outcome<Self::Tx<'_>, Error> {
Outcome::Ok(MockTx)
}
async fn begin_with(
&self,
_cx: &Cx,
_isolation: IsolationLevel,
) -> Outcome<Self::Tx<'_>, Error> {
Outcome::Ok(MockTx)
}
async fn prepare(&self, _cx: &Cx, _sql: &str) -> Outcome<PreparedStatement, Error> {
Outcome::Ok(PreparedStatement::new(1, String::new(), 0))
}
async fn query_prepared(
&self,
_cx: &Cx,
_stmt: &PreparedStatement,
_params: &[Value],
) -> Outcome<Vec<Row>, Error> {
Outcome::Ok(vec![])
}
async fn execute_prepared(
&self,
_cx: &Cx,
_stmt: &PreparedStatement,
_params: &[Value],
) -> Outcome<u64, Error> {
Outcome::Ok(0)
}
async fn ping(&self, _cx: &Cx) -> Outcome<(), Error> {
if self.ping_should_fail.load(Ordering::Relaxed) {
Outcome::Err(Error::Connection(ConnectionError {
kind: ConnectionErrorKind::Disconnected,
message: "mock ping failed".to_string(),
source: None,
}))
} else {
Outcome::Ok(())
}
}
async fn close(self, _cx: &Cx) -> Result<(), Error> {
Ok(())
}
async fn close_for_pool(self, _cx: &Cx) -> Result<(), Error> {
if let (Some(pool_shared), Some(pool_lock_was_free)) =
(self.pool_shared.as_ref(), self.pool_lock_was_free.as_ref())
{
let mutex_is_available = pool_shared
.upgrade()
.is_none_or(|shared| shared.inner.try_lock().is_ok());
pool_lock_was_free.store(mutex_is_available, Ordering::Relaxed);
}
self.pool_close_calls.fetch_add(1, Ordering::Relaxed);
if self.pool_close_pending {
std::future::pending::<()>().await;
}
if self.pool_close_should_fail {
return Err(Error::Custom("mock pool close failure".to_string()));
}
Ok(())
}
}
#[test]
fn test_config_default() {
let config = PoolConfig::default();
assert_eq!(config.min_connections, 1);
assert_eq!(config.max_connections, 10);
assert_eq!(config.idle_timeout_ms, 600_000);
assert_eq!(config.acquire_timeout_ms, 30_000);
assert_eq!(config.max_lifetime_ms, 1_800_000);
assert!(config.test_on_checkout);
}
#[test]
fn test_config_builder() {
let config = PoolConfig::new(20)
.min_connections(5)
.idle_timeout(60_000)
.acquire_timeout(5_000)
.max_lifetime(300_000)
.test_on_checkout(false);
assert_eq!(config.min_connections, 5);
assert_eq!(config.max_connections, 20);
assert_eq!(config.idle_timeout_ms, 60_000);
assert_eq!(config.acquire_timeout_ms, 5_000);
assert_eq!(config.max_lifetime_ms, 300_000);
assert!(!config.test_on_checkout);
}
#[test]
fn test_config_clone() {
let config = PoolConfig::new(15).min_connections(3);
let cloned = config.clone();
assert_eq!(config.max_connections, cloned.max_connections);
assert_eq!(config.min_connections, cloned.min_connections);
}
#[test]
fn test_stats_default() {
let stats = PoolStats::default();
assert_eq!(stats.total_connections, 0);
assert_eq!(stats.idle_connections, 0);
assert_eq!(stats.active_connections, 0);
assert_eq!(stats.pending_requests, 0);
assert_eq!(stats.connections_created, 0);
assert_eq!(stats.connections_closed, 0);
assert_eq!(stats.acquires, 0);
assert_eq!(stats.timeouts, 0);
}
#[test]
fn test_stats_clone() {
let stats = PoolStats {
total_connections: 5,
acquires: 100,
..Default::default()
};
let cloned = stats.clone();
assert_eq!(stats.total_connections, cloned.total_connections);
assert_eq!(stats.acquires, cloned.acquires);
}
#[test]
fn test_connection_meta_timing() {
struct DummyConn;
let (clock, driver) = virtual_clock();
let meta = ConnectionMeta::new(DummyConn, driver);
let initial_age = meta.age();
assert_eq!(initial_age, Duration::ZERO);
clock.advance(10_000_000);
assert!(meta.age() >= Duration::from_millis(10));
assert!(meta.idle_time() >= Duration::from_millis(10));
}
#[test]
fn test_connection_meta_touch() {
struct DummyConn;
let (clock, driver) = virtual_clock();
let mut meta = ConnectionMeta::new(DummyConn, driver);
clock.advance(10_000_000); let idle_before_touch = meta.idle_time();
assert!(idle_before_touch >= Duration::from_millis(10));
meta.touch();
let idle_after_touch = meta.idle_time();
assert_eq!(idle_after_touch, Duration::ZERO);
assert!(idle_after_touch < idle_before_touch);
assert!(meta.age() >= Duration::from_millis(10));
}
#[test]
fn test_pool_new() {
let config = PoolConfig::new(5);
let pool: Pool<MockConnection> = Pool::new(config);
assert_eq!(pool.idle_count(), 0);
assert_eq!(pool.active_count(), 0);
assert_eq!(pool.total_count(), 0);
assert!(!pool.is_closed());
assert!(!pool.at_capacity());
}
#[test]
fn test_pool_config() {
let config = PoolConfig::new(7).min_connections(2);
let pool: Pool<MockConnection> = Pool::new(config);
let retrieved_config = pool.config();
assert_eq!(retrieved_config.max_connections, 7);
assert_eq!(retrieved_config.min_connections, 2);
}
#[test]
fn test_pool_stats_initial() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
let stats = pool.stats();
assert_eq!(stats.total_connections, 0);
assert_eq!(stats.idle_connections, 0);
assert_eq!(stats.active_connections, 0);
assert_eq!(stats.pending_requests, 0);
assert_eq!(stats.connections_created, 0);
assert_eq!(stats.connections_closed, 0);
assert_eq!(stats.acquires, 0);
assert_eq!(stats.timeouts, 0);
}
#[test]
fn test_pool_close() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
assert!(!pool.is_closed());
pool.close();
assert!(pool.is_closed());
}
#[test]
fn test_close_and_drain_zero_active_completes_immediately() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
let runtime = RuntimeBuilder::current_thread()
.build()
.expect("build test runtime");
let cx = Cx::for_testing();
let outcome = runtime.block_on(pool.close_and_drain(&cx));
assert!(matches!(outcome, Outcome::Ok(())));
assert!(pool.is_closed());
assert_eq!(pool.active_count(), 0);
assert_eq!(pool.total_count(), 0);
}
#[test]
fn test_close_and_drain_surfaces_exact_idle_retirement_error() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(1));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.idle.push_back(ConnectionMeta::new(
MockConnection::with_failing_pool_close(1),
test_clock(),
));
}
let runtime = RuntimeBuilder::current_thread()
.build()
.expect("build test runtime");
let cx = Cx::for_testing();
let outcome = runtime.block_on(pool.close_and_drain(&cx));
match outcome {
Outcome::Err(Error::Custom(message)) => {
assert_eq!(message, "mock pool close failure");
}
other => panic!("expected exact driver close error, got {other:?}"),
}
assert!(pool.is_closed());
assert_eq!(pool.active_count(), 0);
assert_eq!(pool.total_count(), 0);
let later_cx = Cx::for_testing();
let later = runtime.block_on(pool.close_and_drain(&later_cx));
match later {
Outcome::Err(Error::Custom(message)) => {
assert_eq!(
message,
"pool retirement failed: pool close-and-drain: mock pool close failure"
);
}
other => panic!("expected persistent retirement error, got {other:?}"),
}
}
#[test]
fn test_checked_out_retirement_error_reaches_every_drainer() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(1));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let pooled = PooledConnection::new(
ConnectionMeta::new(MockConnection::with_failing_pool_close(1), test_clock()),
Arc::downgrade(&pool.shared),
);
let first_cx = Cx::for_testing();
let second_cx = Cx::for_testing();
let mut first_drain = Box::pin(pool.close_and_drain(&first_cx));
let mut second_drain = Box::pin(pool.close_and_drain(&second_cx));
let mut task_cx = Context::from_waker(Waker::noop());
assert!(matches!(
first_drain.as_mut().poll(&mut task_cx),
Poll::Pending
));
assert!(matches!(
second_drain.as_mut().poll(&mut task_cx),
Poll::Pending
));
drop(pooled);
let first_message = match first_drain.as_mut().poll(&mut task_cx) {
Poll::Ready(Outcome::Err(Error::Custom(message))) => message,
other => panic!("first drainer did not fail closed: {other:?}"),
};
let second_message = match second_drain.as_mut().poll(&mut task_cx) {
Poll::Ready(Outcome::Err(Error::Custom(message))) => message,
other => panic!("second drainer did not fail closed: {other:?}"),
};
assert_eq!(first_message, second_message);
assert_eq!(
first_message,
"pool retirement failed: pooled connection drop closed pool: \
mock pool close failure"
);
assert!(pool.is_closed());
assert_eq!(pool.active_count(), 0);
assert_eq!(pool.total_count(), 0);
}
#[test]
fn test_close_and_drain_waits_for_active_return_and_explicit_close() {
let pool_close_calls = Arc::new(AtomicUsize::new(0));
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(1));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let pooled = PooledConnection::new(
ConnectionMeta::new(
MockConnection::with_pool_close_counter(1, Arc::clone(&pool_close_calls)),
test_clock(),
),
Arc::downgrade(&pool.shared),
);
let cx = Cx::for_testing();
let mut drain = Box::pin(pool.close_and_drain(&cx));
let mut task_cx = Context::from_waker(Waker::noop());
assert!(matches!(drain.as_mut().poll(&mut task_cx), Poll::Pending));
assert!(pool.is_closed());
assert_eq!(pool.active_count(), 1);
assert_eq!(pool_close_calls.load(Ordering::Relaxed), 0);
drop(pooled);
assert_eq!(pool_close_calls.load(Ordering::Relaxed), 1);
assert!(matches!(
drain.as_mut().poll(&mut task_cx),
Poll::Ready(Outcome::Ok(()))
));
assert_eq!(pool.active_count(), 0);
assert_eq!(pool.total_count(), 0);
}
#[test]
fn test_close_and_drain_multiple_handles_and_drainers_share_final_wake() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(2));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 2;
inner.active_count = 2;
}
let first = PooledConnection::new(
ConnectionMeta::new(MockConnection::new(1), test_clock()),
Arc::downgrade(&pool.shared),
);
let second = PooledConnection::new(
ConnectionMeta::new(MockConnection::new(2), test_clock()),
Arc::downgrade(&pool.shared),
);
let first_cx = Cx::for_testing();
let second_cx = Cx::for_testing();
let mut first_drain = Box::pin(pool.close_and_drain(&first_cx));
let mut second_drain = Box::pin(pool.close_and_drain(&second_cx));
let mut task_cx = Context::from_waker(Waker::noop());
assert!(matches!(
first_drain.as_mut().poll(&mut task_cx),
Poll::Pending
));
assert!(matches!(
second_drain.as_mut().poll(&mut task_cx),
Poll::Pending
));
drop(first);
assert_eq!(pool.active_count(), 1);
assert!(matches!(
first_drain.as_mut().poll(&mut task_cx),
Poll::Pending
));
assert!(matches!(
second_drain.as_mut().poll(&mut task_cx),
Poll::Pending
));
drop(second);
assert_eq!(pool.active_count(), 0);
assert!(matches!(
first_drain.as_mut().poll(&mut task_cx),
Poll::Ready(Outcome::Ok(()))
));
assert!(matches!(
second_drain.as_mut().poll(&mut task_cx),
Poll::Ready(Outcome::Ok(()))
));
}
#[test]
fn test_close_and_drain_cancellation_keeps_pool_closed() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(1));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let pooled = PooledConnection::new(
ConnectionMeta::new(MockConnection::new(1), test_clock()),
Arc::downgrade(&pool.shared),
);
let cx = Cx::for_testing();
let mut drain = Box::pin(pool.close_and_drain(&cx));
let wake_counter = Arc::new(WakeCounter::default());
let waker = Waker::from(Arc::clone(&wake_counter));
let mut task_cx = Context::from_waker(&waker);
assert!(matches!(drain.as_mut().poll(&mut task_cx), Poll::Pending));
cx.set_cancel_requested(true);
assert!(
wake_counter.wakes.load(Ordering::Relaxed) > 0,
"Cx cancellation must wake the registered close-and-drain waiter"
);
assert!(matches!(
drain.as_mut().poll(&mut task_cx),
Poll::Ready(Outcome::Cancelled(_))
));
assert!(pool.is_closed());
assert_eq!(pool.active_count(), 1);
drop(drain);
let resume_cx = Cx::for_testing();
let mut resumed_drain = Box::pin(pool.close_and_drain(&resume_cx));
let mut resumed_task_cx = Context::from_waker(Waker::noop());
assert!(matches!(
resumed_drain.as_mut().poll(&mut resumed_task_cx),
Poll::Pending
));
assert!(pool.is_closed());
drop(pooled);
assert!(matches!(
resumed_drain.as_mut().poll(&mut resumed_task_cx),
Poll::Ready(Outcome::Ok(()))
));
assert!(pool.is_closed());
assert_eq!(pool.active_count(), 0);
}
#[test]
fn test_close_and_drain_expired_deadline_keeps_pool_closed() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(1));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let pooled = PooledConnection::new(
ConnectionMeta::new(MockConnection::new(1), test_clock()),
Arc::downgrade(&pool.shared),
);
let cx = Cx::for_testing_with_budget(Budget::new().with_deadline(Time::ZERO));
let mut drain = Box::pin(pool.close_and_drain(&cx));
let mut task_cx = Context::from_waker(Waker::noop());
assert!(matches!(
drain.as_mut().poll(&mut task_cx),
Poll::Ready(Outcome::Cancelled(_))
));
assert!(pool.is_closed());
assert_eq!(pool.active_count(), 1);
drop(drain);
drop(pooled);
assert_eq!(pool.active_count(), 0);
}
#[test]
fn test_dropped_in_flight_factory_releases_reserved_slot() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(1));
let cx = Cx::for_testing();
let mut acquire = Box::pin(pool.acquire(&cx, || {
std::future::pending::<Outcome<MockConnection, Error>>()
}));
let mut task_cx = Context::from_waker(Waker::noop());
assert!(matches!(acquire.as_mut().poll(&mut task_cx), Poll::Pending));
assert_eq!(pool.active_count(), 1);
assert_eq!(pool.total_count(), 1);
drop(acquire);
assert_eq!(pool.active_count(), 0);
assert_eq!(pool.total_count(), 0);
}
#[test]
fn test_dropped_multi_expired_idle_close_releases_all_retirement_slots() {
let pool_close_calls = Arc::new(AtomicUsize::new(0));
let pool: Pool<MockConnection> =
Pool::new(PoolConfig::new(3).max_lifetime(1).test_on_checkout(false));
let mut first_expired = ConnectionMeta::new(
MockConnection::with_pending_pool_close(1, Arc::clone(&pool_close_calls)),
pool.shared.clock.clone(),
);
backdate(&mut first_expired, Duration::from_secs(1));
let mut second_expired = ConnectionMeta::new(
MockConnection::with_pool_close_counter(2, Arc::clone(&pool_close_calls)),
pool.shared.clock.clone(),
);
backdate(&mut second_expired, Duration::from_secs(1));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 2;
inner.idle.push_back(first_expired);
inner.idle.push_back(second_expired);
}
let cx = Cx::for_testing();
let mut acquire =
Box::pin(pool.acquire(&cx, || async { Outcome::Ok(MockConnection::new(3)) }));
let mut task_cx = Context::from_waker(Waker::noop());
assert!(matches!(acquire.as_mut().poll(&mut task_cx), Poll::Pending));
assert_eq!(pool_close_calls.load(Ordering::Relaxed), 1);
assert_eq!(pool.active_count(), 2);
assert_eq!(pool.total_count(), 2);
let drain_cx = Cx::for_testing();
let mut drain = Box::pin(pool.close_and_drain(&drain_cx));
assert!(matches!(drain.as_mut().poll(&mut task_cx), Poll::Pending));
drop(acquire);
assert!(matches!(
drain.as_mut().poll(&mut task_cx),
Poll::Ready(Outcome::Ok(()))
));
assert!(pool.is_closed());
assert_eq!(pool.active_count(), 0);
assert_eq!(pool.total_count(), 0);
}
#[test]
fn test_dropped_pending_idle_drain_releases_resource_for_other_drainer() {
let pool_close_calls = Arc::new(AtomicUsize::new(0));
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(1));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.idle.push_back(ConnectionMeta::new(
MockConnection::with_pending_pool_close(1, Arc::clone(&pool_close_calls)),
test_clock(),
));
}
let first_cx = Cx::for_testing();
let second_cx = Cx::for_testing();
let mut first_drain = Box::pin(pool.close_and_drain(&first_cx));
let mut second_drain = Box::pin(pool.close_and_drain(&second_cx));
let mut task_cx = Context::from_waker(Waker::noop());
assert!(matches!(
first_drain.as_mut().poll(&mut task_cx),
Poll::Pending
));
assert!(pool.is_closed());
assert_eq!(pool_close_calls.load(Ordering::Relaxed), 1);
assert_eq!(pool.active_count(), 1);
assert_eq!(pool.total_count(), 1);
assert!(matches!(
second_drain.as_mut().poll(&mut task_cx),
Poll::Pending
));
first_cx.set_cancel_requested(true);
assert!(matches!(
first_drain.as_mut().poll(&mut task_cx),
Poll::Ready(Outcome::Cancelled(_))
));
drop(first_drain);
assert_eq!(pool.active_count(), 0);
assert_eq!(pool.total_count(), 0);
assert!(matches!(
second_drain.as_mut().poll(&mut task_cx),
Poll::Ready(Outcome::Ok(()))
));
assert!(pool.is_closed());
}
#[test]
fn test_dropped_validation_close_releases_armed_active_slot() {
let pool_close_calls = Arc::new(AtomicUsize::new(0));
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(1).test_on_checkout(true));
let failed = MockConnection::with_pending_pool_close(1, Arc::clone(&pool_close_calls));
failed.ping_should_fail.store(true, Ordering::Relaxed);
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner
.idle
.push_back(ConnectionMeta::new(failed, test_clock()));
}
let cx = Cx::for_testing();
let mut acquire =
Box::pin(pool.acquire(&cx, || async { Outcome::Ok(MockConnection::new(2)) }));
let mut task_cx = Context::from_waker(Waker::noop());
assert!(matches!(acquire.as_mut().poll(&mut task_cx), Poll::Pending));
assert_eq!(pool_close_calls.load(Ordering::Relaxed), 1);
assert_eq!(pool.active_count(), 1);
assert_eq!(pool.total_count(), 1);
pool.close();
drop(acquire);
assert_eq!(pool.active_count(), 0);
assert_eq!(pool.total_count(), 0);
let drain_cx = Cx::for_testing();
let mut drain = Box::pin(pool.close_and_drain(&drain_cx));
assert!(matches!(
drain.as_mut().poll(&mut task_cx),
Poll::Ready(Outcome::Ok(()))
));
}
#[test]
fn test_in_flight_factory_cannot_publish_after_close() {
let pool_close_calls = Arc::new(AtomicUsize::new(0));
let factory_ready = Arc::new(AtomicBool::new(false));
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(1));
let cx = Cx::for_testing();
let mut acquire = Box::pin(pool.acquire(&cx, || GatedFactory {
ready: Arc::clone(&factory_ready),
conn: Some(MockConnection::with_pool_close_counter(
1,
Arc::clone(&pool_close_calls),
)),
}));
let mut task_cx = Context::from_waker(Waker::noop());
assert!(matches!(acquire.as_mut().poll(&mut task_cx), Poll::Pending));
assert_eq!(pool.active_count(), 1);
pool.close();
factory_ready.store(true, Ordering::Release);
assert!(matches!(
acquire.as_mut().poll(&mut task_cx),
Poll::Ready(Outcome::Err(Error::Pool(PoolError {
kind: PoolErrorKind::Closed,
..
})))
));
assert_eq!(pool_close_calls.load(Ordering::Relaxed), 1);
assert_eq!(pool.active_count(), 0);
assert_eq!(pool.total_count(), 0);
}
#[test]
fn test_close_wakes_blocked_acquirer_to_observe_closed_state() {
use std::sync::mpsc;
use std::thread;
let pool = Arc::new(Pool::<MockConnection>::new(PoolConfig::new(1)));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let pooled = PooledConnection::new(
ConnectionMeta::new(MockConnection::new(1), test_clock()),
Arc::downgrade(&pool.shared),
);
let waiting_pool = Arc::clone(&pool);
let (started_tx, started_rx) = mpsc::sync_channel(0);
let waiter = thread::spawn(move || {
let runtime = RuntimeBuilder::current_thread()
.build()
.expect("build waiter runtime");
let cx = Cx::for_testing();
started_tx.send(()).expect("signal waiter start");
matches!(
runtime.block_on(
waiting_pool.acquire(&cx, || async { Outcome::Ok(MockConnection::new(2)) })
),
Outcome::Err(Error::Pool(PoolError {
kind: PoolErrorKind::Closed,
..
}))
)
});
started_rx.recv().expect("waiter thread should start");
let mut observed_waiter = false;
for _ in 0..100_000 {
if pool.stats().pending_requests == 1 {
observed_waiter = true;
break;
}
thread::yield_now();
}
pool.close();
let observed_closed = waiter.join().expect("waiter thread should not panic");
drop(pooled);
assert!(
observed_waiter,
"acquirer never registered as a pool waiter"
);
assert!(
observed_closed,
"blocked acquirer did not observe pool close"
);
}
#[test]
fn test_pool_close_routes_through_close_for_pool() {
let pool_close_calls = Arc::new(AtomicUsize::new(0));
let pool_lock_was_free = Arc::new(AtomicBool::new(false));
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(2));
{
let mut inner = pool
.shared
.inner
.lock()
.expect("pool mutex should not be poisoned");
inner.total_count = 1;
inner.idle.push_back(ConnectionMeta::new(
MockConnection::with_pool_close_probe(
1,
Arc::clone(&pool_close_calls),
Arc::downgrade(&pool.shared),
Arc::clone(&pool_lock_was_free),
),
test_clock(),
));
}
pool.close();
assert!(pool.is_closed());
assert_eq!(pool_close_calls.load(Ordering::Relaxed), 1);
assert!(pool_lock_was_free.load(Ordering::Relaxed));
}
#[test]
fn test_expired_idle_connection_routes_through_close_for_pool() {
let pool_close_calls = Arc::new(AtomicUsize::new(0));
let pool: Pool<MockConnection> =
Pool::new(PoolConfig::new(2).max_lifetime(1).test_on_checkout(false));
let mut expired = ConnectionMeta::new(
MockConnection::with_pool_close_counter(1, Arc::clone(&pool_close_calls)),
pool.shared.clock.clone(),
);
backdate(&mut expired, Duration::from_secs(1));
{
let mut inner = pool
.shared
.inner
.lock()
.expect("pool mutex should not be poisoned");
inner.total_count = 1;
inner.idle.push_back(expired);
}
let runtime = RuntimeBuilder::current_thread()
.build()
.expect("build test runtime");
let cx = Cx::for_testing();
let acquired =
runtime.block_on(pool.acquire(&cx, || async { Outcome::Ok(MockConnection::new(2)) }));
assert!(matches!(acquired, Outcome::Ok(_)));
assert_eq!(pool_close_calls.load(Ordering::Relaxed), 1);
}
#[test]
fn test_failed_validation_routes_through_close_for_pool() {
let pool_close_calls = Arc::new(AtomicUsize::new(0));
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(2).test_on_checkout(true));
let failed = MockConnection::with_pool_close_counter(1, Arc::clone(&pool_close_calls));
failed.ping_should_fail.store(true, Ordering::Relaxed);
{
let mut inner = pool
.shared
.inner
.lock()
.expect("pool mutex should not be poisoned");
inner.total_count = 1;
inner
.idle
.push_back(ConnectionMeta::new(failed, test_clock()));
}
let runtime = RuntimeBuilder::current_thread()
.build()
.expect("build test runtime");
let cx = Cx::for_testing();
let acquired =
runtime.block_on(pool.acquire(&cx, || async { Outcome::Ok(MockConnection::new(2)) }));
assert!(
matches!(acquired, Outcome::Ok(_)),
"acquire replaces the dead idle connection"
);
assert_eq!(pool_close_calls.load(Ordering::Relaxed), 1);
let stats = pool.stats();
assert_eq!(stats.connections_closed, 1);
assert_eq!(stats.connections_created, 1);
assert_eq!(stats.total_connections, 1);
drop(acquired);
assert_eq!(pool.stats().idle_connections, 1);
}
#[test]
fn waiters_yield_instead_of_blocking_a_single_threaded_runtime() {
let pool: Arc<Pool<MockConnection>> = Arc::new(Pool::new(
PoolConfig::new(1)
.min_connections(0)
.acquire_timeout(2_000)
.test_on_checkout(false),
));
let runtime = RuntimeBuilder::current_thread()
.build()
.expect("build test runtime");
let handle = runtime.handle();
let tasks: Vec<_> = (0..4)
.map(|i| {
let pool = Arc::clone(&pool);
handle.spawn(async move {
let cx = Cx::for_testing();
match pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(1)) })
.await
{
Outcome::Ok(lease) => {
asupersync::time::sleep(cx.now(), Duration::from_millis(20)).await;
drop(lease);
Ok(i)
}
Outcome::Err(e) => Err(e.to_string()),
Outcome::Cancelled(_) | Outcome::Panicked(_) => Err("cancelled".into()),
}
})
})
.collect();
let results = runtime.block_on(async {
let mut results = Vec::new();
for task in tasks {
results.push(task.await);
}
results
});
assert!(results.iter().all(Result::is_ok), "{results:?}");
let stats = pool.stats();
assert_eq!(stats.timeouts, 0, "{stats:?}");
assert_eq!(stats.connections_created, 1, "{stats:?}");
assert_eq!(stats.acquires, 4, "{stats:?}");
assert_eq!(stats.active_connections, 0, "{stats:?}");
}
#[test]
fn test_pool_inner_can_create_new() {
let mut inner = PoolInner::<MockConnection>::new(PoolConfig::new(3));
assert!(inner.can_create_new());
inner.total_count = 3;
assert!(!inner.can_create_new());
inner.total_count = 2;
assert!(inner.can_create_new());
inner.closed = true;
assert!(!inner.can_create_new());
}
#[test]
fn test_pool_inner_stats() {
let mut inner = PoolInner::<MockConnection>::new(PoolConfig::new(10));
inner.total_count = 5;
inner.active_count = 3;
inner.waiter_count = 2;
inner
.idle
.push_back(ConnectionMeta::new(MockConnection::new(1), test_clock()));
inner
.idle
.push_back(ConnectionMeta::new(MockConnection::new(2), test_clock()));
let stats = inner.stats();
assert_eq!(stats.total_connections, 5);
assert_eq!(stats.idle_connections, 2);
assert_eq!(stats.active_connections, 3);
assert_eq!(stats.pending_requests, 2);
}
#[test]
fn test_pooled_connection_age_and_idle_time() {
let (clock, driver) = virtual_clock();
let pool: Pool<MockConnection> =
Pool::with_timer_driver(PoolConfig::new(5).test_on_checkout(false), driver);
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let meta = ConnectionMeta::new(MockConnection::new(1), pool.shared.clock.clone());
let pooled = PooledConnection::new(meta, Arc::downgrade(&pool.shared));
assert_eq!(pooled.age(), Duration::ZERO);
assert_eq!(pooled.idle_time(), Duration::ZERO);
clock.advance(5_000_000); assert!(pooled.age() >= Duration::from_millis(5));
assert!(pooled.idle_time() >= Duration::from_millis(5));
}
#[test]
fn test_pooled_connection_detach() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let meta = ConnectionMeta::new(MockConnection::new(42), test_clock());
let pooled = PooledConnection::new(meta, Arc::downgrade(&pool.shared));
assert_eq!(pool.total_count(), 1);
assert_eq!(pool.active_count(), 1);
let conn = pooled.detach();
assert_eq!(conn.id, 42);
assert_eq!(pool.total_count(), 0);
assert_eq!(pool.active_count(), 0);
let stats = pool.stats();
assert_eq!(stats.connections_closed, 1);
}
#[test]
fn test_pooled_connection_drop_returns_to_pool() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let meta = ConnectionMeta::new(MockConnection::new(1), test_clock());
let pooled = PooledConnection::new(meta, Arc::downgrade(&pool.shared));
assert_eq!(pool.active_count(), 1);
assert_eq!(pool.idle_count(), 0);
drop(pooled);
assert_eq!(pool.active_count(), 0);
assert_eq!(pool.idle_count(), 1);
assert_eq!(pool.total_count(), 1); }
#[test]
fn test_pooled_connection_drop_when_pool_closed() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let meta = ConnectionMeta::new(MockConnection::new(1), test_clock());
let pooled = PooledConnection::new(meta, Arc::downgrade(&pool.shared));
pool.close();
drop(pooled);
assert_eq!(pool.idle_count(), 0);
assert_eq!(pool.active_count(), 0);
assert_eq!(pool.total_count(), 0);
assert_eq!(pool.stats().connections_closed, 1);
}
#[test]
fn test_pooled_connection_deref() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let meta = ConnectionMeta::new(MockConnection::new(99), test_clock());
let pooled = PooledConnection::new(meta, Arc::downgrade(&pool.shared));
assert_eq!(pooled.id, 99);
}
#[test]
fn test_pooled_connection_deref_mut() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let meta = ConnectionMeta::new(MockConnection::new(1), test_clock());
let mut pooled = PooledConnection::new(meta, Arc::downgrade(&pool.shared));
pooled.id = 50;
assert_eq!(pooled.id, 50);
}
#[test]
fn test_pooled_connection_debug() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let meta = ConnectionMeta::new(MockConnection::new(1), test_clock());
let pooled = PooledConnection::new(meta, Arc::downgrade(&pool.shared));
let debug_str = format!("{:?}", pooled);
assert!(debug_str.contains("PooledConnection"));
assert!(debug_str.contains("age"));
}
#[test]
fn test_pool_at_capacity() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(2));
assert!(!pool.at_capacity());
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
}
assert!(!pool.at_capacity());
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 2;
}
assert!(pool.at_capacity());
}
#[test]
fn test_acquire_action_enum() {
let retire: AcquireAction<MockConnection> = AcquireAction::RetireAndRetry;
assert!(matches!(retire, AcquireAction::RetireAndRetry));
let closed: AcquireAction<MockConnection> = AcquireAction::PoolClosed;
assert!(matches!(closed, AcquireAction::PoolClosed));
let create: AcquireAction<MockConnection> = AcquireAction::CreateNew;
assert!(matches!(create, AcquireAction::CreateNew));
let wait: AcquireAction<MockConnection> = AcquireAction::Wait;
assert!(matches!(wait, AcquireAction::Wait));
let meta = ConnectionMeta::new(MockConnection::new(1), test_clock());
let validate: AcquireAction<MockConnection> = AcquireAction::ValidateExisting(meta);
assert!(matches!(validate, AcquireAction::ValidateExisting(_)));
}
#[test]
fn test_pool_shared_atomic_counters() {
let shared = PoolShared::<MockConnection>::new(PoolConfig::new(5), test_clock());
assert_eq!(shared.connections_created.load(Ordering::Relaxed), 0);
assert_eq!(shared.connections_closed.load(Ordering::Relaxed), 0);
assert_eq!(shared.acquires.load(Ordering::Relaxed), 0);
assert_eq!(shared.timeouts.load(Ordering::Relaxed), 0);
shared.connections_created.fetch_add(1, Ordering::Relaxed);
shared.connections_closed.fetch_add(2, Ordering::Relaxed);
shared.acquires.fetch_add(10, Ordering::Relaxed);
shared.timeouts.fetch_add(3, Ordering::Relaxed);
assert_eq!(shared.connections_created.load(Ordering::Relaxed), 1);
assert_eq!(shared.connections_closed.load(Ordering::Relaxed), 2);
assert_eq!(shared.acquires.load(Ordering::Relaxed), 10);
assert_eq!(shared.timeouts.load(Ordering::Relaxed), 3);
}
#[test]
fn test_pool_close_clears_idle() {
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 3;
inner
.idle
.push_back(ConnectionMeta::new(MockConnection::new(1), test_clock()));
inner
.idle
.push_back(ConnectionMeta::new(MockConnection::new(2), test_clock()));
inner
.idle
.push_back(ConnectionMeta::new(MockConnection::new(3), test_clock()));
}
assert_eq!(pool.idle_count(), 3);
assert_eq!(pool.total_count(), 3);
pool.close();
assert_eq!(pool.idle_count(), 0);
assert_eq!(pool.total_count(), 0);
assert!(pool.is_closed());
assert_eq!(pool.stats().connections_closed, 3);
}
fn poison_pool_mutex() -> Pool<MockConnection> {
use std::panic;
use std::thread;
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 2;
inner.active_count = 1;
inner
.idle
.push_back(ConnectionMeta::new(MockConnection::new(1), test_clock()));
}
let shared_clone = Arc::clone(&pool.shared);
let handle = thread::spawn(move || {
let _guard = shared_clone.inner.lock().unwrap();
panic!("intentional panic to poison mutex");
});
let _ = handle.join();
assert!(pool.shared.inner.lock().is_err());
pool
}
#[test]
fn test_config_after_poisoning_returns_valid_data() {
let pool = poison_pool_mutex();
let config = pool.config();
assert_eq!(config.max_connections, 5);
}
#[test]
fn test_stats_after_poisoning_returns_valid_data() {
let pool = poison_pool_mutex();
let stats = pool.stats();
assert_eq!(stats.total_connections, 2);
assert_eq!(stats.active_connections, 1);
assert_eq!(stats.idle_connections, 1);
}
#[test]
fn test_at_capacity_after_poisoning() {
let pool = poison_pool_mutex();
assert!(!pool.at_capacity());
}
#[test]
fn test_is_closed_after_poisoning() {
let pool = poison_pool_mutex();
assert!(!pool.is_closed());
}
#[test]
fn test_idle_count_after_poisoning() {
let pool = poison_pool_mutex();
assert_eq!(pool.idle_count(), 1);
}
#[test]
fn test_active_count_after_poisoning() {
let pool = poison_pool_mutex();
assert_eq!(pool.active_count(), 1);
}
#[test]
fn test_total_count_after_poisoning() {
let pool = poison_pool_mutex();
assert_eq!(pool.total_count(), 2);
}
#[test]
fn test_lock_or_error_returns_error_when_poisoned() {
use std::thread;
let shared = Arc::new(PoolShared::<MockConnection>::new(
PoolConfig::new(5),
test_clock(),
));
let shared_clone = Arc::clone(&shared);
let handle = thread::spawn(move || {
let _guard = shared_clone.inner.lock().unwrap();
panic!("intentional panic to poison mutex");
});
let _ = handle.join();
let result = shared.lock_or_error("test_operation");
match result {
Err(Error::Pool(pool_err)) => {
assert!(matches!(pool_err.kind, PoolErrorKind::Poisoned));
assert!(pool_err.message.contains("poisoned"));
}
Err(other) => panic!("Expected Pool error, got: {:?}", other),
Ok(_) => panic!("Expected error, got Ok"),
}
}
#[test]
fn test_lock_or_recover_succeeds_when_poisoned() {
use std::thread;
let shared = Arc::new(PoolShared::<MockConnection>::new(
PoolConfig::new(5),
test_clock(),
));
{
let mut inner = shared.inner.lock().unwrap();
inner.total_count = 42;
}
let shared_clone = Arc::clone(&shared);
let handle = thread::spawn(move || {
let _guard = shared_clone.inner.lock().unwrap();
panic!("intentional panic to poison mutex");
});
let _ = handle.join();
assert!(shared.inner.lock().is_err());
let inner = shared.lock_or_recover();
assert_eq!(inner.total_count, 42);
}
#[test]
fn test_close_after_poisoning_recovers_and_closes() {
let pool = poison_pool_mutex();
pool.close();
assert!(pool.is_closed());
assert_eq!(pool.idle_count(), 0);
}
#[test]
fn test_poisoned_pool_return_completes_drain_accounting() {
use std::thread;
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(1));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let pooled = PooledConnection::new(
ConnectionMeta::new(MockConnection::new(1), test_clock()),
Arc::downgrade(&pool.shared),
);
let shared = Arc::clone(&pool.shared);
let poisoner = thread::spawn(move || {
let _guard = shared.inner.lock().unwrap();
panic!("intentional panic to poison drain accounting");
});
let _ = poisoner.join();
let cx = Cx::for_testing();
let mut drain = Box::pin(pool.close_and_drain(&cx));
let mut task_cx = Context::from_waker(Waker::noop());
assert!(matches!(drain.as_mut().poll(&mut task_cx), Poll::Pending));
assert!(pool.is_closed());
assert_eq!(pool.active_count(), 1);
drop(pooled);
assert_eq!(pool.active_count(), 0);
assert_eq!(pool.total_count(), 0);
assert!(
pool.shared.active_drained.get().is_some(),
"poison-aware final release must publish the drain latch"
);
assert!(matches!(
drain.as_mut().poll(&mut task_cx),
Poll::Ready(Outcome::Err(Error::Pool(PoolError {
kind: PoolErrorKind::Poisoned,
..
})))
));
}
#[test]
fn test_drop_pooled_connection_after_poisoning_does_not_panic() {
use std::panic;
use std::thread;
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let meta = ConnectionMeta::new(MockConnection::new(1), test_clock());
let pooled = PooledConnection::new(meta, Arc::downgrade(&pool.shared));
let shared_clone = Arc::clone(&pool.shared);
let handle = thread::spawn(move || {
let _guard = shared_clone.inner.lock().unwrap();
panic!("intentional panic to poison mutex");
});
let _ = handle.join();
assert!(pool.shared.inner.lock().is_err());
let drop_result = panic::catch_unwind(panic::AssertUnwindSafe(|| {
drop(pooled);
}));
assert!(
drop_result.is_ok(),
"Dropping PooledConnection after mutex poisoning should not panic"
);
}
#[test]
fn test_detach_after_poisoning_does_not_panic() {
use std::panic;
use std::thread;
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let meta = ConnectionMeta::new(MockConnection::new(42), test_clock());
let pooled = PooledConnection::new(meta, Arc::downgrade(&pool.shared));
let shared_clone = Arc::clone(&pool.shared);
let handle = thread::spawn(move || {
let _guard = shared_clone.inner.lock().unwrap();
panic!("intentional panic to poison mutex");
});
let _ = handle.join();
assert!(pool.shared.inner.lock().is_err());
let detach_result = panic::catch_unwind(panic::AssertUnwindSafe(|| pooled.detach()));
assert!(
detach_result.is_ok(),
"detach() after mutex poisoning should not panic"
);
let conn = detach_result.unwrap();
assert_eq!(conn.id, 42);
}
#[test]
fn test_pool_survives_thread_panic_during_acquire() {
use std::thread;
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
let pool_arc = Arc::new(pool);
let pool_clone = Arc::clone(&pool_arc);
let handle = thread::spawn(move || {
{
let mut inner = pool_clone.shared.inner.lock().unwrap();
inner.total_count = 1;
inner.active_count = 1;
}
let _guard = pool_clone.shared.inner.lock().unwrap();
panic!("simulated panic during database operation");
});
let _ = handle.join();
assert_eq!(pool_arc.total_count(), 1);
assert_eq!(pool_arc.config().max_connections, 5);
let stats = pool_arc.stats();
assert_eq!(stats.total_connections, 1);
}
#[test]
fn test_pool_close_after_thread_panic() {
use std::thread;
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.total_count = 2;
inner
.idle
.push_back(ConnectionMeta::new(MockConnection::new(1), test_clock()));
inner
.idle
.push_back(ConnectionMeta::new(MockConnection::new(2), test_clock()));
}
let shared_clone = Arc::clone(&pool.shared);
let handle = thread::spawn(move || {
let _guard = shared_clone.inner.lock().unwrap();
panic!("intentional panic");
});
let _ = handle.join();
pool.close();
assert!(pool.is_closed());
assert_eq!(pool.idle_count(), 0);
}
#[test]
fn test_multiple_reads_after_poisoning() {
let pool = poison_pool_mutex();
for _ in 0..10 {
let _ = pool.config();
let _ = pool.stats();
let _ = pool.at_capacity();
let _ = pool.is_closed();
let _ = pool.idle_count();
let _ = pool.active_count();
let _ = pool.total_count();
}
assert_eq!(pool.total_count(), 2);
}
#[test]
fn test_waiters_count_after_poisoning() {
use std::thread;
let pool: Pool<MockConnection> = Pool::new(PoolConfig::new(5));
{
let mut inner = pool.shared.inner.lock().unwrap();
inner.waiter_count = 3;
}
let shared_clone = Arc::clone(&pool.shared);
let handle = thread::spawn(move || {
let _guard = shared_clone.inner.lock().unwrap();
panic!("intentional panic");
});
let _ = handle.join();
let stats = pool.stats();
assert_eq!(stats.pending_requests, 3);
}
fn lab_pool(config: PoolConfig) -> (LabRuntime, Arc<Pool<MockConnection>>, RegionId) {
let mut runtime = LabRuntime::new(LabConfig::new(0x50f1).max_steps(500_000));
let driver = runtime
.state
.timer_driver_handle()
.expect("lab runtime installs a virtual-clock timer driver");
let pool = Arc::new(Pool::with_timer_driver(config, driver));
let region = runtime.state.create_root_region(Budget::INFINITE);
(runtime, pool, region)
}
#[test]
fn lab_idle_timeout_retires_exactly_the_expired_idle_connections() {
let (mut runtime, pool, region) = lab_pool(
PoolConfig::new(3)
.idle_timeout(100)
.test_on_checkout(false)
.acquire_timeout(5_000),
);
let stats_after = Arc::new(OnceCell::<PoolStats>::new());
let stats_recorder = Arc::clone(&stats_after);
let acquired_id = Arc::new(OnceCell::<u32>::new());
let id_recorder = Arc::clone(&acquired_id);
let (task, _) = runtime
.state
.create_task(region, Budget::INFINITE, async move {
let cx = Cx::current().expect("lab task has a cx");
let a = pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(1)) })
.await
.expect("acquire a");
let b = pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(2)) })
.await
.expect("acquire b");
let c = pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(3)) })
.await
.expect("acquire c");
drop(c);
drop(b);
drop(a);
asupersync::time::sleep(cx.now(), Duration::from_millis(250)).await;
let d = pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(4)) })
.await
.expect("acquire d");
let _ = id_recorder.set(d.id);
let _ = stats_recorder.set(pool.stats());
})
.expect("spawn lab task");
runtime.scheduler.lock().schedule(task, 0);
runtime.run_with_auto_advance();
assert_eq!(acquired_id.get().copied(), Some(4));
let stats = stats_after.get().expect("task ran to completion");
assert_eq!(stats.connections_created, 4);
assert_eq!(stats.connections_closed, 3);
assert_eq!(stats.acquires, 4);
assert_eq!(stats.idle_connections, 0);
assert_eq!(stats.active_connections, 1);
assert_eq!(stats.total_connections, 1);
}
#[test]
fn lab_max_lifetime_replaces_a_connection_exactly_when_it_passes() {
let (mut runtime, pool, region) = lab_pool(
PoolConfig::new(2)
.max_lifetime(500)
.idle_timeout(100_000)
.test_on_checkout(false)
.acquire_timeout(5_000),
);
let first_id = Arc::new(OnceCell::<u32>::new());
let first_recorder = Arc::clone(&first_id);
let second = Arc::new(OnceCell::<(u32, u64, u64)>::new());
let second_recorder = Arc::clone(&second);
let (task, _) = runtime
.state
.create_task(region, Budget::INFINITE, async move {
let cx = Cx::current().expect("lab task has a cx");
let a = pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(1)) })
.await
.expect("acquire a");
drop(a);
asupersync::time::sleep(cx.now(), Duration::from_millis(400)).await;
let b = pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(2)) })
.await
.expect("acquire b");
let _ = first_recorder.set(b.id);
drop(b);
asupersync::time::sleep(cx.now(), Duration::from_millis(101)).await;
let c = pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(3)) })
.await
.expect("acquire c");
let stats = pool.stats();
let _ = second_recorder.set((
c.id,
stats.connections_created,
stats.connections_closed,
));
})
.expect("spawn lab task");
runtime.scheduler.lock().schedule(task, 0);
runtime.run_with_auto_advance();
assert_eq!(
first_id.get().copied(),
Some(1),
"not retired before its lifetime passes"
);
assert_eq!(
second.get().copied(),
Some((3, 2, 1)),
"replaced exactly once the lifetime has passed"
);
}
#[test]
fn lab_acquire_timeout_fires_at_the_exact_virtual_deadline() {
let (mut runtime, pool, region) = lab_pool(
PoolConfig::new(1)
.acquire_timeout(200)
.test_on_checkout(false),
);
let holder_result = Arc::new(OnceCell::<u64>::new());
let holder_recorder = Arc::clone(&holder_result);
let waiter_result = Arc::new(OnceCell::<(u64, bool)>::new());
let waiter_recorder = Arc::clone(&waiter_result);
let holder_pool = Arc::clone(&pool);
let (holder, _) = runtime
.state
.create_task(region, Budget::INFINITE, async move {
let cx = Cx::current().expect("lab task has a cx");
let lease = holder_pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(1)) })
.await
.expect("holder acquires the only lease");
asupersync::time::sleep(cx.now(), Duration::from_millis(10_000)).await;
drop(lease);
let _ = holder_recorder.set(cx.now().as_millis());
})
.expect("spawn holder");
runtime.scheduler.lock().schedule(holder, 0);
let waiter_pool = Arc::clone(&pool);
let (waiter, _) = runtime
.state
.create_task(region, Budget::INFINITE, async move {
let cx = Cx::current().expect("lab task has a cx");
let outcome = waiter_pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(2)) })
.await;
let instant_ms = cx.now().as_millis();
let timed_out = matches!(
&outcome,
Outcome::Err(Error::Pool(PoolError {
kind: PoolErrorKind::Timeout,
..
}))
);
let _ = waiter_recorder.set((instant_ms, timed_out));
})
.expect("spawn waiter");
runtime.scheduler.lock().schedule(waiter, 0);
runtime.run_with_auto_advance();
assert_eq!(
waiter_result.get().copied(),
Some((200, true)),
"acquire timeout fires at the exact virtual deadline"
);
assert_eq!(pool.stats().timeouts, 1);
assert_eq!(holder_result.get().copied(), Some(10_000));
}
#[test]
fn lab_unhealthy_connection_is_replaced_on_the_first_checkout_after_failure() {
let (mut runtime, pool, region) = lab_pool(
PoolConfig::new(2)
.test_on_checkout(true)
.idle_timeout(100_000)
.acquire_timeout(5_000),
);
let flag = Arc::new(AtomicBool::new(false));
let checkout = Arc::new(OnceCell::<(u32, u64, u64)>::new());
let checkout_recorder = Arc::clone(&checkout);
let ping_flag = Arc::clone(&flag);
let (task, _) = runtime
.state
.create_task(region, Budget::INFINITE, async move {
let cx = Cx::current().expect("lab task has a cx");
let factory_flag = Arc::clone(&ping_flag);
let a = pool
.acquire(&cx, move || {
let flag_for_conn = Arc::clone(&factory_flag);
async move {
Outcome::Ok(MockConnection::with_ping_behavior(1, flag_for_conn))
}
})
.await
.expect("acquire healthy a");
drop(a);
ping_flag.store(true, Ordering::Relaxed);
asupersync::time::sleep(cx.now(), Duration::from_millis(1)).await;
let b = pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(2)) })
.await
.expect("acquire after failure");
let stats = pool.stats();
let _ = checkout_recorder.set((
b.id,
stats.connections_created,
stats.connections_closed,
));
})
.expect("spawn lab task");
runtime.scheduler.lock().schedule(task, 0);
runtime.run_with_auto_advance();
assert_eq!(
checkout.get().copied(),
Some((2, 2, 1)),
"first checkout after the health flip gets a replacement"
);
}
#[test]
fn lab_close_and_drain_wakes_waiters_immediately_and_completes_at_last_release() {
let (mut runtime, pool, region) = lab_pool(
PoolConfig::new(2)
.test_on_checkout(false)
.acquire_timeout(5_000),
);
let holder_done = Arc::new(OnceCell::<u64>::new());
let holder_recorder = Arc::clone(&holder_done);
let drainer_done = Arc::new(OnceCell::<(u64, bool)>::new());
let drainer_recorder = Arc::clone(&drainer_done);
let waiter_done = Arc::new(OnceCell::<(u64, bool)>::new());
let waiter_recorder = Arc::clone(&waiter_done);
let holder_pool = Arc::clone(&pool);
let (holder, _) = runtime
.state
.create_task(region, Budget::INFINITE, async move {
let cx = Cx::current().expect("lab task has a cx");
let lease = holder_pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(1)) })
.await
.expect("holder acquires");
asupersync::time::sleep(cx.now(), Duration::from_millis(10_000)).await;
drop(lease);
let _ = holder_recorder.set(cx.now().as_millis());
})
.expect("spawn holder");
runtime.scheduler.lock().schedule(holder, 0);
let drainer_pool = Arc::clone(&pool);
let (drainer, _) = runtime
.state
.create_task(region, Budget::INFINITE, async move {
let cx = Cx::current().expect("lab task has a cx");
let outcome = drainer_pool.close_and_drain(&cx).await;
let _ = drainer_recorder.set((cx.now().as_millis(), outcome.is_ok()));
})
.expect("spawn drainer");
runtime.scheduler.lock().schedule(drainer, 0);
let waiter_pool = Arc::clone(&pool);
let (waiter, _) = runtime
.state
.create_task(region, Budget::INFINITE, async move {
let cx = Cx::current().expect("lab task has a cx");
let outcome = waiter_pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(2)) })
.await;
let instant_ms = cx.now().as_millis();
let closed = matches!(
&outcome,
Outcome::Err(Error::Pool(PoolError {
kind: PoolErrorKind::Closed,
..
}))
);
let _ = waiter_recorder.set((instant_ms, closed));
})
.expect("spawn waiter");
runtime.run_until_idle();
assert!(waiter_done.get().is_none(), "waiter has not run yet");
assert!(
drainer_done.get().is_none(),
"drain waits for the active lease"
);
runtime.scheduler.lock().schedule(waiter, 0);
runtime.run_until_idle();
assert_eq!(
waiter_done.get().copied(),
Some((0, true)),
"waiter observes Closed immediately at close"
);
runtime.advance_to_next_timer();
runtime.run_until_idle();
assert_eq!(holder_done.get().copied(), Some(10_000));
assert_eq!(
drainer_done.get().copied(),
Some((10_000, true)),
"drain completes exactly at the last release"
);
}
#[derive(Default)]
struct Violations {
list: std::sync::Mutex<Vec<String>>,
}
impl Violations {
fn lock(&self) -> std::sync::MutexGuard<'_, Vec<String>> {
self.list
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
fn record(&self, message: String) {
self.lock().push(message);
}
fn take(&self, label: &str, seeds: &[u64]) {
let violations = self.lock();
assert!(
violations.is_empty(),
"{label}: {} oracle violation(s) across explored schedules (seeds {seeds:?}):\n{}",
violations.len(),
violations.join("\n")
);
}
}
fn exploration_seed(base: u64, runs: usize) -> (u64, usize) {
match std::env::var("SQLMODEL_DPOR_SEED") {
Ok(seed) => {
let seed = seed.parse::<u64>().unwrap_or_else(|error| {
panic!("SQLMODEL_DPOR_SEED must be an integer: {error}")
});
(seed, 1)
}
Err(_) => (base, runs),
}
}
#[test]
fn dpor_pool_close_and_drain_never_loses_wakeups_or_resurrects() {
let (base_seed, runs) = exploration_seed(0x50F1_0001, 48);
let mut explorer = DporExplorer::new(
ExplorerConfig::new(base_seed, runs)
.worker_count(1)
.max_steps(2_000),
);
let violations = Arc::new(Violations::default());
let report = explorer.explore(|runtime| {
let pool = Arc::new(Pool::with_timer_driver(
PoolConfig::new(2)
.test_on_checkout(false)
.acquire_timeout(60_000),
runtime
.state
.timer_driver_handle()
.expect("lab runtime timer driver"),
));
let region = runtime.state.create_root_region(Budget::INFINITE);
let violations = Arc::new(Violations::default());
let waiter_outcome = Arc::new(AtomicU64::new(0));
let ok_after_close = Arc::new(AtomicU64::new(0));
let holders_done = Arc::new(AtomicU64::new(0));
let waiter_done = Arc::new(AtomicBool::new(false));
let drain_done = Arc::new(AtomicBool::new(false));
for holder in 0..2u64 {
let (pool, violations, holders_done, ok_after_close) = (
Arc::clone(&pool),
Arc::clone(&violations),
Arc::clone(&holders_done),
Arc::clone(&ok_after_close),
);
let hold_ms = if holder == 0 { 5 } else { 3 };
let (task, _) = runtime
.state
.create_task(region, Budget::INFINITE, async move {
let cx = Cx::current().expect("lab task cx");
match pool
.acquire(&cx, || async {
Outcome::Ok(MockConnection::new(
u32::try_from(holder).expect("holder fits u32") + 1,
))
})
.await
{
Outcome::Ok(lease) => {
if pool.is_closed() {
ok_after_close.fetch_add(1, Ordering::Relaxed);
}
asupersync::time::sleep(cx.now(), Duration::from_millis(hold_ms))
.await;
drop(lease);
}
other => violations.record(format!(
"holder {holder}: acquire before close returned {other:?}"
)),
}
holders_done.fetch_add(1, Ordering::Relaxed);
})
.expect("spawn holder");
runtime.scheduler.lock().schedule(task, 0);
}
{
let (pool, violations, ok_after_close) = (
Arc::clone(&pool),
Arc::clone(&violations),
Arc::clone(&ok_after_close),
);
let (waiter_outcome, waiter_done) =
(Arc::clone(&waiter_outcome), Arc::clone(&waiter_done));
let (task, _) = runtime
.state
.create_task(region, Budget::INFINITE, async move {
let cx = Cx::current().expect("lab task cx");
asupersync::time::sleep(cx.now(), Duration::from_millis(1)).await;
match pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(9)) })
.await
{
Outcome::Ok(lease) => {
if pool.is_closed() {
ok_after_close.fetch_add(1, Ordering::Relaxed);
}
drop(lease);
waiter_outcome.store(1, Ordering::Relaxed);
}
Outcome::Err(Error::Pool(PoolError {
kind: PoolErrorKind::Closed,
..
})) => waiter_outcome.store(2, Ordering::Relaxed),
other => {
waiter_outcome.store(3, Ordering::Relaxed);
violations.record(format!("waiter: unexpected outcome {other:?}"));
}
}
waiter_done.store(true, Ordering::Relaxed);
})
.expect("spawn waiter");
runtime.scheduler.lock().schedule(task, 0);
}
{
let (pool, violations, drain_done) = (
Arc::clone(&pool),
Arc::clone(&violations),
Arc::clone(&drain_done),
);
let (task, _) = runtime
.state
.create_task(region, Budget::INFINITE, async move {
let cx = Cx::current().expect("lab task cx");
let outcome = pool.close_and_drain(&cx).await;
if !matches!(outcome, Outcome::Ok(())) {
violations.record(format!("drain: expected Ok, got {outcome:?}"));
}
drain_done.store(true, Ordering::Relaxed);
})
.expect("spawn drainer");
runtime.scheduler.lock().schedule(task, 0);
}
runtime.run_with_auto_advance();
if holders_done.load(Ordering::Relaxed) != 2 {
violations.record(format!(
"holders did not all finish (hang or step exhaustion): {}",
holders_done.load(Ordering::Relaxed)
));
}
if !waiter_done.load(Ordering::Relaxed) {
violations
.record("waiter did not finish (blocked forever after close?)".to_owned());
}
if waiter_outcome.load(Ordering::Relaxed) == 0 {
violations.record("waiter outcome never recorded".to_owned());
}
if !drain_done.load(Ordering::Relaxed) {
violations.record("close_and_drain never completed".to_owned());
}
if !pool.is_closed() {
violations.record("pool not closed after drain".to_owned());
}
if ok_after_close.load(Ordering::Relaxed) > 0 {
violations.record("a lease was handed out after the pool closed".to_owned());
}
let stats = pool.stats();
if stats.total_connections != 0
|| stats.idle_connections != 0
|| stats.active_connections != 0
{
violations.record(format!("stats: pool did not drain empty: {stats:?}"));
}
if stats.connections_created != 2 || stats.connections_closed != 2 {
violations.record(format!(
"stats: created/closed unbalanced (2 expected): {stats:?}"
));
}
});
eprintln!(
"dpor pool drain: {} runs, {} schedule classes",
report.total_runs, report.unique_classes
);
assert!(
report.total_runs >= 8,
"exploration must actually sweep schedules, ran {}",
report.total_runs
);
let seeds = report.violation_seeds();
violations.take("dpor pool close_and_drain", &seeds);
assert!(
!report.has_violations(),
"lab runtime invariants violated (seeds {seeds:?})"
);
}
struct FlakyBegin {
inner: MockConnection,
failures_left: Arc<AtomicU64>,
}
impl FlakyBegin {
fn new(failures: u64) -> Self {
Self {
inner: MockConnection::new(70 + u32::try_from(failures).expect("failures fit u32")),
failures_left: Arc::new(AtomicU64::new(failures)),
}
}
}
impl Connection for FlakyBegin {
type Tx<'conn>
= <MockConnection as Connection>::Tx<'conn>
where
Self: 'conn;
fn dialect(&self) -> Dialect {
self.inner.dialect()
}
async fn begin_with(
&self,
cx: &Cx,
isolation: IsolationLevel,
) -> Outcome<Self::Tx<'_>, Error> {
if self.failures_left.load(Ordering::Relaxed) > 0 {
self.failures_left.fetch_sub(1, Ordering::Relaxed);
return Outcome::Err(Error::Query(QueryError {
kind: QueryErrorKind::Serialization,
sql: Some("BEGIN".to_owned()),
sqlstate: Some("40001".to_owned()),
message: "injected serialization failure".to_owned(),
detail: None,
hint: None,
position: None,
source: None,
}));
}
self.inner.begin_with(cx, isolation).await
}
async fn query(&self, cx: &Cx, sql: &str, params: &[Value]) -> Outcome<Vec<Row>, Error> {
self.inner.query(cx, sql, params).await
}
async fn query_one(
&self,
cx: &Cx,
sql: &str,
params: &[Value],
) -> Outcome<Option<Row>, Error> {
self.inner.query_one(cx, sql, params).await
}
async fn execute(&self, cx: &Cx, sql: &str, params: &[Value]) -> Outcome<u64, Error> {
self.inner.execute(cx, sql, params).await
}
async fn insert(&self, cx: &Cx, sql: &str, params: &[Value]) -> Outcome<i64, Error> {
self.inner.insert(cx, sql, params).await
}
async fn batch(
&self,
cx: &Cx,
statements: &[(String, Vec<Value>)],
) -> Outcome<Vec<u64>, Error> {
self.inner.batch(cx, statements).await
}
async fn begin(&self, cx: &Cx) -> Outcome<Self::Tx<'_>, Error> {
self.begin_with(cx, IsolationLevel::ReadCommitted).await
}
fn supports_transaction_mode(&self, mode: TransactionMode) -> bool {
self.inner.supports_transaction_mode(mode)
}
async fn prepare(&self, cx: &Cx, sql: &str) -> Outcome<PreparedStatement, Error> {
self.inner.prepare(cx, sql).await
}
async fn query_prepared(
&self,
cx: &Cx,
stmt: &PreparedStatement,
params: &[Value],
) -> Outcome<Vec<Row>, Error> {
self.inner.query_prepared(cx, stmt, params).await
}
async fn execute_prepared(
&self,
cx: &Cx,
stmt: &PreparedStatement,
params: &[Value],
) -> Outcome<u64, Error> {
self.inner.execute_prepared(cx, stmt, params).await
}
async fn ping(&self, cx: &Cx) -> Outcome<(), Error> {
self.inner.ping(cx).await
}
async fn close(self, cx: &Cx) -> Result<(), Error> {
self.inner.close(cx).await
}
async fn close_for_pool(self, cx: &Cx) -> Result<(), Error>
where
Self: Sized,
{
self.inner.close_for_pool(cx).await
}
}
#[test]
fn dpor_retry_combinator_callers_always_terminate() {
for (failures, expect_ok) in [(2, true), (11, false)] {
let (base_seed, runs) = exploration_seed(0x50F1_0002 + failures, 32);
let mut explorer = DporExplorer::new(
ExplorerConfig::new(base_seed, runs)
.worker_count(1)
.max_steps(4_000),
);
let violations = Arc::new(Violations::default());
let report = explorer.explore(|runtime| {
let conn = Arc::new(FlakyBegin::new(failures));
let region = runtime.state.create_root_region(Budget::INFINITE);
let violations = Arc::new(Violations::default());
let callers_done = Arc::new(AtomicU64::new(0));
let oks = Arc::new(AtomicU64::new(0));
let exhausted = Arc::new(AtomicU64::new(0));
for caller in 0..2u64 {
let (conn, violations, callers_done, oks, exhausted) = (
Arc::clone(&conn),
Arc::clone(&violations),
Arc::clone(&callers_done),
Arc::clone(&oks),
Arc::clone(&exhausted),
);
let (task, _) = runtime
.state
.create_task(region, Budget::INFINITE, async move {
let cx = Cx::current().expect("lab task cx");
let outcome = retry_transaction(
&cx,
conn.as_ref(),
TransactionOptions::new(),
&RetryPolicy::default(),
async |cx: &Cx, tx| {
tx.execute(cx, "INSERT INTO t VALUES (1)", &[])
.await
.map(|_| ())
},
)
.await;
match outcome {
Outcome::Ok(()) => {
oks.fetch_add(1, Ordering::Relaxed);
}
Outcome::Err(Error::Transaction(e))
if e.kind == TransactionErrorKind::RetriesExhausted =>
{
exhausted.fetch_add(1, Ordering::Relaxed);
}
Outcome::Err(e) => violations.record(format!(
"retry caller {caller}: unexpected Err {e:?}"
)),
Outcome::Cancelled(r) => violations.record(format!(
"retry caller {caller}: unexpected Cancelled {r:?}"
)),
Outcome::Panicked(p) => violations.record(format!(
"retry caller {caller}: panicked {p:?}"
)),
}
callers_done.fetch_add(1, Ordering::Relaxed);
})
.expect("spawn retry caller");
runtime.scheduler.lock().schedule(task, 0);
}
runtime.run_with_auto_advance();
if callers_done.load(Ordering::Relaxed) != 2 {
violations.record(format!(
"retry callers did not finish (hung future?): {}",
callers_done.load(Ordering::Relaxed)
));
}
let oks = oks.load(Ordering::Relaxed);
let exhausted = exhausted.load(Ordering::Relaxed);
if expect_ok && oks != 2 {
violations.record(format!(
"with {failures} injected failures both callers must eventually commit, got {oks} ok"
));
}
if !expect_ok && oks + exhausted != 2 {
violations.record(format!(
"with {failures} injected failures every caller must end Ok or exhausted, got {oks} ok / {exhausted} exhausted"
));
}
});
eprintln!(
"dpor retry ({failures} injected failures): {} runs, {} classes",
report.total_runs, report.unique_classes
);
assert!(
report.total_runs >= 8,
"exploration must actually sweep schedules, ran {}",
report.total_runs
);
let seeds = report.violation_seeds();
violations.take(
&format!("dpor retry ({failures} injected failures)"),
&seeds,
);
assert!(
!report.has_violations(),
"lab runtime invariants violated (seeds {seeds:?})"
);
let _ = expect_ok;
}
}
#[test]
fn lab_acquire_budget_deadline_wins_over_acquire_timeout() {
let mut runtime = LabRuntime::new(LabConfig::new(0x50F1_0004).max_steps(200_000));
let driver = runtime
.state
.timer_driver_handle()
.expect("lab runtime timer driver");
let pool = Arc::new(Pool::with_timer_driver(
PoolConfig::new(1)
.acquire_timeout(5_000)
.test_on_checkout(false),
driver,
));
let outcome_record = Arc::new(std::sync::Mutex::new(None::<(u64, bool, bool)>));
let region = runtime.state.create_root_region(Budget::INFINITE);
let holder_pool = Arc::clone(&pool);
let (holder_task, _) = runtime
.state
.create_task(region, Budget::INFINITE, async move {
let cx = Cx::current().expect("lab task cx");
let lease = holder_pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(1)) })
.await
.expect("holder acquires the only lease");
asupersync::time::sleep(cx.now(), Duration::from_secs(30)).await;
drop(lease);
})
.expect("spawn holder");
runtime.scheduler.lock().schedule(holder_task, 0);
let waiter_pool = Arc::clone(&pool);
let waiter_budget = Budget::new().with_deadline(Time::from_millis(2_000));
let outcome_writer = Arc::clone(&outcome_record);
let (waiter_task, _) = runtime
.state
.create_task(region, waiter_budget, async move {
let cx = Cx::current().expect("lab task cx");
let outcome = waiter_pool
.acquire(&cx, || async { Outcome::Ok(MockConnection::new(2)) })
.await;
let instant = cx.now().as_millis();
match &outcome {
Outcome::Err(Error::Pool(PoolError {
kind: PoolErrorKind::Timeout,
message,
..
})) => {
let budget_limited = message.contains("budget deadline");
*outcome_writer.lock().unwrap() = Some((instant, true, budget_limited));
}
_ => *outcome_writer.lock().unwrap() = Some((instant, false, false)),
}
})
.expect("spawn waiter");
runtime.scheduler.lock().schedule(waiter_task, 0);
runtime.run_with_auto_advance();
let (instant, timed_out, budget_limited) = outcome_record
.lock()
.unwrap()
.expect("waiter recorded an outcome");
assert_eq!(instant, 2_000, "budget deadline fires at exactly T+2000ms");
assert!(timed_out, "waiter must observe PoolErrorKind::Timeout");
assert!(
budget_limited,
"the timeout message must attribute the budget"
);
assert_eq!(pool.stats().timeouts, 1);
}
}