From b63afa0ca069f1529a4bd8be3ffe3e764b61c7a0 Mon Sep 17 00:00:00 2001 From: GordonYuanyc Date: Sun, 1 Mar 2026 23:33:55 -0700 Subject: [PATCH 1/7] Octo structure with a rough benchmark, performance issue: not seeing the expected speedup --- Cargo.lock | 49 +++ Cargo.toml | 7 + benches/octo_speedup.rs | 275 +++++++++++++++++ src/sketch_framework/mod.rs | 6 + src/sketch_framework/octo.rs | 566 +++++++++++++++++++++++++++++++++++ src/sketches/count.rs | 68 +++++ src/sketches/countmin.rs | 67 +++++ src/sketches/hll.rs | 54 ++++ src/sketches/mod.rs | 7 + src/sketches/octo_delta.rs | 39 +++ 10 files changed, 1138 insertions(+) create mode 100644 benches/octo_speedup.rs create mode 100644 src/sketch_framework/octo.rs create mode 100644 src/sketches/octo_delta.rs diff --git a/Cargo.lock b/Cargo.lock index a1b0dbd..c29dfc0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -204,6 +204,17 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" +[[package]] +name = "core_affinity" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a034b3a7b624016c6e13f5df875747cc25f884156aad2abd12b6c46797971342" +dependencies = [ + "libc", + "num_cpus", + "winapi", +] + [[package]] name = "criterion" version = "0.5.1" @@ -474,6 +485,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "once_cell" version = "1.21.3" @@ -557,6 +578,21 @@ dependencies = [ "plotters-backend", ] +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "portable-atomic-util" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a9db96d7fa8782dd8c15ce32ffe8680bbd1e978a43bf51a34d39483540495f5" +dependencies = [ + "portable-atomic", +] + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -730,6 +766,17 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3160422bbd54dd5ecfdca71e5fd59b7b8fe2b1697ab2baf64f6d05dcc66d298" +[[package]] +name = "ringbuf" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe47b720588c8702e34b5979cb3271a8b1842c7cb6f57408efa70c779363488c" +dependencies = [ + "crossbeam-utils", + "portable-atomic", + "portable-atomic-util", +] + [[package]] name = "rmp" version = "0.8.14" @@ -850,11 +897,13 @@ version = "0.1.0" dependencies = [ "bytes", "clap", + "core_affinity", "criterion", "pcap", "prost", "prost-build", "rand", + "ringbuf", "rmp-serde", "serde", "serde-big-array", diff --git a/Cargo.toml b/Cargo.toml index 2de61d3..7d054cc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,8 @@ clap = { version = "4.0", features = ["derive"] } smallvec = "1.13.2" prost = "0.13" bytes = "1" +core_affinity = "0.8" +ringbuf = "0.4" [dev-dependencies] criterion = "0.5" @@ -63,6 +65,11 @@ harness = false name = "median_bench" harness = false +[[bench]] +name = "octo_speedup" +harness = false +path = "benches/octo_speedup.rs" + #[profile.release] #debug = true #strip = false diff --git a/benches/octo_speedup.rs b/benches/octo_speedup.rs new file mode 100644 index 0000000..70824aa --- /dev/null +++ b/benches/octo_speedup.rs @@ -0,0 +1,275 @@ +use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use rand::{rngs::StdRng, Rng, SeedableRng}; +use sketchlib_rust::{ + hash64_seeded, run_octo, CmOctoParent, CmOctoWorker, Count, CountMin, CountOctoParent, + CountOctoWorker, HllOctoParent, HllOctoWorker, HyperLogLog, OctoConfig, Regular, RegularPath, + SketchInput, Vector2D, +}; +use std::thread; + +const RNG_SEED: u64 = 0x0c70_2026_5eed_1234; +const CM_COUNT_ROWS: usize = 3; +const CM_COUNT_COLS: usize = 4096; +const CM_COUNT_INPUTS: usize = 1_000_000; +const HLL_INPUTS: usize = 500_000; +const DOMAIN_MASK: u64 = (1 << 20) - 1; +const QUEUE_CAPACITY: usize = 65_536; +const WORKER_COUNTS: [usize; 4] = [1, 2, 4, 8]; +const PARTITION_SEED: usize = 19; + +fn build_inputs(sample_count: usize) -> Vec> { + let mut rng = StdRng::seed_from_u64(RNG_SEED ^ (sample_count as u64)); + (0..sample_count) + .map(|_| SketchInput::U64(rng.random::() & DOMAIN_MASK)) + .collect() +} + +fn partition_indices(inputs: &[SketchInput<'_>], num_workers: usize) -> Vec> { + let mut partitions: Vec> = (0..num_workers).map(|_| Vec::new()).collect(); + for (idx, item) in inputs.iter().enumerate() { + let h = hash64_seeded(PARTITION_SEED, item) as usize; + partitions[h % num_workers].push(idx); + } + partitions +} + +fn run_regular_countmin_merge( + inputs: &[SketchInput<'_>], + num_workers: usize, +) -> CountMin, RegularPath> { + let partitions = partition_indices(inputs, num_workers); + let mut workers = Vec::with_capacity(num_workers); + + thread::scope(|s| { + let mut handles = Vec::with_capacity(num_workers); + for partition in partitions.iter() { + handles.push(s.spawn(move || { + let mut child = CountMin::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ); + for &idx in partition { + child.insert(&inputs[idx]); + } + child + })); + } + + for h in handles { + workers.push(h.join().expect("regular CountMin worker thread panicked")); + } + }); + + let mut parent = + CountMin::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS); + for worker in workers.iter() { + parent.merge(worker); + } + parent +} + +fn run_regular_count_merge( + inputs: &[SketchInput<'_>], + num_workers: usize, +) -> Count, RegularPath> { + let partitions = partition_indices(inputs, num_workers); + let mut workers = Vec::with_capacity(num_workers); + + thread::scope(|s| { + let mut handles = Vec::with_capacity(num_workers); + for partition in partitions.iter() { + handles.push(s.spawn(move || { + let mut child = Count::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ); + for &idx in partition { + child.insert(&inputs[idx]); + } + child + })); + } + + for h in handles { + workers.push(h.join().expect("regular Count worker thread panicked")); + } + }); + + let mut parent = + Count::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS); + for worker in workers.iter() { + parent.merge(worker); + } + parent +} + +fn run_regular_hll_merge(inputs: &[SketchInput<'_>], num_workers: usize) -> HyperLogLog { + let partitions = partition_indices(inputs, num_workers); + let mut workers = Vec::with_capacity(num_workers); + + thread::scope(|s| { + let mut handles = Vec::with_capacity(num_workers); + for partition in partitions.iter() { + handles.push(s.spawn(move || { + let mut child = HyperLogLog::::default(); + for &idx in partition { + child.insert(&inputs[idx]); + } + child + })); + } + + for h in handles { + workers.push(h.join().expect("regular HLL worker thread panicked")); + } + }); + + let mut parent = HyperLogLog::::default(); + for worker in workers.iter() { + parent.merge(worker); + } + parent +} + +fn bench_countmin_merge_compare(c: &mut Criterion) { + let inputs = build_inputs(CM_COUNT_INPUTS); + let mut group = c.benchmark_group("countmin_merge_compare"); + group.sample_size(10); + group.throughput(Throughput::Elements(inputs.len() as u64)); + + for &workers in &WORKER_COUNTS { + group.bench_with_input( + BenchmarkId::new("octo_style", workers), + &workers, + |b, &w| { + b.iter(|| { + let config = OctoConfig { + num_workers: w, + pin_cores: false, + queue_capacity: QUEUE_CAPACITY, + }; + let result = run_octo( + &inputs, + &config, + |_| CmOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), + || CmOctoParent { + sketch: CountMin::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), + }, + ); + black_box(result.parent.sketch); + }); + }, + ); + + group.bench_with_input( + BenchmarkId::new("regular_full_merge", workers), + &workers, + |b, &w| { + b.iter(|| { + let merged = run_regular_countmin_merge(&inputs, w); + black_box(merged); + }); + }, + ); + } + + group.finish(); +} + +fn bench_count_merge_compare(c: &mut Criterion) { + let inputs = build_inputs(CM_COUNT_INPUTS); + let mut group = c.benchmark_group("count_merge_compare"); + group.sample_size(10); + group.throughput(Throughput::Elements(inputs.len() as u64)); + + for &workers in &WORKER_COUNTS { + group.bench_with_input( + BenchmarkId::new("octo_style", workers), + &workers, + |b, &w| { + b.iter(|| { + let config = OctoConfig { + num_workers: w, + pin_cores: false, + queue_capacity: QUEUE_CAPACITY, + }; + let result = run_octo( + &inputs, + &config, + |_| CountOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), + || CountOctoParent { + sketch: Count::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), + }, + ); + black_box(result.parent.sketch); + }); + }, + ); + + group.bench_with_input( + BenchmarkId::new("regular_full_merge", workers), + &workers, + |b, &w| { + b.iter(|| { + let merged = run_regular_count_merge(&inputs, w); + black_box(merged); + }); + }, + ); + } + + group.finish(); +} + +fn bench_hll_merge_compare(c: &mut Criterion) { + let inputs = build_inputs(HLL_INPUTS); + let mut group = c.benchmark_group("hll_merge_compare"); + group.sample_size(10); + group.throughput(Throughput::Elements(inputs.len() as u64)); + + for &workers in &WORKER_COUNTS { + group.bench_with_input( + BenchmarkId::new("octo_style", workers), + &workers, + |b, &w| { + b.iter(|| { + let config = OctoConfig { + num_workers: w, + pin_cores: false, + queue_capacity: QUEUE_CAPACITY, + }; + let result = run_octo( + &inputs, + &config, + |_| HllOctoWorker::new(), + || HllOctoParent { + sketch: HyperLogLog::::default(), + }, + ); + black_box(result.parent.sketch); + }); + }, + ); + + group.bench_with_input( + BenchmarkId::new("regular_full_merge", workers), + &workers, + |b, &w| { + b.iter(|| { + let merged = run_regular_hll_merge(&inputs, w); + black_box(merged); + }); + }, + ); + } + + group.finish(); +} + +criterion_group!( + octo_speedup_benches, + bench_countmin_merge_compare, + bench_count_merge_compare, + bench_hll_merge_compare +); +criterion_main!(octo_speedup_benches); diff --git a/src/sketch_framework/mod.rs b/src/sketch_framework/mod.rs index 62db908..52a8b5f 100644 --- a/src/sketch_framework/mod.rs +++ b/src/sketch_framework/mod.rs @@ -35,6 +35,12 @@ pub use nitro::{NitroBatch, NitroEstimate, NitroTarget}; pub mod eh_univ_optimized; pub use eh_univ_optimized::{EHMapBucket, EHUnivMonBucket, EHUnivOptimized, EHUnivQueryResult}; +pub mod octo; +pub use octo::{ + CmOctoParent, CmOctoWorker, CountOctoParent, CountOctoWorker, HllOctoParent, HllOctoWorker, + OctoConfig, OctoParent, OctoResult, OctoWorker, run_octo, +}; + pub mod tumbling; pub use tumbling::{ FoldCMSConfig, FoldCSConfig, KLLConfig, SketchPool, TumblingWindow, TumblingWindowSketch, diff --git a/src/sketch_framework/octo.rs b/src/sketch_framework/octo.rs new file mode 100644 index 0000000..e9a6ec0 --- /dev/null +++ b/src/sketch_framework/octo.rs @@ -0,0 +1,566 @@ +//! OctoSketch multi-threaded sketch framework. +//! +//! Implements the parent-child delta-promotion architecture from OctoSketch (NSDI 2024). +//! Worker threads maintain lightweight child sketches with small counters and emit +//! compact delta entries via SPSC queues when counters overflow a promotion threshold. +//! An aggregator thread applies deltas to a full-precision parent sketch continuously. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::thread; + +use ringbuf::traits::{Consumer, Producer, Split}; +use ringbuf::HeapRb; + +use crate::{ + CountChild, CountDelta, CountMin, CountMinChild, CmDelta, Count, HllChild, HllDelta, + HyperLogLog, Regular, RegularPath, SketchInput, Vector2D, hash64_seeded, +}; + +/// Default SPSC ring buffer capacity per worker. +const DEFAULT_QUEUE_CAPACITY: usize = 65536; + +// --------------------------------------------------------------------------- +// Traits +// --------------------------------------------------------------------------- + +/// Worker-side trait: processes inputs and emits deltas. +pub trait OctoWorker: Send { + type Delta: Copy + Send + 'static; + + /// Process one input, emitting zero or more deltas via `emit`. + fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(Self::Delta)); +} + +/// Parent-side trait: absorbs deltas into a full-precision sketch. +pub trait OctoParent: Send { + type Delta: Copy + Send + 'static; + + /// Apply a single delta to the parent sketch. + fn apply(&mut self, delta: Self::Delta); +} + +// --------------------------------------------------------------------------- +// Configuration +// --------------------------------------------------------------------------- + +/// Configuration for `run_octo`. +pub struct OctoConfig { + /// Number of worker threads (default: 4). + pub num_workers: usize, + /// Pin worker threads to cores (default: true). + /// Worker i is pinned to core i, aggregator to core num_workers. + /// Silently skipped if pinning fails. + pub pin_cores: bool, + /// SPSC ring buffer capacity per worker (default: 65536). + pub queue_capacity: usize, +} + +impl Default for OctoConfig { + fn default() -> Self { + Self { + num_workers: 4, + pin_cores: true, + queue_capacity: DEFAULT_QUEUE_CAPACITY, + } + } +} + +/// Result of an `run_octo` execution. +pub struct OctoResult

{ + pub parent: P, +} + +// --------------------------------------------------------------------------- +// Concrete worker/parent implementations +// --------------------------------------------------------------------------- + +// -- CountMin -- + +/// OctoSketch worker backed by a `CountMinChild`. +pub struct CmOctoWorker { + child: CountMinChild, +} + +impl CmOctoWorker { + pub fn new(rows: usize, cols: usize) -> Self { + Self { + child: CountMinChild::with_dimensions(rows, cols), + } + } +} + +impl OctoWorker for CmOctoWorker { + type Delta = CmDelta; + + #[inline(always)] + fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(CmDelta)) { + self.child.insert_and_emit(input, emit); + } +} + +/// OctoSketch parent wrapping a full-precision `CountMin`. +pub struct CmOctoParent { + pub sketch: CountMin, RegularPath>, +} + +impl OctoParent for CmOctoParent { + type Delta = CmDelta; + + #[inline(always)] + fn apply(&mut self, delta: CmDelta) { + self.sketch.apply_delta(delta); + } +} + +// -- Count Sketch -- + +/// OctoSketch worker backed by a `CountChild`. +pub struct CountOctoWorker { + child: CountChild, +} + +impl CountOctoWorker { + pub fn new(rows: usize, cols: usize) -> Self { + Self { + child: CountChild::with_dimensions(rows, cols), + } + } +} + +impl OctoWorker for CountOctoWorker { + type Delta = CountDelta; + + #[inline(always)] + fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(CountDelta)) { + self.child.insert_and_emit(input, emit); + } +} + +/// OctoSketch parent wrapping a full-precision `Count`. +pub struct CountOctoParent { + pub sketch: Count, RegularPath>, +} + +impl OctoParent for CountOctoParent { + type Delta = CountDelta; + + #[inline(always)] + fn apply(&mut self, delta: CountDelta) { + self.sketch.apply_delta(delta); + } +} + +// -- HyperLogLog -- + +/// OctoSketch worker backed by an `HllChild`. +pub struct HllOctoWorker { + child: HllChild, +} + +impl HllOctoWorker { + pub fn new() -> Self { + Self { + child: HllChild::default(), + } + } +} + +impl Default for HllOctoWorker { + fn default() -> Self { + Self::new() + } +} + +impl OctoWorker for HllOctoWorker { + type Delta = HllDelta; + + #[inline(always)] + fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(HllDelta)) { + self.child.insert_and_emit(input, emit); + } +} + +/// OctoSketch parent wrapping a full-precision `HyperLogLog`. +pub struct HllOctoParent { + pub sketch: HyperLogLog, +} + +impl OctoParent for HllOctoParent { + type Delta = HllDelta; + + #[inline(always)] + fn apply(&mut self, delta: HllDelta) { + self.sketch.apply_delta(delta); + } +} + +// --------------------------------------------------------------------------- +// Core execution engine +// --------------------------------------------------------------------------- + +/// Hash-based input partitioning seed index (uses the last seed in SEEDLIST +/// to avoid collision with sketch-internal seeds which use indices 0..~5). +const PARTITION_SEED: usize = 19; + +/// Runs the OctoSketch multi-threaded insert protocol. +/// +/// 1. Partitions `inputs` across `config.num_workers` workers by hash. +/// 2. Each worker maintains a child sketch, emitting deltas via SPSC queues. +/// 3. The aggregator drains all queues continuously, applying deltas to the parent. +/// 4. Returns the fully-merged parent sketch. +/// +/// Uses `std::thread::scope` so `inputs` can have any lifetime (no `'static` needed). +pub fn run_octo( + inputs: &[SketchInput<'_>], + config: &OctoConfig, + worker_factory: impl Fn(usize) -> W + Send + Sync, + parent_factory: impl FnOnce() -> P, +) -> OctoResult

