diff --git a/sqlx-core/Cargo.toml b/sqlx-core/Cargo.toml index 90ed446b4b..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,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 } 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..2946b73343 100644 --- a/sqlx-core/src/sync.rs +++ b/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,7 +15,7 @@ pub struct AsyncSemaphore { ), not(feature = "_rt-tokio") ))] - inner: futures_intrusive::sync::Semaphore, + inner: asyncband::semaphore::Semaphore, #[cfg(feature = "_rt-tokio")] inner: tokio::sync::Semaphore, @@ -32,14 +23,14 @@ pub struct AsyncSemaphore { 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(()); 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<()> {