subetha_cxc/
shared_count_min_sketch.rs

1//! `SharedCountMinSketch` - cross-process probabilistic frequency
2//! estimator.
3//!
4//! `d` hash functions, `w` cells per row. `insert(item)` does `d`
5//! atomic `fetch_add(1)` ops; `estimate_count(item)` reads `d`
6//! cells and returns the minimum (the tightest unbiased upper
7//! bound on true count).
8//!
9//! # Safety properties
10//!
11//! - Pure `fetch_add` writes - no underflow possible.
12//! - No spin loops, no CAS retries, no RAII guards.
13//! - Bounded memory at create time (`d * w` cells of u64).
14//! - Hash collisions OVERCOUNT, never undercount. Estimate is a
15//!   guaranteed upper bound on true count.
16//!
17//! # Error bound
18//!
19//! With probability `1 - delta`, estimate <= true + `epsilon * N`
20//! where N is total inserts. Sizing: `w >= e/epsilon`,
21//! `d >= ln(1/delta)`. Standard config (epsilon=0.001, delta=0.001):
22//! w=2718, d=7, mem ~152 KB.
23//!
24//! # Hash family
25//!
26//! `d` independent hash positions derived from one FNV-1a + fmix64
27//! hash, then `pos[i] = (h + i * h2) % w` using two-hash
28//! double-hashing (Kirsch-Mitzenmacher). Same technique as
29//! [`SharedBloomFilter`](crate::SharedBloomFilter).
30
31use std::fs::{File, OpenOptions};
32use std::mem::size_of;
33use std::path::Path;
34use std::sync::atomic::{AtomicU64, Ordering};
35
36use memmap2::{MmapMut, MmapOptions};
37
38pub const CMS_MAGIC: u64 = 0x4150_434D_5330_3031;
39
40#[repr(C, align(64))]
41pub struct CMSHeader {
42    pub magic: u64,
43    pub d: u32,
44    pub w: u32,
45    pub total_inserts: AtomicU64,
46    _pad: [u8; 40],
47}
48
49const _: () = {
50    assert!(size_of::<CMSHeader>() == 64);
51};
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub enum CMSError {
55    InvalidConfig,
56    LayoutMismatch,
57    IoError(std::io::ErrorKind),
58}
59
60impl From<std::io::Error> for CMSError {
61    fn from(e: std::io::Error) -> Self { Self::IoError(e.kind()) }
62}
63
64pub fn cms_file_size(d: u32, w: u32) -> usize {
65    size_of::<CMSHeader>() + (d as usize) * (w as usize) * size_of::<AtomicU64>()
66}
67
68const FNV_OFFSET_BASIS_1: u64 = 0xcbf2_9ce4_8422_2325;
69const FNV_OFFSET_BASIS_2: u64 = 0x8422_2325_cbf2_9ce4;
70const FNV_PRIME: u64 = 0x100_0000_01b3;
71
72#[inline]
73fn fnv1a(bytes: &[u8], basis: u64) -> u64 {
74    let mut h = basis;
75    for &b in bytes {
76        h ^= b as u64;
77        h = h.wrapping_mul(FNV_PRIME);
78    }
79    h
80}
81
82#[inline]
83fn fmix64(mut h: u64) -> u64 {
84    h ^= h >> 33;
85    h = h.wrapping_mul(0xff51_afd7_ed55_8ccd);
86    h ^= h >> 33;
87    h = h.wrapping_mul(0xc4ce_b9fe_1a85_ec53);
88    h ^= h >> 33;
89    h
90}
91
92pub struct SharedCountMinSketch {
93    _file: File,
94    mmap: MmapMut,
95    d: u32,
96    w: u32,
97    header_sidecar: subetha_core::HandshakeHeader,
98    ring_sidecar: Box<subetha_core::ObservationRing>,
99}
100
101unsafe impl Send for SharedCountMinSketch {}
102unsafe impl Sync for SharedCountMinSketch {}
103
104impl subetha_sidecar::AdaptiveInstance for SharedCountMinSketch {
105    fn header(&self) -> &subetha_core::HandshakeHeader { &self.header_sidecar }
106    fn ring(&self) -> &subetha_core::ObservationRing { &self.ring_sidecar }
107    fn make_policy(&self) -> Box<dyn subetha_sidecar::Policy> {
108        Box::new(subetha_sidecar::NoMigrationPolicy)
109    }
110}
111
112impl SharedCountMinSketch {
113    /// Suggest (d, w) for target (epsilon, delta) error bounds.
114    /// Estimate <= true + epsilon * N with probability 1 - delta.
115    pub fn suggest_config(epsilon: f64, delta: f64) -> (u32, u32) {
116        assert!(epsilon > 0.0 && epsilon < 1.0);
117        assert!(delta > 0.0 && delta < 1.0);
118        let w = (std::f64::consts::E / epsilon).ceil() as u32;
119        let d = (1.0 / delta).ln().ceil() as u32;
120        (d.max(1), w.max(1))
121    }
122
123    /// Obtain the sketch at `path`, initializing an empty one if the
124    /// path does not yet exist and attaching to it if it does.
125    /// Attaching leaves live counts in place; a region built with a
126    /// different `(d, w)` is a `LayoutMismatch`. The in-place
127    /// [`reset`](Self::reset) zeroes a live sketch.
128    pub fn create(
129        path: impl AsRef<Path>, d: u32, w: u32,
130    ) -> Result<Self, CMSError> {
131        if d == 0 || w == 0 {
132            return Err(CMSError::InvalidConfig);
133        }
134        let (file, mmap) = crate::mmf_attach::create_or_attach(
135            path.as_ref(),
136            cms_file_size(d, w),
137            |ptr| unsafe { Self::init_region(ptr, d, w) },
138            |ptr| unsafe { (*(ptr as *const CMSHeader)).magic == CMS_MAGIC },
139        )?;
140        Self::from_region(file, mmap, d, w)
141    }
142
143    /// Lay out an empty sketch: config first, magic last, because
144    /// attachers spin on it. The zeroed region is already the zero
145    /// counter matrix.
146    ///
147    /// # Safety
148    /// `ptr` addresses at least `cms_file_size(d, w)` writable zeroed
149    /// bytes.
150    unsafe fn init_region(ptr: *mut u8, d: u32, w: u32) {
151        let hdr = ptr as *mut CMSHeader;
152        unsafe {
153            (*hdr).d = d;
154            (*hdr).w = w;
155            std::ptr::write_volatile(&raw mut (*hdr).magic, CMS_MAGIC);
156        }
157    }
158
159    /// Wrap an initialized region, refusing one built with a different
160    /// config.
161    fn from_region(file: File, mmap: MmapMut, d: u32, w: u32) -> Result<Self, CMSError> {
162        let hdr = unsafe { &*(mmap.as_ptr() as *const CMSHeader) };
163        if hdr.magic != CMS_MAGIC || hdr.d != d || hdr.w != w {
164            return Err(CMSError::LayoutMismatch);
165        }
166        Ok(Self {
167            _file: file, mmap, d, w,
168            header_sidecar: subetha_core::HandshakeHeader::new(),
169            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
170        })
171    }
172
173    pub fn open(
174        path: impl AsRef<Path>, expected_d: u32, expected_w: u32,
175    ) -> Result<Self, CMSError> {
176        let total = cms_file_size(expected_d, expected_w);
177        let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
178        if file.metadata()?.len() < total as u64 {
179            return Err(CMSError::LayoutMismatch);
180        }
181        let mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
182        Self::from_region(file, mmap, expected_d, expected_w)
183    }
184
185    #[inline]
186    pub fn d(&self) -> u32 { self.d }
187    #[inline]
188    pub fn w(&self) -> u32 { self.w }
189    #[inline]
190    pub fn total_inserts(&self) -> u64 {
191        unsafe { (*(self.mmap.as_ptr() as *const CMSHeader)).total_inserts.load(Ordering::Acquire) }
192    }
193
194    fn cell(&self, row: u32, col: u32) -> &AtomicU64 {
195        let idx = (row as usize) * (self.w as usize) + (col as usize);
196        let base = unsafe { self.mmap.as_ptr().add(size_of::<CMSHeader>()) };
197        unsafe { &*(base.add(idx * size_of::<AtomicU64>()) as *const AtomicU64) }
198    }
199
200    /// Invoke `f(row, col)` for each of the `d` hash positions of `item`,
201    /// without allocating. Two FNV-1a + fmix64 hashes feed
202    /// Kirsch-Mitzenmacher double-hashing `h1 + row * h2`; the column is
203    /// reduced mod `w`, using a bit-mask when `w` is a power of two (the
204    /// runtime divisor otherwise compiles to a hardware DIV per row).
205    #[inline]
206    fn for_each_position(&self, item: &[u8], mut f: impl FnMut(u32, u32)) {
207        let h1 = fmix64(fnv1a(item, FNV_OFFSET_BASIS_1));
208        let h2 = fmix64(fnv1a(item, FNV_OFFSET_BASIS_2));
209        let w = self.w as u64;
210        let pow2 = self.w.is_power_of_two();
211        let mask = w.wrapping_sub(1);
212        for row in 0..self.d {
213            let raw = h1.wrapping_add((row as u64).wrapping_mul(h2));
214            let col = if pow2 {
215                (raw & mask) as u32
216            } else {
217                // Lemire fastrange: multiply-shift maps into [0, w) with no
218                // hardware DIV. h1/h2 are fmix64-avalanched, so the high
219                // bits the shift keys on are well-distributed.
220                ((raw as u128 * w as u128) >> 64) as u32
221            };
222            f(row, col);
223        }
224    }
225
226    /// Insert one observation of `item`. `d` atomic increments.
227    pub fn insert(&self, item: &[u8]) {
228        self.for_each_position(item, |row, col| {
229            self.cell(row, col).fetch_add(1, Ordering::AcqRel);
230        });
231        unsafe {
232            (*(self.mmap.as_ptr() as *const CMSHeader))
233                .total_inserts.fetch_add(1, Ordering::AcqRel);
234        }
235        self.ring_sidecar
236            .push_op(crate::sidecar_ops::sketch::OP_INSERT, 0);
237    }
238
239    /// Insert `count` observations at once (bulk increment).
240    pub fn insert_n(&self, item: &[u8], count: u64) {
241        if count == 0 { return; }
242        self.for_each_position(item, |row, col| {
243            self.cell(row, col).fetch_add(count, Ordering::AcqRel);
244        });
245        unsafe {
246            (*(self.mmap.as_ptr() as *const CMSHeader))
247                .total_inserts.fetch_add(count, Ordering::AcqRel);
248        }
249        self.ring_sidecar
250            .push_op(crate::sidecar_ops::sketch::OP_INSERT, 0);
251    }
252
253    /// Estimate the frequency of `item`. Guaranteed upper bound on
254    /// true count.
255    pub fn estimate_count(&self, item: &[u8]) -> u64 {
256        let mut v = u64::MAX;
257        self.for_each_position(item, |row, col| {
258            let c = self.cell(row, col).load(Ordering::Acquire);
259            if c < v { v = c; }
260        });
261        let v = if self.d == 0 { 0 } else { v };
262        self.ring_sidecar.push_op(
263            crate::sidecar_ops::sketch::OP_QUERY,
264            if v == 0 { 2 } else { 0 },
265        );
266        v
267    }
268
269    /// Reset all cells to zero.
270    pub fn reset(&self) {
271        for row in 0..self.d {
272            for col in 0..self.w {
273                self.cell(row, col).store(0, Ordering::Release);
274            }
275        }
276        unsafe {
277            (*(self.mmap.as_ptr() as *const CMSHeader))
278                .total_inserts.store(0, Ordering::Release);
279        }
280        self.ring_sidecar
281            .push_op(crate::sidecar_ops::sketch::OP_CLEAR, 0);
282    }
283
284    pub fn flush(&self) -> Result<(), CMSError> {
285        self.mmap.flush()?;
286        Ok(())
287    }
288    pub fn flush_async(&self) -> Result<(), CMSError> {
289        self.mmap.flush_async()?;
290        Ok(())
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297    use std::sync::Arc;
298    use std::thread;
299
300    fn tmp(name: &str) -> std::path::PathBuf {
301        let mut p = std::env::temp_dir();
302        let pid = std::process::id();
303        p.push(format!("subetha-cms-{name}-{pid}.bin"));
304        p
305    }
306
307    #[test]
308    fn create_initial_state_zero() {
309        let p = tmp("init");
310        let cms = SharedCountMinSketch::create(&p, 4, 256).unwrap();
311        assert_eq!(cms.d(), 4);
312        assert_eq!(cms.w(), 256);
313        assert_eq!(cms.total_inserts(), 0);
314        assert_eq!(cms.estimate_count(b"anything"), 0);
315        std::fs::remove_file(&p).ok();
316    }
317
318    /// A second create attaches with live counts in place; the
319    /// in-place reset is what zeroes them.
320    #[test]
321    fn second_create_attaches_and_keeps_counts() {
322        let p = tmp("attach");
323        std::fs::remove_file(&p).ok();
324        let cms = SharedCountMinSketch::create(&p, 4, 256).unwrap();
325        cms.insert(b"key");
326        cms.insert(b"key");
327
328        let cms2 = SharedCountMinSketch::create(&p, 4, 256).unwrap();
329        assert_eq!(cms2.estimate_count(b"key"), 2, "attach zeroed live counts");
330        assert!(matches!(
331            SharedCountMinSketch::create(&p, 2, 256),
332            Err(CMSError::LayoutMismatch),
333        ));
334
335        cms2.reset();
336        assert_eq!(cms.estimate_count(b"key"), 0, "reset did not zero for every handle");
337        drop(cms);
338        drop(cms2);
339        std::fs::remove_file(&p).ok();
340    }
341
342    #[test]
343    fn invalid_config_rejected() {
344        let p = tmp("invalid");
345        assert_eq!(
346            SharedCountMinSketch::create(&p, 0, 256).err(),
347            Some(CMSError::InvalidConfig)
348        );
349        assert_eq!(
350            SharedCountMinSketch::create(&p, 4, 0).err(),
351            Some(CMSError::InvalidConfig)
352        );
353        std::fs::remove_file(&p).ok();
354    }
355
356    #[test]
357    fn insert_then_estimate_returns_at_least_true_count() {
358        let p = tmp("insert");
359        let cms = SharedCountMinSketch::create(&p, 4, 1024).unwrap();
360        for _ in 0..5 { cms.insert(b"foo"); }
361        for _ in 0..3 { cms.insert(b"bar"); }
362        // Guaranteed: estimate >= true.
363        assert!(cms.estimate_count(b"foo") >= 5);
364        assert!(cms.estimate_count(b"bar") >= 3);
365        assert_eq!(cms.total_inserts(), 8);
366        std::fs::remove_file(&p).ok();
367    }
368
369    #[test]
370    fn insert_n_is_equivalent_to_n_inserts() {
371        let p = tmp("insert-n");
372        let a = SharedCountMinSketch::create(tmp("insert-n-a"), 4, 1024).unwrap();
373        let b = SharedCountMinSketch::create(tmp("insert-n-b"), 4, 1024).unwrap();
374        for _ in 0..100 { a.insert(b"item"); }
375        b.insert_n(b"item", 100);
376        assert_eq!(a.estimate_count(b"item"), b.estimate_count(b"item"));
377        assert_eq!(a.total_inserts(), b.total_inserts());
378        let _p = p;
379        std::fs::remove_file(tmp("insert-n-a")).ok();
380        std::fs::remove_file(tmp("insert-n-b")).ok();
381    }
382
383    #[test]
384    fn estimate_for_absent_item_is_low() {
385        // Probability of false-positive (estimate > 0) on absent
386        // item is bounded by (n / w)^d. For n=100 inserts, w=1024,
387        // d=4: (100/1024)^4 = 9e-5 - very low.
388        let p = tmp("absent");
389        let cms = SharedCountMinSketch::create(&p, 4, 1024).unwrap();
390        for i in 0..100u32 {
391            cms.insert(format!("inserted-{i}").as_bytes());
392        }
393        // Try 100 absent items; expect very few false positives.
394        let mut fp = 0u32;
395        for i in 0..100u32 {
396            if cms.estimate_count(format!("absent-{i}").as_bytes()) > 0 {
397                fp += 1;
398            }
399        }
400        // Allow up to 10% (well above theoretical bound).
401        assert!(fp < 10, "expected < 10 false positives, got {fp}");
402        std::fs::remove_file(&p).ok();
403    }
404
405    #[test]
406    fn heavy_hitter_detection() {
407        // Insert 1 heavy item (10000 times) + 1000 single-insert
408        // items. Heavy item should clearly stand out.
409        let p = tmp("heavy");
410        let cms = SharedCountMinSketch::create(&p, 5, 2048).unwrap();
411        for _ in 0..10_000 { cms.insert(b"HEAVY"); }
412        for i in 0..1000u32 {
413            cms.insert(format!("light-{i}").as_bytes());
414        }
415        let heavy = cms.estimate_count(b"HEAVY");
416        // Heavy is at least 10000; error <= epsilon * N where
417        // N = 11000 and epsilon ~ e/2048 ~ 0.00133. So error <= ~15.
418        assert!(heavy >= 10_000);
419        assert!(heavy <= 10_100, "heavy estimate {heavy} should be very close to 10000");
420        // Light items should estimate near 1.
421        for i in 0..10u32 {
422            let light = cms.estimate_count(format!("light-{i}").as_bytes());
423            assert!(light < 20, "light item {i} estimate {light} should be small");
424        }
425        std::fs::remove_file(&p).ok();
426    }
427
428    #[test]
429    fn reset_zeroes_everything() {
430        let p = tmp("reset");
431        let cms = SharedCountMinSketch::create(&p, 4, 256).unwrap();
432        for _ in 0..50 { cms.insert(b"x"); }
433        assert!(cms.estimate_count(b"x") >= 50);
434        cms.reset();
435        assert_eq!(cms.estimate_count(b"x"), 0);
436        assert_eq!(cms.total_inserts(), 0);
437        std::fs::remove_file(&p).ok();
438    }
439
440    #[test]
441    fn cross_handle_visibility() {
442        let p = tmp("cross-handle");
443        let w = SharedCountMinSketch::create(&p, 4, 256).unwrap();
444        let r = SharedCountMinSketch::open(&p, 4, 256).unwrap();
445        w.insert(b"shared");
446        w.insert(b"shared");
447        assert!(r.estimate_count(b"shared") >= 2);
448        assert_eq!(r.total_inserts(), 2);
449        std::fs::remove_file(&p).ok();
450    }
451
452    #[test]
453    fn config_mismatch_at_open_rejected() {
454        let p = tmp("mismatch");
455        let _w = SharedCountMinSketch::create(&p, 4, 256).unwrap();
456        assert!(matches!(
457            SharedCountMinSketch::open(&p, 5, 256),
458            Err(CMSError::LayoutMismatch)
459        ));
460        assert!(matches!(
461            SharedCountMinSketch::open(&p, 4, 512),
462            Err(CMSError::LayoutMismatch)
463        ));
464        std::fs::remove_file(&p).ok();
465    }
466
467    #[test]
468    fn suggest_config_returns_sensible_values() {
469        // epsilon=0.01, delta=0.01 -> w ~ e/0.01 ~ 272, d ~ ln(100) ~ 5.
470        let (d, w) = SharedCountMinSketch::suggest_config(0.01, 0.01);
471        assert!((250..=300).contains(&w));
472        assert!((4..=6).contains(&d));
473    }
474
475    #[test]
476    fn concurrent_inserters_accurate() {
477        // 4 threads each insert "shared-item" 1000 times. After
478        // join, estimate should be at least 4000.
479        let p = tmp("concurrent");
480        let cms = Arc::new(SharedCountMinSketch::create(&p, 4, 1024).unwrap());
481        let mut handles = vec![];
482        for _ in 0..4 {
483            let cms = cms.clone();
484            handles.push(thread::spawn(move || {
485                for _ in 0..1000 { cms.insert(b"shared-item"); }
486            }));
487        }
488        for h in handles { h.join().unwrap(); }
489        let est = cms.estimate_count(b"shared-item");
490        assert!(est >= 4000, "concurrent estimate {est} should be >= 4000");
491        assert_eq!(cms.total_inserts(), 4000);
492        std::fs::remove_file(&p).ok();
493    }
494
495    #[test]
496    fn disk_persistence_survives_reopen() {
497        let p = tmp("disk");
498        {
499            let cms = SharedCountMinSketch::create(&p, 4, 256).unwrap();
500            for _ in 0..50 { cms.insert(b"persisted"); }
501            cms.flush().unwrap();
502        }
503        let cms2 = SharedCountMinSketch::open(&p, 4, 256).unwrap();
504        assert!(cms2.estimate_count(b"persisted") >= 50);
505        assert_eq!(cms2.total_inserts(), 50);
506        std::fs::remove_file(&p).ok();
507    }
508}
subetha_cxc/shared_count_min_sketch.rs

subetha_cxc/
shared_count_min_sketch.rs
