WIP refactor: replace sharding with single connection set

Author: abonander

sqlx-core/src/ext/future.rs

@@ -0,0 +1,38 @@
+use pin_project_lite::pin_project;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+pin_project! {
+    #[project = RaceProject]
+    pub struct Race<L, R> {
+        #[pin]
+        left: L,
+        #[pin]
+        right: R,
+    }
+}
+
+impl<L, R> Future for Race<L, R>
+where
+    L: Future,
+    R: Future,
+{
+    type Output = Result<L::Output, R::Output>;
+
+    #[inline(always)]
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let mut this = self.project();
+
+        if let Poll::Ready(left) = this.left.as_mut().poll(cx) {
+            return Poll::Ready(Ok(left));
+        }
+
+        this.right.as_mut().poll(cx).map(Err)
+    }
+}
+
+#[inline(always)]
+pub fn race<L, R>(left: L, right: R) -> Race<L, R> {
+    Race { left, right }
+}

sqlx-core/src/ext/mod.rs

@@ -2,3 +2,5 @@ pub mod ustr;
 
 #[macro_use]
 pub mod async_stream;
+
+pub mod future;

sqlx-core/src/pool/connect.rs

@@ -14,7 +14,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex, RwLock};
 use std::time::Instant;
 
-use crate::pool::shard::DisconnectedSlot;
+use crate::pool::connection_set::DisconnectedSlot;
 #[cfg(doc)]
 use crate::pool::PoolOptions;
 use crate::sync::{AsyncMutex, AsyncMutexGuard};
