Skip to main content

subetha_cxc/
shared_hyper_log_log.rs

1//! `SharedHyperLogLog` - cross-process probabilistic distinct-count
2//! estimator.
3//!
4//! 2^p AtomicU8 registers, each storing the maximum-observed rank
5//! (leading-zero count + 1) of items hashed to that register.
6//! Estimate via harmonic mean with bias correction.
7//!
8//! # Why this is safe
9//!
10//! Each `insert` is exactly one `fetch_max` on one AtomicU8. No
11//! CAS loops, no Drop guards, no spin waits. The cache-line
12//! contention is bounded to one register per insert.
13//!
14//! # Accuracy
15//!
16//! Standard error ~= 1.04 / sqrt(m) where m = 2^p.
17//!
18//! | p  | m      | std err | size  |
19//! |----|--------|---------|-------|
20//! | 8  | 256    | 6.5%    | 256B  |
21//! | 10 | 1024   | 3.3%    | 1 KB  |
22//! | 12 | 4096   | 1.6%    | 4 KB  |
23//! | 14 | 16384  | 0.8%    | 16 KB |
24//! | 16 | 65536  | 0.4%    | 64 KB |
25//!
26//! # Encoding
27//!
28//! Hash item to u64 h. Register index = top p bits of h. Rank =
29//! (leading_zeros of (h << p) | (1 << (63-p))) + 1, clamped to
30//! 64. (The OR ensures rank is bounded even when low bits are 0.)
31
32use std::fs::{File, OpenOptions};
33use std::mem::size_of;
34use std::path::Path;
35use std::sync::atomic::{AtomicU8, Ordering};
36
37use memmap2::{MmapMut, MmapOptions};
38
39pub const HLL_MAGIC: u64 = 0x4150_484C_4C56_3031;
40/// Minimum precision (smaller = less memory but worse accuracy).
41pub const MIN_PRECISION: u8 = 4;
42/// Maximum precision (larger = more memory).
43pub const MAX_PRECISION: u8 = 16;
44
45#[repr(C, align(64))]
46pub struct HLLHeader {
47    pub magic: u64,
48    pub precision: u32,
49    pub m: u32,
50    _pad: [u8; 48],
51}
52
53const _: () = {
54    assert!(size_of::<HLLHeader>() == 64);
55};
56
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub enum HLLError {
59    InvalidPrecision,
60    LayoutMismatch,
61    IoError(std::io::ErrorKind),
62}
63
64impl From<std::io::Error> for HLLError {
65    fn from(e: std::io::Error) -> Self { Self::IoError(e.kind()) }
66}
67
68const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
69const FNV_PRIME: u64 = 0x100_0000_01b3;
70
71#[inline]
72fn fnv1a_64(bytes: &[u8]) -> u64 {
73    let mut h = FNV_OFFSET_BASIS;
74    for &b in bytes {
75        h ^= b as u64;
76        h = h.wrapping_mul(FNV_PRIME);
77    }
78    h
79}
80
81/// MurmurHash3 fmix64 finalizer. Bit-mixes a u64 to give excellent
82/// distribution properties (avalanche: 1-bit input flip => ~50%
83/// of output bits flip). Critical for HLL because the register
84/// index is extracted from the TOP bits of the hash, and raw FNV-1a
85/// on short inputs has poor top-bit distribution.
86#[inline]
87fn fmix64(mut h: u64) -> u64 {
88    h ^= h >> 33;
89    h = h.wrapping_mul(0xff51_afd7_ed55_8ccd);
90    h ^= h >> 33;
91    h = h.wrapping_mul(0xc4ce_b9fe_1a85_ec53);
92    h ^= h >> 33;
93    h
94}
95
96#[inline]
97fn hash_for_hll(item: &[u8]) -> u64 {
98    fmix64(fnv1a_64(item))
99}
100
101pub fn hll_file_size(precision: u8) -> usize {
102    size_of::<HLLHeader>() + (1usize << precision)
103}
104
105pub struct SharedHyperLogLog {
106    _file: File,
107    mmap: MmapMut,
108    precision: u8,
109    m: u32,
110    header_sidecar: subetha_core::HandshakeHeader,
111    ring_sidecar: Box<subetha_core::ObservationRing>,
112}
113
114unsafe impl Send for SharedHyperLogLog {}
115unsafe impl Sync for SharedHyperLogLog {}
116
117impl subetha_sidecar::AdaptiveInstance for SharedHyperLogLog {
118    fn header(&self) -> &subetha_core::HandshakeHeader { &self.header_sidecar }
119    fn ring(&self) -> &subetha_core::ObservationRing { &self.ring_sidecar }
120    fn make_policy(&self) -> Box<dyn subetha_sidecar::Policy> {
121        Box::new(subetha_sidecar::NoMigrationPolicy)
122    }
123}
124
125impl SharedHyperLogLog {
126    /// Obtain the estimator at `path`, initializing an empty one if
127    /// the path does not yet exist and attaching to it if it does.
128    /// Attaching leaves live registers in place; a region built with a
129    /// different precision is a `LayoutMismatch`. The in-place
130    /// [`reset`](Self::reset) zeroes a live estimator.
131    pub fn create(
132        path: impl AsRef<Path>, precision: u8,
133    ) -> Result<Self, HLLError> {
134        if !(MIN_PRECISION..=MAX_PRECISION).contains(&precision) {
135            return Err(HLLError::InvalidPrecision);
136        }
137        let (file, mmap) = crate::mmf_attach::create_or_attach(
138            path.as_ref(),
139            hll_file_size(precision),
140            |ptr| unsafe { Self::init_region(ptr, precision) },
141            |ptr| unsafe { (*(ptr as *const HLLHeader)).magic == HLL_MAGIC },
142        )?;
143        Self::from_region(file, mmap, precision)
144    }
145
146    /// Lay out an empty estimator: precision and register count first,
147    /// magic last, because attachers spin on it. The zeroed region is
148    /// already the all-zero register array.
149    ///
150    /// # Safety
151    /// `ptr` addresses at least `hll_file_size(precision)` writable
152    /// zeroed bytes.
153    unsafe fn init_region(ptr: *mut u8, precision: u8) {
154        let hdr = ptr as *mut HLLHeader;
155        unsafe {
156            (*hdr).precision = precision as u32;
157            (*hdr).m = 1u32 << precision;
158            std::ptr::write_volatile(&raw mut (*hdr).magic, HLL_MAGIC);
159        }
160    }
161
162    /// Wrap an initialized region, refusing one built with a different
163    /// precision.
164    fn from_region(file: File, mmap: MmapMut, precision: u8) -> Result<Self, HLLError> {
165        let hdr = unsafe { &*(mmap.as_ptr() as *const HLLHeader) };
166        if hdr.magic != HLL_MAGIC || hdr.precision != precision as u32 {
167            return Err(HLLError::LayoutMismatch);
168        }
169        let m = hdr.m;
170        Ok(Self {
171            _file: file, mmap, precision, m,
172            header_sidecar: subetha_core::HandshakeHeader::new(),
173            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
174        })
175    }
176
177    pub fn open(
178        path: impl AsRef<Path>, expected_precision: u8,
179    ) -> Result<Self, HLLError> {
180        let total = hll_file_size(expected_precision);
181        let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
182        if file.metadata()?.len() < total as u64 {
183            return Err(HLLError::LayoutMismatch);
184        }
185        let mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
186        Self::from_region(file, mmap, expected_precision)
187    }
188
189    #[inline]
190    pub fn precision(&self) -> u8 { self.precision }
191    #[inline]
192    pub fn n_registers(&self) -> u32 { self.m }
193
194    fn register(&self, idx: usize) -> &AtomicU8 {
195        let base = unsafe { self.mmap.as_ptr().add(size_of::<HLLHeader>()) };
196        unsafe { &*(base.add(idx) as *const AtomicU8) }
197    }
198
199    /// Insert an item. One fetch_max on one register; no spinning.
200    pub fn insert(&self, item: &[u8]) {
201        let h = hash_for_hll(item);
202        let p = self.precision as u32;
203        let reg_idx = (h >> (64 - p)) as usize;
204        // Compute rank: position of leftmost 1-bit in the low
205        // (64 - p) bits. We OR in a guard bit so rank is always
206        // bounded by 64 - p + 1.
207        let w = (h << p) | (1u64 << (p.saturating_sub(1)));
208        let rank = (w.leading_zeros() as u8) + 1;
209        self.register(reg_idx).fetch_max(rank, Ordering::AcqRel);
210        self.ring_sidecar
211            .push_op(crate::sidecar_ops::sketch::OP_INSERT, 0);
212    }
213
214    /// Estimate cardinality (distinct count) via harmonic mean with
215    /// bias correction.
216    pub fn estimate(&self) -> u64 {
217        let m = self.m as f64;
218        let alpha = match self.m {
219            16 => 0.673,
220            32 => 0.697,
221            64 => 0.709,
222            _ => 0.7213 / (1.0 + 1.079 / m),
223        };
224        let mut sum = 0.0f64;
225        let mut zeros = 0u32;
226        for i in 0..self.m as usize {
227            let r = self.register(i).load(Ordering::Acquire);
228            if r == 0 { zeros += 1; }
229            sum += 2f64.powi(-(r as i32));
230        }
231        let raw = alpha * m * m / sum;
232        // Small-range correction: if estimate < 2.5 * m and there
233        // are empty registers, use linear counting.
234        let v = if raw <= 2.5 * m && zeros > 0 {
235            let z = zeros as f64;
236            (m * (m / z).ln()).round() as u64
237        } else {
238            raw.round() as u64
239        };
240        self.ring_sidecar
241            .push_op(crate::sidecar_ops::sketch::OP_QUERY, 0);
242        v
243    }
244
245    /// Reset all registers to zero (the empty state).
246    pub fn reset(&self) {
247        for i in 0..self.m as usize {
248            self.register(i).store(0, Ordering::Release);
249        }
250        self.ring_sidecar
251            .push_op(crate::sidecar_ops::sketch::OP_CLEAR, 0);
252    }
253
254    pub fn flush(&self) -> Result<(), HLLError> {
255        self.mmap.flush()?;
256        Ok(())
257    }
258    pub fn flush_async(&self) -> Result<(), HLLError> {
259        self.mmap.flush_async()?;
260        Ok(())
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267    use std::sync::Arc;
268    use std::thread;
269
270    fn tmp(name: &str) -> std::path::PathBuf {
271        let mut p = std::env::temp_dir();
272        let pid = std::process::id();
273        p.push(format!("subetha-hll-{name}-{pid}.bin"));
274        p
275    }
276
277    /// A second create attaches with live registers in place; the
278    /// in-place reset is what zeroes them.
279    #[test]
280    fn second_create_attaches_and_keeps_registers() {
281        let p = tmp("attach");
282        std::fs::remove_file(&p).ok();
283        let h = SharedHyperLogLog::create(&p, 12).unwrap();
284        h.insert(b"one");
285
286        let h2 = SharedHyperLogLog::create(&p, 12).unwrap();
287        assert_eq!(h2.estimate(), 1, "attach zeroed live registers");
288        assert!(matches!(
289            SharedHyperLogLog::create(&p, 10),
290            Err(HLLError::LayoutMismatch),
291        ));
292
293        h2.reset();
294        assert_eq!(h.estimate(), 0, "reset did not zero for every handle");
295        drop(h);
296        drop(h2);
297        std::fs::remove_file(&p).ok();
298    }
299
300    #[test]
301    fn create_initial_empty_estimate_is_zero() {
302        let p = tmp("init");
303        let h = SharedHyperLogLog::create(&p, 12).unwrap();
304        assert_eq!(h.precision(), 12);
305        assert_eq!(h.n_registers(), 4096);
306        assert_eq!(h.estimate(), 0);
307        std::fs::remove_file(&p).ok();
308    }
309
310    #[test]
311    fn invalid_precision_rejected() {
312        let p = tmp("invalid");
313        assert_eq!(
314            SharedHyperLogLog::create(&p, 3).err(),
315            Some(HLLError::InvalidPrecision)
316        );
317        assert_eq!(
318            SharedHyperLogLog::create(&p, 17).err(),
319            Some(HLLError::InvalidPrecision)
320        );
321        std::fs::remove_file(&p).ok();
322    }
323
324    #[test]
325    fn single_insert_estimate_is_one() {
326        let p = tmp("single");
327        let h = SharedHyperLogLog::create(&p, 12).unwrap();
328        h.insert(b"hello");
329        let est = h.estimate();
330        // Linear-counting small-range correction handles single
331        // insert; estimate should be exactly 1.
332        assert_eq!(est, 1, "single insert should estimate 1, got {est}");
333        std::fs::remove_file(&p).ok();
334    }
335
336    #[test]
337    fn idempotent_inserts_same_item() {
338        let p = tmp("idempotent");
339        let h = SharedHyperLogLog::create(&p, 12).unwrap();
340        for _ in 0..100 { h.insert(b"same-item"); }
341        let est = h.estimate();
342        assert_eq!(est, 1, "100 inserts of same item should estimate 1, got {est}");
343        std::fs::remove_file(&p).ok();
344    }
345
346    #[test]
347    fn distinct_count_within_error_bound_p12() {
348        // p=12 -> m=4096, std err 1.6%. For 1000 distinct items,
349        // expect estimate in roughly [950, 1050] (2 sigma).
350        let p = tmp("distinct-1k");
351        let h = SharedHyperLogLog::create(&p, 12).unwrap();
352        for i in 0..1000u32 {
353            h.insert(format!("item-{i:05}").as_bytes());
354        }
355        let est = h.estimate();
356        assert!(
357            (900..=1100).contains(&est),
358            "estimate {est} should be within 10% of 1000 (got {est})",
359        );
360        std::fs::remove_file(&p).ok();
361    }
362
363    #[test]
364    fn distinct_count_within_error_bound_p14() {
365        // p=14 -> m=16384, std err 0.8%. For 10000 distinct items,
366        // expect estimate within ~4% (2.5 sigma).
367        let p = tmp("distinct-10k");
368        let h = SharedHyperLogLog::create(&p, 14).unwrap();
369        for i in 0..10_000u32 {
370            h.insert(format!("k{i:06}").as_bytes());
371        }
372        let est = h.estimate();
373        assert!(
374            (9600..=10400).contains(&est),
375            "estimate {est} should be within 4% of 10000",
376        );
377        std::fs::remove_file(&p).ok();
378    }
379
380    #[test]
381    fn reset_clears_registers() {
382        let p = tmp("reset");
383        let h = SharedHyperLogLog::create(&p, 10).unwrap();
384        for i in 0..100u32 { h.insert(format!("k{i}").as_bytes()); }
385        assert!(h.estimate() > 0);
386        h.reset();
387        assert_eq!(h.estimate(), 0);
388        std::fs::remove_file(&p).ok();
389    }
390
391    #[test]
392    fn cross_handle_visibility() {
393        let p = tmp("cross-handle");
394        let w = SharedHyperLogLog::create(&p, 10).unwrap();
395        let r = SharedHyperLogLog::open(&p, 10).unwrap();
396        for i in 0..50u32 { w.insert(format!("k{i}").as_bytes()); }
397        let est_w = w.estimate();
398        let est_r = r.estimate();
399        assert_eq!(est_w, est_r);
400        // Roughly 50.
401        assert!((40..=60).contains(&est_r), "estimate {est_r} should be near 50");
402        std::fs::remove_file(&p).ok();
403    }
404
405    #[test]
406    fn config_mismatch_at_open_rejected() {
407        let p = tmp("mismatch");
408        let _w = SharedHyperLogLog::create(&p, 10).unwrap();
409        assert!(matches!(
410            SharedHyperLogLog::open(&p, 12),
411            Err(HLLError::LayoutMismatch)
412        ));
413        std::fs::remove_file(&p).ok();
414    }
415
416    #[test]
417    fn concurrent_inserters_correct_estimate() {
418        // 4 threads each insert 1000 distinct items (disjoint key
419        // spaces). Total = 4000. Estimate should be near 4000.
420        let p = tmp("concurrent");
421        let h: Arc<SharedHyperLogLog>
422            = Arc::new(SharedHyperLogLog::create(&p, 12).unwrap());
423        let mut handles = vec![];
424        for t in 0..4u32 {
425            let h = h.clone();
426            handles.push(thread::spawn(move || {
427                for i in 0..1000u32 {
428                    h.insert(format!("t{t}-i{i:05}").as_bytes());
429                }
430            }));
431        }
432        for h in handles { h.join().unwrap(); }
433        let est = h.estimate();
434        assert!(
435            (3700..=4300).contains(&est),
436            "concurrent inserts: estimate {est} should be within ~7% of 4000",
437        );
438        std::fs::remove_file(&p).ok();
439    }
440
441    #[test]
442    fn disk_persistence_survives_reopen() {
443        let p = tmp("disk");
444        {
445            let h = SharedHyperLogLog::create(&p, 10).unwrap();
446            for i in 0..100u32 { h.insert(format!("k{i}").as_bytes()); }
447            h.flush().unwrap();
448        }
449        let h2 = SharedHyperLogLog::open(&p, 10).unwrap();
450        let est = h2.estimate();
451        assert!((80..=120).contains(&est),
452            "reopened estimate {est} should be near 100");
453        std::fs::remove_file(&p).ok();
454    }
455}