Skip to content

Commit

Permalink
add tcp user timeout config
Browse files Browse the repository at this point in the history
  • Loading branch information
pH14 committed Feb 22, 2023
1 parent 9c90862 commit 18ee7eb
Show file tree
Hide file tree
Showing 7 changed files with 37 additions and 2 deletions.
2 changes: 1 addition & 1 deletion tokio-postgres/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"] }

Expand Down
1 change: 1 addition & 0 deletions tokio-postgres/src/cancel_query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ where
&config.host,
config.port,
config.connect_timeout,
config.user_timeout,
config.keepalive.as_ref(),
)
.await?;
Expand Down
1 change: 1 addition & 0 deletions tokio-postgres/src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
}

Expand Down
14 changes: 14 additions & 0 deletions tokio-postgres/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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`.
Expand Down
2 changes: 2 additions & 0 deletions tokio-postgres/src/connect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ where
host,
port,
config.connect_timeout,
config.user_timeout,
if config.keepalives {
Some(&config.keepalive_config)
} else {
Expand Down Expand Up @@ -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 {
Expand Down
14 changes: 13 additions & 1 deletion tokio-postgres/src/connect_socket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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)?;
}
Expand Down
5 changes: 5 additions & 0 deletions tokio-postgres/src/error/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 18ee7eb

Please sign in to comment.