From 2f1e22a950da8ae1214899882093fa2d9dcd49cd Mon Sep 17 00:00:00 2001 From: Lewis Date: Sun, 19 Jul 2026 17:19:22 +0300 Subject: [PATCH] store: add record-by-cid key shapes & chunked scan helpers Lewis: May this revision serve well! --- crates/tranquil-store/src/metastore/keys.rs | 4 + .../tranquil-store/src/metastore/records.rs | 159 ++++++++++++++++++ crates/tranquil-store/src/metastore/scan.rs | 140 +++++++++++++++ 3 files changed, 303 insertions(+) diff --git a/crates/tranquil-store/src/metastore/keys.rs b/crates/tranquil-store/src/metastore/keys.rs index 9363a95..2b64151 100644 --- a/crates/tranquil-store/src/metastore/keys.rs +++ b/crates/tranquil-store/src/metastore/keys.rs @@ -53,6 +53,8 @@ impl KeyTag { pub const RECORD_BLOBS: Self = Self(0x30); pub const BACKLINK_BY_USER: Self = Self(0x31); + pub const RECORD_BY_CID: Self = Self(0x32); + pub const RECORD_BY_CID_BUILT: Self = Self(0x33); pub const USER_PRIMARY: Self = Self(0x40); pub const USER_BY_HANDLE: Self = Self(0x41); @@ -185,6 +187,8 @@ mod tests { KeyTag::DID_EVENTS, KeyTag::RECORD_BLOBS, KeyTag::BACKLINK_BY_USER, + KeyTag::RECORD_BY_CID, + KeyTag::RECORD_BY_CID_BUILT, KeyTag::USER_PRIMARY, KeyTag::USER_BY_HANDLE, KeyTag::USER_BY_EMAIL, diff --git a/crates/tranquil-store/src/metastore/records.rs b/crates/tranquil-store/src/metastore/records.rs index 8f5c824..6a9896c 100644 --- a/crates/tranquil-store/src/metastore/records.rs +++ b/crates/tranquil-store/src/metastore/records.rs @@ -59,6 +59,63 @@ pub fn records_prefix() -> SmallVec<[u8; 128]> { KeyBuilder::new().tag(KeyTag::RECORDS).build() } +pub fn record_by_cid_key( + user_hash: UserHash, + cid_bytes: &[u8], + collection: &Nsid, + rkey: &Rkey, +) -> SmallVec<[u8; 128]> { + record_by_cid_prefix(user_hash, cid_bytes) + .into_iter() + .chain(record_by_cid_suffix(collection, rkey)) + .collect() +} + +pub fn record_by_cid_suffix(collection: &Nsid, rkey: &Rkey) -> SmallVec<[u8; 128]> { + KeyBuilder::new().string(collection).string(rkey).build() +} + +pub fn record_by_cid_key_from_suffix( + user_hash: UserHash, + cid_bytes: &[u8], + suffix: &[u8], +) -> SmallVec<[u8; 128]> { + record_by_cid_prefix(user_hash, cid_bytes) + .into_iter() + .chain(suffix.iter().copied()) + .collect() +} + +pub fn record_key_user_hash_and_suffix(key_bytes: &[u8]) -> Option<(UserHash, &[u8])> { + let mut reader = super::encoding::KeyReader::new(key_bytes); + let _tag = reader.tag()?; + let hash = reader.u64()?; + Some((UserHash::from_raw(hash), reader.remaining())) +} + +pub fn record_by_cid_prefix(user_hash: UserHash, cid_bytes: &[u8]) -> SmallVec<[u8; 128]> { + KeyBuilder::new() + .tag(KeyTag::RECORD_BY_CID) + .u64(user_hash.raw()) + .bytes(cid_bytes) + .build() +} + +pub fn record_by_cid_user_prefix(user_hash: UserHash) -> SmallVec<[u8; 128]> { + KeyBuilder::new() + .tag(KeyTag::RECORD_BY_CID) + .u64(user_hash.raw()) + .build() +} + +pub fn record_by_cid_index_prefix() -> SmallVec<[u8; 128]> { + KeyBuilder::new().tag(KeyTag::RECORD_BY_CID).build() +} + +pub fn record_by_cid_built_key() -> SmallVec<[u8; 128]> { + KeyBuilder::new().tag(KeyTag::RECORD_BY_CID_BUILT).build() +} + #[cfg(test)] mod tests { use super::*; @@ -168,4 +225,106 @@ mod tests { let pfx = records_prefix(); assert_eq!(pfx.as_slice(), &[KeyTag::RECORDS.raw()]); } + + #[test] + fn record_by_cid_key_roundtrip() { + let hash = UserHash::from_raw(0xDEAD_BEEF_CAFE_BABE); + let cid = [0x01, 0x71, 0x12, 0x20, 0xAB]; + let key = record_by_cid_key(hash, &cid, &nsid("app.bsky.feed.post"), &rkey("3k2abcd")); + let mut reader = KeyReader::new(&key); + assert_eq!(reader.tag(), Some(KeyTag::RECORD_BY_CID.raw())); + assert_eq!(reader.u64(), Some(0xDEAD_BEEF_CAFE_BABE)); + assert_eq!(reader.bytes(), Some(cid.to_vec())); + assert_eq!(reader.string(), Some("app.bsky.feed.post".to_string())); + assert_eq!(reader.string(), Some("3k2abcd".to_string())); + assert!(reader.is_empty()); + } + + #[test] + fn the_raw_suffix_of_a_records_key_rebuilds_the_same_reverse_key() { + let hash = UserHash::from_raw(0xDEAD_BEEF_CAFE_BABE); + let cid = [0x01, 0x71, 0x12, 0x20, 0xAB]; + let collection = nsid("app.bsky.feed.post"); + let rkey = rkey("3k2abcd"); + + let forward = record_key(hash, &collection, &rkey); + let (parsed_hash, suffix) = record_key_user_hash_and_suffix(&forward).unwrap(); + + assert_eq!(parsed_hash, hash); + assert_eq!(suffix, record_by_cid_suffix(&collection, &rkey).as_slice()); + assert_eq!( + record_by_cid_key_from_suffix(parsed_hash, &cid, suffix).as_slice(), + record_by_cid_key(hash, &cid, &collection, &rkey).as_slice() + ); + } + + #[test] + fn a_records_key_suffix_is_copied_without_revalidating_it() { + let hash = UserHash::from_raw(5); + let cid = [0x01, 0x71]; + let forward = KeyBuilder::new() + .tag(KeyTag::RECORDS) + .u64(hash.raw()) + .string("not/an/nsid") + .string("not a valid rkey") + .build(); + + let (parsed_hash, suffix) = record_key_user_hash_and_suffix(&forward) + .expect("a structurally valid records key parses regardless of its string contents"); + let reverse = record_by_cid_key_from_suffix(parsed_hash, &cid, suffix); + + assert!(reverse.starts_with(record_by_cid_prefix(hash, &cid).as_slice())); + assert_eq!(&reverse[record_by_cid_prefix(hash, &cid).len()..], suffix); + } + + #[test] + fn record_by_cid_key_splits_into_prefix_and_suffix() { + let hash = UserHash::from_raw(0xDEAD_BEEF_CAFE_BABE); + let cid = [0x01, 0x71, 0x12, 0x20, 0xAB]; + let collection = nsid("app.bsky.feed.post"); + let rkey = rkey("3k2abcd"); + let key = record_by_cid_key(hash, &cid, &collection, &rkey); + let prefix = record_by_cid_prefix(hash, &cid); + let suffix = record_by_cid_suffix(&collection, &rkey); + assert_eq!(&key[..prefix.len()], prefix.as_slice()); + assert_eq!(&key[prefix.len()..], suffix.as_slice()); + } + + #[test] + fn record_by_cid_prefix_covers_only_that_cid() { + let hash = UserHash::from_raw(7); + let cid = [0x01, 0x02, 0x03]; + let other = [0x01, 0x02, 0x04]; + let prefix = record_by_cid_prefix(hash, &cid); + let matching = record_by_cid_key(hash, &cid, &nsid("app.bsky.feed.post"), &rkey("a")); + let different = record_by_cid_key(hash, &other, &nsid("app.bsky.feed.post"), &rkey("a")); + assert!(matching.as_slice().starts_with(prefix.as_slice())); + assert!(!different.as_slice().starts_with(prefix.as_slice())); + } + + #[test] + fn record_by_cid_prefix_is_not_confused_by_a_longer_cid() { + let hash = UserHash::from_raw(7); + let short = [0x01, 0x02]; + let long = [0x01, 0x02, 0x03]; + let prefix = record_by_cid_prefix(hash, &short); + let longer = record_by_cid_key(hash, &long, &nsid("app.bsky.feed.post"), &rkey("a")); + assert!(!longer.as_slice().starts_with(prefix.as_slice())); + } + + #[test] + fn record_by_cid_user_prefix_isolates_users() { + let h1 = UserHash::from_raw(1); + let h2 = UserHash::from_raw(2); + let cid = [0x09]; + let key = record_by_cid_key(h2, &cid, &nsid("app.bsky.feed.post"), &rkey("a")); + assert!( + key.as_slice() + .starts_with(record_by_cid_user_prefix(h2).as_slice()) + ); + assert!( + !key.as_slice() + .starts_with(record_by_cid_user_prefix(h1).as_slice()) + ); + } } diff --git a/crates/tranquil-store/src/metastore/scan.rs b/crates/tranquil-store/src/metastore/scan.rs index 8f4a217..bba6fc7 100644 --- a/crates/tranquil-store/src/metastore/scan.rs +++ b/crates/tranquil-store/src/metastore/scan.rs @@ -21,6 +21,52 @@ pub fn delete_all_by_prefix( }) } +pub fn stage_in_chunks( + db: &fjall::Database, + mut items: impl Iterator>, + commit_every: usize, + mut stage: impl FnMut(&mut fjall::OwnedWriteBatch, T) -> Result<(), MetastoreError>, +) -> Result { + let (batch, staged) = items.try_fold((db.batch(), 0usize), |(mut batch, staged), item| { + stage(&mut batch, item?)?; + match (staged + 1).is_multiple_of(commit_every) { + true => { + batch.commit()?; + Ok::<_, MetastoreError>((db.batch(), staged + 1)) + } + false => Ok((batch, staged + 1)), + } + })?; + batch.commit()?; + Ok(staged) +} + +pub fn prefix_entries<'a>( + keyspace: &'a Keyspace, + prefix: &[u8], +) -> impl Iterator> + 'a { + keyspace + .prefix(prefix) + .map(|guard| guard.into_inner().map_err(MetastoreError::Fjall)) +} + +pub fn delete_all_by_prefix_chunked( + db: &fjall::Database, + keyspace: &Keyspace, + prefix: &[u8], + commit_every: usize, +) -> Result { + stage_in_chunks( + db, + prefix_entries(keyspace, prefix), + commit_every, + |batch, (key_bytes, _)| { + batch.remove(keyspace, key_bytes.as_ref()); + Ok(()) + }, + ) +} + pub fn point_lookup( keyspace: &Keyspace, key: &[u8], @@ -34,3 +80,97 @@ pub fn point_lookup( None => Ok(None), } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::metastore::{Metastore, MetastoreConfig, Partition}; + + const PREFIX: &[u8] = &[0xE0]; + const NEIGHBOR: &[u8] = &[0xE1]; + + fn open_fresh() -> (tempfile::TempDir, Metastore) { + let dir = tempfile::TempDir::new().unwrap(); + let ms = Metastore::open(dir.path(), MetastoreConfig::default()).unwrap(); + (dir, ms) + } + + fn seed(ms: &Metastore, keyspace: &Keyspace, prefix: &[u8], count: usize) { + let mut batch = ms.database().batch(); + (0..count).for_each(|i| { + let key: Vec = prefix + .iter() + .copied() + .chain((i as u32).to_be_bytes()) + .collect(); + batch.insert(keyspace, key.as_slice(), []); + }); + batch.commit().unwrap(); + } + + fn count(keyspace: &Keyspace, prefix: &[u8]) -> usize { + keyspace.prefix(prefix).count() + } + + #[test] + fn chunked_delete_removes_every_key_across_many_mid_iteration_commits() { + let (_dir, ms) = open_fresh(); + let keyspace = ms.partition(Partition::RepoData).clone(); + let total = 1_000usize; + + seed(&ms, &keyspace, PREFIX, total); + seed(&ms, &keyspace, NEIGHBOR, 8); + + assert_eq!( + delete_all_by_prefix_chunked(ms.database(), &keyspace, PREFIX, 100).unwrap(), + total, + "every key under the prefix must be visited despite committing mid-iteration" + ); + assert_eq!( + count(&keyspace, PREFIX), + 0, + "deleting from the prefix being iterated must not skip keys" + ); + assert_eq!( + count(&keyspace, NEIGHBOR), + 8, + "chunked delete must not reach past its prefix" + ); + } + + #[test] + fn chunked_delete_commits_a_remainder_smaller_than_one_chunk() { + let (_dir, ms) = open_fresh(); + let keyspace = ms.partition(Partition::RepoData).clone(); + + seed(&ms, &keyspace, PREFIX, 101); + + assert_eq!( + delete_all_by_prefix_chunked(ms.database(), &keyspace, PREFIX, 100).unwrap(), + 101 + ); + assert_eq!(count(&keyspace, PREFIX), 0); + } + + #[test] + fn stage_in_chunks_propagates_a_failing_item() { + let (_dir, ms) = open_fresh(); + let keyspace = ms.partition(Partition::RepoData).clone(); + + let items = (0..10usize).map(|i| match i { + 7 => Err(MetastoreError::CorruptData("bad item")), + _ => Ok(i), + }); + let result = stage_in_chunks(ms.database(), items, 4, |batch, i| { + batch.insert(&keyspace, [0xE2, i as u8].as_slice(), []); + Ok(()) + }); + + assert!(matches!(result, Err(MetastoreError::CorruptData(_)))); + assert_eq!( + count(&keyspace, &[0xE2]), + 4, + "chunks committed before the failure stay committed" + ); + } +}