searchi/crates/searchi-core/src/metrics.rs

//! Monotonic request metrics. Stage names are static labels, never queries/paths.
use serde::{Deserialize, Serialize};
use std::{
    collections::BTreeMap,
    sync::atomic::{AtomicU64, Ordering},
    time::Instant,
};
const MAX_SPANS: usize = 512;
static COLLECTOR_ID: AtomicU64 = AtomicU64::new(1);

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchResponse<T> {
    pub results: Vec<T>,
    pub metrics: SearchMetrics,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchMetrics {
    pub total_ms: f64,
    pub result_count: usize,
    pub trace: Option<Vec<SearchSpan>>,
    #[serde(default)]
    pub dropped_spans: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchSpan {
    pub name: String,
    pub start_ms: f64,
    pub duration_ms: f64,
    pub parent: Option<usize>,
    pub counters: BTreeMap<String, u64>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MetricsMode {
    Basic,
    Debug,
}
#[derive(Debug, Clone, Copy)]
pub struct SpanToken {
    collector: u64,
    index: usize,
}
impl SpanToken {
    pub fn index(self) -> usize {
        self.index
    }
}
struct OpenSpan {
    value: SearchSpan,
    started: Instant,
    ended: bool,
}
struct DebugTrace {
    id: u64,
    spans: Vec<OpenSpan>,
    dropped: u64,
}
/// Basic mode stores only one Instant: no span Vec or stage/counter allocation.
pub struct TraceCollector {
    started: Instant,
    debug: Option<DebugTrace>,
}
impl TraceCollector {
    pub fn new(mode: MetricsMode) -> Self {
        Self {
            started: Instant::now(),
            debug: match mode {
                MetricsMode::Basic => None,
                MetricsMode::Debug => Some(DebugTrace {
                    id: COLLECTOR_ID.fetch_add(1, Ordering::Relaxed),
                    spans: Vec::new(),
                    dropped: 0,
                }),
            },
        }
    }
    pub fn elapsed_ms(&self) -> f64 {
        self.started.elapsed().as_secs_f64() * 1000.0
    }
    pub fn is_debug(&self) -> bool {
        self.debug.is_some()
    }
    pub fn begin(&mut self, name: &'static str, parent: Option<SpanToken>) -> Option<SpanToken> {
        let debug = self.debug.as_mut()?;
        let parent_index = match parent {
            Some(token)
                if token.collector == debug.id
                    && debug.spans.get(token.index).is_some_and(|span| !span.ended) =>
            {
                Some(token.index)
            }
            Some(_) => {
                debug.dropped = debug.dropped.saturating_add(1);
                return None;
            }
            None => None,
        };
        if debug.spans.len() >= MAX_SPANS {
            debug.dropped = debug.dropped.saturating_add(1);
            return None;
        }
        let started = Instant::now();
        let index = debug.spans.len();
        debug.spans.push(OpenSpan {
            value: SearchSpan {
                name: name.chars().take(96).collect(),
                start_ms: started.duration_since(self.started).as_secs_f64() * 1000.0,
                duration_ms: 0.0,
                parent: parent_index,
                counters: BTreeMap::new(),
            },
            started,
            ended: false,
        });
        Some(SpanToken {
            collector: debug.id,
            index,
        })
    }
    /// Invalid/previously ended tokens are harmless. Closing a parent also
    /// closes open descendants at the same instant, preserving containment.
    pub fn end(&mut self, token: Option<SpanToken>, counters: BTreeMap<String, u64>) {
        let Some(token) = token else {
            return;
        };
        let Some(debug) = self.debug.as_mut() else {
            return;
        };
        if token.collector != debug.id
            || token.index >= debug.spans.len()
            || debug.spans[token.index].ended
        {
            return;
        }
        let now = Instant::now();
        for index in token.index..debug.spans.len() {
            let mut ancestor = Some(index);
            let mut contained = false;
            while let Some(current) = ancestor {
                if current == token.index {
                    contained = true;
                    break;
                }
                ancestor = debug.spans[current].value.parent;
            }
            if contained && !debug.spans[index].ended {
                let span = &mut debug.spans[index];
                span.value.duration_ms = now.duration_since(span.started).as_secs_f64() * 1000.0;
                span.ended = true;
            }
        }
        // At most32 distinct bounded keys; callers should use static diagnostic
        // names and counts, never source text or paths.
        let stored = &mut debug.spans[token.index].value.counters;
        for (key, value) in counters.into_iter().take(32) {
            let key: String = key.chars().take(96).collect();
            if stored.len() < 32 || stored.contains_key(&key) {
                stored.insert(key, value);
            }
        }
    }
    pub fn counter(&mut self, token: Option<SpanToken>, name: &'static str, value: u64) {
        let Some(token) = token else {
            return;
        };
        let Some(debug) = self.debug.as_mut() else {
            return;
        };
        if token.collector != debug.id {
            return;
        }
        let Some(span) = debug.spans.get_mut(token.index) else {
            return;
        };
        if span.value.counters.len() < 32 || span.value.counters.contains_key(name) {
            span.value
                .counters
                .insert(name.chars().take(96).collect(), value);
        }
    }
    /// Return the closure's value unchanged, including any Result error.
    /// The closure receives its parent token for safe nested instrumentation.
    pub fn timed<R>(
        &mut self,
        name: &'static str,
        parent: Option<SpanToken>,
        run: impl FnOnce(&mut Self, Option<SpanToken>) -> R,
    ) -> R {
        let token = self.begin(name, parent);
        let value = run(self, token);
        self.end(token, BTreeMap::new());
        value
    }
    pub fn finish(mut self, result_count: usize) -> SearchMetrics {
        let now = Instant::now();
        if let Some(debug) = self.debug.as_mut() {
            for span in &mut debug.spans {
                if !span.ended {
                    span.value.duration_ms =
                        now.duration_since(span.started).as_secs_f64() * 1000.0;
                    span.ended = true;
                }
            }
        }
        let (trace, dropped_spans) = match self.debug {
            Some(debug) => (
                Some(debug.spans.into_iter().map(|span| span.value).collect()),
                debug.dropped,
            ),
            None => (None, 0),
        };
        SearchMetrics {
            total_ms: now.duration_since(self.started).as_secs_f64() * 1000.0,
            result_count,
            trace,
            dropped_spans,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn nesting_contains_children_and_preserves_error() {
        let mut trace = TraceCollector::new(MetricsMode::Debug);
        let result: Result<(), &str> = trace.timed("request", None, |trace, parent| {
            trace.timed("index", parent, |trace, child| {
                trace.counter(child, "rows", 7);
                Err("failure")
            })
        });
        assert_eq!(result, Err("failure"));
        let metrics = trace.finish(0);
        let spans = metrics.trace.unwrap();
        assert_eq!(spans[1].parent, Some(0));
        assert_eq!(spans[1].counters["rows"], 7);
        assert!(spans[1].start_ms >= spans[0].start_ms);
        assert!(
            spans[1].start_ms + spans[1].duration_ms
                <= spans[0].start_ms + spans[0].duration_ms + 1e-6
        );
        assert!(metrics.total_ms >= spans[0].start_ms + spans[0].duration_ms - 1e-6);
    }
    #[test]
    fn basic_disables_stage_storage_and_preserves_counts() {
        let mut trace = TraceCollector::new(MetricsMode::Basic);
        assert!(!trace.is_debug());
        assert!(trace.begin("unused", None).is_none());
        assert_eq!(trace.timed("unused", None, |_, _| 42), 42);
        let metrics = trace.finish(3);
        assert_eq!(metrics.result_count, 3);
        assert!(metrics.trace.is_none());
        assert_eq!(metrics.dropped_spans, 0);
    }
    #[test]
    fn trace_is_capped_and_reports_dropped_spans() {
        let mut trace = TraceCollector::new(MetricsMode::Debug);
        for index in 0..520 {
            let token = trace.begin("stage", None);
            if index == 0 {
                trace.counter(token, "existing", 99);
            }
            let counters = if index == 0 {
                (0..40).map(|key| (format!("counter_{key}"), key)).collect()
            } else {
                BTreeMap::new()
            };
            trace.end(token, counters);
        }
        let metrics = trace.finish(2);
        let spans = metrics.trace.unwrap();
        assert_eq!(spans.len(), 512);
        assert_eq!(spans[0].counters.len(), 32);
        assert_eq!(spans[0].counters["existing"], 99);
        assert_eq!(metrics.dropped_spans, 8);
    }
    #[test]
    fn foreign_tokens_and_early_parent_close_are_safe() {
        let mut one = TraceCollector::new(MetricsMode::Debug);
        let mut two = TraceCollector::new(MetricsMode::Debug);
        let root = one.begin("root", None);
        let child = one.begin("child", root);
        assert!(two.begin("foreign", root).is_none());
        one.end(root, BTreeMap::new());
        one.end(child, BTreeMap::new());
        let spans = one.finish(0).trace.unwrap();
        assert!(
            spans[1].start_ms + spans[1].duration_ms
                <= spans[0].start_ms + spans[0].duration_ms + 1e-6
        );
    }
}