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 parking_lot::Mutex;
12use rand::Rng;
13use serde::Serialize;
14use std::collections::HashMap;
15use std::fmt;
16use std::sync::Arc;
17use std::time::Duration;
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        match self.context.try_lock_for(Duration::from_secs(5)) {
61            Some(mut ctx) => {
62                *ctx = *context;
63            }
64            None => {
65                log_w!(TAG, "Failed to set context: Failed to lock context");
66            }
67        }
68    }
69
70    pub fn get_markers(&self, context_type: &ContextType) -> Option<Vec<Marker>> {
71        match self.marker_map.try_lock_for(Duration::from_secs(5)) {
72            Some(map) => {
73                if let Some(markers) = map.get(context_type) {
74                    return Some(markers.clone());
75                }
76            }
77            None => {
78                log_w!(TAG, "Failed to get markers: Failed to lock marker_map");
79            }
80        }
81
82        None
83    }
84
85    pub fn add_marker(&self, context_type: Option<&ContextType>, marker: Marker) {
86        let context_type = self.get_context(context_type);
87
88        match self.marker_map.try_lock_for(Duration::from_secs(5)) {
89            Some(mut map) => {
90                let entry = map.entry(context_type).or_insert_with(Vec::new);
91                if entry.len() < MAX_MARKER_COUNT {
92                    entry.push(marker);
93                }
94            }
95            None => {
96                log_w!(TAG, "Failed to add marker: Failed to lock marker_map");
97            }
98        }
99    }
100
101    pub fn clear_markers(&self, context_type: &ContextType) {
102        match self.marker_map.try_lock_for(Duration::from_secs(5)) {
103            Some(mut map) => {
104                if let Some(markers) = map.get_mut(context_type) {
105                    markers.clear();
106                }
107            }
108            None => {
109                log_w!(TAG, "Failed to clear markers: Failed to lock marker_map");
110            }
111        }
112    }
113
114    pub fn enqueue_diagnostics_event(
115        &self,
116        context_type: Option<&ContextType>,
117        key: Option<KeyType>,
118    ) {
119        let context_type: ContextType = self.get_context(context_type);
120        let markers = match self.get_markers(&context_type) {
121            Some(m) => m,
122            None => return,
123        };
124
125        if markers.is_empty() {
126            return;
127        }
128
129        if !self.should_sample(&context_type, key) {
130            self.clear_markers(&context_type);
131            return;
132        }
133
134        let metadata = match DiagnosticsUtils::format_diagnostics_metadata(&context_type, &markers)
135        {
136            Ok(data) => data,
137            Err(err) => {
138                log_w!(TAG, "Failed to format diagnostics metadata: {}", err);
139                return;
140            }
141        };
142
143        self.event_logger.enqueue(EnqueuePassthroughOp {
144            event: StatsigEventInternal::new_diagnostic_event(metadata),
145        });
146        self.clear_markers(&context_type);
147    }
148
149    pub fn should_sample(&self, context: &ContextType, key: Option<KeyType>) -> bool {
150        fn check_sampling_rate(sampling_rate: Option<&f64>) -> bool {
151            let mut rng = rand::thread_rng();
152            let rand_value = rng.gen::<f64>() * MAX_SAMPLING_RATE;
153
154            match sampling_rate {
155                Some(sampling_rate) => rand_value < *sampling_rate,
156                None => rand_value < DEFAULT_SAMPLING_RATE,
157            }
158        }
159
160        if *context == ContextType::Initialize {
161            return self
162                .global_configs
163                .use_diagnostics_sampling_rate("initialize", check_sampling_rate);
164        }
165
166        if let Some(key) = key {
167            match key {
168                KeyType::GetIDListSources => {
169                    return self
170                        .global_configs
171                        .use_diagnostics_sampling_rate("get_id_list", check_sampling_rate);
172                }
173                KeyType::DownloadConfigSpecs => {
174                    return self
175                        .global_configs
176                        .use_diagnostics_sampling_rate("dcs", check_sampling_rate);
177                }
178                _ => {}
179            }
180        }
181
182        check_sampling_rate(None)
183    }
184
185    fn get_context(&self, maybe_context: Option<&ContextType>) -> ContextType {
186        maybe_context
187            .copied()
188            .or_else(|| match self.context.try_lock_for(Duration::from_secs(5)) {
189                Some(c) => Some(*c),
190                None => {
191                    log_w!(TAG, "Failed to lock context");
192                    None
193                }
194            })
195            .unwrap_or(ContextType::Unknown)
196    }
197}
statsig_rust/sdk_diagnostics/diagnostics.rs

statsig_rust/sdk_diagnostics/
diagnostics.rs
