Skip to main content

wp_knowledge/
telemetry.rs

1use std::sync::atomic::{AtomicBool, Ordering};
2use std::sync::{Arc, OnceLock, RwLock};
3use std::time::Duration;
4
5use crate::loader::ProviderKind;
6use crate::runtime::QueryModeTag;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum CacheLayer {
10    Local,
11    Result,
12    Metadata,
13}
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum CacheOutcome {
17    Hit,
18    Miss,
19}
20
21#[derive(Debug, Clone)]
22pub struct CacheTelemetryEvent {
23    pub layer: CacheLayer,
24    pub outcome: CacheOutcome,
25    pub provider_kind: Option<ProviderKind>,
26}
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub enum ReloadOutcome {
30    Success,
31    Failure,
32}
33
34#[derive(Debug, Clone)]
35pub struct ReloadTelemetryEvent {
36    pub outcome: ReloadOutcome,
37    pub provider_kind: ProviderKind,
38}
39
40#[derive(Debug, Clone)]
41pub struct QueryTelemetryEvent {
42    pub provider_kind: ProviderKind,
43    pub mode: QueryModeTag,
44    pub success: bool,
45    pub elapsed: Duration,
46}
47
48pub trait KnowledgeTelemetry: Send + Sync {
49    fn is_noop(&self) -> bool {
50        false
51    }
52    fn on_cache(&self, _event: &CacheTelemetryEvent) {}
53    fn on_reload(&self, _event: &ReloadTelemetryEvent) {}
54    fn on_query(&self, _event: &QueryTelemetryEvent) {}
55}
56
57#[derive(Debug, Default)]
58pub struct NoopTelemetry;
59
60impl KnowledgeTelemetry for NoopTelemetry {
61    fn is_noop(&self) -> bool {
62        true
63    }
64}
65
66fn telemetry_slot() -> &'static RwLock<Arc<dyn KnowledgeTelemetry>> {
67    static SLOT: OnceLock<RwLock<Arc<dyn KnowledgeTelemetry>>> = OnceLock::new();
68    SLOT.get_or_init(|| RwLock::new(Arc::new(NoopTelemetry)))
69}
70
71fn telemetry_enabled_flag() -> &'static AtomicBool {
72    static FLAG: OnceLock<AtomicBool> = OnceLock::new();
73    FLAG.get_or_init(|| AtomicBool::new(false))
74}
75
76pub fn telemetry() -> Arc<dyn KnowledgeTelemetry> {
77    telemetry_slot()
78        .read()
79        .expect("knowledge telemetry lock poisoned")
80        .clone()
81}
82
83pub fn telemetry_enabled() -> bool {
84    telemetry_enabled_flag().load(Ordering::Relaxed)
85}
86
87pub fn install_telemetry(telemetry: Arc<dyn KnowledgeTelemetry>) -> Arc<dyn KnowledgeTelemetry> {
88    let mut guard = telemetry_slot()
89        .write()
90        .expect("knowledge telemetry lock poisoned");
91    let enabled = !telemetry.is_noop();
92    let previous = std::mem::replace(&mut *guard, telemetry);
93    telemetry_enabled_flag().store(enabled, Ordering::Relaxed);
94    previous
95}
96
97pub fn reset_telemetry() -> Arc<dyn KnowledgeTelemetry> {
98    install_telemetry(Arc::new(NoopTelemetry))
99}