Skip to content

Commit

Permalink
Add support for tcp keepalive through socket2
Browse files Browse the repository at this point in the history
Not fully implemented, probably should be using a custom
structure for `TcpKeepalive` to avoid exposing `socket2` publicly.

In that case, should we have the same feature flag as the `socket2` library?
Or should we allow each setting and just warn when one is not available for the
given platform ?
  • Loading branch information
sosthene-nitrokey committed Feb 3, 2025
1 parent 414c515 commit d278529
Show file tree
Hide file tree
Showing 4 changed files with 50 additions and 0 deletions.
11 changes: 11 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ gzip = ["dep:flate2"]
brotli = ["dep:brotli-decompressor"]
charset = ["dep:encoding_rs"]
json = ["dep:serde", "dep:serde_json", "cookie_store?/serde_json"]
keepalive = ["dep:socket2", "socket2/all"]

######## UNSTABLE FEATURES.
# Might be removed or changed in a minor version.
Expand Down Expand Up @@ -90,6 +91,7 @@ encoding_rs = { version = "0.8.34", optional = true }

serde = { version = "1.0.138", optional = true, default-features = false, features = ["std"] }
serde_json = { version = "1.0.120", optional = true, default-features = false, features = ["std"] }
socket2 = { version = "0.5.8", optional = true }

[dev-dependencies]
env_logger = "0.11.6"
Expand Down
21 changes: 21 additions & 0 deletions src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,8 @@ pub struct Config {
max_idle_connections: usize,
max_idle_connections_per_host: usize,
max_idle_age: Duration,
#[cfg(feature = "keepalive")]
tcp_keepalive: Option<socket2::TcpKeepalive>,

// Chain built for middleware.
pub(crate) middleware: MiddlewareChain,
Expand Down Expand Up @@ -394,6 +396,12 @@ impl Config {
pub fn max_idle_age(&self) -> Duration {
self.max_idle_age
}

#[cfg(feature = "keepalive")]
/// TCP Keepalive configuration
pub fn tcp_keepalive(&self) -> &Option<socket2::TcpKeepalive> {
&self.tcp_keepalive
}
}

/// Builder of [`Config`]
Expand Down Expand Up @@ -712,6 +720,17 @@ impl<Scope: private::ConfigScope> ConfigBuilder<Scope> {
self.config().timeouts.recv_body = v;
self
}

/// Configuration of TCP keepalive
///
/// If `None`, keepalive is not configured for the connections
///
/// Defaults to `None`
#[cfg(feature = "keepalive")]
pub fn tcp_keepalive(mut self, v: Option<socket2::TcpKeepalive>) -> Self {
self.config().tcp_keepalive = v;
self
}
}

/// Possible config values for headers.
Expand Down Expand Up @@ -849,6 +868,8 @@ impl Default for Config {
max_idle_connections_per_host: 3,
max_idle_age: Duration::from_secs(15),
middleware: MiddlewareChain::default(),
#[cfg(feature = "keepalive")]
tcp_keepalive: None,
force_send_body: false,
}
}
Expand Down
16 changes: 16 additions & 0 deletions src/unversioned/transport/tcp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,8 @@ impl<In: Transport> Connector<In> for TcpConnector {

let config = &details.config;
let stream = try_connect(&details.addrs, details.timeout, config)?;
#[cfg(feature = "keepalive")]
let stream = set_keepalive(stream, details.config.tcp_keepalive().as_ref())?;

let buffers = LazyBuffers::new(config.input_buffer_size(), config.output_buffer_size());
let transport = TcpTransport::new(stream, buffers);
Expand All @@ -41,6 +43,20 @@ impl<In: Transport> Connector<In> for TcpConnector {
}
}

#[cfg(feature = "keepalive")]
fn set_keepalive(
stream: TcpStream,
keepalive: Option<&socket2::TcpKeepalive>,
) -> Result<TcpStream, crate::Error> {
use socket2::Socket;
let socket: Socket = stream.into();
if let Some(keepalive) = keepalive {
socket.set_keepalive(true)?;
socket.set_tcp_keepalive(keepalive)?;
}
Ok(socket.into())
}

fn try_connect(
addrs: &ResolvedSocketAddrs,
timeout: NextTimeout,
Expand Down

0 comments on commit d278529

Please sign in to comment.