store: add record-by-cid key shapes & chunked scan helpers

Lewis: May this revision serve well! <lu5a@proton.me>

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,

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())
+        );
+    }
 }

crates/tranquil-store/src/metastore/scan.rs

@@ -21,6 +21,52 @@ pub fn delete_all_by_prefix(
     })
 }
 
+pub fn stage_in_chunks<T>(
+    db: &fjall::Database,
+    mut items: impl Iterator<Item = Result<T, MetastoreError>>,
+    commit_every: usize,
+    mut stage: impl FnMut(&mut fjall::OwnedWriteBatch, T) -> Result<(), MetastoreError>,
+) -> Result<usize, MetastoreError> {
+    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<Item = Result<(fjall::Slice, fjall::Slice), MetastoreError>> + '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<usize, MetastoreError> {
+    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<T>(
     keyspace: &Keyspace,
     key: &[u8],
@@ -34,3 +80,97 @@ pub fn point_lookup<T>(
         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<u8> = 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"
+        );
+    }
+}