statsig_rust/sdk_diagnostics/
diagnostics.rs

1use super::diagnostics_utils::DiagnosticsUtils;
2use super::marker::{KeyType, Marker};
3
4use crate::event_logging::event_logger::{EventLogger, QueuedEventPayload};
5use crate::event_logging::{
6    statsig_event::StatsigEvent, statsig_event_internal::StatsigEventInternal,
7};
8
9use crate::global_configs::{GlobalConfigs, MAX_SAMPLING_RATE};
10
11use crate::log_w;
12
13use rand::Rng;
14use serde::Serialize;
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, Clone, Serialize, Debug, Copy)]
23pub enum ContextType {
24    Initialize,
25    ConfigSync,
26}
27
28impl fmt::Display for ContextType {
29    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
30        match self {
31            ContextType::Initialize => write!(f, "initialize"),
32            ContextType::ConfigSync => write!(f, "config_sync"),
33        }
34    }
35}
36
37const TAG: &str = stringify!(Diagnostics);
38const DEFAULT_SAMPLING_RATE: f64 = 100.0;
39
40pub struct Diagnostics {
41    marker_map: Mutex<HashMap<ContextType, Vec<Marker>>>,
42    event_logger: Arc<EventLogger>,
43    global_configs: Arc<GlobalConfigs>,
44    context: Mutex<ContextType>,
45}
46
47impl Diagnostics {
48    pub fn new(event_logger: Arc<EventLogger>, sdk_key: &str) -> Self {
49        Self {
50            event_logger,
51            marker_map: Mutex::new(HashMap::new()),
52            global_configs: GlobalConfigs::get_instance(sdk_key),
53            context: Mutex::new(ContextType::Initialize),
54        }
55    }
56
57    pub fn set_context(&self, context: &ContextType) {
58        if let Ok(mut ctx) = self.context.lock() {
59            *ctx = *context;
60        }
61    }
62
63    pub fn get_markers(&self, context_type: &ContextType) -> Option<Vec<Marker>> {
64        if let Ok(map) = self.marker_map.lock() {
65            if let Some(markers) = map.get(context_type) {
66                return Some(markers.clone());
67            }
68        }
69        None
70    }
71
72    pub fn add_marker(&self, context_type: Option<&ContextType>, marker: Marker) {
73        let context_type = self.get_context(context_type);
74        if let Ok(mut map) = self.marker_map.lock() {
75            let entry = map.entry(context_type).or_insert_with(Vec::new);
76            if entry.len() < MAX_MARKER_COUNT {
77                entry.push(marker);
78            }
79        }
80    }
81
82    pub fn clear_markers(&self, context_type: &ContextType) {
83        if let Ok(mut map) = self.marker_map.lock() {
84            if let Some(markers) = map.get_mut(context_type) {
85                markers.clear();
86            }
87        }
88    }
89
90    pub fn enqueue_diagnostics_event(
91        &self,
92        context_type: Option<&ContextType>,
93        key: Option<KeyType>,
94    ) {
95        let context_type: ContextType = self.get_context(context_type);
96        let markers = match self.get_markers(&context_type) {
97            Some(m) => m,
98            None => return,
99        };
100
101        if markers.is_empty() {
102            return;
103        }
104
105        let metadata = match DiagnosticsUtils::format_diagnostics_metadata(&context_type, &markers)
106        {
107            Ok(data) => data,
108            Err(err) => {
109                log_w!(TAG, "Failed to format diagnostics metadata: {}", err);
110                return;
111            }
112        };
113
114        let event = StatsigEventInternal::new_diagnostic_event(StatsigEvent {
115            event_name: DIAGNOSTICS_EVENT.to_string(),
116            value: None,
117            metadata: Some(metadata),
118            statsig_metadata: None,
119        });
120
121        if !self.should_sample(&context_type, key) {
122            self.clear_markers(&context_type);
123            return;
124        }
125
126        self.event_logger
127            .enqueue(QueuedEventPayload::CustomEvent(event));
128
129        self.clear_markers(&context_type);
130    }
131
132    pub fn should_sample(&self, context: &ContextType, key: Option<KeyType>) -> bool {
133        let mut rng = rand::thread_rng();
134        let rand_value = rng.gen::<f64>() * MAX_SAMPLING_RATE;
135
136        let sampling_rates = self.global_configs.get_diagnostics_sampling_rate();
137
138        if *context == ContextType::Initialize {
139            return rand_value
140                < *sampling_rates
141                    .get("initialize")
142                    .unwrap_or(&DEFAULT_SAMPLING_RATE);
143        }
144
145        if let Some(key) = key {
146            match key {
147                KeyType::GetIDListSources => {
148                    return rand_value
149                        < *sampling_rates
150                            .get("get_id_list")
151                            .unwrap_or(&DEFAULT_SAMPLING_RATE);
152                }
153                KeyType::DownloadConfigSpecs => {
154                    return rand_value
155                        < *sampling_rates.get("dcs").unwrap_or(&DEFAULT_SAMPLING_RATE);
156                }
157                _ => {}
158            }
159        }
160
161        rand_value < DEFAULT_SAMPLING_RATE
162    }
163
164    fn get_context(&self, maybe_context: Option<&ContextType>) -> ContextType {
165        let context_type = match maybe_context {
166            Some(ctx) => *ctx,
167            None => *self.context.lock().unwrap(),
168        };
169        context_type
170    }
171}