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    pub fn create(
127        path: impl AsRef<Path>, precision: u8,
128    ) -> Result<Self, HLLError> {
129        if !(MIN_PRECISION..=MAX_PRECISION).contains(&precision) {
130            return Err(HLLError::InvalidPrecision);
131        }
132        let m = 1u32 << precision;
133        let total = hll_file_size(precision);
134        let file = OpenOptions::new()
135            .read(true).write(true).create(true).truncate(true)
136            .open(path.as_ref())?;
137        file.set_len(total as u64)?;
138        let mut mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
139        let hdr = mmap.as_mut_ptr() as *mut HLLHeader;
140        unsafe {
141            std::ptr::write_bytes(hdr as *mut u8, 0, size_of::<HLLHeader>());
142            (*hdr).magic = HLL_MAGIC;
143            (*hdr).precision = precision as u32;
144            (*hdr).m = m;
145        }
146        // Registers are zero from set_len + map_mut.
147        Ok(Self {
148            _file: file, mmap, precision, m,
149            header_sidecar: subetha_core::HandshakeHeader::new(),
150            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
151        })
152    }
153
154    pub fn open(
155        path: impl AsRef<Path>, expected_precision: u8,
156    ) -> Result<Self, HLLError> {
157        let total = hll_file_size(expected_precision);
158        let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
159        if file.metadata()?.len() < total as u64 {
160            return Err(HLLError::LayoutMismatch);
161        }
162        let mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
163        let hdr = unsafe { &*(mmap.as_ptr() as *const HLLHeader) };
164        if hdr.magic != HLL_MAGIC || hdr.precision != expected_precision as u32 {
165            return Err(HLLError::LayoutMismatch);
166        }
167        let m = hdr.m;
168        Ok(Self {
169            _file: file, mmap, precision: expected_precision, m,
170            header_sidecar: subetha_core::HandshakeHeader::new(),
171            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
172        })
173    }
174
175    #[inline]
176    pub fn precision(&self) -> u8 { self.precision }
177    #[inline]
178    pub fn n_registers(&self) -> u32 { self.m }
179
180    fn register(&self, idx: usize) -> &AtomicU8 {
181        let base = unsafe { self.mmap.as_ptr().add(size_of::<HLLHeader>()) };
182        unsafe { &*(base.add(idx) as *const AtomicU8) }
183    }
184
185    /// Insert an item. One fetch_max on one register; no spinning.
186    pub fn insert(&self, item: &[u8]) {
187        let h = hash_for_hll(item);
188        let p = self.precision as u32;
189        let reg_idx = (h >> (64 - p)) as usize;
190        // Compute rank: position of leftmost 1-bit in the low
191        // (64 - p) bits. We OR in a guard bit so rank is always
192        // bounded by 64 - p + 1.
193        let w = (h << p) | (1u64 << (p.saturating_sub(1)));
194        let rank = (w.leading_zeros() as u8) + 1;
195        self.register(reg_idx).fetch_max(rank, Ordering::AcqRel);
196        self.ring_sidecar
197            .push_op(crate::sidecar_ops::sketch::OP_INSERT, 0);
198    }
199
200    /// Estimate cardinality (distinct count) via harmonic mean with
201    /// bias correction.
202    pub fn estimate(&self) -> u64 {
203        let m = self.m as f64;
204        let alpha = match self.m {
205            16 => 0.673,
206            32 => 0.697,
207            64 => 0.709,
208            _ => 0.7213 / (1.0 + 1.079 / m),
209        };
210        let mut sum = 0.0f64;
211        let mut zeros = 0u32;
212        for i in 0..self.m as usize {
213            let r = self.register(i).load(Ordering::Acquire);
214            if r == 0 { zeros += 1; }
215            sum += 2f64.powi(-(r as i32));
216        }
217        let raw = alpha * m * m / sum;
218        // Small-range correction: if estimate < 2.5 * m and there
219        // are empty registers, use linear counting.
220        let v = if raw <= 2.5 * m && zeros > 0 {
221            let z = zeros as f64;
222            (m * (m / z).ln()).round() as u64
223        } else {
224            raw.round() as u64
225        };
226        self.ring_sidecar
227            .push_op(crate::sidecar_ops::sketch::OP_QUERY, 0);
228        v
229    }
230
231    /// Reset all registers to zero (the empty state).
232    pub fn reset(&self) {
233        for i in 0..self.m as usize {
234            self.register(i).store(0, Ordering::Release);
235        }
236        self.ring_sidecar
237            .push_op(crate::sidecar_ops::sketch::OP_CLEAR, 0);
238    }
239
240    pub fn flush(&self) -> Result<(), HLLError> {
241        self.mmap.flush()?;
242        Ok(())
243    }
244    pub fn flush_async(&self) -> Result<(), HLLError> {
245        self.mmap.flush_async()?;
246        Ok(())
247    }
248}
249
250#[cfg(test)]
251mod tests {
252    use super::*;
253    use std::sync::Arc;
254    use std::thread;
255
256    fn tmp(name: &str) -> std::path::PathBuf {
257        let mut p = std::env::temp_dir();
258        let pid = std::process::id();
259        p.push(format!("subetha-hll-{name}-{pid}.bin"));
260        p
261    }
262
263    #[test]
264    fn create_initial_empty_estimate_is_zero() {
265        let p = tmp("init");
266        let h = SharedHyperLogLog::create(&p, 12).unwrap();
267        assert_eq!(h.precision(), 12);
268        assert_eq!(h.n_registers(), 4096);
269        assert_eq!(h.estimate(), 0);
270        std::fs::remove_file(&p).ok();
271    }
272
273    #[test]
274    fn invalid_precision_rejected() {
275        let p = tmp("invalid");
276        assert_eq!(
277            SharedHyperLogLog::create(&p, 3).err(),
278            Some(HLLError::InvalidPrecision)
279        );
280        assert_eq!(
281            SharedHyperLogLog::create(&p, 17).err(),
282            Some(HLLError::InvalidPrecision)
283        );
284        std::fs::remove_file(&p).ok();
285    }
286
287    #[test]
288    fn single_insert_estimate_is_one() {
289        let p = tmp("single");
290        let h = SharedHyperLogLog::create(&p, 12).unwrap();
291        h.insert(b"hello");
292        let est = h.estimate();
293        // Linear-counting small-range correction handles single
294        // insert; estimate should be exactly 1.
295        assert_eq!(est, 1, "single insert should estimate 1, got {est}");
296        std::fs::remove_file(&p).ok();
297    }
298
299    #[test]
300    fn idempotent_inserts_same_item() {
301        let p = tmp("idempotent");
302        let h = SharedHyperLogLog::create(&p, 12).unwrap();
303        for _ in 0..100 { h.insert(b"same-item"); }
304        let est = h.estimate();
305        assert_eq!(est, 1, "100 inserts of same item should estimate 1, got {est}");
306        std::fs::remove_file(&p).ok();
307    }
308
309    #[test]
310    fn distinct_count_within_error_bound_p12() {
311        // p=12 -> m=4096, std err 1.6%. For 1000 distinct items,
312        // expect estimate in roughly [950, 1050] (2 sigma).
313        let p = tmp("distinct-1k");
314        let h = SharedHyperLogLog::create(&p, 12).unwrap();
315        for i in 0..1000u32 {
316            h.insert(format!("item-{i:05}").as_bytes());
317        }
318        let est = h.estimate();
319        assert!(
320            (900..=1100).contains(&est),
321            "estimate {est} should be within 10% of 1000 (got {est})",
322        );
323        std::fs::remove_file(&p).ok();
324    }
325
326    #[test]
327    fn distinct_count_within_error_bound_p14() {
328        // p=14 -> m=16384, std err 0.8%. For 10000 distinct items,
329        // expect estimate within ~4% (2.5 sigma).
330        let p = tmp("distinct-10k");
331        let h = SharedHyperLogLog::create(&p, 14).unwrap();
332        for i in 0..10_000u32 {
333            h.insert(format!("k{i:06}").as_bytes());
334        }
335        let est = h.estimate();
336        assert!(
337            (9600..=10400).contains(&est),
338            "estimate {est} should be within 4% of 10000",
339        );
340        std::fs::remove_file(&p).ok();
341    }
342
343    #[test]
344    fn reset_clears_registers() {
345        let p = tmp("reset");
346        let h = SharedHyperLogLog::create(&p, 10).unwrap();
347        for i in 0..100u32 { h.insert(format!("k{i}").as_bytes()); }
348        assert!(h.estimate() > 0);
349        h.reset();
350        assert_eq!(h.estimate(), 0);
351        std::fs::remove_file(&p).ok();
352    }
353
354    #[test]
355    fn cross_handle_visibility() {
356        let p = tmp("cross-handle");
357        let w = SharedHyperLogLog::create(&p, 10).unwrap();
358        let r = SharedHyperLogLog::open(&p, 10).unwrap();
359        for i in 0..50u32 { w.insert(format!("k{i}").as_bytes()); }
360        let est_w = w.estimate();
361        let est_r = r.estimate();
362        assert_eq!(est_w, est_r);
363        // Roughly 50.
364        assert!((40..=60).contains(&est_r), "estimate {est_r} should be near 50");
365        std::fs::remove_file(&p).ok();
366    }
367
368    #[test]
369    fn config_mismatch_at_open_rejected() {
370        let p = tmp("mismatch");
371        let _w = SharedHyperLogLog::create(&p, 10).unwrap();
372        assert!(matches!(
373            SharedHyperLogLog::open(&p, 12),
374            Err(HLLError::LayoutMismatch)
375        ));
376        std::fs::remove_file(&p).ok();
377    }
378
379    #[test]
380    fn concurrent_inserters_correct_estimate() {
381        // 4 threads each insert 1000 distinct items (disjoint key
382        // spaces). Total = 4000. Estimate should be near 4000.
383        let p = tmp("concurrent");
384        let h: Arc<SharedHyperLogLog>
385            = Arc::new(SharedHyperLogLog::create(&p, 12).unwrap());
386        let mut handles = vec![];
387        for t in 0..4u32 {
388            let h = h.clone();
389            handles.push(thread::spawn(move || {
390                for i in 0..1000u32 {
391                    h.insert(format!("t{t}-i{i:05}").as_bytes());
392                }
393            }));
394        }
395        for h in handles { h.join().unwrap(); }
396        let est = h.estimate();
397        assert!(
398            (3700..=4300).contains(&est),
399            "concurrent inserts: estimate {est} should be within ~7% of 4000",
400        );
401        std::fs::remove_file(&p).ok();
402    }
403
404    #[test]
405    fn disk_persistence_survives_reopen() {
406        let p = tmp("disk");
407        {
408            let h = SharedHyperLogLog::create(&p, 10).unwrap();
409            for i in 0..100u32 { h.insert(format!("k{i}").as_bytes()); }
410            h.flush().unwrap();
411        }
412        let h2 = SharedHyperLogLog::open(&p, 10).unwrap();
413        let est = h2.estimate();
414        assert!((80..=120).contains(&est),
415            "reopened estimate {est} should be near 100");
416        std::fs::remove_file(&p).ok();
417    }
418}