use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
#[derive(Debug, Clone)]
pub struct SessionConfig {
pub timeout: Duration,
pub max_bytes: u64,
pub max_concurrent_requests: usize,
pub request_timeout: Duration,
pub max_retries: usize,
pub backoff_base: Duration,
}
impl Default for SessionConfig {
fn default() -> Self {
Self {
timeout: Duration::from_secs(30),
max_bytes: 1024 * 1024 * 1024,
max_concurrent_requests: 16,
request_timeout: Duration::from_secs(10),
max_retries: 3,
backoff_base: Duration::from_millis(200),
}
}
}
impl SessionConfig {
pub fn high_throughput() -> Self {
Self {
timeout: Duration::from_secs(120),
max_bytes: 10 * 1024 * 1024 * 1024,
max_concurrent_requests: 64,
request_timeout: Duration::from_secs(30),
max_retries: 5,
backoff_base: Duration::from_millis(100),
}
}
pub fn low_latency() -> Self {
Self {
timeout: Duration::from_secs(5),
max_bytes: 100 * 1024 * 1024,
max_concurrent_requests: 8,
request_timeout: Duration::from_secs(2),
max_retries: 1,
backoff_base: Duration::from_millis(50),
}
}
pub fn backoff_for_retry(&self, retry: usize) -> Duration {
let exp = retry.min(6) as u32;
self.backoff_base.saturating_mul(2u32.pow(exp))
}
}
#[derive(Debug, Clone)]
pub struct SessionMetricsSnapshot {
pub session_id: String,
pub bytes_transferred: u64,
pub blocks_requested: u64,
pub blocks_received: u64,
pub blocks_failed: u64,
pub total_retries: u64,
pub elapsed_ms: u64,
pub throughput_bps: f64,
}
pub struct SessionMetrics {
pub session_id: String,
bytes_transferred: AtomicU64,
blocks_requested: AtomicU64,
blocks_received: AtomicU64,
blocks_failed: AtomicU64,
total_retries: AtomicU64,
started_at_ms: u64,
}
impl SessionMetrics {
pub fn new(session_id: impl Into<String>, started_at_ms: u64) -> Arc<Self> {
Arc::new(Self {
session_id: session_id.into(),
bytes_transferred: AtomicU64::new(0),
blocks_requested: AtomicU64::new(0),
blocks_received: AtomicU64::new(0),
blocks_failed: AtomicU64::new(0),
total_retries: AtomicU64::new(0),
started_at_ms,
})
}
pub fn record_bytes(&self, bytes: u64) {
self.bytes_transferred.fetch_add(bytes, Ordering::Relaxed);
}
pub fn record_block_requested(&self) {
self.blocks_requested.fetch_add(1, Ordering::Relaxed);
}
pub fn record_block_received(&self) {
self.blocks_received.fetch_add(1, Ordering::Relaxed);
}
pub fn record_block_failed(&self) {
self.blocks_failed.fetch_add(1, Ordering::Relaxed);
}
pub fn record_retry(&self) {
self.total_retries.fetch_add(1, Ordering::Relaxed);
}
pub fn snapshot(&self, now_ms: u64) -> SessionMetricsSnapshot {
let bytes_transferred = self.bytes_transferred.load(Ordering::Acquire);
let elapsed_ms = now_ms.saturating_sub(self.started_at_ms);
let elapsed_secs = elapsed_ms as f64 / 1_000.0;
let throughput_bps = if elapsed_secs > 0.0 {
bytes_transferred as f64 / elapsed_secs
} else {
0.0
};
SessionMetricsSnapshot {
session_id: self.session_id.clone(),
bytes_transferred,
blocks_requested: self.blocks_requested.load(Ordering::Acquire),
blocks_received: self.blocks_received.load(Ordering::Acquire),
blocks_failed: self.blocks_failed.load(Ordering::Acquire),
total_retries: self.total_retries.load(Ordering::Acquire),
elapsed_ms,
throughput_bps,
}
}
pub fn bytes_transferred(&self) -> u64 {
self.bytes_transferred.load(Ordering::Relaxed)
}
pub fn blocks_received(&self) -> u64 {
self.blocks_received.load(Ordering::Relaxed)
}
}
impl std::fmt::Debug for SessionMetrics {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SessionMetrics")
.field("session_id", &self.session_id)
.field(
"bytes_transferred",
&self.bytes_transferred.load(Ordering::Relaxed),
)
.field(
"blocks_requested",
&self.blocks_requested.load(Ordering::Relaxed),
)
.field(
"blocks_received",
&self.blocks_received.load(Ordering::Relaxed),
)
.field("blocks_failed", &self.blocks_failed.load(Ordering::Relaxed))
.field("total_retries", &self.total_retries.load(Ordering::Relaxed))
.field("started_at_ms", &self.started_at_ms)
.finish()
}
}
pub struct SessionMetricsStore {
sessions: RwLock<HashMap<String, Arc<SessionMetrics>>>,
max_sessions: usize,
completed: RwLock<Vec<SessionMetricsSnapshot>>,
}
impl SessionMetricsStore {
pub fn new(max_sessions: usize) -> Arc<Self> {
Arc::new(Self {
sessions: RwLock::new(HashMap::new()),
max_sessions,
completed: RwLock::new(Vec::new()),
})
}
pub fn create_session(
&self,
session_id: impl Into<String>,
started_at_ms: u64,
) -> Arc<SessionMetrics> {
let id: String = session_id.into();
let metrics = SessionMetrics::new(id.clone(), started_at_ms);
self.sessions.write().insert(id, Arc::clone(&metrics));
metrics
}
pub fn get_session(&self, session_id: &str) -> Option<Arc<SessionMetrics>> {
self.sessions.read().get(session_id).cloned()
}
pub fn complete_session(&self, session_id: &str, now_ms: u64) -> bool {
let removed = self.sessions.write().remove(session_id);
match removed {
None => false,
Some(metrics) => {
let snap = metrics.snapshot(now_ms);
let mut completed = self.completed.write();
completed.push(snap);
while completed.len() > self.max_sessions {
completed.remove(0);
}
true
}
}
}
pub fn active_count(&self) -> usize {
self.sessions.read().len()
}
pub fn completed_snapshots(&self) -> Vec<SessionMetricsSnapshot> {
self.completed.read().clone()
}
pub fn total_bytes_transferred(&self) -> u64 {
self.sessions
.read()
.values()
.map(|m| m.bytes_transferred())
.sum()
}
}
impl std::fmt::Debug for SessionMetricsStore {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SessionMetricsStore")
.field("active_count", &self.active_count())
.field("max_sessions", &self.max_sessions)
.field("completed_count", &self.completed.read().len())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_config_values() {
let cfg = SessionConfig::default();
assert_eq!(cfg.timeout, Duration::from_secs(30));
assert_eq!(cfg.max_bytes, 1024 * 1024 * 1024);
assert_eq!(cfg.max_concurrent_requests, 16);
assert_eq!(cfg.request_timeout, Duration::from_secs(10));
assert_eq!(cfg.max_retries, 3);
assert_eq!(cfg.backoff_base, Duration::from_millis(200));
}
#[test]
fn test_high_throughput_config() {
let cfg = SessionConfig::high_throughput();
assert_eq!(cfg.timeout, Duration::from_secs(120));
assert_eq!(cfg.max_bytes, 10 * 1024 * 1024 * 1024);
assert_eq!(cfg.max_concurrent_requests, 64);
assert_eq!(cfg.request_timeout, Duration::from_secs(30));
assert_eq!(cfg.max_retries, 5);
assert_eq!(cfg.backoff_base, Duration::from_millis(100));
}
#[test]
fn test_low_latency_config() {
let cfg = SessionConfig::low_latency();
assert_eq!(cfg.timeout, Duration::from_secs(5));
assert_eq!(cfg.max_bytes, 100 * 1024 * 1024);
assert_eq!(cfg.max_concurrent_requests, 8);
assert_eq!(cfg.request_timeout, Duration::from_secs(2));
assert_eq!(cfg.max_retries, 1);
assert_eq!(cfg.backoff_base, Duration::from_millis(50));
}
#[test]
fn test_backoff_exponential() {
let cfg = SessionConfig::default(); assert_eq!(cfg.backoff_for_retry(0), Duration::from_millis(200)); assert_eq!(cfg.backoff_for_retry(1), Duration::from_millis(400)); assert_eq!(cfg.backoff_for_retry(2), Duration::from_millis(800)); assert_eq!(cfg.backoff_for_retry(3), Duration::from_millis(1_600));
assert_eq!(cfg.backoff_for_retry(4), Duration::from_millis(3_200));
assert_eq!(cfg.backoff_for_retry(5), Duration::from_millis(6_400));
assert_eq!(cfg.backoff_for_retry(6), Duration::from_millis(12_800));
assert_eq!(cfg.backoff_for_retry(7), Duration::from_millis(12_800));
assert_eq!(cfg.backoff_for_retry(100), Duration::from_millis(12_800));
}
#[test]
fn test_session_metrics_record_bytes() {
let m = SessionMetrics::new("sess-bytes", 0);
assert_eq!(m.bytes_transferred(), 0);
m.record_bytes(512);
m.record_bytes(1024);
assert_eq!(m.bytes_transferred(), 1_536);
}
#[test]
fn test_session_metrics_record_blocks() {
let m = SessionMetrics::new("sess-blocks", 0);
m.record_block_requested();
m.record_block_requested();
m.record_block_requested();
m.record_block_received();
m.record_block_received();
m.record_block_failed();
m.record_retry();
m.record_retry();
let snap = m.snapshot(1_000);
assert_eq!(snap.blocks_requested, 3);
assert_eq!(snap.blocks_received, 2);
assert_eq!(snap.blocks_failed, 1);
assert_eq!(snap.total_retries, 2);
assert_eq!(m.blocks_received(), 2);
}
#[test]
fn test_session_metrics_snapshot_throughput() {
let m = SessionMetrics::new("sess-tput", 0);
let mib: u64 = 1024 * 1024;
m.record_bytes(mib);
let snap = m.snapshot(2_000);
assert_eq!(snap.elapsed_ms, 2_000);
assert_eq!(snap.bytes_transferred, mib);
let expected = mib as f64 / 2.0;
let diff = (snap.throughput_bps - expected).abs();
assert!(
diff < 1.0,
"throughput mismatch: got {}, expected {}",
snap.throughput_bps,
expected
);
}
#[test]
fn test_session_metrics_store_create_and_get() {
let store = SessionMetricsStore::new(10);
assert_eq!(store.active_count(), 0);
let m = store.create_session("alpha", 1_000);
m.record_bytes(256);
assert_eq!(store.active_count(), 1);
let fetched = store.get_session("alpha").expect("session should exist");
assert_eq!(fetched.bytes_transferred(), 256);
assert!(store.get_session("missing").is_none());
}
#[test]
fn test_session_metrics_store_complete() {
let store = SessionMetricsStore::new(10);
store.create_session("beta", 0).record_bytes(4_096);
assert_eq!(store.active_count(), 1);
assert!(store.completed_snapshots().is_empty());
let found = store.complete_session("beta", 1_000);
assert!(found, "complete_session should return true");
assert_eq!(store.active_count(), 0);
let snaps = store.completed_snapshots();
assert_eq!(snaps.len(), 1);
assert_eq!(snaps[0].session_id, "beta");
assert_eq!(snaps[0].bytes_transferred, 4_096);
assert!(!store.complete_session("ghost", 2_000));
}
#[test]
fn test_session_metrics_store_total_bytes() {
let store = SessionMetricsStore::new(10);
store.create_session("s1", 0).record_bytes(100);
store.create_session("s2", 0).record_bytes(200);
store.create_session("s3", 0).record_bytes(300);
assert_eq!(store.total_bytes_transferred(), 600);
store.complete_session("s2", 1_000);
assert_eq!(store.total_bytes_transferred(), 400);
}
#[test]
fn test_max_sessions_cap() {
let store = SessionMetricsStore::new(3);
for i in 0..6u64 {
let id = format!("sess-{}", i);
store.create_session(&id, i * 100).record_bytes(i * 10);
store.complete_session(&id, (i + 1) * 100);
}
assert_eq!(store.active_count(), 0);
let snaps = store.completed_snapshots();
assert_eq!(snaps.len(), 3);
let ids: Vec<&str> = snaps.iter().map(|s| s.session_id.as_str()).collect();
assert!(ids.contains(&"sess-3"));
assert!(ids.contains(&"sess-4"));
assert!(ids.contains(&"sess-5"));
let extra = store.create_session("extra-1", 9_000);
extra.record_bytes(999);
assert_eq!(store.active_count(), 1);
assert_eq!(extra.bytes_transferred(), 999);
}
}