use crate::hnsw::DistanceMetric;
use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
#[derive(Debug, Clone)]
pub struct QueryMetrics {
pub duration: Duration,
pub result_count: usize,
pub cache_hit: bool,
pub metric: DistanceMetric,
pub ef_search: usize,
pub k: usize,
}
#[derive(Debug, Clone)]
pub struct AnalyticsSummary {
pub total_queries: usize,
pub cache_hits: usize,
pub cache_hit_rate: f32,
pub avg_duration: Duration,
pub p50_latency: Duration,
pub p90_latency: Duration,
pub p99_latency: Duration,
pub top_k_values: Vec<(usize, usize)>, pub qps: f32,
}
#[derive(Debug, Clone)]
pub struct DetectedPattern {
pub embedding_hash: u64,
pub frequency: usize,
pub avg_duration: Duration,
}
pub struct AnalyticsTracker {
query_history: Arc<RwLock<Vec<(Instant, QueryMetrics)>>>,
query_patterns: Arc<RwLock<HashMap<u64, DetectedPattern>>>,
max_history_size: usize,
start_time: Instant,
}
impl AnalyticsTracker {
pub fn new(max_history_size: usize) -> Self {
Self {
query_history: Arc::new(RwLock::new(Vec::new())),
query_patterns: Arc::new(RwLock::new(HashMap::new())),
max_history_size,
start_time: Instant::now(),
}
}
pub fn with_defaults() -> Self {
Self::new(10000) }
pub fn record_query(&self, embedding: &[f32], metrics: QueryMetrics) {
let now = Instant::now();
let hash = Self::hash_embedding(embedding);
{
let mut history = self.query_history.write();
history.push((now, metrics.clone()));
if history.len() > self.max_history_size {
let remove_count = history.len() - self.max_history_size;
history.drain(0..remove_count);
}
}
{
let mut patterns = self.query_patterns.write();
patterns
.entry(hash)
.and_modify(|pattern| {
pattern.frequency += 1;
let total = pattern.avg_duration.as_nanos() as f64
* (pattern.frequency - 1) as f64
+ metrics.duration.as_nanos() as f64;
pattern.avg_duration =
Duration::from_nanos((total / pattern.frequency as f64) as u64);
})
.or_insert(DetectedPattern {
embedding_hash: hash,
frequency: 1,
avg_duration: metrics.duration,
});
}
}
pub fn get_summary(&self, window: Option<Duration>) -> AnalyticsSummary {
let history = self.query_history.read();
let now = Instant::now();
let filtered: Vec<&QueryMetrics> = if let Some(duration) = window {
history
.iter()
.filter(|(timestamp, _)| now.duration_since(*timestamp) <= duration)
.map(|(_, metrics)| metrics)
.collect()
} else {
history.iter().map(|(_, metrics)| metrics).collect()
};
if filtered.is_empty() {
return AnalyticsSummary {
total_queries: 0,
cache_hits: 0,
cache_hit_rate: 0.0,
avg_duration: Duration::from_secs(0),
p50_latency: Duration::from_secs(0),
p90_latency: Duration::from_secs(0),
p99_latency: Duration::from_secs(0),
top_k_values: Vec::new(),
qps: 0.0,
};
}
let total_queries = filtered.len();
let cache_hits = filtered.iter().filter(|m| m.cache_hit).count();
let cache_hit_rate = cache_hits as f32 / total_queries as f32;
let total_duration: u128 = filtered.iter().map(|m| m.duration.as_nanos()).sum();
let avg_duration = Duration::from_nanos((total_duration / total_queries as u128) as u64);
let mut durations: Vec<Duration> = filtered.iter().map(|m| m.duration).collect();
durations.sort();
let p50_latency = durations[total_queries * 50 / 100];
let p90_latency = durations[total_queries * 90 / 100];
let p99_latency = durations[total_queries * 99 / 100];
let mut k_counts: HashMap<usize, usize> = HashMap::new();
for metrics in &filtered {
*k_counts.entry(metrics.k).or_insert(0) += 1;
}
let mut top_k_values: Vec<(usize, usize)> = k_counts.into_iter().collect();
top_k_values.sort_by_key(|a| std::cmp::Reverse(a.1)); top_k_values.truncate(5);
let elapsed = self.start_time.elapsed().as_secs_f32();
let qps = if elapsed > 0.0 {
total_queries as f32 / elapsed
} else {
0.0
};
AnalyticsSummary {
total_queries,
cache_hits,
cache_hit_rate,
avg_duration,
p50_latency,
p90_latency,
p99_latency,
top_k_values,
qps,
}
}
pub fn get_top_patterns(&self, limit: usize) -> Vec<DetectedPattern> {
let patterns = self.query_patterns.read();
let mut sorted: Vec<DetectedPattern> = patterns.values().cloned().collect();
sorted.sort_by_key(|a| std::cmp::Reverse(a.frequency));
sorted.truncate(limit);
sorted
}
pub fn clear(&self) {
self.query_history.write().clear();
self.query_patterns.write().clear();
}
pub fn total_queries(&self) -> usize {
self.query_history.read().len()
}
fn hash_embedding(embedding: &[f32]) -> u64 {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let mut hasher = DefaultHasher::new();
for (i, &val) in embedding.iter().enumerate().step_by(8) {
(i, (val * 1000.0) as i32).hash(&mut hasher);
}
hasher.finish()
}
}
pub struct QueryTimer {
start: Instant,
embedding: Vec<f32>,
k: usize,
ef_search: usize,
metric: DistanceMetric,
cache_hit: bool,
}
impl QueryTimer {
pub fn start(embedding: Vec<f32>, k: usize, ef_search: usize, metric: DistanceMetric) -> Self {
Self {
start: Instant::now(),
embedding,
k,
ef_search,
metric,
cache_hit: false,
}
}
pub fn set_cache_hit(&mut self, hit: bool) {
self.cache_hit = hit;
}
pub fn finish(self, tracker: &AnalyticsTracker, result_count: usize) {
let duration = self.start.elapsed();
let metrics = QueryMetrics {
duration,
result_count,
cache_hit: self.cache_hit,
metric: self.metric,
ef_search: self.ef_search,
k: self.k,
};
tracker.record_query(&self.embedding, metrics);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tracker_creation() {
let tracker = AnalyticsTracker::with_defaults();
assert_eq!(tracker.total_queries(), 0);
}
#[test]
fn test_record_query() {
let tracker = AnalyticsTracker::with_defaults();
let embedding = vec![0.5; 128];
let metrics = QueryMetrics {
duration: Duration::from_millis(10),
result_count: 5,
cache_hit: false,
metric: DistanceMetric::Cosine,
ef_search: 50,
k: 10,
};
tracker.record_query(&embedding, metrics);
assert_eq!(tracker.total_queries(), 1);
}
#[test]
fn test_analytics_summary() {
let tracker = AnalyticsTracker::with_defaults();
let embedding = vec![0.5; 128];
for i in 0..10 {
let metrics = QueryMetrics {
duration: Duration::from_millis(i * 10),
result_count: 5,
cache_hit: i % 2 == 0, metric: DistanceMetric::Cosine,
ef_search: 50,
k: 10,
};
tracker.record_query(&embedding, metrics);
}
let summary = tracker.get_summary(None);
assert_eq!(summary.total_queries, 10);
assert_eq!(summary.cache_hits, 5);
assert!((summary.cache_hit_rate - 0.5).abs() < 0.01);
}
#[test]
fn test_query_patterns() {
let tracker = AnalyticsTracker::with_defaults();
let embedding1 = vec![0.5; 128];
for _ in 0..5 {
let metrics = QueryMetrics {
duration: Duration::from_millis(10),
result_count: 5,
cache_hit: false,
metric: DistanceMetric::Cosine,
ef_search: 50,
k: 10,
};
tracker.record_query(&embedding1, metrics);
}
let embedding2 = vec![0.8; 128];
for _ in 0..3 {
let metrics = QueryMetrics {
duration: Duration::from_millis(20),
result_count: 5,
cache_hit: false,
metric: DistanceMetric::Cosine,
ef_search: 50,
k: 10,
};
tracker.record_query(&embedding2, metrics);
}
let patterns = tracker.get_top_patterns(2);
assert_eq!(patterns.len(), 2);
assert_eq!(patterns[0].frequency, 5); }
#[test]
fn test_query_timer() {
let tracker = AnalyticsTracker::with_defaults();
let embedding = vec![0.5; 128];
let timer = QueryTimer::start(embedding, 10, 50, DistanceMetric::Cosine);
std::thread::sleep(Duration::from_millis(10));
timer.finish(&tracker, 5);
assert_eq!(tracker.total_queries(), 1);
let summary = tracker.get_summary(None);
assert!(summary.avg_duration >= Duration::from_millis(10));
}
#[test]
fn test_top_k_values() {
let tracker = AnalyticsTracker::with_defaults();
let embedding = vec![0.5; 128];
for k in &[5, 10, 10, 10, 20] {
let metrics = QueryMetrics {
duration: Duration::from_millis(10),
result_count: 5,
cache_hit: false,
metric: DistanceMetric::Cosine,
ef_search: 50,
k: *k,
};
tracker.record_query(&embedding, metrics);
}
let summary = tracker.get_summary(None);
assert_eq!(summary.top_k_values[0].0, 10); assert_eq!(summary.top_k_values[0].1, 3); }
#[test]
fn test_clear_analytics() {
let tracker = AnalyticsTracker::with_defaults();
let embedding = vec![0.5; 128];
let metrics = QueryMetrics {
duration: Duration::from_millis(10),
result_count: 5,
cache_hit: false,
metric: DistanceMetric::Cosine,
ef_search: 50,
k: 10,
};
tracker.record_query(&embedding, metrics);
assert_eq!(tracker.total_queries(), 1);
tracker.clear();
assert_eq!(tracker.total_queries(), 0);
}
#[test]
fn test_time_window_filtering() {
let tracker = AnalyticsTracker::with_defaults();
let embedding = vec![0.5; 128];
let metrics = QueryMetrics {
duration: Duration::from_millis(10),
result_count: 5,
cache_hit: false,
metric: DistanceMetric::Cosine,
ef_search: 50,
k: 10,
};
tracker.record_query(&embedding, metrics);
let summary = tracker.get_summary(Some(Duration::from_secs(1)));
assert_eq!(summary.total_queries, 1);
std::thread::sleep(Duration::from_millis(100));
let summary = tracker.get_summary(Some(Duration::from_millis(10)));
assert_eq!(summary.total_queries, 0);
}
}