use std::time::{Duration, Instant};
use http::{HeaderMap, Version, header};
#[derive(Debug, Clone, Copy)]
pub struct KeepAliveConfig {
pub enabled: bool,
pub idle_timeout: Duration,
pub max_requests: Option<u32>,
pub tcp_keepalive: Option<Duration>,
}
impl Default for KeepAliveConfig {
fn default() -> Self {
Self {
enabled: true,
idle_timeout: Duration::from_secs(90),
max_requests: Some(100),
tcp_keepalive: Some(Duration::from_secs(60)),
}
}
}
impl KeepAliveConfig {
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn disabled() -> Self {
Self {
enabled: false,
..Default::default()
}
}
#[must_use]
pub fn with_idle_timeout(mut self, timeout: Duration) -> Self {
self.idle_timeout = timeout;
self
}
#[must_use]
pub fn with_max_requests(mut self, max: u32) -> Self {
self.max_requests = Some(max);
self
}
#[must_use]
pub fn unlimited_requests(mut self) -> Self {
self.max_requests = None;
self
}
#[must_use]
pub fn with_tcp_keepalive(mut self, interval: Duration) -> Self {
self.tcp_keepalive = Some(interval);
self
}
}
#[derive(Debug)]
pub struct ConnectionState {
pub created_at: Instant,
pub last_used: Instant,
pub request_count: u32,
pub in_use: bool,
config: KeepAliveConfig,
}
impl ConnectionState {
pub fn new(config: KeepAliveConfig) -> Self {
let now = Instant::now();
Self {
created_at: now,
last_used: now,
request_count: 0,
in_use: false,
config,
}
}
pub fn use_connection(&mut self) {
self.last_used = Instant::now();
self.request_count += 1;
self.in_use = true;
}
pub fn release(&mut self) {
self.last_used = Instant::now();
self.in_use = false;
}
pub fn should_keep_alive(&self) -> bool {
if !self.config.enabled {
return false;
}
if let Some(max) = self.config.max_requests {
if self.request_count >= max {
return false;
}
}
if self.last_used.elapsed() > self.config.idle_timeout {
return false;
}
true
}
pub fn is_expired(&self) -> bool {
self.last_used.elapsed() > self.config.idle_timeout
}
pub fn time_to_expiry(&self) -> Duration {
self.config
.idle_timeout
.saturating_sub(self.last_used.elapsed())
}
pub fn age(&self) -> Duration {
self.created_at.elapsed()
}
pub fn idle_time(&self) -> Duration {
self.last_used.elapsed()
}
}
#[derive(Debug, Clone)]
pub struct KeepAliveHints {
pub keep_alive: bool,
pub timeout: Option<Duration>,
pub max: Option<u32>,
}
impl KeepAliveHints {
pub fn from_headers(headers: &HeaderMap, version: Version) -> Self {
let mut keep_alive = version == Version::HTTP_11;
if let Some(conn) = headers.get(header::CONNECTION) {
if let Ok(value) = conn.to_str() {
let value_lower = value.to_lowercase();
if value_lower.contains("close") {
keep_alive = false;
} else if value_lower.contains("keep-alive") {
keep_alive = true;
}
}
}
let mut timeout = None;
let mut max = None;
if let Some(ka) = headers.get("keep-alive") {
if let Ok(value) = ka.to_str() {
for param in value.split(',') {
let param = param.trim();
if let Some(rest) = param.strip_prefix("timeout=") {
if let Ok(secs) = rest.trim().parse::<u64>() {
timeout = Some(Duration::from_secs(secs));
}
} else if let Some(rest) = param.strip_prefix("max=") {
if let Ok(n) = rest.trim().parse::<u32>() {
max = Some(n);
}
}
}
}
}
Self {
keep_alive,
timeout,
max,
}
}
pub fn should_close(&self) -> bool {
!self.keep_alive
}
}
pub fn connection_header_value(keep_alive: bool) -> &'static str {
if keep_alive { "keep-alive" } else { "close" }
}
pub fn keep_alive_header_value(config: &KeepAliveConfig) -> String {
let mut parts = vec![format!("timeout={}", config.idle_timeout.as_secs())];
if let Some(max) = config.max_requests {
parts.push(format!("max={}", max));
}
parts.join(", ")
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used, clippy::expect_used)]
use super::*;
use http::header::HeaderValue;
#[test]
fn test_keep_alive_config() {
let config = KeepAliveConfig::new();
assert!(config.enabled);
assert_eq!(config.idle_timeout, Duration::from_secs(90));
assert_eq!(config.max_requests, Some(100));
}
#[test]
fn test_connection_state() {
let config = KeepAliveConfig::new();
let mut state = ConnectionState::new(config);
assert!(!state.in_use);
assert_eq!(state.request_count, 0);
assert!(state.should_keep_alive());
state.use_connection();
assert!(state.in_use);
assert_eq!(state.request_count, 1);
state.release();
assert!(!state.in_use);
}
#[test]
fn test_keep_alive_hints_http11() {
let mut headers = HeaderMap::new();
headers.insert(header::CONNECTION, HeaderValue::from_static("keep-alive"));
headers.insert("keep-alive", HeaderValue::from_static("timeout=30, max=50"));
let hints = KeepAliveHints::from_headers(&headers, Version::HTTP_11);
assert!(hints.keep_alive);
assert_eq!(hints.timeout, Some(Duration::from_secs(30)));
assert_eq!(hints.max, Some(50));
}
#[test]
fn test_keep_alive_hints_close() {
let mut headers = HeaderMap::new();
headers.insert(header::CONNECTION, HeaderValue::from_static("close"));
let hints = KeepAliveHints::from_headers(&headers, Version::HTTP_11);
assert!(!hints.keep_alive);
assert!(hints.should_close());
}
#[test]
fn test_connection_header_value() {
assert_eq!(connection_header_value(true), "keep-alive");
assert_eq!(connection_header_value(false), "close");
}
}