feat: add multimodal search and granular performance diagnostics

crates/searchi-core/src/metrics.rs

@@ -0,0 +1,292 @@
+//! 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
+        );
+    }
+}