subms-hdr-histogram 0.10.0

//! Histogram with per-recording 1-byte tags.
//!
//! Useful when you want to slice latencies by category at query
//! time (request method, endpoint, tenant tier) without paying for
//! N separate histograms.
//!
//! Storage: each base bucket carries a small per-tag map. The total
//! count is the sum across all tags. Per-tag percentiles iterate
//! over buckets and accumulate only the matching tag's count.
//!
//! Trade-off vs N separate histograms:
//! - Cheaper when the tag cardinality is small but the value range
//!   is wide (one bucket array shared across tags).
//! - More expensive per-tag percentile read (must walk all buckets).
//! - Constant overhead per untagged write is one extra map lookup.

use crate::{index_of, value_from_index};
use std::collections::HashMap;

/// Histogram with per-bucket tag-keyed sub-counters.
pub struct TaggedHdrHistogram {
    sub_count_bits: u32,
    /// Outer: bucket index. Inner: tag -> count.
    buckets: Vec<HashMap<u8, u64>>,
    total: u64,
    high_index: usize,
    /// Per-tag running totals so `count_for_tag()` is O(1).
    per_tag_total: HashMap<u8, u64>,
}

impl TaggedHdrHistogram {
    /// New tagged histogram with the given significant-digit precision.
    pub fn new(significant_digits: u32) -> Self {
        let sig = significant_digits.clamp(1, 5);
        let target = 2u32 * 10u32.pow(sig);
        let sub_count_bits = (32 - target.leading_zeros()).max(1);
        let sub_count = 1u32 << sub_count_bits;
        Self {
            sub_count_bits,
            buckets: (0..sub_count).map(|_| HashMap::new()).collect(),
            total: 0,
            high_index: 0,
            per_tag_total: HashMap::new(),
        }
    }

    /// Record a value tagged with the given byte.
    pub fn record(&mut self, value: u64, tag: u8) {
        let idx = index_of(value, self.sub_count_bits) as usize;
        if idx >= self.buckets.len() {
            self.buckets.resize_with(idx + 1, HashMap::new);
        }
        *self.buckets[idx].entry(tag).or_insert(0) += 1;
        self.total += 1;
        if idx > self.high_index {
            self.high_index = idx;
        }
        *self.per_tag_total.entry(tag).or_insert(0) += 1;
    }

    pub fn count(&self) -> u64 {
        self.total
    }

    pub fn count_for_tag(&self, tag: u8) -> u64 {
        self.per_tag_total.get(&tag).copied().unwrap_or(0)
    }

    pub fn max(&self) -> u64 {
        if self.total == 0 {
            return 0;
        }
        value_from_index(self.high_index, self.sub_count_bits)
    }

    /// Quantile across all tags (matches the base histogram's
    /// behaviour).
    pub fn value_at_percentile(&self, q: f64) -> u64 {
        if self.total == 0 {
            return 0;
        }
        let target = ((q.clamp(0.0, 1.0) * self.total as f64) as u64).max(1);
        let mut cum = 0u64;
        let end = (self.high_index + 1).min(self.buckets.len());
        for i in 0..end {
            let bucket_total: u64 = self.buckets[i].values().sum();
            cum += bucket_total;
            if cum >= target {
                return value_from_index(i, self.sub_count_bits);
            }
        }
        value_from_index(self.high_index, self.sub_count_bits)
    }

    /// Quantile restricted to one tag.
    pub fn value_at_percentile_for_tag(&self, q: f64, tag: u8) -> u64 {
        let tag_total = self.count_for_tag(tag);
        if tag_total == 0 {
            return 0;
        }
        let target = ((q.clamp(0.0, 1.0) * tag_total as f64) as u64).max(1);
        let mut cum = 0u64;
        let end = (self.high_index + 1).min(self.buckets.len());
        for i in 0..end {
            if let Some(&c) = self.buckets[i].get(&tag) {
                cum += c;
                if cum >= target {
                    return value_from_index(i, self.sub_count_bits);
                }
            }
        }
        value_from_index(self.high_index, self.sub_count_bits)
    }

    /// List of tags seen in any recording, in unspecified order.
    pub fn tags(&self) -> Vec<u8> {
        self.per_tag_total.keys().copied().collect()
    }
}

#[cfg(test)]
#[path = "value_tagging_tests.rs"]
mod tests;