refactor(pool): replace futures-intrusive semaphore with asyncband

sqlx-core/Cargo.toml

@@ -20,11 +20,11 @@ any = []
 json = ["serde", "serde_json"]
 
 # for conditional compilation
-_rt-async-global-executor = ["async-global-executor", "_rt-async-io", "_rt-async-task"]
+_rt-async-global-executor = ["async-global-executor", "asyncband", "_rt-async-io", "_rt-async-task"]
 _rt-async-io = ["async-io", "async-fs"] # see note at async-fs declaration
-_rt-async-std = ["async-std", "_rt-async-io"]
+_rt-async-std = ["async-std", "asyncband", "_rt-async-io"]
 _rt-async-task = ["async-task"]
-_rt-smol = ["smol", "_rt-async-io", "_rt-async-task"]
+_rt-smol = ["smol", "asyncband", "_rt-async-io", "_rt-async-task"]
 _rt-tokio = ["tokio", "tokio-stream"]
 
 _tls-native-tls = ["native-tls"]
@@ -74,6 +74,7 @@ uuid = { workspace = true, optional = true }
 async-fs = { version = "2.1", optional = true }
 async-io = { version = "2.4.1", optional = true }
 async-task = { version = "4.7.1", optional = true }
+asyncband = { version = "0.7.1", features = ["semaphore"], optional = true }
 
 base64.workspace = true
 bytes = "1.2.0"
@@ -83,7 +84,6 @@ crossbeam-queue = "0.3.2"
 either = "1.6.1"
 futures-core = { version = "0.3.32", default-features = false }
 futures-io = "0.3.32"
-futures-intrusive = "0.5.0"
 futures-util = { version = "0.3.32", default-features = false, features = ["alloc", "sink", "io"] }
 log = { version = "0.4.18", default-features = false }
 memchr = { version = "2.5.0", default-features = false }

sqlx-core/src/pool/inner.rs

