fix(pool): tweaks and fixes

Author: abonander

sqlx-core/src/pool/connect.rs

@@ -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;
@@ -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;
         }
     }
 
@@ -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")

sqlx-core/src/pool/idle.rs

@@ -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
@@ -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) {

sqlx-core/src/pool/inner.rs

@@ -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;
 
@@ -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,
+                    }
                 }
             }
         }