Skip to main content

codemem_mcp/
metrics.rs

1//! In-memory metrics collector for MCP tool calls.
2
3use codemem_core::{LatencyStats, Metrics, MetricsSnapshot};
4use std::collections::HashMap;
5use std::sync::Mutex;
6
7/// In-memory metrics collector.
8///
9/// Collects latency samples, counter increments, and gauge values.
10/// Thread-safe via internal `Mutex`.
11pub struct InMemoryMetrics {
12    inner: Mutex<Inner>,
13}
14
15struct Inner {
16    /// Raw latency samples per operation (kept for percentile calculation).
17    latency_samples: HashMap<String, Vec<f64>>,
18    /// Cumulative counters.
19    counters: HashMap<String, u64>,
20    /// Point-in-time gauges.
21    gauges: HashMap<String, f64>,
22}
23
24impl InMemoryMetrics {
25    pub fn new() -> Self {
26        Self {
27            inner: Mutex::new(Inner {
28                latency_samples: HashMap::new(),
29                counters: HashMap::new(),
30                gauges: HashMap::new(),
31            }),
32        }
33    }
34
35    /// Take a snapshot of all collected metrics.
36    pub fn snapshot(&self) -> MetricsSnapshot {
37        let inner = match self.inner.lock() {
38            Ok(guard) => guard,
39            Err(e) => {
40                tracing::warn!("Metrics lock poisoned: {e}");
41                return MetricsSnapshot::default();
42            }
43        };
44
45        let latencies: HashMap<String, LatencyStats> = inner
46            .latency_samples
47            .iter()
48            .map(|(name, samples)| {
49                let stats = compute_latency_stats(samples);
50                (name.clone(), stats)
51            })
52            .collect();
53
54        MetricsSnapshot {
55            latencies,
56            counters: inner.counters.clone(),
57            gauges: inner.gauges.clone(),
58        }
59    }
60}
61
62impl Default for InMemoryMetrics {
63    fn default() -> Self {
64        Self::new()
65    }
66}
67
68impl Metrics for InMemoryMetrics {
69    fn record_latency(&self, operation: &str, duration_ms: f64) {
70        if let Ok(mut inner) = self.inner.lock() {
71            inner
72                .latency_samples
73                .entry(operation.to_string())
74                .or_default()
75                .push(duration_ms);
76        }
77    }
78
79    fn increment_counter(&self, name: &str, delta: u64) {
80        if let Ok(mut inner) = self.inner.lock() {
81            *inner.counters.entry(name.to_string()).or_insert(0) += delta;
82        }
83    }
84
85    fn record_gauge(&self, name: &str, value: f64) {
86        if let Ok(mut inner) = self.inner.lock() {
87            inner.gauges.insert(name.to_string(), value);
88        }
89    }
90}
91
92/// Compute percentile-based statistics from a slice of samples.
93fn compute_latency_stats(samples: &[f64]) -> LatencyStats {
94    if samples.is_empty() {
95        return LatencyStats::default();
96    }
97
98    let mut sorted = samples.to_vec();
99    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
100
101    let count = sorted.len() as u64;
102    let total: f64 = sorted.iter().sum();
103    let min = sorted[0];
104    let max = sorted[sorted.len() - 1];
105
106    let p50 = percentile(&sorted, 50.0);
107    let p95 = percentile(&sorted, 95.0);
108    let p99 = percentile(&sorted, 99.0);
109
110    LatencyStats {
111        count,
112        total_ms: total,
113        min_ms: min,
114        max_ms: max,
115        p50_ms: p50,
116        p95_ms: p95,
117        p99_ms: p99,
118    }
119}
120
121fn percentile(sorted: &[f64], pct: f64) -> f64 {
122    if sorted.is_empty() {
123        return 0.0;
124    }
125    let idx = (pct / 100.0 * (sorted.len() as f64 - 1.0)).round() as usize;
126    sorted[idx.min(sorted.len() - 1)]
127}
128
129#[cfg(test)]
130#[path = "tests/metrics_tests.rs"]
131mod tests;