use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
#[derive(Debug, Clone)]
pub struct SearchStats {
pub total_queries: u64,
pub qps: f64,
pub p50_latency: Duration,
pub p95_latency: Duration,
pub p99_latency: Duration,
pub avg_latency: Duration,
pub min_latency: Duration,
pub max_latency: Duration,
}
impl Default for SearchStats {
fn default() -> Self {
Self {
total_queries: 0,
qps: 0.0,
p50_latency: Duration::ZERO,
p95_latency: Duration::ZERO,
p99_latency: Duration::ZERO,
avg_latency: Duration::ZERO,
min_latency: Duration::MAX,
max_latency: Duration::ZERO,
}
}
}
#[derive(Debug, Clone)]
pub struct IndexStats {
pub num_vectors: usize,
pub dimensions: usize,
pub build_time: Duration,
pub memory_bytes: usize,
}
impl Default for IndexStats {
fn default() -> Self {
Self {
num_vectors: 0,
dimensions: 0,
build_time: Duration::ZERO,
memory_bytes: 0,
}
}
}
#[derive(Clone)]
pub struct Metrics {
search_metrics: Arc<Mutex<SearchMetrics>>,
index_stats: Arc<Mutex<IndexStats>>,
}
impl Metrics {
pub fn new() -> Self {
Self {
search_metrics: Arc::new(Mutex::new(SearchMetrics::new())),
index_stats: Arc::new(Mutex::new(IndexStats::default())),
}
}
pub fn record_search_latency(&self, latency: Duration) {
let mut metrics = self.search_metrics.lock().unwrap();
metrics.record_latency(latency);
}
pub fn get_search_stats(&self) -> SearchStats {
let metrics = self.search_metrics.lock().unwrap();
metrics.compute_stats()
}
pub fn set_index_stats(&self, stats: IndexStats) {
let mut index_stats = self.index_stats.lock().unwrap();
*index_stats = stats;
}
pub fn get_index_stats(&self) -> IndexStats {
let index_stats = self.index_stats.lock().unwrap();
index_stats.clone()
}
pub fn reset(&self) {
let mut search_metrics = self.search_metrics.lock().unwrap();
*search_metrics = SearchMetrics::new();
let mut index_stats = self.index_stats.lock().unwrap();
*index_stats = IndexStats::default();
}
}
impl Default for Metrics {
fn default() -> Self {
Self::new()
}
}
struct SearchMetrics {
latencies: Vec<Duration>,
start_time: Instant,
}
impl SearchMetrics {
fn new() -> Self {
Self {
latencies: Vec::new(),
start_time: Instant::now(),
}
}
fn record_latency(&mut self, latency: Duration) {
self.latencies.push(latency);
}
fn compute_stats(&self) -> SearchStats {
if self.latencies.is_empty() {
return SearchStats::default();
}
let mut sorted_latencies = self.latencies.clone();
sorted_latencies.sort();
let total_queries = sorted_latencies.len() as u64;
let p50_idx = ((total_queries as f64 * 0.50).ceil() as usize).saturating_sub(1);
let p95_idx = ((total_queries as f64 * 0.95).ceil() as usize).saturating_sub(1);
let p99_idx = ((total_queries as f64 * 0.99).ceil() as usize).saturating_sub(1);
let p50_latency = sorted_latencies
.get(p50_idx)
.copied()
.unwrap_or(Duration::ZERO);
let p95_latency = sorted_latencies
.get(p95_idx)
.copied()
.unwrap_or(Duration::ZERO);
let p99_latency = sorted_latencies
.get(p99_idx)
.copied()
.unwrap_or(Duration::ZERO);
let total_micros: u128 = sorted_latencies.iter().map(|d| d.as_micros()).sum();
let avg_micros = total_micros / total_queries as u128;
let avg_latency = Duration::from_micros(avg_micros as u64);
let min_latency = *sorted_latencies.first().unwrap();
let max_latency = *sorted_latencies.last().unwrap();
let elapsed = self.start_time.elapsed().as_secs_f64();
let qps = if elapsed > 0.0 {
total_queries as f64 / elapsed
} else {
0.0
};
SearchStats {
total_queries,
qps,
p50_latency,
p95_latency,
p99_latency,
avg_latency,
min_latency,
max_latency,
}
}
}
pub struct LatencyTimer {
start: Instant,
metrics: Option<Metrics>,
}
impl LatencyTimer {
pub fn new(metrics: Option<Metrics>) -> Self {
Self {
start: Instant::now(),
metrics,
}
}
pub fn finish(self) -> Duration {
let latency = self.start.elapsed();
if let Some(metrics) = self.metrics {
metrics.record_search_latency(latency);
}
latency
}
pub fn elapsed(&self) -> Duration {
self.start.elapsed()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::thread;
use std::time::Duration;
#[test]
fn test_metrics_creation() {
let metrics = Metrics::new();
let stats = metrics.get_search_stats();
assert_eq!(stats.total_queries, 0);
assert_eq!(stats.qps, 0.0);
}
#[test]
fn test_record_search_latency() {
let metrics = Metrics::new();
metrics.record_search_latency(Duration::from_micros(100));
metrics.record_search_latency(Duration::from_micros(200));
metrics.record_search_latency(Duration::from_micros(300));
let stats = metrics.get_search_stats();
assert_eq!(stats.total_queries, 3);
assert_eq!(stats.min_latency, Duration::from_micros(100));
assert_eq!(stats.max_latency, Duration::from_micros(300));
assert_eq!(stats.p50_latency, Duration::from_micros(200));
}
#[test]
fn test_percentiles() {
let metrics = Metrics::new();
for i in 1..=100 {
metrics.record_search_latency(Duration::from_micros(i));
}
let stats = metrics.get_search_stats();
assert_eq!(stats.total_queries, 100);
assert!(stats.p50_latency >= Duration::from_micros(49));
assert!(stats.p50_latency <= Duration::from_micros(51));
assert!(stats.p95_latency >= Duration::from_micros(94));
assert!(stats.p95_latency <= Duration::from_micros(96));
assert!(stats.p99_latency >= Duration::from_micros(98));
assert!(stats.p99_latency <= Duration::from_micros(100));
}
#[test]
fn test_average_latency() {
let metrics = Metrics::new();
metrics.record_search_latency(Duration::from_micros(100));
metrics.record_search_latency(Duration::from_micros(200));
metrics.record_search_latency(Duration::from_micros(300));
let stats = metrics.get_search_stats();
assert_eq!(stats.avg_latency, Duration::from_micros(200));
}
#[test]
fn test_qps_calculation() {
let metrics = Metrics::new();
for _ in 0..10 {
metrics.record_search_latency(Duration::from_micros(100));
}
thread::sleep(Duration::from_millis(100));
let stats = metrics.get_search_stats();
assert!(stats.qps > 0.0);
assert_eq!(stats.total_queries, 10);
}
#[test]
fn test_index_stats() {
let metrics = Metrics::new();
let index_stats = IndexStats {
num_vectors: 1000,
dimensions: 768,
build_time: Duration::from_secs(5),
memory_bytes: 1024 * 1024, };
metrics.set_index_stats(index_stats.clone());
let retrieved = metrics.get_index_stats();
assert_eq!(retrieved.num_vectors, 1000);
assert_eq!(retrieved.dimensions, 768);
assert_eq!(retrieved.build_time, Duration::from_secs(5));
assert_eq!(retrieved.memory_bytes, 1024 * 1024);
}
#[test]
fn test_metrics_reset() {
let metrics = Metrics::new();
metrics.record_search_latency(Duration::from_micros(100));
metrics.record_search_latency(Duration::from_micros(200));
let stats = metrics.get_search_stats();
assert_eq!(stats.total_queries, 2);
metrics.reset();
let stats = metrics.get_search_stats();
assert_eq!(stats.total_queries, 0);
}
#[test]
fn test_latency_timer() {
let metrics = Metrics::new();
let timer = LatencyTimer::new(Some(metrics.clone()));
thread::sleep(Duration::from_millis(10));
let latency = timer.finish();
assert!(latency >= Duration::from_millis(10));
let stats = metrics.get_search_stats();
assert_eq!(stats.total_queries, 1);
assert!(stats.min_latency >= Duration::from_millis(10));
}
#[test]
fn test_latency_timer_without_metrics() {
let timer = LatencyTimer::new(None);
thread::sleep(Duration::from_millis(5));
let latency = timer.finish();
assert!(latency >= Duration::from_millis(5));
}
#[test]
fn test_latency_timer_elapsed() {
let timer = LatencyTimer::new(None);
thread::sleep(Duration::from_millis(5));
let elapsed = timer.elapsed();
assert!(elapsed >= Duration::from_millis(5));
}
#[test]
fn test_thread_safety() {
let metrics = Metrics::new();
let metrics_clone = metrics.clone();
let handle = thread::spawn(move || {
for _ in 0..100 {
metrics_clone.record_search_latency(Duration::from_micros(100));
}
});
for _ in 0..100 {
metrics.record_search_latency(Duration::from_micros(200));
}
handle.join().unwrap();
let stats = metrics.get_search_stats();
assert_eq!(stats.total_queries, 200);
}
}