From be545ed1c609480325b054615538c657111867bb Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Fri, 24 Jul 2026 21:10:19 +0000 Subject: [PATCH] feat: TCP keepalive for Postgres and MySQL connections Adds `TcpKeepalive` and `net::connect_tcp_with_keepalive`, plus `tcp_keepalive()` on `PgConnectOptions` and `MySqlConnectOptions` and libpq-compatible URL parameters for Postgres (`keepalives`, `keepalives_idle`, `keepalives_interval`, `keepalives_count`). Keepalive is off by default, so nothing changes for existing users. Motivation: without it, a connection whose server disappeared *without* closing the socket never finds out. A failover, a killed container, or a dropped NAT mapping leaves the client blocked reading a response that will never arrive; there is nothing left to retransmit, so no RST is ever provoked and the read waits forever. `TCP_NODELAY`, which is all SQLx sets today, does not help, and a server-side `statement_timeout` cannot fire on a server that is gone. This is the case @abonander allowed for in https://github.com/launchbadge/sqlx/pull/3559#issuecomment-2415299442: "I suppose that's still preferable to it hanging forever on a read that will never complete." Worth adding on the objection raised there, that a keepalive timeout is only noticed the next time the socket is used: that is not true of a blocked reader. When the probes are exhausted the kernel sets `sk_err` to `ETIMEDOUT` and wakes anyone parked on the socket, so the pending read fails rather than waiting for someone to poke it. We hit this in production: a maintenance loop that had a query in flight when its Postgres instance went away stopped doing work permanently, while other loops in the same process recovered in seconds. Reproduced by holding an `ACCESS EXCLUSIVE` lock so the query blocks server-side, then removing the server from the network and restarting it. Implementation follows #3559 by @xuehaonan27, rebased onto the current `connect_tcp` and reduced to an additive API: `connect_tcp` keeps its signature and delegates, so external drivers are unaffected. --- sqlx-core/Cargo.toml | 1 + sqlx-core/src/net/mod.rs | 3 +- sqlx-core/src/net/socket/mod.rs | 73 ++++++++++++++-- sqlx-core/src/net/socket/tcp_keepalive.rs | 100 ++++++++++++++++++++++ sqlx-mysql/src/connection/establish.rs | 10 ++- sqlx-mysql/src/options/mod.rs | 28 ++++++ sqlx-postgres/src/connection/stream.rs | 10 ++- sqlx-postgres/src/options/mod.rs | 32 +++++++ sqlx-postgres/src/options/parse.rs | 62 ++++++++++++++ src/lib.rs | 1 + 10 files changed, 309 insertions(+), 11 deletions(-) create mode 100644 sqlx-core/src/net/socket/tcp_keepalive.rs diff --git a/sqlx-core/Cargo.toml b/sqlx-core/Cargo.toml index 90ed446b4b..d39ef4c1e4 100644 --- a/sqlx-core/Cargo.toml +++ b/sqlx-core/Cargo.toml @@ -100,6 +100,7 @@ url = { version = "2.2.2" } bstr = { version = "1.0.1", default-features = false, features = ["std"], optional = true } hashlink = "0.11.0" indexmap = "2.0" +socket2 = { version = "0.6", features = ["all"] } event-listener = "5.2.0" hashbrown = "0.16.0" diff --git a/sqlx-core/src/net/mod.rs b/sqlx-core/src/net/mod.rs index f9c43668ab..65ede3db68 100644 --- a/sqlx-core/src/net/mod.rs +++ b/sqlx-core/src/net/mod.rs @@ -2,5 +2,6 @@ mod socket; pub mod tls; pub use socket::{ - connect_tcp, connect_uds, BufferedSocket, Socket, SocketIntoBox, WithSocket, WriteBuffer, + connect_tcp, connect_tcp_with_keepalive, connect_uds, BufferedSocket, Socket, SocketIntoBox, + TcpKeepalive, WithSocket, WriteBuffer, }; diff --git a/sqlx-core/src/net/socket/mod.rs b/sqlx-core/src/net/socket/mod.rs index 0f9aae61b4..e20e1205aa 100644 --- a/sqlx-core/src/net/socket/mod.rs +++ b/sqlx-core/src/net/socket/mod.rs @@ -7,10 +7,12 @@ use std::task::{ready, Context, Poll}; pub use buffered::{BufferedSocket, WriteBuffer}; use bytes::BufMut; use cfg_if::cfg_if; +pub use tcp_keepalive::TcpKeepalive; use crate::io::ReadBuf; mod buffered; +mod tcp_keepalive; pub trait Socket: Send + Sync + Unpin + 'static { fn try_read(&mut self, buf: &mut dyn ReadBuf) -> io::Result; @@ -185,30 +187,78 @@ pub async fn connect_tcp( host: &str, port: u16, with_socket: Ws, +) -> crate::Result { + connect_tcp_with_keepalive(host, port, with_socket, None).await +} + +/// Open a TCP socket to `host` and `port`, optionally configuring TCP keepalive on it. +/// +/// Without keepalive, a connection whose server disappeared without closing the socket +/// (a failover, a killed container, a dropped NAT mapping) is only discovered the next +/// time the client writes to it. A connection blocked reading a response waits forever. +pub async fn connect_tcp_with_keepalive( + host: &str, + port: u16, + with_socket: Ws, + keepalive: Option<&TcpKeepalive>, ) -> crate::Result { #[cfg(feature = "_rt-tokio")] if crate::rt::rt_tokio::available() { - return Ok(with_socket - .with_socket(tokio::net::TcpStream::connect((host, port)).await?) - .await); + let stream = tokio::net::TcpStream::connect((host, port)).await?; + set_tcp_keepalive(&stream, keepalive)?; + + return Ok(with_socket.with_socket(stream).await); } cfg_if! { if #[cfg(feature = "_rt-async-io")] { - Ok(with_socket.with_socket(connect_tcp_async_io(host, port).await?).await) + // Keepalive is applied inside, on the concrete socket type. + Ok(with_socket.with_socket(connect_tcp_async_io(host, port, keepalive).await?).await) } else { - crate::rt::missing_rt((host, port, with_socket)) + crate::rt::missing_rt((host, port, with_socket, keepalive)) } } } +#[cfg(all(unix, any(feature = "_rt-tokio", feature = "_rt-async-io")))] +fn set_tcp_keepalive( + stream: &S, + keepalive: Option<&TcpKeepalive>, +) -> crate::Result<()> { + let Some(keepalive) = keepalive else { + return Ok(()); + }; + + socket2::SockRef::from(stream).set_tcp_keepalive(&keepalive.to_socket2())?; + + Ok(()) +} + +#[cfg(all(windows, any(feature = "_rt-tokio", feature = "_rt-async-io")))] +fn set_tcp_keepalive( + stream: &S, + keepalive: Option<&TcpKeepalive>, +) -> crate::Result<()> { + let Some(keepalive) = keepalive else { + return Ok(()); + }; + + socket2::SockRef::from(stream).set_tcp_keepalive(&keepalive.to_socket2())?; + + Ok(()) +} + /// Open a TCP socket to `host` and `port`. /// /// If `host` is a hostname, attempt to connect to each address it resolves to. /// /// This implements the same behavior as [`tokio::net::TcpStream::connect()`]. #[cfg(feature = "_rt-async-io")] -async fn connect_tcp_async_io(host: &str, port: u16) -> crate::Result { +async fn connect_tcp_async_io( + host: &str, + port: u16, + keepalive: Option<&TcpKeepalive>, +) -> crate::Result { use async_io::Async; use std::net::{IpAddr, TcpStream, ToSocketAddrs}; @@ -216,7 +266,10 @@ async fn connect_tcp_async_io(host: &str, port: u16) -> crate::Result() { - return Ok(Async::::connect((addr, port)).await?); + let stream = Async::::connect((addr, port)).await?; + set_tcp_keepalive(&stream, keepalive)?; + + return Ok(stream); } let host = host.to_string(); @@ -232,7 +285,11 @@ async fn connect_tcp_async_io(host: &str, port: u16) -> crate::Result::connect(socket_addr).await { - Ok(stream) => return Ok(stream), + Ok(stream) => { + set_tcp_keepalive(&stream, keepalive)?; + + return Ok(stream); + } Err(e) => last_err = Some(e), } } diff --git a/sqlx-core/src/net/socket/tcp_keepalive.rs b/sqlx-core/src/net/socket/tcp_keepalive.rs new file mode 100644 index 0000000000..50dc267e55 --- /dev/null +++ b/sqlx-core/src/net/socket/tcp_keepalive.rs @@ -0,0 +1,100 @@ +use std::time::Duration; + +/// TCP keepalive parameters for a connection's socket. +/// +/// Keepalive is what lets a client notice that its server is gone when the +/// server disappeared without closing the socket: a failover, a killed +/// container, a NAT table that dropped the mapping. Without it, a connection +/// blocked reading a response it will never receive waits forever, because +/// there is nothing left to retransmit and so nothing to time out. +/// +/// The parameters mirror libpq's `keepalives_idle`, `keepalives_interval` and +/// `keepalives_count`, and are applied with `setsockopt` after connecting. +/// +/// Support varies by platform: `interval` is ignored on OpenBSD and Solaris, +/// and `retries` is ignored on OpenBSD, Solaris, watchOS and tvOS. `idle` is +/// supported everywhere keepalive itself is. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct TcpKeepalive { + /// Idle time after which the first keepalive probe is sent. + pub idle: Option, + /// Time between probes once the first one has been sent. + pub interval: Option, + /// Number of unacknowledged probes before the connection is dropped. + pub retries: Option, +} + +impl TcpKeepalive { + /// Keepalive with no parameters overridden, leaving the system defaults in + /// place. On Linux those are 7200s idle, 75s interval, 9 retries. + pub fn new() -> Self { + Self::default() + } + + /// Sets the idle time after which the first keepalive probe is sent. + pub fn with_idle(mut self, idle: Duration) -> Self { + self.idle = Some(idle); + self + } + + /// Sets the time between keepalive probes. + pub fn with_interval(mut self, interval: Duration) -> Self { + self.interval = Some(interval); + self + } + + /// Sets the number of unacknowledged probes before the connection is dropped. + pub fn with_retries(mut self, retries: u32) -> Self { + self.retries = Some(retries); + self + } + + #[cfg(any(feature = "_rt-tokio", feature = "_rt-async-io"))] + pub(crate) fn to_socket2(self) -> socket2::TcpKeepalive { + let mut keepalive = socket2::TcpKeepalive::new(); + + if let Some(idle) = self.idle { + keepalive = keepalive.with_time(idle); + } + + #[cfg(not(any(target_os = "openbsd", target_os = "solaris")))] + if let Some(interval) = self.interval { + keepalive = keepalive.with_interval(interval); + } + + #[cfg(not(any( + target_os = "openbsd", + target_os = "solaris", + target_os = "watchos", + target_os = "tvos", + )))] + if let Some(retries) = self.retries { + keepalive = keepalive.with_retries(retries); + } + + keepalive + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn keepalive_is_disabled_by_default() { + assert_eq!(TcpKeepalive::new(), TcpKeepalive::default()); + assert!(TcpKeepalive::new().idle.is_none()); + } + + #[test] + fn builders_set_each_parameter() { + let keepalive = TcpKeepalive::new() + .with_idle(Duration::from_secs(30)) + .with_interval(Duration::from_secs(10)) + .with_retries(3); + + assert_eq!(keepalive.idle, Some(Duration::from_secs(30))); + assert_eq!(keepalive.interval, Some(Duration::from_secs(10))); + assert_eq!(keepalive.retries, Some(3)); + } +} diff --git a/sqlx-mysql/src/connection/establish.rs b/sqlx-mysql/src/connection/establish.rs index f61654d876..ae3441a1c6 100644 --- a/sqlx-mysql/src/connection/establish.rs +++ b/sqlx-mysql/src/connection/establish.rs @@ -17,7 +17,15 @@ impl MySqlConnection { let handshake = match &options.socket { Some(path) => crate::net::connect_uds(path, do_handshake).await?, - None => crate::net::connect_tcp(&options.host, options.port, do_handshake).await?, + None => { + crate::net::connect_tcp_with_keepalive( + &options.host, + options.port, + do_handshake, + options.tcp_keepalive.as_ref(), + ) + .await? + } }; let stream = handshake?; diff --git a/sqlx-mysql/src/options/mod.rs b/sqlx-mysql/src/options/mod.rs index 421bfb700e..038fd72895 100644 --- a/sqlx-mysql/src/options/mod.rs +++ b/sqlx-mysql/src/options/mod.rs @@ -4,6 +4,7 @@ mod connect; mod parse; mod ssl_mode; +use crate::net::TcpKeepalive; use crate::{connection::LogSettings, net::tls::CertificateInput}; pub use ssl_mode::MySqlSslMode; @@ -80,6 +81,7 @@ pub struct MySqlConnectOptions { pub(crate) no_engine_substitution: bool, pub(crate) timezone: Option, pub(crate) set_names: bool, + pub(crate) tcp_keepalive: Option, } impl Default for MySqlConnectOptions { @@ -105,6 +107,7 @@ impl MySqlConnectOptions { ssl_client_cert: None, ssl_client_key: None, statement_cache_capacity: 100, + tcp_keepalive: None, log_settings: Default::default(), pipes_as_concat: true, enable_cleartext_plugin: false, @@ -286,6 +289,31 @@ impl MySqlConnectOptions { self } + /// Configure TCP keepalive on the connection's socket. + /// + /// Disabled by default, matching the socket default. Enable it when connections + /// are long-lived and the server may disappear without closing the socket (a + /// failover, a killed container, a dropped NAT mapping): without keepalive, a + /// connection blocked reading a response that will never arrive waits forever. + /// + /// # Example + /// + /// ```rust + /// # use std::time::Duration; + /// # use sqlx_core::net::TcpKeepalive; + /// # use sqlx_mysql::MySqlConnectOptions; + /// let options = MySqlConnectOptions::new().tcp_keepalive( + /// TcpKeepalive::new() + /// .with_idle(Duration::from_secs(30)) + /// .with_interval(Duration::from_secs(10)) + /// .with_retries(3), + /// ); + /// ``` + pub fn tcp_keepalive(mut self, keepalive: TcpKeepalive) -> Self { + self.tcp_keepalive = Some(keepalive); + self + } + /// Sets the capacity of the connection's statement cache in a number of stored /// distinct statements. Caching is handled using LRU, meaning when the /// amount of queries hits the defined limit, the oldest statement will get diff --git a/sqlx-postgres/src/connection/stream.rs b/sqlx-postgres/src/connection/stream.rs index e8a1aedc47..81fb8bfacc 100644 --- a/sqlx-postgres/src/connection/stream.rs +++ b/sqlx-postgres/src/connection/stream.rs @@ -44,7 +44,15 @@ impl PgStream { pub(super) async fn connect(options: &PgConnectOptions) -> Result { let socket_result = match options.fetch_socket() { Some(ref path) => net::connect_uds(path, MaybeUpgradeTls(options)).await?, - None => net::connect_tcp(&options.host, options.port, MaybeUpgradeTls(options)).await?, + None => { + net::connect_tcp_with_keepalive( + &options.host, + options.port, + MaybeUpgradeTls(options), + options.tcp_keepalive.as_ref(), + ) + .await? + } }; let socket = socket_result?; diff --git a/sqlx-postgres/src/options/mod.rs b/sqlx-postgres/src/options/mod.rs index 21e6628cae..94ecf83b14 100644 --- a/sqlx-postgres/src/options/mod.rs +++ b/sqlx-postgres/src/options/mod.rs @@ -5,6 +5,7 @@ use std::path::{Path, PathBuf}; pub use ssl_mode::PgSslMode; +use crate::net::TcpKeepalive; use crate::{connection::LogSettings, net::tls::CertificateInput}; mod connect; @@ -30,6 +31,7 @@ pub struct PgConnectOptions { pub(crate) log_settings: LogSettings, pub(crate) extra_float_digits: Option>, pub(crate) options: Option, + pub(crate) tcp_keepalive: Option, } impl Default for PgConnectOptions { @@ -95,6 +97,7 @@ impl PgConnectOptions { statement_cache_capacity: 100, application_name: var("PGAPPNAME").ok(), extra_float_digits: Some("2".into()), + tcp_keepalive: None, log_settings: Default::default(), options: var("PGOPTIONS").ok(), } @@ -421,6 +424,35 @@ impl PgConnectOptions { self } + /// Configure TCP keepalive on the connection's socket. + /// + /// Disabled by default, matching the socket default. Enable it when connections + /// are long-lived and the server may disappear without closing the socket (a + /// failover, a killed container, a dropped NAT mapping): without keepalive, a + /// connection blocked reading a response that will never arrive waits forever. + /// + /// The equivalent libpq parameters (`keepalives`, `keepalives_idle`, + /// `keepalives_interval`, `keepalives_count`) are also accepted in the + /// connection URL. + /// + /// # Example + /// + /// ```rust + /// # use std::time::Duration; + /// # use sqlx_core::net::TcpKeepalive; + /// # use sqlx_postgres::PgConnectOptions; + /// let options = PgConnectOptions::new().tcp_keepalive( + /// TcpKeepalive::new() + /// .with_idle(Duration::from_secs(30)) + /// .with_interval(Duration::from_secs(10)) + /// .with_retries(3), + /// ); + /// ``` + pub fn tcp_keepalive(mut self, keepalive: TcpKeepalive) -> Self { + self.tcp_keepalive = Some(keepalive); + self + } + /// Set additional startup options for the connection as a list of key-value pairs. /// /// Escapes the options’ backslash and space characters as per diff --git a/sqlx-postgres/src/options/parse.rs b/sqlx-postgres/src/options/parse.rs index e911305698..1f53c8ad75 100644 --- a/sqlx-postgres/src/options/parse.rs +++ b/sqlx-postgres/src/options/parse.rs @@ -4,6 +4,7 @@ use sqlx_core::percent_encoding::{percent_decode_str, utf8_percent_encode, NON_A use sqlx_core::Url; use std::net::IpAddr; use std::str::FromStr; +use std::time::Duration; impl PgConnectOptions { pub(crate) fn parse_from_url(url: &Url) -> Result { @@ -104,6 +105,40 @@ impl PgConnectOptions { } } + // libpq-compatible keepalive parameters. + "keepalives" => { + options.tcp_keepalive = match &*value { + "0" => None, + _ => Some(options.tcp_keepalive.unwrap_or_default()), + }; + } + + "keepalives_idle" => { + let idle = Duration::from_secs(value.parse().map_err(Error::config)?); + options.tcp_keepalive = + Some(options.tcp_keepalive.unwrap_or_default().with_idle(idle)); + } + + "keepalives_interval" => { + let interval = Duration::from_secs(value.parse().map_err(Error::config)?); + options.tcp_keepalive = Some( + options + .tcp_keepalive + .unwrap_or_default() + .with_interval(interval), + ); + } + + "keepalives_count" => { + let retries = value.parse().map_err(Error::config)?; + options.tcp_keepalive = Some( + options + .tcp_keepalive + .unwrap_or_default() + .with_retries(retries), + ); + } + _ => tracing::warn!(%key, %value, "ignoring unrecognized connect parameter"), } } @@ -188,6 +223,33 @@ fn it_parses_socket_correctly_from_parameter() { assert_eq!(Some("/var/run/postgres/".into()), opts.socket); } +#[test] +fn it_parses_libpq_keepalive_parameters() { + let url = + "postgres://user@localhost/db?keepalives_idle=30&keepalives_interval=10&keepalives_count=3"; + let opts = PgConnectOptions::from_str(url).unwrap(); + let keepalive = opts.tcp_keepalive.expect("keepalive should be enabled"); + + assert_eq!(keepalive.idle, Some(Duration::from_secs(30))); + assert_eq!(keepalive.interval, Some(Duration::from_secs(10))); + assert_eq!(keepalive.retries, Some(3)); +} + +#[test] +fn keepalive_is_off_unless_asked_for() { + let opts = PgConnectOptions::from_str("postgres://user@localhost/db").unwrap(); + assert_eq!(opts.tcp_keepalive, None); + + let opts = PgConnectOptions::from_str("postgres://user@localhost/db?keepalives=1").unwrap(); + assert_eq!(opts.tcp_keepalive, Some(Default::default())); + + // `keepalives=0` disables it again, even after another parameter turned it on. + let opts = + PgConnectOptions::from_str("postgres://user@localhost/db?keepalives_idle=30&keepalives=0") + .unwrap(); + assert_eq!(opts.tcp_keepalive, None); +} + #[test] fn it_parses_host_correctly_from_parameter() { let url = "postgres:///?host=google.database.com"; diff --git a/src/lib.rs b/src/lib.rs index 438463210d..d373afb4e7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,6 +19,7 @@ pub use sqlx_core::database::{self, Database}; pub use sqlx_core::describe::Describe; pub use sqlx_core::executor::{Execute, Executor}; pub use sqlx_core::from_row::FromRow; +pub use sqlx_core::net::TcpKeepalive; pub use sqlx_core::pool::{self, Pool}; #[doc(hidden)] pub use sqlx_core::query::query_with_result as __query_with_result;