diff --git a/sqlx-core/src/pool/inner.rs b/sqlx-core/src/pool/inner.rs index b698dc9df0..1e42721b27 100644 --- a/sqlx-core/src/pool/inner.rs +++ b/sqlx-core/src/pool/inner.rs @@ -289,7 +289,37 @@ impl PoolInner { }; // Attempt to connect... - return self.connect(deadline, guard).await; + // + // If `connect_timeout` is set, spawn the connection as a separate task + // so that cancellation of `acquire()` doesn't abort the connection attempt. + // If the connection succeeds but no one picks it up, it's returned to the pool. + if let Some(connect_timeout) = self.options.connect_timeout { + let connect_deadline = acquire_started_at + self.options.acquire_timeout + connect_timeout; + let pool = (*self).clone(); + + let (tx, rx) = futures_intrusive::channel::shared::oneshot_channel(); + + crate::rt::spawn(async move { + let result = pool.connect(connect_deadline, guard).await; + + // If `acquire()` is still waiting, send the result. + // If `acquire()` was cancelled (rx dropped), send() will fail + // and we return the connection to the pool. + match tx.send(result) { + Ok(()) => {} + Err(futures_intrusive::channel::ChannelSendError(result)) => { + // acquire() was cancelled, return connection to idle queue + if let Ok(live) = result { + pool.release(live); + } + } + } + }); + + return rx.receive().await.ok_or(Error::PoolTimedOut)?; + } else { + return self.connect(deadline, guard).await; + } } } ) diff --git a/sqlx-core/src/pool/options.rs b/sqlx-core/src/pool/options.rs index 3d048f1795..40beb4ff93 100644 --- a/sqlx-core/src/pool/options.rs +++ b/sqlx-core/src/pool/options.rs @@ -83,6 +83,7 @@ pub struct PoolOptions { pub(crate) max_lifetime: Option, pub(crate) idle_timeout: Option, pub(crate) fair: bool, + pub(crate) connect_timeout: Option, pub(crate) parent_pool: Option>, } @@ -106,6 +107,7 @@ impl Clone for PoolOptions { max_lifetime: self.max_lifetime, idle_timeout: self.idle_timeout, fair: self.fair, + connect_timeout: self.connect_timeout, parent_pool: self.parent_pool.clone(), } } @@ -161,6 +163,7 @@ impl PoolOptions { idle_timeout: Some(Duration::from_secs(10 * 60)), max_lifetime: Some(Duration::from_secs(30 * 60)), fair: true, + connect_timeout: None, parent_pool: None, } } @@ -307,6 +310,30 @@ impl PoolOptions { self.idle_timeout } + /// Set a separate timeout for establishing new connections. + /// + /// When set, this timeout is used specifically for the connection establishment phase + /// of [`Pool::acquire()`], separate from [`acquire_timeout`][Self::acquire_timeout]. + /// This allows the connection to have its own timeout that doesn't share the budget + /// with semaphore acquisition and idle connection checks. + /// + /// If not set, the connection establishment uses the remaining time from `acquire_timeout`. + /// + /// This is useful when: + /// * You want a longer timeout for establishing connections than for acquiring from the pool. + /// * You want connection attempts to survive cancellation of the `acquire()` future. + /// When `acquire()` is cancelled, the spawned connection task continues in the background. + /// If it succeeds, the connection is returned to the pool's idle queue. + pub fn connect_timeout(mut self, timeout: Duration) -> Self { + self.connect_timeout = Some(timeout); + self + } + + /// Get the connect timeout, if set. + pub fn get_connect_timeout(&self) -> Option { + self.connect_timeout + } + /// If true, the health of a connection will be verified by a call to [`Connection::ping`] /// before returning the connection. /// @@ -587,9 +614,10 @@ impl Debug for PoolOptions { f.debug_struct("PoolOptions") .field("max_connections", &self.max_connections) .field("min_connections", &self.min_connections) - .field("connect_timeout", &self.acquire_timeout) + .field("acquire_timeout", &self.acquire_timeout) .field("max_lifetime", &self.max_lifetime) .field("idle_timeout", &self.idle_timeout) + .field("connect_timeout", &self.connect_timeout) .field("test_before_acquire", &self.test_before_acquire) .finish() }