From c3735abf034689bb1fab9d7bf9f040e492de28cf Mon Sep 17 00:00:00 2001 From: KimSoungRyoul <24240623+KimSoungRyoul@users.noreply.github.com> Date: Mon, 14 Sep 2026 14:51:44 +0900 Subject: [PATCH] fix(pool): bound the entire connection return operation --- Cargo.toml | 5 + sqlx-core/src/pool/connection.rs | 12 +- sqlx-core/src/pool/options.rs | 4 + tests/postgres/pool.rs | 219 +++++++++++++++++++++++++++++++ 4 files changed, 239 insertions(+), 1 deletion(-) create mode 100644 tests/postgres/pool.rs diff --git a/Cargo.toml b/Cargo.toml index 3ad344c3f9..669c20f2df 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -429,6 +429,11 @@ name = "postgres" path = "tests/postgres/postgres.rs" required-features = ["postgres"] +[[test]] +name = "postgres-pool" +path = "tests/postgres/pool.rs" +required-features = ["postgres", "runtime-tokio"] + [[test]] name = "postgres-types" path = "tests/postgres/types.rs" diff --git a/sqlx-core/src/pool/connection.rs b/sqlx-core/src/pool/connection.rs index 7912b12aa1..2cb2dc1d3a 100644 --- a/sqlx-core/src/pool/connection.rs +++ b/sqlx-core/src/pool/connection.rs @@ -14,6 +14,7 @@ use super::inner::{is_beyond_max_lifetime, DecrementSizeGuard, PoolInner}; use crate::pool::options::PoolConnectionMetadata; const CLOSE_ON_DROP_TIMEOUT: Duration = Duration::from_secs(5); +const RETURN_TO_POOL_TIMEOUT: Duration = Duration::from_secs(5); /// A connection managed by a [`Pool`][crate::pool::Pool]. /// @@ -143,7 +144,16 @@ impl PoolConnection { async move { let returned_to_pool = if let Some(floating) = floating { - floating.return_to_pool().await + // Bound the whole return, including callbacks and connection shutdown. + // On timeout, dropping the future drops the connection and its size guard, + // releasing the permit without awaiting any further connection I/O. + match crate::rt::timeout(RETURN_TO_POOL_TIMEOUT, floating.return_to_pool()).await { + Ok(returned) => returned, + Err(_) => { + tracing::warn!("timed out while returning a connection to the pool"); + false + } + } } else { false }; diff --git a/sqlx-core/src/pool/options.rs b/sqlx-core/src/pool/options.rs index 3d048f1795..20004ae360 100644 --- a/sqlx-core/src/pool/options.rs +++ b/sqlx-core/src/pool/options.rs @@ -453,6 +453,10 @@ impl PoolOptions { /// and the connection is closed, allowing a task waiting on [`Pool::acquire`] to /// open a new one in its place. /// + /// Returning a connection to the pool, including this callback and the subsequent + /// health check or connection shutdown, has a five-second timeout. If it expires, + /// the operation is cancelled and the connection is discarded. + /// /// # Example (Postgres): Close Memory-Hungry Connections /// Instead of relying on [`max_lifetime`][Self::max_lifetime] to close connections, /// we can monitor their memory usage directly and close any that have allocated too much. diff --git a/tests/postgres/pool.rs b/tests/postgres/pool.rs new file mode 100644 index 0000000000..f47283e03c --- /dev/null +++ b/tests/postgres/pool.rs @@ -0,0 +1,219 @@ +//! Pool regression tests using an in-process server; no DATABASE_URL is needed. + +use std::future; +use std::io; +use std::net::SocketAddr; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use sqlx::pool::PoolConnection; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions, PgSslMode}; +use sqlx::{PgPool, Postgres}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::mpsc; +use tokio::task::{JoinHandle, JoinSet}; +use tokio::time::{sleep, timeout}; + +// Allow ample time for the pool's five-second return timeout on busy CI workers. +const TEST_TIMEOUT: Duration = Duration::from_secs(20); + +struct TestServer { + addr: SocketAddr, + accepted: Arc, + disconnected: mpsc::UnboundedReceiver<()>, + task: JoinHandle<()>, +} + +impl TestServer { + async fn start(respond_to_ping: bool) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let accepted = Arc::new(AtomicUsize::new(0)); + let (disconnect_tx, disconnected) = mpsc::unbounded_channel(); + + let task = tokio::spawn({ + let accepted = accepted.clone(); + async move { + let mut connections = JoinSet::new(); + loop { + tokio::select! { + result = listener.accept() => { + let (socket, _) = result.unwrap(); + accepted.fetch_add(1, Ordering::SeqCst); + let disconnect_tx = disconnect_tx.clone(); + connections.spawn(async move { + serve(socket, respond_to_ping).await.unwrap(); + let _ = disconnect_tx.send(()); + }); + } + result = connections.join_next(), if !connections.is_empty() => { + result.unwrap().unwrap(); + } + } + } + } + }); + + Self { + addr, + accepted, + disconnected, + task, + } + } + + fn connect_options(&self) -> PgConnectOptions { + PgConnectOptions::new() + .host("127.0.0.1") + .port(self.addr.port()) + .username("test") + .database("test") + .ssl_mode(PgSslMode::Disable) + } + + async fn wait_for_disconnect(&mut self) { + timeout(TEST_TIMEOUT, self.disconnected.recv()) + .await + .expect("timed-out return must close the socket") + .expect("server stopped before the socket closed"); + } +} + +impl Drop for TestServer { + fn drop(&mut self) { + // Dropping the server's JoinSet also aborts any remaining connections. + self.task.abort(); + } +} + +async fn serve(mut socket: TcpStream, respond_to_ping: bool) -> io::Result<()> { + // StartupMessage has a length prefix but no message-type byte. + let length = socket.read_u32().await?; + let mut body = vec![0; usize::try_from(length.checked_sub(4).unwrap()).unwrap()]; + socket.read_exact(&mut body).await?; + + // AuthenticationOk followed by ReadyForQuery(idle). + socket.write_all(b"R\0\0\0\x08\0\0\0\0Z\0\0\0\x05I").await?; + + if !respond_to_ping { + // Keep accepting client bytes, but never answer the release ping. + // Unlike a reset or EOF, this leaves the client's read pending forever. + tokio::io::copy(&mut socket, &mut tokio::io::sink()).await?; + return Ok(()); + } + + loop { + let mut tag = [0]; + if socket.read(&mut tag).await? == 0 { + return Ok(()); + } + let length = socket.read_u32().await?; + body.resize(usize::try_from(length.checked_sub(4).unwrap()).unwrap(), 0); + socket.read_exact(&mut body).await?; + + match tag[0] { + b'S' => socket.write_all(b"Z\0\0\0\x05I").await?, // Sync -> ReadyForQuery + b'X' => return Ok(()), // Terminate + tag => panic!("unexpected client message: {tag}"), + } + } +} + +fn pool_options() -> PgPoolOptions { + PgPoolOptions::new() + .max_connections(1) + .acquire_timeout(TEST_TIMEOUT) + .max_lifetime(None) + .idle_timeout(None) +} + +async fn assert_capacity_recovers( + server: &mut TestServer, + pool: &PgPool, + conn: PoolConnection, +) { + drop(conn); + server.wait_for_disconnect().await; + assert_eq!( + pool.size(), + 0, + "the discarded connection must leave the pool" + ); + assert_eq!( + pool.num_idle(), + 0, + "the stalled connection must not be reused" + ); + + let replacement = pool.acquire().await.expect("pool permit must be restored"); + assert_eq!(server.accepted.load(Ordering::SeqCst), 2); + assert_eq!(pool.size(), 1); + + drop(replacement.detach()); + pool.close().await; +} + +#[tokio::test] +async fn drop_recovers_capacity_after_unresponsive_release_ping() { + let mut server = TestServer::start(false).await; + let pool = pool_options().connect_lazy_with(server.connect_options()); + let conn = pool.acquire().await.unwrap(); + + assert_capacity_recovers(&mut server, &pool, conn).await; +} + +#[tokio::test] +async fn drop_recovers_capacity_after_hanging_after_release() { + let mut server = TestServer::start(true).await; + let pool = pool_options() + .after_release(|_, _| Box::pin(future::pending())) + .connect_lazy_with(server.connect_options()); + let conn = pool.acquire().await.unwrap(); + + // A timeout around only ping() cannot recover from a stalled callback. + assert_capacity_recovers(&mut server, &pool, conn).await; +} + +#[tokio::test] +async fn return_timeout_replenishes_min_connections() { + let mut server = TestServer::start(true).await; + let pool = pool_options() + .min_connections(1) + .after_release(|_, _| Box::pin(future::pending())) + .connect_lazy_with(server.connect_options()); + let conn = pool.acquire().await.unwrap(); + drop(conn); + + server.wait_for_disconnect().await; + timeout(TEST_TIMEOUT, async { + while pool.num_idle() != 1 { + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("return timeout must still run min_connections maintenance"); + assert_eq!(server.accepted.load(Ordering::SeqCst), 2); + assert_eq!(pool.size(), 1); + + pool.close().await; +} + +#[tokio::test] +async fn healthy_connection_returns_to_idle_queue() { + let server = TestServer::start(true).await; + let pool = pool_options() + .after_release(|_, _| Box::pin(async { Ok(true) })) + .connect_lazy_with(server.connect_options()); + let mut conn = pool.acquire().await.unwrap(); + + timeout(TEST_TIMEOUT, conn.return_to_pool()).await.unwrap(); + assert_eq!(pool.num_idle(), 1); + assert_eq!(pool.size(), 1); + + let replacement = pool.acquire().await.unwrap(); + assert_eq!(server.accepted.load(Ordering::SeqCst), 1); + drop(replacement.detach()); + pool.close().await; +}