Skip to main content

vv_agent/
tracing.rs

1use std::fs::{File, OpenOptions};
2use std::io::Write;
3use std::path::Path;
4use std::sync::Mutex;
5use std::time::{SystemTime, UNIX_EPOCH};
6
7use serde::{Deserialize, Serialize};
8use serde_json::json;
9
10pub trait TraceSink: Send + Sync {
11    fn on_span_start(&self, span: &Span);
12    fn on_span_end(&self, span: &Span);
13
14    fn flush(&self) -> Result<(), String> {
15        Ok(())
16    }
17}
18
19#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
20pub struct Span {
21    pub trace_id: String,
22    pub span_id: String,
23    pub name: String,
24    pub run_id: String,
25    pub agent_name: Option<String>,
26}
27
28impl Span {
29    pub fn new(
30        run_id: impl Into<String>,
31        name: impl Into<String>,
32        agent_name: Option<String>,
33    ) -> Self {
34        let run_id = run_id.into();
35        let name = name.into();
36        Self {
37            trace_id: run_id.clone(),
38            span_id: format!("{}_{}", name, timestamp_millis()),
39            name,
40            run_id,
41            agent_name,
42        }
43    }
44}
45
46pub struct JsonlTraceExporter {
47    file: Mutex<File>,
48}
49
50impl JsonlTraceExporter {
51    pub fn new(path: impl AsRef<Path>) -> Result<Self, String> {
52        let file = OpenOptions::new()
53            .create(true)
54            .append(true)
55            .open(path)
56            .map_err(|error| error.to_string())?;
57        Ok(Self {
58            file: Mutex::new(file),
59        })
60    }
61
62    fn write_event(&self, event: &str, span: &Span) {
63        if let Ok(mut file) = self.file.lock() {
64            let _ = writeln!(
65                file,
66                "{}",
67                json!({
68                    "event": event,
69                    "timestamp_ms": timestamp_millis(),
70                    "span": span,
71                })
72            );
73        }
74    }
75}
76
77impl TraceSink for JsonlTraceExporter {
78    fn on_span_start(&self, span: &Span) {
79        self.write_event("span_start", span);
80    }
81
82    fn on_span_end(&self, span: &Span) {
83        self.write_event("span_end", span);
84    }
85
86    fn flush(&self) -> Result<(), String> {
87        self.file
88            .lock()
89            .map_err(|_| "trace exporter lock poisoned".to_string())?
90            .flush()
91            .map_err(|error| error.to_string())
92    }
93}
94
95fn timestamp_millis() -> u128 {
96    SystemTime::now()
97        .duration_since(UNIX_EPOCH)
98        .map(|duration| duration.as_millis())
99        .unwrap_or_default()
100}