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