use std::collections::HashMap;
use std::sync::RwLock;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
#[derive(Debug, Default)]
pub struct ServerMetrics {
total_requests: AtomicU64,
successful_requests: AtomicU64,
failed_requests: AtomicU64,
total_latency_us: AtomicU64,
method_counts: RwLock<HashMap<String, AtomicU64>>,
method_errors: RwLock<HashMap<String, AtomicU64>>,
method_latency_us: RwLock<HashMap<String, AtomicU64>>,
}
impl ServerMetrics {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn record_request(&self, method: &str, duration: Duration, success: bool) {
let latency_us = duration.as_micros() as u64;
self.total_requests.fetch_add(1, Ordering::Relaxed);
self.total_latency_us
.fetch_add(latency_us, Ordering::Relaxed);
if success {
self.successful_requests.fetch_add(1, Ordering::Relaxed);
} else {
self.failed_requests.fetch_add(1, Ordering::Relaxed);
}
self.increment_method_counter(&self.method_counts, method);
self.add_method_latency(method, latency_us);
if !success {
self.increment_method_counter(&self.method_errors, method);
}
}
pub fn record_success(&self, method: &str, duration: Duration) {
self.record_request(method, duration, true);
}
pub fn record_failure(&self, method: &str, duration: Duration) {
self.record_request(method, duration, false);
}
#[must_use]
pub fn snapshot(&self) -> MetricsSnapshot {
let method_counts = self
.method_counts
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let method_errors = self
.method_errors
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let method_latency = self
.method_latency_us
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let per_method: HashMap<String, MethodStats> = method_counts
.iter()
.map(|(method, count)| {
let requests = count.load(Ordering::Relaxed);
let errors = method_errors
.get(method)
.map_or(0, |c| c.load(Ordering::Relaxed));
let latency_us = method_latency
.get(method)
.map_or(0, |c| c.load(Ordering::Relaxed));
(
method.clone(),
MethodStats {
requests,
errors,
avg_latency_ms: if requests > 0 {
(latency_us as f64 / requests as f64) / 1000.0
} else {
0.0
},
},
)
})
.collect();
let total = self.total_requests.load(Ordering::Relaxed);
let total_latency = self.total_latency_us.load(Ordering::Relaxed);
MetricsSnapshot {
total_requests: total,
successful_requests: self.successful_requests.load(Ordering::Relaxed),
failed_requests: self.failed_requests.load(Ordering::Relaxed),
avg_latency_ms: if total > 0 {
(total_latency as f64 / total as f64) / 1000.0
} else {
0.0
},
per_method,
}
}
pub fn reset(&self) {
self.total_requests.store(0, Ordering::Relaxed);
self.successful_requests.store(0, Ordering::Relaxed);
self.failed_requests.store(0, Ordering::Relaxed);
self.total_latency_us.store(0, Ordering::Relaxed);
if let Ok(mut counts) = self.method_counts.write() {
counts.clear();
}
if let Ok(mut errors) = self.method_errors.write() {
errors.clear();
}
if let Ok(mut latency) = self.method_latency_us.write() {
latency.clear();
}
}
fn increment_method_counter(&self, map: &RwLock<HashMap<String, AtomicU64>>, method: &str) {
if let Ok(counts) = map.read() {
if let Some(counter) = counts.get(method) {
counter.fetch_add(1, Ordering::Relaxed);
return;
}
}
if let Ok(mut counts) = map.write() {
counts
.entry(method.to_string())
.or_insert_with(|| AtomicU64::new(0))
.fetch_add(1, Ordering::Relaxed);
}
}
fn add_method_latency(&self, method: &str, latency_us: u64) {
if let Ok(latencies) = self.method_latency_us.read() {
if let Some(counter) = latencies.get(method) {
counter.fetch_add(latency_us, Ordering::Relaxed);
return;
}
}
if let Ok(mut latencies) = self.method_latency_us.write() {
latencies
.entry(method.to_string())
.or_insert_with(|| AtomicU64::new(0))
.fetch_add(latency_us, Ordering::Relaxed);
}
}
}
#[derive(Debug, Clone)]
pub struct MetricsSnapshot {
pub total_requests: u64,
pub successful_requests: u64,
pub failed_requests: u64,
pub avg_latency_ms: f64,
pub per_method: HashMap<String, MethodStats>,
}
impl MetricsSnapshot {
#[must_use]
pub fn error_rate(&self) -> f64 {
if self.total_requests == 0 {
0.0
} else {
self.failed_requests as f64 / self.total_requests as f64
}
}
#[must_use]
pub fn success_rate(&self) -> f64 {
1.0 - self.error_rate()
}
#[must_use]
pub fn method(&self, name: &str) -> Option<&MethodStats> {
self.per_method.get(name)
}
#[must_use]
pub fn top_methods(&self, limit: usize) -> Vec<(&String, &MethodStats)> {
let mut methods: Vec<_> = self.per_method.iter().collect();
methods.sort_by_key(|m| std::cmp::Reverse(m.1.requests));
methods.into_iter().take(limit).collect()
}
#[must_use]
pub fn most_errors(&self, limit: usize) -> Vec<(&String, &MethodStats)> {
let mut methods: Vec<_> = self
.per_method
.iter()
.filter(|(_, s)| s.errors > 0)
.collect();
methods.sort_by(|a, b| b.1.error_rate().total_cmp(&a.1.error_rate()));
methods.into_iter().take(limit).collect()
}
#[must_use]
pub fn slowest_methods(&self, limit: usize) -> Vec<(&String, &MethodStats)> {
let mut methods: Vec<_> = self.per_method.iter().collect();
methods.sort_by(|a, b| b.1.avg_latency_ms.total_cmp(&a.1.avg_latency_ms));
methods.into_iter().take(limit).collect()
}
}
#[derive(Debug, Clone)]
pub struct MethodStats {
pub requests: u64,
pub errors: u64,
pub avg_latency_ms: f64,
}
impl MethodStats {
#[must_use]
pub fn error_rate(&self) -> f64 {
if self.requests == 0 {
0.0
} else {
self.errors as f64 / self.requests as f64
}
}
#[must_use]
pub fn success_rate(&self) -> f64 {
1.0 - self.error_rate()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_basic_metrics() {
let metrics = ServerMetrics::new();
metrics.record_success("tools/call", Duration::from_millis(50));
metrics.record_success("tools/call", Duration::from_millis(100));
metrics.record_failure("tools/call", Duration::from_millis(25));
metrics.record_success("resources/read", Duration::from_millis(10));
let snapshot = metrics.snapshot();
assert_eq!(snapshot.total_requests, 4);
assert_eq!(snapshot.successful_requests, 3);
assert_eq!(snapshot.failed_requests, 1);
assert!((snapshot.error_rate() - 0.25).abs() < 0.01);
}
#[test]
fn test_per_method_stats() -> Result<(), Box<dyn std::error::Error>> {
let metrics = ServerMetrics::new();
metrics.record_success("tools/call", Duration::from_millis(100));
metrics.record_success("tools/call", Duration::from_millis(100));
metrics.record_failure("tools/call", Duration::from_millis(100));
let snapshot = metrics.snapshot();
let tools_stats = snapshot.method("tools/call").ok_or("Method not found")?;
assert_eq!(tools_stats.requests, 3);
assert_eq!(tools_stats.errors, 1);
assert!((tools_stats.avg_latency_ms - 100.0).abs() < 1.0);
Ok(())
}
#[test]
fn test_reset() {
let metrics = ServerMetrics::new();
metrics.record_success("test", Duration::from_millis(50));
assert_eq!(metrics.snapshot().total_requests, 1);
metrics.reset();
assert_eq!(metrics.snapshot().total_requests, 0);
}
#[test]
fn test_top_methods() {
let metrics = ServerMetrics::new();
for _ in 0..10 {
metrics.record_success("tools/call", Duration::from_millis(10));
}
for _ in 0..5 {
metrics.record_success("resources/read", Duration::from_millis(10));
}
for _ in 0..3 {
metrics.record_success("prompts/get", Duration::from_millis(10));
}
let snapshot = metrics.snapshot();
let top = snapshot.top_methods(2);
assert_eq!(top.len(), 2);
assert_eq!(top[0].0, "tools/call");
assert_eq!(top[1].0, "resources/read");
}
}