feat: TCP keepalive for Postgres and MySQL connections

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"
 

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,
 };

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<usize>;
@@ -185,38 +187,89 @@ pub async fn connect_tcp<Ws: WithSocket>(
     host: &str,
     port: u16,
     with_socket: Ws,
+) -> crate::Result<Ws::Output> {
+    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<Ws: WithSocket>(
+    host: &str,
+    port: u16,
+    with_socket: Ws,
+    keepalive: Option<&TcpKeepalive>,
 ) -> crate::Result<Ws::Output> {
     #[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<S: std::os::fd::AsFd>(
+    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<S: std::os::windows::io::AsSocket>(
+    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<impl Socket> {
+async fn connect_tcp_async_io(
+    host: &str,
+    port: u16,
+    keepalive: Option<&TcpKeepalive>,
+) -> crate::Result<impl Socket> {
     use async_io::Async;
     use std::net::{IpAddr, TcpStream, ToSocketAddrs};
 
     // IPv6 addresses in URLs will be wrapped in brackets and the `url` crate doesn't trim those.
     let host = host.trim_matches(&['[', ']'][..]);
 
     if let Ok(addr) = host.parse::<IpAddr>() {
-        return Ok(Async::<TcpStream>::connect((addr, port)).await?);
+        let stream = Async::<TcpStream>::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<impl Socke
     // Loop through all the Socket Addresses that the hostname resolves to
     for socket_addr in addresses {
         match Async::<TcpStream>::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),
         }
     }

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<Duration>,
+    /// Time between probes once the first one has been sent.
+    pub interval: Option<Duration>,
+    /// Number of unacknowledged probes before the connection is dropped.
+    pub retries: Option<u32>,
+}
+
+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));
+    }
+}

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?;

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<String>,
     pub(crate) set_names: bool,
+    pub(crate) tcp_keepalive: Option<TcpKeepalive>,
 }
 
 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

sqlx-postgres/src/connection/stream.rs

@@ -44,7 +44,15 @@ impl PgStream {
     pub(super) async fn connect(options: &PgConnectOptions) -> Result<Self, Error> {
         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?;

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<Cow<'static, str>>,
     pub(crate) options: Option<String>,
+    pub(crate) tcp_keepalive: Option<TcpKeepalive>,
 }
 
 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