Skip to main content

subetha_cxc/
shared_histogram.rs

1//! `SharedHistogram` - cross-process bucketed counter for
2//! distribution tracking.
3//!
4//! Fixed-bucket histogram: caller supplies N bucket boundaries at
5//! create time. `record(value)` finds the right bucket via binary
6//! search and atomically increments its counter. Useful for
7//! latency distributions, request-size distributions, queue-depth
8//! sampling - anything where N distributed processes need to
9//! aggregate "how many in each bucket" into one shared view.
10//!
11//! # Bucket semantics
12//!
13//! For boundaries `[b0, b1, b2, ..., bN-1]`:
14//! - Bucket 0: values `value < b0`
15//! - Bucket i (1..N-1): values `b{i-1} <= value < bi`
16//! - Bucket N: values `value >= b{N-1}` (the overflow bucket)
17//!
18//! So a histogram with K boundaries has K+1 buckets.
19//!
20//! # Layout
21//!
22//! Single MMF file:
23//!
24//! ```text
25//! +---------------------------+
26//! | HistogramHeader (64B)     |
27//! |   magic, n_boundaries     |
28//! |   total_count: AtomicU64  |
29//! +---------------------------+
30//! | boundaries [u64; N]       |  ascending; verified at open
31//! +---------------------------+
32//! | counters [AtomicU64; N+1] |  one per bucket
33//! +---------------------------+
34//! ```
35//!
36//! # Concurrency
37//!
38//! Each bucket's counter is its own AtomicU64. `record` uses
39//! `fetch_add(1, AcqRel)` to atomically increment; multiple
40//! recorders contend only on the SAME bucket's cache line
41//! (different buckets are fully concurrent).
42//!
43//! # Percentile estimation
44//!
45//! `percentile(p)` walks buckets accumulating counts until p of the
46//! total is covered, then linearly interpolates within the target
47//! bucket. For coarse boundaries the estimate has bucket-width
48//! granularity; for log-spaced boundaries that's typically <1
49//! decade error which suffices for latency dashboards.
50
51use std::fs::{File, OpenOptions};
52use std::mem::size_of;
53use std::path::Path;
54use std::sync::atomic::{AtomicU64, Ordering};
55
56use memmap2::{MmapMut, MmapOptions};
57
58pub const HISTOGRAM_MAGIC: u64 = 0x4150_4849_5354_3031;
59
60#[repr(C, align(64))]
61pub struct HistogramHeader {
62    pub magic: u64,
63    pub n_boundaries: u32,
64    _pad1: u32,
65    pub total_count: AtomicU64,
66    _pad2: [u8; 40],
67}
68
69const _: () = {
70    assert!(size_of::<HistogramHeader>() == 64);
71};
72
73#[derive(Debug, Clone, Copy, PartialEq, Eq)]
74pub enum HistogramError {
75    EmptyBoundaries,
76    NonMonotonicBoundaries,
77    LayoutMismatch,
78    OutOfBounds,
79    IoError(std::io::ErrorKind),
80}
81
82impl From<std::io::Error> for HistogramError {
83    fn from(e: std::io::Error) -> Self { Self::IoError(e.kind()) }
84}
85
86pub const fn histogram_file_size(n_boundaries: usize) -> usize {
87    size_of::<HistogramHeader>()
88        + n_boundaries * size_of::<u64>()
89        + (n_boundaries + 1) * size_of::<AtomicU64>()
90}
91
92pub struct SharedHistogram {
93    _file: File,
94    mmap: MmapMut,
95    n_boundaries: usize,
96    boundaries_offset: usize,
97    counters_offset: usize,
98    header_sidecar: subetha_core::HandshakeHeader,
99    ring_sidecar: Box<subetha_core::ObservationRing>,
100}
101
102unsafe impl Send for SharedHistogram {}
103unsafe impl Sync for SharedHistogram {}
104
105impl subetha_sidecar::AdaptiveInstance for SharedHistogram {
106    fn header(&self) -> &subetha_core::HandshakeHeader { &self.header_sidecar }
107    fn ring(&self) -> &subetha_core::ObservationRing { &self.ring_sidecar }
108    fn make_policy(&self) -> Box<dyn subetha_sidecar::Policy> {
109        Box::new(subetha_sidecar::NoMigrationPolicy)
110    }
111}
112
113impl SharedHistogram {
114    /// Obtain the histogram at `path`, initializing empty buckets if
115    /// the path does not yet exist and attaching to it if it does.
116    /// Attaching leaves live counts in place; a region built with
117    /// different boundaries is a `LayoutMismatch`.
118    /// [`reset`](Self::reset) zeroes a live histogram in place.
119    pub fn create(
120        path: impl AsRef<Path>, boundaries: &[u64],
121    ) -> Result<Self, HistogramError> {
122        Self::check_boundaries(boundaries)?;
123        let (file, mmap) = crate::mmf_attach::create_or_attach(
124            path.as_ref(),
125            histogram_file_size(boundaries.len()),
126            |ptr| unsafe { Self::init_region(ptr, boundaries) },
127            |ptr| unsafe { (*(ptr as *const HistogramHeader)).magic == HISTOGRAM_MAGIC },
128        )?;
129        Self::from_region(file, mmap, boundaries)
130    }
131
132    fn check_boundaries(boundaries: &[u64]) -> Result<(), HistogramError> {
133        if boundaries.is_empty() {
134            return Err(HistogramError::EmptyBoundaries);
135        }
136        for w in boundaries.windows(2) {
137            if w[0] >= w[1] {
138                return Err(HistogramError::NonMonotonicBoundaries);
139            }
140        }
141        Ok(())
142    }
143
144    /// Lay out empty buckets: boundary count and the boundary array
145    /// first, magic last, because attachers spin on it. The zeroed
146    /// region is already the zero counters and zero total.
147    ///
148    /// # Safety
149    /// `ptr` addresses at least `histogram_file_size(boundaries.len())`
150    /// writable zeroed bytes.
151    unsafe fn init_region(ptr: *mut u8, boundaries: &[u64]) {
152        let hdr = ptr as *mut HistogramHeader;
153        unsafe {
154            (*hdr).n_boundaries = boundaries.len() as u32;
155            let dst = ptr.add(size_of::<HistogramHeader>()) as *mut u64;
156            std::ptr::copy_nonoverlapping(boundaries.as_ptr(), dst, boundaries.len());
157            std::ptr::write_volatile(&raw mut (*hdr).magic, HISTOGRAM_MAGIC);
158        }
159    }
160
161    /// Wrap an initialized region, refusing one built with different
162    /// boundaries.
163    fn from_region(
164        file: File,
165        mmap: MmapMut,
166        boundaries: &[u64],
167    ) -> Result<Self, HistogramError> {
168        let n_boundaries = boundaries.len();
169        let hdr = unsafe { &*(mmap.as_ptr() as *const HistogramHeader) };
170        if hdr.magic != HISTOGRAM_MAGIC || hdr.n_boundaries != n_boundaries as u32 {
171            return Err(HistogramError::LayoutMismatch);
172        }
173        let boundaries_offset = size_of::<HistogramHeader>();
174        let counters_offset = boundaries_offset + std::mem::size_of_val(boundaries);
175        let stored = unsafe {
176            std::slice::from_raw_parts(
177                mmap.as_ptr().add(boundaries_offset) as *const u64,
178                n_boundaries,
179            )
180        };
181        if stored != boundaries {
182            return Err(HistogramError::LayoutMismatch);
183        }
184        Ok(Self {
185            _file: file, mmap, n_boundaries,
186            boundaries_offset, counters_offset,
187            header_sidecar: subetha_core::HandshakeHeader::new(),
188            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
189        })
190    }
191
192    pub fn open(
193        path: impl AsRef<Path>, expected_boundaries: &[u64],
194    ) -> Result<Self, HistogramError> {
195        let n_boundaries = expected_boundaries.len();
196        if n_boundaries == 0 { return Err(HistogramError::EmptyBoundaries); }
197        let total = histogram_file_size(n_boundaries);
198        let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
199        if file.metadata()?.len() < total as u64 {
200            return Err(HistogramError::LayoutMismatch);
201        }
202        let mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
203        Self::from_region(file, mmap, expected_boundaries)
204    }
205
206    fn header(&self) -> &HistogramHeader {
207        unsafe { &*(self.mmap.as_ptr() as *const HistogramHeader) }
208    }
209
210    fn boundaries(&self) -> &[u64] {
211        unsafe {
212            std::slice::from_raw_parts(
213                self.mmap.as_ptr().add(self.boundaries_offset) as *const u64,
214                self.n_boundaries,
215            )
216        }
217    }
218
219    fn counter(&self, bucket_idx: usize) -> &AtomicU64 {
220        let base = unsafe { self.mmap.as_ptr().add(self.counters_offset) };
221        unsafe { &*(base.add(bucket_idx * size_of::<AtomicU64>()) as *const AtomicU64) }
222    }
223
224    /// Number of buckets (n_boundaries + 1).
225    pub fn n_buckets(&self) -> usize { self.n_boundaries + 1 }
226
227    /// Total count across all buckets.
228    pub fn total_count(&self) -> u64 {
229        self.header().total_count.load(Ordering::Acquire)
230    }
231
232    /// Find the bucket index for `value`. Binary search on boundaries.
233    pub fn bucket_for(&self, value: u64) -> usize {
234        let bounds = self.boundaries();
235        // partition_point returns the first index where the predicate
236        // is false. With `|&b| b <= value`, it returns the first
237        // boundary GREATER than value, i.e., the bucket index.
238        bounds.partition_point(|&b| b <= value)
239    }
240
241    /// Record one observation of `value`. Atomically increments the
242    /// matching bucket counter and the total. Returns the bucket
243    /// index it landed in.
244    pub fn record(&self, value: u64) -> usize {
245        let idx = self.bucket_for(value);
246        self.counter(idx).fetch_add(1, Ordering::AcqRel);
247        self.header().total_count.fetch_add(1, Ordering::AcqRel);
248        self.ring_sidecar
249            .push_op(crate::sidecar_ops::histogram::OP_RECORD, 0);
250        idx
251    }
252
253    /// Read a specific bucket's count.
254    pub fn count(&self, bucket_idx: usize) -> Result<u64, HistogramError> {
255        if bucket_idx >= self.n_buckets() {
256            self.ring_sidecar
257                .push_op(crate::sidecar_ops::histogram::OP_COUNT, 1);
258            return Err(HistogramError::OutOfBounds);
259        }
260        let v = self.counter(bucket_idx).load(Ordering::Acquire);
261        self.ring_sidecar
262            .push_op(crate::sidecar_ops::histogram::OP_COUNT, 0);
263        Ok(v)
264    }
265
266    /// Snapshot all bucket counts as a Vec.
267    pub fn counts(&self) -> Vec<u64> {
268        (0..self.n_buckets())
269            .map(|i| self.counter(i).load(Ordering::Acquire))
270            .collect()
271    }
272
273    /// Get the boundaries vector (copy).
274    pub fn boundaries_vec(&self) -> Vec<u64> {
275        self.boundaries().to_vec()
276    }
277
278    /// Estimate the p-th percentile (p in 0.0..=1.0). Walks buckets
279    /// accumulating counts until p of total is covered, then linearly
280    /// interpolates within the target bucket. Returns 0 if total is 0.
281    pub fn percentile(&self, p: f64) -> u64 {
282        let p = p.clamp(0.0, 1.0);
283        let total = self.total_count();
284        self.ring_sidecar.push_op(
285            crate::sidecar_ops::histogram::OP_PERCENTILE,
286            if total == 0 { 2 } else { 0 },
287        );
288        if total == 0 { return 0; }
289        let target = (total as f64 * p).round() as u64;
290        let mut acc = 0u64;
291        let counts = self.counts();
292        let bounds = self.boundaries();
293        for (i, &c) in counts.iter().enumerate() {
294            let new_acc = acc.saturating_add(c);
295            if new_acc >= target {
296                // Target falls in bucket i. Interpolate within.
297                let lo = if i == 0 { 0 } else { bounds[i - 1] };
298                let hi = if i < bounds.len() { bounds[i] } else { lo.saturating_mul(2) };
299                if c == 0 { return lo; }
300                let frac = (target - acc) as f64 / c as f64;
301                return lo + ((hi - lo) as f64 * frac) as u64;
302            }
303            acc = new_acc;
304        }
305        // Shouldn't reach; return last boundary as fallback.
306        bounds.last().copied().unwrap_or(0)
307    }
308
309    /// Reset all counters to 0. Not concurrency-coordinated; expect
310    /// transient race with concurrent recorders.
311    pub fn reset(&self) {
312        for i in 0..self.n_buckets() {
313            self.counter(i).store(0, Ordering::Release);
314        }
315        self.header().total_count.store(0, Ordering::Release);
316    }
317
318    pub fn flush(&self) -> Result<(), HistogramError> {
319        self.mmap.flush()?;
320        Ok(())
321    }
322
323    pub fn flush_async(&self) -> Result<(), HistogramError> {
324        self.mmap.flush_async()?;
325        Ok(())
326    }
327}
328
329#[cfg(test)]
330mod tests {
331    use super::*;
332    use std::sync::Arc;
333    use std::thread;
334
335    fn tmp(name: &str) -> std::path::PathBuf {
336        let mut p = std::env::temp_dir();
337        let pid = std::process::id();
338        p.push(format!("subetha-histogram-{name}-{pid}.bin"));
339        p
340    }
341
342    #[test]
343    fn create_initial_state_is_empty() {
344        let p = tmp("init");
345        let h = SharedHistogram::create(&p, &[10, 100, 1000]).unwrap();
346        assert_eq!(h.n_buckets(), 4);  // 3 boundaries -> 4 buckets
347        assert_eq!(h.total_count(), 0);
348        for i in 0..h.n_buckets() {
349            assert_eq!(h.count(i).unwrap(), 0);
350        }
351        std::fs::remove_file(&p).ok();
352    }
353
354    /// A second create attaches with live counts in place; the
355    /// in-place reset is what zeroes them.
356    #[test]
357    fn second_create_attaches_and_keeps_counts() {
358        let p = tmp("attach");
359        std::fs::remove_file(&p).ok();
360        let h = SharedHistogram::create(&p, &[10, 100]).unwrap();
361        h.record(50);
362        h.record(5);
363
364        let h2 = SharedHistogram::create(&p, &[10, 100]).unwrap();
365        assert_eq!(h2.total_count(), 2, "attach zeroed live counts");
366        assert!(matches!(
367            SharedHistogram::create(&p, &[10, 999]),
368            Err(HistogramError::LayoutMismatch),
369        ));
370
371        h2.reset();
372        assert_eq!(h.total_count(), 0, "reset did not zero for every handle");
373        drop(h);
374        drop(h2);
375        std::fs::remove_file(&p).ok();
376    }
377
378    #[test]
379    fn empty_boundaries_rejected() {
380        let p = tmp("empty");
381        assert_eq!(
382            SharedHistogram::create(&p, &[]).err(),
383            Some(HistogramError::EmptyBoundaries)
384        );
385        std::fs::remove_file(&p).ok();
386    }
387
388    #[test]
389    fn non_monotonic_boundaries_rejected() {
390        let p = tmp("non-mono");
391        assert_eq!(
392            SharedHistogram::create(&p, &[10, 5, 20]).err(),
393            Some(HistogramError::NonMonotonicBoundaries)
394        );
395        assert_eq!(
396            SharedHistogram::create(&p, &[10, 10]).err(),
397            Some(HistogramError::NonMonotonicBoundaries)
398        );
399        std::fs::remove_file(&p).ok();
400    }
401
402    #[test]
403    fn bucket_assignment_correct() {
404        let p = tmp("bucket-assign");
405        // Boundaries: [10, 100, 1000]
406        // Bucket 0: value < 10
407        // Bucket 1: 10 <= value < 100
408        // Bucket 2: 100 <= value < 1000
409        // Bucket 3: value >= 1000
410        let h = SharedHistogram::create(&p, &[10, 100, 1000]).unwrap();
411        assert_eq!(h.bucket_for(0), 0);
412        assert_eq!(h.bucket_for(9), 0);
413        assert_eq!(h.bucket_for(10), 1);
414        assert_eq!(h.bucket_for(99), 1);
415        assert_eq!(h.bucket_for(100), 2);
416        assert_eq!(h.bucket_for(999), 2);
417        assert_eq!(h.bucket_for(1000), 3);
418        assert_eq!(h.bucket_for(u64::MAX), 3);
419        std::fs::remove_file(&p).ok();
420    }
421
422    #[test]
423    fn record_increments_correct_bucket() {
424        let p = tmp("record");
425        let h = SharedHistogram::create(&p, &[10, 100, 1000]).unwrap();
426        let inputs = [5, 5, 50, 500, 5000, 50_000];
427        for &v in &inputs {
428            h.record(v);
429        }
430        // Bucket 0 (< 10): 2 (the two 5s)
431        // Bucket 1 (10..100): 1 (the 50)
432        // Bucket 2 (100..1000): 1 (the 500)
433        // Bucket 3 (>= 1000): 2 (5000 and 50_000)
434        assert_eq!(h.count(0).unwrap(), 2);
435        assert_eq!(h.count(1).unwrap(), 1);
436        assert_eq!(h.count(2).unwrap(), 1);
437        assert_eq!(h.count(3).unwrap(), 2);
438        assert_eq!(h.total_count(), 6);
439        std::fs::remove_file(&p).ok();
440    }
441
442    #[test]
443    fn record_returns_bucket_index() {
444        let p = tmp("record-idx");
445        let h = SharedHistogram::create(&p, &[10, 100]).unwrap();
446        assert_eq!(h.record(5), 0);
447        assert_eq!(h.record(50), 1);
448        assert_eq!(h.record(500), 2);
449        std::fs::remove_file(&p).ok();
450    }
451
452    #[test]
453    fn counts_snapshot_returns_all_buckets() {
454        let p = tmp("counts");
455        let h = SharedHistogram::create(&p, &[10, 100]).unwrap();
456        h.record(5);
457        h.record(50);
458        h.record(50);
459        h.record(500);
460        h.record(500);
461        h.record(500);
462        let counts = h.counts();
463        assert_eq!(counts, vec![1, 2, 3]);
464        std::fs::remove_file(&p).ok();
465    }
466
467    #[test]
468    fn percentile_basic() {
469        let p = tmp("p-basic");
470        let h = SharedHistogram::create(&p, &[10, 100, 1000]).unwrap();
471        // 100 values uniformly in 0..10. All land in bucket 0.
472        for v in 0..100u64 { h.record(v % 10); }
473        // p50 should be ~5 (within bucket 0 interpolation).
474        let p50 = h.percentile(0.5);
475        assert!(p50 <= 10, "p50 {p50} should be in bucket 0 (<10)");
476        std::fs::remove_file(&p).ok();
477    }
478
479    #[test]
480    fn percentile_zero_total_returns_zero() {
481        let p = tmp("p-zero");
482        let h = SharedHistogram::create(&p, &[10]).unwrap();
483        assert_eq!(h.percentile(0.5), 0);
484        assert_eq!(h.percentile(0.99), 0);
485        std::fs::remove_file(&p).ok();
486    }
487
488    #[test]
489    fn reset_clears_all_buckets() {
490        let p = tmp("reset");
491        let h = SharedHistogram::create(&p, &[10, 100]).unwrap();
492        for _ in 0..5 { h.record(50); }
493        assert_eq!(h.total_count(), 5);
494        h.reset();
495        assert_eq!(h.total_count(), 0);
496        for i in 0..h.n_buckets() {
497            assert_eq!(h.count(i).unwrap(), 0);
498        }
499        std::fs::remove_file(&p).ok();
500    }
501
502    #[test]
503    fn cross_handle_visibility() {
504        let p = tmp("cross-handle");
505        let writer = SharedHistogram::create(&p, &[10, 100, 1000]).unwrap();
506        let reader = SharedHistogram::open(&p, &[10, 100, 1000]).unwrap();
507        writer.record(50);
508        writer.record(500);
509        assert_eq!(reader.count(1).unwrap(), 1);
510        assert_eq!(reader.count(2).unwrap(), 1);
511        assert_eq!(reader.total_count(), 2);
512        std::fs::remove_file(&p).ok();
513    }
514
515    #[test]
516    fn open_with_wrong_boundaries_rejected() {
517        let p = tmp("wrong-bounds");
518        let _w = SharedHistogram::create(&p, &[10, 100]).unwrap();
519        assert!(matches!(
520            SharedHistogram::open(&p, &[10, 200]),
521            Err(HistogramError::LayoutMismatch)
522        ));
523        std::fs::remove_file(&p).ok();
524    }
525
526    #[test]
527    fn concurrent_recorders_no_lost_updates() {
528        let p = tmp("concurrent");
529        let h: Arc<SharedHistogram>
530            = Arc::new(SharedHistogram::create(&p, &[10, 100, 1000]).unwrap());
531        let n_threads = 4;
532        let per_thread = 250;
533        let mut handles = vec![];
534        for t in 0..n_threads {
535            let h = h.clone();
536            handles.push(thread::spawn(move || {
537                for i in 0..per_thread {
538                    // Distribute across buckets via modulo.
539                    let value = match (t * per_thread + i) % 4 {
540                        0 => 5,    // bucket 0
541                        1 => 50,   // bucket 1
542                        2 => 500,  // bucket 2
543                        _ => 5000, // bucket 3
544                    };
545                    h.record(value);
546                }
547            }));
548        }
549        for h in handles { h.join().unwrap(); }
550        let total = n_threads * per_thread;
551        assert_eq!(h.total_count() as usize, total);
552        let counts = h.counts();
553        // Each bucket should have ~total/4 records.
554        let expected_per = (total / 4) as u64;
555        for (i, &c) in counts.iter().enumerate() {
556            assert_eq!(c, expected_per,
557                "bucket {i} count {c} should be {expected_per}");
558        }
559        std::fs::remove_file(&p).ok();
560    }
561
562    #[test]
563    fn disk_persistence_survives_reopen() {
564        let p = tmp("disk");
565        let bounds = vec![10u64, 100, 1000];
566        {
567            let h = SharedHistogram::create(&p, &bounds).unwrap();
568            h.record(5);
569            h.record(50);
570            h.record(50);
571            h.record(5000);
572            h.flush().unwrap();
573        }
574        let h2 = SharedHistogram::open(&p, &bounds).unwrap();
575        assert_eq!(h2.count(0).unwrap(), 1);
576        assert_eq!(h2.count(1).unwrap(), 2);
577        assert_eq!(h2.count(2).unwrap(), 0);
578        assert_eq!(h2.count(3).unwrap(), 1);
579        assert_eq!(h2.total_count(), 4);
580        std::fs::remove_file(&p).ok();
581    }
582
583    #[test]
584    fn latency_distribution_pattern() {
585        // Realistic latency histogram: log-spaced boundaries in us.
586        let p = tmp("latency");
587        let bounds = vec![10u64, 100, 1_000, 10_000, 100_000, 1_000_000];
588        let h = SharedHistogram::create(&p, &bounds).unwrap();
589        // Simulate 1000 measurements; mostly fast, some tail.
590        for i in 0..1000u64 {
591            let latency_us = match i % 100 {
592                0..=80 => 5 + (i % 5),    // 81% under 10us
593                81..=95 => 50 + (i % 50), // 15% in 10..100us
594                _ => 500 + (i * 10),      // tail
595            };
596            h.record(latency_us);
597        }
598        // p50 should be in bucket 0 (< 10us).
599        let p50 = h.percentile(0.5);
600        assert!(p50 < 10, "p50 {p50} should be under 10us");
601        // p99 should be much higher (in the tail).
602        let p99 = h.percentile(0.99);
603        assert!(p99 > 100, "p99 {p99} should be over 100us");
604        std::fs::remove_file(&p).ok();
605    }
606
607    #[test]
608    fn count_out_of_bounds_rejected() {
609        let p = tmp("oob");
610        let h = SharedHistogram::create(&p, &[10]).unwrap();
611        // 2 buckets total (indices 0 and 1).
612        assert_eq!(h.count(2).err(), Some(HistogramError::OutOfBounds));
613        std::fs::remove_file(&p).ok();
614    }
615}