fix(pool): bound each step of returning a connection to the pool

Fixes #4349.

`Floating::return_to_pool` holds the pool's `DecrementSizeGuard` while it
does I/O on the connection: an `after_release` hook, a `ping()` to check
the connection is still usable, or a graceful `close()`. None of those were
bounded.

That matters because the I/O can be on a socket that is dead in a way the
socket cannot report. If the server disappeared without closing the
connection, and the client has nothing left to retransmit, no RST is ever
provoked and reads never complete. The returning task then holds its permit
forever and the pool shrinks by one connection. After `max_connections` of
those, every `acquire()` fails with `PoolTimedOut` and the pool never
recovers, even though the server is back.

Each step now gets `RETURN_TO_POOL_TIMEOUT` (5s, matching the existing
`CLOSE_ON_DROP_TIMEOUT`), and on expiry the connection is dropped via
`close_hard()`, which does no I/O. Cancelling `close()` likewise drops the
connection and releases the permit.

The regression test uses an `after_release` hook that never completes, which
is a deterministic stand-in for a socket that never answers: before this
change it exhausts `acquire_timeout`, after it the pool is usable again once
the bound expires.

I realise #3582 reworks this area and includes a timeout of its own; this is
meant for main until that lands.

Author: adriangb

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<DB: Database> Floating<DB, Live<DB>> {
     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<DB: Database> Floating<DB, Live<DB>> {
         // 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;

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::<Postgres>().await?;