statsig_rust/sdk_diagnostics/
diagnostics.rs

1use super::{
2    diagnostics_utils::DiagnosticsUtils,
3    marker::{ActionType, KeyType, Marker, StepType},
4};
5use crate::log_w;
6use crate::{
7    evaluation::evaluation_details::EvaluationDetails,
8    event_logging::{statsig_event::StatsigEvent, statsig_event_internal::StatsigEventInternal},
9};
10use crate::{
11    event_logging::event_logger::{EventLogger, QueuedEventPayload},
12    read_lock_or_else, SpecStore,
13};
14use chrono::Utc;
15use std::collections::HashMap;
16use std::fmt;
17use std::sync::{Arc, Mutex};
18
19const MAX_MARKER_COUNT: usize = 50;
20pub const DIAGNOSTICS_EVENT: &str = "statsig::diagnostics";
21
22#[derive(Eq, Hash, PartialEq)]
23pub enum ContextType {
24    Initialize, // we only care about initialize for now
25}
26
27impl fmt::Display for ContextType {
28    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
29        match self {
30            ContextType::Initialize => write!(f, "initialize"),
31        }
32    }
33}
34
35const TAG: &str = stringify!(Diagnostics);
36
37pub struct Diagnostics {
38    marker_map: Mutex<HashMap<ContextType, Vec<Marker>>>,
39    event_logger: Arc<EventLogger>,
40    spec_store: Arc<SpecStore>,
41}
42
43impl Diagnostics {
44    #[must_use]
45    pub fn new(event_logger: Arc<EventLogger>, spec_store: Arc<SpecStore>) -> Self {
46        Self {
47            event_logger,
48            marker_map: Mutex::new(HashMap::new()),
49            spec_store,
50        }
51    }
52
53    pub fn get_markers(&self, context_type: &ContextType) -> Option<Vec<Marker>> {
54        if let Ok(map) = self.marker_map.lock() {
55            if let Some(markers) = map.get(context_type) {
56                return Some(markers.clone());
57            }
58        }
59        None
60    }
61
62    pub fn add_marker(&self, context_type: ContextType, marker: Marker) {
63        if let Ok(mut map) = self.marker_map.lock() {
64            let entry = map.entry(context_type).or_insert_with(Vec::new);
65            if entry.len() < MAX_MARKER_COUNT {
66                entry.push(marker);
67            }
68        }
69    }
70
71    pub fn clear_markers(&self, context_type: &ContextType) {
72        if let Ok(mut map) = self.marker_map.lock() {
73            if let Some(markers) = map.get_mut(context_type) {
74                markers.clear();
75            }
76        }
77    }
78
79    pub fn mark_init_overall_start(&self) {
80        let init_marker = Marker::new(
81            KeyType::Overall,
82            ActionType::Start,
83            Some(StepType::Process),
84            Utc::now().timestamp_millis() as u64,
85        );
86        self.add_marker(ContextType::Initialize, init_marker);
87    }
88
89    pub fn mark_init_overall_end(
90        &self,
91        success: bool,
92        error_message: Option<String>,
93        evaluation_details: EvaluationDetails,
94    ) {
95        let mut init_marker = Marker::new(
96            KeyType::Overall,
97            ActionType::End,
98            Some(StepType::Process),
99            Utc::now().timestamp_millis() as u64,
100        )
101        .with_is_success(success)
102        .with_eval_details(evaluation_details);
103
104        if let Some(msg) = error_message {
105            init_marker = init_marker.with_message(msg);
106        }
107        self.add_marker(ContextType::Initialize, init_marker);
108        self.enqueue_diagnostics_event(ContextType::Initialize);
109    }
110
111    pub fn enqueue_diagnostics_event(&self, context_type: ContextType) {
112        let markers = match self.get_markers(&context_type) {
113            Some(m) => m,
114            None => return,
115        };
116
117        if markers.is_empty() {
118            return;
119        }
120
121        let metadata = match DiagnosticsUtils::format_diagnostics_metadata(&context_type, &markers)
122        {
123            Ok(data) => data,
124            Err(err) => {
125                log_w!(TAG, "Failed to format diagnostics metadata: {}", err);
126                return;
127            }
128        };
129
130        let event = StatsigEventInternal::new_diagnostic_event(StatsigEvent {
131            event_name: DIAGNOSTICS_EVENT.to_string(),
132            value: None,
133            metadata: Some(metadata),
134            statsig_metadata: None,
135        });
136
137        let data = read_lock_or_else!(self.spec_store.data, {
138            log_w!(TAG, "Failed to acquire read lock for diagnostics event");
139            return;
140        });
141
142        let diagnostics = &data.values.diagnostics;
143
144        if let Some(diagnostics) = diagnostics {
145            if let Some(sample_rate) = diagnostics.get(&context_type.to_string()) {
146                if !DiagnosticsUtils::should_sample(*sample_rate) {
147                    self.clear_markers(&context_type);
148                    return;
149                }
150            }
151        }
152
153        self.event_logger
154            .enqueue(QueuedEventPayload::CustomEvent(event));
155
156        self.clear_markers(&context_type);
157    }
158}
statsig_rust/sdk_diagnostics/diagnostics.rs

statsig_rust/sdk_diagnostics/
diagnostics.rs
