glean_core/histogram/
mod.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4
5//! A simple histogram implementation for exponential histograms.
6
7use std::any::TypeId;
8use std::collections::HashMap;
9
10use malloc_size_of_derive::MallocSizeOf;
11use once_cell::sync::OnceCell;
12use serde::{Deserialize, Serialize};
13
14use crate::error::{Error, ErrorKind};
15
16pub use exponential::PrecomputedExponential;
17pub use functional::Functional;
18pub use linear::PrecomputedLinear;
19
20mod exponential;
21mod functional;
22mod linear;
23
24/// Different kinds of histograms.
25#[derive(Debug, Clone, Copy, Serialize, Deserialize, MallocSizeOf)]
26#[serde(rename_all = "lowercase")]
27pub enum HistogramType {
28    /// A histogram with linear distributed buckets.
29    Linear,
30    /// A histogram with exponential distributed buckets.
31    Exponential,
32}
33
34impl TryFrom<i32> for HistogramType {
35    type Error = Error;
36
37    fn try_from(value: i32) -> Result<HistogramType, Self::Error> {
38        match value {
39            0 => Ok(HistogramType::Linear),
40            1 => Ok(HistogramType::Exponential),
41            e => Err(ErrorKind::HistogramType(e).into()),
42        }
43    }
44}
45
46/// A histogram.
47///
48/// Stores the counts per bucket and tracks the count of added samples and the total sum.
49/// The bucketing algorithm can be changed.
50///
51/// ## Example
52///
53/// ```rust,ignore
54/// let mut hist = Histogram::exponential(1, 500, 10);
55///
56/// for i in 1..=10 {
57///     hist.accumulate(i);
58/// }
59///
60/// assert_eq!(10, hist.count());
61/// assert_eq!(55, hist.sum());
62/// ```
63#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, MallocSizeOf)]
64pub struct Histogram<B> {
65    /// Mapping bucket's minimum to sample count.
66    values: HashMap<u64, u64>,
67
68    /// The count of samples added.
69    count: u64,
70    /// The total sum of samples.
71    sum: u64,
72
73    /// The bucketing algorithm used.
74    bucketing: B,
75}
76
77/// A bucketing algorithm for histograms.
78///
79/// It's responsible to calculate the bucket a sample goes into.
80/// It can calculate buckets on-the-fly or pre-calculate buckets and re-use that when needed.
81pub trait Bucketing {
82    /// Get the bucket's minimum value the sample falls into.
83    fn sample_to_bucket_minimum(&self, sample: u64) -> u64;
84
85    /// The computed bucket ranges for this bucketing algorithm.
86    fn ranges(&self) -> &[u64];
87}
88
89impl<B: Bucketing> Histogram<B> {
90    /// Gets the number of buckets in this histogram.
91    pub fn bucket_count(&self) -> usize {
92        self.values.len()
93    }
94
95    /// Adds a single value to this histogram.
96    pub fn accumulate(&mut self, sample: u64) {
97        let bucket_min = self.bucketing.sample_to_bucket_minimum(sample);
98        let entry = self.values.entry(bucket_min).or_insert(0);
99        *entry += 1;
100        self.sum = self.sum.saturating_add(sample);
101        self.count += 1;
102    }
103
104    /// Gets the total sum of values recorded in this histogram.
105    pub fn sum(&self) -> u64 {
106        self.sum
107    }
108
109    /// Gets the total count of values recorded in this histogram.
110    pub fn count(&self) -> u64 {
111        self.count
112    }
113
114    /// Gets the filled values.
115    pub fn values(&self) -> &HashMap<u64, u64> {
116        &self.values
117    }
118
119    /// Checks if this histogram recorded any values.
120    pub fn is_empty(&self) -> bool {
121        self.count() == 0
122    }
123
124    /// Gets a snapshot of all values from the first bucket until one past the last filled bucket,
125    /// filling in empty buckets with 0.
126    pub fn snapshot_values(&self) -> HashMap<u64, u64> {
127        self.values.clone()
128    }
129
130    /// Clear this histogram.
131    pub fn clear(&mut self) {
132        self.sum = 0;
133        self.count = 0;
134        self.values.clear();
135    }
136}
137
138/// Either linear or exponential histogram bucketing
139///
140/// This is to be used as a single type to avoid generic use in the buffered API.
141pub enum LinearOrExponential {
142    Linear(PrecomputedLinear),
143    Exponential(PrecomputedExponential),
144}
145
146impl Histogram<LinearOrExponential> {
147    /// A histogram using linear bucketing.
148    ///
149    /// _Note:_ Special naming to avoid needing to use extensive type annotations in other parts.
150    /// This type is only used for the buffered API.
151    pub fn _linear(min: u64, max: u64, bucket_count: usize) -> Histogram<LinearOrExponential> {
152        Histogram {
153            values: HashMap::new(),
154            count: 0,
155            sum: 0,
156            bucketing: LinearOrExponential::Linear(PrecomputedLinear {
157                bucket_ranges: OnceCell::new(),
158                min,
159                max,
160                bucket_count,
161            }),
162        }
163    }
164
165    /// A histogram using expontential bucketing.
166    ///
167    /// _Note:_ Special naming to avoid needing to use extensive type annotations in other parts.
168    /// This type is only used for the buffered API.
169    pub fn _exponential(min: u64, max: u64, bucket_count: usize) -> Histogram<LinearOrExponential> {
170        Histogram {
171            values: HashMap::new(),
172            count: 0,
173            sum: 0,
174            bucketing: LinearOrExponential::Exponential(PrecomputedExponential {
175                bucket_ranges: OnceCell::new(),
176                min,
177                max,
178                bucket_count,
179            }),
180        }
181    }
182}
183
184impl Bucketing for LinearOrExponential {
185    fn sample_to_bucket_minimum(&self, sample: u64) -> u64 {
186        use LinearOrExponential::*;
187        match self {
188            Linear(lin) => lin.sample_to_bucket_minimum(sample),
189            Exponential(exp) => exp.sample_to_bucket_minimum(sample),
190        }
191    }
192
193    fn ranges(&self) -> &[u64] {
194        use LinearOrExponential::*;
195        match self {
196            Linear(lin) => lin.ranges(),
197            Exponential(exp) => exp.ranges(),
198        }
199    }
200}
201
202impl<B> Histogram<B>
203where
204    B: Bucketing,
205    B: std::fmt::Debug,
206    B: PartialEq,
207{
208    /// Merges data from one histogram into the other.
209    ///
210    /// ## Panics
211    ///
212    /// Panics if the two histograms don't use the same bucketing.
213    pub fn merge(&mut self, other: &Self) {
214        assert_eq!(self.bucketing, other.bucketing);
215
216        self.sum = self.sum.saturating_add(other.sum);
217        self.count = self.count.saturating_add(other.count);
218        for (&bucket, &count) in &other.values {
219            let entry = self.values.entry(bucket).or_insert(0);
220            *entry = entry.saturating_add(count)
221        }
222    }
223}
224
225impl<B> Histogram<B>
226where
227    B: Bucketing + 'static,
228    B: std::fmt::Debug,
229    B: PartialEq,
230{
231    /// Merges data from one histogram into the other.
232    ///
233    /// ## Panics
234    ///
235    /// Panics if the two histograms don't use the same bucketing.
236    /// Note that the `other` side can be either linear or exponential
237    /// and we only merge if it matches `self`'s bucketing.
238    // _Note:_ Unfortunately this needs a separate name from the above, otherwise it's a conflicting
239    // method.
240    // We only use it internally for the buffered API, and can guarantee correct usage that way.
241    pub fn _merge(&mut self, other: &Histogram<LinearOrExponential>) {
242        #[rustfmt::skip]
243        assert!(
244            (
245                TypeId::of::<B>() == TypeId::of::<PrecomputedLinear>()
246                && matches!(other.bucketing, LinearOrExponential::Linear(_))
247            ) ||
248            (
249                TypeId::of::<B>() == TypeId::of::<PrecomputedExponential>()
250                && matches!(other.bucketing, LinearOrExponential::Exponential(_))
251            )
252        );
253        self.sum = self.sum.saturating_add(other.sum);
254        self.count = self.count.saturating_add(other.count);
255        for (&bucket, &count) in &other.values {
256            let entry = self.values.entry(bucket).or_insert(0);
257            *entry = entry.saturating_add(count);
258        }
259    }
260}