@@ -44,7 +44,6 @@ impl<DB: Database> PoolInner<DB> {
 
         let semaphore_capacity = if let Some(parent) = &options.parent_pool {
             assert!(options.max_connections <= parent.options().max_connections);
-            assert_eq!(options.fair, parent.options().fair);
             // The child pool must steal permits from the parent
             0
         } else {
@@ -54,7 +53,7 @@ impl<DB: Database> PoolInner<DB> {
         let pool = Self {
             connect_options: RwLock::new(Arc::new(connect_options)),
             idle_conns: ArrayQueue::new(capacity),
-            semaphore: AsyncSemaphore::new(options.fair, semaphore_capacity),
+            semaphore: AsyncSemaphore::new(semaphore_capacity),
             size: AtomicU32::new(0),
             num_idle: AtomicUsize::new(0),
             is_closed: AtomicBool::new(false),

sqlx-core/src/pool/options.rs

@@ -82,7 +82,6 @@ pub struct PoolOptions<DB: Database> {
     pub(crate) min_connections: u32,
     pub(crate) max_lifetime: Option<Duration>,
     pub(crate) idle_timeout: Option<Duration>,
-    pub(crate) fair: bool,
 
     pub(crate) parent_pool: Option<Pool<DB>>,
 }
@@ -105,7 +104,6 @@ impl<DB: Database> Clone for PoolOptions<DB> {
             min_connections: self.min_connections,
             max_lifetime: self.max_lifetime,
             idle_timeout: self.idle_timeout,
-            fair: self.fair,
             parent_pool: self.parent_pool.clone(),
         }
     }
@@ -160,7 +158,6 @@ impl<DB: Database> PoolOptions<DB> {
             acquire_timeout: Duration::from_secs(30),
             idle_timeout: Some(Duration::from_secs(10 * 60)),
             max_lifetime: Some(Duration::from_secs(30 * 60)),
-            fair: true,
             parent_pool: None,
         }
     }
@@ -321,24 +318,6 @@ impl<DB: Database> PoolOptions<DB> {
         self.test_before_acquire
     }
 
-    /// If set to `true`, calls to `acquire()` are fair and connections  are issued
-    /// in first-come-first-serve order. If `false`, "drive-by" tasks may steal idle connections
-    /// ahead of tasks that have been waiting.
-    ///
-    /// According to `sqlx-bench/benches/pg_pool` this may slightly increase time
-    /// to `acquire()` at low pool contention but at very high contention it helps
-    /// avoid tasks at the head of the waiter queue getting repeatedly preempted by
-    /// these "drive-by" tasks and tasks further back in the queue timing out because
-    /// the queue isn't moving.
-    ///
-    /// Currently only exposed for benchmarking; `fair = true` seems to be the superior option
-    /// in most cases.
-    #[doc(hidden)]
-    pub fn __fair(mut self, fair: bool) -> Self {
-        self.fair = fair;
-        self
-    }
-
     /// Perform an asynchronous action after connecting to the database.
     ///
     /// If the operation returns with an error then the error is logged, the connection is closed
@@ -505,8 +484,7 @@ impl<DB: Database> PoolOptions<DB> {
     /// This is currently an internal-only API.
     ///
     /// ### Panics
-    /// If `self.max_connections` is greater than the setting the given pool was created with,
-    /// or `self.fair` differs from the setting the given pool was created with.
+    /// If `self.max_connections` is greater than the setting the given pool was created with.
     #[doc(hidden)]
     pub fn parent(mut self, pool: Pool<DB>) -> Self {
         self.parent_pool = Some(pool);

sqlx-core/src/sync.rs

@@ -3,19 +3,10 @@ use cfg_if::cfg_if;
 // For types with identical signatures that don't require runtime support,
 // we can just arbitrarily pick one to use based on what's enabled.
 //
-// We'll generally lean towards Tokio's types as those are more featureful
-// (including `tokio-console` support) and more widely deployed.
+// Prefer Tokio's types when enabled for cooperative scheduling and
+// `tracing`/`tokio-console` integration.
 
 pub struct AsyncSemaphore {
-    // We use the semaphore from futures-intrusive as the one from async-lock
-    // is missing the ability to add arbitrary permits, and is not guaranteed to be fair:
-    // * https://github.com/smol-rs/async-lock/issues/22
-    // * https://github.com/smol-rs/async-lock/issues/23
-    //
-    // We're on the look-out for a replacement, however, as futures-intrusive is not maintained
-    // and there are some soundness concerns (although it turns out any intrusive future is unsound
-    // in MIRI due to the necessitated mutable aliasing):
-    // https://github.com/launchbadge/sqlx/issues/1668
     #[cfg(all(
         any(
             feature = "_rt-async-global-executor",
@@ -24,22 +15,22 @@ pub struct AsyncSemaphore {
         ),
         not(feature = "_rt-tokio")
     ))]
-    inner: futures_intrusive::sync::Semaphore,
+    inner: asyncband::semaphore::Semaphore,
 
     #[cfg(feature = "_rt-tokio")]
     inner: tokio::sync::Semaphore,
 }
 
 impl AsyncSemaphore {
     #[track_caller]
-    pub fn new(fair: bool, permits: usize) -> Self {
+    pub fn new(permits: usize) -> Self {
         if cfg!(not(any(
             feature = "_rt-async-global-executor",
             feature = "_rt-async-std",
             feature = "_rt-smol",
             feature = "_rt-tokio"
         ))) {
-            crate::rt::missing_rt((fair, permits));
+            crate::rt::missing_rt(permits);
         }
 
         AsyncSemaphore {
@@ -51,27 +42,20 @@ impl AsyncSemaphore {
                 ),
                 not(feature = "_rt-tokio")
             ))]
-            inner: futures_intrusive::sync::Semaphore::new(fair, permits),
+            inner: asyncband::semaphore::Semaphore::new(permits),
             #[cfg(feature = "_rt-tokio")]
-            inner: {
-                debug_assert!(fair, "Tokio only has fair permits");
-                tokio::sync::Semaphore::new(permits)
-            },
+            inner: tokio::sync::Semaphore::new(permits),
         }
     }
 
     pub fn permits(&self) -> usize {
         cfg_if! {
-            if #[cfg(all(
-                any(
-                    feature = "_rt-async-global-executor",
-                    feature = "_rt-async-std",
-                    feature = "_rt-smol"
-                ),
-                not(feature = "_rt-tokio")
+            if #[cfg(any(
+                feature = "_rt-async-global-executor",
+                feature = "_rt-async-std",
+                feature = "_rt-smol",
+                feature = "_rt-tokio"
             ))] {
-                self.inner.permits()
-            } else if #[cfg(feature = "_rt-tokio")] {
                 self.inner.available_permits()
             } else {
                 crate::rt::missing_rt(())
@@ -152,15 +136,6 @@ impl AsyncSemaphore {
 }
 
 pub struct AsyncSemaphoreReleaser<'a> {
-    // We use the semaphore from futures-intrusive as the one from async-std
-    // is missing the ability to add arbitrary permits, and is not guaranteed to be fair:
-    // * https://github.com/smol-rs/async-lock/issues/22
-    // * https://github.com/smol-rs/async-lock/issues/23
-    //
-    // We're on the look-out for a replacement, however, as futures-intrusive is not maintained
-    // and there are some soundness concerns (although it turns out any intrusive future is unsound
-    // in MIRI due to the necessitated mutable aliasing):
-    // https://github.com/launchbadge/sqlx/issues/1668
     #[cfg(all(
         any(
             feature = "_rt-async-global-executor",
@@ -169,7 +144,7 @@ pub struct AsyncSemaphoreReleaser<'a> {
         ),
         not(feature = "_rt-tokio")
     ))]
