use crate::Histogram;
use crate::label::LabelEnum;
use std::marker::PhantomData;
pub struct LabeledHistogram<L: LabelEnum> {
histograms: Vec<Histogram>,
_phantom: PhantomData<L>,
}
impl<L: LabelEnum> LabeledHistogram<L> {
pub fn new(bounds: &[u64], shard_count: usize) -> Self {
let histograms = (0..L::CARDINALITY)
.map(|_| Histogram::new(bounds, shard_count))
.collect();
Self {
histograms,
_phantom: PhantomData,
}
}
pub fn with_latency_buckets(shard_count: usize) -> Self {
let histograms = (0..L::CARDINALITY)
.map(|_| Histogram::with_latency_buckets(shard_count))
.collect();
Self {
histograms,
_phantom: PhantomData,
}
}
#[inline]
pub fn record(&self, label: L, value: u64) {
let idx = label.as_index();
debug_assert!(idx < self.histograms.len(), "label index out of bounds");
if cfg!(debug_assertions) {
self.histograms[idx].record(value);
} else {
unsafe { self.histograms.get_unchecked(idx) }.record(value);
}
}
pub fn get(&self, label: L) -> &Histogram {
let idx = label.as_index();
if cfg!(debug_assertions) {
&self.histograms[idx]
} else {
unsafe { self.histograms.get_unchecked(idx) }
}
}
pub fn iter(&self) -> impl Iterator<Item = (L, &'_ Histogram)> + '_ {
self.histograms
.iter()
.enumerate()
.map(|(idx, histogram)| (L::from_index(idx), histogram))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[derive(Copy, Clone, Debug, PartialEq)]
enum TestLabel {
Fast,
Medium,
Slow,
}
impl LabelEnum for TestLabel {
const CARDINALITY: usize = 3;
const LABEL_NAME: &'static str = "speed";
fn as_index(self) -> usize {
self as usize
}
fn from_index(index: usize) -> Self {
match index {
0 => Self::Fast,
1 => Self::Medium,
_ => Self::Slow,
}
}
fn variant_name(self) -> &'static str {
match self {
Self::Fast => "fast",
Self::Medium => "medium",
Self::Slow => "slow",
}
}
}
#[test]
fn test_basic_recording() {
let histogram: LabeledHistogram<TestLabel> = LabeledHistogram::new(&[10, 100, 1000], 4);
histogram.record(TestLabel::Fast, 5);
histogram.record(TestLabel::Fast, 8);
histogram.record(TestLabel::Medium, 50);
histogram.record(TestLabel::Slow, 5000);
assert_eq!(histogram.get(TestLabel::Fast).count(), 2);
assert_eq!(histogram.get(TestLabel::Medium).count(), 1);
assert_eq!(histogram.get(TestLabel::Slow).count(), 1);
}
#[test]
fn test_iteration() {
let histogram: LabeledHistogram<TestLabel> = LabeledHistogram::new(&[100], 4);
histogram.record(TestLabel::Fast, 50);
histogram.record(TestLabel::Medium, 150);
let data: Vec<_> = histogram.iter().collect();
assert_eq!(data.len(), 3);
assert_eq!(data[0].0, TestLabel::Fast);
assert_eq!(data[0].1.count(), 1);
assert_eq!(data[1].0, TestLabel::Medium);
assert_eq!(data[1].1.count(), 1);
assert_eq!(data[2].0, TestLabel::Slow);
assert_eq!(data[2].1.count(), 0);
}
#[test]
fn test_latency_buckets() {
let histogram: LabeledHistogram<TestLabel> = LabeledHistogram::with_latency_buckets(4);
histogram.record(TestLabel::Fast, 5); histogram.record(TestLabel::Fast, 100); histogram.record(TestLabel::Medium, 1_000);
assert_eq!(histogram.get(TestLabel::Fast).count(), 2);
assert_eq!(histogram.get(TestLabel::Medium).count(), 1);
}
#[test]
fn test_concurrent_recording() {
let histogram: LabeledHistogram<TestLabel> = LabeledHistogram::new(&[100], 4);
std::thread::scope(|s| {
for _ in 0..4 {
s.spawn(|| {
for i in 0..1000 {
histogram.record(TestLabel::Fast, i);
}
});
}
});
assert_eq!(histogram.get(TestLabel::Fast).count(), 4000);
}
}