+where + W: OctoWorker, + P: OctoParent, +{ + let num_workers = config.num_workers.max(1); + let capacity = config.queue_capacity.max(1024); + + // Step 1: Partition inputs by hash into per-worker index lists. + let mut partitions: Vec> = (0..num_workers).map(|_| Vec::new()).collect(); + for (idx, item) in inputs.iter().enumerate() { + let h = hash64_seeded(PARTITION_SEED, item) as usize; + partitions[h % num_workers].push(idx); + } + + // Step 2: Create SPSC ring buffers. + let mut producers = Vec::with_capacity(num_workers); + let mut consumers = Vec::with_capacity(num_workers); + for _ in 0..num_workers { + let rb = HeapRb::::new(capacity); + let (prod, cons) = rb.split(); + producers.push(prod); + consumers.push(cons); + } + + // Step 3: Finish flags (one per worker). + let finished: Vec = (0..num_workers).map(|_| AtomicBool::new(false)).collect(); + + // Step 4: Run workers + aggregator inside a scoped thread block. + let mut parent = parent_factory(); + + thread::scope(|s| { + // Pin the aggregator (calling thread) to core num_workers if requested. + if config.pin_cores { + let _ = core_affinity::set_for_current(core_affinity::CoreId { id: num_workers }); + } + + // Spawn worker threads. + let wf = &worker_factory; + for (worker_id, (partition, producer)) in partitions + .iter() + .zip(producers.into_iter()) + .enumerate() + { + let finished_flag = &finished[worker_id]; + let pin_cores = config.pin_cores; + + s.spawn(move || { + // Pin to core. + if pin_cores { + let _ = core_affinity::set_for_current(core_affinity::CoreId { + id: worker_id, + }); + } + + let mut worker = wf(worker_id); + let mut prod = producer; + + for &input_idx in partition { + let input = &inputs[input_idx]; + worker.process(input, &mut |delta| { + // Spin-wait if the ring buffer is full (backpressure). + while prod.try_push(delta).is_err() { + std::hint::spin_loop(); + } + }); + } + + finished_flag.store(true, Ordering::Release); + }); + } + + // Aggregator collect loop: drain all consumer queues. + let all_done = || finished.iter().all(|f| f.load(Ordering::Acquire)); + + while !all_done() { + for cons in consumers.iter_mut() { + while let Some(delta) = cons.try_pop() { + parent.apply(delta); + } + } + } + + // Final drain pass after all workers have finished. + for cons in consumers.iter_mut() { + while let Some(delta) = cons.try_pop() { + parent.apply(delta); + } + } + }); + + OctoResult { parent } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::SketchInput; + + // ----------------------------------------------------------------------- + // Layer 2 unit tests: child sketches + apply_delta + // ----------------------------------------------------------------------- + + #[test] + fn cm_child_emits_delta_at_threshold() { + let mut child = CountMinChild::with_dimensions(3, 64); + let key = SketchInput::U64(42); + let mut deltas: Vec = Vec::new(); + + // Insert CM_PROMASK-1 times: no delta yet. + for _ in 0..(crate::CM_PROMASK - 1) { + child.insert_and_emit(&key, |d| deltas.push(d)); + } + assert!(deltas.is_empty(), "should not emit before threshold"); + + // One more insert triggers the delta. + child.insert_and_emit(&key, |d| deltas.push(d)); + assert_eq!(deltas.len(), 3, "should emit one delta per row (3 rows)"); + for d in &deltas { + assert_eq!(d.value, crate::CM_PROMASK); + } + } + + #[test] + fn cm_apply_delta_increments_parent() { + let mut parent = CountMin::, RegularPath>::with_dimensions(3, 64); + let delta = CmDelta { + row: 1, + col: 5, + value: 100, + }; + parent.apply_delta(delta); + assert_eq!(parent.as_storage().query_one_counter(1, 5), 100); + + parent.apply_delta(delta); + assert_eq!(parent.as_storage().query_one_counter(1, 5), 200); + } + + #[test] + fn count_child_emits_delta_at_threshold() { + let mut child = CountChild::with_dimensions(3, 64); + let key = SketchInput::U64(99); + let mut deltas: Vec = Vec::new(); + + // Insert enough times to trigger at least one delta. + for _ in 0..200 { + child.insert_and_emit(&key, |d| deltas.push(d)); + } + // With COUNT_PROMASK = 63 and 3 rows, after 200 inserts we + // should have emitted at least 3 deltas (one per row per threshold). + assert!( + deltas.len() >= 3, + "expected at least 3 deltas, got {}", + deltas.len() + ); + } + + #[test] + fn count_apply_delta_increments_parent() { + let mut parent = Count::, RegularPath>::with_dimensions(3, 64); + let delta = CountDelta { + row: 0, + col: 10, + value: -50, + }; + parent.apply_delta(delta); + assert_eq!(parent.as_storage().query_one_counter(0, 10), -50); + } + + #[test] + fn hll_child_emits_delta_on_improvement() { + let mut child = HllChild::default(); + let mut deltas: Vec = Vec::new(); + + // First insert should always emit (register goes from 0 to something > 0). + child.insert_and_emit(&SketchInput::U64(1), |d| deltas.push(d)); + assert_eq!(deltas.len(), 1, "first insert should emit a delta"); + + // Inserting the same value again should NOT emit (register unchanged). + let len_before = deltas.len(); + child.insert_and_emit(&SketchInput::U64(1), |d| deltas.push(d)); + assert_eq!(deltas.len(), len_before, "duplicate should not emit"); + } + + #[test] + fn hll_apply_delta_updates_register() { + let mut parent = HyperLogLog::::default(); + let delta = HllDelta { + pos: 100, + value: 10, + }; + parent.apply_delta(delta); + + // Apply a smaller value — should not change. + let smaller = HllDelta { + pos: 100, + value: 5, + }; + parent.apply_delta(smaller); + + // Apply a larger value — should update. + let larger = HllDelta { + pos: 100, + value: 15, + }; + parent.apply_delta(larger); + } + + // ----------------------------------------------------------------------- + // Layer 3 integration tests: run_octo + // ----------------------------------------------------------------------- + + #[test] + fn run_octo_cm_matches_single_thread() { + let rows = 3; + let cols = 4096; + let n = 100_000u64; + + let inputs: Vec> = (0..n).map(|i| SketchInput::U64(i % 1024)).collect(); + + // Single-threaded reference. + let mut reference = CountMin::, RegularPath>::with_dimensions(rows, cols); + for input in &inputs { + reference.insert(input); + } + + // Multi-threaded via OctoSketch. + let config = OctoConfig { + num_workers: 4, + pin_cores: false, // don't pin in tests (CI may have few cores) + queue_capacity: 8192, + }; + let result = run_octo( + &inputs, + &config, + |_| CmOctoWorker::new(rows, cols), + || CmOctoParent { + sketch: CountMin::with_dimensions(rows, cols), + }, + ); + + // Verify estimates are close. + // OctoSketch child uses u8 counters, so there may be slight differences + // due to the promotion threshold batching, but the total counts should match. + for key_val in 0u64..1024 { + let key = SketchInput::U64(key_val); + let ref_est = reference.estimate(&key); + let octo_est = result.parent.sketch.estimate(&key); + // The octo estimate should be <= reference (some counts may be + // lost in the u8 remainder that didn't reach PROMASK). + assert!( + (ref_est - octo_est).abs() < 200, + "key {key_val}: ref={ref_est}, octo={octo_est}, diff={}", + (ref_est - octo_est).abs() + ); + } + } + + #[test] + fn run_octo_hll_cardinality() { + let n = 50_000u64; + let inputs: Vec> = (0..n).map(SketchInput::U64).collect(); + + let config = OctoConfig { + num_workers: 4, + pin_cores: false, + queue_capacity: 16384, + }; + + let result = run_octo( + &inputs, + &config, + |_| HllOctoWorker::new(), + || HllOctoParent { + sketch: HyperLogLog::::default(), + }, + ); + + let estimate = result.parent.sketch.estimate(); + let truth = n as f64; + let error = (estimate as f64 - truth).abs() / truth; + assert!( + error < 0.05, + "HLL cardinality error {error:.4} exceeded 5% (truth {truth}, estimate {estimate})" + ); + } + + #[test] + fn run_octo_count_sketch_basic() { + let rows = 3; + let cols = 4096; + let n = 50_000u64; + + let inputs: Vec> = (0..n).map(|i| SketchInput::U64(i % 512)).collect(); + + let config = OctoConfig { + num_workers: 2, + pin_cores: false, + queue_capacity: 8192, + }; + + let result = run_octo( + &inputs, + &config, + |_| CountOctoWorker::new(rows, cols), + || CountOctoParent { + sketch: Count::with_dimensions(rows, cols), + }, + ); + + // Each key appears ~97 times. Check a sample of keys. + let expected_per_key = (n / 512) as f64; + for key_val in [0u64, 100, 255, 511] { + let key = SketchInput::U64(key_val); + let est = result.parent.sketch.estimate(&key); + assert!( + (est - expected_per_key).abs() < expected_per_key * 0.5, + "key {key_val}: estimate={est}, expected~{expected_per_key}" + ); + } + } + + #[test] + fn run_octo_single_worker() { + // Edge case: single worker should still work. + let inputs: Vec> = (0..1000u64).map(SketchInput::U64).collect(); + + let config = OctoConfig { + num_workers: 1, + pin_cores: false, + queue_capacity: 4096, + }; + + let result = run_octo( + &inputs, + &config, + |_| HllOctoWorker::new(), + || HllOctoParent { + sketch: HyperLogLog::::default(), + }, + ); + + let estimate = result.parent.sketch.estimate(); + assert!( + estimate > 900 && estimate < 1100, + "single-worker HLL estimate {estimate} out of range for 1000 distinct keys" + ); + } +} diff --git a/src/sketches/count.rs b/src/sketches/count.rs index 2e0a0ad..059ddcd 100644 --- a/src/sketches/count.rs +++ b/src/sketches/count.rs @@ -716,6 +716,74 @@ impl CountL2HH { } } +// --------------------------------------------------------------------------- +// OctoSketch child sketch for multi-threaded delta-based promotion. +// --------------------------------------------------------------------------- + +use crate::octo_delta::{COUNT_PROMASK, CountDelta}; + +/// Lightweight Count sketch child with i8 counters. +/// Used by OctoSketch workers: accumulates signed counts and emits +/// `CountDelta` entries when |counter| reaches the promotion threshold. +pub struct CountChild { + counts: Vector2D, +} + +impl CountChild { + /// Creates a child sketch matching the parent's dimensions. + pub fn with_dimensions(rows: usize, cols: usize) -> Self { + let mut counts = Vector2D::init(rows, cols); + counts.fill(0i8); + CountChild { counts } + } + + pub fn rows(&self) -> usize { + self.counts.rows() + } + + pub fn cols(&self) -> usize { + self.counts.cols() + } + + /// Insert a key with the Count sketch sign convention. + /// Emits `CountDelta` entries via `emit` when |counter| >= `COUNT_PROMASK`. + #[inline(always)] + pub fn insert_and_emit(&mut self, value: &SketchInput, mut emit: impl FnMut(CountDelta)) { + let rows = self.counts.rows(); + let cols = self.counts.cols(); + let data = self.counts.as_mut_slice(); + for r in 0..rows { + let hashed = hash64_seeded(r, value); + let col = ((hashed & LOWER_32_MASK) as usize) % cols; + let sign: i8 = if ((hashed >> 63) & 1) == 1 { 1 } else { -1 }; + let cell = &mut data[r * cols + col]; + *cell = cell.saturating_add(sign); + if cell.unsigned_abs() >= COUNT_PROMASK { + emit(CountDelta { + row: r as u16, + col: col as u16, + value: *cell, + }); + *cell = 0; + } + } + } +} + +/// Apply a `CountDelta` to a full-precision parent Count sketch. +impl Count +where + S::Counter: Copy + std::ops::AddAssign + From, +{ + pub fn apply_delta(&mut self, delta: CountDelta) { + self.counts.increment_by_row( + delta.row as usize, + delta.col as usize, + S::Counter::from(delta.value as i32), + ); + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/sketches/countmin.rs b/src/sketches/countmin.rs index 30be6cd..4a38073 100644 --- a/src/sketches/countmin.rs +++ b/src/sketches/countmin.rs @@ -469,6 +469,73 @@ impl NitroTarget for CountMin, FastPath, H> { } } +// --------------------------------------------------------------------------- +// OctoSketch child sketch for multi-threaded delta-based promotion. +// --------------------------------------------------------------------------- + +use crate::octo_delta::{CM_PROMASK, CmDelta}; + +/// Lightweight CountMin child sketch with u8 counters. +/// Used by OctoSketch workers: accumulates counts in u8 cells and emits +/// `CmDelta` entries when a counter reaches the promotion threshold. +pub struct CountMinChild { + counts: Vector2D, +} + +impl CountMinChild { + /// Creates a child sketch matching the parent's dimensions. + pub fn with_dimensions(rows: usize, cols: usize) -> Self { + let mut counts = Vector2D::init(rows, cols); + counts.fill(0u8); + CountMinChild { counts } + } + + pub fn rows(&self) -> usize { + self.counts.rows() + } + + pub fn cols(&self) -> usize { + self.counts.cols() + } + + /// Insert a key, emitting `CmDelta` entries via `emit` when a counter + /// reaches `CM_PROMASK`. The counter is reset to 0 after emission. + #[inline(always)] + pub fn insert_and_emit(&mut self, value: &SketchInput, mut emit: impl FnMut(CmDelta)) { + let rows = self.counts.rows(); + let cols = self.counts.cols(); + let data = self.counts.as_mut_slice(); + for r in 0..rows { + let hashed = hash64_seeded(r, value); + let col = ((hashed & LOWER_32_MASK) as usize) % cols; + let cell = &mut data[r * cols + col]; + *cell = cell.wrapping_add(1); + if *cell >= CM_PROMASK { + emit(CmDelta { + row: r as u16, + col: col as u16, + value: *cell, + }); + *cell = 0; + } + } + } +} + +/// Apply a `CmDelta` to a full-precision parent CountMin sketch. +impl CountMin +where + S::Counter: Copy + std::ops::AddAssign + From, +{ + pub fn apply_delta(&mut self, delta: CmDelta) { + self.counts.increment_by_row( + delta.row as usize, + delta.col as usize, + S::Counter::from(delta.value as i32), + ); + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/sketches/hll.rs b/src/sketches/hll.rs index 2d0ab79..4eaa02c 100644 --- a/src/sketches/hll.rs +++ b/src/sketches/hll.rs @@ -331,6 +331,60 @@ impl HyperLogLogHIP { } } } +// --------------------------------------------------------------------------- +// OctoSketch child sketch for multi-threaded delta-based promotion. +// --------------------------------------------------------------------------- + +use crate::octo_delta::HllDelta; + +/// Lightweight HLL child sketch backed by a fixed register array. +/// Emits `HllDelta` whenever a register is improved (PROMASK = 0, +/// meaning every improvement is propagated immediately). +pub struct HllChild { + registers: Box<[u8; NUM_REGISTERS]>, +} + +impl Default for HllChild { + fn default() -> Self { + Self { + registers: Box::new([0u8; NUM_REGISTERS]), + } + } +} + +impl HllChild { + /// Insert using a pre-computed hash, emitting `HllDelta` on register improvement. + #[inline(always)] + pub fn insert_with_hash_and_emit(&mut self, hashed_val: u64, mut emit: impl FnMut(HllDelta)) { + let bucket_num = ((hashed_val >> HLL_Q) & HLL_P_MASK) as usize; + let leading_zero = ((hashed_val << HLL_P) + HLL_P_MASK).leading_zeros() as u8 + 1; + if leading_zero > self.registers[bucket_num] { + self.registers[bucket_num] = leading_zero; + emit(HllDelta { + pos: bucket_num as u16, + value: leading_zero, + }); + } + } + + /// Insert a key, emitting `HllDelta` on register improvement. + #[inline(always)] + pub fn insert_and_emit(&mut self, obj: &SketchInput, emit: impl FnMut(HllDelta)) { + let hashed_val = hash64_seeded(CANONICAL_HASH_SEED, obj); + self.insert_with_hash_and_emit(hashed_val, emit); + } +} + +/// Apply an `HllDelta` to a full-precision parent HyperLogLog sketch. +impl HyperLogLog { + pub fn apply_delta(&mut self, delta: HllDelta) { + let pos = delta.pos as usize; + if delta.value > self.registers[pos] { + self.registers[pos] = delta.value; + } + } +} + #[cfg(test)] mod tests { diff --git a/src/sketches/mod.rs b/src/sketches/mod.rs index e504fa7..a174355 100644 --- a/src/sketches/mod.rs +++ b/src/sketches/mod.rs @@ -44,6 +44,13 @@ pub use cms_heap::CMSHeap; pub mod cs_heap; pub use cs_heap::CSHeap; +pub mod octo_delta; +pub use octo_delta::{CM_PROMASK, COUNT_PROMASK, CmDelta, CountDelta, HLL_PROMASK, HllDelta}; + +pub use countmin::CountMinChild; +pub use count::CountChild; +pub use hll::HllChild; + pub mod fold_cms; pub use fold_cms::{FoldCMS, FoldCell, FoldEntry}; diff --git a/src/sketches/octo_delta.rs b/src/sketches/octo_delta.rs new file mode 100644 index 0000000..04a3fb6 --- /dev/null +++ b/src/sketches/octo_delta.rs @@ -0,0 +1,39 @@ +//! Delta entry types for OctoSketch-style multi-threaded sketch updates. +//! +//! Each delta represents an accumulated counter change emitted by a child +//! worker sketch when a local counter overflows the promotion threshold (PROMASK). + +/// CountMin promotion threshold: emit delta when u8 counter >= 127. +pub const CM_PROMASK: u8 = 0x7f; + +/// Count sketch promotion threshold: emit delta when |i8 counter| >= 63. +pub const COUNT_PROMASK: u8 = 0x3f; + +/// HLL promotion threshold: 0 means every register improvement is emitted. +pub const HLL_PROMASK: u8 = 0; + +/// Delta emitted by a CountMin child worker. +/// Represents an accumulated unsigned count for a single cell. +#[derive(Clone, Copy, Debug)] +pub struct CmDelta { + pub row: u16, + pub col: u16, + pub value: u8, +} + +/// Delta emitted by a Count sketch child worker. +/// Represents a signed accumulated count for a single cell. +#[derive(Clone, Copy, Debug)] +pub struct CountDelta { + pub row: u16, + pub col: u16, + pub value: i8, +} + +/// Delta emitted by an HLL child worker. +/// Represents a register improvement (max-register semantics). +#[derive(Clone, Copy, Debug)] +pub struct HllDelta { + pub pos: u16, + pub value: u8, +} From ca15b5d297d127c1341ad0c3ead64befce5e215d Mon Sep 17 00:00:00 2001 From: GordonYuanyc Date: Mon, 2 Mar 2026 15:02:23 -0700 Subject: [PATCH 2/7] remove ringbuf and aggregator polling; use MPSC for worker and aggregator communication; don't see speedup when increasing of thread num --- Cargo.lock | 27 -- Cargo.toml | 1 - benches/octo_speedup.rs | 192 ++++++++++-- src/sketch_framework/mod.rs | 2 +- src/sketch_framework/octo.rs | 574 ++++++++++++++++++++++++++++++----- src/sketches/mod.rs | 2 +- 6 files changed, 665 insertions(+), 133 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c29dfc0..c2773a0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -578,21 +578,6 @@ dependencies = [ "plotters-backend", ] -[[package]] -name = "portable-atomic" -version = "1.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" - -[[package]] -name = "portable-atomic-util" -version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a9db96d7fa8782dd8c15ce32ffe8680bbd1e978a43bf51a34d39483540495f5" -dependencies = [ - "portable-atomic", -] - [[package]] name = "ppv-lite86" version = "0.2.21" @@ -766,17 +751,6 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3160422bbd54dd5ecfdca71e5fd59b7b8fe2b1697ab2baf64f6d05dcc66d298" -[[package]] -name = "ringbuf" -version = "0.4.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe47b720588c8702e34b5979cb3271a8b1842c7cb6f57408efa70c779363488c" -dependencies = [ - "crossbeam-utils", - "portable-atomic", - "portable-atomic-util", -] - [[package]] name = "rmp" version = "0.8.14" @@ -903,7 +877,6 @@ dependencies = [ "prost", "prost-build", "rand", - "ringbuf", "rmp-serde", "serde", "serde-big-array", diff --git a/Cargo.toml b/Cargo.toml index 7d054cc..e437ba6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,7 +15,6 @@ smallvec = "1.13.2" prost = "0.13" bytes = "1" core_affinity = "0.8" -ringbuf = "0.4" [dev-dependencies] criterion = "0.5" diff --git a/benches/octo_speedup.rs b/benches/octo_speedup.rs index 70824aa..a47fa9c 100644 --- a/benches/octo_speedup.rs +++ b/benches/octo_speedup.rs @@ -1,10 +1,11 @@ -use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; -use rand::{rngs::StdRng, Rng, SeedableRng}; +use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; +use rand::{Rng, SeedableRng, rngs::StdRng}; use sketchlib_rust::{ - hash64_seeded, run_octo, CmOctoParent, CmOctoWorker, Count, CountMin, CountOctoParent, - CountOctoWorker, HllOctoParent, HllOctoWorker, HyperLogLog, OctoConfig, Regular, RegularPath, - SketchInput, Vector2D, + CmOctoParent, CmOctoWorker, Count, CountMin, CountOctoParent, CountOctoWorker, HllOctoParent, + HllOctoWorker, HyperLogLog, OctoConfig, Regular, RegularPath, SketchInput, Vector2D, + hash64_seeded, run_octo, }; +use std::sync::mpsc; use std::thread; const RNG_SEED: u64 = 0x0c70_2026_5eed_1234; @@ -16,6 +17,7 @@ const DOMAIN_MASK: u64 = (1 << 20) - 1; const QUEUE_CAPACITY: usize = 65_536; const WORKER_COUNTS: [usize; 4] = [1, 2, 4, 8]; const PARTITION_SEED: usize = 19; +const MERGE_INTERVAL: usize = 10_000; fn build_inputs(sample_count: usize) -> Vec> { let mut rng = StdRng::seed_from_u64(RNG_SEED ^ (sample_count as u64)); @@ -24,39 +26,66 @@ fn build_inputs(sample_count: usize) -> Vec> { .collect() } -fn partition_indices(inputs: &[SketchInput<'_>], num_workers: usize) -> Vec> { - let mut partitions: Vec> = (0..num_workers).map(|_| Vec::new()).collect(); - for (idx, item) in inputs.iter().enumerate() { - let h = hash64_seeded(PARTITION_SEED, item) as usize; - partitions[h % num_workers].push(idx); - } - partitions +#[inline(always)] +fn should_merge(number: usize, interval: usize, thread_num: usize, thread_id: usize) -> bool { + number != 0 && number % (interval * thread_num) == interval * thread_id } fn run_regular_countmin_merge( inputs: &[SketchInput<'_>], num_workers: usize, ) -> CountMin, RegularPath> { - let partitions = partition_indices(inputs, num_workers); - let mut workers = Vec::with_capacity(num_workers); + let mut workers = Vec::new(); thread::scope(|s| { + let mut worker_txs = Vec::with_capacity(num_workers); let mut handles = Vec::with_capacity(num_workers); - for partition in partitions.iter() { + for worker_id in 0..num_workers { + let (tx, rx) = mpsc::channel::>(); + worker_txs.push(tx); handles.push(s.spawn(move || { let mut child = CountMin::, RegularPath>::with_dimensions( CM_COUNT_ROWS, CM_COUNT_COLS, ); - for &idx in partition { - child.insert(&inputs[idx]); + let mut partials = Vec::new(); + let mut number = 0usize; + while let Ok(msg) = rx.recv() { + match msg { + Some(idx) => { + child.insert(&inputs[idx]); + number += 1; + if should_merge(number, MERGE_INTERVAL, num_workers, worker_id) { + partials.push(child); + child = CountMin::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ); + } + } + None => break, + } } - child + partials.push(child); + partials })); } + s.spawn(move || { + for (idx, item) in inputs.iter().enumerate() { + let h = hash64_seeded(PARTITION_SEED, item) as usize; + let worker_id = h % num_workers; + worker_txs[worker_id] + .send(Some(idx)) + .expect("regular CountMin worker receiver dropped unexpectedly"); + } + for tx in worker_txs { + let _ = tx.send(None); + } + }); + for h in handles { - workers.push(h.join().expect("regular CountMin worker thread panicked")); + workers.extend(h.join().expect("regular CountMin worker thread panicked")); } }); @@ -72,26 +101,57 @@ fn run_regular_count_merge( inputs: &[SketchInput<'_>], num_workers: usize, ) -> Count, RegularPath> { - let partitions = partition_indices(inputs, num_workers); - let mut workers = Vec::with_capacity(num_workers); + let mut workers = Vec::new(); thread::scope(|s| { + let mut worker_txs = Vec::with_capacity(num_workers); let mut handles = Vec::with_capacity(num_workers); - for partition in partitions.iter() { + for worker_id in 0..num_workers { + let (tx, rx) = mpsc::channel::>(); + worker_txs.push(tx); handles.push(s.spawn(move || { let mut child = Count::, RegularPath>::with_dimensions( CM_COUNT_ROWS, CM_COUNT_COLS, ); - for &idx in partition { - child.insert(&inputs[idx]); + let mut partials = Vec::new(); + let mut number = 0usize; + while let Ok(msg) = rx.recv() { + match msg { + Some(idx) => { + child.insert(&inputs[idx]); + number += 1; + if should_merge(number, MERGE_INTERVAL, num_workers, worker_id) { + partials.push(child); + child = Count::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ); + } + } + None => break, + } } - child + partials.push(child); + partials })); } + s.spawn(move || { + for (idx, item) in inputs.iter().enumerate() { + let h = hash64_seeded(PARTITION_SEED, item) as usize; + let worker_id = h % num_workers; + worker_txs[worker_id] + .send(Some(idx)) + .expect("regular Count worker receiver dropped unexpectedly"); + } + for tx in worker_txs { + let _ = tx.send(None); + } + }); + for h in handles { - workers.push(h.join().expect("regular Count worker thread panicked")); + workers.extend(h.join().expect("regular Count worker thread panicked")); } }); @@ -104,23 +164,51 @@ fn run_regular_count_merge( } fn run_regular_hll_merge(inputs: &[SketchInput<'_>], num_workers: usize) -> HyperLogLog { - let partitions = partition_indices(inputs, num_workers); - let mut workers = Vec::with_capacity(num_workers); + let mut workers = Vec::new(); thread::scope(|s| { + let mut worker_txs = Vec::with_capacity(num_workers); let mut handles = Vec::with_capacity(num_workers); - for partition in partitions.iter() { + for worker_id in 0..num_workers { + let (tx, rx) = mpsc::channel::>(); + worker_txs.push(tx); handles.push(s.spawn(move || { let mut child = HyperLogLog::::default(); - for &idx in partition { - child.insert(&inputs[idx]); + let mut partials = Vec::new(); + let mut number = 0usize; + while let Ok(msg) = rx.recv() { + match msg { + Some(idx) => { + child.insert(&inputs[idx]); + number += 1; + if should_merge(number, MERGE_INTERVAL, num_workers, worker_id) { + partials.push(child); + child = HyperLogLog::::default(); + } + } + None => break, + } } - child + partials.push(child); + partials })); } + s.spawn(move || { + for (idx, item) in inputs.iter().enumerate() { + let h = hash64_seeded(PARTITION_SEED, item) as usize; + let worker_id = h % num_workers; + worker_txs[worker_id] + .send(Some(idx)) + .expect("regular HLL worker receiver dropped unexpectedly"); + } + for tx in worker_txs { + let _ = tx.send(None); + } + }); + for h in handles { - workers.push(h.join().expect("regular HLL worker thread panicked")); + workers.extend(h.join().expect("regular HLL worker thread panicked")); } }); @@ -137,6 +225,23 @@ fn bench_countmin_merge_compare(c: &mut Criterion) { group.sample_size(10); group.throughput(Throughput::Elements(inputs.len() as u64)); + group.bench_function("single_thread_insert", |b| { + b.iter_with_setup( + || { + CountMin::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ) + }, + |mut sketch| { + for input in &inputs { + sketch.insert(input); + } + black_box(sketch); + }, + ); + }); + for &workers in &WORKER_COUNTS { group.bench_with_input( BenchmarkId::new("octo_style", workers), @@ -182,6 +287,18 @@ fn bench_count_merge_compare(c: &mut Criterion) { group.sample_size(10); group.throughput(Throughput::Elements(inputs.len() as u64)); + group.bench_function("single_thread_insert", |b| { + b.iter_with_setup( + || Count::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), + |mut sketch| { + for input in &inputs { + sketch.insert(input); + } + black_box(sketch); + }, + ); + }); + for &workers in &WORKER_COUNTS { group.bench_with_input( BenchmarkId::new("octo_style", workers), @@ -227,6 +344,15 @@ fn bench_hll_merge_compare(c: &mut Criterion) { group.sample_size(10); group.throughput(Throughput::Elements(inputs.len() as u64)); + group.bench_function("single_thread_insert", |b| { + b.iter_with_setup(HyperLogLog::::default, |mut sketch| { + for input in &inputs { + sketch.insert(input); + } + black_box(sketch); + }); + }); + for &workers in &WORKER_COUNTS { group.bench_with_input( BenchmarkId::new("octo_style", workers), diff --git a/src/sketch_framework/mod.rs b/src/sketch_framework/mod.rs index 52a8b5f..7bc87a7 100644 --- a/src/sketch_framework/mod.rs +++ b/src/sketch_framework/mod.rs @@ -38,7 +38,7 @@ pub use eh_univ_optimized::{EHMapBucket, EHUnivMonBucket, EHUnivOptimized, EHUni pub mod octo; pub use octo::{ CmOctoParent, CmOctoWorker, CountOctoParent, CountOctoWorker, HllOctoParent, HllOctoWorker, - OctoConfig, OctoParent, OctoResult, OctoWorker, run_octo, + OctoConfig, OctoParent, OctoResult, OctoRuntime, OctoWorker, OwnedSketchInput, run_octo, }; pub mod tumbling; diff --git a/src/sketch_framework/octo.rs b/src/sketch_framework/octo.rs index e9a6ec0..0c3b9fc 100644 --- a/src/sketch_framework/octo.rs +++ b/src/sketch_framework/octo.rs @@ -2,22 +2,22 @@ //! //! Implements the parent-child delta-promotion architecture from OctoSketch (NSDI 2024). //! Worker threads maintain lightweight child sketches with small counters and emit -//! compact delta entries via SPSC queues when counters overflow a promotion threshold. -//! An aggregator thread applies deltas to a full-precision parent sketch continuously. +//! compact delta entries via an MPSC channel when counters overflow a promotion +//! threshold. An aggregator thread applies deltas to a full-precision parent sketch. -use std::sync::atomic::{AtomicBool, Ordering}; +use std::marker::PhantomData; +use std::sync::Arc; +use std::sync::mpsc; use std::thread; -use ringbuf::traits::{Consumer, Producer, Split}; -use ringbuf::HeapRb; - use crate::{ - CountChild, CountDelta, CountMin, CountMinChild, CmDelta, Count, HllChild, HllDelta, + CmDelta, Count, CountChild, CountDelta, CountMin, CountMinChild, HllChild, HllDelta, HyperLogLog, Regular, RegularPath, SketchInput, Vector2D, hash64_seeded, }; -/// Default SPSC ring buffer capacity per worker. +/// Legacy queue capacity default retained for config compatibility. const DEFAULT_QUEUE_CAPACITY: usize = 65536; +const WORKER_DISPATCH_BATCH_SIZE: usize = 64; // --------------------------------------------------------------------------- // Traits @@ -51,7 +51,8 @@ pub struct OctoConfig { /// Worker i is pinned to core i, aggregator to core num_workers. /// Silently skipped if pinning fails. pub pin_cores: bool, - /// SPSC ring buffer capacity per worker (default: 65536). + /// Legacy queue capacity retained for compatibility (default: 65536). + /// Currently unused by the unbounded MPSC transport. pub queue_capacity: usize, } @@ -70,6 +71,274 @@ pub struct OctoResult

{ pub parent: P, } +/// Owned variant of `SketchInput` for cross-thread transport. +#[derive(Clone, Debug)] +pub enum OwnedSketchInput { + I8(i8), + I16(i16), + I32(i32), + I64(i64), + I128(i128), + ISIZE(isize), + U8(u8), + U16(u16), + U32(u32), + U64(u64), + U128(u128), + USIZE(usize), + F32(f32), + F64(f64), + String(String), + Bytes(Vec), +} + +impl OwnedSketchInput { + fn as_sketch_input(&self) -> SketchInput<'_> { + match self { + OwnedSketchInput::I8(v) => SketchInput::I8(*v), + OwnedSketchInput::I16(v) => SketchInput::I16(*v), + OwnedSketchInput::I32(v) => SketchInput::I32(*v), + OwnedSketchInput::I64(v) => SketchInput::I64(*v), + OwnedSketchInput::I128(v) => SketchInput::I128(*v), + OwnedSketchInput::ISIZE(v) => SketchInput::ISIZE(*v), + OwnedSketchInput::U8(v) => SketchInput::U8(*v), + OwnedSketchInput::U16(v) => SketchInput::U16(*v), + OwnedSketchInput::U32(v) => SketchInput::U32(*v), + OwnedSketchInput::U64(v) => SketchInput::U64(*v), + OwnedSketchInput::U128(v) => SketchInput::U128(*v), + OwnedSketchInput::USIZE(v) => SketchInput::USIZE(*v), + OwnedSketchInput::F32(v) => SketchInput::F32(*v), + OwnedSketchInput::F64(v) => SketchInput::F64(*v), + OwnedSketchInput::String(s) => SketchInput::Str(s.as_str()), + OwnedSketchInput::Bytes(b) => SketchInput::Bytes(b.as_slice()), + } + } +} + +impl From<&SketchInput<'_>> for OwnedSketchInput { + fn from(value: &SketchInput<'_>) -> Self { + match value { + SketchInput::I8(v) => OwnedSketchInput::I8(*v), + SketchInput::I16(v) => OwnedSketchInput::I16(*v), + SketchInput::I32(v) => OwnedSketchInput::I32(*v), + SketchInput::I64(v) => OwnedSketchInput::I64(*v), + SketchInput::I128(v) => OwnedSketchInput::I128(*v), + SketchInput::ISIZE(v) => OwnedSketchInput::ISIZE(*v), + SketchInput::U8(v) => OwnedSketchInput::U8(*v), + SketchInput::U16(v) => OwnedSketchInput::U16(*v), + SketchInput::U32(v) => OwnedSketchInput::U32(*v), + SketchInput::U64(v) => OwnedSketchInput::U64(*v), + SketchInput::U128(v) => OwnedSketchInput::U128(*v), + SketchInput::USIZE(v) => OwnedSketchInput::USIZE(*v), + SketchInput::F32(v) => OwnedSketchInput::F32(*v), + SketchInput::F64(v) => OwnedSketchInput::F64(*v), + SketchInput::Str(s) => OwnedSketchInput::String((*s).to_string()), + SketchInput::String(s) => OwnedSketchInput::String(s.clone()), + SketchInput::Bytes(b) => OwnedSketchInput::Bytes(b.to_vec()), + } + } +} + +enum IngressMsg { + Data(OwnedSketchInput), + End, +} + +enum WorkerMsg { + Batch(Vec), + End, +} + +#[inline(always)] +fn worker_index(hash: u64, num_workers: usize) -> usize { + if num_workers.is_power_of_two() { + (hash as usize) & (num_workers - 1) + } else { + (hash as usize) % num_workers + } +} + +#[inline] +fn flush_worker_batch( + worker_input_txs: &[mpsc::Sender], + pending: &mut [Vec], + worker_id: usize, +) { + if pending[worker_id].is_empty() { + return; + } + let batch = std::mem::take(&mut pending[worker_id]); + worker_input_txs[worker_id] + .send(WorkerMsg::Batch(batch)) + .expect("worker receiver dropped unexpectedly"); +} + +#[inline] +fn flush_all_worker_batches( + worker_input_txs: &[mpsc::Sender], + pending: &mut [Vec], +) { + for worker_id in 0..pending.len() { + flush_worker_batch(worker_input_txs, pending, worker_id); + } +} + +#[inline] +fn send_end_to_all_workers(worker_input_txs: &[mpsc::Sender]) { + for tx in worker_input_txs { + let _ = tx.send(WorkerMsg::End); + } +} + +/// Streaming Octo runtime that accepts incremental inserts and finalizes into a parent sketch. +pub struct OctoRuntime +where + W: OctoWorker + 'static, + P: OctoParent + Send + 'static, +{ + ingress_tx: mpsc::Sender, + dispatcher_handle: Option>, + worker_handles: Vec>, + aggregator_handle: Option>, + _worker_marker: PhantomData, +} + +impl OctoRuntime +where + W: OctoWorker + 'static, + P: OctoParent + Send + 'static, +{ + pub fn new(config: &OctoConfig, worker_factory: F, parent_factory: PF) -> Self + where + F: Fn(usize) -> W + Send + Sync + 'static, + PF: FnOnce() -> P + Send + 'static, + { + let num_workers = config.num_workers.max(1); + let (ingress_tx, ingress_rx) = mpsc::channel::(); + let (delta_tx, delta_rx) = mpsc::channel::(); + let pin_cores = config.pin_cores; + + let aggregator_handle = thread::spawn(move || { + if pin_cores { + let _ = core_affinity::set_for_current(core_affinity::CoreId { id: num_workers }); + } + let mut parent = parent_factory(); + for delta in delta_rx { + parent.apply(delta); + } + parent + }); + + let wf = Arc::new(worker_factory); + let mut worker_input_txs = Vec::with_capacity(num_workers); + let mut worker_handles = Vec::with_capacity(num_workers); + for worker_id in 0..num_workers { + let (worker_tx, worker_rx) = mpsc::channel::(); + worker_input_txs.push(worker_tx); + let wf_clone = Arc::clone(&wf); + let delta_tx_worker = delta_tx.clone(); + let pin_cores = config.pin_cores; + worker_handles.push(thread::spawn(move || { + if pin_cores { + let _ = core_affinity::set_for_current(core_affinity::CoreId { id: worker_id }); + } + let mut worker = wf_clone(worker_id); + while let Ok(msg) = worker_rx.recv() { + match msg { + WorkerMsg::Batch(batch) => { + for input in batch { + let borrowed = input.as_sketch_input(); + worker.process(&borrowed, &mut |delta| { + delta_tx_worker.send(delta).expect( + "aggregator receiver dropped while workers still running", + ); + }); + } + } + WorkerMsg::End => break, + } + } + })); + } + drop(delta_tx); + + let dispatcher_handle = thread::spawn(move || { + let mut sent_end = false; + let mut pending = vec![Vec::new(); num_workers]; + while let Ok(msg) = ingress_rx.recv() { + match msg { + IngressMsg::Data(input) => { + let borrowed = input.as_sketch_input(); + let worker_id = + worker_index(hash64_seeded(PARTITION_SEED, &borrowed), num_workers); + pending[worker_id].push(input); + if pending[worker_id].len() >= WORKER_DISPATCH_BATCH_SIZE { + flush_worker_batch(&worker_input_txs, &mut pending, worker_id); + } + } + IngressMsg::End => { + flush_all_worker_batches(&worker_input_txs, &mut pending); + send_end_to_all_workers(&worker_input_txs); + sent_end = true; + break; + } + } + } + + if !sent_end { + flush_all_worker_batches(&worker_input_txs, &mut pending); + send_end_to_all_workers(&worker_input_txs); + } + }); + + Self { + ingress_tx, + dispatcher_handle: Some(dispatcher_handle), + worker_handles, + aggregator_handle: Some(aggregator_handle), + _worker_marker: PhantomData, + } + } + + pub fn insert(&mut self, input: SketchInput<'_>) { + self.ingress_tx + .send(IngressMsg::Data(OwnedSketchInput::from(&input))) + .expect("dispatcher receiver dropped while runtime is active"); + } + + pub fn insert_batch(&mut self, inputs: &[SketchInput<'_>]) { + for input in inputs { + self.insert(input.clone()); + } + } + + pub fn finish(mut self) -> OctoResult

{ + self.ingress_tx + .send(IngressMsg::End) + .expect("dispatcher receiver dropped before finish"); + drop(self.ingress_tx); + + if let Some(dispatcher) = self.dispatcher_handle.take() { + dispatcher + .join() + .expect("dispatcher thread panicked during finish"); + } + + for handle in self.worker_handles { + handle.join().expect("worker thread panicked during finish"); + } + + let parent = self + .aggregator_handle + .take() + .expect("aggregator handle missing") + .join() + .expect("aggregator thread panicked during finish"); + + OctoResult { parent } + } +} + // --------------------------------------------------------------------------- // Concrete worker/parent implementations // --------------------------------------------------------------------------- @@ -204,9 +473,9 @@ const PARTITION_SEED: usize = 19; /// Runs the OctoSketch multi-threaded insert protocol. /// -/// 1. Partitions `inputs` across `config.num_workers` workers by hash. -/// 2. Each worker maintains a child sketch, emitting deltas via SPSC queues. -/// 3. The aggregator drains all queues continuously, applying deltas to the parent. +/// 1. Dispatches `inputs` across workers online by hash through channels. +/// 2. Each worker maintains a child sketch, emitting deltas via an MPSC channel. +/// 3. The aggregator blocks on the channel and applies deltas to the parent. /// 4. Returns the fully-merged parent sketch. /// /// Uses `std::thread::scope` so `inputs` can have any lifetime (no `'static` needed). @@ -221,29 +490,12 @@ where P: OctoParent, { let num_workers = config.num_workers.max(1); - let capacity = config.queue_capacity.max(1024); - - // Step 1: Partition inputs by hash into per-worker index lists. - let mut partitions: Vec> = (0..num_workers).map(|_| Vec::new()).collect(); - for (idx, item) in inputs.iter().enumerate() { - let h = hash64_seeded(PARTITION_SEED, item) as usize; - partitions[h % num_workers].push(idx); - } - - // Step 2: Create SPSC ring buffers. - let mut producers = Vec::with_capacity(num_workers); - let mut consumers = Vec::with_capacity(num_workers); - for _ in 0..num_workers { - let rb = HeapRb::::new(capacity); - let (prod, cons) = rb.split(); - producers.push(prod); - consumers.push(cons); - } + let _capacity = config.queue_capacity.max(1024); - // Step 3: Finish flags (one per worker). - let finished: Vec = (0..num_workers).map(|_| AtomicBool::new(false)).collect(); + // Step 1: Create an unbounded MPSC channel for worker-to-aggregator deltas. + let (tx, rx) = mpsc::channel::(); - // Step 4: Run workers + aggregator inside a scoped thread block. + // Step 2: Run dispatcher + workers + aggregator inside a scoped thread block. let mut parent = parent_factory(); thread::scope(|s| { @@ -254,55 +506,61 @@ where // Spawn worker threads. let wf = &worker_factory; - for (worker_id, (partition, producer)) in partitions - .iter() - .zip(producers.into_iter()) - .enumerate() - { - let finished_flag = &finished[worker_id]; + let mut worker_input_txs = Vec::with_capacity(num_workers); + for worker_id in 0..num_workers { + let (worker_tx, worker_rx) = mpsc::channel::(); + worker_input_txs.push(worker_tx); + let tx_worker = tx.clone(); let pin_cores = config.pin_cores; s.spawn(move || { // Pin to core. if pin_cores { - let _ = core_affinity::set_for_current(core_affinity::CoreId { - id: worker_id, - }); + let _ = core_affinity::set_for_current(core_affinity::CoreId { id: worker_id }); } let mut worker = wf(worker_id); - let mut prod = producer; - - for &input_idx in partition { - let input = &inputs[input_idx]; - worker.process(input, &mut |delta| { - // Spin-wait if the ring buffer is full (backpressure). - while prod.try_push(delta).is_err() { - std::hint::spin_loop(); + while let Ok(msg) = worker_rx.recv() { + match msg { + WorkerMsg::Batch(batch) => { + for input in batch { + let borrowed = input.as_sketch_input(); + worker.process(&borrowed, &mut |delta| { + tx_worker.send(delta).expect( + "aggregator receiver dropped while workers still running", + ); + }); + } } - }); + WorkerMsg::End => break, + } } - - finished_flag.store(true, Ordering::Release); }); } - // Aggregator collect loop: drain all consumer queues. - let all_done = || finished.iter().all(|f| f.load(Ordering::Acquire)); - - while !all_done() { - for cons in consumers.iter_mut() { - while let Some(delta) = cons.try_pop() { - parent.apply(delta); + // Spawn dispatcher thread: hash-route each input to a worker channel. + s.spawn(move || { + let mut pending = vec![Vec::new(); num_workers]; + for item in inputs { + let owned = OwnedSketchInput::from(item); + let borrowed = owned.as_sketch_input(); + let worker_id = worker_index(hash64_seeded(PARTITION_SEED, &borrowed), num_workers); + pending[worker_id].push(owned); + if pending[worker_id].len() >= WORKER_DISPATCH_BATCH_SIZE { + flush_worker_batch(&worker_input_txs, &mut pending, worker_id); } } - } - // Final drain pass after all workers have finished. - for cons in consumers.iter_mut() { - while let Some(delta) = cons.try_pop() { - parent.apply(delta); - } + flush_all_worker_batches(&worker_input_txs, &mut pending); + send_end_to_all_workers(&worker_input_txs); + }); + + // Drop the main sender so the receiver closes once all worker clones drop. + drop(tx); + + // Aggregator collect loop: block until a delta arrives; exits when senders are dropped. + for delta in rx { + parent.apply(delta); } }); @@ -409,10 +667,7 @@ mod tests { parent.apply_delta(delta); // Apply a smaller value — should not change. - let smaller = HllDelta { - pos: 100, - value: 5, - }; + let smaller = HllDelta { pos: 100, value: 5 }; parent.apply_delta(smaller); // Apply a larger value — should update. @@ -563,4 +818,183 @@ mod tests { "single-worker HLL estimate {estimate} out of range for 1000 distinct keys" ); } + + #[test] + fn octo_runtime_streaming_cm_matches_batch_api() { + let rows = 3; + let cols = 4096; + let n = 30_000u64; + let inputs: Vec> = (0..n).map(|i| SketchInput::U64(i % 1024)).collect(); + let config = OctoConfig { + num_workers: 4, + pin_cores: false, + queue_capacity: 8192, + }; + + let batch_result = run_octo( + &inputs, + &config, + |_| CmOctoWorker::new(rows, cols), + || CmOctoParent { + sketch: CountMin::with_dimensions(rows, cols), + }, + ); + + let mut runtime = OctoRuntime::new( + &config, + move |_| CmOctoWorker::new(rows, cols), + move || CmOctoParent { + sketch: CountMin::with_dimensions(rows, cols), + }, + ); + for input in &inputs { + runtime.insert(input.clone()); + } + let streaming_result = runtime.finish(); + + for key_val in 0u64..128 { + let key = SketchInput::U64(key_val); + let batch_est = batch_result.parent.sketch.estimate(&key); + let stream_est = streaming_result.parent.sketch.estimate(&key); + assert_eq!(batch_est, stream_est, "key {key_val} mismatch"); + } + } + + #[test] + fn octo_runtime_mixed_insert_and_batch_hll() { + let config = OctoConfig { + num_workers: 2, + pin_cores: false, + queue_capacity: 4096, + }; + let mut runtime = OctoRuntime::new( + &config, + |_| HllOctoWorker::new(), + || HllOctoParent { + sketch: HyperLogLog::::default(), + }, + ); + + runtime.insert(SketchInput::U64(1)); + runtime.insert(SketchInput::U64(2)); + let batch: Vec> = (3..2000).map(SketchInput::U64).collect(); + runtime.insert_batch(&batch); + let result = runtime.finish(); + let estimate = result.parent.sketch.estimate(); + assert!( + estimate > 1700 && estimate < 2300, + "runtime mixed insert+batch estimate {estimate} is out of expected range" + ); + } + + #[test] + fn octo_runtime_empty_stream_finishes() { + let config = OctoConfig { + num_workers: 4, + pin_cores: false, + queue_capacity: 4096, + }; + let runtime = OctoRuntime::new( + &config, + |_| HllOctoWorker::new(), + || HllOctoParent { + sketch: HyperLogLog::::default(), + }, + ); + let result = runtime.finish(); + let estimate = result.parent.sketch.estimate(); + assert!( + estimate == 0, + "empty runtime should estimate 0 cardinality, got {estimate}" + ); + } + + struct CountingWorker; + + impl OctoWorker for CountingWorker { + type Delta = u64; + + fn process(&mut self, _input: &SketchInput, emit: &mut dyn FnMut(Self::Delta)) { + emit(1); + } + } + + struct CountingParent { + total: u64, + } + + impl OctoParent for CountingParent { + type Delta = u64; + + fn apply(&mut self, delta: Self::Delta) { + self.total += delta; + } + } + + #[test] + fn worker_index_is_deterministic_for_pow2_and_non_pow2_workers() { + let key = SketchInput::U64(0xdead_beef); + let hash = hash64_seeded(PARTITION_SEED, &key); + + let idx_pow2_a = worker_index(hash, 8); + let idx_pow2_b = worker_index(hash, 8); + assert_eq!(idx_pow2_a, idx_pow2_b); + assert_eq!(idx_pow2_a, (hash as usize) & 7); + + let idx_non_pow2_a = worker_index(hash, 6); + let idx_non_pow2_b = worker_index(hash, 6); + assert_eq!(idx_non_pow2_a, idx_non_pow2_b); + assert_eq!(idx_non_pow2_a, (hash as usize) % 6); + } + + #[test] + fn run_octo_batch_flush_boundary_matches_reference_count() { + let config = OctoConfig { + num_workers: 3, + pin_cores: false, + queue_capacity: 4096, + }; + let sizes = [ + WORKER_DISPATCH_BATCH_SIZE - 1, + WORKER_DISPATCH_BATCH_SIZE, + WORKER_DISPATCH_BATCH_SIZE + 1, + ]; + + for &n in &sizes { + let inputs: Vec> = (0..n as u64) + .map(|i| SketchInput::U64(i ^ 0x1234)) + .collect(); + let result = run_octo( + &inputs, + &config, + |_| CountingWorker, + || CountingParent { total: 0 }, + ); + assert_eq!( + result.parent.total, n as u64, + "batch boundary size {n} should process all items" + ); + } + } + + #[test] + fn octo_runtime_end_flush_preserves_tail_items() { + let config = OctoConfig { + num_workers: 4, + pin_cores: false, + queue_capacity: 4096, + }; + let n = WORKER_DISPATCH_BATCH_SIZE + 7; + let mut runtime = + OctoRuntime::new(&config, |_| CountingWorker, || CountingParent { total: 0 }); + + for i in 0..n as u64 { + runtime.insert(SketchInput::U64(i + 42)); + } + let result = runtime.finish(); + assert_eq!( + result.parent.total, n as u64, + "runtime finish should flush partial worker batches" + ); + } } diff --git a/src/sketches/mod.rs b/src/sketches/mod.rs index a174355..cbe7e96 100644 --- a/src/sketches/mod.rs +++ b/src/sketches/mod.rs @@ -47,8 +47,8 @@ pub use cs_heap::CSHeap; pub mod octo_delta; pub use octo_delta::{CM_PROMASK, COUNT_PROMASK, CmDelta, CountDelta, HLL_PROMASK, HllDelta}; -pub use countmin::CountMinChild; pub use count::CountChild; +pub use countmin::CountMinChild; pub use hll::HllChild; pub mod fold_cms; From 838d6a75f93c1efbb100b01d229b3080dea04cbb Mon Sep 17 00:00:00 2001 From: GordonYuanyc Date: Mon, 2 Mar 2026 23:25:08 -0700 Subject: [PATCH 3/7] remove data dispatch from octo.rs --- benches/octo_speedup.rs | 474 ++++++++++++++++---------- src/sketch_framework/mod.rs | 2 +- src/sketch_framework/octo.rs | 638 ++++++++++++++++------------------- src/sketches/count.rs | 19 +- src/sketches/countmin.rs | 20 +- src/sketches/hll.rs | 23 +- 6 files changed, 642 insertions(+), 534 deletions(-) diff --git a/benches/octo_speedup.rs b/benches/octo_speedup.rs index a47fa9c..7d445d0 100644 --- a/benches/octo_speedup.rs +++ b/benches/octo_speedup.rs @@ -2,11 +2,12 @@ use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, use rand::{Rng, SeedableRng, rngs::StdRng}; use sketchlib_rust::{ CmOctoParent, CmOctoWorker, Count, CountMin, CountOctoParent, CountOctoWorker, HllOctoParent, - HllOctoWorker, HyperLogLog, OctoConfig, Regular, RegularPath, SketchInput, Vector2D, - hash64_seeded, run_octo, + HllOctoWorker, HyperLogLog, OctoAggregator, OctoConfig, OctoWorker, Regular, RegularPath, + SketchInput, Vector2D, }; -use std::sync::mpsc; +use std::sync::{Once, mpsc}; use std::thread; +use std::time::Duration; const RNG_SEED: u64 = 0x0c70_2026_5eed_1234; const CM_COUNT_ROWS: usize = 3; @@ -16,9 +17,10 @@ const HLL_INPUTS: usize = 500_000; const DOMAIN_MASK: u64 = (1 << 20) - 1; const QUEUE_CAPACITY: usize = 65_536; const WORKER_COUNTS: [usize; 4] = [1, 2, 4, 8]; -const PARTITION_SEED: usize = 19; const MERGE_INTERVAL: usize = 10_000; +static SANITY_ONCE: Once = Once::new(); + fn build_inputs(sample_count: usize) -> Vec> { let mut rng = StdRng::seed_from_u64(RNG_SEED ^ (sample_count as u64)); (0..sample_count) @@ -26,203 +28,304 @@ fn build_inputs(sample_count: usize) -> Vec> { .collect() } -#[inline(always)] -fn should_merge(number: usize, interval: usize, thread_num: usize, thread_id: usize) -> bool { - number != 0 && number % (interval * thread_num) == interval * thread_id -} - -fn run_regular_countmin_merge( - inputs: &[SketchInput<'_>], +fn build_round_robin_shards( + inputs: &[SketchInput<'static>], num_workers: usize, -) -> CountMin, RegularPath> { - let mut workers = Vec::new(); +) -> Vec>> { + let mut shards = vec![Vec::new(); num_workers]; + for (idx, input) in inputs.iter().enumerate() { + shards[idx % num_workers].push(input.clone()); + } + shards +} +fn run_octo_sharded( + shards: &[Vec>], + config: &OctoConfig, + worker_factory: impl Fn(usize) -> W, + parent_factory: impl FnOnce() -> P + Send, +) -> P +where + W: OctoWorker + 'static, + P: OctoAggregator + Send + 'static, +{ + let num_workers = shards.len(); thread::scope(|s| { - let mut worker_txs = Vec::with_capacity(num_workers); - let mut handles = Vec::with_capacity(num_workers); - for worker_id in 0..num_workers { - let (tx, rx) = mpsc::channel::>(); - worker_txs.push(tx); - handles.push(s.spawn(move || { - let mut child = CountMin::, RegularPath>::with_dimensions( - CM_COUNT_ROWS, - CM_COUNT_COLS, - ); - let mut partials = Vec::new(); - let mut number = 0usize; - while let Ok(msg) = rx.recv() { - match msg { - Some(idx) => { - child.insert(&inputs[idx]); - number += 1; - if should_merge(number, MERGE_INTERVAL, num_workers, worker_id) { - partials.push(child); - child = CountMin::, RegularPath>::with_dimensions( - CM_COUNT_ROWS, - CM_COUNT_COLS, - ); - } - } - None => break, - } - } - partials.push(child); - partials - })); - } - - s.spawn(move || { - for (idx, item) in inputs.iter().enumerate() { - let h = hash64_seeded(PARTITION_SEED, item) as usize; - let worker_id = h % num_workers; - worker_txs[worker_id] - .send(Some(idx)) - .expect("regular CountMin worker receiver dropped unexpectedly"); + let (delta_tx, delta_rx) = mpsc::channel::(); + let aggregator = s.spawn(move || { + let mut parent = parent_factory(); + if config.pin_cores { + let _ = core_affinity::set_for_current(core_affinity::CoreId { id: num_workers }); } - for tx in worker_txs { - let _ = tx.send(None); + for delta in delta_rx { + parent.apply(delta); } + parent }); - for h in handles { - workers.extend(h.join().expect("regular CountMin worker thread panicked")); + for (worker_id, shard) in shards.iter().enumerate() { + let delta_tx_worker = delta_tx.clone(); + let mut worker = worker_factory(worker_id); + let pin_cores = config.pin_cores; + s.spawn(move || { + if pin_cores { + let _ = core_affinity::set_for_current(core_affinity::CoreId { id: worker_id }); + } + for input in shard { + worker.process(input, &mut |delta| { + delta_tx_worker + .send(delta) + .expect("octo sharded aggregator dropped unexpectedly"); + }); + } + }); } - }); + drop(delta_tx); - let mut parent = - CountMin::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS); - for worker in workers.iter() { - parent.merge(worker); - } - parent + aggregator + .join() + .expect("octo sharded aggregator thread panicked") + }) } -fn run_regular_count_merge( - inputs: &[SketchInput<'_>], - num_workers: usize, -) -> Count, RegularPath> { - let mut workers = Vec::new(); +enum CountMinMergeMsg { + Snapshot(CountMin, RegularPath>), + Done, +} + +enum CountMergeMsg { + Snapshot(Count, RegularPath>), + Done, +} + +enum HllMergeMsg { + Snapshot(HyperLogLog), + Done, +} +fn run_periodic_countmin_full_merge_sharded( + shards: &[Vec>], +) -> CountMin, RegularPath> { + let num_workers = shards.len(); thread::scope(|s| { - let mut worker_txs = Vec::with_capacity(num_workers); - let mut handles = Vec::with_capacity(num_workers); - for worker_id in 0..num_workers { - let (tx, rx) = mpsc::channel::>(); - worker_txs.push(tx); - handles.push(s.spawn(move || { - let mut child = Count::, RegularPath>::with_dimensions( - CM_COUNT_ROWS, - CM_COUNT_COLS, - ); - let mut partials = Vec::new(); - let mut number = 0usize; - while let Ok(msg) = rx.recv() { - match msg { - Some(idx) => { - child.insert(&inputs[idx]); - number += 1; - if should_merge(number, MERGE_INTERVAL, num_workers, worker_id) { - partials.push(child); - child = Count::, RegularPath>::with_dimensions( - CM_COUNT_ROWS, - CM_COUNT_COLS, - ); - } - } - None => break, + let (merge_tx, merge_rx) = mpsc::channel::(); + let aggregator = s.spawn(move || { + let mut parent = + CountMin::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS); + let mut finished_workers = 0usize; + while finished_workers < num_workers { + match merge_rx.recv().expect("CountMin merge channel closed unexpectedly") { + CountMinMergeMsg::Snapshot(snapshot) => parent.merge(&snapshot), + CountMinMergeMsg::Done => finished_workers += 1, + } + } + parent + }); + + for shard in shards { + let merge_tx_worker = merge_tx.clone(); + s.spawn(move || { + let mut child = + CountMin::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS); + let mut local_inserts = 0usize; + for input in shard { + child.insert(input); + local_inserts += 1; + if local_inserts % MERGE_INTERVAL == 0 { + merge_tx_worker + .send(CountMinMergeMsg::Snapshot(child)) + .expect("CountMin aggregator dropped unexpectedly"); + child = CountMin::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ); } } - partials.push(child); - partials - })); + merge_tx_worker + .send(CountMinMergeMsg::Snapshot(child)) + .expect("CountMin aggregator dropped unexpectedly"); + merge_tx_worker + .send(CountMinMergeMsg::Done) + .expect("CountMin aggregator dropped unexpectedly"); + }); } + drop(merge_tx); - s.spawn(move || { - for (idx, item) in inputs.iter().enumerate() { - let h = hash64_seeded(PARTITION_SEED, item) as usize; - let worker_id = h % num_workers; - worker_txs[worker_id] - .send(Some(idx)) - .expect("regular Count worker receiver dropped unexpectedly"); - } - for tx in worker_txs { - let _ = tx.send(None); + aggregator + .join() + .expect("CountMin periodic merge aggregator panicked") + }) +} + +fn run_periodic_count_full_merge_sharded( + shards: &[Vec>], +) -> Count, RegularPath> { + let num_workers = shards.len(); + thread::scope(|s| { + let (merge_tx, merge_rx) = mpsc::channel::(); + let aggregator = s.spawn(move || { + let mut parent = + Count::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS); + let mut finished_workers = 0usize; + while finished_workers < num_workers { + match merge_rx.recv().expect("Count merge channel closed unexpectedly") { + CountMergeMsg::Snapshot(snapshot) => parent.merge(&snapshot), + CountMergeMsg::Done => finished_workers += 1, + } } + parent }); - for h in handles { - workers.extend(h.join().expect("regular Count worker thread panicked")); + for shard in shards { + let merge_tx_worker = merge_tx.clone(); + s.spawn(move || { + let mut child = + Count::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS); + let mut local_inserts = 0usize; + for input in shard { + child.insert(input); + local_inserts += 1; + if local_inserts % MERGE_INTERVAL == 0 { + merge_tx_worker + .send(CountMergeMsg::Snapshot(child)) + .expect("Count aggregator dropped unexpectedly"); + child = Count::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ); + } + } + merge_tx_worker + .send(CountMergeMsg::Snapshot(child)) + .expect("Count aggregator dropped unexpectedly"); + merge_tx_worker + .send(CountMergeMsg::Done) + .expect("Count aggregator dropped unexpectedly"); + }); } - }); + drop(merge_tx); - let mut parent = - Count::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS); - for worker in workers.iter() { - parent.merge(worker); - } - parent + aggregator + .join() + .expect("Count periodic merge aggregator panicked") + }) } -fn run_regular_hll_merge(inputs: &[SketchInput<'_>], num_workers: usize) -> HyperLogLog { - let mut workers = Vec::new(); - +fn run_periodic_hll_full_merge_sharded( + shards: &[Vec>], +) -> HyperLogLog { + let num_workers = shards.len(); thread::scope(|s| { - let mut worker_txs = Vec::with_capacity(num_workers); - let mut handles = Vec::with_capacity(num_workers); - for worker_id in 0..num_workers { - let (tx, rx) = mpsc::channel::>(); - worker_txs.push(tx); - handles.push(s.spawn(move || { + let (merge_tx, merge_rx) = mpsc::channel::(); + let aggregator = s.spawn(move || { + let mut parent = HyperLogLog::::default(); + let mut finished_workers = 0usize; + while finished_workers < num_workers { + match merge_rx.recv().expect("HLL merge channel closed unexpectedly") { + HllMergeMsg::Snapshot(snapshot) => parent.merge(&snapshot), + HllMergeMsg::Done => finished_workers += 1, + } + } + parent + }); + + for shard in shards { + let merge_tx_worker = merge_tx.clone(); + s.spawn(move || { let mut child = HyperLogLog::::default(); - let mut partials = Vec::new(); - let mut number = 0usize; - while let Ok(msg) = rx.recv() { - match msg { - Some(idx) => { - child.insert(&inputs[idx]); - number += 1; - if should_merge(number, MERGE_INTERVAL, num_workers, worker_id) { - partials.push(child); - child = HyperLogLog::::default(); - } - } - None => break, + let mut local_inserts = 0usize; + for input in shard { + child.insert(input); + local_inserts += 1; + if local_inserts % MERGE_INTERVAL == 0 { + merge_tx_worker + .send(HllMergeMsg::Snapshot(child)) + .expect("HLL aggregator dropped unexpectedly"); + child = HyperLogLog::::default(); } } - partials.push(child); - partials - })); + merge_tx_worker + .send(HllMergeMsg::Snapshot(child)) + .expect("HLL aggregator dropped unexpectedly"); + merge_tx_worker + .send(HllMergeMsg::Done) + .expect("HLL aggregator dropped unexpectedly"); + }); } + drop(merge_tx); - s.spawn(move || { - for (idx, item) in inputs.iter().enumerate() { - let h = hash64_seeded(PARTITION_SEED, item) as usize; - let worker_id = h % num_workers; - worker_txs[worker_id] - .send(Some(idx)) - .expect("regular HLL worker receiver dropped unexpectedly"); - } - for tx in worker_txs { - let _ = tx.send(None); - } - }); + aggregator + .join() + .expect("HLL periodic merge aggregator panicked") + }) +} - for h in handles { - workers.extend(h.join().expect("regular HLL worker thread panicked")); +fn sanity_check_periodic_baselines() { + SANITY_ONCE.call_once(|| { + let inputs = build_inputs(20_000); + let workers = 4; + let shards = build_round_robin_shards(&inputs, workers); + + let config = OctoConfig { + num_workers: workers, + pin_cores: false, + queue_capacity: QUEUE_CAPACITY, + }; + + let octo_cm = run_octo_sharded( + &shards, + &config, + |_| CmOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), + || CmOctoParent { + sketch: CountMin::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), + }, + ) + .sketch; + let periodic_cm = run_periodic_countmin_full_merge_sharded(&shards); + for key in [1_u64, 7, 42, 4097] { + let input = SketchInput::U64(key); + assert!(octo_cm.estimate(&input) <= periodic_cm.estimate(&input)); } - }); - let mut parent = HyperLogLog::::default(); - for worker in workers.iter() { - parent.merge(worker); - } - parent + let octo_count = run_octo_sharded( + &shards, + &config, + |_| CountOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), + || CountOctoParent { + sketch: Count::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), + }, + ) + .sketch; + let periodic_count = run_periodic_count_full_merge_sharded(&shards); + for key in [1_u64, 7, 42, 4097] { + let input = SketchInput::U64(key); + let octo = octo_count.estimate(&input); + let full = periodic_count.estimate(&input); + assert!(octo.abs() <= full.abs() + 500.0); + } + + let octo_hll = run_octo_sharded( + &shards, + &config, + |_| HllOctoWorker::new(), + || HllOctoParent { + sketch: HyperLogLog::::default(), + }, + ) + .sketch; + let periodic_hll = run_periodic_hll_full_merge_sharded(&shards); + assert!(octo_hll.estimate() <= periodic_hll.estimate()); + }); } fn bench_countmin_merge_compare(c: &mut Criterion) { + sanity_check_periodic_baselines(); + let inputs = build_inputs(CM_COUNT_INPUTS); let mut group = c.benchmark_group("countmin_merge_compare"); group.sample_size(10); + group.warm_up_time(Duration::from_secs(2)); + group.measurement_time(Duration::from_secs(8)); group.throughput(Throughput::Elements(inputs.len() as u64)); group.bench_function("single_thread_insert", |b| { @@ -243,6 +346,7 @@ fn bench_countmin_merge_compare(c: &mut Criterion) { }); for &workers in &WORKER_COUNTS { + let shards = build_round_robin_shards(&inputs, workers); group.bench_with_input( BenchmarkId::new("octo_style", workers), &workers, @@ -253,25 +357,25 @@ fn bench_countmin_merge_compare(c: &mut Criterion) { pin_cores: false, queue_capacity: QUEUE_CAPACITY, }; - let result = run_octo( - &inputs, + let parent = run_octo_sharded( + &shards, &config, |_| CmOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), || CmOctoParent { sketch: CountMin::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), }, ); - black_box(result.parent.sketch); + black_box(parent.sketch); }); }, ); group.bench_with_input( - BenchmarkId::new("regular_full_merge", workers), + BenchmarkId::new("regular_periodic_full_merge", workers), &workers, - |b, &w| { + |b, _| { b.iter(|| { - let merged = run_regular_countmin_merge(&inputs, w); + let merged = run_periodic_countmin_full_merge_sharded(&shards); black_box(merged); }); }, @@ -282,9 +386,13 @@ fn bench_countmin_merge_compare(c: &mut Criterion) { } fn bench_count_merge_compare(c: &mut Criterion) { + sanity_check_periodic_baselines(); + let inputs = build_inputs(CM_COUNT_INPUTS); let mut group = c.benchmark_group("count_merge_compare"); group.sample_size(10); + group.warm_up_time(Duration::from_secs(2)); + group.measurement_time(Duration::from_secs(8)); group.throughput(Throughput::Elements(inputs.len() as u64)); group.bench_function("single_thread_insert", |b| { @@ -300,6 +408,7 @@ fn bench_count_merge_compare(c: &mut Criterion) { }); for &workers in &WORKER_COUNTS { + let shards = build_round_robin_shards(&inputs, workers); group.bench_with_input( BenchmarkId::new("octo_style", workers), &workers, @@ -310,25 +419,25 @@ fn bench_count_merge_compare(c: &mut Criterion) { pin_cores: false, queue_capacity: QUEUE_CAPACITY, }; - let result = run_octo( - &inputs, + let parent = run_octo_sharded( + &shards, &config, |_| CountOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), || CountOctoParent { sketch: Count::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), }, ); - black_box(result.parent.sketch); + black_box(parent.sketch); }); }, ); group.bench_with_input( - BenchmarkId::new("regular_full_merge", workers), + BenchmarkId::new("regular_periodic_full_merge", workers), &workers, - |b, &w| { + |b, _| { b.iter(|| { - let merged = run_regular_count_merge(&inputs, w); + let merged = run_periodic_count_full_merge_sharded(&shards); black_box(merged); }); }, @@ -339,9 +448,13 @@ fn bench_count_merge_compare(c: &mut Criterion) { } fn bench_hll_merge_compare(c: &mut Criterion) { + sanity_check_periodic_baselines(); + let inputs = build_inputs(HLL_INPUTS); let mut group = c.benchmark_group("hll_merge_compare"); group.sample_size(10); + group.warm_up_time(Duration::from_secs(2)); + group.measurement_time(Duration::from_secs(8)); group.throughput(Throughput::Elements(inputs.len() as u64)); group.bench_function("single_thread_insert", |b| { @@ -354,6 +467,7 @@ fn bench_hll_merge_compare(c: &mut Criterion) { }); for &workers in &WORKER_COUNTS { + let shards = build_round_robin_shards(&inputs, workers); group.bench_with_input( BenchmarkId::new("octo_style", workers), &workers, @@ -364,25 +478,25 @@ fn bench_hll_merge_compare(c: &mut Criterion) { pin_cores: false, queue_capacity: QUEUE_CAPACITY, }; - let result = run_octo( - &inputs, + let parent = run_octo_sharded( + &shards, &config, |_| HllOctoWorker::new(), || HllOctoParent { sketch: HyperLogLog::::default(), }, ); - black_box(result.parent.sketch); + black_box(parent.sketch); }); }, ); group.bench_with_input( - BenchmarkId::new("regular_full_merge", workers), + BenchmarkId::new("regular_periodic_full_merge", workers), &workers, - |b, &w| { + |b, _| { b.iter(|| { - let merged = run_regular_hll_merge(&inputs, w); + let merged = run_periodic_hll_full_merge_sharded(&shards); black_box(merged); }); }, diff --git a/src/sketch_framework/mod.rs b/src/sketch_framework/mod.rs index 7bc87a7..fd2488c 100644 --- a/src/sketch_framework/mod.rs +++ b/src/sketch_framework/mod.rs @@ -38,7 +38,7 @@ pub use eh_univ_optimized::{EHMapBucket, EHUnivMonBucket, EHUnivOptimized, EHUni pub mod octo; pub use octo::{ CmOctoParent, CmOctoWorker, CountOctoParent, CountOctoWorker, HllOctoParent, HllOctoWorker, - OctoConfig, OctoParent, OctoResult, OctoRuntime, OctoWorker, OwnedSketchInput, run_octo, + OctoAggregator, OctoConfig, OctoReadHandle, OctoResult, OctoRuntime, OctoWorker, run_octo, }; pub mod tumbling; diff --git a/src/sketch_framework/octo.rs b/src/sketch_framework/octo.rs index 0c3b9fc..c89776d 100644 --- a/src/sketch_framework/octo.rs +++ b/src/sketch_framework/octo.rs @@ -6,18 +6,18 @@ //! threshold. An aggregator thread applies deltas to a full-precision parent sketch. use std::marker::PhantomData; -use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, RwLock, Weak}; use std::sync::mpsc; use std::thread; use crate::{ CmDelta, Count, CountChild, CountDelta, CountMin, CountMinChild, HllChild, HllDelta, - HyperLogLog, Regular, RegularPath, SketchInput, Vector2D, hash64_seeded, + HyperLogLog, Regular, RegularPath, SketchInput, Vector2D, }; /// Legacy queue capacity default retained for config compatibility. const DEFAULT_QUEUE_CAPACITY: usize = 65536; -const WORKER_DISPATCH_BATCH_SIZE: usize = 64; // --------------------------------------------------------------------------- // Traits @@ -27,12 +27,12 @@ const WORKER_DISPATCH_BATCH_SIZE: usize = 64; pub trait OctoWorker: Send { type Delta: Copy + Send + 'static; - /// Process one input, emitting zero or more deltas via `emit`. + /// Process one input and emit zero or more deltas. fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(Self::Delta)); } /// Parent-side trait: absorbs deltas into a full-precision sketch. -pub trait OctoParent: Send { +pub trait OctoAggregator: Send { type Delta: Copy + Send + 'static; /// Apply a single delta to the parent sketch. @@ -71,190 +71,147 @@ pub struct OctoResult

