WIP refactor: new #[sqlx::test] architecture

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

sqlx-core/src/pool/connection.rs

@@ -127,11 +127,15 @@ impl<DB: Database> PoolConnection<DB> {
         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<Output = ()> + Send + 'static {
+    pub fn release(&mut self) -> impl Future<Output = ()> + 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<DB: Database> PoolConnection<DB> {
 
         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<DB: Database> PoolConnection<DB> {
             }
         }
     }
-
-    fn take_and_close(&mut self) -> impl Future<Output = ()> + 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<DB> {
@@ -198,14 +192,9 @@ impl<'c, DB: Database> crate::acquire::Acquire<'c> for &'c mut PoolConnection<DB
 /// Returns the connection to the [`Pool`][crate::pool::Pool] it was checked-out from.
 impl<DB: Database> Drop for PoolConnection<DB> {
     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());
         }
     }
 }

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<T> {
+    #[cfg(feature = "_rt-tokio")]
+    inner: tokio::sync::OnceCell<T>,
+
+    #[cfg(all(feature = "_rt-async-io", not(feature = "_rt-tokio")))]
+    inner: async_lock::OnceCell<T>,
+
+    #[cfg(not(any(feature = "_rt-async-std", feature = "_rt-tokio")))]
+    phantom: std::marker::PhantomData<T>,
+}
+
+impl<T> AsyncOnceCell<T> {
+    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<F, Fut, E>(&self, f: F) -> Result<&T, E>
+    where
+        F: FnOnce() -> Fut,
+        Fut: Future<Output = Result<T, E>>,
+    {
+        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<T> Default for AsyncOnceCell<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}

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 () {