From 02a634b20206adcc43de2a9f93586a498c1897f0 Mon Sep 17 00:00:00 2001 From: Cole Helbling Date: Wed, 5 Aug 2026 19:02:24 -0700 Subject: [PATCH] Bloom filter Assisted-by: Amp --- Cargo.lock | 361 +++++++++++- Cargo.toml | 13 + benches/snapshot_build.rs | 135 +++++ src/bloom/mod.rs | 1090 +++++++++++++++++++++++++++++++++++++ src/bloom/source/mod.rs | 352 ++++++++++++ src/bloom/source/tests.rs | 184 +++++++ src/bloom/tests.rs | 330 +++++++++++ src/bloom/wire/mod.rs | 162 ++++++ src/bloom/wire/tests.rs | 108 ++++ src/lib.rs | 2 + 10 files changed, 2732 insertions(+), 5 deletions(-) create mode 100644 benches/snapshot_build.rs create mode 100644 src/bloom/mod.rs create mode 100644 src/bloom/source/mod.rs create mode 100644 src/bloom/source/tests.rs create mode 100644 src/bloom/tests.rs create mode 100644 src/bloom/wire/mod.rs create mode 100644 src/bloom/wire/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 95160f7..eef3a09 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,12 +11,33 @@ dependencies = [ "memchr", ] +[[package]] +name = "alloca" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" +dependencies = [ + "cc", +] + +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + [[package]] name = "anstyle" version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + [[package]] name = "block-buffer" version = "0.12.1" @@ -26,6 +47,18 @@ dependencies = [ "hybrid-array", ] +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "cc" version = "1.4.0" @@ -42,6 +75,33 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + [[package]] name = "clap" version = "4.6.5" @@ -94,6 +154,45 @@ dependencies = [ "libc", ] +[[package]] +name = "criterion" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3" +dependencies = [ + "alloca", + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "itertools", + "num-traits", + "oorandom", + "page_size", + "regex", + "serde", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" +dependencies = [ + "cast", + "itertools", +] + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "crypto-common" version = "0.2.2" @@ -130,7 +229,7 @@ dependencies = [ "proc-macro2", "quote", "scratch", - "syn", + "syn 3.0.3", ] [[package]] @@ -144,7 +243,7 @@ dependencies = [ "indexmap", "proc-macro2", "quote", - "syn", + "syn 3.0.3", ] [[package]] @@ -162,7 +261,7 @@ dependencies = [ "indexmap", "proc-macro2", "quote", - "syn", + "syn 3.0.3", ] [[package]] @@ -184,9 +283,15 @@ checksum = "c6232dd377dcc64799954cbd3a9bb882e9cdc1308ccd87b1c098f1fb2eaf82a8" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.3", ] +[[package]] +name = "either" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" + [[package]] name = "equivalent" version = "1.0.2" @@ -203,8 +308,10 @@ checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" name = "flakehub-cache-types" version = "0.1.0" dependencies = [ + "criterion", "cxx", "displaydoc", + "flakehub-cache-types", "hex", "lazy_static", "nix-base32", @@ -212,6 +319,8 @@ dependencies = [ "serde", "serde_json", "sha2", + "uuid", + "zerocopy", ] [[package]] @@ -220,6 +329,41 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" +[[package]] +name = "futures-core" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" + +[[package]] +name = "futures-task" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" + +[[package]] +name = "futures-util" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" +dependencies = [ + "futures-core", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.17.1" @@ -251,12 +395,32 @@ dependencies = [ "hashbrown", ] +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -290,6 +454,43 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d2628953ed836273ee4262e3708a8ef63ca38bd8a922070626eef7f9e5d8d536" +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + +[[package]] +name = "page_size" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" +dependencies = [ + "libc", + "winapi", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + [[package]] name = "proc-macro2" version = "1.0.107" @@ -337,6 +538,21 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "scratch" version = "1.0.9" @@ -370,7 +586,7 @@ checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.3", ] [[package]] @@ -403,12 +619,29 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + [[package]] name = "strsim" version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "syn" version = "3.0.3" @@ -429,6 +662,16 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "typenum" version = "1.20.1" @@ -447,6 +690,88 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" +[[package]] +name = "uuid" +version = "1.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239" +dependencies = [ + "js-sys", + "serde_core", + "wasm-bindgen", +] + +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn 2.0.119", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + [[package]] name = "winapi-util" version = "0.1.11" @@ -456,6 +781,12 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + [[package]] name = "windows-link" version = "0.2.1" @@ -471,6 +802,26 @@ dependencies = [ "windows-link", ] +[[package]] +name = "zerocopy" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5a105cd7b140f6eeec8acff2ea38135d3cab283ada58540f629fe51e46696eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fe976fb70c78cd64cccfe3a6fc142244e8a77b70959b30faf9d0ac37ee228eb" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "zmij" version = "1.0.23" diff --git a/Cargo.toml b/Cargo.toml index 266cab3..d0dcb71 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,9 @@ default = [] # Enables `From` for `StoreError`. cxx = ["dep:cxx"] +# Enables the bloom filter types +bloom = ["dep:uuid", "dep:zerocopy"] + [dependencies] cxx = { version = "1.0", optional = true } displaydoc = "0.2.4" @@ -27,4 +30,14 @@ regex = "1.8.3" serde = { version = "1.0.163", features = ["derive"] } serde_json = "1.0.96" sha2 = "0.11.0" +uuid = { version = "1", features = ["serde"], optional = true } +zerocopy = { version = "0.8.55", features = ["derive"], optional = true } + +[dev-dependencies] +criterion = { version = "0.8.2", default-features = false, features = ["cargo_bench_support"] } +flakehub-cache-types = { path = ".", features = ["bloom"] } +[[bench]] +name = "snapshot_build" +harness = false +required-features = ["bloom"] diff --git a/benches/snapshot_build.rs b/benches/snapshot_build.rs new file mode 100644 index 0000000..be3c5f4 --- /dev/null +++ b/benches/snapshot_build.rs @@ -0,0 +1,135 @@ +use std::hint::black_box; +use std::time::Duration; + +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use flakehub_cache_types::bloom::wire::ProbePositions; +use flakehub_cache_types::bloom::{ConcurrentBloomFilter, M_BITS, PROBE_COUNT, SnapshotManifest}; +use flakehub_cache_types::nix_store::StorePathHash; +use sha2::{Digest as _, Sha256}; +use uuid::Uuid; + +const BUILT_AT: u64 = 1_700_000_000; +const REPLAY_FROM: u64 = 1_699_999_700; + +fn snapshot_fixture() -> (ConcurrentBloomFilter, Uuid) { + let mut snapshot = ConcurrentBloomFilter::new(M_BITS, PROBE_COUNT, BUILT_AT).unwrap(); + for id in 0_u64..7_642 { + let digest = Sha256::digest(id.to_le_bytes()); + let encoded = nix_base32::to_nix_base32(&digest[..20]); + let hash = StorePathHash::new(&encoded).unwrap(); + snapshot + .insert_positions(&ProbePositions::of(&hash).unwrap()) + .unwrap(); + } + snapshot.set_source_max_row_id(456); + + let generation = "019bf5a7-f6e8-7ac0-b973-8536596bdb45".parse().unwrap(); + (snapshot, generation) +} + +fn snapshot_build(c: &mut Criterion) { + let (snapshot, generation) = snapshot_fixture(); + + // Warm the sparse, physically backed body before measuring snapshot preparation. + black_box(snapshot.stats()); + + let mut group = c.benchmark_group("snapshot_build_256_mib"); + group + .sample_size(20) + .warm_up_time(Duration::from_secs(3)) + .measurement_time(Duration::from_secs(15)) + .throughput(Throughput::Bytes(snapshot.body_len() as u64)); + + group.bench_function("legacy_repeated_passes", |b| { + b.iter(|| { + let fill_ratio = snapshot.fill_ratio(); + let estimated_fpr = snapshot.estimated_false_positive_rate(); + let estimated_distinct_items = snapshot.estimated_distinct_items(); + let encoded = snapshot.encode(); + let manifest = SnapshotManifest::for_snapshot(&snapshot, generation, REPLAY_FROM); + black_box(( + fill_ratio, + estimated_fpr, + estimated_distinct_items, + encoded, + manifest, + )) + }); + }); + + group.bench_function("prepare_snapshot", |b| { + b.iter(|| { + let stats = snapshot.stats(); + let encoded = snapshot.encode_with_checksum(); + let manifest = + SnapshotManifest::for_encoded_snapshot(&encoded, generation, REPLAY_FROM); + black_box((stats, encoded, manifest)) + }); + }); + + group.finish(); +} + +fn snapshot_decode(c: &mut Criterion) { + let (snapshot, generation) = snapshot_fixture(); + let encoded = snapshot.encode_with_checksum(); + let manifest = SnapshotManifest::for_encoded_snapshot(&encoded, generation, REPLAY_FROM); + + let mut group = c.benchmark_group("snapshot_decode_256_mib"); + group + .sample_size(20) + .warm_up_time(Duration::from_secs(3)) + .measurement_time(Duration::from_secs(20)) + .throughput(Throughput::Bytes(snapshot.body_len() as u64)); + + group.bench_function("decode", |b| { + b.iter(|| { + let decoded = ConcurrentBloomFilter::decode(black_box(encoded.as_bytes())).unwrap(); + black_box(decoded) + }); + }); + + group.bench_function("decode_and_validate_manifest", |b| { + b.iter(|| { + let decoded = ConcurrentBloomFilter::decode(black_box(encoded.as_bytes())).unwrap(); + manifest.validate_snapshot(&decoded).unwrap(); + black_box(decoded) + }); + }); + + group.bench_function("decode_validate_and_calculate_stats", |b| { + b.iter(|| { + let decoded = ConcurrentBloomFilter::decode(black_box(encoded.as_bytes())).unwrap(); + manifest.validate_snapshot(&decoded).unwrap(); + let stats = decoded.stats(); + black_box((decoded, stats)) + }); + }); + + group.bench_function("manifest_decode_to_concurrent", |b| { + b.iter(|| { + let decoded = ConcurrentBloomFilter::decode_with_manifest( + &manifest, + std::iter::once(black_box(encoded.as_bytes())), + ) + .unwrap(); + black_box(decoded) + }); + }); + + group.bench_function("manifest_decode_chunks_to_concurrent", |b| { + b.iter(|| { + let decoded = ConcurrentBloomFilter::decode_with_manifest( + &manifest, + black_box(encoded.as_bytes()).chunks(64 * 1024), + ) + .unwrap(); + black_box(decoded) + }); + }); + + group.finish(); +} + +criterion_group!(benches, snapshot_build, snapshot_decode); +criterion_main!(benches); diff --git a/src/bloom/mod.rs b/src/bloom/mod.rs new file mode 100644 index 0000000..ab09f4d --- /dev/null +++ b/src/bloom/mod.rs @@ -0,0 +1,1090 @@ +//! Shared Bloom-filter snapshot format for origin builders and edge readers. +//! +//! There are three kinds of parameters: +//! +//! - The [`FilterDefinition`] version changes how bytes are interpreted. A version bump creates +//! a new manifest namespace, so a new edge fails open until an explicit origin build populates +//! that namespace. +//! - `m_bits` and `k` are carried by each snapshot. Readers validate and use them rather than +//! assuming the current sizing policy. The [`M_BITS`] and [`PROBE_COUNT`] constants below define +//! the policy for newly built snapshots and path-created stream events. +//! - `built_at` and `source_max_row_id` describe one build and do not affect compatibility. +//! +//! The number of inserted items is deliberately not part of the format: what matters for +//! correctness is the false-positive rate, and both it and the distinct-item estimate are +//! calculated from the set bits. + +#![deny( + asm_sub_register, + deprecated, + missing_abi, + unsafe_code, + unused_macros, + unused_must_use, + unused_unsafe +)] +#![deny(clippy::from_over_into, clippy::needless_question_mark)] +#![cfg_attr( + not(debug_assertions), + deny(unused_imports, unused_mut, unused_variables) +)] + +pub mod source; +#[cfg(test)] +mod tests; +pub mod wire; + +use std::path::{Component, Path, PathBuf}; +use std::sync::atomic::{AtomicBool, AtomicU8, Ordering}; +use std::time::Duration; + +use serde::{Deserialize, Serialize}; +use sha2::{Digest as _, Sha256}; +use uuid::Uuid; +use zerocopy::byteorder::little_endian::U64; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Unaligned}; + +use crate::nix_store::StorePathHash; + +/// Magic bytes at the start of a Bloom snapshot, stored verbatim on the wire. +pub const SNAPSHOT_MAGIC: [u8; 4] = *b"ABF1"; +pub const FILTER_EXT: &str = "abf"; + +/// Fixed upper bound on the source-sync heartbeat cadence. +/// +/// A consumer that has caught up and seen no heartbeat for this long may treat publishing as +/// unavailable and fail open until it installs a newer snapshot. Run the source sync more often +/// than this deadline. +pub const MAX_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(6 * 60 * 60); + +/// Allowance for disagreement between the clocks of the publisher, the consumers, and Kinesis. +/// +/// Cross-machine safety comparisons subtract or add this slack: the snapshot build rewinds +/// `replay_from` by it, the consumer widens the heartbeat detection window by it, and a persisted +/// local filter rewinds its retention anchor by it. Transactions and publish paths run in +/// milliseconds, so five minutes is a large safety margin. +pub const CLOCK_SKEW_SLACK: Duration = Duration::from_secs(5 * 60); + +/// Maximum false-positive rate that builders and readers accept. +/// +/// This value is a safety limit. It is not a compatibility parameter. Readers calculate the rate +/// from the set bits in each snapshot. If you decrease this value, a reader can reject an old +/// snapshot that is above the new limit. The reader continues to decode valid snapshot bytes in the +/// same way. +pub const MAX_ACCEPTED_FALSE_POSITIVE_RATE: f64 = 0.1 / 100.; + +/// Arbitrary upper bound to protect against excessive CPU utilization. +pub const MAX_PROBE_COUNT: u8 = 64; + +/// The number of bits in one Bloom filter snapshot. This is the `m` parameter. +/// +/// The builder writes this value into each snapshot header. A reader gets the value from the +/// header, so a reader can decode snapshots of a different size. +/// +/// A larger value makes each snapshot larger and makes the false-positive rate smaller. A smaller +/// value makes each snapshot smaller and makes the false-positive rate larger. Increase this value +/// when the estimated false-positive rate of new snapshots comes near +/// [`MAX_ACCEPTED_FALSE_POSITIVE_RATE`]. The origin builder logs a recommended size when a build +/// exceeds the limit. Readers obtain the size from each snapshot header, so changing this policy +/// constant does not by itself make the snapshot format incompatible. +/// +/// The value must be a power of two, and it must not be smaller than 8. +/// +/// 2^31 gives us approximately 149 million distinct items before the false positive rate exceeds +/// [`MAX_ACCEPTED_FALSE_POSITIVE_RATE`]. +pub const M_BITS: u64 = 1 << 31; + +/// The number of bit positions that one item sets or tests. This is the `k` parameter. +/// +/// The builder writes this value into each snapshot header, and a reader gets it from the header. A +/// change to this value does not change the wire format, and it does not change the snapshot size. +/// A new snapshot with a different value replaces an old snapshot safely. +/// +/// To calculate the optimum value, use the target false-positive rate `p`: +/// +/// ```text +/// k = log2(1 / p), rounded to the nearest integer +/// ``` +/// +/// The value 10 comes from the target rate 0.001 ([`MAX_ACCEPTED_FALSE_POSITIVE_RATE`]): `log2(1 / +/// (0.1 / 100)) = 9.97`, thus 10. The formula works because a filter at its designed capacity has +/// half of its bits set. A false positive occurs when all `k` probes find a set bit, so `p = +/// 0.5^k`. +/// +/// A larger value makes each probe slower and fills the filter more quickly. A smaller value +/// makes the false-positive rate larger at the same fill ratio. +/// +/// Do not calculate `k` from the filter size with `k = (m / n) * ln(2)`. That formula gives the +/// optimum for a filter with no headroom. With a large [`M_BITS`], it gives a large `k` that makes +/// each lookup slower and sets more bits for each item. Calculate `k` from the target rate only. +/// Then calculate the required [`M_BITS`] from `k` and the expected item count `n`: +/// +/// ```text +/// target_fill = p ^ (1 / k) // the fill ratio that gives rate p +/// m = -k * n / ln(1 - target_fill) // bits required for n items +/// ``` +/// +/// Round `m` up to the next power of two, and add headroom for growth. The origin builder does +/// this calculation in `recommended_m_bits` and logs the result when a build goes above the rate +/// limit. Example: for `p = 0.001`, `k = 10`, and `n = 37,000,000`, the target fill is 0.501 and +/// `m` is approximately `5.3e8`, thus `2^29` bits before headroom. +/// +/// The value must be between 1 and [`MAX_PROBE_COUNT`]. +pub const PROBE_COUNT: u8 = 10; + +const _: () = { + assert!(M_BITS >= 8 && M_BITS.is_power_of_two()); + assert!(M_BITS / 8 <= usize::MAX as u64); + assert!(PROBE_COUNT >= 1 && PROBE_COUNT <= MAX_PROBE_COUNT); +}; + +/// The fixed header size for the current snapshot format, derived from the header struct's layout. +pub const SNAPSHOT_HEADER_LEN: usize = std::mem::size_of::(); +const _: () = assert!(SNAPSHOT_HEADER_LEN == 62); + +const MANIFEST_PREFIX: &str = "bloom"; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct FilterDefinition { + pub version: u8, +} + +impl FilterDefinition { + pub const CURRENT: Self = Self { version: 1 }; + + /// Relative directory containing snapshots compatible with this definition. + pub fn object_prefix(self) -> PathBuf { + Path::new(MANIFEST_PREFIX).join(format!("v{}", self.version)) + } + + /// Relative key of the mutable manifest for this definition. + pub fn manifest_key(self) -> PathBuf { + self.object_prefix().join("latest") + } + + fn snapshot_object_key( + self, + source_generation: Uuid, + built_at: u64, + source_max_row_id: u64, + ) -> PathBuf { + self.object_prefix() + .join(source_generation.to_string()) + .join(built_at.to_string()) + .join(source_max_row_id.to_string()) + .join("snapshot") + .with_extension(FILTER_EXT) + } +} + +#[derive(Clone, Copy, FromBytes, Immutable, IntoBytes, KnownLayout, Unaligned)] +#[repr(C)] +struct SnapshotHeader { + magic: [u8; 4], + version: u8, + m_bits: U64, + k: u8, + built_at: U64, + source_max_row_id: U64, + body_checksum: [u8; 32], +} + +#[derive(Clone, Copy, Debug)] +struct DecodedSnapshotHeader { + dims: Dimensions, + built_at: u64, + source_max_row_id: u64, + body_checksum: [u8; 32], +} + +impl DecodedSnapshotHeader { + fn decode(header: &SnapshotHeader) -> Result { + if header.magic != SNAPSHOT_MAGIC { + return Err(BloomError::InvalidMagic { + actual: header.magic, + }); + } + + if header.version != FilterDefinition::CURRENT.version { + return Err(BloomError::IncompatibleDefinition { + expected: FilterDefinition::CURRENT.version, + actual: header.version, + }); + } + + Ok(Self { + dims: Dimensions::new(header.m_bits.get(), header.k)?, + built_at: header.built_at.get(), + source_max_row_id: header.source_max_row_id.get(), + body_checksum: header.body_checksum, + }) + } +} + +/// Validated filter dimensions. +#[derive(Clone, Copy, Debug)] +struct Dimensions { + m_bits: u64, + k: u8, + body_len: usize, +} + +impl Dimensions { + /// The dimensions of the current sizing policy, [`M_BITS`] and [`PROBE_COUNT`]. The const + /// assertions next to those constants make an invalid policy a compile error, so this value + /// needs no runtime validation. + const CURRENT_POLICY: Self = Self { + m_bits: M_BITS, + k: PROBE_COUNT, + body_len: (M_BITS / 8) as usize, + }; + + fn new(m_bits: u64, k: u8) -> Result { + // Each body byte stores eight filter bits. Thus, `m_bits` must be divisible by eight. The + // probe calculation also requires `m_bits` to be a power of two. Each power of two that is + // at least eight is divisible by eight. This check makes sure that `m_bits / 8` gives the + // exact body length and does not discard a remainder. + if m_bits < 8 || !m_bits.is_power_of_two() { + return Err(BloomError::InvalidBitCount { m_bits }); + } + + if k == 0 || k > MAX_PROBE_COUNT { + return Err(BloomError::InvalidProbeCount { k }); + } + + let body_len = usize::try_from(m_bits / 8).map_err(|_| BloomError::TooLarge { m_bits })?; + + Ok(Self { + m_bits, + k, + body_len, + }) + } + + /// Calculates the probe positions for one store path hash using double hashing. Probe `i` is + /// `(h1 + i * h2) mod m_bits`. This gives `k` positions from one 20-byte hash. + /// + /// The store path hash is a cryptographic digest. Thus its bytes have a uniform distribution, + /// as required by the bloom filter. The method divides the digest into two hash values. `h1` is + /// bytes 0..8. `h2` is bytes 8..16. The method reads each value as a little-endian `u64`. + /// + /// Two operations use the property that `m_bits` is a power of two: + /// + /// - `h2 | 1` makes the stride an odd number. An odd number and a power of two have no common + /// divisor greater than one. Thus, the sequence can visit each bit position before it repeats. + /// Without this operation, the stride can share a divisor with `m_bits`. The sequence can then + /// repeat after it visits only a small group of positions. For example, with `m_bits = 16` and + /// a stride of 4, the sequence visits 0, 4, 8, and 12, and then returns to 0. + /// - `& mask`, with `mask = m_bits - 1`, does a modulo operation by `m_bits`. When you subtract + /// 1 from a power of two, all the lower bits of the result are set. Example: `1024 - 1 = + /// 0b11_1111_1111`. The AND operation keeps only the low `log2(m_bits)` bits. The result is the + /// same as `% m_bits`, but a division is not necessary. + /// + /// `wrapping_add` and `wrapping_mul` give the overflow a defined behavior. The mask discards + /// the high bits that overflow. + fn probe_positions( + self, + store_path_hash: &StorePathHash, + ) -> Result, BloomError> { + let decoded = decode_store_path_hash(store_path_hash)?; + let h1 = u64::from_le_bytes(decoded[0..8].try_into().expect("slice length is 8")); + let h2 = u64::from_le_bytes(decoded[8..16].try_into().expect("slice length is 8")) | 1; + let mask = self.m_bits - 1; + + Ok((0..u64::from(self.k)).map(move |i| h1.wrapping_add(i.wrapping_mul(h2)) & mask)) + } +} + +/// Occupancy-derived statistics for a Bloom snapshot. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct BloomSnapshotStats { + pub fill_ratio: f64, + pub estimated_false_positive_rate: f64, + pub estimated_distinct_items: Option, +} + +impl BloomSnapshotStats { + fn from_set_bits(dims: Dimensions, set_bits: u64) -> Self { + let fill_ratio = set_bits as f64 / dims.m_bits as f64; + let estimated_false_positive_rate = fill_ratio.powi(i32::from(dims.k)); + let estimated_distinct_items = if fill_ratio >= 1.0 { + None + } else { + let n = -(dims.m_bits as f64 / f64::from(dims.k)) * (1.0 - fill_ratio).ln(); + Some(n) + }; + + Self { + fill_ratio, + estimated_false_positive_rate, + estimated_distinct_items, + } + } + + fn validate(self) -> Result { + if self.estimated_false_positive_rate > MAX_ACCEPTED_FALSE_POSITIVE_RATE { + return Err(BloomError::ExcessiveFalsePositiveRate { + estimated_fpr: self.estimated_false_positive_rate, + limit: MAX_ACCEPTED_FALSE_POSITIVE_RATE, + }); + } + + Ok(self) + } +} + +/// Encoded snapshot bytes and the metadata captured in their header. +#[derive(Debug)] +pub struct EncodedBloomSnapshot { + bytes: Vec, + body_checksum: [u8; 32], + built_at: u64, + source_max_row_id: u64, +} + +impl EncodedBloomSnapshot { + pub fn as_bytes(&self) -> &[u8] { + &self.bytes + } + + pub fn into_bytes(self) -> Vec { + self.bytes + } + + pub fn body_checksum_hex(&self) -> String { + hex::encode(self.body_checksum) + } +} + +/// A Bloom filter that accepts concurrent probes and inserts without locks. +/// +/// The bits live in atomic bytes, so any number of threads can probe and insert at the same time. +/// Every bit access uses relaxed ordering: the bits are independent of each other, a set bit never +/// becomes unset, and a probe that races an insert is indistinguishable from a probe that ran one +/// moment earlier. No access establishes an ordering that another location depends on. +#[derive(Debug)] +pub struct ConcurrentBloomFilter { + dims: Dimensions, + built_at: u64, + source_max_row_id: u64, + bits: Vec, + decoded_body_checksum: Option, +} + +/// A checksum verified while decoding, reusable until the first insertion. +#[derive(Debug)] +struct DecodedBodyChecksum { + value: [u8; 32], + invalidated: AtomicBool, +} + +impl DecodedBodyChecksum { + fn new(value: [u8; 32]) -> Self { + Self { + value, + invalidated: AtomicBool::new(false), + } + } + + fn current(&self) -> Option<[u8; 32]> { + if self.invalidated.load(Ordering::SeqCst) { + None + } else { + Some(self.value) + } + } + + fn invalidate(&self) { + if !self.invalidated.load(Ordering::Relaxed) { + self.invalidated.store(true, Ordering::SeqCst); + } + } +} + +impl ConcurrentBloomFilter { + pub fn new(m_bits: u64, k: u8, built_at: u64) -> Result { + let dims = Dimensions::new(m_bits, k)?; + let bits = std::iter::repeat_with(|| AtomicU8::new(0)) + .take(dims.body_len) + .collect(); + + Ok(Self { + dims, + built_at, + source_max_row_id: 0, + bits, + decoded_body_checksum: None, + }) + } + + /// Sets probe positions from a path-created event. + /// + /// The [`wire::ProbePositions`] type guarantees positions that are in range for a power-of-two + /// `m_bits`, so the only failure left is a size mismatch. Both sizes are powers of two, and + /// each probe position is `(h1 + i * h2) & (m_bits - 1)`. Masking such a position by a smaller + /// filter's `m_bits - 1` keeps only the low bits, which gives exactly the position the smaller + /// filter calculates itself. Because of this, a filter accepts positions from a publisher with + /// an equal or larger size and masks them down. A smaller publisher size has already discarded + /// high bits that this filter needs, so that record is rejected and leaves the filter + /// unchanged. + pub fn insert_positions(&self, positions: &wire::ProbePositions) -> Result<(), BloomError> { + if positions.m_bits() < self.dims.m_bits { + return Err(BloomError::IncompatiblePositions { + event_m_bits: positions.m_bits(), + filter_m_bits: self.dims.m_bits, + }); + } + + // Position `i` depends only on `i` (`h1 + i*h2`), so a filter with a smaller `k` probes a + // prefix of the event's positions. An event with fewer positions than this filter's `k` + // leaves required probe bits unset, which would create a false negative. + let k = usize::from(self.dims.k); + if positions.positions().len() < k { + return Err(BloomError::InsufficientPositions { + count: positions.positions().len(), + k: self.dims.k, + }); + } + + // A decoded filter starts with a checksum that matches its bits. Mark that checksum stale + // before setting the first bit so it is never used for modified data. New filters have no + // cached checksum, and later inserts find an already-stale checksum. + if let Some(checksum) = &self.decoded_body_checksum { + checksum.invalidate(); + } + + let mask = self.dims.m_bits - 1; + for position in positions.positions().iter().take(k) { + // A bit position points to byte `position / 8`, and to bit `position % 8` in that byte. + // The bit order is least-significant bit first. `1 << (position % 8)` makes a mask that + // has one bit set. + let position = position & mask; + let byte_pos = (position / 8) as usize; + let bit_pos = position % 8; + self.bits[byte_pos].fetch_or(1 << bit_pos, Ordering::Relaxed); + } + + Ok(()) + } + + /// Tests one hash against this snapshot. A negative means the required bits are not all set; + /// whether that proves absence from the origin depends on the snapshot's external coverage and + /// freshness guarantees. + pub fn contains(&self, store_path_hash: &StorePathHash) -> Result { + // This uses the same byte and bit addresses as `insert_positions`. The AND operation with + // the one-bit mask isolates the probed bit. The result is not zero only when that bit is + // set. The key is possibly in the filter only if all `k` probed bits are set. If one probed + // bit is not set, the key was never inserted. + Ok(self.dims.probe_positions(store_path_hash)?.all(|position| { + let byte_pos = (position / 8) as usize; + let bit_pos = position % 8; + self.bits[byte_pos].load(Ordering::Relaxed) & (1 << bit_pos) != 0 + })) + } + + pub fn m_bits(&self) -> u64 { + self.dims.m_bits + } + + pub fn k(&self) -> u8 { + self.dims.k + } + + pub fn built_at(&self) -> u64 { + self.built_at + } + + /// Returns the ordinary source-sync cursor recorded in this snapshot. Bootstrap overlay shards + /// do not advance it. + pub fn source_max_row_id(&self) -> u64 { + self.source_max_row_id + } + + /// Sets the ordinary source-sync cursor. + pub fn set_source_max_row_id(&mut self, source_max_row_id: u64) { + self.source_max_row_id = source_max_row_id; + } + + pub fn body_checksum(&self) -> [u8; 32] { + if let Some(checksum) = self + .decoded_body_checksum + .as_ref() + .and_then(|cs| cs.current()) + { + return checksum; + } + + // Copy atomic bytes in bounded chunks instead of allocating a second full filter body. + let mut hasher = Sha256::new(); + let mut copied = Vec::with_capacity(64 * 1024); + for chunk in self.bits.chunks(64 * 1024) { + copied.clear(); + copied.extend(chunk.iter().map(|byte| byte.load(Ordering::Relaxed))); + hasher.update(&copied); + } + + hasher.finalize().into() + } + + pub fn body_checksum_hex(&self) -> String { + hex::encode(self.body_checksum()) + } + + pub fn body_len(&self) -> usize { + self.bits.len() + } + + /// Calculates all occupancy-derived statistics with one scan of the snapshot body. + pub fn stats(&self) -> BloomSnapshotStats { + let set_bits: u64 = self + .bits + .iter() + .map(|byte| u64::from(byte.load(Ordering::Relaxed).count_ones())) + .sum(); + BloomSnapshotStats::from_set_bits(self.dims, set_bits) + } + + /// The fraction of body bits that are set. This is the fill ratio `p` that the two estimator + /// methods below use. + pub fn fill_ratio(&self) -> f64 { + self.stats().fill_ratio + } + + /// Estimates the false-positive probability from the measured bit occupancy. + /// + /// A false positive occurs when a key was never inserted, but all its `k` probed bits are set. + /// For a key that was never inserted, the probe positions are effectively random. Thus each + /// probe finds a set bit with a probability that is equal to the fill ratio `p`. The `k` probes + /// are approximately independent. Thus the probability that all `k` probes find a set bit is + /// `p^k`. + /// + /// This estimate uses the measured fill ratio, not a theoretical one. Thus the estimate is + /// correct also when the filter contains more items than its sizing policy planned. + pub fn estimated_false_positive_rate(&self) -> f64 { + self.stats().estimated_false_positive_rate + } + + /// Estimates the number of distinct inserted keys from the measured bit occupancy. + /// + /// The filter does not store the keys, and a duplicate insert does not change the bits. But the + /// fill ratio `p` is a known function of the number of distinct keys `n`. Thus you can + /// calculate `n` from the measured `p`. + /// + /// The forward direction: each insert sets `k` of the `m_bits` positions. After `n` distinct + /// inserts, one given bit is 0 with probability `(1 - 1/m_bits)^(k * n)`, which is + /// approximately `e^(-k * n / m_bits)`. Thus the expected fill ratio is `p = 1 - e^(-k * n / + /// m_bits)`. + /// + /// The reverse direction, solved for `n`: + /// + /// `n = -(m_bits / k) * ln(1 - p)` + /// + /// When `p` is 1, all the bits are set. The formula then contains `ln(0)` and cannot give a + /// finite estimate. The method returns `None` for this condition. + pub fn estimated_distinct_items(&self) -> Option { + self.stats().estimated_distinct_items + } + + fn header(&self, body_checksum: [u8; 32]) -> SnapshotHeader { + let definition = FilterDefinition::CURRENT; + SnapshotHeader { + magic: SNAPSHOT_MAGIC, + version: definition.version, + m_bits: self.m_bits().into(), + k: self.k(), + built_at: self.built_at().into(), + source_max_row_id: self.source_max_row_id().into(), + body_checksum, + } + } + + /// Serialize the fixed little-endian header followed by the bit array. + pub fn encode(&self) -> Vec { + self.encode_with_checksum().into_bytes() + } + + /// Serializes the snapshot and returns the body checksum calculated for its header. + pub fn encode_with_checksum(&self) -> EncodedBloomSnapshot { + let mut encoded = Vec::with_capacity(SNAPSHOT_HEADER_LEN + self.body_len()); + encoded.resize(SNAPSHOT_HEADER_LEN, 0); + encoded.extend(self.bits.iter().map(|byte| byte.load(Ordering::Relaxed))); + + let body_checksum = checksum(&encoded[SNAPSHOT_HEADER_LEN..]); + let header = self.header(body_checksum); + encoded[..SNAPSHOT_HEADER_LEN].copy_from_slice(header.as_bytes()); + + EncodedBloomSnapshot { + bytes: encoded, + body_checksum, + built_at: header.built_at.get(), + source_max_row_id: header.source_max_row_id.get(), + } + } + + /// Writes a point-in-time copy of the filter. The output is identical to [`Self::encode`]. + pub fn write_to(&self, writer: &mut W) -> std::io::Result<()> { + writer.write_all(self.encode_with_checksum().as_bytes()) + } + + #[cfg(test)] + fn copied_bits(&self) -> Vec { + self.bits + .iter() + .map(|byte| byte.load(Ordering::Relaxed)) + .collect() + } + + /// Decode and fully validate a snapshot. + pub fn decode(encoded: &[u8]) -> Result { + let (header, body) = + SnapshotHeader::ref_from_prefix(encoded).map_err(|_| BloomError::Truncated { + expected_at_least: SNAPSHOT_HEADER_LEN, + actual: encoded.len(), + })?; + + let header = DecodedSnapshotHeader::decode(header)?; + if body.len() != header.dims.body_len { + return Err(BloomError::InvalidBodyLength { + expected: header.dims.body_len, + actual: body.len(), + }); + } + + let actual_checksum = checksum(body); + if actual_checksum != header.body_checksum { + return Err(BloomError::ChecksumMismatch { + expected: hex::encode(header.body_checksum), + actual: hex::encode(actual_checksum), + }); + } + + let snapshot = Self { + dims: header.dims, + built_at: header.built_at, + source_max_row_id: header.source_max_row_id, + bits: body.iter().copied().map(AtomicU8::new).collect(), + decoded_body_checksum: Some(DecodedBodyChecksum::new(actual_checksum)), + }; + + snapshot.stats().validate()?; + + Ok(snapshot) + } + + /// Decodes segmented snapshot bytes directly into the concurrent representation and validates + /// the final filter against both the snapshot header and its manifest. + /// + /// Chunks may split the fixed header or body at any byte. The filter is not returned until its + /// dimensions, body length, checksum, manifest identity, and false-positive rate all pass + /// validation. Statistics are calculated while the chunks are copied into atomic bytes. + pub fn decode_with_manifest( + manifest: &SnapshotManifest, + encoded_chunks: I, + ) -> Result<(Self, BloomSnapshotStats), BloomError> + where + I: IntoIterator, + C: AsRef<[u8]>, + { + manifest.validate_definition()?; + + let mut header_bytes = [0; SNAPSHOT_HEADER_LEN]; + let mut header_len = 0; + let mut total_len: usize = 0; + let mut decoded_header = None; + let mut bits = None; + let mut body_hasher = Sha256::new(); + let mut set_bits = 0_u64; + + for chunk in encoded_chunks { + let mut chunk = chunk.as_ref(); + total_len = total_len.saturating_add(chunk.len()); + + if header_len < SNAPSHOT_HEADER_LEN { + let copy_len = (SNAPSHOT_HEADER_LEN - header_len).min(chunk.len()); + header_bytes[header_len..header_len + copy_len].copy_from_slice(&chunk[..copy_len]); + header_len += copy_len; + chunk = &chunk[copy_len..]; + + if header_len < SNAPSHOT_HEADER_LEN { + continue; + } + + let (header, remainder) = SnapshotHeader::ref_from_prefix(&header_bytes) + .expect("the complete fixed-size header is valid"); + + debug_assert!(remainder.is_empty()); + + let header = DecodedSnapshotHeader::decode(header)?; + let mut body = Vec::new(); + body.try_reserve_exact(header.dims.body_len) + .map_err(|_| BloomError::TooLarge { + m_bits: header.dims.m_bits, + })?; + + bits = Some(body); + decoded_header = Some(header); + } + + if chunk.is_empty() { + continue; + } + + let header = decoded_header.expect("a body chunk follows a decoded header"); + let bits = bits.as_mut().expect("a decoded header allocates its body"); + let actual_len = bits.len().saturating_add(chunk.len()); + if actual_len > header.dims.body_len { + return Err(BloomError::InvalidBodyLength { + expected: header.dims.body_len, + actual: actual_len, + }); + } + + body_hasher.update(chunk); + set_bits += chunk + .iter() + .map(|byte| u64::from(byte.count_ones())) + .sum::(); + bits.extend(chunk.iter().copied().map(AtomicU8::new)); + } + + if header_len < SNAPSHOT_HEADER_LEN { + return Err(BloomError::Truncated { + expected_at_least: SNAPSHOT_HEADER_LEN, + actual: total_len, + }); + } + + let header = decoded_header.expect("a complete header was decoded"); + let bits = bits.expect("a complete header allocated its body"); + if bits.len() != header.dims.body_len { + return Err(BloomError::InvalidBodyLength { + expected: header.dims.body_len, + actual: bits.len(), + }); + } + + let actual_checksum: [u8; 32] = body_hasher.finalize().into(); + if actual_checksum != header.body_checksum { + return Err(BloomError::ChecksumMismatch { + expected: hex::encode(header.body_checksum), + actual: hex::encode(actual_checksum), + }); + } + + manifest.validate_snapshot_parts( + header.built_at, + header.source_max_row_id, + actual_checksum, + )?; + + let stats = BloomSnapshotStats::from_set_bits(header.dims, set_bits).validate()?; + + Ok(( + Self { + dims: header.dims, + built_at: header.built_at, + source_max_row_id: header.source_max_row_id, + bits, + decoded_body_checksum: Some(DecodedBodyChecksum::new(actual_checksum)), + }, + stats, + )) + } + + /// Copies the current bits into a serializable snapshot. + /// + /// The copy loads each atomic byte independently. A concurrent multi-bit insert may therefore + /// be only partly represented. Callers that persist a live filter must save a stream + /// continuation before copying and replay after it when restoring the copy. The caller supplies + /// `built_at` as the copy's freshness metadata. + pub fn to_snapshot(&self, built_at: u64) -> Self { + Self { + dims: self.dims, + built_at, + source_max_row_id: self.source_max_row_id, + bits: self + .bits + .iter() + .map(|byte| AtomicU8::new(byte.load(Ordering::Relaxed))) + .collect(), + decoded_body_checksum: None, + } + } +} + +/// A mutable pointer to one immutable snapshot object. +/// +/// Concurrent builds publish distinct immutable snapshot objects and race only on this pointer, so +/// no build lease is needed. The S3 publisher moves the pointer through an ETag compare-and-swap +/// that orders manifests by generation and then by `source_max_row_id`, so a slow build never +/// replaces a newer manifest. +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub struct SnapshotManifest { + pub version: u8, + pub object_key: PathBuf, + pub body_checksum: String, + pub source_generation: Uuid, + + /// The snapshot's ordinary source-sync cursor. Within one generation this is the authoritative + /// measure of build progress, so the publish compare-and-swap orders manifests by it. Bootstrap + /// overlay shards do not advance this value. + pub source_max_row_id: u64, + + /// Unix timestamp used to start an `AT_TIMESTAMP` replay when no saved continuation exists and + /// to prove that the required stream history is still retained. + /// + /// The build takes the later of the pre-deployment bootstrap barrier and the ordinary watermark + /// row's `created_at` rewound by [`CLOCK_SKEW_SLACK`]. With no ordinary watermark row, it uses + /// the barrier. + pub replay_from: u64, +} + +impl SnapshotManifest { + pub fn for_snapshot( + snapshot: &ConcurrentBloomFilter, + source_generation: Uuid, + replay_from: u64, + ) -> Self { + Self::new( + snapshot.built_at(), + snapshot.source_max_row_id(), + snapshot.body_checksum_hex(), + source_generation, + replay_from, + ) + } + + pub fn for_encoded_snapshot( + encoded: &EncodedBloomSnapshot, + source_generation: Uuid, + replay_from: u64, + ) -> Self { + Self::new( + encoded.built_at, + encoded.source_max_row_id, + encoded.body_checksum_hex(), + source_generation, + replay_from, + ) + } + + fn new( + built_at: u64, + source_max_row_id: u64, + body_checksum: String, + source_generation: Uuid, + replay_from: u64, + ) -> Self { + let definition = FilterDefinition::CURRENT; + + Self { + version: definition.version, + object_key: definition.snapshot_object_key( + source_generation, + built_at, + source_max_row_id, + ), + body_checksum, + source_generation, + source_max_row_id, + replay_from, + } + } + + pub fn encode(&self) -> Result, BloomError> { + let mut encoded = serde_json::to_vec(self).map_err(BloomError::ManifestJson)?; + encoded.push(b'\n'); + Ok(encoded) + } + + pub fn decode(encoded: &[u8]) -> Result { + let manifest: Self = serde_json::from_slice(encoded).map_err(BloomError::ManifestJson)?; + manifest.validate_definition()?; + Ok(manifest) + } + + pub fn validate_snapshot(&self, snapshot: &ConcurrentBloomFilter) -> Result<(), BloomError> { + self.validate_snapshot_parts( + snapshot.built_at(), + snapshot.source_max_row_id(), + snapshot.body_checksum(), + ) + } + + fn validate_snapshot_parts( + &self, + built_at: u64, + source_max_row_id: u64, + body_checksum: [u8; 32], + ) -> Result<(), BloomError> { + self.validate_definition()?; + + let expected_object_key = FilterDefinition::CURRENT.snapshot_object_key( + self.source_generation, + built_at, + source_max_row_id, + ); + if self.body_checksum != hex::encode(body_checksum) + || self.object_key != expected_object_key + || self.source_max_row_id != source_max_row_id + { + return Err(BloomError::ManifestSnapshotMismatch); + } + + Ok(()) + } + + /// Validates the version and the shape of the object key. + /// + /// The manifest carries the watermark but not the build time, so the exact key is only + /// recomputable against the downloaded snapshot ([`Self::validate_snapshot`]). This check pins + /// everything else: the key lies in this definition's namespace, names this manifest's + /// generation, carries two decimal components, and ends in the snapshot file name. + fn validate_definition(&self) -> Result<(), BloomError> { + let definition = FilterDefinition::CURRENT; + if self.version != definition.version { + return Err(BloomError::IncompatibleDefinition { + expected: definition.version, + actual: self.version, + }); + } + + let invalid = || BloomError::InvalidObjectKey { + object_key: self.object_key.clone(), + }; + + let relative = self + .object_key + .strip_prefix(definition.object_prefix()) + .map_err(|_| invalid())?; + let mut components = relative.components(); + let mut next = || match components.next() { + Some(Component::Normal(component)) => component.to_str().ok_or_else(invalid), + _ => Err(invalid()), + }; + + let generation: Uuid = next()?.parse().map_err(|_| invalid())?; + if generation != self.source_generation { + return Err(invalid()); + } + + let _start = next()?.parse::().map_err(|_| invalid())?; + let _end = next()?.parse::().map_err(|_| invalid())?; + + let file_name = next()?; + if file_name != format!("snapshot.{FILTER_EXT}") || components.next().is_some() { + return Err(invalid()); + } + + Ok(()) + } +} + +fn checksum(bytes: &[u8]) -> [u8; 32] { + Sha256::digest(bytes).into() +} + +fn decode_store_path_hash(store_path_hash: &StorePathHash) -> Result<[u8; 20], BloomError> { + nix_base32::from_nix_base32(store_path_hash.as_str()) + .and_then(|decoded| decoded.try_into().ok()) + .ok_or_else(|| BloomError::InvalidStorePathHash { + hash: store_path_hash.to_string(), + }) +} + +#[derive(Debug, displaydoc::Display)] +pub enum BloomError { + /// invalid Bloom magic: {actual:02x?} + InvalidMagic { actual: [u8; 4] }, + + /// truncated Bloom snapshot: expected at least {expected_at_least} bytes, got {actual} + Truncated { + expected_at_least: usize, + actual: usize, + }, + + /// incompatible Bloom definition version: expected {expected}, got {actual} + IncompatibleDefinition { expected: u8, actual: u8 }, + + /// Bloom bit count must be a byte-aligned power of two, got {m_bits} + InvalidBitCount { m_bits: u64 }, + + /// Bloom probe count must be between 1 and 64, got {k} + InvalidProbeCount { k: u8 }, + + /// Bloom filter with {m_bits} bits does not fit in memory + TooLarge { m_bits: u64 }, + + /// invalid Bloom body length: expected {expected} bytes, got {actual} + InvalidBodyLength { expected: usize, actual: usize }, + + /// Bloom body checksum mismatch: expected {expected}, got {actual} + ChecksumMismatch { expected: String, actual: String }, + + /// Bloom snapshot estimated FPR {estimated_fpr:.6} exceeds safety limit {limit:.6} + ExcessiveFalsePositiveRate { estimated_fpr: f64, limit: f64 }, + + /// could not Nix-base32 decode store path hash {hash} to 20 bytes + InvalidStorePathHash { hash: String }, + + /// streamed Bloom positions use {event_m_bits} bits, smaller than this filter's {filter_m_bits} + IncompatiblePositions { + event_m_bits: u64, + filter_m_bits: u64, + }, + + /// streamed Bloom event carries {count} probe positions, fewer than this filter's {k} + InsufficientPositions { count: usize, k: u8 }, + + /// streamed Bloom position count must be between 1 and 64, got {count} + InvalidPositionCount { count: usize }, + + /// streamed Bloom position {position} is out of range for {m_bits} bits + PositionOutOfRange { position: u64, m_bits: u64 }, + + /// invalid Bloom manifest JSON: {0} + ManifestJson(serde_json::Error), + + /// invalid Bloom snapshot object key {object_key:?} + InvalidObjectKey { object_key: PathBuf }, + + /// Bloom manifest does not describe the decoded snapshot + ManifestSnapshotMismatch, + + /// invalid Bloom source shard key {key} + InvalidShardKey { key: PathBuf }, + + /// invalid Bloom source shard bounds {first_id}..={last_id} + InvalidShardBounds { first_id: u64, last_id: u64 }, + + /// invalid Bloom source shard JSON: {0} + ShardJson(serde_json::Error), + + /// Bloom source shards must not be empty + EmptyShard, + + /// Bloom source shard has {count} rows, more than the flush size + OversizedShard { count: usize }, + + /// Bloom source shard row IDs are not strictly increasing + NonMonotonicShardRows, + + /// Bloom source shard contents do not match the bounds of key {key} + ShardBoundsMismatch { key: PathBuf }, + + /// Bloom source shard {key} overlaps another shard of its generation + ShardOverlap { key: PathBuf }, + + /// Bloom source shard {key} does not belong to generation {expected} + WrongShardGeneration { key: PathBuf, expected: String }, +} + +impl std::error::Error for BloomError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::ManifestJson(error) | Self::ShardJson(error) => Some(error), + _ => None, + } + } +} diff --git a/src/bloom/source/mod.rs b/src/bloom/source/mod.rs new file mode 100644 index 0000000..fba9d68 --- /dev/null +++ b/src/bloom/source/mod.rs @@ -0,0 +1,352 @@ +//! Parsed types for the immutable source-cache shards in S3. +//! +//! The source sync streams object rows out of the database into immutable JSON shards under +//! `bloom/source/v1/{generation}/{kind}/{first_id}/{last_id}/{anchor_ts}.json`. Ordinary shards +//! checkpoint source-sync progress (i.e. source-sync will continue from the last checkpoint); +//! bootstrap shards are non-checkpointing overlays whose ranges may overlap with ordinary shards. +//! The snapshot build lists those shards and folds them into a Bloom filter. The ordinary sync +//! cursor is derived from the listed shard keys; there is no separate mutable checkpoint object. + +#[cfg(test)] +mod tests; + +use std::path::{Component, Path, PathBuf}; + +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::bloom::BloomError; + +/// Maximum rows in one source shard; the source sync flushes at this size. +pub const SOURCE_SHARD_FLUSH_ROWS: usize = 100_000; + +const SHARD_EXT: &str = "json"; + +/// One object row carried by a source shard. +#[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct SourceRecord { + /// The `object` table row ID. Always positive. + pub id: u64, + /// The store path hash of the row, as stored in the database. The hash is validated only when + /// it is inserted into a filter, so one invalid row does not fail a whole shard. + pub store_path_hash: String, +} + +/// Whether a source shard advances the checkpoint or overlays bootstrap data. +#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +pub enum ShardKind { + Ordinary, + Bootstrap, +} + +impl ShardKind { + fn as_str(self) -> &'static str { + match self { + Self::Ordinary => "ordinary", + Self::Bootstrap => "bootstrap", + } + } + + fn parse(value: &str) -> Option { + match value { + "ordinary" => Some(Self::Ordinary), + "bootstrap" => Some(Self::Bootstrap), + _ => None, + } + } +} + +/// The parsed key of one immutable source shard: +/// `bloom/source/v1/{generation}/{kind}/{first_id}/{last_id}/{anchor_ts}.json`. +/// +/// Invariants, enforced at construction and at parse: +/// +/// - `first_id` is at least 1, +/// - `last_id` is at least `first_id`, +/// - the string form round-trips exactly (no leading zeros, no aliases). +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ShardKey { + generation: Uuid, + kind: ShardKind, + first_id: u64, + last_id: u64, + anchor_ts: u64, +} + +impl ShardKey { + pub fn new( + generation: Uuid, + kind: ShardKind, + first_id: u64, + last_id: u64, + anchor_ts: u64, + ) -> Result { + if first_id == 0 || last_id < first_id { + return Err(BloomError::InvalidShardBounds { first_id, last_id }); + } + + Ok(Self { + generation, + kind, + first_id, + last_id, + anchor_ts, + }) + } + + /// Root S3 prefix of the source cache, spanning every generation. + pub fn source_prefix() -> PathBuf { + Path::new("bloom").join("source").join("v1") + } + + /// Key of the ready marker of a generation. + /// + /// The source sync writes the marker after the generation's ordinary cursor and shards are durable. + /// A snapshot build reads only marked generations, so an interrupted total rebuild stays invisible + /// to builds until its next sync completes. + pub fn ready_marker_key(generation: Uuid) -> PathBuf { + Self::source_prefix() + .join(generation.to_string()) + .join("ready") + } + + pub fn generation(&self) -> Uuid { + self.generation + } + + pub fn kind(&self) -> ShardKind { + self.kind + } + + /// The smallest row ID in the shard. + pub fn first_id(&self) -> u64 { + self.first_id + } + + /// The largest row ID in the shard: the watermark contribution of this shard. + pub fn last_id(&self) -> u64 { + self.last_id + } + + /// The `created_at` of the watermark row (the last row of the shard), in Unix seconds. The + /// snapshot build derives its ordinary replay candidate from the newest shard's anchor. + pub fn anchor_ts(&self) -> u64 { + self.anchor_ts + } + + /// The canonical S3 key of this shard. + pub fn to_key(&self) -> PathBuf { + Self::source_prefix() + .join(self.generation().to_string()) + .join(self.kind().as_str()) + .join(self.first_id().to_string()) + .join(self.last_id().to_string()) + .join(self.anchor_ts().to_string()) + .with_extension(SHARD_EXT) + } + + /// Parses one S3 key. A key that does not round-trip through [`Self::to_key`] fails, so every + /// parsed key has exactly one string form. + pub fn parse(key: &Path) -> Result { + let invalid = || BloomError::InvalidShardKey { + key: key.to_owned(), + }; + + let relative = key + .strip_prefix(Self::source_prefix()) + .map_err(|_| invalid())?; + let mut components = relative.components(); + let mut next = || match components.next() { + Some(Component::Normal(component)) => component.to_str().ok_or_else(invalid), + _ => Err(invalid()), + }; + + let generation: Uuid = next()?.parse().map_err(|_| invalid())?; + let kind = ShardKind::parse(next()?).ok_or_else(invalid)?; + let first_id: u64 = next()?.parse().map_err(|_| invalid())?; + let last_id: u64 = next()?.parse().map_err(|_| invalid())?; + let file_name = next()?; + if components.next().is_some() { + return Err(invalid()); + } + + let anchor_ts: u64 = file_name + .strip_suffix(&format!(".{SHARD_EXT}")) + .ok_or_else(invalid)? + .parse() + .map_err(|_| invalid())?; + + let parsed = + Self::new(generation, kind, first_id, last_id, anchor_ts).map_err(|_| invalid())?; + if parsed.to_key() != key { + return Err(invalid()); + } + + Ok(parsed) + } +} + +/// The decoded contents of one source shard, validated against its key. +/// +/// Invariants: at least one and at most [`SOURCE_SHARD_FLUSH_ROWS`] records, strictly increasing +/// row IDs, and first/last IDs that equal the key's bounds. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SourceShard { + key: ShardKey, + records: Vec, +} + +impl SourceShard { + /// Builds a shard from freshly scanned records. The key bounds come from the records + /// themselves; `anchor_ts` is the `created_at` of the last record. + pub fn new( + generation: Uuid, + kind: ShardKind, + anchor_ts: u64, + records: Vec, + ) -> Result { + let (first_id, last_id) = match (records.first(), records.last()) { + (Some(first), Some(last)) => (first.id, last.id), + _ => return Err(BloomError::EmptyShard), + }; + + let key = ShardKey::new(generation, kind, first_id, last_id, anchor_ts)?; + validate_records(&records, &key)?; + + Ok(Self { key, records }) + } + + /// Decodes shard contents and validates them against the parsed key. + pub fn decode(key: ShardKey, encoded: &[u8]) -> Result { + let records: Vec = + serde_json::from_slice(encoded).map_err(BloomError::ShardJson)?; + validate_records(&records, &key)?; + + Ok(Self { key, records }) + } + + pub fn encode(&self) -> Result, BloomError> { + let mut encoded = serde_json::to_vec(&self.records).map_err(BloomError::ShardJson)?; + encoded.push(b'\n'); + Ok(encoded) + } + + pub fn key(&self) -> &ShardKey { + &self.key + } + + pub fn kind(&self) -> ShardKind { + self.key.kind() + } + + pub fn records(&self) -> &[SourceRecord] { + &self.records + } +} + +fn validate_records(records: &[SourceRecord], key: &ShardKey) -> Result<(), BloomError> { + if records.is_empty() { + return Err(BloomError::EmptyShard); + } + + if records.len() > SOURCE_SHARD_FLUSH_ROWS { + return Err(BloomError::OversizedShard { + count: records.len(), + }); + } + + if records.windows(2).any(|pair| pair[0].id >= pair[1].id) { + return Err(BloomError::NonMonotonicShardRows); + } + + let first = records.first().map(|record| record.id); + let last = records.last().map(|record| record.id); + if first != Some(key.first_id()) || last != Some(key.last_id()) { + return Err(BloomError::ShardBoundsMismatch { key: key.to_key() }); + } + + Ok(()) +} + +/// All shards of one generation, as listed from S3. +/// +/// Construction rejects overlapping ordinary row ranges. Bootstrap shards are overlays and may +/// overlap any shard; duplicate Bloom insertions are harmless, and retries may change bootstrap +/// chunk boundaries. Gaps between ordinary shards are legal because row IDs have natural gaps. +#[derive(Clone, Debug)] +pub struct ShardIndex { + generation: Uuid, + /// Deterministically sorted by range, kind, and anchor. Ordinary ranges are pairwise disjoint. + shards: Vec, +} + +impl ShardIndex { + pub fn new(generation: Uuid, mut shards: Vec) -> Result { + if let Some(foreign) = shards.iter().find(|shard| shard.generation() != generation) { + return Err(BloomError::WrongShardGeneration { + key: foreign.to_key(), + expected: generation.to_string(), + }); + } + + shards.sort_unstable_by_key(|shard| { + ( + shard.first_id(), + shard.last_id(), + shard.kind(), + shard.anchor_ts(), + ) + }); + + if let Some(overlapping) = shards + .iter() + .filter(|shard| shard.kind() == ShardKind::Ordinary) + .collect::>() + .windows(2) + .find(|pair| pair[1].first_id() <= pair[0].last_id()) + { + return Err(BloomError::ShardOverlap { + key: overlapping[1].to_key(), + }); + } + + Ok(Self { generation, shards }) + } + + pub fn generation(&self) -> Uuid { + self.generation + } + + /// The largest row ID any ordinary shard covers, or 0 when there are no ordinary shards. + /// + /// The next source sync scans rows above this value, and the next snapshot records it. + pub fn watermark(&self) -> u64 { + self.shards + .iter() + .filter(|shard| shard.kind() == ShardKind::Ordinary) + .map(ShardKey::last_id) + .max() + .unwrap_or(0) + } + + /// The anchor timestamp of the ordinary shard that holds the watermark row, or `None` when + /// there are no ordinary shards. The snapshot build combines this value with the bootstrap + /// barrier to derive `replay_from`. + pub fn anchor_ts(&self) -> Option { + self.shards + .iter() + .filter(|shard| shard.kind() == ShardKind::Ordinary) + .max_by_key(|shard| shard.last_id()) + .map(ShardKey::anchor_ts) + } + + /// The shards in ascending row-ID order. + pub fn shards(&self) -> &[ShardKey] { + &self.shards + } + + pub fn is_empty(&self) -> bool { + self.shards.is_empty() + } +} diff --git a/src/bloom/source/tests.rs b/src/bloom/source/tests.rs new file mode 100644 index 0000000..6c4184e --- /dev/null +++ b/src/bloom/source/tests.rs @@ -0,0 +1,184 @@ +use super::*; + +const GENERATION: Uuid = uuid::uuid!("019bf5a7-f6e8-7ac0-b973-8536596bdb45"); + +fn record(id: u64) -> SourceRecord { + SourceRecord { + id, + store_path_hash: "ia70ss13m22znbl8khrf2hq72qmh5drr".to_owned(), + } +} + +#[test] +fn shard_keys_round_trip() { + let key = ShardKey::new(GENERATION, ShardKind::Bootstrap, 12, 345, 1_700_000_000).unwrap(); + assert_eq!( + key.to_key(), + PathBuf::from( + "bloom/source/v1/019bf5a7-f6e8-7ac0-b973-8536596bdb45/bootstrap/12/345/1700000000.json" + ) + ); + assert_eq!(key.kind(), ShardKind::Bootstrap); + assert_eq!(ShardKey::parse(&key.to_key()).unwrap(), key); +} + +/// The ready marker lives under the generation prefix but is not a shard: it must never parse as +/// one, so a listing that forgets to skip it fails loudly instead of misreading it. +#[test] +fn ready_markers_are_not_shard_keys() { + let key = ShardKey::ready_marker_key(GENERATION); + assert_eq!( + key, + PathBuf::from("bloom/source/v1/019bf5a7-f6e8-7ac0-b973-8536596bdb45/ready") + ); + assert!(ShardKey::parse(&key).is_err()); +} + +#[test] +fn malformed_shard_keys_fail_to_parse() { + for key in [ + // A key without a shard kind is incomplete. + "bloom/source/v1/019bf5a7-f6e8-7ac0-b973-8536596bdb45/12/345.json", + "bloom/source/v1/019bf5a7-f6e8-7ac0-b973-8536596bdb45/12/345/1700000000.json", + "bloom/source/v1/019bf5a7-f6e8-7ac0-b973-8536596bdb45/other/12/345/1700000000.json", + "bloom/source/v1/019bf5a7-f6e8-7ac0-b973-8536596bdb45/ordinary/12/345/1700000000", + "bloom/source/v1/019bf5a7-f6e8-7ac0-b973-8536596bdb45/ordinary/12/345/1700000000.json/extra", + "bloom/source/v1/not-a-uuid/ordinary/12/345/1700000000.json", + "bloom/source/v2/019bf5a7-f6e8-7ac0-b973-8536596bdb45/ordinary/12/345/1700000000.json", + // Bounds violations. + "bloom/source/v1/019bf5a7-f6e8-7ac0-b973-8536596bdb45/ordinary/0/345/1700000000.json", + "bloom/source/v1/019bf5a7-f6e8-7ac0-b973-8536596bdb45/ordinary/12/11/1700000000.json", + // Non-canonical numerals must not alias a canonical key. + "bloom/source/v1/019bf5a7-f6e8-7ac0-b973-8536596bdb45/ordinary/012/345/1700000000.json", + ] { + assert!(ShardKey::parse(Path::new(key)).is_err(), "{key}"); + } +} + +#[test] +fn shard_contents_validate_against_the_key() { + let shard = SourceShard::new( + GENERATION, + ShardKind::Bootstrap, + 1_700_000_000, + vec![record(1), record(4), record(9)], + ) + .unwrap(); + assert_eq!(shard.key().first_id(), 1); + assert_eq!(shard.key().last_id(), 9); + assert_eq!(shard.kind(), ShardKind::Bootstrap); + + let decoded = SourceShard::decode(shard.key().clone(), &shard.encode().unwrap()).unwrap(); + assert_eq!(decoded, shard); + + // Contents whose bounds disagree with the key are rejected. + let wrong_key = ShardKey::new(GENERATION, ShardKind::Bootstrap, 1, 10, 1_700_000_000).unwrap(); + assert!(matches!( + SourceShard::decode(wrong_key, &shard.encode().unwrap()), + Err(BloomError::ShardBoundsMismatch { .. }) + )); +} + +#[test] +fn shards_reject_bad_orders_and_sizes() { + assert!(matches!( + SourceShard::new(GENERATION, ShardKind::Ordinary, 0, vec![]), + Err(BloomError::EmptyShard) + )); + assert!(matches!( + SourceShard::new( + GENERATION, + ShardKind::Ordinary, + 0, + vec![record(2), record(2)] + ), + Err(BloomError::NonMonotonicShardRows) + )); + assert!(matches!( + SourceShard::new( + GENERATION, + ShardKind::Ordinary, + 0, + (1..=(SOURCE_SHARD_FLUSH_ROWS as u64 + 1)) + .map(record) + .collect() + ), + Err(BloomError::OversizedShard { .. }) + )); +} + +#[test] +fn the_index_rejects_ordinary_overlap_and_exposes_the_watermark() { + let shard = + |kind, first, last, anchor| ShardKey::new(GENERATION, kind, first, last, anchor).unwrap(); + + let index = ShardIndex::new( + GENERATION, + vec![ + shard(ShardKind::Ordinary, 200, 250, 40), + shard(ShardKind::Ordinary, 1, 100, 10), + shard(ShardKind::Ordinary, 101, 199, 20), + shard(ShardKind::Bootstrap, 50, 300, 99), + shard(ShardKind::Bootstrap, 50, 150, 98), + ], + ) + .unwrap(); + assert_eq!(index.watermark(), 250); + assert_eq!(index.anchor_ts(), Some(40)); + assert_eq!(index.shards().len(), 5); + assert_eq!(index.shards()[1].kind(), ShardKind::Bootstrap); + assert_eq!(index.shards()[2].kind(), ShardKind::Bootstrap); + assert_eq!( + index + .shards() + .iter() + .map(ShardKey::first_id) + .collect::>(), + vec![1, 50, 50, 101, 200] + ); + + assert!(matches!( + ShardIndex::new( + GENERATION, + vec![ + shard(ShardKind::Ordinary, 1, 100, 10), + shard(ShardKind::Ordinary, 100, 150, 20) + ] + ), + Err(BloomError::ShardOverlap { .. }) + )); + + let other_generation: Uuid = uuid::uuid!("019bf5a7-f6e8-7ac0-b973-8536596bdb46"); + assert!(matches!( + ShardIndex::new( + other_generation, + vec![shard(ShardKind::Bootstrap, 1, 100, 10)] + ), + Err(BloomError::WrongShardGeneration { .. }) + )); +} + +#[test] +fn an_empty_index_has_watermark_zero_and_no_anchor() { + let index = ShardIndex::new(GENERATION, Vec::new()).unwrap(); + assert!(index.is_empty()); + assert_eq!(index.watermark(), 0); + assert_eq!(index.anchor_ts(), None); +} + +#[test] +fn an_overlay_only_index_is_not_empty_but_has_no_checkpoint() { + let index = ShardIndex::new( + GENERATION, + vec![ + ShardKey::new(GENERATION, ShardKind::Bootstrap, 1, 100, 10).unwrap(), + ShardKey::new(GENERATION, ShardKind::Bootstrap, 50, 150, 20).unwrap(), + ], + ) + .unwrap(); + + assert!(!index.is_empty()); + assert_eq!(index.shards().len(), 2); + assert_eq!(index.watermark(), 0); + assert_eq!(index.anchor_ts(), None); +} diff --git a/src/bloom/tests.rs b/src/bloom/tests.rs new file mode 100644 index 0000000..f09edca --- /dev/null +++ b/src/bloom/tests.rs @@ -0,0 +1,330 @@ +use std::sync::LazyLock; + +use super::*; + +static HASH_A: LazyLock = + LazyLock::new(|| StorePathHash::new("ia70ss13m22znbl8khrf2hq72qmh5drr").unwrap()); +static HASH_B: LazyLock = + LazyLock::new(|| StorePathHash::new("00000000000000000000000000000000").unwrap()); + +#[test] +fn inserted_items_are_never_negative_and_duplicates_are_free() { + let filter = ConcurrentBloomFilter::new(1024, 10, 123).unwrap(); + let positions = wire::ProbePositions::of(&HASH_A).unwrap(); + filter.insert_positions(&positions).unwrap(); + let after_first = filter.copied_bits(); + filter.insert_positions(&positions).unwrap(); + + assert!(filter.contains(&HASH_A).unwrap()); + assert!(!filter.contains(&HASH_B).unwrap()); + assert_eq!(filter.copied_bits(), after_first); +} + +/// The concurrent filter accepts inserts through a shared reference. +#[test] +fn concurrent_filter_probes_and_inserts() { + let filter = ConcurrentBloomFilter::new(1024, 10, 123).unwrap(); + let positions = wire::ProbePositions::of(&HASH_A).unwrap(); + filter.insert_positions(&positions).unwrap(); + + assert!(filter.contains(&HASH_A).unwrap()); + assert!(!filter.contains(&HASH_B).unwrap()); + assert_eq!(filter.built_at(), 123); + + let positions = wire::ProbePositions::of(&HASH_B).unwrap(); + filter.insert_positions(&positions).unwrap(); + assert!(filter.contains(&HASH_B).unwrap()); +} + +#[test] +fn insertion_invalidates_a_decoded_snapshot_checksum() { + let source = ConcurrentBloomFilter::new(1024, 10, 123).unwrap(); + let decoded = ConcurrentBloomFilter::decode(&source.encode()).unwrap(); + let empty_checksum = decoded.body_checksum(); + + let positions = wire::ProbePositions::of(&HASH_A).unwrap(); + decoded.insert_positions(&positions).unwrap(); + + assert_ne!(decoded.body_checksum(), empty_checksum); +} + +/// Streamed positions from an equal or larger publisher filter mask down to exactly the positions +/// this filter calculates itself, so a probe of the same hash turns positive. +#[test] +fn streamed_positions_mask_down_to_local_probes() { + let filter = ConcurrentBloomFilter::new(1024, 10, 123).unwrap(); + + // The publisher runs the current sizing policy, which is larger than this filter. + let positions = wire::ProbePositions::of(&HASH_A).unwrap(); + + filter.insert_positions(&positions).unwrap(); + assert!(filter.contains(&HASH_A).unwrap()); + assert!(!filter.contains(&HASH_B).unwrap()); +} + +/// Positions from a publisher filter smaller than this one have discarded high bits, so the record +/// is rejected and no bit changes. +#[test] +fn positions_from_a_smaller_publisher_are_rejected() { + let filter = ConcurrentBloomFilter::new(1024, 10, 123).unwrap(); + let empty_bits = filter.copied_bits(); + + // Valid in itself, but calculated against a smaller filter. + let positions: wire::ProbePositions = + serde_json::from_value(serde_json::json!({ "m_bits": 8, "positions": [3] })).unwrap(); + + assert!(matches!( + filter.insert_positions(&positions), + Err(BloomError::IncompatiblePositions { + event_m_bits: 8, + filter_m_bits: 1024, + }) + )); + assert_eq!(filter.copied_bits(), empty_bits); +} + +#[test] +fn positions_with_fewer_probes_than_the_filter_are_rejected() { + let filter = ConcurrentBloomFilter::new(1024, 10, 123).unwrap(); + let empty_bits = filter.copied_bits(); + + // Valid in itself, but calculated with a smaller probe count: applying it would leave probe + // bits 10 and higher unset, leading to a false negative. + let positions: wire::ProbePositions = serde_json::from_value(serde_json::json!({ + "m_bits": M_BITS, + "positions": [1, 2, 3, 4], + })) + .unwrap(); + + assert!(matches!( + filter.insert_positions(&positions), + Err(BloomError::InsufficientPositions { count: 4, k: 10 }) + )); + assert_eq!(filter.copied_bits(), empty_bits); +} + +/// A concurrent filter round-trips through a snapshot: the copy keeps the source bits and the live +/// inserts, carries the caller's freshness time, and encodes into decodable bytes. +#[test] +fn concurrent_filter_round_trips_through_a_snapshot() { + let mut filter = ConcurrentBloomFilter::new(1024, 10, 123).unwrap(); + let positions = wire::ProbePositions::of(&HASH_A).unwrap(); + filter.insert_positions(&positions).unwrap(); + filter.set_source_max_row_id(456); + + let positions = wire::ProbePositions::of(&HASH_B).unwrap(); + filter.insert_positions(&positions).unwrap(); + + let copy = filter.to_snapshot(789); + assert_eq!(copy.built_at(), 789); + assert_eq!(copy.source_max_row_id(), 456); + + let decoded = ConcurrentBloomFilter::decode(©.encode()).unwrap(); + assert!(decoded.contains(&HASH_A).unwrap()); + assert!(decoded.contains(&HASH_B).unwrap()); + + let mut streamed = Vec::new(); + copy.write_to(&mut streamed).unwrap(); + assert_eq!(streamed, copy.encode()); +} + +#[test] +fn full_filter_has_no_distinct_item_estimate() { + let snapshot = ConcurrentBloomFilter::new(8, 1, 123).unwrap(); + for byte in &snapshot.bits { + byte.store(u8::MAX, Ordering::Relaxed); + } + + assert_eq!(snapshot.estimated_distinct_items(), None); +} + +#[test] +fn snapshot_and_manifest_round_trip() { + let mut snapshot = ConcurrentBloomFilter::new(1024, 10, 123).unwrap(); + let positions = wire::ProbePositions::of(&HASH_A).unwrap(); + snapshot.insert_positions(&positions).unwrap(); + snapshot.set_source_max_row_id(456); + + let encoded = snapshot.encode_with_checksum(); + assert!(encoded.as_bytes().starts_with(&SNAPSHOT_MAGIC)); + assert_eq!(encoded.as_bytes(), snapshot.encode()); + + let decoded = ConcurrentBloomFilter::decode(encoded.as_bytes()).unwrap(); + // The encoded value owns the metadata from its header, so later changes to the source filter + // cannot produce a manifest that disagrees with the encoded bytes. + snapshot.set_source_max_row_id(789); + let generation: Uuid = "019bf5a7-f6e8-7ac0-b973-8536596bdb45".parse().unwrap(); + let manifest = SnapshotManifest::for_encoded_snapshot(&encoded, generation, 1_700_000_000); + let decoded_manifest = SnapshotManifest::decode(&manifest.encode().unwrap()).unwrap(); + decoded_manifest.validate_snapshot(&decoded).unwrap(); + let (concurrent, stats) = ConcurrentBloomFilter::decode_with_manifest( + &decoded_manifest, + encoded.as_bytes().chunks(7), + ) + .unwrap(); + + assert!(decoded.contains(&HASH_A).unwrap()); + assert!(concurrent.contains(&HASH_A).unwrap()); + assert_eq!(stats, snapshot.stats()); + assert_eq!(decoded.built_at(), 123); + assert_eq!(concurrent.built_at(), 123); + assert_eq!(decoded.source_max_row_id(), 456); + assert_eq!(concurrent.source_max_row_id(), 456); + assert_eq!( + decoded_manifest.object_key, + FilterDefinition::CURRENT + .object_prefix() + .join("019bf5a7-f6e8-7ac0-b973-8536596bdb45") + .join("123") + .join("456") + .join("snapshot.abf") + ); + assert_eq!(decoded_manifest.source_generation, generation); + assert_eq!(decoded_manifest.source_max_row_id, 456); + assert_eq!(decoded_manifest.replay_from, 1_700_000_000); + + // A manifest without a replay timestamp is rejected: every build derives one from the source + // cache. + let mut old_manifest: serde_json::Value = + serde_json::from_slice(&manifest.encode().unwrap()).unwrap(); + old_manifest + .as_object_mut() + .unwrap() + .remove("replay_from") + .unwrap(); + SnapshotManifest::decode(&serde_json::to_vec(&old_manifest).unwrap()).unwrap_err(); +} + +#[test] +fn corruption_is_rejected() { + let snapshot = ConcurrentBloomFilter::new(1024, 10, 123).unwrap(); + let generation: Uuid = "019bf5a7-f6e8-7ac0-b973-8536596bdb45".parse().unwrap(); + let manifest = SnapshotManifest::for_snapshot(&snapshot, generation, 1_700_000_000); + let mut encoded = snapshot.encode(); + encoded[SNAPSHOT_HEADER_LEN] ^= 1; + + assert!(matches!( + ConcurrentBloomFilter::decode(&encoded), + Err(BloomError::ChecksumMismatch { .. }) + )); + assert!(matches!( + ConcurrentBloomFilter::decode_with_manifest(&manifest, std::iter::once(encoded.as_slice())), + Err(BloomError::ChecksumMismatch { .. }) + )); +} + +#[test] +fn manifest_aware_decode_rejects_wrong_manifest_and_lengths() { + let mut snapshot = ConcurrentBloomFilter::new(1024, 10, 123).unwrap(); + snapshot.set_source_max_row_id(456); + let generation: Uuid = "019bf5a7-f6e8-7ac0-b973-8536596bdb45".parse().unwrap(); + let encoded = snapshot.encode_with_checksum(); + let manifest = SnapshotManifest::for_encoded_snapshot(&encoded, generation, 1_700_000_000); + + let mut wrong_manifest = manifest.clone(); + wrong_manifest.body_checksum = "0".repeat(64); + assert!(matches!( + ConcurrentBloomFilter::decode_with_manifest( + &wrong_manifest, + std::iter::once(encoded.as_bytes()) + ), + Err(BloomError::ManifestSnapshotMismatch) + )); + + assert!(matches!( + ConcurrentBloomFilter::decode_with_manifest( + &manifest, + std::iter::once(&encoded.as_bytes()[..SNAPSHOT_HEADER_LEN - 1]) + ), + Err(BloomError::Truncated { .. }) + )); + assert!(matches!( + ConcurrentBloomFilter::decode_with_manifest( + &manifest, + std::iter::once(&encoded.as_bytes()[..encoded.as_bytes().len() - 1]) + ), + Err(BloomError::InvalidBodyLength { .. }) + )); + + let mut overlong = encoded.as_bytes().to_vec(); + overlong.push(0); + assert!(matches!( + ConcurrentBloomFilter::decode_with_manifest(&manifest, overlong.chunks(13)), + Err(BloomError::InvalidBodyLength { .. }) + )); + + let mut oversized = encoded.as_bytes().to_vec(); + oversized[5..13].copy_from_slice(&(1_u64 << 63).to_le_bytes()); + assert!(matches!( + ConcurrentBloomFilter::decode_with_manifest(&manifest, std::iter::once(oversized)), + Err(BloomError::TooLarge { .. }) + )); +} + +#[test] +fn excessive_false_positive_rate_is_rejected() { + let snapshot = ConcurrentBloomFilter::new(1024, 10, 123).unwrap(); + for byte in &snapshot.bits { + byte.store(u8::MAX, Ordering::Relaxed); + } + let generation: Uuid = "019bf5a7-f6e8-7ac0-b973-8536596bdb45".parse().unwrap(); + let manifest = SnapshotManifest::for_snapshot(&snapshot, generation, 1_700_000_000); + let encoded = snapshot.encode(); + + assert!(matches!( + ConcurrentBloomFilter::decode(&encoded), + Err(BloomError::ExcessiveFalsePositiveRate { .. }) + )); + assert!(matches!( + ConcurrentBloomFilter::decode_with_manifest(&manifest, std::iter::once(encoded.as_slice())), + Err(BloomError::ExcessiveFalsePositiveRate { .. }) + )); +} + +#[test] +fn incompatible_definition_is_rejected() { + let snapshot = ConcurrentBloomFilter::new(1024, 10, 123).unwrap(); + let mut encoded = snapshot.encode(); + // Byte 4 holds the version, directly after the 4-byte magic. + encoded[4] = FilterDefinition::CURRENT.version + 1; + + assert!(matches!( + ConcurrentBloomFilter::decode(&encoded), + Err(BloomError::IncompatibleDefinition { .. }) + )); +} + +// XXX: These golden values pin the format contract. If a code change makes you update this test, +// you changed the format: bump `FilterDefinition::CURRENT.version` in the same change, which moves +// new snapshots to a new manifest namespace. +#[test] +fn definition_and_probe_behavior_are_golden() { + assert_eq!( + FilterDefinition::CURRENT.object_prefix(), + Path::new("bloom/v1") + ); + assert_eq!( + decode_store_path_hash(&HASH_A).unwrap(), + [ + 57, 183, 2, 43, 22, 7, 67, 225, 50, 156, 136, 46, 251, 133, 168, 35, 104, 13, 142, 138, + ] + ); + assert_eq!( + Dimensions::new(1024, 10) + .unwrap() + .probe_positions(&HASH_A) + .unwrap() + .collect::>(), + vec![825, 876, 927, 978, 5, 56, 107, 158, 209, 260] + ); +} + +#[test] +fn compatibility_change_uses_a_new_namespace() { + let current = FilterDefinition::CURRENT; + let changed = FilterDefinition { + version: current.version + 1, + }; + + assert_ne!(changed.object_prefix(), current.object_prefix()); +} diff --git a/src/bloom/wire/mod.rs b/src/bloom/wire/mod.rs new file mode 100644 index 0000000..63bc780 --- /dev/null +++ b/src/bloom/wire/mod.rs @@ -0,0 +1,162 @@ +//! Wire types for Bloom filter live updates. +//! +//! The origin publishes one [`BloomStreamEvent`] per created object row, and each scheduled source +//! sync publishes one empty heartbeat to the stream shard. Each PoP consumes the stream directly +//! through its own enhanced fan-out consumer. The publisher and the consumers share these types, so +//! the wire format has one definition. + +#[cfg(test)] +mod tests; + +use serde::{Deserialize, Serialize}; + +use crate::bloom::{BloomError, Dimensions, M_BITS, MAX_PROBE_COUNT}; +use crate::nix_store::StorePathHash; + +/// Validated probe positions of one store path, as carried by a path-created record. +/// +/// Every value of this type upholds these invariants, whether it was calculated by +/// [`ProbePositions::of`] or deserialized from a record: +/// +/// - `m_bits` is a power of two and at least 8, +/// - there is at least one and at most [`MAX_PROBE_COUNT`] positions, +/// - every position is below `m_bits`. +#[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)] +#[serde(try_from = "UncheckedPositions")] +pub struct ProbePositions { + m_bits: u64, + positions: Vec, +} + +/// The deserialization shape of [`ProbePositions`], before validation. +#[derive(Deserialize)] +struct UncheckedPositions { + m_bits: u64, + positions: Vec, +} + +impl TryFrom for ProbePositions { + type Error = BloomError; + + fn try_from(unchecked: UncheckedPositions) -> Result { + let UncheckedPositions { m_bits, positions } = unchecked; + + if m_bits < 8 || !m_bits.is_power_of_two() { + return Err(BloomError::InvalidBitCount { m_bits }); + } + if positions.is_empty() || positions.len() > usize::from(MAX_PROBE_COUNT) { + return Err(BloomError::InvalidPositionCount { + count: positions.len(), + }); + } + if let Some(&position) = positions.iter().find(|&&position| position >= m_bits) { + return Err(BloomError::PositionOutOfRange { position, m_bits }); + } + + Ok(Self { m_bits, positions }) + } +} + +impl ProbePositions { + /// Calculates the probe positions of a store path hash under the current sizing policy, + /// [`M_BITS`] and [`PROBE_COUNT`](crate::bloom::PROBE_COUNT). + /// + /// The result upholds the type's invariants by construction: the policy is validated at compile + /// time, and each calculated position is masked below [`M_BITS`]. The only failure is a store + /// path hash that does not Nix-base32 decode to 20 bytes. + pub fn of(store_path_hash: &StorePathHash) -> Result { + Ok(Self { + m_bits: M_BITS, + positions: Dimensions::CURRENT_POLICY + .probe_positions(store_path_hash)? + .collect(), + }) + } + + /// The filter size the positions were calculated against. + pub fn m_bits(&self) -> u64 { + self.m_bits + } + + /// The bit positions to set, each below [`Self::m_bits`]. + pub fn positions(&self) -> &[u64] { + &self.positions + } +} + +/// One record on the Bloom update stream. +/// +/// The encoding is JSON with two top-level discriminators: `version` and `kind`. +/// [`BloomStreamEvent::decode`] reads the version before the rest of the record, so a consumer can +/// recognize and skip a record of an unsupported version without understanding its fields. +#[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum BloomStreamEvent { + /// An object row was created for a store path whose probe positions are these. + /// + /// The record carries the bit positions to set rather than the raw store path hash. The origin + /// calculates them with double hashing modulo `m_bits`, the size of its filter ([`M_BITS`] at + /// publish time). Construct this variant with + /// [`BloomStreamEvent::path_created`]. + PathCreated(ProbePositions), + + /// A source sync run completed. + /// + /// The event has no payload. Consumers use its Kinesis arrival time to detect a stopped + /// publisher after catching up and to save a conservative retention anchor with a local filter. + /// The sequence continuation is the actual resume position; replay-horizon proof comes from + /// probing the shard's oldest surviving record against the retention anchor. + Heartbeat, +} + +/// The result of decoding one well-formed record. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum DecodedEvent { + Event(BloomStreamEvent), + + /// The record carries a version this consumer does not support. The fields of the record stay + /// uninterpreted, so a future version can change them freely. + UnsupportedVersion(u8), +} + +#[derive(Serialize)] +struct VersionedEvent<'a> { + version: u8, + #[serde(flatten)] + event: &'a BloomStreamEvent, +} + +#[derive(Deserialize)] +struct VersionOnly { + version: u8, +} + +impl BloomStreamEvent { + /// The version this code encodes, and the only version it decodes. + pub const CURRENT_VERSION: u8 = 1; + + /// Builds a path-created event from a store path hash. + /// + /// The positions come from the current sizing policy, [`M_BITS`] and + /// [`PROBE_COUNT`](crate::bloom::PROBE_COUNT). The hash itself stays out of the event. + pub fn path_created(store_path_hash: &StorePathHash) -> Result { + Ok(Self::PathCreated(ProbePositions::of(store_path_hash)?)) + } + + pub fn encode(&self) -> Result, serde_json::Error> { + serde_json::to_vec(&VersionedEvent { + version: Self::CURRENT_VERSION, + event: self, + }) + } + + /// Decodes one record, reading the version before the event itself. + pub fn decode(encoded: &[u8]) -> Result { + let VersionOnly { version } = serde_json::from_slice(encoded)?; + if version != Self::CURRENT_VERSION { + return Ok(DecodedEvent::UnsupportedVersion(version)); + } + + Ok(DecodedEvent::Event(serde_json::from_slice(encoded)?)) + } +} diff --git a/src/bloom/wire/tests.rs b/src/bloom/wire/tests.rs new file mode 100644 index 0000000..5ea2893 --- /dev/null +++ b/src/bloom/wire/tests.rs @@ -0,0 +1,108 @@ +use std::sync::LazyLock; + +use super::*; +use crate::bloom::PROBE_COUNT; + +static HASH: LazyLock = + LazyLock::new(|| StorePathHash::new("s66mzxpvicwk07gjbjfw9izjfa797vsw").unwrap()); + +#[test] +fn events_round_trip() { + for event in [ + BloomStreamEvent::path_created(&HASH).unwrap(), + BloomStreamEvent::Heartbeat, + ] { + let encoded = event.encode().unwrap(); + assert_eq!( + BloomStreamEvent::decode(&encoded).unwrap(), + DecodedEvent::Event(event) + ); + } +} + +#[test] +fn encoding_carries_version_and_kind() { + let event = BloomStreamEvent::path_created(&HASH).unwrap(); + let encoded: serde_json::Value = serde_json::from_slice(&event.encode().unwrap()).unwrap(); + + assert_eq!(encoded["version"], BloomStreamEvent::CURRENT_VERSION); + assert_eq!(encoded["kind"], "path_created"); +} + +/// A path-created event never carries the store path hash, in any encoding of it. It carries the +/// probe positions of the current sizing policy instead. +#[test] +fn path_created_carries_positions_and_not_the_hash() { + let positions = ProbePositions::of(&HASH).unwrap(); + assert_eq!(positions.m_bits(), M_BITS); + assert_eq!(positions.positions().len(), usize::from(PROBE_COUNT)); + + let event = BloomStreamEvent::PathCreated(positions); + let encoded = String::from_utf8(event.encode().unwrap()).unwrap(); + + // The record carries exactly these fields; no hash field rides along. + let value: serde_json::Value = serde_json::from_str(&encoded).unwrap(); + let object = value.as_object().unwrap(); + assert_eq!(object.len(), 4, "{encoded}"); + for key in ["version", "kind", "m_bits", "positions"] { + assert!(object.contains_key(key), "{key} is missing: {encoded}"); + } + + // The raw record text never contains the hash, in any position. + assert!(!encoded.contains(HASH.as_str()), "{encoded}"); +} + +/// A record whose positions violate an invariant fails to decode, so no consumer ever holds an +/// invalid [`ProbePositions`] value. +#[test] +fn invalid_positions_fail_to_decode() { + for (reason, m_bits, positions) in [ + ("not a power of two", 1000, vec![3]), + ("no positions at all", 1024, vec![]), + ("a position beyond the claimed size", 1024, vec![1024]), + ( + "more positions than any publisher sets", + 1024, + (0..65).collect(), + ), + ] { + let record = serde_json::json!({ + "version": BloomStreamEvent::CURRENT_VERSION, + "kind": "path_created", + "m_bits": m_bits, + "positions": positions, + }); + let encoded = serde_json::to_vec(&record).unwrap(); + assert!(BloomStreamEvent::decode(&encoded).is_err(), "{reason}"); + } +} + +/// Heartbeats carry only the wire discriminators; liveness timing comes from their Kinesis arrival +/// timestamps and the shared fixed deadline. +#[test] +fn heartbeat_has_no_payload() { + let encoded = BloomStreamEvent::Heartbeat.encode().unwrap(); + let value: serde_json::Value = serde_json::from_slice(&encoded).unwrap(); + let object = value.as_object().unwrap(); + + assert_eq!(object.len(), 2); + assert_eq!(object["version"], BloomStreamEvent::CURRENT_VERSION); + assert_eq!(object["kind"], "heartbeat"); +} + +/// An unsupported version is recognized from the version alone; the other fields of the record are +/// free to change shape. +#[test] +fn unsupported_version_is_recognized_without_the_fields() { + let encoded = serde_json::to_vec(&serde_json::json!({ + "version": 255, + "kind": "path_created", + "store_path_hash": { "reshaped": true }, + })) + .unwrap(); + + assert_eq!( + BloomStreamEvent::decode(&encoded).unwrap(), + DecodedEvent::UnsupportedVersion(255) + ); +} diff --git a/src/lib.rs b/src/lib.rs index 7aa4fe2..f74e564 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -13,6 +13,8 @@ deny(unused_imports, unused_mut, unused_variables) )] +#[cfg(feature = "bloom")] +pub mod bloom; pub mod error; pub mod hash; pub mod nix_store;