-    inner: futures_intrusive::sync::SemaphoreReleaser<'a>,
+    inner: asyncband::semaphore::SemaphorePermit<'a>,
 
     #[cfg(feature = "_rt-tokio")]
     inner: tokio::sync::SemaphorePermit<'a>,
@@ -186,17 +161,12 @@ pub struct AsyncSemaphoreReleaser<'a> {
 impl AsyncSemaphoreReleaser<'_> {
     pub fn disarm(self) {
         cfg_if! {
-            if #[cfg(all(
-                any(
-                    feature = "_rt-async-global-executor",
-                    feature = "_rt-async-std",
-                    feature = "_rt-smol"
-                ),
-                not(feature = "_rt-tokio")
+            if #[cfg(any(
+                feature = "_rt-async-global-executor",
+                feature = "_rt-async-std",
+                feature = "_rt-smol",
+                feature = "_rt-tokio"
             ))] {
-                let mut this = self;
-                this.inner.disarm();
-            } else if #[cfg(feature = "_rt-tokio")] {
                 self.inner.forget();
             } else {
                 crate::rt::missing_rt(());

tests/any/pool.rs

@@ -1,10 +1,12 @@
 use sqlx::any::{AnyConnectOptions, AnyPoolOptions};
 use sqlx::Executor;
 use sqlx_core::sql_str::AssertSqlSafe;
+use std::future::Future;
 use std::sync::{
     atomic::{AtomicI32, AtomicUsize, Ordering},
     Arc, Mutex,
 };
+use std::task::{Context, Waker};
 use std::time::Duration;
 
 #[sqlx_macros::test]
@@ -205,6 +207,107 @@ async fn test_pool_callbacks() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[sqlx_macros::test]
+async fn pool_acquires_in_order_after_cancelling_a_waiter() -> anyhow::Result<()> {
+    sqlx::any::install_default_drivers();
+    let pool = AnyPoolOptions::new()
+        .max_connections(1)
+        .test_before_acquire(false)
+        .connect(&dotenvy::var("DATABASE_URL")?)
+        .await?;
+    let mut held = pool.acquire().await?;
+    let mut cancelled = Box::pin(pool.acquire());
+    let mut first = Box::pin(pool.acquire());
+    let mut second = Box::pin(pool.acquire());
+    let mut cx = Context::from_waker(Waker::noop());
+
+    assert!(cancelled.as_mut().poll(&mut cx).is_pending());
+    assert!(first.as_mut().poll(&mut cx).is_pending());
+    assert!(second.as_mut().poll(&mut cx).is_pending());
+    drop(cancelled);
+    held.return_to_pool().await;
+
+    assert!(pool.try_acquire().is_none());
+    assert!(second.as_mut().poll(&mut cx).is_pending());
+    let mut first = first.await?;
+    assert!(pool.try_acquire().is_none());
+    first.return_to_pool().await;
+    let mut second = second.await?;
+    second.execute("SELECT 1").await?;
+    second.return_to_pool().await;
+
+    assert_eq!(pool.size(), 1);
+    assert_eq!(pool.num_idle(), 1);
+    pool.close().await;
+    Ok(())
+}
+
+#[sqlx_macros::test]
+async fn pool_close_waits_for_all_connections_and_wakes_waiters() -> anyhow::Result<()> {
+    sqlx::any::install_default_drivers();
+    let pool = AnyPoolOptions::new()
+        .max_connections(2)
+        .test_before_acquire(false)
+        .connect(&dotenvy::var("DATABASE_URL")?)
+        .await?;
+    let mut first = pool.acquire().await?;
+    let mut second = pool.acquire().await?;
+    let mut waiter = Box::pin(pool.acquire());
+    let mut cx = Context::from_waker(Waker::noop());
+
+    assert!(waiter.as_mut().poll(&mut cx).is_pending());
+    let mut close = Box::pin(pool.close());
+    assert!(close.as_mut().poll(&mut cx).is_pending());
+    assert!(matches!(waiter.await, Err(sqlx::Error::PoolClosed)));
+    first.return_to_pool().await;
+    assert!(close.as_mut().poll(&mut cx).is_pending());
+
+    // Cancelling a partially satisfied close must return its reserved permits.
+    drop(close);
+    let mut close = Box::pin(pool.close());
+    assert!(close.as_mut().poll(&mut cx).is_pending());
+    second.return_to_pool().await;
+    close.await;
+
+    assert!(pool.is_closed());
+    assert_eq!(pool.size(), 0);
+    assert_eq!(pool.num_idle(), 0);
+    Ok(())
+}
+
+#[sqlx_macros::test]
+async fn child_pool_returns_permits_to_parent_on_drop() -> anyhow::Result<()> {
+    sqlx::any::install_default_drivers();
+    let url = dotenvy::var("DATABASE_URL")?;
+    let parent = AnyPoolOptions::new()
+        .max_connections(2)
+        .test_before_acquire(false)
+        .connect(&url)
+        .await?;
+    let mut held = parent.acquire().await?;
+    let child = AnyPoolOptions::new()
+        .max_connections(1)
+        .test_before_acquire(false)
+        .parent(parent.clone())
+        .connect(&url)
+        .await?;
+    let mut child_conn = child.acquire().await?;
+    let mut waiting_on_parent = Box::pin(parent.acquire());
+    let mut cx = Context::from_waker(Waker::noop());
+
+    assert!(waiting_on_parent.as_mut().poll(&mut cx).is_pending());
+    child_conn.return_to_pool().await;
+    drop(child_conn);
+    assert!(waiting_on_parent.as_mut().poll(&mut cx).is_pending());
+    drop(child);
+    let mut returned = waiting_on_parent.await?;
+    returned.execute("SELECT 1").await?;
+    returned.return_to_pool().await;
+    held.return_to_pool().await;
+    parent.close().await;
+    Ok(())
+}
+
 #[ignore]
 #[sqlx_macros::test]
 async fn test_connection_maintenance() -> anyhow::Result<()> {