{ pub parent: P, } -/// Owned variant of `SketchInput` for cross-thread transport. -#[derive(Clone, Debug)] -pub enum OwnedSketchInput { - I8(i8), - I16(i16), - I32(i32), - I64(i64), - I128(i128), - ISIZE(isize), - U8(u8), - U16(u16), - U32(u32), - U64(u64), - U128(u128), - USIZE(usize), - F32(f32), - F64(f64), - String(String), - Bytes(Vec), -} - -impl OwnedSketchInput { - fn as_sketch_input(&self) -> SketchInput<'_> { - match self { - OwnedSketchInput::I8(v) => SketchInput::I8(*v), - OwnedSketchInput::I16(v) => SketchInput::I16(*v), - OwnedSketchInput::I32(v) => SketchInput::I32(*v), - OwnedSketchInput::I64(v) => SketchInput::I64(*v), - OwnedSketchInput::I128(v) => SketchInput::I128(*v), - OwnedSketchInput::ISIZE(v) => SketchInput::ISIZE(*v), - OwnedSketchInput::U8(v) => SketchInput::U8(*v), - OwnedSketchInput::U16(v) => SketchInput::U16(*v), - OwnedSketchInput::U32(v) => SketchInput::U32(*v), - OwnedSketchInput::U64(v) => SketchInput::U64(*v), - OwnedSketchInput::U128(v) => SketchInput::U128(*v), - OwnedSketchInput::USIZE(v) => SketchInput::USIZE(*v), - OwnedSketchInput::F32(v) => SketchInput::F32(*v), - OwnedSketchInput::F64(v) => SketchInput::F64(*v), - OwnedSketchInput::String(s) => SketchInput::Str(s.as_str()), - OwnedSketchInput::Bytes(b) => SketchInput::Bytes(b.as_slice()), - } - } +enum WorkerMsg { + Data(SketchInput<'static>), + End, } -impl From<&SketchInput<'_>> for OwnedSketchInput { - fn from(value: &SketchInput<'_>) -> Self { - match value { - SketchInput::I8(v) => OwnedSketchInput::I8(*v), - SketchInput::I16(v) => OwnedSketchInput::I16(*v), - SketchInput::I32(v) => OwnedSketchInput::I32(*v), - SketchInput::I64(v) => OwnedSketchInput::I64(*v), - SketchInput::I128(v) => OwnedSketchInput::I128(*v), - SketchInput::ISIZE(v) => OwnedSketchInput::ISIZE(*v), - SketchInput::U8(v) => OwnedSketchInput::U8(*v), - SketchInput::U16(v) => OwnedSketchInput::U16(*v), - SketchInput::U32(v) => OwnedSketchInput::U32(*v), - SketchInput::U64(v) => OwnedSketchInput::U64(*v), - SketchInput::U128(v) => OwnedSketchInput::U128(*v), - SketchInput::USIZE(v) => OwnedSketchInput::USIZE(*v), - SketchInput::F32(v) => OwnedSketchInput::F32(*v), - SketchInput::F64(v) => OwnedSketchInput::F64(*v), - SketchInput::Str(s) => OwnedSketchInput::String((*s).to_string()), - SketchInput::String(s) => OwnedSketchInput::String(s.clone()), - SketchInput::Bytes(b) => OwnedSketchInput::Bytes(b.to_vec()), - } - } +/// Extends a `SketchInput` lifetime to `'static` for cross-thread transport in +/// streaming mode. Caller must ensure all borrowed data outlives worker processing. +#[inline(always)] +unsafe fn assume_input_static(input: SketchInput<'_>) -> SketchInput<'static> { + // SAFETY: enforced by caller contract described above. + unsafe { std::mem::transmute::, SketchInput<'static>>(input) } } -enum IngressMsg { - Data(OwnedSketchInput), - End, +/// Streaming Octo runtime that accepts incremental inserts and finalizes into a parent sketch. +pub struct OctoRuntime +where + W: OctoWorker + 'static, + P: OctoAggregator + Send + Sync + 'static, +{ + core: Option>, + _worker_marker: PhantomData, } -enum WorkerMsg { - Batch(Vec), - End, +/// Read-only handle for querying the live aggregator state while runtime is active. +pub struct OctoReadHandle

