#[cfg(not(feature = "turmoil"))]
use std::time::Duration;
use std::{future::Future, io::Result};
#[cfg(not(feature = "turmoil"))]
use socket2::{SockRef, TcpKeepalive};
use tokio::io::{AsyncRead, AsyncWrite};
#[cfg(not(feature = "turmoil"))]
pub use tokio::net::{TcpListener, TcpStream};
#[cfg(feature = "turmoil")]
pub use turmoil::net::{TcpListener, TcpStream};
pub trait TcpConnector: Send + Sync {
type Stream: AsyncRead + AsyncWrite + Send + Unpin + 'static;
fn connect(&self, addr: &str) -> impl Future<Output = Result<Self::Stream>> + Send;
}
#[derive(Default, Clone, Debug)]
pub struct RealTcpConnector;
impl TcpConnector for RealTcpConnector {
type Stream = TcpStream;
fn connect(&self, addr: &str) -> impl Future<Output = Result<Self::Stream>> + Send {
TcpStream::connect(addr.to_string())
}
}
pub(crate) fn apply_socket_options(stream: &TcpStream) {
if let Err(e) = stream.set_nodelay(true) {
log::warn!("Failed to enable TCP_NODELAY: {e}");
}
apply_keepalive(stream);
}
#[cfg(not(feature = "turmoil"))]
const KEEPALIVE_TIME: Duration = Duration::from_secs(20);
#[cfg(not(feature = "turmoil"))]
const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(10);
#[cfg(not(feature = "turmoil"))]
const KEEPALIVE_RETRIES: u32 = 3;
#[cfg(all(not(feature = "turmoil"), target_os = "linux"))]
const UNACKED_DATA_TIMEOUT: Duration = Duration::from_mins(1);
#[cfg(not(feature = "turmoil"))]
fn apply_keepalive(stream: &TcpStream) {
let socket = SockRef::from(stream);
let keepalive = TcpKeepalive::new()
.with_time(KEEPALIVE_TIME)
.with_interval(KEEPALIVE_INTERVAL)
.with_retries(KEEPALIVE_RETRIES);
if let Err(e) = socket.set_tcp_keepalive(&keepalive) {
log::warn!("Failed to enable TCP keepalive: {e}");
}
#[cfg(target_os = "linux")]
if let Err(e) = socket.set_tcp_user_timeout(Some(UNACKED_DATA_TIMEOUT)) {
log::warn!("Failed to set TCP_USER_TIMEOUT: {e}");
}
}
#[cfg(feature = "turmoil")]
const fn apply_keepalive(_stream: &TcpStream) {}
#[cfg(all(test, not(feature = "turmoil")))]
mod tests {
use rstest::rstest;
use tokio::net::TcpListener;
use super::*;
#[rstest]
#[tokio::test]
async fn test_apply_socket_options_sets_nodelay_and_keepalive() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let accept = tokio::spawn(async move { listener.accept().await.unwrap() });
let stream = TcpStream::connect(addr).await.unwrap();
let _accepted = accept.await.unwrap();
apply_socket_options(&stream);
let socket = SockRef::from(&stream);
assert!(stream.nodelay().unwrap());
assert!(socket.keepalive().unwrap());
assert_eq!(socket.tcp_keepalive_time().unwrap(), KEEPALIVE_TIME);
#[cfg(target_os = "linux")]
assert_eq!(
socket.tcp_user_timeout().unwrap(),
Some(UNACKED_DATA_TIMEOUT)
);
}
#[cfg(target_os = "linux")]
#[rstest]
fn test_unacked_timeout_covers_keepalive_probe_budget() {
let probe_budget = KEEPALIVE_TIME + KEEPALIVE_INTERVAL * KEEPALIVE_RETRIES;
assert!(UNACKED_DATA_TIMEOUT >= probe_budget);
}
}