Skip to main content

vtcode_core/telemetry/
perf.rs

1use hashbrown::HashMap;
2use std::sync::Arc;
3use std::sync::atomic::{AtomicBool, Ordering};
4use std::time::{Duration, Instant};
5
6use once_cell::sync::OnceCell;
7use tokio::runtime::Handle;
8use tracing::debug;
9
10use crate::config::TelemetryConfig;
11use crate::telemetry::{TelemetryEvent, TelemetryPipeline};
12
13#[derive(Debug)]
14pub struct PerfRecorder {
15    enabled: AtomicBool,
16    pipeline: Arc<TelemetryPipeline>,
17}
18
19static PERF_RECORDER: OnceCell<PerfRecorder> = OnceCell::new();
20
21pub fn initialize_perf_telemetry(config: &TelemetryConfig) {
22    let recorder = PerfRecorder {
23        enabled: AtomicBool::new(config.perf_events),
24        pipeline: Arc::new(TelemetryPipeline::new(config.clone())),
25    };
26    let _ = PERF_RECORDER.set(recorder);
27}
28
29pub fn enabled() -> bool {
30    PERF_RECORDER
31        .get()
32        .map(|recorder| recorder.enabled.load(Ordering::Relaxed))
33        .unwrap_or(false)
34}
35
36pub fn record_duration(name: &'static str, duration: Duration, tags: HashMap<String, String>) {
37    record_value(name, duration.as_secs_f64() * 1000.0, tags);
38}
39
40pub fn record_value(name: &'static str, value: f64, tags: HashMap<String, String>) {
41    let Some(recorder) = PERF_RECORDER.get() else {
42        return;
43    };
44    if !recorder.enabled.load(Ordering::Relaxed) {
45        return;
46    }
47
48    let mut event = TelemetryEvent::new(name, value);
49    event.tags = tags;
50
51    if let Ok(handle) = Handle::try_current() {
52        let pipeline = Arc::clone(&recorder.pipeline);
53        handle.spawn(async move {
54            let _ = pipeline.record(event).await;
55        });
56    } else {
57        debug!(name, "Skipping perf event (no runtime)");
58    }
59}
60
61pub struct PerfSpan {
62    name: &'static str,
63    start: Instant,
64    tags: HashMap<String, String>,
65    enabled: bool,
66}
67
68impl PerfSpan {
69    pub fn new(name: &'static str) -> Self {
70        Self {
71            name,
72            start: Instant::now(),
73            tags: HashMap::new(),
74            enabled: enabled(),
75        }
76    }
77
78    pub fn tag(&mut self, key: impl Into<String>, value: impl Into<String>) {
79        if self.enabled {
80            self.tags.insert(key.into(), value.into());
81        }
82    }
83}
84
85impl Drop for PerfSpan {
86    fn drop(&mut self) {
87        if self.enabled {
88            record_duration(
89                self.name,
90                self.start.elapsed(),
91                std::mem::take(&mut self.tags),
92            );
93        }
94    }
95}