From 18ee7eb1cd0dc0e5bc220606342a261f677a5c95 Mon Sep 17 00:00:00 2001
From: Paul Hemberger <paul@materialize.com>
Date: Wed, 22 Feb 2023 16:13:54 -0500
Subject: [PATCH] add tcp user timeout config

---
 tokio-postgres/Cargo.toml            |  2 +-
 tokio-postgres/src/cancel_query.rs   |  1 +
 tokio-postgres/src/client.rs         |  1 +
 tokio-postgres/src/config.rs         | 14 ++++++++++++++
 tokio-postgres/src/connect.rs        |  2 ++
 tokio-postgres/src/connect_socket.rs | 14 +++++++++++++-
 tokio-postgres/src/error/mod.rs      |  5 +++++
 7 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/tokio-postgres/Cargo.toml b/tokio-postgres/Cargo.toml
index 4020a96bf..a413de829 100644
--- a/tokio-postgres/Cargo.toml
+++ b/tokio-postgres/Cargo.toml
@@ -56,7 +56,7 @@ phf = "0.11"
 postgres-protocol = { version = "0.6.4", path = "../postgres-protocol" }
 postgres-types = { version = "0.2.4", path = "../postgres-types" }
 serde = { version = "1.0", optional = true }
-socket2 = "0.4"
+socket2 = { version = "0.4.7", features = ["all"] }
 tokio = { version = "1.0", features = ["io-util"] }
 tokio-util = { version = "0.7", features = ["codec"] }
 
diff --git a/tokio-postgres/src/cancel_query.rs b/tokio-postgres/src/cancel_query.rs
index b02729f85..022278804 100644
--- a/tokio-postgres/src/cancel_query.rs
+++ b/tokio-postgres/src/cancel_query.rs
@@ -38,6 +38,7 @@ where
         &config.host,
         config.port,
         config.connect_timeout,
+        config.user_timeout,
         config.keepalive.as_ref(),
     )
     .await?;
diff --git a/tokio-postgres/src/client.rs b/tokio-postgres/src/client.rs
index eea779f77..7093b4165 100644
--- a/tokio-postgres/src/client.rs
+++ b/tokio-postgres/src/client.rs
@@ -158,6 +158,7 @@ pub(crate) struct SocketConfig {
     pub host: Host,
     pub port: u16,
     pub connect_timeout: Option<Duration>,
+    pub user_timeout: Option<Duration>,
     pub keepalive: Option<KeepaliveConfig>,
 }
 
diff --git a/tokio-postgres/src/config.rs b/tokio-postgres/src/config.rs
index 75cb896e4..eb9eb0408 100644
--- a/tokio-postgres/src/config.rs
+++ b/tokio-postgres/src/config.rs
@@ -183,6 +183,7 @@ pub struct Config {
     pub(crate) host: Vec<Host>,
     pub(crate) port: Vec<u16>,
     pub(crate) connect_timeout: Option<Duration>,
+    pub(crate) user_timeout: Option<Duration>,
     pub(crate) keepalives: bool,
     pub(crate) keepalive_config: KeepaliveConfig,
     pub(crate) target_session_attrs: TargetSessionAttrs,
@@ -217,6 +218,7 @@ impl Config {
             host: vec![],
             port: vec![],
             connect_timeout: None,
+            user_timeout: None,
             keepalives: true,
             keepalive_config,
             target_session_attrs: TargetSessionAttrs::Any,
@@ -407,6 +409,18 @@ impl Config {
         self.connect_timeout.as_ref()
     }
 
+    /// Sets the TCP user timeout.
+    pub fn user_timeout(&mut self, user_timeout: Duration) -> &mut Config {
+        self.user_timeout = Some(user_timeout);
+        self
+    }
+
+    /// Gets the TCP user timeout, if one has been set with the
+    /// `user_timeout` method.
+    pub fn get_user_timeout(&self) -> Option<&Duration> {
+        self.user_timeout.as_ref()
+    }
+
     /// Controls the use of TCP keepalive.
     ///
     /// This is ignored for Unix domain socket connections. Defaults to `true`.
diff --git a/tokio-postgres/src/connect.rs b/tokio-postgres/src/connect.rs
index 97a00c812..0006be154 100644
--- a/tokio-postgres/src/connect.rs
+++ b/tokio-postgres/src/connect.rs
@@ -65,6 +65,7 @@ where
         host,
         port,
         config.connect_timeout,
+        config.user_timeout,
         if config.keepalives {
             Some(&config.keepalive_config)
         } else {
@@ -118,6 +119,7 @@ where
         host: host.clone(),
         port,
         connect_timeout: config.connect_timeout,
+        user_timeout: config.user_timeout,
         keepalive: if config.keepalives {
             Some(config.keepalive_config.clone())
         } else {
diff --git a/tokio-postgres/src/connect_socket.rs b/tokio-postgres/src/connect_socket.rs
index 19d01d87a..1a8331beb 100644
--- a/tokio-postgres/src/connect_socket.rs
+++ b/tokio-postgres/src/connect_socket.rs
@@ -10,10 +10,12 @@ use tokio::net::UnixStream;
 use tokio::net::{self, TcpStream};
 use tokio::time;
 
+#[allow(unused_variables)]
 pub(crate) async fn connect_socket(
     host: &Host,
     port: u16,
     connect_timeout: Option<Duration>,
+    user_timeout: Option<Duration>,
     keepalive_config: Option<&KeepaliveConfig>,
 ) -> Result<Socket, Error> {
     match host {
@@ -35,8 +37,18 @@ pub(crate) async fn connect_socket(
                     };
 
                 stream.set_nodelay(true).map_err(Error::connect)?;
+
+                let sock_ref = SockRef::from(&stream);
+
+                #[cfg(target_os = "linux")]
+                {
+                    sock_ref
+                        .set_tcp_user_timeout(user_timeout)
+                        .map_err(Error::timeout)?;
+                }
+
                 if let Some(keepalive_config) = keepalive_config {
-                    SockRef::from(&stream)
+                    sock_ref
                         .set_tcp_keepalive(&TcpKeepalive::from(keepalive_config))
                         .map_err(Error::connect)?;
                 }
diff --git a/tokio-postgres/src/error/mod.rs b/tokio-postgres/src/error/mod.rs
index 47a31e793..34a1e1e3d 100644
--- a/tokio-postgres/src/error/mod.rs
+++ b/tokio-postgres/src/error/mod.rs
@@ -503,6 +503,11 @@ impl Error {
         Error::new(Kind::Connect, Some(Box::new(e)))
     }
 
+    #[allow(dead_code)]
+    pub(crate) fn timeout(e: io::Error) -> Error {
+        Error::new(Kind::Timeout, Some(Box::new(e)))
+    }
+
     #[doc(hidden)]
     pub fn __private_api_timeout() -> Error {
         Error::new(Kind::Timeout, None)