diff --git a/Cargo.toml b/Cargo.toml index b7ed7ac2cb..de52af8afe 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -424,6 +424,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..590d75c5f5 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_PING_TIMEOUT: Duration = Duration::from_secs(5); /// A connection managed by a [`Pool`][crate::pool::Pool]. /// @@ -311,19 +312,32 @@ 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 + match crate::rt::timeout(RETURN_TO_POOL_PING_TIMEOUT, self.raw.ping()).await { + 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!( + timeout = ?RETURN_TO_POOL_PING_TIMEOUT, + "timed out while testing the connection on-release", + ); + + // The connection is unresponsive, so avoid all async connection I/O here. + // Dropping `self` synchronously releases the pool guard and discards the socket. + false + } } } diff --git a/tests/postgres/pool.rs b/tests/postgres/pool.rs new file mode 100644 index 0000000000..bb030b7874 --- /dev/null +++ b/tests/postgres/pool.rs @@ -0,0 +1,111 @@ +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, +}; +use std::time::Duration; + +use sqlx::postgres::{PgConnectOptions, PgPoolOptions, PgSslMode}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; + +const AUTHENTICATION_OK: &[u8] = b"R\0\0\0\x08\0\0\0\0"; +const BACKEND_KEY_DATA: &[u8] = b"K\0\0\0\x0c\0\0\0\x01\0\0\0\x02"; +const READY_FOR_QUERY: &[u8] = b"Z\0\0\0\x05I"; + +#[tokio::test] +async fn return_to_pool_ping_timeout_recovers_pool_capacity() -> anyhow::Result<()> { + let server = FakePostgresServer::bind().await?; + let options = PgConnectOptions::new() + .host("127.0.0.1") + .port(server.port()) + .username("postgres") + .database("postgres") + .ssl_mode(PgSslMode::Disable); + + let pool = PgPoolOptions::new() + .min_connections(0) + .max_connections(1) + .acquire_timeout(Duration::from_secs(8)) + .test_before_acquire(false) + .connect_with(options) + .await?; + + let conn = pool.acquire().await?; + assert_eq!(server.connection_count(), 1); + + drop(conn); + + let conn = pool.acquire().await?; + assert_eq!(server.connection_count(), 2); + + conn.close().await?; + pool.close().await; + + Ok(()) +} + +struct FakePostgresServer { + port: u16, + connection_count: Arc, +} + +impl FakePostgresServer { + async fn bind() -> std::io::Result { + let listener = TcpListener::bind(("127.0.0.1", 0)).await?; + let port = listener.local_addr()?.port(); + let connection_count = Arc::new(AtomicUsize::new(0)); + + tokio::spawn(accept_connections(listener, Arc::clone(&connection_count))); + + Ok(Self { + port, + connection_count, + }) + } + + fn port(&self) -> u16 { + self.port + } + + fn connection_count(&self) -> usize { + self.connection_count.load(Ordering::SeqCst) + } +} + +async fn accept_connections(listener: TcpListener, connection_count: Arc) { + while let Ok((socket, _)) = listener.accept().await { + connection_count.fetch_add(1, Ordering::SeqCst); + + tokio::spawn(async move { + let _ = handle_connection(socket).await; + }); + } +} + +async fn handle_connection(mut socket: TcpStream) -> std::io::Result<()> { + read_startup_message(&mut socket).await?; + + socket.write_all(AUTHENTICATION_OK).await?; + socket.write_all(BACKEND_KEY_DATA).await?; + socket.write_all(READY_FOR_QUERY).await?; + socket.flush().await?; + + let mut buf = [0_u8; 1024]; + + loop { + if socket.read(&mut buf).await? == 0 { + return Ok(()); + } + } +} + +async fn read_startup_message(socket: &mut TcpStream) -> std::io::Result<()> { + let mut len = [0_u8; 4]; + socket.read_exact(&mut len).await?; + + let len = u32::from_be_bytes(len) as usize; + let mut body = vec![0_u8; len.saturating_sub(4)]; + socket.read_exact(&mut body).await?; + + Ok(()) +}