use std::sync::Arc;
use std::time::Duration;
use crate::error::{KrafkaError, Result};
use super::connection::{ConnectionConfig, ConnectionConfigBuilder};
use super::pool::{ConnectionPool, DEFAULT_MAX_IDLE};
#[derive(Debug, Clone)]
pub struct TransportConfig {
pub(crate) tcp_nodelay: bool,
pub(crate) tcp_keepalive: Option<Duration>,
pub(crate) max_response_size: usize,
pub(crate) max_in_flight_requests: usize,
pub(crate) high_priority_channel_capacity: usize,
pub(crate) normal_priority_channel_capacity: usize,
pub(crate) max_high_priority_bypasses_per_round: usize,
pub(crate) connection_attempt_delay: Duration,
pub(crate) connections_max_idle: Option<Duration>,
pub(crate) max_connections: Option<usize>,
#[cfg(feature = "socks5")]
#[cfg_attr(docsrs, doc(cfg(feature = "socks5")))]
pub(crate) proxy: Option<crate::network::ProxyConfig>,
pub(crate) socket_send_buffer: Option<usize>,
pub(crate) socket_receive_buffer: Option<usize>,
pub(crate) tls_reload_interval: Option<Duration>,
}
impl Default for TransportConfig {
fn default() -> Self {
Self {
tcp_nodelay: true,
tcp_keepalive: Some(Duration::from_secs(60)),
max_response_size: crate::protocol::MAX_MESSAGE_SIZE,
max_in_flight_requests: 10,
high_priority_channel_capacity: 64,
normal_priority_channel_capacity: 256,
max_high_priority_bypasses_per_round: 4,
connection_attempt_delay: Duration::from_millis(250),
connections_max_idle: Some(DEFAULT_MAX_IDLE),
max_connections: None,
#[cfg(feature = "socks5")]
proxy: None,
socket_send_buffer: None,
socket_receive_buffer: None,
tls_reload_interval: None,
}
}
}
impl TransportConfig {
pub fn builder() -> TransportConfigBuilder {
TransportConfigBuilder(Self::default())
}
#[inline]
#[must_use]
pub fn tcp_nodelay(&self) -> bool {
self.tcp_nodelay
}
#[inline]
#[must_use]
pub fn tcp_keepalive(&self) -> Option<Duration> {
self.tcp_keepalive
}
#[inline]
#[must_use]
pub fn max_response_size(&self) -> usize {
self.max_response_size
}
#[inline]
#[must_use]
pub fn max_in_flight_requests(&self) -> usize {
self.max_in_flight_requests
}
#[cfg(feature = "socks5")]
#[cfg_attr(docsrs, doc(cfg(feature = "socks5")))]
#[inline]
#[must_use]
pub fn proxy(&self) -> Option<&crate::network::ProxyConfig> {
self.proxy.as_ref()
}
#[inline]
#[must_use]
pub fn socket_send_buffer(&self) -> Option<usize> {
self.socket_send_buffer
}
#[inline]
#[must_use]
pub fn socket_receive_buffer(&self) -> Option<usize> {
self.socket_receive_buffer
}
#[inline]
#[must_use]
pub fn high_priority_channel_capacity(&self) -> usize {
self.high_priority_channel_capacity
}
#[inline]
#[must_use]
pub fn normal_priority_channel_capacity(&self) -> usize {
self.normal_priority_channel_capacity
}
#[inline]
#[must_use]
pub fn max_high_priority_bypasses_per_round(&self) -> usize {
self.max_high_priority_bypasses_per_round
}
#[inline]
#[must_use]
pub fn connection_attempt_delay(&self) -> Duration {
self.connection_attempt_delay
}
#[inline]
#[must_use]
pub fn connections_max_idle(&self) -> Option<Duration> {
self.connections_max_idle
}
#[inline]
#[must_use]
pub fn max_connections(&self) -> Option<usize> {
self.max_connections
}
#[inline]
#[must_use]
pub fn tls_reload_interval(&self) -> Option<Duration> {
self.tls_reload_interval
}
pub(crate) fn apply(&self, builder: ConnectionConfigBuilder) -> ConnectionConfigBuilder {
let builder = builder
.nodelay(self.tcp_nodelay)
.tcp_keepalive(self.tcp_keepalive)
.max_response_size(self.max_response_size)
.max_in_flight_requests(self.max_in_flight_requests)
.high_priority_channel_capacity(self.high_priority_channel_capacity)
.normal_priority_channel_capacity(self.normal_priority_channel_capacity)
.max_high_priority_bypasses_per_round(self.max_high_priority_bypasses_per_round)
.connection_attempt_delay(self.connection_attempt_delay)
.socket_send_buffer(self.socket_send_buffer)
.socket_receive_buffer(self.socket_receive_buffer);
#[cfg(feature = "socks5")]
let builder = match self.proxy.clone() {
Some(proxy) => builder.proxy(proxy),
None => builder,
};
builder
}
pub(crate) fn build_pool(&self, config: ConnectionConfig) -> Arc<ConnectionPool> {
let pool = Arc::new(
ConnectionPool::new(config)
.with_max_idle(self.connections_max_idle)
.with_max_total_connections(self.max_connections),
);
pool.start_idle_evictor();
if let Some(interval) = self.tls_reload_interval {
pool.start_tls_reload(interval);
}
pool
}
}
#[must_use = "builders do nothing until .build() is called"]
#[derive(Debug)]
pub struct TransportConfigBuilder(TransportConfig);
impl Default for TransportConfigBuilder {
fn default() -> Self {
TransportConfig::builder()
}
}
impl TransportConfigBuilder {
pub fn tcp_nodelay(mut self, enabled: bool) -> Self {
self.0.tcp_nodelay = enabled;
self
}
pub fn tcp_keepalive(mut self, interval: Option<Duration>) -> Self {
self.0.tcp_keepalive = interval;
self
}
pub fn max_response_size(mut self, bytes: usize) -> Self {
self.0.max_response_size = bytes;
self
}
pub fn max_in_flight_requests(mut self, max: usize) -> Self {
self.0.max_in_flight_requests = max;
self
}
#[cfg(feature = "socks5")]
#[cfg_attr(docsrs, doc(cfg(feature = "socks5")))]
pub fn proxy(mut self, proxy: crate::network::ProxyConfig) -> Self {
self.0.proxy = Some(proxy);
self
}
pub fn socket_send_buffer(mut self, bytes: Option<usize>) -> Self {
self.0.socket_send_buffer = bytes;
self
}
pub fn socket_receive_buffer(mut self, bytes: Option<usize>) -> Self {
self.0.socket_receive_buffer = bytes;
self
}
pub fn high_priority_channel_capacity(mut self, capacity: usize) -> Self {
self.0.high_priority_channel_capacity = capacity;
self
}
pub fn normal_priority_channel_capacity(mut self, capacity: usize) -> Self {
self.0.normal_priority_channel_capacity = capacity;
self
}
pub fn max_high_priority_bypasses_per_round(mut self, n: usize) -> Self {
self.0.max_high_priority_bypasses_per_round = n;
self
}
pub fn connection_attempt_delay(mut self, delay: Duration) -> Self {
self.0.connection_attempt_delay = delay;
self
}
pub fn connections_max_idle(mut self, max_idle: Option<Duration>) -> Self {
self.0.connections_max_idle = max_idle;
self
}
pub fn max_connections(mut self, limit: Option<usize>) -> Self {
self.0.max_connections = limit;
self
}
pub fn tls_reload_interval(mut self, interval: Option<Duration>) -> Self {
self.0.tls_reload_interval = interval;
self
}
pub fn build(self) -> Result<TransportConfig> {
if self.0.max_in_flight_requests == 0 {
return Err(KrafkaError::config(
"max_in_flight_requests must be >= 1; 0 would block every request forever",
));
}
const MIN_RESPONSE_SIZE: usize = 1024;
if self.0.max_response_size < MIN_RESPONSE_SIZE {
return Err(KrafkaError::config(format!(
"max_response_size is {} B; the minimum is {MIN_RESPONSE_SIZE} B",
self.0.max_response_size
)));
}
if self.0.connections_max_idle == Some(Duration::ZERO) {
return Err(KrafkaError::config(
"connections_max_idle must be > 0; pass None to disable idle eviction",
));
}
if self.0.tls_reload_interval == Some(Duration::ZERO) {
return Err(KrafkaError::config(
"tls_reload_interval must be > 0; pass None to disable automatic TLS reloading",
));
}
if self.0.max_connections == Some(0) {
return Err(KrafkaError::config(
"max_connections must be >= 1; pass None for unlimited",
));
}
Ok(self.0)
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
mod tests {
use super::*;
#[test]
fn defaults_match_historical_connection_config() {
let transport = TransportConfig::default();
let legacy = ConnectionConfig::default();
assert_eq!(transport.tcp_nodelay, legacy.nodelay());
assert_eq!(transport.max_response_size, legacy.max_response_size());
assert_eq!(
transport.max_in_flight_requests,
legacy.max_in_flight_requests()
);
assert_eq!(
transport.high_priority_channel_capacity,
legacy.high_priority_channel_capacity()
);
assert_eq!(
transport.normal_priority_channel_capacity,
legacy.normal_priority_channel_capacity()
);
assert_eq!(
transport.connection_attempt_delay,
legacy.connection_attempt_delay()
);
assert_eq!(transport.connections_max_idle, Some(DEFAULT_MAX_IDLE));
assert_eq!(transport.max_connections, None);
assert_eq!(transport.tls_reload_interval, None);
}
#[test]
fn socket_buffer_sizes_reach_the_connection_config() {
let transport = TransportConfig::builder()
.socket_send_buffer(Some(4 * 1024 * 1024))
.socket_receive_buffer(Some(2 * 1024 * 1024))
.build()
.expect("socket buffer sizes are always valid");
assert_eq!(transport.socket_send_buffer(), Some(4 * 1024 * 1024));
assert_eq!(transport.socket_receive_buffer(), Some(2 * 1024 * 1024));
let applied = transport
.apply(ConnectionConfig::builder())
.build()
.expect("a valid transport config yields a valid connection config");
assert_eq!(
applied.send_buffer_size(),
Some(4 * 1024 * 1024),
"SO_SNDBUF must reach the socket, not stop at the transport config"
);
assert_eq!(applied.recv_buffer_size(), Some(2 * 1024 * 1024));
}
#[cfg(feature = "socks5")]
#[test]
fn a_proxy_on_the_transport_config_reaches_the_connection() {
use crate::network::ProxyConfig;
let transport = TransportConfig::builder()
.proxy(ProxyConfig::new("bastion:1080"))
.build()
.expect("a proxy address is not validated at build time");
assert_eq!(transport.proxy().map(|p| p.address()), Some("bastion:1080"));
let applied = transport
.apply(ConnectionConfig::builder())
.build()
.expect("a valid transport config yields a valid connection config");
assert_eq!(
applied.proxy().map(|p| p.address()),
Some("bastion:1080"),
"the SOCKS5 route must survive TransportConfig -> ConnectionConfig"
);
}
#[test]
fn socket_buffer_sizes_default_to_the_os() {
let applied = TransportConfig::default()
.apply(ConnectionConfig::builder())
.build()
.unwrap();
assert_eq!(applied.send_buffer_size(), None);
assert_eq!(applied.recv_buffer_size(), None);
}
#[test]
fn applying_defaults_changes_nothing() {
let applied = TransportConfig::default()
.apply(ConnectionConfig::builder())
.build()
.unwrap();
let plain = ConnectionConfig::default();
assert_eq!(applied.nodelay(), plain.nodelay());
assert_eq!(applied.max_response_size(), plain.max_response_size());
assert_eq!(
applied.max_in_flight_requests(),
plain.max_in_flight_requests()
);
assert_eq!(
applied.high_priority_channel_capacity(),
plain.high_priority_channel_capacity()
);
assert_eq!(
applied.normal_priority_channel_capacity(),
plain.normal_priority_channel_capacity()
);
assert_eq!(
applied.connection_attempt_delay(),
plain.connection_attempt_delay()
);
}
#[test]
fn every_connection_field_reaches_the_connection_config() {
let transport = TransportConfig::builder()
.tcp_nodelay(false)
.tcp_keepalive(Some(Duration::from_secs(17)))
.max_response_size(7 * 1024 * 1024)
.max_in_flight_requests(3)
.high_priority_channel_capacity(128)
.normal_priority_channel_capacity(512)
.max_high_priority_bypasses_per_round(9)
.connection_attempt_delay(Duration::from_millis(400))
.socket_send_buffer(Some(1024 * 1024))
.socket_receive_buffer(Some(512 * 1024))
.build()
.unwrap();
let config = transport
.apply(ConnectionConfig::builder())
.build()
.unwrap();
assert!(!config.nodelay());
assert_eq!(config.max_response_size(), 7 * 1024 * 1024);
assert_eq!(config.max_in_flight_requests(), 3);
assert_eq!(config.high_priority_channel_capacity(), 128);
assert_eq!(config.normal_priority_channel_capacity(), 512);
assert_eq!(
config.connection_attempt_delay(),
Duration::from_millis(400)
);
assert_eq!(config.send_buffer_size(), Some(1024 * 1024));
assert_eq!(config.recv_buffer_size(), Some(512 * 1024));
}
#[test]
fn rejects_zero_in_flight() {
let err = TransportConfig::builder()
.max_in_flight_requests(0)
.build()
.unwrap_err()
.to_string();
assert!(err.contains("max_in_flight_requests"), "got: {err}");
}
#[test]
fn rejects_tiny_response_ceiling() {
let err = TransportConfig::builder()
.max_response_size(64)
.build()
.unwrap_err()
.to_string();
assert!(err.contains("max_response_size"), "got: {err}");
}
#[test]
fn rejects_zero_intervals() {
assert!(
TransportConfig::builder()
.connections_max_idle(Some(Duration::ZERO))
.build()
.is_err()
);
assert!(
TransportConfig::builder()
.tls_reload_interval(Some(Duration::ZERO))
.build()
.is_err()
);
assert!(
TransportConfig::builder()
.max_connections(Some(0))
.build()
.is_err()
);
}
#[test]
fn none_disables_rather_than_erroring() {
let config = TransportConfig::builder()
.connections_max_idle(None)
.tls_reload_interval(None)
.max_connections(None)
.tcp_keepalive(None)
.build()
.unwrap();
assert_eq!(config.connections_max_idle(), None);
assert_eq!(config.tls_reload_interval(), None);
assert_eq!(config.max_connections(), None);
assert_eq!(config.tcp_keepalive(), None);
}
#[tokio::test]
async fn pool_level_fields_reach_the_pool() {
let transport = TransportConfig::builder()
.connections_max_idle(Some(Duration::from_secs(120)))
.max_connections(Some(42))
.build()
.unwrap();
let pool = transport.build_pool(ConnectionConfig::default());
assert_eq!(pool.max_idle(), Some(Duration::from_secs(120)));
assert_eq!(pool.max_total_connections(), Some(42));
pool.close_all().await;
}
}