Skip to content

Commit

Permalink
fix(pool): tweaks and fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
abonander committed Nov 11, 2024
1 parent 80831e8 commit 8420b43
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 11 deletions.
10 changes: 6 additions & 4 deletions sqlx-core/src/pool/connect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use crate::pool::PoolConnection;
use crate::rt::JoinHandle;
use crate::Error;
use ease_off::EaseOff;
use event_listener::Event;
use event_listener::{listener, Event};
use std::fmt::{Display, Formatter};
use std::future::Future;
use std::ptr;
Expand Down Expand Up @@ -318,7 +318,8 @@ impl ConnectionCounter {

pub async fn drain(&self) {
while self.count.load(Ordering::Acquire) > 0 {
self.connect_available.listen().await;
listener!(self.connect_available => permit_released);
permit_released.await;
}
}

Expand Down Expand Up @@ -386,13 +387,14 @@ impl ConnectionCounter {
return acquired;
}

self.connect_available.listen().await;

if attempt == 2 {
tracing::warn!(
"unable to acquire a connect permit after sleeping; this may indicate a bug"
);
}

listener!(self.connect_available => connect_available);
connect_available.await;
}

panic!("BUG: was never able to acquire a connection despite waking many times")
Expand Down
5 changes: 4 additions & 1 deletion sqlx-core/src/pool/idle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ use futures_util::FutureExt;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

use event_listener::listener;

pub struct IdleQueue<DB: Database> {
queue: ArrayQueue<Idle<DB>>,
// Keep a separate count because `ArrayQueue::len()` loops until the head and tail pointers
Expand Down Expand Up @@ -36,7 +38,8 @@ impl<DB: Database> IdleQueue<DB> {

for attempt in 1usize.. {
if should_wait {
self.release_event.listen().await;
listener!(self.release_event => release_event);
release_event.await;
}

if let Some(conn) = self.try_acquire(pool) {
Expand Down
17 changes: 11 additions & 6 deletions sqlx-core/src/pool/inner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use crate::rt::JoinHandle;
use crate::{private_tracing_dynamic_event, rt};
use either::Either;
use futures_util::future::{self, OptionFuture};
use futures_util::{select, FutureExt};
use futures_util::{FutureExt};
use std::time::{Duration, Instant};
use tracing::Level;

Expand Down Expand Up @@ -77,14 +77,19 @@ impl<DB: Database> PoolInner<DB> {

// Keep clearing the idle queue as connections are released until the count reaches zero.
async move {
let mut drained = pin!(self.counter.drain()).fuse();
let mut drained = pin!(self.counter.drain());

loop {
select! {
idle = self.idle.acquire(self) => {
let mut acquire_idle = pin!(self.idle.acquire(self));

// Not using `futures::select!{}` here because it requires a proc-macro dep,
// and frankly it's a little broken.
match future::select(drained.as_mut(), acquire_idle.as_mut()).await {
// *not* `either::Either`; they rolled their own
future::Either::Left(_) => break,
future::Either::Right((idle, _)) => {
idle.close().await;
},
() = drained.as_mut() => break,
}
}
}
}
Expand Down

0 comments on commit 8420b43

Please sign in to comment.