diff --git a/Cargo.lock b/Cargo.lock index 82fb074..f122713 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -574,13 +574,13 @@ dependencies = [ [[package]] name = "cardinality-estimator-safe" -version = "2.1.1" +version = "4.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50c14632b90cb42ff2174d2a544ca553c1bccfab54b848ae9ab9e004b90243bf" +checksum = "b41ec0cd313b46ba3b508377544b25aa1d56d05ce9e657e77dfb001d5e726e53" dependencies = [ + "digest", "enum_dispatch", "serde", - "wyhash", ] [[package]] @@ -3225,6 +3225,17 @@ dependencies = [ "digest", ] +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "sharded-slab" version = "0.1.7" @@ -3816,6 +3827,7 @@ dependencies = [ "semver", "serde", "serde_json", + "sha2", "tempfile", "thiserror 2.0.12", "tikv-jemallocator", diff --git a/ufos/Cargo.toml b/ufos/Cargo.toml index fb2bb88..6e0a4b6 100644 --- a/ufos/Cargo.toml +++ b/ufos/Cargo.toml @@ -7,7 +7,7 @@ edition = "2021" anyhow = "1.0.97" async-trait = "0.1.88" bincode = { version = "2.0.1", features = ["serde"] } -cardinality-estimator-safe = { version = "2.1.1", features = ["with_serde"] } +cardinality-estimator-safe = { version = "4.0.1", features = ["with_serde", "with_digest"] } clap = { version = "4.5.31", features = ["derive"] } dropshot = "0.16.0" env_logger = "0.11.7" @@ -20,6 +20,7 @@ schemars = { version = "0.8.22", features = ["raw_value"] } semver = "1.0.26" serde = "1.0.219" serde_json = "1.0.140" +sha2 = "0.10.9" thiserror = "2.0.12" tokio = { version = "1.44.2", features = ["full", "sync", "time"] } diff --git a/ufos/src/lib.rs b/ufos/src/lib.rs index 52d893c..dcaf274 100644 --- a/ufos/src/lib.rs +++ b/ufos/src/lib.rs @@ -10,24 +10,29 @@ pub mod storage_mem; pub mod store_types; use crate::error::BatchInsertError; -use cardinality_estimator_safe::CardinalityEstimator; +use cardinality_estimator_safe::{Element, Sketch}; use error::FirehoseEventError; use jetstream::events::{CommitEvent, CommitOp, Cursor}; use jetstream::exports::{Did, Nsid, RecordKey}; use schemars::JsonSchema; use serde::Serialize; use serde_json::value::RawValue; +use sha2::Sha256; use std::collections::HashMap; #[derive(Debug, Default, Clone)] pub struct CollectionCommits { pub total_seen: usize, - pub dids_estimate: CardinalityEstimator, + pub dids_estimate: Sketch<14>, pub commits: Vec, head: usize, non_creates: usize, } +fn did_element(did: &Did) -> Element<14> { + Element::from_digest_oneshot::(did.as_bytes()) +} + impl CollectionCommits { fn advance_head(&mut self) { self.head += 1; @@ -66,7 +71,7 @@ impl CollectionCommits { if is_create { self.total_seen += 1; - self.dids_estimate.insert(&did); + self.dids_estimate.insert(did_element(&did)); } else { self.non_creates += 1; } @@ -181,7 +186,7 @@ impl EventBatch { self.account_removes.len() } pub fn estimate_dids(&self) -> usize { - let mut estimator = CardinalityEstimator::::new(); + let mut estimator = Sketch::<14>::default(); for commits in self.commits_by_nsid.values() { estimator.merge(&commits.dids_estimate); } diff --git a/ufos/src/store_types.rs b/ufos/src/store_types.rs index 359856f..c9afd8d 100644 --- a/ufos/src/store_types.rs +++ b/ufos/src/store_types.rs @@ -3,7 +3,7 @@ use crate::db_types::{ }; use crate::{Cursor, Did, Nsid, PutAction, RecordKey, UFOsCommit}; use bincode::{Decode, Encode}; -use cardinality_estimator_safe::CardinalityEstimator; +use cardinality_estimator_safe::Sketch; use std::ops::Range; macro_rules! static_str { @@ -199,7 +199,7 @@ pub struct TotalRecordsValue(pub u64); impl UseBincodePlz for TotalRecordsValue {} #[derive(Debug, PartialEq, serde::Serialize, serde::Deserialize)] -pub struct EstimatedDidsValue(pub CardinalityEstimator); +pub struct EstimatedDidsValue(pub Sketch<14>); impl SerdeBytes for EstimatedDidsValue {} impl DbBytes for EstimatedDidsValue { #[cfg(test)] @@ -223,7 +223,7 @@ impl DbBytes for EstimatedDidsValue { pub type CountsValue = DbConcat; impl CountsValue { - pub fn new(total: u64, dids: CardinalityEstimator) -> Self { + pub fn new(total: u64, dids: Sketch<14>) -> Self { Self { prefix: TotalRecordsValue(total), suffix: EstimatedDidsValue(dids), @@ -232,7 +232,7 @@ impl CountsValue { pub fn records(&self) -> u64 { self.prefix.0 } - pub fn dids(&self) -> &CardinalityEstimator { + pub fn dids(&self) -> &Sketch<14> { &self.suffix.0 } pub fn merge(&mut self, other: &Self) { @@ -244,7 +244,7 @@ impl Default for CountsValue { fn default() -> Self { Self { prefix: TotalRecordsValue(0), - suffix: EstimatedDidsValue(CardinalityEstimator::new()), + suffix: EstimatedDidsValue(Sketch::<14>::default()), } } } @@ -433,10 +433,12 @@ pub type WeekTruncatedCursor = TruncatedCursor; #[cfg(test)] mod test { use super::{ - CardinalityEstimator, CountsValue, Cursor, Did, EncodingError, HourTruncatedCursor, - HourlyRollupKey, Nsid, HOUR_IN_MICROS, + CountsValue, Cursor, Did, EncodingError, HourTruncatedCursor, HourlyRollupKey, Nsid, + Sketch, HOUR_IN_MICROS, }; use crate::db_types::DbBytes; + use cardinality_estimator_safe::Element; + use sha2::Sha256; #[test] fn test_by_hourly_rollup_key() -> Result<(), EncodingError> { @@ -456,9 +458,14 @@ mod test { #[test] fn test_by_hourly_rollup_value() -> Result<(), EncodingError> { - let mut estimator = CardinalityEstimator::new(); + let mut estimator = Sketch::<14>::default(); + fn to_element(d: Did) -> Element<14> { + Element::from_digest_oneshot::(d.to_string().as_bytes()) + } for i in 0..10 { - estimator.insert(&Did::new(format!("did:plc:inze6wrmsm7pjl7yta3oig7{i}")).unwrap()); + estimator.insert(to_element( + Did::new(format!("did:plc:inze6wrmsm7pjl7yta3oig7{i}")).unwrap(), + )); } let original = CountsValue::new(123, estimator.clone()); let serialized = original.to_db_bytes()?; @@ -467,7 +474,9 @@ mod test { assert_eq!(bytes_consumed, serialized.len()); for i in 10..1_000 { - estimator.insert(&Did::new(format!("did:plc:inze6wrmsm7pjl7yta3oig{i}")).unwrap()); + estimator.insert(to_element( + Did::new(format!("did:plc:inze6wrmsm7pjl7yta3oig{i}")).unwrap(), + )); } let original = CountsValue::new(123, estimator); let serialized = original.to_db_bytes()?; -- 2.54.0 From 0f3c2f87d75d57f0376ed28cd0f076966ba787cf Mon Sep 17 00:00:00 2001 From: phil Date: Fri, 23 May 2025 15:50:39 -0400 Subject: [PATCH 2/2] HLL sketch: secret hash prefix --- Cargo.lock | 20 ++-- ufos/Cargo.toml | 1 + ufos/src/consumer.rs | 14 ++- ufos/src/db_types.rs | 2 + ufos/src/file_consumer.rs | 8 +- ufos/src/lib.rs | 234 ++++++++++++++++++++++---------------- ufos/src/main.rs | 12 +- ufos/src/storage.rs | 3 +- ufos/src/storage_fjall.rs | 74 +++++++++--- ufos/src/storage_mem.rs | 17 +-- ufos/src/store_types.rs | 4 + 11 files changed, 248 insertions(+), 141 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f122713..ba290f0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1360,9 +1360,9 @@ dependencies = [ [[package]] name = "getrandom" -version = "0.3.2" +version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73fea8450eea4bac3940448fb7ae50d91f034f941199fcd9d909a5a07aa455f0" +checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4" dependencies = [ "cfg-if", "libc", @@ -1909,7 +1909,7 @@ version = "0.1.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "38f262f097c174adebe41eb73d66ae9c06b2844fb0da69969647bbddd9b0538a" dependencies = [ - "getrandom 0.3.2", + "getrandom 0.3.3", "libc", ] @@ -2697,13 +2697,12 @@ dependencies = [ [[package]] name = "rand" -version = "0.9.0" +version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94" +checksum = "9fbfd9d094a40bf3ae768db9361049ace4c0e04a4fd6b359518bd7b73a73dd97" dependencies = [ "rand_chacha 0.9.0", "rand_core 0.9.3", - "zerocopy 0.8.24", ] [[package]] @@ -2741,7 +2740,7 @@ version = "0.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38" dependencies = [ - "getrandom 0.3.2", + "getrandom 0.3.3", ] [[package]] @@ -3426,7 +3425,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7437ac7763b9b123ccf33c338a5cc1bac6f69b45a136c19bdd8a65e3916435bf" dependencies = [ "fastrand", - "getrandom 0.3.2", + "getrandom 0.3.3", "once_cell", "rustix 1.0.5", "windows-sys 0.59.0", @@ -3794,7 +3793,7 @@ dependencies = [ "httparse", "log", "native-tls", - "rand 0.9.0", + "rand 0.9.1", "sha1", "thiserror 2.0.12", "url", @@ -3819,6 +3818,7 @@ dependencies = [ "dropshot", "env_logger", "fjall", + "getrandom 0.3.3", "http", "jetstream", "log", @@ -3917,7 +3917,7 @@ version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "458f7a779bf54acc9f347480ac654f68407d3aab21269a6e3c9f922acd9e2da9" dependencies = [ - "getrandom 0.3.2", + "getrandom 0.3.3", "serde", ] diff --git a/ufos/Cargo.toml b/ufos/Cargo.toml index 6e0a4b6..2340285 100644 --- a/ufos/Cargo.toml +++ b/ufos/Cargo.toml @@ -12,6 +12,7 @@ clap = { version = "4.5.31", features = ["derive"] } dropshot = "0.16.0" env_logger = "0.11.7" fjall = { version = "2.8.0", features = ["lz4"] } +getrandom = "0.3.3" http = "1.3.1" jetstream = { path = "../jetstream" } log = "0.4.26" diff --git a/ufos/src/consumer.rs b/ufos/src/consumer.rs index 3165e4e..b93bc33 100644 --- a/ufos/src/consumer.rs +++ b/ufos/src/consumer.rs @@ -1,3 +1,4 @@ +use crate::store_types::SketchSecretPrefix; use jetstream::{ events::{Cursor, EventKind, JetstreamEvent}, exports::{Did, Nsid}, @@ -32,12 +33,14 @@ pub struct Batcher { jetstream_receiver: JetstreamReceiver, batch_sender: Sender, current_batch: CurrentBatch, + sketch_secret: SketchSecretPrefix, } pub async fn consume( jetstream_endpoint: &str, cursor: Option, no_compress: bool, + sketch_secret: SketchSecretPrefix, ) -> anyhow::Result> { let endpoint = DefaultJetstreamEndpoints::endpoint_or_shortcut(jetstream_endpoint); if endpoint == jetstream_endpoint { @@ -60,17 +63,22 @@ pub async fn consume( .connect_cursor(cursor) .await?; let (batch_sender, batch_reciever) = channel::(BATCH_QUEUE_SIZE); - let mut batcher = Batcher::new(jetstream_receiver, batch_sender); + let mut batcher = Batcher::new(jetstream_receiver, batch_sender, sketch_secret); tokio::task::spawn(async move { batcher.run().await }); Ok(batch_reciever) } impl Batcher { - pub fn new(jetstream_receiver: JetstreamReceiver, batch_sender: Sender) -> Self { + pub fn new( + jetstream_receiver: JetstreamReceiver, + batch_sender: Sender, + sketch_secret: SketchSecretPrefix, + ) -> Self { Self { jetstream_receiver, batch_sender, current_batch: Default::default(), + sketch_secret, } } @@ -129,6 +137,7 @@ impl Batcher { &collection, commit, MAX_BATCHED_COLLECTIONS, + &self.sketch_secret, ); if let Err(BatchInsertError::BatchFull(commit)) = optimistic_res { @@ -137,6 +146,7 @@ impl Batcher { &collection, commit, MAX_BATCHED_COLLECTIONS, + &self.sketch_secret, )?; } else { optimistic_res?; diff --git a/ufos/src/db_types.rs b/ufos/src/db_types.rs index 99339b3..8d0f71d 100644 --- a/ufos/src/db_types.rs +++ b/ufos/src/db_types.rs @@ -224,6 +224,8 @@ pub trait SerdeBytes: serde::Serialize + for<'a> serde::Deserialize<'a> { ////// +impl UseBincodePlz for [u8; N] {} + impl DbBytes for Vec { fn to_db_bytes(&self) -> Result, EncodingError> { Ok(self.to_vec()) diff --git a/ufos/src/file_consumer.rs b/ufos/src/file_consumer.rs index fb0bfed..b5ecbaa 100644 --- a/ufos/src/file_consumer.rs +++ b/ufos/src/file_consumer.rs @@ -1,4 +1,5 @@ use crate::consumer::{Batcher, LimitedBatch, BATCH_QUEUE_SIZE}; +use crate::store_types::SketchSecretPrefix; use anyhow::Result; use jetstream::{error::JetstreamEventError, events::JetstreamEvent}; use std::path::PathBuf; @@ -21,11 +22,14 @@ async fn read_jsonl(f: File, sender: Sender) -> Result<()> { Ok(()) } -pub async fn consume(p: PathBuf) -> Result> { +pub async fn consume( + p: PathBuf, + sketch_secret: SketchSecretPrefix, +) -> Result> { let f = File::open(p).await?; let (jsonl_sender, jsonl_receiver) = channel::(16); let (batch_sender, batch_reciever) = channel::(BATCH_QUEUE_SIZE); - let mut batcher = Batcher::new(jsonl_receiver, batch_sender); + let mut batcher = Batcher::new(jsonl_receiver, batch_sender, sketch_secret); tokio::task::spawn(async move { read_jsonl(f, jsonl_sender).await }); tokio::task::spawn(async move { batcher.run().await }); Ok(batch_reciever) diff --git a/ufos/src/lib.rs b/ufos/src/lib.rs index dcaf274..2af46ec 100644 --- a/ufos/src/lib.rs +++ b/ufos/src/lib.rs @@ -10,6 +10,7 @@ pub mod storage_mem; pub mod store_types; use crate::error::BatchInsertError; +use crate::store_types::SketchSecretPrefix; use cardinality_estimator_safe::{Element, Sketch}; use error::FirehoseEventError; use jetstream::events::{CommitEvent, CommitOp, Cursor}; @@ -20,6 +21,10 @@ use serde_json::value::RawValue; use sha2::Sha256; use std::collections::HashMap; +fn did_element(sketch_secret: &SketchSecretPrefix, did: &Did) -> Element<14> { + Element::from_digest_with_prefix::(sketch_secret, did.as_bytes()) +} + #[derive(Debug, Default, Clone)] pub struct CollectionCommits { pub total_seen: usize, @@ -29,10 +34,6 @@ pub struct CollectionCommits { non_creates: usize, } -fn did_element(did: &Did) -> Element<14> { - Element::from_digest_oneshot::(did.as_bytes()) -} - impl CollectionCommits { fn advance_head(&mut self) { self.head += 1; @@ -40,7 +41,11 @@ impl CollectionCommits { self.head = 0; } } - pub fn truncating_insert(&mut self, commit: UFOsCommit) -> Result<(), BatchInsertError> { + pub fn truncating_insert( + &mut self, + commit: UFOsCommit, + sketch_secret: &SketchSecretPrefix, + ) -> Result<(), BatchInsertError> { if self.non_creates == LIMIT { return Err(BatchInsertError::BatchFull(commit)); } @@ -71,7 +76,7 @@ impl CollectionCommits { if is_create { self.total_seen += 1; - self.dids_estimate.insert(did_element(&did)); + self.dids_estimate.insert(did_element(sketch_secret, &did)); } else { self.non_creates += 1; } @@ -163,6 +168,7 @@ impl EventBatch { collection: &Nsid, commit: UFOsCommit, max_collections: usize, + sketch_secret: &SketchSecretPrefix, ) -> Result<(), BatchInsertError> { let map = &mut self.commits_by_nsid; if !map.contains_key(collection) && map.len() >= max_collections { @@ -170,7 +176,7 @@ impl EventBatch { } map.entry(collection.clone()) .or_default() - .truncating_insert(commit)?; + .truncating_insert(commit, sketch_secret)?; Ok(()) } pub fn total_records(&self) -> usize { @@ -313,38 +319,47 @@ mod tests { fn test_truncating_insert_truncates() -> anyhow::Result<()> { let mut commits: CollectionCommits<2> = Default::default(); - commits.truncating_insert(UFOsCommit { - cursor: Cursor::from_raw_u64(100), - did: Did::new("did:plc:whatever".to_string()).unwrap(), - rkey: RecordKey::new("rkey-asdf-a".to_string()).unwrap(), - rev: "rev-asdf".to_string(), - action: CommitAction::Put(PutAction { - record: RawValue::from_string("{}".to_string())?, - is_update: false, - }), - })?; - - commits.truncating_insert(UFOsCommit { - cursor: Cursor::from_raw_u64(101), - did: Did::new("did:plc:whatever".to_string()).unwrap(), - rkey: RecordKey::new("rkey-asdf-b".to_string()).unwrap(), - rev: "rev-asdg".to_string(), - action: CommitAction::Put(PutAction { - record: RawValue::from_string("{}".to_string())?, - is_update: false, - }), - })?; - - commits.truncating_insert(UFOsCommit { - cursor: Cursor::from_raw_u64(102), - did: Did::new("did:plc:whatever".to_string()).unwrap(), - rkey: RecordKey::new("rkey-asdf-c".to_string()).unwrap(), - rev: "rev-asdh".to_string(), - action: CommitAction::Put(PutAction { - record: RawValue::from_string("{}".to_string())?, - is_update: false, - }), - })?; + commits.truncating_insert( + UFOsCommit { + cursor: Cursor::from_raw_u64(100), + did: Did::new("did:plc:whatever".to_string()).unwrap(), + rkey: RecordKey::new("rkey-asdf-a".to_string()).unwrap(), + rev: "rev-asdf".to_string(), + action: CommitAction::Put(PutAction { + record: RawValue::from_string("{}".to_string())?, + is_update: false, + }), + }, + &[0u8; 16], + )?; + + commits.truncating_insert( + UFOsCommit { + cursor: Cursor::from_raw_u64(101), + did: Did::new("did:plc:whatever".to_string()).unwrap(), + rkey: RecordKey::new("rkey-asdf-b".to_string()).unwrap(), + rev: "rev-asdg".to_string(), + action: CommitAction::Put(PutAction { + record: RawValue::from_string("{}".to_string())?, + is_update: false, + }), + }, + &[0u8; 16], + )?; + + commits.truncating_insert( + UFOsCommit { + cursor: Cursor::from_raw_u64(102), + did: Did::new("did:plc:whatever".to_string()).unwrap(), + rkey: RecordKey::new("rkey-asdf-c".to_string()).unwrap(), + rev: "rev-asdh".to_string(), + action: CommitAction::Put(PutAction { + record: RawValue::from_string("{}".to_string())?, + is_update: false, + }), + }, + &[0u8; 16], + )?; assert_eq!(commits.total_seen, 3); assert_eq!(commits.dids_estimate.estimate(), 1); @@ -373,35 +388,44 @@ mod tests { fn test_truncating_insert_does_not_truncate_deletes() -> anyhow::Result<()> { let mut commits: CollectionCommits<2> = Default::default(); - commits.truncating_insert(UFOsCommit { - cursor: Cursor::from_raw_u64(100), - did: Did::new("did:plc:whatever".to_string()).unwrap(), - rkey: RecordKey::new("rkey-asdf-a".to_string()).unwrap(), - rev: "rev-asdf".to_string(), - action: CommitAction::Cut, - })?; - - commits.truncating_insert(UFOsCommit { - cursor: Cursor::from_raw_u64(101), - did: Did::new("did:plc:whatever".to_string()).unwrap(), - rkey: RecordKey::new("rkey-asdf-b".to_string()).unwrap(), - rev: "rev-asdg".to_string(), - action: CommitAction::Put(PutAction { - record: RawValue::from_string("{}".to_string())?, - is_update: false, - }), - })?; - - commits.truncating_insert(UFOsCommit { - cursor: Cursor::from_raw_u64(102), - did: Did::new("did:plc:whatever".to_string()).unwrap(), - rkey: RecordKey::new("rkey-asdf-c".to_string()).unwrap(), - rev: "rev-asdh".to_string(), - action: CommitAction::Put(PutAction { - record: RawValue::from_string("{}".to_string())?, - is_update: false, - }), - })?; + commits.truncating_insert( + UFOsCommit { + cursor: Cursor::from_raw_u64(100), + did: Did::new("did:plc:whatever".to_string()).unwrap(), + rkey: RecordKey::new("rkey-asdf-a".to_string()).unwrap(), + rev: "rev-asdf".to_string(), + action: CommitAction::Cut, + }, + &[0u8; 16], + )?; + + commits.truncating_insert( + UFOsCommit { + cursor: Cursor::from_raw_u64(101), + did: Did::new("did:plc:whatever".to_string()).unwrap(), + rkey: RecordKey::new("rkey-asdf-b".to_string()).unwrap(), + rev: "rev-asdg".to_string(), + action: CommitAction::Put(PutAction { + record: RawValue::from_string("{}".to_string())?, + is_update: false, + }), + }, + &[0u8; 16], + )?; + + commits.truncating_insert( + UFOsCommit { + cursor: Cursor::from_raw_u64(102), + did: Did::new("did:plc:whatever".to_string()).unwrap(), + rkey: RecordKey::new("rkey-asdf-c".to_string()).unwrap(), + rev: "rev-asdh".to_string(), + action: CommitAction::Put(PutAction { + record: RawValue::from_string("{}".to_string())?, + is_update: false, + }), + }, + &[0u8; 16], + )?; assert_eq!(commits.total_seen, 2); assert_eq!(commits.dids_estimate.estimate(), 1); @@ -436,46 +460,58 @@ mod tests { let mut commits: CollectionCommits<2> = Default::default(); commits - .truncating_insert(UFOsCommit { - cursor: Cursor::from_raw_u64(100), - did: Did::new("did:plc:whatever".to_string()).unwrap(), - rkey: RecordKey::new("rkey-asdf-a".to_string()).unwrap(), - rev: "rev-asdf".to_string(), - action: CommitAction::Cut, - }) + .truncating_insert( + UFOsCommit { + cursor: Cursor::from_raw_u64(100), + did: Did::new("did:plc:whatever".to_string()).unwrap(), + rkey: RecordKey::new("rkey-asdf-a".to_string()).unwrap(), + rev: "rev-asdf".to_string(), + action: CommitAction::Cut, + }, + &[0u8; 16], + ) .unwrap(); // this create will just be discarded commits - .truncating_insert(UFOsCommit { - cursor: Cursor::from_raw_u64(80), - did: Did::new("did:plc:whatever".to_string()).unwrap(), - rkey: RecordKey::new("rkey-asdf-zzz".to_string()).unwrap(), - rev: "rev-asdzzz".to_string(), - action: CommitAction::Put(PutAction { - record: RawValue::from_string("{}".to_string())?, - is_update: false, - }), - }) + .truncating_insert( + UFOsCommit { + cursor: Cursor::from_raw_u64(80), + did: Did::new("did:plc:whatever".to_string()).unwrap(), + rkey: RecordKey::new("rkey-asdf-zzz".to_string()).unwrap(), + rev: "rev-asdzzz".to_string(), + action: CommitAction::Put(PutAction { + record: RawValue::from_string("{}".to_string())?, + is_update: false, + }), + }, + &[0u8; 16], + ) .unwrap(); commits - .truncating_insert(UFOsCommit { - cursor: Cursor::from_raw_u64(101), - did: Did::new("did:plc:whatever".to_string()).unwrap(), - rkey: RecordKey::new("rkey-asdf-b".to_string()).unwrap(), - rev: "rev-asdg".to_string(), - action: CommitAction::Cut, - }) + .truncating_insert( + UFOsCommit { + cursor: Cursor::from_raw_u64(101), + did: Did::new("did:plc:whatever".to_string()).unwrap(), + rkey: RecordKey::new("rkey-asdf-b".to_string()).unwrap(), + rev: "rev-asdg".to_string(), + action: CommitAction::Cut, + }, + &[0u8; 16], + ) .unwrap(); - let res = commits.truncating_insert(UFOsCommit { - cursor: Cursor::from_raw_u64(102), - did: Did::new("did:plc:whatever".to_string()).unwrap(), - rkey: RecordKey::new("rkey-asdf-c".to_string()).unwrap(), - rev: "rev-asdh".to_string(), - action: CommitAction::Cut, - }); + let res = commits.truncating_insert( + UFOsCommit { + cursor: Cursor::from_raw_u64(102), + did: Did::new("did:plc:whatever".to_string()).unwrap(), + rkey: RecordKey::new("rkey-asdf-c".to_string()).unwrap(), + rev: "rev-asdh".to_string(), + action: CommitAction::Cut, + }, + &[0u8; 16], + ); assert!(res.is_err()); let overflowed = match res { diff --git a/ufos/src/main.rs b/ufos/src/main.rs index 73ea985..5c77ad0 100644 --- a/ufos/src/main.rs +++ b/ufos/src/main.rs @@ -7,6 +7,7 @@ use ufos::server; use ufos::storage::{StorageWhatever, StoreBackground, StoreReader, StoreWriter}; use ufos::storage_fjall::FjallStorage; use ufos::storage_mem::MemStorage; +use ufos::store_types::SketchSecretPrefix; #[cfg(not(target_env = "msvc"))] use tikv_jemallocator::Jemalloc; @@ -57,7 +58,7 @@ async fn main() -> anyhow::Result<()> { let args = Args::parse(); let jetstream = args.jetstream.clone(); if args.in_mem { - let (read_store, write_store, cursor) = MemStorage::init( + let (read_store, write_store, cursor, sketch_secret) = MemStorage::init( args.data, jetstream, args.jetstream_force, @@ -70,10 +71,11 @@ async fn main() -> anyhow::Result<()> { read_store, write_store, cursor, + sketch_secret, ) .await?; } else { - let (read_store, write_store, cursor) = FjallStorage::init( + let (read_store, write_store, cursor, sketch_secret) = FjallStorage::init( args.data, jetstream, args.jetstream_force, @@ -86,6 +88,7 @@ async fn main() -> anyhow::Result<()> { read_store, write_store, cursor, + sketch_secret, ) .await?; } @@ -100,6 +103,7 @@ async fn go( read_store: impl StoreReader + 'static, mut write_store: impl StoreWriter + 'static, cursor: Option, + sketch_secret: SketchSecretPrefix, ) -> anyhow::Result<()> { println!("starting server with storage..."); let serving = server::serve(read_store); @@ -112,13 +116,13 @@ async fn go( let batches = if jetstream_fixture { log::info!("starting with jestream file fixture: {jetstream:?}"); - file_consumer::consume(jetstream.into()).await? + file_consumer::consume(jetstream.into(), sketch_secret).await? } else { log::info!( "starting consumer with cursor: {cursor:?} from {:?} ago", cursor.map(|c| c.elapsed()) ); - consumer::consume(&jetstream, cursor, false).await? + consumer::consume(&jetstream, cursor, false, sketch_secret).await? }; let rolling = write_store.background_tasks()?.run(); diff --git a/ufos/src/storage.rs b/ufos/src/storage.rs index 103d1b5..8ab709e 100644 --- a/ufos/src/storage.rs +++ b/ufos/src/storage.rs @@ -1,3 +1,4 @@ +use crate::store_types::SketchSecretPrefix; use crate::{ error::StorageError, ConsumerInfo, Count, Cursor, EventBatch, QueryPeriod, TopCollections, UFOsRecord, @@ -16,7 +17,7 @@ pub trait StorageWhatever, B: StoreBackground, endpoint: String, force_endpoint: bool, config: C, - ) -> StorageResult<(R, W, Option)> + ) -> StorageResult<(R, W, Option, SketchSecretPrefix)> where Self: Sized; } diff --git a/ufos/src/storage_fjall.rs b/ufos/src/storage_fjall.rs index dc2622d..5ef4db1 100644 --- a/ufos/src/storage_fjall.rs +++ b/ufos/src/storage_fjall.rs @@ -6,9 +6,9 @@ use crate::store_types::{ DeleteAccountQueueVal, HourTruncatedCursor, HourlyDidsKey, HourlyRecordsKey, HourlyRollupKey, JetstreamCursorKey, JetstreamCursorValue, JetstreamEndpointKey, JetstreamEndpointValue, LiveCountsKey, NewRollupCursorKey, NewRollupCursorValue, NsidRecordFeedKey, NsidRecordFeedVal, - RecordLocationKey, RecordLocationMeta, RecordLocationVal, RecordRawValue, TakeoffKey, - TakeoffValue, TrimCollectionCursorKey, WeekTruncatedCursor, WeeklyDidsKey, WeeklyRecordsKey, - WeeklyRollupKey, + RecordLocationKey, RecordLocationMeta, RecordLocationVal, RecordRawValue, SketchSecretKey, + SketchSecretPrefix, TakeoffKey, TakeoffValue, TrimCollectionCursorKey, WeekTruncatedCursor, + WeeklyDidsKey, WeeklyRecordsKey, WeeklyRollupKey, }; use crate::{ CommitAction, ConsumerInfo, Count, Did, EventBatch, Nsid, QueryPeriod, TopCollections, @@ -46,6 +46,10 @@ const MAX_BATCHED_ROLLUP_COUNTS: usize = 256; /// - key: "takeoff" (literal) /// - val: u64 (micros timestamp, not from jetstream for now so not precise) /// +/// - Cardinality estimator secret +/// - key: "sketch_secret" (literal) +/// - val: [u8; 16] +/// /// - Rollup cursor (bg work: roll stats into hourlies, delete accounts, old record deletes) /// - key: "rollup_cursor" (literal) /// - val: u64 (tracks behind js_cursor) @@ -141,7 +145,7 @@ impl StorageWhatever for endpoint: String, force_endpoint: bool, _config: FjallConfig, - ) -> StorageResult<(FjallReader, FjallWriter, Option)> { + ) -> StorageResult<(FjallReader, FjallWriter, Option, SketchSecretPrefix)> { let keyspace = { let config = Config::new(path); @@ -159,14 +163,21 @@ impl StorageWhatever for let js_cursor = get_static_neu::(&global)?; - if js_cursor.is_some() { + let sketch_secret = if js_cursor.is_some() { let stored_endpoint = get_static_neu::(&global)?; - let JetstreamEndpointValue(stored) = stored_endpoint.ok_or(StorageError::InitError( "found cursor but missing js_endpoint, refusing to start.".to_string(), ))?; + let Some(stored_secret) = + get_static_neu::(&global)? + else { + return Err(StorageError::InitError( + "found cursor but missing sketch_secret, refusing to start.".to_string(), + )); + }; + if stored != endpoint { if force_endpoint { log::warn!("forcing a jetstream switch from {stored:?} to {endpoint:?}"); @@ -179,14 +190,28 @@ impl StorageWhatever for "stored js_endpoint {stored:?} differs from provided {endpoint:?}, refusing to start."))); } } + stored_secret } else { - insert_static_neu::( + log::info!("initializing a fresh db!"); + init_static_neu::( &global, JetstreamEndpointValue(endpoint.to_string()), )?; - insert_static_neu::(&global, Cursor::at(SystemTime::now()))?; - insert_static_neu::(&global, Cursor::from_start())?; - } + + log::info!("generating new secret for cardinality sketches..."); + let mut sketch_secret: SketchSecretPrefix = [0u8; 16]; + getrandom::fill(&mut sketch_secret).map_err(|e| { + StorageError::InitError(format!( + "failed to get a random secret for cardinality sketches: {e:?}" + )) + })?; + init_static_neu::(&global, sketch_secret)?; + + init_static_neu::(&global, Cursor::at(SystemTime::now()))?; + init_static_neu::(&global, Cursor::from_start())?; + + sketch_secret + }; let reader = FjallReader { keyspace: keyspace.clone(), @@ -204,7 +229,7 @@ impl StorageWhatever for rollups, queues, }; - Ok((reader, writer, js_cursor)) + Ok((reader, writer, js_cursor, sketch_secret)) } } @@ -1089,6 +1114,25 @@ fn insert_static_neu( Ok(()) } +/// Set a value to a fixed key, erroring if the value already exists +/// +/// Intended for single-threaded init: not safe under concurrency, since there +/// is no transaction between checking if the already exists and writing it. +fn init_static_neu( + global: &PartitionHandle, + value: impl DbBytes, +) -> StorageResult<()> { + let key_bytes = DbStaticStr::::default().to_db_bytes()?; + if global.get(&key_bytes)?.is_some() { + return Err(StorageError::InitError(format!( + "init failed: value for key {key_bytes:?} already exists" + ))); + } + let value_bytes = value.to_db_bytes()?; + global.insert(&key_bytes, &value_bytes)?; + Ok(()) +} + /// Set a value to a fixed key fn insert_batch_static_neu( batch: &mut FjallBatch, @@ -1132,7 +1176,7 @@ mod tests { use serde_json::value::RawValue; fn fjall_db() -> (FjallReader, FjallWriter) { - let (read, write, _) = FjallStorage::init( + let (read, write, _, _) = FjallStorage::init( tempfile::tempdir().unwrap(), "offline test (no real jetstream endpoint)".to_string(), false, @@ -1187,7 +1231,7 @@ mod tests { .commits_by_nsid .entry(collection.clone()) .or_default() - .truncating_insert(commit) + .truncating_insert(commit, &[0u8; 16]) .unwrap(); collection @@ -1229,7 +1273,7 @@ mod tests { .commits_by_nsid .entry(collection.clone()) .or_default() - .truncating_insert(commit) + .truncating_insert(commit, &[0u8; 16]) .unwrap(); collection @@ -1261,7 +1305,7 @@ mod tests { .commits_by_nsid .entry(collection.clone()) .or_default() - .truncating_insert(commit) + .truncating_insert(commit, &[0u8; 16]) .unwrap(); collection diff --git a/ufos/src/storage_mem.rs b/ufos/src/storage_mem.rs index 23e41ee..8b26614 100644 --- a/ufos/src/storage_mem.rs +++ b/ufos/src/storage_mem.rs @@ -9,8 +9,8 @@ use crate::store_types::{ HourTruncatedCursor, HourlyRollupKey, JetstreamCursorKey, JetstreamCursorValue, JetstreamEndpointKey, JetstreamEndpointValue, LiveCountsKey, NewRollupCursorKey, NewRollupCursorValue, NsidRecordFeedKey, NsidRecordFeedVal, RecordLocationKey, - RecordLocationMeta, RecordLocationVal, RecordRawValue, TakeoffKey, TakeoffValue, - WeekTruncatedCursor, WeeklyRollupKey, + RecordLocationMeta, RecordLocationVal, RecordRawValue, SketchSecretPrefix, TakeoffKey, + TakeoffValue, WeekTruncatedCursor, WeeklyRollupKey, }; use crate::{ CommitAction, ConsumerInfo, Count, Did, EventBatch, Nsid, QueryPeriod, TopCollections, @@ -257,7 +257,7 @@ impl StorageWhatever for MemStor endpoint: String, force_endpoint: bool, _config: MemConfig, - ) -> StorageResult<(MemReader, MemWriter, Option)> { + ) -> StorageResult<(MemReader, MemWriter, Option, SketchSecretPrefix)> { let keyspace = MemKeyspace::open(); let global = keyspace.open_partition("global")?; @@ -312,7 +312,8 @@ impl StorageWhatever for MemStor rollups, queues, }; - Ok((reader, writer, js_cursor)) + let secret_prefix = [0u8; 16]; // in-mem store is always deterministic: no secret + Ok((reader, writer, js_cursor, secret_prefix)) } } @@ -1106,7 +1107,7 @@ mod tests { use serde_json::value::RawValue; fn fjall_db() -> (MemReader, MemWriter) { - let (read, write, _) = MemStorage::init( + let (read, write, _, _) = MemStorage::init( tempfile::tempdir().unwrap(), "offline test (no real jetstream endpoint)".to_string(), false, @@ -1161,7 +1162,7 @@ mod tests { .commits_by_nsid .entry(collection.clone()) .or_default() - .truncating_insert(commit) + .truncating_insert(commit, &[0u8; 16]) .unwrap(); collection @@ -1203,7 +1204,7 @@ mod tests { .commits_by_nsid .entry(collection.clone()) .or_default() - .truncating_insert(commit) + .truncating_insert(commit, &[0u8; 16]) .unwrap(); collection @@ -1235,7 +1236,7 @@ mod tests { .commits_by_nsid .entry(collection.clone()) .or_default() - .truncating_insert(commit) + .truncating_insert(commit, &[0u8; 16]) .unwrap(); collection diff --git a/ufos/src/store_types.rs b/ufos/src/store_types.rs index c9afd8d..8356a73 100644 --- a/ufos/src/store_types.rs +++ b/ufos/src/store_types.rs @@ -22,6 +22,10 @@ macro_rules! static_str { static_str!("js_cursor", JetstreamCursorKey); pub type JetstreamCursorValue = Cursor; +// key format: ["sketch_secret"] +static_str!("sketch_secret", SketchSecretKey); +pub type SketchSecretPrefix = [u8; 16]; + // key format: ["rollup_cursor"] static_str!("rollup_cursor", NewRollupCursorKey); // pub type NewRollupCursorKey = DbStaticStr<_NewRollupCursorKey>;