WIP refactor: replace sharding with single connection set (6)

Author: abonander

benches/any/pool.rs

@@ -3,7 +3,8 @@ use criterion::{criterion_group, criterion_main, Bencher, BenchmarkId, Criterion
 use sqlx_core::any::AnyPoolOptions;
 use std::fmt::{Display, Formatter};
 use std::future::Future;
-use std::time::Instant;
+use std::thread;
+use std::time::{Duration, Instant};
 use tracing::Instrument;
 
 #[derive(Debug)]
@@ -131,7 +132,9 @@ fn bench_pool_with(b: &mut Bencher, input: &Input, database_url: &str) {
         }
     });
 
-    drop(pool.close());
+    runtime.block_on(pool.close());
+    // Give the server a second to clean up
+    thread::sleep(Duration::from_millis(50));
 }
 
 criterion_group!(benches, bench_pool,);

sqlx-core/src/pool/connection_set.rs

@@ -238,13 +238,13 @@ impl<C> ConnectionSet<C> {
             closing.push(async {
                 let locked = slot.lock().await;
 
-                let DisconnectedSlot(mut guard) = match locked.try_connected() {
+                let slot = match locked.try_connected() {
                     Ok(connected) => close(connected).await,
                     Err(disconnected) => disconnected,
                 };
 
                 // The pool is shutting down; don't wake any tasks that might have been interested
-                guard.drop_without_notify();
+                slot.leak();
             });
         }
 
@@ -461,8 +461,9 @@ impl<C> DisconnectedSlot<C> {
         ConnectedSlot(self.0)
     }
 
-    pub fn leak(self) {
+    pub fn leak(mut self) {
         self.0.slot.leaked.store(true, Ordering::Release);
+        self.0.drop_without_notify();
     }
 }
 

sqlx-core/src/pool/inner.rs

@@ -79,10 +79,7 @@ impl<DB: Database> PoolInner<DB> {
 
         // Keep clearing the idle queue as connections are released until the count reaches zero.
         self.connections.drain(async |slot| {
-            let (conn, slot) = ConnectedSlot::take(slot);
-
-            let _ = rt::timeout(GRACEFUL_CLOSE_TIMEOUT, conn.raw.close()).await;
-
+            let (_res, slot) = connection::close(slot).await;
             slot
         })
     }