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

Author: abonander

sqlx-core/src/pool/connect.rs

@@ -663,7 +663,7 @@ async fn connect_with_backoff<DB: Database>(
                 })));
             }
             ControlFlow::Break(Err(e)) => {
-                tracing::warn!(
+                tracing::error!(
                     target: "sqlx::pool::connect",
                     %connection_id,
                     attempt,

sqlx-core/src/pool/connection_set.rs

@@ -127,6 +127,8 @@ impl<C> ConnectionSet<C> {
 
         let preferred_slot = current_thread_id() % self.slots.len();
 
+        tracing::trace!(preferred_slot, ?pref, "acquire_inner");
+
         // Always try to lock the connection associated with our thread ID
         let mut acquire_preferred = pin!(self.slots[preferred_slot].acquire(pref));
 
@@ -174,9 +176,8 @@ impl<C> ConnectionSet<C> {
                     }
 
                     listen_global.as_mut().set(self.global.listen(pref));
+                    continue;
                 }
-
-                continue;
             }
 
             return Poll::Pending;
@@ -263,15 +264,20 @@ impl<C> ConnectionSet<C> {
 
 impl AcquirePreference {
     #[inline(always)]
-    fn wants_connected(&self) -> bool {
-        matches!(self, Self::Connected | Self::Either)
+    fn wants_connected(&self, is_connected: bool) -> bool {
+        match (self, is_connected) {
+            (Self::Connected, true) => true,
+            (Self::Disconnected, false) => true,
+            (Self::Either, _) => true,
+            _ => false,
+        }
     }
 }
 
 impl<C> Slot<C> {
     #[inline(always)]
     fn matches_pref(&self, pref: AcquirePreference) -> bool {
-        !self.is_leaked() && self.is_connected() == pref.wants_connected()
+        !self.is_leaked() && pref.wants_connected(self.is_connected())
     }
 
     #[inline(always)]
@@ -308,21 +314,30 @@ impl<C> Slot<C> {
     async fn acquire(self: &Arc<Self>, pref: AcquirePreference) -> SlotGuard<C> {
         loop {
             if self.matches_pref(pref) {
+                tracing::trace!(slot_index=%self.index, "waiting for lock");
+
                 let locked = self.lock().await;
 
                 if locked.matches_pref(pref) {
                     return locked;
                 }
             }
 
-            let event = if pref.wants_connected() {
-                &self.unlock_event
-            } else {
-                &self.disconnect_event
-            };
-
-            listener!(event => listener);
-            listener.await;
+            match pref {
+                AcquirePreference::Connected => {
+                    listener!(self.unlock_event => listener);
+                    listener.await;
+                }
+                AcquirePreference::Disconnected => {
+                    listener!(self.disconnect_event => listener);
+                    listener.await
+                }
+                AcquirePreference::Either => {
+                    listener!(self.unlock_event => unlock_listener);
+                    listener!(self.disconnect_event => disconnect_listener);
+                    race(unlock_listener, disconnect_listener).await.ok();
+                }
+            }
         }
     }
 
@@ -374,7 +389,7 @@ impl<C> SlotGuard<C> {
 
     #[inline(always)]
     fn matches_pref(&self, pref: AcquirePreference) -> bool {
-        !self.slot.is_leaked() && self.is_connected() == pref.wants_connected()
+        !self.slot.is_leaked() && pref.wants_connected(self.is_connected())
     }
 
     #[inline(always)]

sqlx-core/src/pool/inner.rs

@@ -181,6 +181,8 @@ impl<DB: Database> PoolInner<DB> {
         self: &Arc<Self>,
         connect_shared: &mut Option<Arc<ConnectTaskShared>>,
     ) -> Result<PoolConnection<DB>, Error> {
+        tracing::trace!("waiting for any connection");
+
         let disconnected = match self.connections.acquire_any().await {
             Ok(conn) => match finish_acquire(self, conn).await {
                 Ok(conn) => return Ok(conn),

sqlx-core/src/pool/mod.rs

@@ -68,7 +68,7 @@ use crate::sql_str::SqlSafeStr;
 use crate::transaction::Transaction;
 use event_listener::EventListener;
 use futures_core::FusedFuture;
-
+use tracing::Instrument;
 pub use self::connect::{PoolConnectMetadata, PoolConnector};
 pub use self::connection::PoolConnection;
 use self::inner::PoolInner;
@@ -363,6 +363,7 @@ impl<DB: Database> Pool<DB> {
     pub fn acquire(&self) -> impl Future<Output = Result<PoolConnection<DB>, Error>> + 'static {
         let shared = self.0.clone();
         async move { shared.acquire().await }
+            .instrument(tracing::error_span!(target: "sqlx::pool", "acquire"))
     }
 
     /// Attempts to retrieve a connection from the pool if there is one available.

sqlx-test/src/lib.rs

@@ -2,13 +2,15 @@ use sqlx::pool::PoolOptions;
 use sqlx::{Connection, Database, Error, Pool};
 use std::env;
 use tracing_subscriber::EnvFilter;
+use tracing_subscriber::fmt::format::FmtSpan;
 
 pub fn setup_if_needed() {
     let _ = dotenvy::dotenv();
     let _ = tracing_subscriber::fmt::Subscriber::builder()
         .with_env_filter(EnvFilter::from_default_env())
-        .with_test_writer()
-        .finish();
+        .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE)
+        // .with_test_writer()
+        .try_init();
 }
 
 // Make a new connection

tests/postgres/postgres.rs

@@ -255,6 +255,10 @@ async fn it_works_with_cache_disabled() -> anyhow::Result<()> {
 
 #[sqlx_macros::test]
 async fn it_executes_with_pool() -> anyhow::Result<()> {
+    setup_if_needed();
+
+    tracing::info!("starting test");
+
     let pool = sqlx_test::pool::<Postgres>().await?;
 
     let rows = pool.fetch_all("SELECT 1; SElECT 2").await?;
@@ -1146,7 +1150,7 @@ async fn test_listener_try_recv_buffered() -> anyhow::Result<()> {
     assert!(listener.next_buffered().is_none());
 
     // Activate connection.
-    sqlx::query!("SELECT 1 AS one")
+    sqlx::query("SELECT 1 AS one")
         .fetch_all(&mut listener)
         .await?;
 
@@ -2086,6 +2090,7 @@ async fn test_issue_3052() {
 }
 
 #[sqlx_macros::test]
+#[cfg(feature = "chrono")]
 async fn test_bind_iter() -> anyhow::Result<()> {
     use sqlx::postgres::PgBindIterExt;
     use sqlx::types::chrono::{DateTime, Utc};