diff --git a/sqlx-core/src/pool/connection.rs b/sqlx-core/src/pool/connection.rs index 7912b12aa1..07e914063e 100644 --- a/sqlx-core/src/pool/connection.rs +++ b/sqlx-core/src/pool/connection.rs @@ -15,6 +15,15 @@ use crate::pool::options::PoolConnectionMetadata; const CLOSE_ON_DROP_TIMEOUT: Duration = Duration::from_secs(5); +/// Bounds each step of returning a connection to the pool. +/// +/// Every step here does I/O on a connection that may be dead in a way the socket +/// cannot report: if the server disappeared without closing it, and the client has +/// nothing left to retransmit, reads never complete. The returning task holds the +/// pool's `DecrementSizeGuard` for as long as it runs, so an unbounded step leaks a +/// permit and the pool shrinks by one connection every time it happens. +const RETURN_TO_POOL_TIMEOUT: Duration = Duration::from_secs(5); + /// A connection managed by a [`Pool`][crate::pool::Pool]. /// /// Will be returned to the pool on-drop. @@ -275,32 +284,40 @@ impl Floating> { async fn return_to_pool(mut self) -> bool { // Immediately close the connection. if self.guard.pool.is_closed() { - self.close().await; + self.close_bounded().await; return false; } // If the connection is beyond max lifetime, close the connection and // immediately create a new connection if is_beyond_max_lifetime(&self.inner, &self.guard.pool.options) { - self.close().await; + self.close_bounded().await; return false; } if let Some(test) = &self.guard.pool.options.after_release { let meta = self.metadata(); - match (test)(&mut self.inner.raw, meta).await { - Ok(true) => (), - Ok(false) => { - self.close().await; + let result = + crate::rt::timeout(RETURN_TO_POOL_TIMEOUT, (test)(&mut self.inner.raw, meta)).await; + + match result { + Ok(Ok(true)) => (), + Ok(Ok(false)) => { + self.close_bounded().await; return false; } - Err(error) => { + Ok(Err(error)) => { tracing::warn!(%error, "error from `after_release`"); // Connection is broken, don't try to gracefully close as // something weird might happen. self.close_hard().await; return false; } + Err(_) => { + tracing::warn!("timed out in `after_release`; discarding the connection"); + self.close_hard().await; + return false; + } } } @@ -311,22 +328,42 @@ impl Floating> { // returned to the pool; also of course, if it was dropped due to an error // this is simply a band-aid as SQLx-next connections should be able // to recover from cancellations - if let Err(error) = self.raw.ping().await { - tracing::warn!( - %error, - "error occurred while testing the connection on-release", - ); - - // Connection is broken, don't try to gracefully close. - self.close_hard().await; - false - } else { - // if the connection is still viable, release it to the pool - self.release(); - true + let result = crate::rt::timeout(RETURN_TO_POOL_TIMEOUT, self.raw.ping()).await; + + match result { + Ok(Ok(())) => { + // if the connection is still viable, release it to the pool + self.release(); + true + } + Ok(Err(error)) => { + tracing::warn!( + %error, + "error occurred while testing the connection on-release", + ); + + // Connection is broken, don't try to gracefully close. + self.close_hard().await; + false + } + Err(_) => { + tracing::warn!("timed out while testing the connection on-release; discarding it",); + + // The socket is not answering; dropping it is the only way to get the + // pool permit back. + self.close_hard().await; + false + } } } + /// Close the connection, giving up on a graceful close if it does not complete + /// promptly. Cancelling `close()` still drops the connection and releases the + /// pool permit, which is the outcome that matters here. + async fn close_bounded(self) { + let _ = crate::rt::timeout(RETURN_TO_POOL_TIMEOUT, self.close()).await; + } + pub async fn close(self) { // This isn't used anywhere that we care about the return value let _ = self.inner.raw.close().await; diff --git a/tests/postgres/postgres.rs b/tests/postgres/postgres.rs index 126771565a..5e2c55b97a 100644 --- a/tests/postgres/postgres.rs +++ b/tests/postgres/postgres.rs @@ -12,7 +12,7 @@ use sqlx_test::{new, pool, setup_if_needed}; use std::env; use std::pin::{pin, Pin}; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; #[sqlx_macros::test] async fn it_connects() -> anyhow::Result<()> { @@ -321,6 +321,34 @@ async fn it_can_fail_and_recover() -> anyhow::Result<()> { Ok(()) } +/// A hook that never completes must not hold the pool permit forever: returning a +/// connection to the pool does I/O on a socket that may be dead in a way it cannot +/// report, so each step is bounded and the connection is discarded on timeout. +#[sqlx_macros::test] +async fn pool_recovers_from_a_hanging_after_release() -> anyhow::Result<()> { + setup_if_needed(); + + let pool = PgPoolOptions::new() + .max_connections(1) + .acquire_timeout(Duration::from_secs(30)) + .after_release(|_conn, _meta| Box::pin(std::future::pending())) + .connect(&env::var("DATABASE_URL")?) + .await?; + + let conn = pool.acquire().await?; + drop(conn); + + // The permit comes back once the bounded return-to-pool gives up on the hook. + let started_at = Instant::now(); + let _conn = pool.acquire().await?; + assert!( + started_at.elapsed() < Duration::from_secs(30), + "acquire waited on a connection whose return never finished", + ); + + Ok(()) +} + #[sqlx_macros::test] async fn it_can_fail_and_recover_with_pool() -> anyhow::Result<()> { let pool = sqlx_test::pool::().await?;