histogram/
atomic.rs

1use crate::{Config, Error, Histogram};
2use core::sync::atomic::{AtomicU64, Ordering};
3
4/// A histogram that uses atomic 64bit counters for each bucket.
5///
6/// Unlike the non-atomic variant, it cannot be used directly to report
7/// percentiles. Instead, a snapshot must be taken which captures the state of
8/// the histogram at a point in time.
9pub struct AtomicHistogram {
10    config: Config,
11    buckets: Box<[AtomicU64]>,
12}
13
14impl AtomicHistogram {
15    /// Construct a new atomic histogram from the provided parameters. See the
16    /// documentation for [`crate::Config`] to understand their meaning.
17    pub fn new(p: u8, n: u8) -> Result<Self, Error> {
18        let config = Config::new(p, n)?;
19
20        Ok(Self::with_config(&config))
21    }
22
23    /// Creates a new atomic histogram using a provided [`crate::Config`].
24    pub fn with_config(config: &Config) -> Self {
25        let mut buckets = Vec::with_capacity(config.total_buckets());
26        buckets.resize_with(config.total_buckets(), || AtomicU64::new(0));
27
28        Self {
29            config: *config,
30            buckets: buckets.into(),
31        }
32    }
33
34    /// Increment the bucket that contains the value by one.
35    pub fn increment(&self, value: u64) -> Result<(), Error> {
36        self.add(value, 1)
37    }
38
39    /// Increment the bucket that contains the value by some count.
40    pub fn add(&self, value: u64, count: u64) -> Result<(), Error> {
41        let index = self.config.value_to_index(value)?;
42        self.buckets[index].fetch_add(count, Ordering::Relaxed);
43        Ok(())
44    }
45
46    // NOTE: once stabilized, `target_has_atomic_load_store` is more correct. https://github.com/rust-lang/rust/issues/94039
47    #[cfg(target_has_atomic = "64")]
48    /// Drains the bucket values into a new Histogram
49    ///
50    /// Unlike [`load`](AtomicHistogram::load), this method will reset all bucket values to zero. This uses [`AtomicU64::swap`] and is not available
51    /// on platforms where [`AtomicU64::swap`] is not available.
52    pub fn drain(&self) -> Histogram {
53        let buckets: Vec<u64> = self
54            .buckets
55            .iter()
56            .map(|bucket| bucket.swap(0, Ordering::Relaxed))
57            .collect();
58
59        Histogram {
60            config: self.config,
61            buckets: buckets.into(),
62        }
63    }
64
65    /// Read the bucket values into a new `Histogram`
66    pub fn load(&self) -> Histogram {
67        let buckets: Vec<u64> = self
68            .buckets
69            .iter()
70            .map(|bucket| bucket.load(Ordering::Relaxed))
71            .collect();
72
73        Histogram {
74            config: self.config,
75            buckets: buckets.into(),
76        }
77    }
78}
79
80#[cfg(test)]
81mod test {
82    use crate::*;
83
84    #[cfg(target_pointer_width = "64")]
85    #[test]
86    fn size() {
87        assert_eq!(std::mem::size_of::<AtomicHistogram>(), 48);
88    }
89
90    #[cfg(target_has_atomic = "64")]
91    #[test]
92    /// Tests that drain properly resets buckets to 0
93    fn drain() {
94        let histogram = AtomicHistogram::new(7, 64).unwrap();
95        for i in 0..=100 {
96            let _ = histogram.increment(i);
97        }
98        let percentiles = histogram.drain();
99        assert_eq!(
100            percentiles.percentile(50.0),
101            Ok(Some(Bucket {
102                count: 1,
103                range: 50..=50,
104            }))
105        );
106        histogram.increment(1000).unwrap();
107        // after another load the map is empty
108        let percentiles = histogram.drain();
109        assert_eq!(
110            percentiles.percentile(50.0),
111            Ok(Some(Bucket {
112                count: 1,
113                range: 1000..=1003,
114            }))
115        );
116    }
117
118    #[test]
119    // Tests percentiles
120    fn percentiles() {
121        let histogram = AtomicHistogram::new(7, 64).unwrap();
122        let percentiles = [25.0, 50.0, 75.0, 90.0, 99.0];
123
124        // check empty
125        assert_eq!(histogram.load().percentiles(&percentiles), Ok(None));
126
127        for percentile in percentiles {
128            assert_eq!(histogram.load().percentile(percentile), Ok(None));
129        }
130
131        // populate and check percentiles
132        for i in 0..=100 {
133            let _ = histogram.increment(i);
134            assert_eq!(
135                histogram.load().percentile(0.0),
136                Ok(Some(Bucket {
137                    count: 1,
138                    range: 0..=0,
139                }))
140            );
141            assert_eq!(
142                histogram.load().percentile(100.0),
143                Ok(Some(Bucket {
144                    count: 1,
145                    range: i..=i,
146                }))
147            );
148        }
149
150        for percentile in percentiles {
151            assert_eq!(
152                histogram
153                    .load()
154                    .percentile(percentile)
155                    .map(|b| b.unwrap().end()),
156                Ok(percentile as u64)
157            );
158        }
159
160        assert_eq!(
161            histogram.load().percentile(99.9).map(|b| b.unwrap().end()),
162            Ok(100)
163        );
164
165        assert_eq!(
166            histogram.load().percentile(-1.0),
167            Err(Error::InvalidPercentile)
168        );
169        assert_eq!(
170            histogram.load().percentile(101.0),
171            Err(Error::InvalidPercentile)
172        );
173
174        let percentiles: Vec<(f64, u64)> = histogram
175            .load()
176            .percentiles(&[50.0, 90.0, 99.0, 99.9])
177            .unwrap()
178            .unwrap()
179            .iter()
180            .map(|(p, b)| (*p, b.end()))
181            .collect();
182
183        assert_eq!(
184            percentiles,
185            vec![(50.0, 50), (90.0, 90), (99.0, 99), (99.9, 100)]
186        );
187
188        let _ = histogram.increment(1024);
189        assert_eq!(
190            histogram.load().percentile(99.9),
191            Ok(Some(Bucket {
192                count: 1,
193                range: 1024..=1031,
194            }))
195        );
196    }
197}
histogram/atomic.rs

histogram/
atomic.rs
