statsig_rust/sdk_diagnostics/
diagnostics.rs

1use super::{
2    diagnostics_utils::DiagnosticsUtils,
3    marker::{ActionType, KeyType, Marker, StepType},
4};
5use crate::event_logging::event_logger::{EventLogger, QueuedEventPayload};
6use crate::log_w;
7use crate::{
8    evaluation::evaluation_details::EvaluationDetails,
9    event_logging::{statsig_event::StatsigEvent, statsig_event_internal::StatsigEventInternal},
10};
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)]
21pub enum ContextType {
22    Initialize, // we only care about initialize for now
23    ConfigSync,
24}
25
26impl fmt::Display for ContextType {
27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28        match self {
29            ContextType::Initialize => write!(f, "initialize"),
30            ContextType::ConfigSync => write!(f, "config_sync"),
31        }
32    }
33}
34
35const TAG: &str = stringify!(Diagnostics);
36const MAX_SAMPLING_RATE: f64 = 10000.0;
37const DEFAULT_SAMPLING_RATE: f64 = 100.0;
38
39pub struct Diagnostics {
40    marker_map: Mutex<HashMap<ContextType, Vec<Marker>>>,
41    event_logger: Arc<EventLogger>,
42    sampling_rates: Mutex<HashMap<String, f64>>,
43}
44
45impl Diagnostics {
46    pub fn new(event_logger: Arc<EventLogger>) -> Self {
47        Self {
48            event_logger,
49            marker_map: Mutex::new(HashMap::new()),
50            sampling_rates: Mutex::new(HashMap::from([
51                ("initialize".to_string(), 10000.0),
52                ("config_sync".to_string(), 1000.0),
53            ])),
54        }
55    }
56
57    pub fn set_sampling_rate(&self, new_sampling_rate: HashMap<std::string::String, f64>) {
58        if let Ok(mut rates) = self.sampling_rates.lock() {
59            for (key, rate) in new_sampling_rate {
60                let clamped_rate = rate.clamp(0.0, MAX_SAMPLING_RATE);
61                rates.insert(key, clamped_rate);
62            }
63        }
64    }
65
66    pub fn get_markers(&self, context_type: &ContextType) -> Option<Vec<Marker>> {
67        if let Ok(map) = self.marker_map.lock() {
68            if let Some(markers) = map.get(context_type) {
69                return Some(markers.clone());
70            }
71        }
72        None
73    }
74
75    pub fn add_marker(&self, context_type: ContextType, marker: Marker) {
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 mark_init_overall_start(&self) {
93        let init_marker = Marker::new(KeyType::Overall, ActionType::Start, Some(StepType::Process));
94        self.add_marker(ContextType::Initialize, init_marker);
95    }
96
97    pub fn mark_init_overall_end(
98        &self,
99        success: bool,
100        error_message: Option<String>,
101        evaluation_details: EvaluationDetails,
102    ) {
103        let mut init_marker =
104            Marker::new(KeyType::Overall, ActionType::End, Some(StepType::Process))
105                .with_is_success(success)
106                .with_eval_details(evaluation_details);
107
108        if let Some(msg) = error_message {
109            init_marker = init_marker.with_message(msg);
110        }
111        self.add_marker(ContextType::Initialize, init_marker);
112        self.enqueue_diagnostics_event(ContextType::Initialize, None);
113    }
114
115    pub fn enqueue_diagnostics_event(&self, context_type: ContextType, key: Option<KeyType>) {
116        let markers = match self.get_markers(&context_type) {
117            Some(m) => m,
118            None => return,
119        };
120
121        if markers.is_empty() {
122            return;
123        }
124
125        let metadata = match DiagnosticsUtils::format_diagnostics_metadata(&context_type, &markers)
126        {
127            Ok(data) => data,
128            Err(err) => {
129                log_w!(TAG, "Failed to format diagnostics metadata: {}", err);
130                return;
131            }
132        };
133
134        let event = StatsigEventInternal::new_diagnostic_event(StatsigEvent {
135            event_name: DIAGNOSTICS_EVENT.to_string(),
136            value: None,
137            metadata: Some(metadata),
138            statsig_metadata: None,
139        });
140
141        if !self.should_sample(&context_type, key) {
142            self.clear_markers(&context_type);
143            return;
144        }
145
146        self.event_logger
147            .enqueue(QueuedEventPayload::CustomEvent(event));
148
149        self.clear_markers(&context_type);
150    }
151
152    pub fn should_sample(&self, context: &ContextType, key: Option<KeyType>) -> bool {
153        let mut rng = rand::thread_rng();
154        let rand_value = rng.gen::<f64>() * MAX_SAMPLING_RATE;
155
156        let sampling_rates = self.sampling_rates.lock().unwrap();
157
158        if *context == ContextType::Initialize {
159            return rand_value
160                < *sampling_rates
161                    .get("initialize")
162                    .unwrap_or(&DEFAULT_SAMPLING_RATE);
163        }
164
165        if let Some(key) = key {
166            if key == KeyType::GetIDList || key == KeyType::GetIDListSources {
167                return rand_value
168                    < *sampling_rates
169                        .get("id_list")
170                        .unwrap_or(&DEFAULT_SAMPLING_RATE);
171            }
172            if key == KeyType::DownloadConfigSpecs {
173                return rand_value < *sampling_rates.get("dcs").unwrap_or(&DEFAULT_SAMPLING_RATE);
174            }
175        }
176
177        rand_value < DEFAULT_SAMPLING_RATE
178    }
179}
statsig_rust/sdk_diagnostics/diagnostics.rs

statsig_rust/sdk_diagnostics/
diagnostics.rs
