diff --git a/sqlx-core/Cargo.toml b/sqlx-core/Cargo.toml index 90ed446b4b..ca34aa4e29 100644 --- a/sqlx-core/Cargo.toml +++ b/sqlx-core/Cargo.toml @@ -21,11 +21,11 @@ json = ["serde", "serde_json"] # for conditional compilation _rt-async-global-executor = ["async-global-executor", "_rt-async-io", "_rt-async-task"] -_rt-async-io = ["async-io", "async-fs"] # see note at async-fs declaration +_rt-async-io = ["async-io", "async-fs", "async-lock"] # see note at async-fs declaration _rt-async-std = ["async-std", "_rt-async-io"] _rt-async-task = ["async-task"] _rt-smol = ["smol", "_rt-async-io", "_rt-async-task"] -_rt-tokio = ["tokio", "tokio-stream"] +_rt-tokio = ["tokio", "tokio-stream", "tokio/rt"] # `rt` feature is almost always going to be enabled anyway _tls-native-tls = ["native-tls"] _tls-rustls-aws-lc-rs = ["_tls-rustls", "rustls/aws-lc-rs", "webpki-roots"] @@ -72,6 +72,7 @@ uuid = { workspace = true, optional = true } # work around bug in async-fs 2.0.0, which references futures-lite dependency wrongly, see https://github.com/launchbadge/sqlx/pull/3791#issuecomment-3043363281 async-fs = { version = "2.1", optional = true } +async-lock = { version = "3.4.2", optional = true } async-io = { version = "2.4.1", optional = true } async-task = { version = "4.7.1", optional = true } @@ -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 } @@ -103,10 +103,11 @@ indexmap = "2.0" event-listener = "5.2.0" hashbrown = "0.16.0" -thiserror.workspace = true +futures-intrusive = "0.5.0" -[dev-dependencies] -tokio = { version = "1.25.0", features = ["rt"] } +pin-project-lite = "0.2.17" + +thiserror.workspace = true [dev-dependencies.sqlx] # FIXME: https://github.com/rust-lang/cargo/issues/15622 diff --git a/sqlx-core/src/pool/connection.rs b/sqlx-core/src/pool/connection.rs index 7912b12aa1..c3d48819ab 100644 --- a/sqlx-core/src/pool/connection.rs +++ b/sqlx-core/src/pool/connection.rs @@ -127,11 +127,15 @@ impl PoolConnection { self.live.take().expect(EXPECT_MSG) } + /// Release the connection. + /// + /// If the connection is to be closed, close it and run pool maintenance. + /// /// Test the connection to make sure it is still live before returning it to the pool. /// /// This effectively runs the drop handler eagerly instead of spawning a task to do it. #[doc(hidden)] - pub fn return_to_pool(&mut self) -> impl Future + Send + 'static { + pub fn release(&mut self) -> impl Future + Send + 'static { // float the connection in the pool before we move into the task // in case the returned `Future` isn't executed, like if it's spawned into a dying runtime // https://github.com/launchbadge/sqlx/issues/1396 @@ -141,9 +145,20 @@ impl PoolConnection { let pool = self.pool.clone(); + let close_on_drop = self.close_on_drop; + async move { let returned_to_pool = if let Some(floating) = floating { - floating.return_to_pool().await + if close_on_drop { + // Don't hold the connection forever if it hangs while trying to close + crate::rt::timeout(CLOSE_ON_DROP_TIMEOUT, floating.close()) + .await + .ok(); + + false + } else { + floating.return_to_pool().await + } } else { false }; @@ -153,27 +168,6 @@ impl PoolConnection { } } } - - fn take_and_close(&mut self) -> impl Future + Send + 'static { - // float the connection in the pool before we move into the task - // in case the returned `Future` isn't executed, like if it's spawned into a dying runtime - // https://github.com/launchbadge/sqlx/issues/1396 - // Type hints seem to be broken by `Option` combinators in IntelliJ Rust right now (6/22). - let floating = self.live.take().map(|live| live.float(self.pool.clone())); - - let pool = self.pool.clone(); - - async move { - if let Some(floating) = floating { - // Don't hold the connection forever if it hangs while trying to close - crate::rt::timeout(CLOSE_ON_DROP_TIMEOUT, floating.close()) - .await - .ok(); - } - - pool.min_connections_maintenance(None).await; - } - } } impl<'c, DB: Database> crate::acquire::Acquire<'c> for &'c mut PoolConnection { @@ -198,14 +192,9 @@ impl<'c, DB: Database> crate::acquire::Acquire<'c> for &'c mut PoolConnection Drop for PoolConnection { fn drop(&mut self) { - if self.close_on_drop { - crate::rt::spawn(self.take_and_close()); - return; - } - // We still need to spawn a task to maintain `min_connections`. if self.live.is_some() || self.pool.options.min_connections > 0 { - crate::rt::spawn(self.return_to_pool()); + crate::rt::spawn(self.release()); } } } diff --git a/sqlx-core/src/sync.rs b/sqlx-core/src/sync.rs index ed082f752c..7fdd360f2d 100644 --- a/sqlx-core/src/sync.rs +++ b/sqlx-core/src/sync.rs @@ -1,4 +1,5 @@ use cfg_if::cfg_if; +use std::future::Future; // For types with identical signatures that don't require runtime support, // we can just arbitrarily pick one to use based on what's enabled. @@ -204,3 +205,70 @@ impl AsyncSemaphoreReleaser<'_> { } } } + +pub struct AsyncOnceCell { + #[cfg(feature = "_rt-tokio")] + inner: tokio::sync::OnceCell, + + #[cfg(all(feature = "_rt-async-io", not(feature = "_rt-tokio")))] + inner: async_lock::OnceCell, + + #[cfg(not(any(feature = "_rt-async-std", feature = "_rt-tokio")))] + phantom: std::marker::PhantomData, +} + +impl AsyncOnceCell { + pub fn new() -> Self { + cfg_if! { + if #[cfg(feature = "_rt-tokio")] { + Self { inner: tokio::sync::OnceCell::new() } + } else if #[cfg(feature = "_rt-async-io")] { + Self { inner: async_lock::OnceCell::new() } + } else { + Self { phantom: std::marker::PhantomData } + } + } + } + + pub const fn const_new() -> Self { + cfg_if! { + if #[cfg(feature = "_rt-tokio")] { + Self { inner: tokio::sync::OnceCell::const_new() } + } else if #[cfg(feature = "_rt-async-io")] { + Self { inner: async_lock::OnceCell::new() } + } else { + Self { phantom: std::marker::PhantomData } + } + } + } + + pub async fn get_or_try_init(&self, f: F) -> Result<&T, E> + where + F: FnOnce() -> Fut, + Fut: Future>, + { + cfg_if! { + if #[cfg(any(feature = "_rt-tokio", feature = "_rt-async-io"))] { + self.inner.get_or_try_init(f).await + } else { + crate::rt::missing_rt(f) + } + } + } + + pub fn get(&self) -> Option<&T> { + cfg_if! { + if #[cfg(any(feature = "_rt-tokio", feature = "_rt-async-io"))] { + self.inner.get() + } else { + crate::rt::missing_rt(()) + } + } + } +} + +impl Default for AsyncOnceCell { + fn default() -> Self { + Self::new() + } +} diff --git a/sqlx-core/src/testing/mod.rs b/sqlx-core/src/testing/mod.rs index 379fac690e..62d75edefa 100644 --- a/sqlx-core/src/testing/mod.rs +++ b/sqlx-core/src/testing/mod.rs @@ -13,6 +13,9 @@ use crate::migrate::{Migrate, Migrator}; use crate::pool::{Pool, PoolConnection, PoolOptions}; mod fixtures; +mod pool; + +pub use pool::{TestMasterConnection, TestMasterPool}; pub trait TestSupport: Database { /// Get parameters to construct a `Pool` suitable for testing. @@ -66,6 +69,7 @@ pub struct TestArgs { pub test_path: &'static str, pub migrator: Option<&'static Migrator>, pub fixtures: &'static [TestFixture], + pub max_connections: u32, } pub trait TestFn { @@ -158,6 +162,7 @@ impl TestArgs { test_path, migrator: None, fixtures: &[], + max_connections: 5, } } @@ -168,6 +173,10 @@ impl TestArgs { pub fn fixtures(&mut self, fixtures: &'static [TestFixture]) { self.fixtures = fixtures; } + + pub fn max_connections(&mut self, max_connections: u32) { + self.max_connections = max_connections; + } } impl TestTermination for () { diff --git a/sqlx-core/src/testing/pool.rs b/sqlx-core/src/testing/pool.rs new file mode 100644 index 0000000000..ab5b36f9db --- /dev/null +++ b/sqlx-core/src/testing/pool.rs @@ -0,0 +1,184 @@ +use crate::connection::Connection; +use crate::database::Database; +use crate::pool::{Pool, PoolConnection, PoolOptions}; +use crate::sync::AsyncOnceCell; +use cfg_if::cfg_if; +use pin_project_lite::pin_project; +use std::future::Future; +use std::ops::{Deref, DerefMut}; +use std::pin::Pin; +use std::task::{Context, Poll}; + +pub struct TestMasterPool { + inner: AsyncOnceCell>, +} + +pub struct TestMasterConnection { + conn: PoolConnection, + + #[cfg(feature = "_rt-tokio")] + handle: tokio::runtime::Handle, +} + +struct Inner { + pool: Pool, + + #[cfg(feature = "_rt-tokio")] + handle: tokio::runtime::Handle, +} + +macro_rules! poll_with_handle( + ($handle:expr, $fut:expr) => { + PollWithHandle { + fut: $fut, + #[cfg(feature = "_rt-tokio")] + handle: &$handle, + #[cfg(not(feature = "_rt-tokio"))] + _marker: std::marker::PhantomData, + } + } +); + +impl TestMasterPool { + pub const fn new() -> Self { + TestMasterPool { + inner: AsyncOnceCell::const_new(), + } + } + + pub async fn connect( + &self, + opts: &::Options, + ) -> crate::Result> { + self.inner + .get_or_try_init::<_, _, crate::Error>(|| { + let opts = opts.clone(); + + async move { + #[cfg(feature = "_rt-tokio")] + let handle = spawn_test_runtime(); + + // Ensure this pool is linked to our master runtime so it can survive an individual + // test runtime shutting down. + let pool = poll_with_handle!( + handle, + PoolOptions::new() + // Tests don't need a master connection for very long + .max_connections(1) + .test_before_acquire(false) + .connect_with(opts) + ) + .await?; + + Ok(Inner { + pool, + #[cfg(feature = "_rt-tokio")] + handle, + }) + } + }) + .await? + .acquire() + .await + } + + /// # Panics + /// If [`Self::connect()`] has not already completed successfully. + pub async fn acquire(&self) -> crate::Result> { + self.inner + .get() + .expect("`TestMasterPool::connect()` has not been called") + .acquire() + .await + } +} + +impl Deref for TestMasterConnection { + type Target = PoolConnection; + + fn deref(&self) -> &Self::Target { + &self.conn + } +} + +impl DerefMut for TestMasterConnection { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.conn + } +} + +impl Drop for TestMasterConnection { + fn drop(&mut self) { + cfg_if!( + if #[cfg(feature = "_rt-tokio")] { + self.handle.spawn(self.conn.release()); + } else { + crate::rt::spawn(self.conn.release()); + } + ); + } +} + +impl Inner { + async fn acquire(&self) -> crate::Result> { + Ok(TestMasterConnection { + // Ostensibly we only need to enter the runtime if the connection isn't already established + conn: poll_with_handle!(self.handle, self.pool.acquire()).await?, + #[cfg(feature = "_rt-tokio")] + handle: self.handle.clone(), + }) + } +} + +// It's likely not advisable to hold an `EnterGuard` across an `.await` point, +// so we need to define an adapter that only enters the alternate runtime when it's polled. +#[cfg(feature = "_rt-tokio")] +pin_project! { + struct PollWithHandle<'a, F> { + #[pin] + fut: F, + handle: &'a tokio::runtime::Handle, + } +} + +#[cfg(not(feature = "_rt-tokio"))] +pin_project! { + struct PollWithHandle<'a, F> { + #[pin] + fut: F, + _marker: std::marker::PhantomData<&'a ()>, + } +} + +impl Future for PollWithHandle<'_, F> { + type Output = F::Output; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.project(); + + #[cfg(feature = "_rt-tokio")] + let _guard = this.handle.enter(); + + this.fut.poll(cx) + } +} + +#[cfg(feature = "_rt-tokio")] +fn spawn_test_runtime() -> tokio::runtime::Handle { + // Instead of forcing the `rt-multi-thread` feature on, + // we just run a current-thread runtime in a background thread that we ourselves spawn. + let rt = tokio::runtime::Builder::new_current_thread() + .name("sqlx-test-master-pool") + .enable_all() + .build() + .expect("failed to spawn master runtime"); + + let handle = rt.handle().clone(); + + std::thread::Builder::new() + .name("sqlx-test-master-pool".into()) + .spawn(move || rt.block_on(std::future::pending::<()>())) + .expect("failed to spawn thread for master runtime"); + + handle +} diff --git a/sqlx-postgres/src/listener.rs b/sqlx-postgres/src/listener.rs index f9b9b98b1f..a2c08a069c 100644 --- a/sqlx-postgres/src/listener.rs +++ b/sqlx-postgres/src/listener.rs @@ -367,7 +367,7 @@ impl Drop for PgListener { // inline the drop handler from `PoolConnection` so it doesn't try to spawn another task // otherwise, it may trigger a panic if this task is dropped because the runtime is going away: // https://github.com/launchbadge/sqlx/issues/1389 - conn.return_to_pool().await; + conn.release().await; }; // Unregister any listeners before returning the connection to the pool. diff --git a/sqlx-postgres/src/testing/mod.rs b/sqlx-postgres/src/testing/mod.rs index 3e1cf0ddf7..05a9b58b78 100644 --- a/sqlx-postgres/src/testing/mod.rs +++ b/sqlx-postgres/src/testing/mod.rs @@ -17,9 +17,7 @@ use crate::{PgConnectOptions, PgConnection, Postgres}; pub(crate) use sqlx_core::testing::*; -// Using a blocking `OnceLock` here because the critical sections are short. -static MASTER_POOL: OnceLock> = OnceLock::new(); -// Automatically delete any databases created before the start of the test binary. +static MASTER_POOL: TestMasterPool = TestMasterPool::new(); impl TestSupport for Postgres { fn test_context( @@ -29,11 +27,7 @@ impl TestSupport for Postgres { } async fn cleanup_test(db_name: &str) -> Result<(), Error> { - let mut conn = MASTER_POOL - .get() - .expect("cleanup_test() invoked outside `#[sqlx::test]`") - .acquire() - .await?; + let mut conn = MASTER_POOL.acquire().await?; do_cleanup(&mut conn, db_name).await } @@ -94,36 +88,7 @@ async fn test_context(args: &TestArgs) -> Result, Error> { let master_opts = PgConnectOptions::from_str(&url).expect("failed to parse DATABASE_URL"); - let pool = PoolOptions::new() - // Postgres' normal connection limit is 100 plus 3 superuser connections - // We don't want to use the whole cap and there may be fuzziness here due to - // concurrently running tests anyway. - .max_connections(20) - // Immediately close master connections. Tokio's I/O streams don't like hopping runtimes. - .after_release(|_conn, _| Box::pin(async move { Ok(false) })) - .connect_lazy_with(master_opts); - - let master_pool = match once_lock_try_insert_polyfill(&MASTER_POOL, pool) { - Ok(inserted) => inserted, - Err((existing, pool)) => { - // Sanity checks. - assert_eq!( - existing.connect_options().host, - pool.connect_options().host, - "DATABASE_URL changed at runtime, host differs" - ); - - assert_eq!( - existing.connect_options().database, - pool.connect_options().database, - "DATABASE_URL changed at runtime, database differs" - ); - - existing - } - }; - - let mut conn = master_pool.acquire().await?; + let mut conn = MASTER_POOL.connect(&master_opts).await?; // language=PostgreSQL conn.execute( @@ -146,7 +111,56 @@ async fn test_context(args: &TestArgs) -> Result, Error> { create index if not exists databases_created_at on _sqlx_test.databases(created_at); - create sequence if not exists _sqlx_test.database_ids; + create table if not exists _sqlx_test.tests ( + test_id int8 primary key generated always as identity, + -- Automatically cleans up leaked test runs as well + db_name text not null references _sqlx_test.databases(db_name) on delete cascade, + required_connections int4 not null + check (required_connections > 0 and required_connections <= max_connections), + -- Each test's `SQLX_TEST_MAX_CONNECTIONS`, ideally all the same + max_connections int4 not null check (max_connections > 0), + started_at timestamptz not null default now() + ); + + create or replace function _sqlx_test.tests_check_max_connections() + returns trigger as + $$ + declare + used_connections int4; + max_required_connections int4; + max_connections int4; + begin + select + sum(required_connections), + max(required_connections), + -- Abide by the highest `SQLX_TEST_MAX_CONNECTIONS` + max(max_connections) + into + used_connections, + max_required_connections, + max_connections + from _sqlx_test.tests; + + if max_required_connections > max_connections then + raise + 'max(required_connections) exceeds min(max_connections) of any process' + using constraint = 'required_connections_exceeds_max'; + elsif max_connections > max_connections then + raise 'not enough spare connections available; used: %i, total: %i', + used_connections, max_connections + using constraint = 'insufficient_connections_available'; + end if; + end; + $$ + language plpgsql; + + -- `create or replace constraint trigger` not supported + drop trigger if exists check_max_connections on _sqlx_test.tests; + + create constraint trigger check_max_connections + after insert on _sqlx_test.tests + for each row + execute function _sqlx_test.tests_check_max_connections(); "#, ) .await?; @@ -161,7 +175,7 @@ async fn test_context(args: &TestArgs) -> Result, Error> { ) .bind(&db_name) .bind(args.test_path) - .execute(&mut *conn) + .execute(&mut **conn) .await?; let create_command = format!("create database {db_name:?}"); @@ -170,18 +184,10 @@ async fn test_context(args: &TestArgs) -> Result, Error> { Ok(TestContext { pool_opts: PoolOptions::new() - // Don't allow a single test to take all the connections. - // Most tests shouldn't require more than 5 connections concurrently, - // or else they're likely doing too much in one test. - .max_connections(5) + .max_connections(args.max_connections) // Close connections ASAP if left in the idle queue. - .idle_timeout(Some(Duration::from_secs(1))) - .parent(master_pool.clone()), - connect_opts: master_pool - .connect_options() - .deref() - .clone() - .database(&db_name), + .idle_timeout(Some(Duration::from_secs(1))), + connect_opts: master_opts.database(&db_name), db_name, }) } diff --git a/tests/any/pool.rs b/tests/any/pool.rs index a4849940b8..22d6a939b0 100644 --- a/tests/any/pool.rs +++ b/tests/any/pool.rs @@ -234,7 +234,7 @@ async fn test_connection_maintenance() -> anyhow::Result<()> { assert_eq!(pool.size(), 5); assert_eq!(pool.num_idle(), 0); for mut conn in conns { - conn.return_to_pool().await; + conn.release().await; } assert_eq!(pool.size(), 5);