@@ -646,26 +646,24 @@ async fn connect_with_backoff<DB: Database>(
 
         match res {
             ControlFlow::Break(Ok(conn)) => {
-                tracing::trace!(
+                tracing::debug!(
                     target: "sqlx::pool::connect",
                     %connection_id,
                     attempt,
                     elapsed_seconds,
                     "connection established",
                 );
 
-                return Ok(PoolConnection::new(
-                    slot.put(ConnectionInner {
-                        raw: conn,
-                        id: connection_id,
-                        created_at: now,
-                        last_released_at: now,
-                    }),
-                    pool.0.clone(),
-                ));
+                return Ok(PoolConnection::new(slot.put(ConnectionInner {
+                    pool: Arc::downgrade(&pool.0),
+                    raw: conn,
+                    id: connection_id,
+                    created_at: now,
+                    last_released_at: now,
+                })));
             }
             ControlFlow::Break(Err(e)) => {
-                tracing::warn!(
+                tracing::error!(
                     target: "sqlx::pool::connect",
                     %connection_id,
                     attempt,

sqlx-core/src/pool/connection.rs

@@ -1,33 +1,35 @@
 use std::fmt::{self, Debug, Formatter};
 use std::future::{self, Future};
+use std::io;
 use std::ops::{Deref, DerefMut};
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 use std::time::{Duration, Instant};
 
 use crate::connection::Connection;
 use crate::database::Database;
 use crate::error::Error;
 
-use super::inner::{is_beyond_max_lifetime, PoolInner};
+use super::inner::PoolInner;
 use crate::pool::connect::{ConnectPermit, ConnectTaskShared, ConnectionId};
+use crate::pool::connection_set::{ConnectedSlot, DisconnectedSlot};
 use crate::pool::options::PoolConnectionMetadata;
-use crate::pool::shard::{ConnectedSlot, DisconnectedSlot};
-use crate::pool::Pool;
+use crate::pool::{Pool, PoolOptions};
 use crate::rt;
 
 const RETURN_TO_POOL_TIMEOUT: Duration = Duration::from_secs(5);
-const CLOSE_ON_DROP_TIMEOUT: Duration = Duration::from_secs(5);
+const CLOSE_TIMEOUT: Duration = Duration::from_secs(5);
 
 /// A connection managed by a [`Pool`][crate::pool::Pool].
 ///
 /// Will be returned to the pool on-drop.
 pub struct PoolConnection<DB: Database> {
     conn: Option<ConnectedSlot<ConnectionInner<DB>>>,
-    pub(crate) pool: Arc<PoolInner<DB>>,
     close_on_drop: bool,
 }
 
 pub(super) struct ConnectionInner<DB: Database> {
+    // Note: must be `Weak` to prevent a reference cycle
+    pub(crate) pool: Weak<PoolInner<DB>>,
     pub(super) raw: DB::Connection,
     pub(super) id: ConnectionId,
     pub(super) created_at: Instant,
@@ -72,11 +74,10 @@ impl<DB: Database> AsMut<DB::Connection> for PoolConnection<DB> {
 }
 
 impl<DB: Database> PoolConnection<DB> {
-    pub(super) fn new(live: ConnectedSlot<ConnectionInner<DB>>, pool: Arc<PoolInner<DB>>) -> Self {
+    pub(super) fn new(live: ConnectedSlot<ConnectionInner<DB>>) -> Self {
         Self {
             conn: Some(live),
             close_on_drop: false,
-            pool,
         }
     }
 
@@ -140,13 +141,16 @@ impl<DB: Database> PoolConnection<DB> {
     #[doc(hidden)]
     pub fn return_to_pool(&mut self) -> impl Future<Output = ()> + Send + 'static {
         let conn = self.conn.take();
-        let pool = self.pool.clone();
 
         async move {
             let Some(conn) = conn else {
                 return;
             };
 
+            let Some(pool) = Weak::upgrade(&conn.pool) else {
+                return;
+            };
+
             rt::timeout(RETURN_TO_POOL_TIMEOUT, return_to_pool(conn, &pool))
                 .await
                 // Dropping of the `slot` will check if the connection must be re-established
@@ -161,7 +165,7 @@ impl<DB: Database> PoolConnection<DB> {
         async move {
             if let Some(conn) = conn {
                 // Don't hold the connection forever if it hangs while trying to close
-                rt::timeout(CLOSE_ON_DROP_TIMEOUT, close(conn)).await.ok();
+                rt::timeout(CLOSE_TIMEOUT, close(conn)).await.ok();
             }
         }
     }
@@ -195,7 +199,7 @@ impl<DB: Database> Drop for PoolConnection<DB> {
         }
 
         // We still need to spawn a task to maintain `min_connections`.
-        if self.conn.is_some() || self.pool.options.min_connections > 0 {
+        if self.conn.is_some() {
             crate::rt::spawn(self.return_to_pool());
         }
     }
@@ -220,6 +224,48 @@ impl<DB: Database> ConnectionInner<DB> {
             idle_for: now.saturating_duration_since(self.last_released_at),
         }
     }
+
+    pub fn is_beyond_max_lifetime(&self, options: &PoolOptions<DB>) -> bool {
+        if let Some(max_lifetime) = options.max_lifetime {
+            let age = self.created_at.elapsed();
+
+            if age > max_lifetime {
+                tracing::info!(
+                    target: "sqlx::pool",
+                    connection_id=%self.id,
+                    ?age,
+                    "connection is beyond `max_lifetime`, closing"
+                );
+
+                return true;
+            }
+        }
+
+        false
+    }
+
+    pub fn is_beyond_idle_timeout(&self, options: &PoolOptions<DB>) -> bool {
+        if let Some(idle_timeout) = options.idle_timeout {
+            let now = Instant::now();
+
+            let age = now.duration_since(self.created_at);
+            let idle_duration = now.duration_since(self.last_released_at);
+
+            if idle_duration > idle_timeout {
+                tracing::info!(
+                    target: "sqlx::pool",
+                    connection_id=%self.id,
+                    ?age,
+                    ?idle_duration,
+                    "connection is beyond `idle_timeout`, closing"
+                );
+
+                return true;
+            }
+        }
+
+        false
+    }
 }
 
 pub(crate) async fn close<DB: Database>(
@@ -231,14 +277,19 @@ pub(crate) async fn close<DB: Database>(
 
     let (conn, slot) = ConnectedSlot::take(conn);
 
-    let res = conn.raw.close().await.inspect_err(|error| {
-        tracing::debug!(
-            target: "sqlx::pool",
-            %connection_id,
-            %error,
-            "error occurred while closing the pool connection"
-        );
-    });
+    let res = rt::timeout(CLOSE_TIMEOUT, conn.raw.close())
+        .await
+        .unwrap_or_else(|_| {
+            Err(io::Error::new(io::ErrorKind::TimedOut, "timed out sending close packet").into())
+        })
+        .inspect_err(|error| {
+            tracing::debug!(
+                target: "sqlx::pool",
+                %connection_id,
+                %error,
+                "error occurred while closing the pool connection"
+            );
+        });
 
     (res, slot)
 }
@@ -255,14 +306,19 @@ pub(crate) async fn close_hard<DB: Database>(
 
     let (conn, slot) = ConnectedSlot::take(conn);
 
-    let res = conn.raw.close_hard().await.inspect_err(|error| {
-        tracing::debug!(
-            target: "sqlx::pool",
-            %connection_id,
-            %error,
-            "error occurred while closing the pool connection"
-        );
-    });
+    let res = rt::timeout(CLOSE_TIMEOUT, conn.raw.close_hard())
+        .await
+        .unwrap_or_else(|_| {
+            Err(io::Error::new(io::ErrorKind::TimedOut, "timed out sending close packet").into())
+        })
+        .inspect_err(|error| {
+            tracing::debug!(
+                target: "sqlx::pool",
+                %connection_id,
+                %error,
+                "error occurred while closing the pool connection"
+            );
+        });
 
     (res, slot)
 }
@@ -282,7 +338,7 @@ async fn return_to_pool<DB: Database>(
 
     // If the connection is beyond max lifetime, close the connection and
     // immediately create a new connection
-    if is_beyond_max_lifetime(&conn, &pool.options) {
+    if conn.is_beyond_max_lifetime(&pool.options) {
         let (_res, slot) = close(conn).await;
         return Err(slot);
     }
@@ -314,6 +370,7 @@ async fn return_to_pool<DB: Database>(
     // to recover from cancellations
     if let Err(error) = conn.raw.ping().await {
         tracing::warn!(
+            target: "sqlx::pool",
             %error,
             "error occurred while testing the connection on-release",
         );