statsig_rust/sdk_diagnostics/
diagnostics.rs

1use super::diagnostics_utils::DiagnosticsUtils;
2use super::marker::{KeyType, Marker};
3
4use crate::event_logging::event_queue::queued_passthrough::EnqueuePassthroughOp;
5use crate::event_logging::statsig_event_internal::StatsigEventInternal;
6use crate::global_configs::{GlobalConfigs, MAX_SAMPLING_RATE};
7
8use crate::log_w;
9
10use crate::event_logging::event_logger::EventLogger;
11use rand::Rng;
12use serde::Serialize;
13use std::collections::HashMap;
14use std::fmt;
15use std::sync::{Arc, Mutex};
16
17const MAX_MARKER_COUNT: usize = 50;
18pub const DIAGNOSTICS_EVENT: &str = "statsig::diagnostics";
19
20#[derive(Eq, Hash, PartialEq, Clone, Serialize, Debug, Copy)]
21pub enum ContextType {
22    Initialize,
23    ConfigSync,
24    Unknown,
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            ContextType::ConfigSync => write!(f, "config_sync"),
32            ContextType::Unknown => write!(f, "unknown"),
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        if !self.should_sample(&context_type, key) {
106            self.clear_markers(&context_type);
107            return;
108        }
109
110        let metadata = match DiagnosticsUtils::format_diagnostics_metadata(&context_type, &markers)
111        {
112            Ok(data) => data,
113            Err(err) => {
114                log_w!(TAG, "Failed to format diagnostics metadata: {}", err);
115                return;
116            }
117        };
118
119        self.event_logger.enqueue(EnqueuePassthroughOp {
120            event: StatsigEventInternal::new_diagnostic_event(metadata),
121        });
122        self.clear_markers(&context_type);
123    }
124
125    pub fn should_sample(&self, context: &ContextType, key: Option<KeyType>) -> bool {
126        fn check_sampling_rate(sampling_rate: Option<&f64>) -> bool {
127            let mut rng = rand::thread_rng();
128            let rand_value = rng.gen::<f64>() * MAX_SAMPLING_RATE;
129
130            match sampling_rate {
131                Some(sampling_rate) => rand_value < *sampling_rate,
132                None => rand_value < DEFAULT_SAMPLING_RATE,
133            }
134        }
135
136        if *context == ContextType::Initialize {
137            return self
138                .global_configs
139                .use_diagnostics_sampling_rate("initialize", check_sampling_rate);
140        }
141
142        if let Some(key) = key {
143            match key {
144                KeyType::GetIDListSources => {
145                    return self
146                        .global_configs
147                        .use_diagnostics_sampling_rate("get_id_list", check_sampling_rate);
148                }
149                KeyType::DownloadConfigSpecs => {
150                    return self
151                        .global_configs
152                        .use_diagnostics_sampling_rate("dcs", check_sampling_rate);
153                }
154                _ => {}
155            }
156        }
157
158        check_sampling_rate(None)
159    }
160
161    fn get_context(&self, maybe_context: Option<&ContextType>) -> ContextType {
162        maybe_context
163            .copied()
164            .or_else(|| self.context.try_lock().ok().map(|c| *c))
165            .unwrap_or(ContextType::Unknown)
166    }
167}