feat: add tracing-flame to benches/any-pool

Author: abonander

Cargo.lock

@@ -241,6 +241,7 @@ criterion = { version = "0.7.0", features = ["async_tokio"] }
 libsqlite3-sys = { version = "0.30.1" }
 
 tracing = "0.1.41"
+tracing-flame = "0.2.0"
 tracing-subscriber = "0.3.20"
 
 # If this is an unconditional dev-dependency then Cargo will *always* try to build `libsqlite3-sys`,

Cargo.toml

@@ -4,6 +4,10 @@ use std::fmt::{Display, Formatter};
 use std::thread;
 use std::time::{Duration, Instant};
 use tracing::Instrument;
+use tracing_flame::FlameLayer;
+use tracing_subscriber::{EnvFilter, Layer};
+use tracing_subscriber::layer::SubscriberExt;
+use tracing_subscriber::util::SubscriberInitExt;
 
 #[derive(Debug)]
 struct Input {
@@ -24,7 +28,27 @@ impl Display for Input {
 
 fn bench_pool(c: &mut Criterion) {
     sqlx::any::install_default_drivers();
-    tracing_subscriber::fmt::try_init().ok();
+
+    let _guard = if let Ok(path) = dotenvy::var("FLAMEGRAPH_OUT") {
+        let (layer, guard) = FlameLayer::with_file(&path)
+            .expect(&format!("error opening path {path:?} (`FLAMEGRAPH_OUT`)"));
+
+        tracing_subscriber::registry()
+            .with(
+                tracing_subscriber::fmt::layer().with_filter(EnvFilter::from_default_env())
+            )
+            .with(layer.with_threads_collapsed(true))
+            .try_init()
+            .ok();
+
+        tracing::info!("Writing flamegraph to {path:?}");
+
+        Some(guard)
+    } else {
+        tracing_subscriber::fmt::try_init().ok();
+
+        None
+    };
 
     let database_url = dotenvy::var("DATABASE_URL").expect("DATABASE_URL must be set");
 
@@ -136,4 +160,4 @@ fn bench_pool_with(b: &mut Bencher, input: &Input, database_url: &str) {
 }
 
 criterion_group!(benches, bench_pool,);
-criterion_main!(benches);
+criterion_main!(benches);

benches/any/pool.rs

@@ -13,6 +13,7 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::task::Poll;
 use std::time::Duration;
+use tracing::Instrument;
 
 pub struct ConnectionSet<C> {
     global: Arc<Global>,
@@ -120,22 +121,37 @@ impl<C> ConnectionSet<C> {
     }
 
     async fn acquire_inner(&self, pref: AcquirePreference) -> SlotGuard<C> {
-                let preferred_slot = current_thread_id() % self.slots.len();
-
-        tracing::trace!(preferred_slot, ?pref, "acquire_inner");
+        let preferred_slot = current_thread_id() % self.slots.len();
 
         // Always try to lock the connection associated with our thread ID
         let mut acquire_preferred = pin!(self.slots[preferred_slot].acquire(pref));
 
+        let alternate_slot = (preferred_slot + 547usize.wrapping_mul(
+            Arc::strong_count(&self.slots[preferred_slot].connection)
+        )) % self.slots.len();
+
+        let mut acquire_alternate = pin!(self.slots[alternate_slot].acquire(pref));
+
         let mut listen_global = pin!(self.global.listen(pref));
 
-        let mut yielded = false;
+        let mut yielded_1 = false;
+        let mut yielded_2 = false;
 
         std::future::poll_fn(|cx| {
             if let Poll::Ready(locked) = acquire_preferred.as_mut().poll(cx) {
                 return Poll::Ready(locked);
             }
 
+            if let Poll::Ready(locked) = acquire_alternate.as_mut().poll(cx) {
+                return Poll::Ready(locked);
+            }
+
+            // if !yielded_1 {
+            //     cx.waker().wake_by_ref();
+            //     yielded_1 = true;
+            //     return Poll::Pending;
+            // }
+
             if let Poll::Ready(slot) = listen_global.as_mut().poll(cx) {
                 if let Some(locked) = self.slots[slot].try_acquire(pref) {
                     return Poll::Ready(locked);
@@ -144,9 +160,9 @@ impl<C> ConnectionSet<C> {
                 listen_global.as_mut().set(self.global.listen(pref));
             }
 
-            if !yielded {
+            if !yielded_2 {
                 cx.waker().wake_by_ref();
-                yielded = true;
+                yielded_2 = true;
                 return Poll::Pending;
             }
 
@@ -156,6 +172,12 @@ impl<C> ConnectionSet<C> {
 
             Poll::Pending
         })
+            .instrument(tracing::trace_span!(
+                target: "sqlx::pool::connection_set",
+                "acquire_inner",
+                preferred_slot,
+                ?pref,
+            ))
         .await
     }
 
@@ -174,14 +196,24 @@ impl<C> ConnectionSet<C> {
     }
 
     fn try_acquire(&self, pref: AcquirePreference) -> Option<SlotGuard<C>> {
-        let mut search_slot = current_thread_id() % self.slots.len();
+        let preferred_slot = current_thread_id() % self.slots.len();
 
-        for _ in 0..self.slots.len() {
-            if let Some(locked) = self.slots[search_slot].try_acquire(pref) {
-                return Some(locked);
+        let (slots_before, slots_after) = self.slots.split_at(preferred_slot);
+
+        let (preferred_slot, slots_after) = slots_after.split_first().unwrap();
+
+        if let Some(locked) = preferred_slot.try_acquire(pref) {
+            return Some(locked);
+        }
+
+        for slot in slots_before.iter().chain(slots_after).rev() {
+            if self.global.locked_set[slot.index].load(Ordering::Relaxed) {
+                continue;
             }
 
-            search_slot = self.next_slot(search_slot);
+            if let Some(locked) = slot.try_acquire(pref) {
+                return Some(locked);
+            }
         }
 
         None
@@ -398,6 +430,7 @@ impl<C> SlotGuard<C> {
             let connected = locked.is_some();
             self.slot.set_is_connected(connected);
             self.slot.locked.store(false, Ordering::Release);
+            self.slot.global.locked_set[self.slot.index].store(false, Ordering::Relaxed);
             connected
         })
     }

sqlx-core/src/pool/connection_set.rs

@@ -20,7 +20,7 @@ use crate::{private_tracing_dynamic_event, rt};
 use event_listener::listener;
 use futures_util::future::{self};
 use std::time::{Duration, Instant};
-use tracing::Level;
+use tracing::{Instrument, Level};
 
 const GRACEFUL_CLOSE_TIMEOUT: Duration = Duration::from_secs(5);
 const TEST_BEFORE_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(60);
@@ -181,7 +181,7 @@ impl<DB: Database> PoolInner<DB> {
         tracing::trace!("waiting for any connection");
 
         let disconnected = match self.connections.acquire_any().await {
-            Ok(conn) => match finish_acquire(self, conn).await {
+            Ok(conn) => match self.finish_acquire(conn).await {
                 Ok(conn) => return Ok(conn),
                 Err(slot) => slot,
             },
@@ -199,14 +199,23 @@ impl<DB: Database> PoolInner<DB> {
             match race(&mut connect_task, self.connections.acquire_connected()).await {
                 Ok(Ok(conn)) => return Ok(conn),
                 Ok(Err(e)) => return Err(e),
-                Err(conn) => match finish_acquire(self, conn).await {
+                Err(conn) => match self.finish_acquire(conn).await {
                     Ok(conn) => return Ok(conn),
                     Err(_) => continue,
                 },
             }
         }
     }
 
+    #[inline(always)]
+    async fn finish_acquire(self: &Arc<Self>, conn: ConnectedSlot<ConnectionInner<DB>>) -> Result<PoolConnection<DB>, DisconnectedSlot<ConnectionInner<DB>>> {
+        let span = tracing::trace_span!(target: "sqlx::pool", "finish_acquire", connection_id=?conn.id);
+
+        finish_acquire(self, conn)
+            .instrument(span)
+            .await
+    }
+
     pub(crate) async fn try_min_connections(
         self: &Arc<Self>,
         deadline: Option<Instant>,