1use crate::types::Degrade;
7use chrono::{Local, SecondsFormat};
8use parking_lot::Mutex;
9use serde::{Deserialize, Serialize};
10use std::collections::BTreeMap;
11use std::sync::Arc;
12use std::time::Instant;
13
14pub type EventSink = Arc<dyn Fn(&LogEvent) + Send + Sync>;
17
18#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct LogEvent {
24 pub ts: String,
26 pub kind: String,
28 pub ms: u64,
30 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
33 pub stages: BTreeMap<String, u64>,
34 #[serde(default, skip_serializing_if = "Option::is_none")]
36 pub candidates: Option<usize>,
37 #[serde(default, skip_serializing_if = "Option::is_none")]
39 pub folded: Option<usize>,
40 #[serde(default, skip_serializing_if = "Option::is_none")]
42 pub rerank_docs: Option<usize>,
43 #[serde(default, skip_serializing_if = "Option::is_none")]
45 pub rerank_tokens: Option<usize>,
46 #[serde(default, skip_serializing_if = "Option::is_none")]
48 pub hits: Option<usize>,
49 #[serde(default, skip_serializing_if = "Option::is_none")]
51 pub documents: Option<usize>,
52 #[serde(default, skip_serializing_if = "Option::is_none")]
54 pub format: Option<String>,
55 #[serde(default, skip_serializing_if = "Vec::is_empty")]
57 pub degraded: Vec<Degrade>,
58}
59
60impl LogEvent {
61 pub fn new(kind: &str) -> Self {
63 Self {
64 ts: Local::now().to_rfc3339_opts(SecondsFormat::Millis, false),
65 kind: kind.to_string(),
66 ms: 0,
67 stages: BTreeMap::new(),
68 candidates: None,
69 folded: None,
70 rerank_docs: None,
71 rerank_tokens: None,
72 hits: None,
73 documents: None,
74 format: None,
75 degraded: Vec::new(),
76 }
77 }
78}
79
80#[derive(Default)]
82pub(crate) struct EventRegistry {
83 sink: Mutex<Option<EventSink>>,
84}
85
86impl EventRegistry {
87 pub fn get(&self) -> Option<EventSink> {
88 self.sink.lock().clone()
89 }
90 pub fn set(&self, sink: EventSink) {
91 *self.sink.lock() = Some(sink);
92 }
93 pub fn clear(&self) -> bool {
94 self.sink.lock().take().is_some()
95 }
96 pub fn is_registered(&self) -> bool {
97 self.sink.lock().is_some()
98 }
99}
100
101pub(crate) struct StageTimer {
103 last: Instant,
104 marks: BTreeMap<String, u64>,
105}
106
107impl StageTimer {
108 pub fn start() -> Self {
109 Self { last: Instant::now(), marks: BTreeMap::new() }
110 }
111 pub fn mark(&mut self, name: &str) {
112 let now = Instant::now();
113 self.marks.insert(name.to_string(), (now - self.last).as_millis() as u64);
114 self.last = now;
115 }
116 pub fn finish(mut self, name: &str) -> BTreeMap<String, u64> {
118 self.mark(name);
119 self.marks
120 }
121}