subms_hdr_histogram/features/

concurrent_writes.rs

//! Lock-free concurrent histogram.
//!
//! Same log-linear bucketing as the base `HdrHistogram` but every
//! counter is an `AtomicU64`. Producers call `record()` from any
//! thread without external synchronisation - the only contention is
//! the per-bucket `fetch_add`. Snapshot reads walk the array with
//! relaxed loads; the value-at-percentile / max / count answers
//! reflect a point in time that may interleave with concurrent
//! writers, but each individual counter is intact.
//!
//! Trade-off vs the base: the array is fixed-size and pre-allocated
//! at construction. A growable layout would need a mutex around the
//! resize, which defeats the lock-free property. Pick the upper bound
//! at construction (number of major buckets) - at 3 sig-digits the
//! default 32 majors is 65536 counters tracking values past 4e12,
//! which is well beyond sub-millisecond latency in nanosecond units.

use crate::{index_of, value_from_index};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};

/// Concurrent HDR histogram. All operations are lock-free.
pub struct ConcurrentHdrHistogram {
    sub_count: u32,
    sub_count_bits: u32,
    counters: Vec<AtomicU64>,
    total: AtomicU64,
    high_index: AtomicUsize,
}

impl ConcurrentHdrHistogram {
    /// New histogram with the given significant-digit precision and
    /// default major-bucket capacity (32, tracking values past 4e12 at
    /// 3 sig-digits).
    pub fn new(significant_digits: u32) -> Self {
        Self::with_majors(significant_digits, 32)
    }

    /// Explicit major-bucket capacity. Counter array length is
    /// `sub_count * majors`. Records that would land past the last
    /// bucket are clamped into the final bucket (no resize possible
    /// in a lock-free design).
    pub fn with_majors(significant_digits: u32, majors: 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;
        let majors = majors.max(1);
        let total_buckets = (sub_count as usize) * (majors as usize);
        let mut counters = Vec::with_capacity(total_buckets);
        for _ in 0..total_buckets {
            counters.push(AtomicU64::new(0));
        }
        Self {
            sub_count,
            sub_count_bits,
            counters,
            total: AtomicU64::new(0),
            high_index: AtomicUsize::new(0),
        }
    }

    pub fn sub_count(&self) -> u32 {
        self.sub_count
    }

    pub fn sub_count_bits(&self) -> u32 {
        self.sub_count_bits
    }

    pub fn count(&self) -> u64 {
        self.total.load(Ordering::Relaxed)
    }

    pub fn max(&self) -> u64 {
        if self.count() == 0 {
            return 0;
        }
        value_from_index(self.high_index.load(Ordering::Relaxed), self.sub_count_bits)
    }

    /// Record a value. Safe from any thread.
    pub fn record(&self, value: u64) {
        let raw = index_of(value, self.sub_count_bits) as usize;
        let idx = raw.min(self.counters.len() - 1);
        self.counters[idx].fetch_add(1, Ordering::Relaxed);
        self.total.fetch_add(1, Ordering::Relaxed);
        // Race on high_index is benign: any thread that wrote a
        // higher index will eventually win the CAS. Worst case the
        // snapshot reader sees a stale low high_index for a few ns.
        let mut cur = self.high_index.load(Ordering::Relaxed);
        while idx > cur {
            match self.high_index.compare_exchange_weak(
                cur,
                idx,
                Ordering::Relaxed,
                Ordering::Relaxed,
            ) {
                Ok(_) => break,
                Err(seen) => cur = seen,
            }
        }
    }

    /// Snapshot-style percentile read. Walks the array with relaxed
    /// loads; not linearisable across all concurrent writers but
    /// each individual counter read is intact.
    pub fn value_at_percentile(&self, q: f64) -> u64 {
        let total = self.count();
        if total == 0 {
            return 0;
        }
        let target = ((q.clamp(0.0, 1.0) * total as f64) as u64).max(1);
        let high = self.high_index.load(Ordering::Relaxed);
        let mut cum = 0u64;
        let end = (high + 1).min(self.counters.len());
        for i in 0..end {
            cum += self.counters[i].load(Ordering::Relaxed);
            if cum >= target {
                return value_from_index(i, self.sub_count_bits);
            }
        }
        value_from_index(high, self.sub_count_bits)
    }

    /// Atomically drain every counter and total into a `Snapshot`,
    /// leaving the histogram empty. Used by `DualRecorder` to harvest
    /// the inactive side. Each per-counter swap is independent, so
    /// concurrent writers may land their increments in EITHER the
    /// drained snapshot OR the now-zeroed live histogram - we never
    /// double-count or lose a write.
    pub fn drain_snapshot(&self) -> Snapshot {
        let mut counts = Vec::with_capacity(self.counters.len());
        let high = self.high_index.load(Ordering::Relaxed);
        let len = (high + 1).min(self.counters.len());
        let mut total = 0u64;
        for i in 0..len {
            let v = self.counters[i].swap(0, Ordering::AcqRel);
            counts.push(v);
            total += v;
        }
        // Subtract what we drained from `total`. Concurrent
        // record()s that fired during the loop will keep `total`
        // ahead of the per-counter sum; that's correct for the
        // live side.
        self.total.fetch_sub(total, Ordering::AcqRel);
        self.high_index.store(0, Ordering::Relaxed);
        Snapshot {
            sub_count_bits: self.sub_count_bits,
            counts,
            total,
        }
    }
}

/// Frozen view of a histogram at the moment of `drain_snapshot()`.
/// Has the same percentile / max APIs as the live histogram but is
/// immutable.
pub struct Snapshot {
    sub_count_bits: u32,
    counts: Vec<u64>,
    total: u64,
}

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

    pub fn max(&self) -> u64 {
        if self.total == 0 {
            return 0;
        }
        // The drained snapshot truncates at high_index, so the last
        // non-zero counter is at most counts.len()-1.
        for i in (0..self.counts.len()).rev() {
            if self.counts[i] > 0 {
                return value_from_index(i, self.sub_count_bits);
            }
        }
        0
    }

    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;
        // `total` is the sum of `counts` by construction (drain_snapshot adds
        // each swapped counter), and target <= total, so the loop always
        // returns; the trailing 0 is only there to satisfy the compiler.
        for (i, &c) in self.counts.iter().enumerate() {
            cum += c;
            if cum >= target {
                return value_from_index(i, self.sub_count_bits);
            }
        }
        0
    }
}

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