1use std::mem;
6use std::sync::Arc;
7
8use crate::common_metric_data::CommonMetricDataInternal;
9use crate::error_recording::{record_error, test_get_num_recorded_errors, ErrorType};
10use crate::histogram::{Bucketing, Histogram, HistogramType, LinearOrExponential};
11use crate::metrics::{DistributionData, Metric, MetricType};
12use crate::storage::StorageManager;
13use crate::CommonMetricData;
14use crate::Glean;
15
16#[derive(Clone, Debug)]
18pub struct CustomDistributionMetric {
19    meta: Arc<CommonMetricDataInternal>,
20    range_min: u64,
21    range_max: u64,
22    bucket_count: u64,
23    histogram_type: HistogramType,
24}
25
26pub(crate) fn snapshot<B: Bucketing>(hist: &Histogram<B>) -> DistributionData {
30    DistributionData {
31        values: hist
32            .snapshot_values()
33            .into_iter()
34            .map(|(k, v)| (k as i64, v as i64))
35            .collect(),
36        sum: hist.sum() as i64,
37        count: hist.count() as i64,
38    }
39}
40
41impl MetricType for CustomDistributionMetric {
42    fn meta(&self) -> &CommonMetricDataInternal {
43        &self.meta
44    }
45
46    fn with_name(&self, name: String) -> Self {
47        let mut meta = (*self.meta).clone();
48        meta.inner.name = name;
49        Self {
50            meta: Arc::new(meta),
51            range_min: self.range_min,
52            range_max: self.range_max,
53            bucket_count: self.bucket_count,
54            histogram_type: self.histogram_type,
55        }
56    }
57
58    fn with_dynamic_label(&self, label: String) -> Self {
59        let mut meta = (*self.meta).clone();
60        meta.inner.dynamic_label = Some(label);
61        Self {
62            meta: Arc::new(meta),
63            range_min: self.range_min,
64            range_max: self.range_max,
65            bucket_count: self.bucket_count,
66            histogram_type: self.histogram_type,
67        }
68    }
69}
70
71impl CustomDistributionMetric {
76    pub fn new(
78        meta: CommonMetricData,
79        range_min: i64,
80        range_max: i64,
81        bucket_count: i64,
82        histogram_type: HistogramType,
83    ) -> Self {
84        Self {
85            meta: Arc::new(meta.into()),
86            range_min: range_min as u64,
87            range_max: range_max as u64,
88            bucket_count: bucket_count as u64,
89            histogram_type,
90        }
91    }
92
93    pub fn accumulate_samples(&self, samples: Vec<i64>) {
109        let metric = self.clone();
110        crate::launch_with_glean(move |glean| metric.accumulate_samples_sync(glean, &samples))
111    }
112
113    pub fn accumulate_single_sample(&self, sample: i64) {
128        let metric = self.clone();
129        crate::launch_with_glean(move |glean| metric.accumulate_samples_sync(glean, &[sample]))
130    }
131
132    #[doc(hidden)]
136    pub fn accumulate_samples_sync(&self, glean: &Glean, samples: &[i64]) {
137        if !self.should_record(glean) {
138            return;
139        }
140
141        let mut num_negative_samples = 0;
142
143        fn accumulate<B: Bucketing, F>(
146            samples: &[i64],
147            mut hist: Histogram<B>,
148            metric: F,
149        ) -> (i32, Metric)
150        where
151            F: Fn(Histogram<B>) -> Metric,
152        {
153            let mut num_negative_samples = 0;
154            for &sample in samples.iter() {
155                if sample < 0 {
156                    num_negative_samples += 1;
157                } else {
158                    let sample = sample as u64;
159                    hist.accumulate(sample);
160                }
161            }
162            (num_negative_samples, metric(hist))
163        }
164
165        glean.storage().record_with(glean, &self.meta, |old_value| {
166            let (num_negative, hist) = match self.histogram_type {
167                HistogramType::Linear => {
168                    let hist = if let Some(Metric::CustomDistributionLinear(hist)) = old_value {
169                        hist
170                    } else {
171                        Histogram::linear(
172                            self.range_min,
173                            self.range_max,
174                            self.bucket_count as usize,
175                        )
176                    };
177                    accumulate(samples, hist, Metric::CustomDistributionLinear)
178                }
179                HistogramType::Exponential => {
180                    let hist = if let Some(Metric::CustomDistributionExponential(hist)) = old_value
181                    {
182                        hist
183                    } else {
184                        Histogram::exponential(
185                            self.range_min,
186                            self.range_max,
187                            self.bucket_count as usize,
188                        )
189                    };
190                    accumulate(samples, hist, Metric::CustomDistributionExponential)
191                }
192            };
193
194            num_negative_samples = num_negative;
195            hist
196        });
197
198        if num_negative_samples > 0 {
199            let msg = format!("Accumulated {} negative samples", num_negative_samples);
200            record_error(
201                glean,
202                &self.meta,
203                ErrorType::InvalidValue,
204                msg,
205                num_negative_samples,
206            );
207        }
208    }
209
210    #[doc(hidden)]
212    pub fn get_value<'a, S: Into<Option<&'a str>>>(
213        &self,
214        glean: &Glean,
215        ping_name: S,
216    ) -> Option<DistributionData> {
217        let queried_ping_name = ping_name
218            .into()
219            .unwrap_or_else(|| &self.meta().inner.send_in_pings[0]);
220
221        match StorageManager.snapshot_metric_for_test(
222            glean.storage(),
223            queried_ping_name,
224            &self.meta.identifier(glean),
225            self.meta.inner.lifetime,
226        ) {
227            Some(Metric::CustomDistributionExponential(hist)) => Some(snapshot(&hist)),
229            Some(Metric::CustomDistributionLinear(hist)) => Some(snapshot(&hist)),
230            _ => None,
231        }
232    }
233
234    pub fn test_get_value(&self, ping_name: Option<String>) -> Option<DistributionData> {
249        crate::block_on_dispatcher();
250        crate::core::with_glean(|glean| self.get_value(glean, ping_name.as_deref()))
251    }
252
253    pub fn test_get_num_recorded_errors(&self, error: ErrorType) -> i32 {
265        crate::block_on_dispatcher();
266
267        crate::core::with_glean(|glean| {
268            test_get_num_recorded_errors(glean, self.meta(), error).unwrap_or(0)
269        })
270    }
271
272    pub fn start_buffer(&self) -> LocalCustomDistribution<'_> {
277        LocalCustomDistribution::new(self)
278    }
279
280    fn commit_histogram(&self, histogram: Histogram<LinearOrExponential>) {
281        let metric = self.clone();
282        crate::launch_with_glean(move |glean| {
283            glean
284                .storage()
285                .record_with(glean, &metric.meta, move |old_value| {
286                    match metric.histogram_type {
287                        HistogramType::Linear => {
288                            let mut hist =
289                                if let Some(Metric::CustomDistributionLinear(hist)) = old_value {
290                                    hist
291                                } else {
292                                    Histogram::linear(
293                                        metric.range_min,
294                                        metric.range_max,
295                                        metric.bucket_count as usize,
296                                    )
297                                };
298
299                            hist._merge(&histogram);
300                            Metric::CustomDistributionLinear(hist)
301                        }
302                        HistogramType::Exponential => {
303                            let mut hist = if let Some(Metric::CustomDistributionExponential(
304                                hist,
305                            )) = old_value
306                            {
307                                hist
308                            } else {
309                                Histogram::exponential(
310                                    metric.range_min,
311                                    metric.range_max,
312                                    metric.bucket_count as usize,
313                                )
314                            };
315
316                            hist._merge(&histogram);
317                            Metric::CustomDistributionExponential(hist)
318                        }
319                    }
320                });
321        });
322    }
323}
324
325pub struct LocalCustomDistribution<'a> {
330    histogram: Histogram<LinearOrExponential>,
331    metric: &'a CustomDistributionMetric,
332}
333
334impl<'a> LocalCustomDistribution<'a> {
335    fn new(metric: &'a CustomDistributionMetric) -> Self {
337        let histogram = match metric.histogram_type {
338            HistogramType::Linear => Histogram::<LinearOrExponential>::_linear(
339                metric.range_min,
340                metric.range_max,
341                metric.bucket_count as usize,
342            ),
343            HistogramType::Exponential => Histogram::<LinearOrExponential>::_exponential(
344                metric.range_min,
345                metric.range_max,
346                metric.bucket_count as usize,
347            ),
348        };
349        Self { histogram, metric }
350    }
351
352    pub fn accumulate(&mut self, sample: u64) {
360        self.histogram.accumulate(sample)
361    }
362
363    pub fn abandon(mut self) {
365        self.histogram.clear();
366    }
367}
368
369impl Drop for LocalCustomDistribution<'_> {
370    fn drop(&mut self) {
371        if self.histogram.is_empty() {
372            return;
373        }
374
375        let empty = Histogram::_linear(0, 0, 0);
378        let buffer = mem::replace(&mut self.histogram, empty);
379        self.metric.commit_histogram(buffer);
380    }
381}