From fe12fbf81dafedbd1e4b7649d0ffd3dc8a65e6d2 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 16:24:25 +0800 Subject: [PATCH 1/2] refactor(pool): use asyncband for all semaphore backends --- sqlx-core/Cargo.toml | 2 +- sqlx-core/src/pool/inner.rs | 3 +- sqlx-core/src/pool/options.rs | 24 +---- sqlx-core/src/sync.rs | 183 +++------------------------------- tests/any/pool.rs | 103 +++++++++++++++++++ 5 files changed, 121 insertions(+), 194 deletions(-) diff --git a/sqlx-core/Cargo.toml b/sqlx-core/Cargo.toml index 90ed446b4b..ca06a8f11f 100644 --- a/sqlx-core/Cargo.toml +++ b/sqlx-core/Cargo.toml @@ -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"] } 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 } diff --git a/sqlx-core/src/pool/inner.rs b/sqlx-core/src/pool/inner.rs index b698dc9df0..5eed762ba4 100644 --- a/sqlx-core/src/pool/inner.rs +++ b/sqlx-core/src/pool/inner.rs @@ -44,7 +44,6 @@ impl PoolInner { 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 PoolInner { 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), diff --git a/sqlx-core/src/pool/options.rs b/sqlx-core/src/pool/options.rs index 3d048f1795..6c632f156b 100644 --- a/sqlx-core/src/pool/options.rs +++ b/sqlx-core/src/pool/options.rs @@ -82,7 +82,6 @@ pub struct PoolOptions { pub(crate) min_connections: u32, pub(crate) max_lifetime: Option, pub(crate) idle_timeout: Option, - pub(crate) fair: bool, pub(crate) parent_pool: Option>, } @@ -105,7 +104,6 @@ impl Clone for PoolOptions { 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 PoolOptions { 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 PoolOptions { 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 PoolOptions { /// 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) -> Self { self.parent_pool = Some(pool); diff --git a/sqlx-core/src/sync.rs b/sqlx-core/src/sync.rs index ed082f752c..dde37dc6e2 100644 --- a/sqlx-core/src/sync.rs +++ b/sqlx-core/src/sync.rs @@ -1,206 +1,53 @@ -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. +use asyncband::semaphore::{Semaphore, SemaphorePermit}; 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", - feature = "_rt-async-std", - feature = "_rt-smol" - ), - not(feature = "_rt-tokio") - ))] - inner: futures_intrusive::sync::Semaphore, - - #[cfg(feature = "_rt-tokio")] - inner: tokio::sync::Semaphore, + inner: 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 { - #[cfg(all( - any( - feature = "_rt-async-global-executor", - feature = "_rt-async-std", - feature = "_rt-smol" - ), - not(feature = "_rt-tokio") - ))] - inner: futures_intrusive::sync::Semaphore::new(fair, permits), - #[cfg(feature = "_rt-tokio")] - inner: { - debug_assert!(fair, "Tokio only has fair permits"); - tokio::sync::Semaphore::new(permits) - }, + Self { + inner: 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") - ))] { - self.inner.permits() - } else if #[cfg(feature = "_rt-tokio")] { - self.inner.available_permits() - } else { - crate::rt::missing_rt(()) - } - } + self.inner.available_permits() } pub async fn acquire(&self, permits: u32) -> AsyncSemaphoreReleaser<'_> { - cfg_if! { - if #[cfg(all( - any( - feature = "_rt-async-global-executor", - feature = "_rt-async-std", - feature = "_rt-smol" - ), - not(feature = "_rt-tokio") - ))] { - AsyncSemaphoreReleaser { - inner: self.inner.acquire(permits as usize).await, - } - } else if #[cfg(feature = "_rt-tokio")] { - AsyncSemaphoreReleaser { - inner: self - .inner - // Weird quirk: `tokio::sync::Semaphore` mostly uses `usize` for permit counts, - // but `u32` for this and `try_acquire_many()`. - .acquire_many(permits) - .await - .expect("BUG: we do not expose the `.close()` method"), - } - } else { - crate::rt::missing_rt(permits) - } + AsyncSemaphoreReleaser { + inner: self.inner.acquire(permits as usize).await, } } pub fn try_acquire(&self, permits: u32) -> Option> { - cfg_if! { - if #[cfg(all( - any( - feature = "_rt-async-global-executor", - feature = "_rt-async-std", - feature = "_rt-smol" - ), - not(feature = "_rt-tokio") - ))] { - Some(AsyncSemaphoreReleaser { - inner: self.inner.try_acquire(permits as usize)?, - }) - } else if #[cfg(feature = "_rt-tokio")] { - Some(AsyncSemaphoreReleaser { - inner: self.inner.try_acquire_many(permits).ok()?, - }) - } else { - crate::rt::missing_rt(permits) - } - } + Some(AsyncSemaphoreReleaser { + inner: self.inner.try_acquire(permits as usize)?, + }) } pub fn release(&self, permits: usize) { - cfg_if! { - if #[cfg(all( - any( - feature = "_rt-async-global-executor", - feature = "_rt-async-std", - feature = "_rt-smol" - ), - not(feature = "_rt-tokio") - ))] { - self.inner.release(permits); - } else if #[cfg(feature = "_rt-tokio")] { - self.inner.add_permits(permits); - } else { - crate::rt::missing_rt(permits); - } - } + self.inner.release(permits); } } 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", - feature = "_rt-async-std", - feature = "_rt-smol" - ), - not(feature = "_rt-tokio") - ))] - inner: futures_intrusive::sync::SemaphoreReleaser<'a>, - - #[cfg(feature = "_rt-tokio")] - inner: tokio::sync::SemaphorePermit<'a>, - - #[cfg(not(any( - feature = "_rt-async-global-executor", - feature = "_rt-async-std", - feature = "_rt-smol", - feature = "_rt-tokio" - )))] - _phantom: std::marker::PhantomData<&'a ()>, + inner: SemaphorePermit<'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") - ))] { - let mut this = self; - this.inner.disarm(); - } else if #[cfg(feature = "_rt-tokio")] { - self.inner.forget(); - } else { - crate::rt::missing_rt(()); - } - } + self.inner.forget(); } } diff --git a/tests/any/pool.rs b/tests/any/pool.rs index a4849940b8..b70be27551 100644 --- a/tests/any/pool.rs +++ b/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<()> { From e7e712129a6cd16b5a27a773b5e9c7564d24a88c Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 21:52:23 +0800 Subject: [PATCH 2/2] refactor(pool): retain Tokio semaphore backend --- sqlx-core/Cargo.toml | 8 +-- sqlx-core/src/sync.rs | 149 ++++++++++++++++++++++++++++++++++++++---- 2 files changed, 140 insertions(+), 17 deletions(-) diff --git a/sqlx-core/Cargo.toml b/sqlx-core/Cargo.toml index ca06a8f11f..c9078b8e3f 100644 --- a/sqlx-core/Cargo.toml +++ b/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,7 +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"] } +asyncband = { version = "0.7.1", features = ["semaphore"], optional = true } base64.workspace = true bytes = "1.2.0" diff --git a/sqlx-core/src/sync.rs b/sqlx-core/src/sync.rs index dde37dc6e2..2946b73343 100644 --- a/sqlx-core/src/sync.rs +++ b/sqlx-core/src/sync.rs @@ -1,7 +1,24 @@ -use asyncband::semaphore::{Semaphore, SemaphorePermit}; +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. +// +// Prefer Tokio's types when enabled for cooperative scheduling and +// `tracing`/`tokio-console` integration. pub struct AsyncSemaphore { - inner: Semaphore, + #[cfg(all( + any( + feature = "_rt-async-global-executor", + feature = "_rt-async-std", + feature = "_rt-smol" + ), + not(feature = "_rt-tokio") + ))] + inner: asyncband::semaphore::Semaphore, + + #[cfg(feature = "_rt-tokio")] + inner: tokio::sync::Semaphore, } impl AsyncSemaphore { @@ -16,38 +33,144 @@ impl AsyncSemaphore { crate::rt::missing_rt(permits); } - Self { - inner: Semaphore::new(permits), + AsyncSemaphore { + #[cfg(all( + any( + feature = "_rt-async-global-executor", + feature = "_rt-async-std", + feature = "_rt-smol" + ), + not(feature = "_rt-tokio") + ))] + inner: asyncband::semaphore::Semaphore::new(permits), + #[cfg(feature = "_rt-tokio")] + inner: tokio::sync::Semaphore::new(permits), } } pub fn permits(&self) -> usize { - self.inner.available_permits() + cfg_if! { + if #[cfg(any( + feature = "_rt-async-global-executor", + feature = "_rt-async-std", + feature = "_rt-smol", + feature = "_rt-tokio" + ))] { + self.inner.available_permits() + } else { + crate::rt::missing_rt(()) + } + } } pub async fn acquire(&self, permits: u32) -> AsyncSemaphoreReleaser<'_> { - AsyncSemaphoreReleaser { - inner: self.inner.acquire(permits as usize).await, + cfg_if! { + if #[cfg(all( + any( + feature = "_rt-async-global-executor", + feature = "_rt-async-std", + feature = "_rt-smol" + ), + not(feature = "_rt-tokio") + ))] { + AsyncSemaphoreReleaser { + inner: self.inner.acquire(permits as usize).await, + } + } else if #[cfg(feature = "_rt-tokio")] { + AsyncSemaphoreReleaser { + inner: self + .inner + // Weird quirk: `tokio::sync::Semaphore` mostly uses `usize` for permit counts, + // but `u32` for this and `try_acquire_many()`. + .acquire_many(permits) + .await + .expect("BUG: we do not expose the `.close()` method"), + } + } else { + crate::rt::missing_rt(permits) + } } } pub fn try_acquire(&self, permits: u32) -> Option> { - Some(AsyncSemaphoreReleaser { - inner: self.inner.try_acquire(permits as usize)?, - }) + cfg_if! { + if #[cfg(all( + any( + feature = "_rt-async-global-executor", + feature = "_rt-async-std", + feature = "_rt-smol" + ), + not(feature = "_rt-tokio") + ))] { + Some(AsyncSemaphoreReleaser { + inner: self.inner.try_acquire(permits as usize)?, + }) + } else if #[cfg(feature = "_rt-tokio")] { + Some(AsyncSemaphoreReleaser { + inner: self.inner.try_acquire_many(permits).ok()?, + }) + } else { + crate::rt::missing_rt(permits) + } + } } pub fn release(&self, permits: usize) { - self.inner.release(permits); + cfg_if! { + if #[cfg(all( + any( + feature = "_rt-async-global-executor", + feature = "_rt-async-std", + feature = "_rt-smol" + ), + not(feature = "_rt-tokio") + ))] { + self.inner.release(permits); + } else if #[cfg(feature = "_rt-tokio")] { + self.inner.add_permits(permits); + } else { + crate::rt::missing_rt(permits); + } + } } } pub struct AsyncSemaphoreReleaser<'a> { - inner: SemaphorePermit<'a>, + #[cfg(all( + any( + feature = "_rt-async-global-executor", + feature = "_rt-async-std", + feature = "_rt-smol" + ), + not(feature = "_rt-tokio") + ))] + inner: asyncband::semaphore::SemaphorePermit<'a>, + + #[cfg(feature = "_rt-tokio")] + inner: tokio::sync::SemaphorePermit<'a>, + + #[cfg(not(any( + feature = "_rt-async-global-executor", + feature = "_rt-async-std", + feature = "_rt-smol", + feature = "_rt-tokio" + )))] + _phantom: std::marker::PhantomData<&'a ()>, } impl AsyncSemaphoreReleaser<'_> { pub fn disarm(self) { - self.inner.forget(); + cfg_if! { + if #[cfg(any( + feature = "_rt-async-global-executor", + feature = "_rt-async-std", + feature = "_rt-smol", + feature = "_rt-tokio" + ))] { + self.inner.forget(); + } else { + crate::rt::missing_rt(()); + } + } } }