fix(pool): don't stop emptying idle queue in .close()

Author: abonander

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::FutureExt;
+use futures_util::{select, FutureExt};
 use std::time::{Duration, Instant};
 use tracing::Level;
 
@@ -75,12 +75,18 @@ impl<DB: Database> PoolInner<DB> {
     pub(super) fn close<'a>(self: &'a Arc<Self>) -> impl Future<Output = ()> + 'a {
         self.mark_closed();
 
+        // Keep clearing the idle queue as connections are released until the count reaches zero.
         async move {
-            while let Some(idle) = self.idle.try_acquire(self) {
-                idle.close().await;
+            let mut drained = pin!(self.counter.drain());
+
+            loop {
+                select! {
+                    idle = self.idle.acquire(self) => {
+                        idle.close().await;
+                    },
+                    () = drained.as_mut() => break,
+                }
             }
-
-            self.counter.drain().await;
         }
     }
 
@@ -116,7 +122,7 @@ impl<DB: Database> PoolInner<DB> {
         let acquire_started_at = Instant::now();
 
         let mut close_event = pin!(self.close_event());
-        let mut deadline = pin!(crate::rt::sleep(self.options.acquire_timeout));
+        let mut deadline = pin!(rt::sleep(self.options.acquire_timeout));
         let mut acquire_idle = pin!(self.idle.acquire(self).fuse());
         let mut before_acquire = OptionFuture::from(None);
         let mut acquire_connect_permit = pin!(OptionFuture::from(Some(
@@ -130,6 +136,9 @@ impl<DB: Database> PoolInner<DB> {
         // * If we acquire a `ConnectPermit`, we begin the connection loop (with backoff)
         //   as implemented by `DynConnector`.
         // * If we acquire an idle connection, we then start polling `check_idle_conn()`.
+        //
+        // This doesn't quite fit into `select!{}` because the set of futures that may be polled
+        // at a given time is dynamic, so it's actually simpler to hand-roll it.
         let acquired = future::poll_fn(|cx| {
             use std::task::Poll::*;