{ + parent: Weak>, } -#[inline(always)] -fn worker_index(hash: u64, num_workers: usize) -> usize { - if num_workers.is_power_of_two() { - (hash as usize) & (num_workers - 1) - } else { - (hash as usize) % num_workers +impl

Clone for OctoReadHandle

{ + fn clone(&self) -> Self { + Self { + parent: Weak::clone(&self.parent), + } } } -#[inline] -fn flush_worker_batch( - worker_input_txs: &[mpsc::Sender], - pending: &mut [Vec], - worker_id: usize, -) { - if pending[worker_id].is_empty() { - return; +impl

OctoReadHandle

{ + /// Executes a read-only closure over the live parent state. + pub fn with_parent(&self, f: impl FnOnce(&P) -> R) -> R { + let parent = self + .parent + .upgrade() + .expect("Octo runtime has been finished and parent state was dropped"); + let guard = parent.read().expect("parent lock poisoned"); + f(&guard) } - let batch = std::mem::take(&mut pending[worker_id]); - worker_input_txs[worker_id] - .send(WorkerMsg::Batch(batch)) - .expect("worker receiver dropped unexpectedly"); } -#[inline] -fn flush_all_worker_batches( - worker_input_txs: &[mpsc::Sender], - pending: &mut [Vec], -) { - for worker_id in 0..pending.len() { - flush_worker_batch(worker_input_txs, pending, worker_id); - } +struct OctoCore

{ + worker_input_txs: Vec>, + next_worker: AtomicUsize, + worker_handles: Vec>, + aggregator_handle: Option>, + parent: Arc>, + closed: AtomicBool, } -#[inline] -fn send_end_to_all_workers(worker_input_txs: &[mpsc::Sender]) { - for tx in worker_input_txs { - let _ = tx.send(WorkerMsg::End); +impl

OctoCore

{ + fn read_handle(&self) -> OctoReadHandle

{ + OctoReadHandle { + parent: Arc::downgrade(&self.parent), + } + } + + fn close(&self) { + if self.closed.swap(true, Ordering::AcqRel) { + return; + } + for tx in &self.worker_input_txs { + let _ = tx.send(WorkerMsg::End); + } } } -/// Streaming Octo runtime that accepts incremental inserts and finalizes into a parent sketch. -pub struct OctoRuntime -where - W: OctoWorker + 'static, - P: OctoParent + Send + 'static, -{ - ingress_tx: mpsc::Sender, - dispatcher_handle: Option>, - worker_handles: Vec>, - aggregator_handle: Option>, - _worker_marker: PhantomData, +impl

OctoCore

{ + fn into_parent(mut self) -> P { + self.close(); + + for handle in self.worker_handles.drain(..) { + handle.join().expect("worker thread panicked during finish"); + } + + if let Some(aggregator) = self.aggregator_handle.take() { + aggregator + .join() + .expect("aggregator thread panicked during finish"); + } + + let parent_lock = match Arc::try_unwrap(self.parent) { + Ok(lock) => lock, + Err(_) => panic!("Octo parent still has external strong references at finish"), + }; + parent_lock.into_inner().expect("parent lock poisoned") + } } -impl OctoRuntime +impl

OctoCore

where - W: OctoWorker + 'static, - P: OctoParent + Send + 'static, + P: Send + Sync + 'static, { - pub fn new(config: &OctoConfig, worker_factory: F, parent_factory: PF) -> Self + fn start(workers: Vec, parent: P, num_workers: usize, pin_cores: bool) -> Self where - F: Fn(usize) -> W + Send + Sync + 'static, - PF: FnOnce() -> P + Send + 'static, + W: OctoWorker + 'static, + P: OctoAggregator, { - let num_workers = config.num_workers.max(1); - let (ingress_tx, ingress_rx) = mpsc::channel::(); + assert_eq!(workers.len(), num_workers); + let (delta_tx, delta_rx) = mpsc::channel::(); - let pin_cores = config.pin_cores; + let parent = Arc::new(RwLock::new(parent)); + let parent_for_aggregator = Arc::clone(&parent); let aggregator_handle = thread::spawn(move || { if pin_cores { let _ = core_affinity::set_for_current(core_affinity::CoreId { id: num_workers }); } - let mut parent = parent_factory(); for delta in delta_rx { - parent.apply(delta); + let mut guard = parent_for_aggregator + .write() + .expect("parent lock poisoned in aggregator"); + guard.apply(delta); } - parent }); - let wf = Arc::new(worker_factory); let mut worker_input_txs = Vec::with_capacity(num_workers); let mut worker_handles = Vec::with_capacity(num_workers); - for worker_id in 0..num_workers { + for (worker_id, mut worker) in workers.into_iter().enumerate() { let (worker_tx, worker_rx) = mpsc::channel::(); worker_input_txs.push(worker_tx); - let wf_clone = Arc::clone(&wf); let delta_tx_worker = delta_tx.clone(); - let pin_cores = config.pin_cores; + let pin_cores = pin_cores; worker_handles.push(thread::spawn(move || { if pin_cores { let _ = core_affinity::set_for_current(core_affinity::CoreId { id: worker_id }); } - let mut worker = wf_clone(worker_id); while let Ok(msg) = worker_rx.recv() { match msg { - WorkerMsg::Batch(batch) => { - for input in batch { - let borrowed = input.as_sketch_input(); - worker.process(&borrowed, &mut |delta| { - delta_tx_worker.send(delta).expect( - "aggregator receiver dropped while workers still running", - ); - }); - } - } + WorkerMsg::Data(input) => worker.process(&input, &mut |delta| { + delta_tx_worker.send(delta).expect( + "aggregator receiver dropped while workers still running", + ); + }), WorkerMsg::End => break, } } @@ -262,48 +219,61 @@ where } drop(delta_tx); - let dispatcher_handle = thread::spawn(move || { - let mut sent_end = false; - let mut pending = vec![Vec::new(); num_workers]; - while let Ok(msg) = ingress_rx.recv() { - match msg { - IngressMsg::Data(input) => { - let borrowed = input.as_sketch_input(); - let worker_id = - worker_index(hash64_seeded(PARTITION_SEED, &borrowed), num_workers); - pending[worker_id].push(input); - if pending[worker_id].len() >= WORKER_DISPATCH_BATCH_SIZE { - flush_worker_batch(&worker_input_txs, &mut pending, worker_id); - } - } - IngressMsg::End => { - flush_all_worker_batches(&worker_input_txs, &mut pending); - send_end_to_all_workers(&worker_input_txs); - sent_end = true; - break; - } - } - } - - if !sent_end { - flush_all_worker_batches(&worker_input_txs, &mut pending); - send_end_to_all_workers(&worker_input_txs); - } - }); - Self { - ingress_tx, - dispatcher_handle: Some(dispatcher_handle), + worker_input_txs, + next_worker: AtomicUsize::new(0), worker_handles, aggregator_handle: Some(aggregator_handle), + parent, + closed: AtomicBool::new(false), + } + } +} + +impl OctoRuntime +where + W: OctoWorker + 'static, + P: OctoAggregator + Send + Sync + 'static, +{ + pub fn new(config: &OctoConfig, worker_factory: F, parent_factory: PF) -> Self + where + F: Fn(usize) -> W, + PF: FnOnce() -> P, + { + let num_workers = config.num_workers.max(1); + let workers: Vec = (0..num_workers).map(worker_factory).collect(); + let parent = parent_factory(); + let core = OctoCore::start(workers, parent, num_workers, config.pin_cores); + + Self { + core: Some(core), _worker_marker: PhantomData, } } + pub fn read_handle(&self) -> OctoReadHandle

{ + self.core + .as_ref() + .expect("runtime core missing") + .read_handle() + } + + pub fn close(&self) { + self.core.as_ref().expect("runtime core missing").close(); + } + pub fn insert(&mut self, input: SketchInput<'_>) { - self.ingress_tx - .send(IngressMsg::Data(OwnedSketchInput::from(&input))) - .expect("dispatcher receiver dropped while runtime is active"); + let core = self.core.as_ref().expect("runtime core missing"); + if core.closed.load(Ordering::Acquire) { + panic!("cannot insert after runtime has been closed"); + } + + let worker_id = core.next_worker.fetch_add(1, Ordering::AcqRel) % core.worker_input_txs.len(); + // SAFETY: caller explicitly guarantees borrowed data lives long enough. + let static_input = unsafe { assume_input_static(input) }; + core.worker_input_txs[worker_id] + .send(WorkerMsg::Data(static_input)) + .expect("worker receiver dropped while runtime is active"); } pub fn insert_batch(&mut self, inputs: &[SketchInput<'_>]) { @@ -313,27 +283,11 @@ where } pub fn finish(mut self) -> OctoResult

{ - self.ingress_tx - .send(IngressMsg::End) - .expect("dispatcher receiver dropped before finish"); - drop(self.ingress_tx); - - if let Some(dispatcher) = self.dispatcher_handle.take() { - dispatcher - .join() - .expect("dispatcher thread panicked during finish"); - } - - for handle in self.worker_handles { - handle.join().expect("worker thread panicked during finish"); - } - let parent = self - .aggregator_handle + .core .take() - .expect("aggregator handle missing") - .join() - .expect("aggregator thread panicked during finish"); + .expect("runtime core missing") + .into_parent(); OctoResult { parent } } @@ -363,7 +317,7 @@ impl OctoWorker for CmOctoWorker { #[inline(always)] fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(CmDelta)) { - self.child.insert_and_emit(input, emit); + self.child.insert(input, emit); } } @@ -372,7 +326,7 @@ pub struct CmOctoParent { pub sketch: CountMin, RegularPath>, } -impl OctoParent for CmOctoParent { +impl OctoAggregator for CmOctoParent { type Delta = CmDelta; #[inline(always)] @@ -401,7 +355,7 @@ impl OctoWorker for CountOctoWorker { #[inline(always)] fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(CountDelta)) { - self.child.insert_and_emit(input, emit); + self.child.insert(input, emit); } } @@ -410,7 +364,7 @@ pub struct CountOctoParent { pub sketch: Count, RegularPath>, } -impl OctoParent for CountOctoParent { +impl OctoAggregator for CountOctoParent { type Delta = CountDelta; #[inline(always)] @@ -445,7 +399,7 @@ impl OctoWorker for HllOctoWorker { #[inline(always)] fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(HllDelta)) { - self.child.insert_and_emit(input, emit); + self.child.insert(input, emit); } } @@ -454,7 +408,7 @@ pub struct HllOctoParent { pub sketch: HyperLogLog, } -impl OctoParent for HllOctoParent { +impl OctoAggregator for HllOctoParent { type Delta = HllDelta; #[inline(always)] @@ -467,104 +421,27 @@ impl OctoParent for HllOctoParent { // Core execution engine // --------------------------------------------------------------------------- -/// Hash-based input partitioning seed index (uses the last seed in SEEDLIST -/// to avoid collision with sketch-internal seeds which use indices 0..~5). -const PARTITION_SEED: usize = 19; - /// Runs the OctoSketch multi-threaded insert protocol. /// -/// 1. Dispatches `inputs` across workers online by hash through channels. +/// 1. Dispatches `inputs` across workers round-robin through per-worker channels. /// 2. Each worker maintains a child sketch, emitting deltas via an MPSC channel. /// 3. The aggregator blocks on the channel and applies deltas to the parent. /// 4. Returns the fully-merged parent sketch. -/// -/// Uses `std::thread::scope` so `inputs` can have any lifetime (no `'static` needed). pub fn run_octo( inputs: &[SketchInput<'_>], config: &OctoConfig, - worker_factory: impl Fn(usize) -> W + Send + Sync, + worker_factory: impl Fn(usize) -> W, parent_factory: impl FnOnce() -> P, ) -> OctoResult

where - W: OctoWorker, - P: OctoParent, + W: OctoWorker + 'static, + P: OctoAggregator + Send + Sync + 'static, { - let num_workers = config.num_workers.max(1); - let _capacity = config.queue_capacity.max(1024); - - // Step 1: Create an unbounded MPSC channel for worker-to-aggregator deltas. - let (tx, rx) = mpsc::channel::(); - - // Step 2: Run dispatcher + workers + aggregator inside a scoped thread block. - let mut parent = parent_factory(); - - thread::scope(|s| { - // Pin the aggregator (calling thread) to core num_workers if requested. - if config.pin_cores { - let _ = core_affinity::set_for_current(core_affinity::CoreId { id: num_workers }); - } - - // Spawn worker threads. - let wf = &worker_factory; - let mut worker_input_txs = Vec::with_capacity(num_workers); - for worker_id in 0..num_workers { - let (worker_tx, worker_rx) = mpsc::channel::(); - worker_input_txs.push(worker_tx); - let tx_worker = tx.clone(); - let pin_cores = config.pin_cores; - - s.spawn(move || { - // Pin to core. - if pin_cores { - let _ = core_affinity::set_for_current(core_affinity::CoreId { id: worker_id }); - } - - let mut worker = wf(worker_id); - while let Ok(msg) = worker_rx.recv() { - match msg { - WorkerMsg::Batch(batch) => { - for input in batch { - let borrowed = input.as_sketch_input(); - worker.process(&borrowed, &mut |delta| { - tx_worker.send(delta).expect( - "aggregator receiver dropped while workers still running", - ); - }); - } - } - WorkerMsg::End => break, - } - } - }); - } - - // Spawn dispatcher thread: hash-route each input to a worker channel. - s.spawn(move || { - let mut pending = vec![Vec::new(); num_workers]; - for item in inputs { - let owned = OwnedSketchInput::from(item); - let borrowed = owned.as_sketch_input(); - let worker_id = worker_index(hash64_seeded(PARTITION_SEED, &borrowed), num_workers); - pending[worker_id].push(owned); - if pending[worker_id].len() >= WORKER_DISPATCH_BATCH_SIZE { - flush_worker_batch(&worker_input_txs, &mut pending, worker_id); - } - } - - flush_all_worker_batches(&worker_input_txs, &mut pending); - send_end_to_all_workers(&worker_input_txs); - }); - - // Drop the main sender so the receiver closes once all worker clones drop. - drop(tx); - - // Aggregator collect loop: block until a delta arrives; exits when senders are dropped. - for delta in rx { - parent.apply(delta); - } - }); - - OctoResult { parent } + let mut runtime = OctoRuntime::new(config, worker_factory, parent_factory); + for input in inputs { + runtime.insert(input.clone()); + } + runtime.finish() } #[cfg(test)] @@ -584,12 +461,12 @@ mod tests { // Insert CM_PROMASK-1 times: no delta yet. for _ in 0..(crate::CM_PROMASK - 1) { - child.insert_and_emit(&key, |d| deltas.push(d)); + child.insert(&key, &mut |d| deltas.push(d)); } assert!(deltas.is_empty(), "should not emit before threshold"); // One more insert triggers the delta. - child.insert_and_emit(&key, |d| deltas.push(d)); + child.insert(&key, &mut |d| deltas.push(d)); assert_eq!(deltas.len(), 3, "should emit one delta per row (3 rows)"); for d in &deltas { assert_eq!(d.value, crate::CM_PROMASK); @@ -619,7 +496,7 @@ mod tests { // Insert enough times to trigger at least one delta. for _ in 0..200 { - child.insert_and_emit(&key, |d| deltas.push(d)); + child.insert(&key, &mut |d| deltas.push(d)); } // With COUNT_PROMASK = 63 and 3 rows, after 200 inserts we // should have emitted at least 3 deltas (one per row per threshold). @@ -648,12 +525,12 @@ mod tests { let mut deltas: Vec = Vec::new(); // First insert should always emit (register goes from 0 to something > 0). - child.insert_and_emit(&SketchInput::U64(1), |d| deltas.push(d)); + child.insert(&SketchInput::U64(1), &mut |d| deltas.push(d)); assert_eq!(deltas.len(), 1, "first insert should emit a delta"); // Inserting the same value again should NOT emit (register unchanged). let len_before = deltas.len(); - child.insert_and_emit(&SketchInput::U64(1), |d| deltas.push(d)); + child.insert(&SketchInput::U64(1), &mut |d| deltas.push(d)); assert_eq!(deltas.len(), len_before, "duplicate should not emit"); } @@ -887,6 +764,74 @@ mod tests { ); } + #[test] + fn octo_runtime_live_read_handle_observes_progress() { + let config = OctoConfig { + num_workers: 2, + pin_cores: false, + queue_capacity: 4096, + }; + let mut runtime = OctoRuntime::new( + &config, + |_| CountingWorker, + || CountingParent { total: 0 }, + ); + let reader = runtime.read_handle(); + + for i in 0..64u64 { + runtime.insert(SketchInput::U64(i)); + } + std::thread::sleep(std::time::Duration::from_millis(5)); + let observed = reader.with_parent(|p| p.total); + assert!(observed <= 64, "live reader should observe a partial or complete total"); + + let result = runtime.finish(); + assert_eq!(result.parent.total, 64, "all inserted items should be accounted for"); + assert!( + result.parent.total >= observed, + "final total should not be less than live snapshot" + ); + } + + #[test] + fn octo_runtime_close_is_idempotent() { + let config = OctoConfig { + num_workers: 2, + pin_cores: false, + queue_capacity: 4096, + }; + let runtime = OctoRuntime::new( + &config, + |_| HllOctoWorker::new(), + || HllOctoParent { + sketch: HyperLogLog::::default(), + }, + ); + runtime.close(); + runtime.close(); + let result = runtime.finish(); + assert_eq!(result.parent.sketch.estimate(), 0); + } + + #[test] + #[should_panic(expected = "cannot insert after runtime has been closed")] + fn octo_runtime_insert_after_close_panics() { + let config = OctoConfig { + num_workers: 2, + pin_cores: false, + queue_capacity: 4096, + }; + let mut runtime = OctoRuntime::new( + &config, + |_| HllOctoWorker::new(), + || HllOctoParent { + sketch: HyperLogLog::::default(), + }, + ); + runtime.close(); + runtime.insert(SketchInput::U64(1)); + } + #[test] fn octo_runtime_empty_stream_finishes() { let config = OctoConfig { @@ -923,7 +868,7 @@ mod tests { total: u64, } - impl OctoParent for CountingParent { + impl OctoAggregator for CountingParent { type Delta = u64; fn apply(&mut self, delta: Self::Delta) { @@ -932,69 +877,76 @@ mod tests { } #[test] - fn worker_index_is_deterministic_for_pow2_and_non_pow2_workers() { - let key = SketchInput::U64(0xdead_beef); - let hash = hash64_seeded(PARTITION_SEED, &key); - - let idx_pow2_a = worker_index(hash, 8); - let idx_pow2_b = worker_index(hash, 8); - assert_eq!(idx_pow2_a, idx_pow2_b); - assert_eq!(idx_pow2_a, (hash as usize) & 7); - - let idx_non_pow2_a = worker_index(hash, 6); - let idx_non_pow2_b = worker_index(hash, 6); - assert_eq!(idx_non_pow2_a, idx_non_pow2_b); - assert_eq!(idx_non_pow2_a, (hash as usize) % 6); - } - - #[test] - fn run_octo_batch_flush_boundary_matches_reference_count() { + fn octo_runtime_close_preserves_queued_items_without_dispatcher() { let config = OctoConfig { - num_workers: 3, + num_workers: 4, pin_cores: false, queue_capacity: 4096, }; - let sizes = [ - WORKER_DISPATCH_BATCH_SIZE - 1, - WORKER_DISPATCH_BATCH_SIZE, - WORKER_DISPATCH_BATCH_SIZE + 1, - ]; - - for &n in &sizes { - let inputs: Vec> = (0..n as u64) - .map(|i| SketchInput::U64(i ^ 0x1234)) - .collect(); - let result = run_octo( - &inputs, - &config, - |_| CountingWorker, - || CountingParent { total: 0 }, - ); - assert_eq!( - result.parent.total, n as u64, - "batch boundary size {n} should process all items" - ); + let n = 257usize; + let mut runtime = OctoRuntime::new( + &config, + |_| CountingWorker, + || CountingParent { total: 0 }, + ); + + for i in 0..n as u64 { + runtime.insert(SketchInput::U64(i + 42)); + } + runtime.close(); + let result = runtime.finish(); + assert_eq!( + result.parent.total, n as u64, + "runtime close should preserve already queued items" + ); + } + + struct WorkerIdEmitter { + worker_id: usize, + } + + impl OctoWorker for WorkerIdEmitter { + type Delta = usize; + + fn process(&mut self, _input: &SketchInput, emit: &mut dyn FnMut(Self::Delta)) { + emit(self.worker_id); + } + } + + struct WorkerLoadParent { + loads: Vec, + } + + impl OctoAggregator for WorkerLoadParent { + type Delta = usize; + + fn apply(&mut self, delta: Self::Delta) { + self.loads[delta] += 1; } } #[test] - fn octo_runtime_end_flush_preserves_tail_items() { + fn octo_runtime_round_robin_selector_distributes_deterministically() { + let num_workers = 3; + let inserts = 10u64; let config = OctoConfig { - num_workers: 4, + num_workers, pin_cores: false, queue_capacity: 4096, }; - let n = WORKER_DISPATCH_BATCH_SIZE + 7; - let mut runtime = - OctoRuntime::new(&config, |_| CountingWorker, || CountingParent { total: 0 }); + let mut runtime = OctoRuntime::new( + &config, + |worker_id| WorkerIdEmitter { worker_id }, + || WorkerLoadParent { + loads: vec![0; num_workers], + }, + ); - for i in 0..n as u64 { - runtime.insert(SketchInput::U64(i + 42)); + for i in 0..inserts { + runtime.insert(SketchInput::U64(i)); } let result = runtime.finish(); - assert_eq!( - result.parent.total, n as u64, - "runtime finish should flush partial worker batches" - ); + + assert_eq!(result.parent.loads, vec![4, 3, 3]); } } diff --git a/src/sketches/count.rs b/src/sketches/count.rs index 059ddcd..35e643c 100644 --- a/src/sketches/count.rs +++ b/src/sketches/count.rs @@ -746,9 +746,9 @@ impl CountChild { } /// Insert a key with the Count sketch sign convention. - /// Emits `CountDelta` entries via `emit` when |counter| >= `COUNT_PROMASK`. + /// Emits `CountDelta` when |counter| >= `COUNT_PROMASK`, then resets the counter. #[inline(always)] - pub fn insert_and_emit(&mut self, value: &SketchInput, mut emit: impl FnMut(CountDelta)) { + pub fn insert(&mut self, value: &SketchInput, emit: &mut dyn FnMut(CountDelta)) { let rows = self.counts.rows(); let cols = self.counts.cols(); let data = self.counts.as_mut_slice(); @@ -793,6 +793,21 @@ mod tests { use crate::{SketchInput, hash64_seeded}; use std::collections::HashMap; + #[test] + fn count_child_insert_emits_at_threshold() { + let mut child = CountChild::with_dimensions(3, 64); + let key = SketchInput::U64(99); + let mut deltas: Vec = Vec::new(); + + for _ in 0..200 { + child.insert(&key, &mut |d| deltas.push(d)); + } + assert!( + deltas.len() >= 3, + "expected at least one promoted delta per row" + ); + } + fn counter_sign(row: usize, key: &SketchInput) -> i32 { let hash = hash64_seeded(row, key); if (hash >> 63) & 1 == 1 { 1 } else { -1 } diff --git a/src/sketches/countmin.rs b/src/sketches/countmin.rs index 4a38073..817a017 100644 --- a/src/sketches/countmin.rs +++ b/src/sketches/countmin.rs @@ -498,10 +498,9 @@ impl CountMinChild { self.counts.cols() } - /// Insert a key, emitting `CmDelta` entries via `emit` when a counter - /// reaches `CM_PROMASK`. The counter is reset to 0 after emission. + /// Insert a key and emit a `CmDelta` each time a counter reaches `CM_PROMASK`. #[inline(always)] - pub fn insert_and_emit(&mut self, value: &SketchInput, mut emit: impl FnMut(CmDelta)) { + pub fn insert(&mut self, value: &SketchInput, emit: &mut dyn FnMut(CmDelta)) { let rows = self.counts.rows(); let cols = self.counts.cols(); let data = self.counts.as_mut_slice(); @@ -546,6 +545,21 @@ mod tests { use core::f64; use std::collections::HashMap; + #[test] + fn countmin_child_insert_emits_at_threshold() { + let mut child = CountMinChild::with_dimensions(3, 64); + let key = SketchInput::U64(42); + let mut deltas: Vec = Vec::new(); + + for _ in 0..(CM_PROMASK - 1) { + child.insert(&key, &mut |d| deltas.push(d)); + } + assert!(deltas.is_empty(), "should not emit before threshold"); + + child.insert(&key, &mut |d| deltas.push(d)); + assert_eq!(deltas.len(), 3, "should emit one delta per row"); + } + fn run_zipf_stream( rows: usize, cols: usize, diff --git a/src/sketches/hll.rs b/src/sketches/hll.rs index 4eaa02c..e3835bb 100644 --- a/src/sketches/hll.rs +++ b/src/sketches/hll.rs @@ -353,9 +353,9 @@ impl Default for HllChild { } impl HllChild { - /// Insert using a pre-computed hash, emitting `HllDelta` on register improvement. + /// Insert using a pre-computed hash and emit `HllDelta` on register improvement. #[inline(always)] - pub fn insert_with_hash_and_emit(&mut self, hashed_val: u64, mut emit: impl FnMut(HllDelta)) { + pub fn insert_with_hash(&mut self, hashed_val: u64, emit: &mut dyn FnMut(HllDelta)) { let bucket_num = ((hashed_val >> HLL_Q) & HLL_P_MASK) as usize; let leading_zero = ((hashed_val << HLL_P) + HLL_P_MASK).leading_zeros() as u8 + 1; if leading_zero > self.registers[bucket_num] { @@ -367,11 +367,11 @@ impl HllChild { } } - /// Insert a key, emitting `HllDelta` on register improvement. + /// Insert a key and emit `HllDelta` on register improvement. #[inline(always)] - pub fn insert_and_emit(&mut self, obj: &SketchInput, emit: impl FnMut(HllDelta)) { + pub fn insert(&mut self, obj: &SketchInput, emit: &mut dyn FnMut(HllDelta)) { let hashed_val = hash64_seeded(CANONICAL_HASH_SEED, obj); - self.insert_with_hash_and_emit(hashed_val, emit); + self.insert_with_hash(hashed_val, emit); } } @@ -395,6 +395,19 @@ mod tests { const ERROR_TOLERANCE: f64 = 0.02; const SERDE_SAMPLE: usize = 100_000; + #[test] + fn hll_child_insert_emits_on_improvement() { + let mut child = HllChild::default(); + let mut deltas: Vec = Vec::new(); + + child.insert(&SketchInput::U64(1), &mut |d| deltas.push(d)); + assert_eq!(deltas.len(), 1, "first insert should improve one register"); + + let before = deltas.len(); + child.insert(&SketchInput::U64(1), &mut |d| deltas.push(d)); + assert_eq!(deltas.len(), before, "duplicate should not emit"); + } + trait HllEstimator: Default { fn push(&mut self, input: &SketchInput); fn insert_with_hash(&mut self, hashed: u64); From 930f18301c20d6aa7d341d7998b90a89659ec419 Mon Sep 17 00:00:00 2001 From: GordonYuanyc Date: Tue, 3 Mar 2026 00:22:11 -0700 Subject: [PATCH 4/7] cleanup, but the thread communication within octo is still too hard --- benches/octo_speedup.rs | 79 +++++++++++++---- src/sketch_framework/mod.rs | 4 +- src/sketch_framework/octo.rs | 81 +++++++++--------- src/sketches/count.rs | 45 +++------- src/sketches/countmin.rs | 160 +++++++++++++++++++---------------- src/sketches/hll.rs | 33 ++------ src/sketches/mod.rs | 4 - 7 files changed, 211 insertions(+), 195 deletions(-) diff --git a/benches/octo_speedup.rs b/benches/octo_speedup.rs index 7d445d0..15b6295 100644 --- a/benches/octo_speedup.rs +++ b/benches/octo_speedup.rs @@ -1,9 +1,9 @@ use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; use rand::{Rng, SeedableRng, rngs::StdRng}; use sketchlib_rust::{ - CmOctoParent, CmOctoWorker, Count, CountMin, CountOctoParent, CountOctoWorker, HllOctoParent, - HllOctoWorker, HyperLogLog, OctoAggregator, OctoConfig, OctoWorker, Regular, RegularPath, - SketchInput, Vector2D, + CmDelta, Count, CountMin, CountOctoParent, CountOctoWorker, HllOctoParent, HllOctoWorker, + HyperLogLog, OctoAggregator, OctoConfig, OctoWorker, Regular, RegularPath, SketchInput, + Vector2D, }; use std::sync::{Once, mpsc}; use std::thread; @@ -21,6 +21,38 @@ const MERGE_INTERVAL: usize = 10_000; static SANITY_ONCE: Once = Once::new(); +struct BenchCmOctoWorker { + sketch: CountMin, RegularPath>, +} + +impl BenchCmOctoWorker { + fn new(rows: usize, cols: usize) -> Self { + Self { + sketch: CountMin::with_dimensions(rows, cols), + } + } +} + +impl OctoWorker for BenchCmOctoWorker { + type Delta = CmDelta; + + fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(Self::Delta)) { + self.sketch.insert_emit_delta(input, emit); + } +} + +struct BenchCmOctoParent { + sketch: CountMin, RegularPath>, +} + +impl OctoAggregator for BenchCmOctoParent { + type Delta = CmDelta; + + fn apply(&mut self, delta: Self::Delta) { + self.sketch.apply_delta(delta); + } +} + fn build_inputs(sample_count: usize) -> Vec> { let mut rng = StdRng::seed_from_u64(RNG_SEED ^ (sample_count as u64)); (0..sample_count) @@ -110,11 +142,16 @@ fn run_periodic_countmin_full_merge_sharded( thread::scope(|s| { let (merge_tx, merge_rx) = mpsc::channel::(); let aggregator = s.spawn(move || { - let mut parent = - CountMin::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS); + let mut parent = CountMin::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ); let mut finished_workers = 0usize; while finished_workers < num_workers { - match merge_rx.recv().expect("CountMin merge channel closed unexpectedly") { + match merge_rx + .recv() + .expect("CountMin merge channel closed unexpectedly") + { CountMinMergeMsg::Snapshot(snapshot) => parent.merge(&snapshot), CountMinMergeMsg::Done => finished_workers += 1, } @@ -125,8 +162,10 @@ fn run_periodic_countmin_full_merge_sharded( for shard in shards { let merge_tx_worker = merge_tx.clone(); s.spawn(move || { - let mut child = - CountMin::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS); + let mut child = CountMin::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ); let mut local_inserts = 0usize; for input in shard { child.insert(input); @@ -168,7 +207,10 @@ fn run_periodic_count_full_merge_sharded( Count::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS); let mut finished_workers = 0usize; while finished_workers < num_workers { - match merge_rx.recv().expect("Count merge channel closed unexpectedly") { + match merge_rx + .recv() + .expect("Count merge channel closed unexpectedly") + { CountMergeMsg::Snapshot(snapshot) => parent.merge(&snapshot), CountMergeMsg::Done => finished_workers += 1, } @@ -179,8 +221,10 @@ fn run_periodic_count_full_merge_sharded( for shard in shards { let merge_tx_worker = merge_tx.clone(); s.spawn(move || { - let mut child = - Count::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS); + let mut child = Count::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ); let mut local_inserts = 0usize; for input in shard { child.insert(input); @@ -221,7 +265,10 @@ fn run_periodic_hll_full_merge_sharded( let mut parent = HyperLogLog::::default(); let mut finished_workers = 0usize; while finished_workers < num_workers { - match merge_rx.recv().expect("HLL merge channel closed unexpectedly") { + match merge_rx + .recv() + .expect("HLL merge channel closed unexpectedly") + { HllMergeMsg::Snapshot(snapshot) => parent.merge(&snapshot), HllMergeMsg::Done => finished_workers += 1, } @@ -275,8 +322,8 @@ fn sanity_check_periodic_baselines() { let octo_cm = run_octo_sharded( &shards, &config, - |_| CmOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), - || CmOctoParent { + |_| BenchCmOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), + || BenchCmOctoParent { sketch: CountMin::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), }, ) @@ -360,8 +407,8 @@ fn bench_countmin_merge_compare(c: &mut Criterion) { let parent = run_octo_sharded( &shards, &config, - |_| CmOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), - || CmOctoParent { + |_| BenchCmOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), + || BenchCmOctoParent { sketch: CountMin::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), }, ); diff --git a/src/sketch_framework/mod.rs b/src/sketch_framework/mod.rs index fd2488c..f5c1285 100644 --- a/src/sketch_framework/mod.rs +++ b/src/sketch_framework/mod.rs @@ -37,8 +37,8 @@ pub use eh_univ_optimized::{EHMapBucket, EHUnivMonBucket, EHUnivOptimized, EHUni pub mod octo; pub use octo::{ - CmOctoParent, CmOctoWorker, CountOctoParent, CountOctoWorker, HllOctoParent, HllOctoWorker, - OctoAggregator, OctoConfig, OctoReadHandle, OctoResult, OctoRuntime, OctoWorker, run_octo, + CountOctoParent, CountOctoWorker, HllOctoParent, HllOctoWorker, OctoAggregator, OctoConfig, + OctoReadHandle, OctoResult, OctoRuntime, OctoWorker, run_octo, }; pub mod tumbling; diff --git a/src/sketch_framework/octo.rs b/src/sketch_framework/octo.rs index c89776d..dc2725e 100644 --- a/src/sketch_framework/octo.rs +++ b/src/sketch_framework/octo.rs @@ -7,13 +7,13 @@ use std::marker::PhantomData; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::sync::{Arc, RwLock, Weak}; use std::sync::mpsc; +use std::sync::{Arc, RwLock, Weak}; use std::thread; use crate::{ - CmDelta, Count, CountChild, CountDelta, CountMin, CountMinChild, HllChild, HllDelta, - HyperLogLog, Regular, RegularPath, SketchInput, Vector2D, + CmDelta, Count, CountDelta, CountMin, HllDelta, HyperLogLog, Regular, RegularPath, SketchInput, + Vector2D, }; /// Legacy queue capacity default retained for config compatibility. @@ -208,9 +208,9 @@ where while let Ok(msg) = worker_rx.recv() { match msg { WorkerMsg::Data(input) => worker.process(&input, &mut |delta| { - delta_tx_worker.send(delta).expect( - "aggregator receiver dropped while workers still running", - ); + delta_tx_worker + .send(delta) + .expect("aggregator receiver dropped while workers still running"); }), WorkerMsg::End => break, } @@ -268,7 +268,8 @@ where panic!("cannot insert after runtime has been closed"); } - let worker_id = core.next_worker.fetch_add(1, Ordering::AcqRel) % core.worker_input_txs.len(); + let worker_id = + core.next_worker.fetch_add(1, Ordering::AcqRel) % core.worker_input_txs.len(); // SAFETY: caller explicitly guarantees borrowed data lives long enough. let static_input = unsafe { assume_input_static(input) }; core.worker_input_txs[worker_id] @@ -299,15 +300,15 @@ where // -- CountMin -- -/// OctoSketch worker backed by a `CountMinChild`. +/// OctoSketch worker backed by `CountMin`. pub struct CmOctoWorker { - child: CountMinChild, + sketch: CountMin, RegularPath>, } impl CmOctoWorker { pub fn new(rows: usize, cols: usize) -> Self { Self { - child: CountMinChild::with_dimensions(rows, cols), + sketch: CountMin::with_dimensions(rows, cols), } } } @@ -317,7 +318,7 @@ impl OctoWorker for CmOctoWorker { #[inline(always)] fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(CmDelta)) { - self.child.insert(input, emit); + self.sketch.insert_emit_delta(input, emit); } } @@ -337,15 +338,15 @@ impl OctoAggregator for CmOctoParent { // -- Count Sketch -- -/// OctoSketch worker backed by a `CountChild`. +/// OctoSketch worker backed by `Count`. pub struct CountOctoWorker { - child: CountChild, + child: Count, RegularPath>, } impl CountOctoWorker { pub fn new(rows: usize, cols: usize) -> Self { Self { - child: CountChild::with_dimensions(rows, cols), + child: Count::with_dimensions(rows, cols), } } } @@ -355,7 +356,7 @@ impl OctoWorker for CountOctoWorker { #[inline(always)] fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(CountDelta)) { - self.child.insert(input, emit); + self.child.insert_emit_delta(input, emit); } } @@ -375,15 +376,15 @@ impl OctoAggregator for CountOctoParent { // -- HyperLogLog -- -/// OctoSketch worker backed by an `HllChild`. +/// OctoSketch worker backed by `HyperLogLog`. pub struct HllOctoWorker { - child: HllChild, + child: HyperLogLog, } impl HllOctoWorker { pub fn new() -> Self { Self { - child: HllChild::default(), + child: HyperLogLog::default(), } } } @@ -399,7 +400,7 @@ impl OctoWorker for HllOctoWorker { #[inline(always)] fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(HllDelta)) { - self.child.insert(input, emit); + self.child.insert_emit_delta(input, emit); } } @@ -454,19 +455,19 @@ mod tests { // ----------------------------------------------------------------------- #[test] - fn cm_child_emits_delta_at_threshold() { - let mut child = CountMinChild::with_dimensions(3, 64); + fn cm_insert_emit_delta_emits_at_threshold() { + let mut worker_sketch = CountMin::with_dimensions(3, 64); let key = SketchInput::U64(42); let mut deltas: Vec = Vec::new(); // Insert CM_PROMASK-1 times: no delta yet. for _ in 0..(crate::CM_PROMASK - 1) { - child.insert(&key, &mut |d| deltas.push(d)); + worker_sketch.insert_emit_delta(&key, &mut |d| deltas.push(d)); } assert!(deltas.is_empty(), "should not emit before threshold"); // One more insert triggers the delta. - child.insert(&key, &mut |d| deltas.push(d)); + worker_sketch.insert_emit_delta(&key, &mut |d| deltas.push(d)); assert_eq!(deltas.len(), 3, "should emit one delta per row (3 rows)"); for d in &deltas { assert_eq!(d.value, crate::CM_PROMASK); @@ -490,13 +491,13 @@ mod tests { #[test] fn count_child_emits_delta_at_threshold() { - let mut child = CountChild::with_dimensions(3, 64); + let mut child = Count::, RegularPath>::with_dimensions(3, 64); let key = SketchInput::U64(99); let mut deltas: Vec = Vec::new(); // Insert enough times to trigger at least one delta. for _ in 0..200 { - child.insert(&key, &mut |d| deltas.push(d)); + child.insert_emit_delta(&key, &mut |d| deltas.push(d)); } // With COUNT_PROMASK = 63 and 3 rows, after 200 inserts we // should have emitted at least 3 deltas (one per row per threshold). @@ -521,16 +522,16 @@ mod tests { #[test] fn hll_child_emits_delta_on_improvement() { - let mut child = HllChild::default(); + let mut child = HyperLogLog::::default(); let mut deltas: Vec = Vec::new(); // First insert should always emit (register goes from 0 to something > 0). - child.insert(&SketchInput::U64(1), &mut |d| deltas.push(d)); + child.insert_emit_delta(&SketchInput::U64(1), &mut |d| deltas.push(d)); assert_eq!(deltas.len(), 1, "first insert should emit a delta"); // Inserting the same value again should NOT emit (register unchanged). let len_before = deltas.len(); - child.insert(&SketchInput::U64(1), &mut |d| deltas.push(d)); + child.insert_emit_delta(&SketchInput::U64(1), &mut |d| deltas.push(d)); assert_eq!(deltas.len(), len_before, "duplicate should not emit"); } @@ -771,11 +772,8 @@ mod tests { pin_cores: false, queue_capacity: 4096, }; - let mut runtime = OctoRuntime::new( - &config, - |_| CountingWorker, - || CountingParent { total: 0 }, - ); + let mut runtime = + OctoRuntime::new(&config, |_| CountingWorker, || CountingParent { total: 0 }); let reader = runtime.read_handle(); for i in 0..64u64 { @@ -783,10 +781,16 @@ mod tests { } std::thread::sleep(std::time::Duration::from_millis(5)); let observed = reader.with_parent(|p| p.total); - assert!(observed <= 64, "live reader should observe a partial or complete total"); + assert!( + observed <= 64, + "live reader should observe a partial or complete total" + ); let result = runtime.finish(); - assert_eq!(result.parent.total, 64, "all inserted items should be accounted for"); + assert_eq!( + result.parent.total, 64, + "all inserted items should be accounted for" + ); assert!( result.parent.total >= observed, "final total should not be less than live snapshot" @@ -884,11 +888,8 @@ mod tests { queue_capacity: 4096, }; let n = 257usize; - let mut runtime = OctoRuntime::new( - &config, - |_| CountingWorker, - || CountingParent { total: 0 }, - ); + let mut runtime = + OctoRuntime::new(&config, |_| CountingWorker, || CountingParent { total: 0 }); for i in 0..n as u64 { runtime.insert(SketchInput::U64(i + 42)); diff --git a/src/sketches/count.rs b/src/sketches/count.rs index 35e643c..e5eb443 100644 --- a/src/sketches/count.rs +++ b/src/sketches/count.rs @@ -2,7 +2,7 @@ use crate::{ DefaultMatrixI32, DefaultMatrixI64, DefaultMatrixI128, DefaultXxHasher, FastPath, FastPathHasher, FixedMatrix, MatrixFastHash, MatrixStorage, NitroTarget, QuickMatrixI64, QuickMatrixI128, RegularPath, SketchHasher, SketchInput, Vector1D, Vector2D, - compute_median_inline_f64, + compute_median_inline_f64, hash64_seeded, }; use rmp_serde::{ decode::Error as RmpDecodeError, encode::Error as RmpEncodeError, from_slice, to_vec_named, @@ -716,53 +716,28 @@ impl CountL2HH { } } -// --------------------------------------------------------------------------- -// OctoSketch child sketch for multi-threaded delta-based promotion. -// --------------------------------------------------------------------------- - use crate::octo_delta::{COUNT_PROMASK, CountDelta}; -/// Lightweight Count sketch child with i8 counters. -/// Used by OctoSketch workers: accumulates signed counts and emits -/// `CountDelta` entries when |counter| reaches the promotion threshold. -pub struct CountChild { - counts: Vector2D, -} - -impl CountChild { - /// Creates a child sketch matching the parent's dimensions. - pub fn with_dimensions(rows: usize, cols: usize) -> Self { - let mut counts = Vector2D::init(rows, cols); - counts.fill(0i8); - CountChild { counts } - } - - pub fn rows(&self) -> usize { - self.counts.rows() - } - - pub fn cols(&self) -> usize { - self.counts.cols() - } - +/// Worker-side update for OctoSketch delta promotion. +impl Count, RegularPath> { /// Insert a key with the Count sketch sign convention. /// Emits `CountDelta` when |counter| >= `COUNT_PROMASK`, then resets the counter. #[inline(always)] - pub fn insert(&mut self, value: &SketchInput, emit: &mut dyn FnMut(CountDelta)) { + pub fn insert_emit_delta(&mut self, value: &SketchInput, emit: &mut dyn FnMut(CountDelta)) { let rows = self.counts.rows(); let cols = self.counts.cols(); let data = self.counts.as_mut_slice(); for r in 0..rows { let hashed = hash64_seeded(r, value); let col = ((hashed & LOWER_32_MASK) as usize) % cols; - let sign: i8 = if ((hashed >> 63) & 1) == 1 { 1 } else { -1 }; + let sign: i32 = if ((hashed >> 63) & 1) == 1 { 1 } else { -1 }; let cell = &mut data[r * cols + col]; - *cell = cell.saturating_add(sign); - if cell.unsigned_abs() >= COUNT_PROMASK { + *cell += sign; + if cell.unsigned_abs() >= COUNT_PROMASK as u32 { emit(CountDelta { row: r as u16, col: col as u16, - value: *cell, + value: *cell as i8, }); *cell = 0; } @@ -795,12 +770,12 @@ mod tests { #[test] fn count_child_insert_emits_at_threshold() { - let mut child = CountChild::with_dimensions(3, 64); + let mut child = Count::, RegularPath>::with_dimensions(3, 64); let key = SketchInput::U64(99); let mut deltas: Vec = Vec::new(); for _ in 0..200 { - child.insert(&key, &mut |d| deltas.push(d)); + child.insert_emit_delta(&key, &mut |d| deltas.push(d)); } assert!( deltas.len() >= 3, diff --git a/src/sketches/countmin.rs b/src/sketches/countmin.rs index 817a017..e247d96 100644 --- a/src/sketches/countmin.rs +++ b/src/sketches/countmin.rs @@ -4,6 +4,7 @@ use rmp_serde::{ use serde::{Deserialize, Serialize}; use std::marker::PhantomData; +use crate::octo_delta::{CM_PROMASK, CmDelta}; use crate::{ DefaultMatrixI32, DefaultMatrixI64, DefaultMatrixI128, DefaultXxHasher, FastPath, FastPathHasher, FixedMatrix, MatrixStorage, NitroTarget, QuickMatrixI64, QuickMatrixI128, @@ -327,6 +328,57 @@ where /// Count-Min sketch with floating-point counters (no integer rounding). pub type CountMinF64 = CountMin, RegularPath, H>; +impl CountMin, RegularPath> { + /// Inserts one key and emits a `CmDelta` whenever a touched counter reaches + /// a `CM_PROMASK` multiple. + #[inline(always)] + pub fn insert_emit_delta(&mut self, value: &SketchInput, emit: &mut dyn FnMut(CmDelta)) { + let rows = self.counts.rows(); + let cols = self.counts.cols(); + let threshold = CM_PROMASK as i32; + for r in 0..rows { + let hashed = hash64_seeded(r, value); + let col = ((hashed & LOWER_32_MASK) as usize) % cols; + self.counts.increment_by_row(r, col, 1); + let current = self.counts.query_one_counter(r, col); + if current % threshold == 0 { + emit(CmDelta { + row: r as u16, + col: col as u16, + value: CM_PROMASK, + }); + } + } + } + /// Apply a `CmDelta` to a full-precision parent CountMin sketch. + pub fn apply_delta(&mut self, delta: CmDelta) { + self.counts + .increment_by_row(delta.row as usize, delta.col as usize, delta.value as i32); + } +} + +// Fast-path hashing adapter for Vector2D. +impl FastPathHasher for Vector2D +where + T: Copy + std::ops::AddAssign, +{ + #[inline(always)] + fn hash_for_matrix(&self, value: &SketchInput) -> MatrixHashType { + Vector2D::hash_for_matrix(self, value) + } +} + +// Fast-path hashing adapter for u64-backed storage. +impl FastPathHasher for S +where + S: MatrixStorage, +{ + #[inline(always)] + fn hash_for_matrix(&self, value: &SketchInput) -> u64 { + hash64_seeded(0, value) + } +} + // SketchInput adapters for the fast-path Count-Min update rule. // Fast-path CountMin operations using precomputed hashes. Uses PartialOrd for f64 support. impl CountMin @@ -469,72 +521,6 @@ impl NitroTarget for CountMin, FastPath, H> { } } -// --------------------------------------------------------------------------- -// OctoSketch child sketch for multi-threaded delta-based promotion. -// --------------------------------------------------------------------------- - -use crate::octo_delta::{CM_PROMASK, CmDelta}; - -/// Lightweight CountMin child sketch with u8 counters. -/// Used by OctoSketch workers: accumulates counts in u8 cells and emits -/// `CmDelta` entries when a counter reaches the promotion threshold. -pub struct CountMinChild { - counts: Vector2D, -} - -impl CountMinChild { - /// Creates a child sketch matching the parent's dimensions. - pub fn with_dimensions(rows: usize, cols: usize) -> Self { - let mut counts = Vector2D::init(rows, cols); - counts.fill(0u8); - CountMinChild { counts } - } - - pub fn rows(&self) -> usize { - self.counts.rows() - } - - pub fn cols(&self) -> usize { - self.counts.cols() - } - - /// Insert a key and emit a `CmDelta` each time a counter reaches `CM_PROMASK`. - #[inline(always)] - pub fn insert(&mut self, value: &SketchInput, emit: &mut dyn FnMut(CmDelta)) { - let rows = self.counts.rows(); - let cols = self.counts.cols(); - let data = self.counts.as_mut_slice(); - for r in 0..rows { - let hashed = hash64_seeded(r, value); - let col = ((hashed & LOWER_32_MASK) as usize) % cols; - let cell = &mut data[r * cols + col]; - *cell = cell.wrapping_add(1); - if *cell >= CM_PROMASK { - emit(CmDelta { - row: r as u16, - col: col as u16, - value: *cell, - }); - *cell = 0; - } - } - } -} - -/// Apply a `CmDelta` to a full-precision parent CountMin sketch. -impl CountMin -where - S::Counter: Copy + std::ops::AddAssign + From, -{ - pub fn apply_delta(&mut self, delta: CmDelta) { - self.counts.increment_by_row( - delta.row as usize, - delta.col as usize, - S::Counter::from(delta.value as i32), - ); - } -} - #[cfg(test)] mod tests { use super::*; @@ -546,18 +532,48 @@ mod tests { use std::collections::HashMap; #[test] - fn countmin_child_insert_emits_at_threshold() { - let mut child = CountMinChild::with_dimensions(3, 64); + fn countmin_insert_emit_delta_emits_at_threshold_and_resets_period() { + let mut sketch = CountMin::, RegularPath>::with_dimensions(3, 64); let key = SketchInput::U64(42); let mut deltas: Vec = Vec::new(); for _ in 0..(CM_PROMASK - 1) { - child.insert(&key, &mut |d| deltas.push(d)); + sketch.insert_emit_delta(&key, &mut |d| deltas.push(d)); + } + assert!( + deltas.is_empty(), + "regular CMS worker path should not emit before threshold" + ); + + sketch.insert_emit_delta(&key, &mut |d| deltas.push(d)); + assert_eq!( + deltas.len(), + 3, + "should emit one delta per row at threshold" + ); + assert!(deltas.iter().all(|d| d.value == CM_PROMASK)); + + for _ in 0..(CM_PROMASK - 1) { + sketch.insert_emit_delta(&key, &mut |d| deltas.push(d)); } - assert!(deltas.is_empty(), "should not emit before threshold"); + assert_eq!(deltas.len(), 3, "no second emission before next threshold"); + sketch.insert_emit_delta(&key, &mut |d| deltas.push(d)); + assert_eq!(deltas.len(), 6, "should emit again on next threshold"); + } - child.insert(&key, &mut |d| deltas.push(d)); - assert_eq!(deltas.len(), 3, "should emit one delta per row"); + #[test] + fn countmin_apply_delta_increments_parent_counter() { + let mut parent = CountMin::, RegularPath>::with_dimensions(3, 64); + let delta = CmDelta { + row: 1, + col: 5, + value: CM_PROMASK, + }; + parent.apply_delta(delta); + assert_eq!( + parent.as_storage().query_one_counter(1, 5), + CM_PROMASK as i32 + ); } fn run_zipf_stream( diff --git a/src/sketches/hll.rs b/src/sketches/hll.rs index e3835bb..191278b 100644 --- a/src/sketches/hll.rs +++ b/src/sketches/hll.rs @@ -331,31 +331,12 @@ impl HyperLogLogHIP { } } } -// --------------------------------------------------------------------------- -// OctoSketch child sketch for multi-threaded delta-based promotion. -// --------------------------------------------------------------------------- - use crate::octo_delta::HllDelta; -/// Lightweight HLL child sketch backed by a fixed register array. -/// Emits `HllDelta` whenever a register is improved (PROMASK = 0, -/// meaning every improvement is propagated immediately). -pub struct HllChild { - registers: Box<[u8; NUM_REGISTERS]>, -} - -impl Default for HllChild { - fn default() -> Self { - Self { - registers: Box::new([0u8; NUM_REGISTERS]), - } - } -} - -impl HllChild { +impl HyperLogLog { /// Insert using a pre-computed hash and emit `HllDelta` on register improvement. #[inline(always)] - pub fn insert_with_hash(&mut self, hashed_val: u64, emit: &mut dyn FnMut(HllDelta)) { + pub fn insert_emit_delta_with_hash(&mut self, hashed_val: u64, emit: &mut dyn FnMut(HllDelta)) { let bucket_num = ((hashed_val >> HLL_Q) & HLL_P_MASK) as usize; let leading_zero = ((hashed_val << HLL_P) + HLL_P_MASK).leading_zeros() as u8 + 1; if leading_zero > self.registers[bucket_num] { @@ -369,9 +350,9 @@ impl HllChild { /// Insert a key and emit `HllDelta` on register improvement. #[inline(always)] - pub fn insert(&mut self, obj: &SketchInput, emit: &mut dyn FnMut(HllDelta)) { + pub fn insert_emit_delta(&mut self, obj: &SketchInput, emit: &mut dyn FnMut(HllDelta)) { let hashed_val = hash64_seeded(CANONICAL_HASH_SEED, obj); - self.insert_with_hash(hashed_val, emit); + self.insert_emit_delta_with_hash(hashed_val, emit); } } @@ -397,14 +378,14 @@ mod tests { #[test] fn hll_child_insert_emits_on_improvement() { - let mut child = HllChild::default(); + let mut child = HyperLogLog::::default(); let mut deltas: Vec = Vec::new(); - child.insert(&SketchInput::U64(1), &mut |d| deltas.push(d)); + child.insert_emit_delta(&SketchInput::U64(1), &mut |d| deltas.push(d)); assert_eq!(deltas.len(), 1, "first insert should improve one register"); let before = deltas.len(); - child.insert(&SketchInput::U64(1), &mut |d| deltas.push(d)); + child.insert_emit_delta(&SketchInput::U64(1), &mut |d| deltas.push(d)); assert_eq!(deltas.len(), before, "duplicate should not emit"); } diff --git a/src/sketches/mod.rs b/src/sketches/mod.rs index cbe7e96..a119757 100644 --- a/src/sketches/mod.rs +++ b/src/sketches/mod.rs @@ -47,10 +47,6 @@ pub use cs_heap::CSHeap; pub mod octo_delta; pub use octo_delta::{CM_PROMASK, COUNT_PROMASK, CmDelta, CountDelta, HLL_PROMASK, HllDelta}; -pub use count::CountChild; -pub use countmin::CountMinChild; -pub use hll::HllChild; - pub mod fold_cms; pub use fold_cms::{FoldCMS, FoldCell, FoldEntry}; From 2d128eac3dc7b412b75cbb42ef4e376bed7ef58f Mon Sep 17 00:00:00 2001 From: GordonYuanyc Date: Mon, 9 Mar 2026 17:29:17 -0600 Subject: [PATCH 5/7] remove dyn FnMut to FnMut and change mpsc to multiple spsc --- Cargo.lock | 10 +++ Cargo.toml | 1 + benches/octo_speedup.rs | 69 +++++++++++---- src/sketch_framework/mod.rs | 4 +- src/sketch_framework/octo.rs | 159 ++++++++++++++++++++++++----------- src/sketches/count.rs | 2 +- src/sketches/countmin.rs | 6 +- src/sketches/hll.rs | 8 +- 8 files changed, 185 insertions(+), 74 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c2773a0..04baa67 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -251,6 +251,15 @@ dependencies = [ "itertools", ] +[[package]] +name = "crossbeam-channel" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -873,6 +882,7 @@ dependencies = [ "clap", "core_affinity", "criterion", + "crossbeam-channel", "pcap", "prost", "prost-build", diff --git a/Cargo.toml b/Cargo.toml index e437ba6..4d35056 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,7 @@ smallvec = "1.13.2" prost = "0.13" bytes = "1" core_affinity = "0.8" +crossbeam-channel = "0.5" [dev-dependencies] criterion = "0.5" diff --git a/benches/octo_speedup.rs b/benches/octo_speedup.rs index 15b6295..0600679 100644 --- a/benches/octo_speedup.rs +++ b/benches/octo_speedup.rs @@ -1,9 +1,10 @@ use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; +use crossbeam_channel::{TryRecvError, bounded}; use rand::{Rng, SeedableRng, rngs::StdRng}; use sketchlib_rust::{ - CmDelta, Count, CountMin, CountOctoParent, CountOctoWorker, HllOctoParent, HllOctoWorker, - HyperLogLog, OctoAggregator, OctoConfig, OctoWorker, Regular, RegularPath, SketchInput, - Vector2D, + CmDelta, Count, CountMin, CountOctoAggregator, CountOctoWorker, HllOctoAggregator, + HllOctoWorker, HyperLogLog, OctoAggregator, OctoConfig, OctoWorker, Regular, RegularPath, + SketchInput, Vector2D, }; use std::sync::{Once, mpsc}; use std::thread; @@ -36,16 +37,19 @@ impl BenchCmOctoWorker { impl OctoWorker for BenchCmOctoWorker { type Delta = CmDelta; - fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(Self::Delta)) { + fn process(&mut self, input: &SketchInput, emit: &mut F) + where + F: FnMut(Self::Delta), + { self.sketch.insert_emit_delta(input, emit); } } -struct BenchCmOctoParent { +struct BenchCmOctoAggregator { sketch: CountMin, RegularPath>, } -impl OctoAggregator for BenchCmOctoParent { +impl OctoAggregator for BenchCmOctoAggregator { type Delta = CmDelta; fn apply(&mut self, delta: Self::Delta) { @@ -83,20 +87,50 @@ where { let num_workers = shards.len(); thread::scope(|s| { - let (delta_tx, delta_rx) = mpsc::channel::(); + let queue_capacity = config.queue_capacity.max(1); + let mut worker_txs = Vec::with_capacity(num_workers); + let mut worker_rxs = Vec::with_capacity(num_workers); + for _ in 0..num_workers { + let (tx, rx) = bounded::(queue_capacity); + worker_txs.push(tx); + worker_rxs.push(Some(rx)); + } + let aggregator = s.spawn(move || { let mut parent = parent_factory(); if config.pin_cores { let _ = core_affinity::set_for_current(core_affinity::CoreId { id: num_workers }); } - for delta in delta_rx { - parent.apply(delta); + + let mut disconnected = 0usize; + while disconnected < num_workers { + let mut made_progress = false; + for rx_slot in &mut worker_rxs { + let Some(rx) = rx_slot else { + continue; + }; + match rx.try_recv() { + Ok(delta) => { + parent.apply(delta); + made_progress = true; + } + Err(TryRecvError::Empty) => {} + Err(TryRecvError::Disconnected) => { + *rx_slot = None; + disconnected += 1; + } + } + } + if !made_progress { + std::hint::spin_loop(); + } } parent }); - for (worker_id, shard) in shards.iter().enumerate() { - let delta_tx_worker = delta_tx.clone(); + for (worker_id, (shard, delta_tx_worker)) in + shards.iter().zip(worker_txs.into_iter()).enumerate() + { let mut worker = worker_factory(worker_id); let pin_cores = config.pin_cores; s.spawn(move || { @@ -112,7 +146,6 @@ where } }); } - drop(delta_tx); aggregator .join() @@ -323,7 +356,7 @@ fn sanity_check_periodic_baselines() { &shards, &config, |_| BenchCmOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), - || BenchCmOctoParent { + || BenchCmOctoAggregator { sketch: CountMin::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), }, ) @@ -338,7 +371,7 @@ fn sanity_check_periodic_baselines() { &shards, &config, |_| CountOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), - || CountOctoParent { + || CountOctoAggregator { sketch: Count::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), }, ) @@ -355,7 +388,7 @@ fn sanity_check_periodic_baselines() { &shards, &config, |_| HllOctoWorker::new(), - || HllOctoParent { + || HllOctoAggregator { sketch: HyperLogLog::::default(), }, ) @@ -408,7 +441,7 @@ fn bench_countmin_merge_compare(c: &mut Criterion) { &shards, &config, |_| BenchCmOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), - || BenchCmOctoParent { + || BenchCmOctoAggregator { sketch: CountMin::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), }, ); @@ -470,7 +503,7 @@ fn bench_count_merge_compare(c: &mut Criterion) { &shards, &config, |_| CountOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), - || CountOctoParent { + || CountOctoAggregator { sketch: Count::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), }, ); @@ -529,7 +562,7 @@ fn bench_hll_merge_compare(c: &mut Criterion) { &shards, &config, |_| HllOctoWorker::new(), - || HllOctoParent { + || HllOctoAggregator { sketch: HyperLogLog::::default(), }, ); diff --git a/src/sketch_framework/mod.rs b/src/sketch_framework/mod.rs index f5c1285..63113fa 100644 --- a/src/sketch_framework/mod.rs +++ b/src/sketch_framework/mod.rs @@ -37,8 +37,8 @@ pub use eh_univ_optimized::{EHMapBucket, EHUnivMonBucket, EHUnivOptimized, EHUni pub mod octo; pub use octo::{ - CountOctoParent, CountOctoWorker, HllOctoParent, HllOctoWorker, OctoAggregator, OctoConfig, - OctoReadHandle, OctoResult, OctoRuntime, OctoWorker, run_octo, + CountOctoAggregator, CountOctoWorker, HllOctoAggregator, HllOctoWorker, OctoAggregator, + OctoConfig, OctoReadHandle, OctoResult, OctoRuntime, OctoWorker, run_octo, }; pub mod tumbling; diff --git a/src/sketch_framework/octo.rs b/src/sketch_framework/octo.rs index dc2725e..17343e7 100644 --- a/src/sketch_framework/octo.rs +++ b/src/sketch_framework/octo.rs @@ -7,10 +7,11 @@ use std::marker::PhantomData; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::sync::mpsc; use std::sync::{Arc, RwLock, Weak}; use std::thread; +use crossbeam_channel::{Receiver, Sender, TryRecvError, bounded}; + use crate::{ CmDelta, Count, CountDelta, CountMin, HllDelta, HyperLogLog, Regular, RegularPath, SketchInput, Vector2D, @@ -28,10 +29,12 @@ pub trait OctoWorker: Send { type Delta: Copy + Send + 'static; /// Process one input and emit zero or more deltas. - fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(Self::Delta)); + fn process(&mut self, input: &SketchInput, emit: &mut F) + where + F: FnMut(Self::Delta); } -/// Parent-side trait: absorbs deltas into a full-precision sketch. +/// Aggregator-side trait: absorbs deltas into a full-precision sketch. pub trait OctoAggregator: Send { type Delta: Copy + Send + 'static; @@ -51,8 +54,7 @@ pub struct OctoConfig { /// Worker i is pinned to core i, aggregator to core num_workers. /// Silently skipped if pinning fails. pub pin_cores: bool, - /// Legacy queue capacity retained for compatibility (default: 65536). - /// Currently unused by the unbounded MPSC transport. + /// Queue capacity for bounded worker-input and worker-delta channels (default: 65536). pub queue_capacity: usize, } @@ -120,7 +122,7 @@ impl

OctoReadHandle

{ } struct OctoCore

{ - worker_input_txs: Vec>, + worker_input_txs: Vec>, next_worker: AtomicUsize, worker_handles: Vec>, aggregator_handle: Option>, @@ -171,35 +173,70 @@ impl

OctoCore

where P: Send + Sync + 'static, { - fn start(workers: Vec, parent: P, num_workers: usize, pin_cores: bool) -> Self + fn start( + workers: Vec, + parent: P, + num_workers: usize, + pin_cores: bool, + queue_capacity: usize, + ) -> Self where W: OctoWorker + 'static, P: OctoAggregator, { assert_eq!(workers.len(), num_workers); + let queue_capacity = queue_capacity.max(1); - let (delta_tx, delta_rx) = mpsc::channel::(); let parent = Arc::new(RwLock::new(parent)); let parent_for_aggregator = Arc::clone(&parent); + let mut delta_txs: Vec> = Vec::with_capacity(num_workers); + let mut delta_rxs: Vec>> = Vec::with_capacity(num_workers); + for _ in 0..num_workers { + let (tx, rx) = bounded::(queue_capacity); + delta_txs.push(tx); + delta_rxs.push(Some(rx)); + } let aggregator_handle = thread::spawn(move || { if pin_cores { let _ = core_affinity::set_for_current(core_affinity::CoreId { id: num_workers }); } - for delta in delta_rx { - let mut guard = parent_for_aggregator - .write() - .expect("parent lock poisoned in aggregator"); - guard.apply(delta); + + let mut disconnected_workers = 0usize; + while disconnected_workers < num_workers { + let mut made_progress = false; + for rx_slot in &mut delta_rxs { + let Some(rx) = rx_slot else { + continue; + }; + match rx.try_recv() { + Ok(delta) => { + let mut guard = parent_for_aggregator + .write() + .expect("parent lock poisoned in aggregator"); + guard.apply(delta); + made_progress = true; + } + Err(TryRecvError::Empty) => {} + Err(TryRecvError::Disconnected) => { + *rx_slot = None; + disconnected_workers += 1; + } + } + } + if !made_progress { + std::hint::spin_loop(); + } } }); let mut worker_input_txs = Vec::with_capacity(num_workers); let mut worker_handles = Vec::with_capacity(num_workers); - for (worker_id, mut worker) in workers.into_iter().enumerate() { - let (worker_tx, worker_rx) = mpsc::channel::(); + for (worker_id, (mut worker, delta_tx_worker)) in + workers.into_iter().zip(delta_txs.into_iter()).enumerate() + { + let (worker_tx, worker_rx) = bounded::(queue_capacity); worker_input_txs.push(worker_tx); - let delta_tx_worker = delta_tx.clone(); let pin_cores = pin_cores; worker_handles.push(thread::spawn(move || { if pin_cores { @@ -217,7 +254,6 @@ where } })); } - drop(delta_tx); Self { worker_input_txs, @@ -243,7 +279,13 @@ where let num_workers = config.num_workers.max(1); let workers: Vec = (0..num_workers).map(worker_factory).collect(); let parent = parent_factory(); - let core = OctoCore::start(workers, parent, num_workers, config.pin_cores); + let core = OctoCore::start( + workers, + parent, + num_workers, + config.pin_cores, + config.queue_capacity, + ); Self { core: Some(core), @@ -317,17 +359,20 @@ impl OctoWorker for CmOctoWorker { type Delta = CmDelta; #[inline(always)] - fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(CmDelta)) { + fn process(&mut self, input: &SketchInput, emit: &mut F) + where + F: FnMut(Self::Delta), + { self.sketch.insert_emit_delta(input, emit); } } /// OctoSketch parent wrapping a full-precision `CountMin`. -pub struct CmOctoParent { +pub struct CmOctoAggregator { pub sketch: CountMin, RegularPath>, } -impl OctoAggregator for CmOctoParent { +impl OctoAggregator for CmOctoAggregator { type Delta = CmDelta; #[inline(always)] @@ -355,17 +400,20 @@ impl OctoWorker for CountOctoWorker { type Delta = CountDelta; #[inline(always)] - fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(CountDelta)) { + fn process(&mut self, input: &SketchInput, emit: &mut F) + where + F: FnMut(Self::Delta), + { self.child.insert_emit_delta(input, emit); } } /// OctoSketch parent wrapping a full-precision `Count`. -pub struct CountOctoParent { +pub struct CountOctoAggregator { pub sketch: Count, RegularPath>, } -impl OctoAggregator for CountOctoParent { +impl OctoAggregator for CountOctoAggregator { type Delta = CountDelta; #[inline(always)] @@ -399,17 +447,20 @@ impl OctoWorker for HllOctoWorker { type Delta = HllDelta; #[inline(always)] - fn process(&mut self, input: &SketchInput, emit: &mut dyn FnMut(HllDelta)) { + fn process(&mut self, input: &SketchInput, emit: &mut F) + where + F: FnMut(Self::Delta), + { self.child.insert_emit_delta(input, emit); } } /// OctoSketch parent wrapping a full-precision `HyperLogLog`. -pub struct HllOctoParent { +pub struct HllOctoAggregator { pub sketch: HyperLogLog, } -impl OctoAggregator for HllOctoParent { +impl OctoAggregator for HllOctoAggregator { type Delta = HllDelta; #[inline(always)] @@ -584,7 +635,7 @@ mod tests { &inputs, &config, |_| CmOctoWorker::new(rows, cols), - || CmOctoParent { + || CmOctoAggregator { sketch: CountMin::with_dimensions(rows, cols), }, ); @@ -621,7 +672,7 @@ mod tests { &inputs, &config, |_| HllOctoWorker::new(), - || HllOctoParent { + || HllOctoAggregator { sketch: HyperLogLog::::default(), }, ); @@ -653,7 +704,7 @@ mod tests { &inputs, &config, |_| CountOctoWorker::new(rows, cols), - || CountOctoParent { + || CountOctoAggregator { sketch: Count::with_dimensions(rows, cols), }, ); @@ -685,7 +736,7 @@ mod tests { &inputs, &config, |_| HllOctoWorker::new(), - || HllOctoParent { + || HllOctoAggregator { sketch: HyperLogLog::::default(), }, ); @@ -713,7 +764,7 @@ mod tests { &inputs, &config, |_| CmOctoWorker::new(rows, cols), - || CmOctoParent { + || CmOctoAggregator { sketch: CountMin::with_dimensions(rows, cols), }, ); @@ -721,7 +772,7 @@ mod tests { let mut runtime = OctoRuntime::new( &config, move |_| CmOctoWorker::new(rows, cols), - move || CmOctoParent { + move || CmOctoAggregator { sketch: CountMin::with_dimensions(rows, cols), }, ); @@ -748,7 +799,7 @@ mod tests { let mut runtime = OctoRuntime::new( &config, |_| HllOctoWorker::new(), - || HllOctoParent { + || HllOctoAggregator { sketch: HyperLogLog::::default(), }, ); @@ -772,8 +823,11 @@ mod tests { pin_cores: false, queue_capacity: 4096, }; - let mut runtime = - OctoRuntime::new(&config, |_| CountingWorker, || CountingParent { total: 0 }); + let mut runtime = OctoRuntime::new( + &config, + |_| CountingWorker, + || CountingAggregator { total: 0 }, + ); let reader = runtime.read_handle(); for i in 0..64u64 { @@ -807,7 +861,7 @@ mod tests { let runtime = OctoRuntime::new( &config, |_| HllOctoWorker::new(), - || HllOctoParent { + || HllOctoAggregator { sketch: HyperLogLog::::default(), }, ); @@ -828,7 +882,7 @@ mod tests { let mut runtime = OctoRuntime::new( &config, |_| HllOctoWorker::new(), - || HllOctoParent { + || HllOctoAggregator { sketch: HyperLogLog::::default(), }, ); @@ -846,7 +900,7 @@ mod tests { let runtime = OctoRuntime::new( &config, |_| HllOctoWorker::new(), - || HllOctoParent { + || HllOctoAggregator { sketch: HyperLogLog::::default(), }, ); @@ -863,16 +917,19 @@ mod tests { impl OctoWorker for CountingWorker { type Delta = u64; - fn process(&mut self, _input: &SketchInput, emit: &mut dyn FnMut(Self::Delta)) { + fn process(&mut self, _input: &SketchInput, emit: &mut F) + where + F: FnMut(Self::Delta), + { emit(1); } } - struct CountingParent { + struct CountingAggregator { total: u64, } - impl OctoAggregator for CountingParent { + impl OctoAggregator for CountingAggregator { type Delta = u64; fn apply(&mut self, delta: Self::Delta) { @@ -888,8 +945,11 @@ mod tests { queue_capacity: 4096, }; let n = 257usize; - let mut runtime = - OctoRuntime::new(&config, |_| CountingWorker, || CountingParent { total: 0 }); + let mut runtime = OctoRuntime::new( + &config, + |_| CountingWorker, + || CountingAggregator { total: 0 }, + ); for i in 0..n as u64 { runtime.insert(SketchInput::U64(i + 42)); @@ -909,16 +969,19 @@ mod tests { impl OctoWorker for WorkerIdEmitter { type Delta = usize; - fn process(&mut self, _input: &SketchInput, emit: &mut dyn FnMut(Self::Delta)) { + fn process(&mut self, _input: &SketchInput, emit: &mut F) + where + F: FnMut(Self::Delta), + { emit(self.worker_id); } } - struct WorkerLoadParent { + struct WorkerLoadAggregator { loads: Vec, } - impl OctoAggregator for WorkerLoadParent { + impl OctoAggregator for WorkerLoadAggregator { type Delta = usize; fn apply(&mut self, delta: Self::Delta) { @@ -938,7 +1001,7 @@ mod tests { let mut runtime = OctoRuntime::new( &config, |worker_id| WorkerIdEmitter { worker_id }, - || WorkerLoadParent { + || WorkerLoadAggregator { loads: vec![0; num_workers], }, ); diff --git a/src/sketches/count.rs b/src/sketches/count.rs index e5eb443..af52c8b 100644 --- a/src/sketches/count.rs +++ b/src/sketches/count.rs @@ -723,7 +723,7 @@ impl Count, RegularPath> { /// Insert a key with the Count sketch sign convention. /// Emits `CountDelta` when |counter| >= `COUNT_PROMASK`, then resets the counter. #[inline(always)] - pub fn insert_emit_delta(&mut self, value: &SketchInput, emit: &mut dyn FnMut(CountDelta)) { + pub fn insert_emit_delta(&mut self, value: &SketchInput, emit: &mut impl FnMut(CountDelta)) { let rows = self.counts.rows(); let cols = self.counts.cols(); let data = self.counts.as_mut_slice(); diff --git a/src/sketches/countmin.rs b/src/sketches/countmin.rs index e247d96..7bdf456 100644 --- a/src/sketches/countmin.rs +++ b/src/sketches/countmin.rs @@ -332,16 +332,16 @@ impl CountMin, RegularPath> { /// Inserts one key and emits a `CmDelta` whenever a touched counter reaches /// a `CM_PROMASK` multiple. #[inline(always)] - pub fn insert_emit_delta(&mut self, value: &SketchInput, emit: &mut dyn FnMut(CmDelta)) { + pub fn insert_emit_delta(&mut self, value: &SketchInput, emit: &mut impl FnMut(CmDelta)) { let rows = self.counts.rows(); let cols = self.counts.cols(); - let threshold = CM_PROMASK as i32; + // let threshold = CM_PROMASK as i32; for r in 0..rows { let hashed = hash64_seeded(r, value); let col = ((hashed & LOWER_32_MASK) as usize) % cols; self.counts.increment_by_row(r, col, 1); let current = self.counts.query_one_counter(r, col); - if current % threshold == 0 { + if current % CM_PROMASK as i32 == 0 { emit(CmDelta { row: r as u16, col: col as u16, diff --git a/src/sketches/hll.rs b/src/sketches/hll.rs index 191278b..848c107 100644 --- a/src/sketches/hll.rs +++ b/src/sketches/hll.rs @@ -336,7 +336,11 @@ use crate::octo_delta::HllDelta; impl HyperLogLog { /// Insert using a pre-computed hash and emit `HllDelta` on register improvement. #[inline(always)] - pub fn insert_emit_delta_with_hash(&mut self, hashed_val: u64, emit: &mut dyn FnMut(HllDelta)) { + pub fn insert_emit_delta_with_hash( + &mut self, + hashed_val: u64, + emit: &mut impl FnMut(HllDelta), + ) { let bucket_num = ((hashed_val >> HLL_Q) & HLL_P_MASK) as usize; let leading_zero = ((hashed_val << HLL_P) + HLL_P_MASK).leading_zeros() as u8 + 1; if leading_zero > self.registers[bucket_num] { @@ -350,7 +354,7 @@ impl HyperLogLog { /// Insert a key and emit `HllDelta` on register improvement. #[inline(always)] - pub fn insert_emit_delta(&mut self, obj: &SketchInput, emit: &mut dyn FnMut(HllDelta)) { + pub fn insert_emit_delta(&mut self, obj: &SketchInput, emit: &mut impl FnMut(HllDelta)) { let hashed_val = hash64_seeded(CANONICAL_HASH_SEED, obj); self.insert_emit_delta_with_hash(hashed_val, emit); } From bfdbc54582be7b13f6645d7d909e5046bb390ace Mon Sep 17 00:00:00 2001 From: GordonYuanyc Date: Mon, 9 Mar 2026 20:56:33 -0600 Subject: [PATCH 6/7] make CMS size larger in benchmark... surprisingly there is a difference --- benches/octo_speedup.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benches/octo_speedup.rs b/benches/octo_speedup.rs index 0600679..fcf7114 100644 --- a/benches/octo_speedup.rs +++ b/benches/octo_speedup.rs @@ -12,7 +12,7 @@ use std::time::Duration; const RNG_SEED: u64 = 0x0c70_2026_5eed_1234; const CM_COUNT_ROWS: usize = 3; -const CM_COUNT_COLS: usize = 4096; +const CM_COUNT_COLS: usize = 65536; const CM_COUNT_INPUTS: usize = 1_000_000; const HLL_INPUTS: usize = 500_000; const DOMAIN_MASK: u64 = (1 << 20) - 1; From 459781d39c1a7f45281f31a38d28fb9fc651ac86 Mon Sep 17 00:00:00 2001 From: GordonYuanyc Date: Mon, 30 Mar 2026 23:46:39 -0600 Subject: [PATCH 7/7] cargo fmt --- src/sketches/countmin.rs | 24 +----------------------- 1 file changed, 1 insertion(+), 23 deletions(-) diff --git a/src/sketches/countmin.rs b/src/sketches/countmin.rs index 7bdf456..7753f0d 100644 --- a/src/sketches/countmin.rs +++ b/src/sketches/countmin.rs @@ -8,7 +8,7 @@ use crate::octo_delta::{CM_PROMASK, CmDelta}; use crate::{ DefaultMatrixI32, DefaultMatrixI64, DefaultMatrixI128, DefaultXxHasher, FastPath, FastPathHasher, FixedMatrix, MatrixStorage, NitroTarget, QuickMatrixI64, QuickMatrixI128, - RegularPath, SketchHasher, SketchInput, Vector2D, + RegularPath, SketchHasher, SketchInput, Vector2D, hash64_seeded, }; const DEFAULT_ROW_NUM: usize = 3; @@ -357,28 +357,6 @@ impl CountMin, RegularPath> { } } -// Fast-path hashing adapter for Vector2D. -impl FastPathHasher for Vector2D -where - T: Copy + std::ops::AddAssign, -{ - #[inline(always)] - fn hash_for_matrix(&self, value: &SketchInput) -> MatrixHashType { - Vector2D::hash_for_matrix(self, value) - } -} - -// Fast-path hashing adapter for u64-backed storage. -impl FastPathHasher for S -where - S: MatrixStorage, -{ - #[inline(always)] - fn hash_for_matrix(&self, value: &SketchInput) -> u64 { - hash64_seeded(0, value) - } -} - // SketchInput adapters for the fast-path Count-Min update rule. // Fast-path CountMin operations using precomputed hashes. Uses PartialOrd for f64 support. impl CountMin