diff --git a/src/backfill/worker/process.rs b/src/backfill/worker/process.rs index f36e661..56a94d3 100644 --- a/src/backfill/worker/process.rs +++ b/src/backfill/worker/process.rs @@ -13,14 +13,11 @@ use jacquard_api::com_atproto::sync::get_repo::GetRepoError; use jacquard_common::IntoStatic; use jacquard_common::types::cid::Cid as AtCid; use jacquard_common::types::did::Did; -use jacquard_repo::BlockStore; -use jacquard_repo::mst::Mst; use crate::backfill::admission::BackfillAdmission; use crate::backfill::client::{ThrottledHttpClient, collect_body_bounded}; use crate::backfill::error::BackfillError; use crate::backfill::sparse::{SparseBackfillResult, process_did_sparse}; -use crate::car::CarBlockStore; use crate::config::{BackfillStrategy, RateTier}; use crate::db::types::{DbAction, DbRkey}; use crate::db::{self, Txn as DbTxn, keys}; @@ -278,22 +275,18 @@ pub(super) async fn process_did( let parsed = tokio::task::spawn_blocking(move || crate::car::parse_car(car_bytes, verify_cids)) .await .into_diagnostic()??; - trace!(elapsed = %start.elapsed().as_secs_f32(), verify_cids, "parsed car"); - - let start = Instant::now(); - let root_cid = parsed.root; - let store = Arc::new(CarBlockStore::new(parsed.blocks)); trace!( - blocks = store.len(), - elapsed = ?start.elapsed(), - "stored blocks in memory" + elapsed = %start.elapsed().as_secs_f32(), + blocks = parsed.blocks.len(), + verify_cids, + "parsed car" ); // 4. parse root commit to get mst root - let root_bytes = store - .get(&root_cid) - .await - .into_diagnostic()? + let blocks = parsed.blocks; + let root_bytes = blocks + .get(&parsed.root) + .cloned() .ok_or_else(|| miette::miette!("root block missing from CAR"))?; let root_commit = @@ -316,20 +309,16 @@ pub(super) async fn process_did( let root_commit = Commit::from(root_commit); - // 5. walk mst and fetch every record block in one batch, off the runtime. the sequential - // walk decodes every node on this blocking thread; `leaves()` would spawn a task per - // subtree onto the runtime workers and spend 3-4x the cpu doing it. + // 5. walk mst and fetch every record block, off the runtime let start = Instant::now(); let mst_root = root_commit.data; - let handle = tokio::runtime::Handle::current(); let records = tokio::task::spawn_blocking(move || { - let mst: Mst = Mst::load(store, mst_root, None); - let leaves = handle.block_on(mst.leaves_sequential()).into_diagnostic()?; - let leaf_cids = leaves.iter().map(|(_, cid)| *cid).collect::>(); - let leaf_blocks = handle - .block_on(mst.storage().get_many(&leaf_cids)) - .into_diagnostic()?; - Ok::<_, miette::Report>(leaves.into_iter().zip(leaf_blocks)) + let leaves = crate::mst::leaves(&blocks, mst_root)?; + let bodies: Vec<_> = leaves + .iter() + .map(|(_, cid)| blocks.get(cid).cloned()) + .collect(); + Ok::<_, miette::Report>(leaves.into_iter().zip(bodies)) }) .await .into_diagnostic()??; diff --git a/src/bin/backfill_strategy_bench.rs b/src/bin/backfill_strategy_bench.rs index f533cf6..21fc7fe 100644 --- a/src/bin/backfill_strategy_bench.rs +++ b/src/bin/backfill_strategy_bench.rs @@ -1,6 +1,9 @@ #[path = "../car.rs"] #[allow(dead_code)] mod car; +#[path = "../mst.rs"] +#[allow(dead_code)] +mod mst; #[path = "../sparse_mst.rs"] mod sparse_mst; diff --git a/src/car.rs b/src/car.rs index 276c1cc..275e764 100644 --- a/src/car.rs +++ b/src/car.rs @@ -1,10 +1,8 @@ use std::collections::{BTreeMap, HashMap}; use std::io::Cursor; -use std::sync::Arc; use bytes::Bytes; use cid::Cid as IpldCid; -use jacquard_repo::{BlockStore, CommitData, RepoError, RepoErrorKind}; use miette::{IntoDiagnostic, Result}; /// blocks of a whole-repo CAR. hashed rather than ordered: the MST walk looks up every @@ -123,58 +121,6 @@ fn validate_block_cid(claimed_cid: &IpldCid, bytes: &[u8]) -> Result<()> { Ok(()) } -/// read-only [`BlockStore`] over the blocks of a parsed CAR, for walking an imported repo. -#[cfg_attr(not(feature = "indexer"), allow(dead_code))] -#[derive(Clone)] -pub(crate) struct CarBlockStore(Arc); - -#[cfg_attr(not(feature = "indexer"), allow(dead_code))] -impl CarBlockStore { - pub(crate) fn new(blocks: CarBlocks) -> Self { - Self(Arc::new(blocks)) - } - - pub(crate) fn len(&self) -> usize { - self.0.len() - } -} - -fn read_only_store() -> RepoError { - RepoError::new( - RepoErrorKind::Storage, - Some("CarBlockStore is read-only".into()), - ) -} - -impl BlockStore for CarBlockStore { - async fn get(&self, cid: &IpldCid) -> jacquard_repo::Result> { - Ok(self.0.get(cid).cloned()) - } - - async fn put(&self, _data: &[u8]) -> jacquard_repo::Result { - Err(read_only_store()) - } - - async fn has(&self, cid: &IpldCid) -> jacquard_repo::Result { - Ok(self.0.contains_key(cid)) - } - - async fn put_many( - &self, - _blocks: impl IntoIterator + Send, - ) -> jacquard_repo::Result<()> { - Err(read_only_store()) - } - - async fn get_many(&self, cids: &[IpldCid]) -> jacquard_repo::Result>> { - Ok(cids.iter().map(|cid| self.0.get(cid).cloned()).collect()) - } - - async fn apply_commit(&self, _commit: CommitData) -> jacquard_repo::Result<()> { - Err(read_only_store()) - } -} - #[cfg_attr(not(feature = "indexer"), allow(dead_code))] fn read_uvarint(data: &[u8], offset: &mut usize) -> Result> { if *offset == data.len() { @@ -302,61 +248,6 @@ mod tests { assert_eq!(blocks[&cid(1)].as_ref(), b"second"); } - #[tokio::test] - async fn car_block_store_serves_reads_and_rejects_writes() { - let present = cid(1); - let store = CarBlockStore::new(CarBlocks::from_iter([( - present, - Bytes::from_static(b"block"), - )])); - - assert_eq!( - store.get(&present).await.unwrap().as_deref(), - Some(&b"block"[..]) - ); - assert_eq!(store.get(&cid(2)).await.unwrap(), None); - assert!(store.has(&present).await.unwrap()); - assert!(!store.has(&cid(2)).await.unwrap()); - assert_eq!( - store.get_many(&[cid(2), present]).await.unwrap(), - vec![None, Some(Bytes::from_static(b"block"))] - ); - assert!(store.put(b"new").await.is_err()); - assert!(store.put_many([(cid(3), Bytes::new())]).await.is_err()); - assert_eq!(store.len(), 1); - } - - #[tokio::test] - async fn car_block_store_walks_a_parsed_repo_like_memory_store() { - use jacquard_repo::{MemoryBlockStore, Mst}; - - let mut mst = Mst::new(Arc::new(MemoryBlockStore::new())); - let mut records = Vec::new(); - for i in 0..500 { - let body = format!("record {i}").into_bytes(); - let record_cid = jacquard_repo::mst::util::compute_cid(&body).unwrap(); - let key = format!("app.bsky.feed.post/{i:013}"); - mst = mst.add(&key, record_cid).await.unwrap(); - records.push((record_cid, body)); - } - let (root, nodes) = mst.collect_blocks().await.unwrap(); - let blocks: Vec<(IpldCid, &[u8])> = nodes - .iter() - .map(|(cid, bytes)| (*cid, bytes.as_ref())) - .chain(records.iter().map(|(cid, body)| (*cid, body.as_slice()))) - .collect(); - let parsed = parse_car(car_with_blocks(&blocks).await, true).unwrap(); - - let loaded = Mst::load(Arc::new(CarBlockStore::new(parsed.blocks)), root, None); - let expected = mst.leaves().await.unwrap(); - assert_eq!(loaded.leaves().await.unwrap(), expected); - // backfill walks with leaves_sequential; pin that it matches the parallel walk - assert_eq!(loaded.leaves_sequential().await.unwrap(), expected); - let leaf_cids: Vec<_> = expected.iter().map(|(_, cid)| *cid).collect(); - let bodies = loaded.storage().get_many(&leaf_cids).await.unwrap(); - assert!(bodies.iter().all(Option::is_some)); - } - #[test] fn validate_block_cids_rejects_mismatched_cid() { let blocks = BTreeMap::from([(cid(1), Bytes::from_static(b"forged"))]); diff --git a/src/lib.rs b/src/lib.rs index 49991ee..1fb3696 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -45,6 +45,8 @@ pub(crate) mod db; pub(crate) mod ingest; #[cfg(feature = "jetstream")] pub(crate) mod jetstream; +#[cfg(feature = "indexer")] +pub(crate) mod mst; pub(crate) mod net; #[cfg(feature = "indexer")] pub(crate) mod ops; diff --git a/src/mst.rs b/src/mst.rs new file mode 100644 index 0000000..2a0c93a --- /dev/null +++ b/src/mst.rs @@ -0,0 +1,309 @@ +//! mst node decoding, shared by the full-repo walk here and the sparse scanner in +//! `sparse_mst`. + +use std::collections::HashSet; + +use cid::Cid as IpldCid; +use jacquard_repo::mst::NodeData; +use miette::{IntoDiagnostic, Result, WrapErr}; +use smol_str::SmolStr; + +use crate::car::CarBlocks; + +#[derive(Debug, Clone)] +pub(crate) enum FlatEntry { + Tree { cid: IpldCid }, + Leaf { key: SmolStr, cid: IpldCid }, +} + +/// a node's subtrees and leaves in key order, with full keys rebuilt from prefix compression. +pub(crate) fn decode_node(bytes: &[u8]) -> Result> { + let node: NodeData = serde_ipld_dagcbor::from_slice(bytes).into_diagnostic()?; + decode_node_entries(&node) +} + +fn decode_node_entries(node: &NodeData) -> Result> { + let mut entries = Vec::with_capacity(node.entries.len() * 2 + 1); + if let Some(cid) = node.left { + entries.push(FlatEntry::Tree { cid }); + } + + // each key is stored as a prefix length into the previous key plus a suffix + let mut key = String::new(); + for entry in &node.entries { + let suffix = std::str::from_utf8(&entry.key_suffix) + .into_diagnostic() + .wrap_err("invalid utf8 in mst key suffix")?; + let prefix_len = usize::from(entry.prefix_len); + if !key.is_char_boundary(prefix_len) { + return Err(miette::miette!( + "invalid mst key prefix length {prefix_len}" + )); + } + key.truncate(prefix_len); + key.push_str(suffix); + + entries.push(FlatEntry::Leaf { + key: SmolStr::new(&key), + cid: entry.value, + }); + if let Some(cid) = entry.tree { + entries.push(FlatEntry::Tree { cid }); + } + } + + Ok(entries) +} + +/// every leaf of the mst rooted at `root`, in key order. +/// +/// the tree comes from an untrusted car, so anything a valid mst cannot contain is an error: +/// a missing node, a node reached twice (which also ends cid cycles), or keys that are not +/// strictly increasing. the walk keeps its own stack, so tree depth cannot overflow the +/// thread's. +pub(crate) fn leaves(blocks: &CarBlocks, root: IpldCid) -> Result> { + let mut found: Vec<(SmolStr, IpldCid)> = Vec::new(); + let mut visited = HashSet::with_hasher(ahash::RandomState::new()); + let mut stack = vec![FlatEntry::Tree { cid: root }]; + + while let Some(entry) = stack.pop() { + match entry { + FlatEntry::Leaf { key, cid } => { + let out_of_order = found.last().filter(|(previous, _)| *previous >= key); + if let Some((previous, _)) = out_of_order { + miette::bail!("mst key {key} does not sort after {previous}"); + } + found.push((key, cid)); + } + FlatEntry::Tree { cid } => { + if !visited.insert(cid) { + miette::bail!("mst node {cid} is reachable more than once"); + } + let bytes = blocks + .get(&cid) + .ok_or_else(|| miette::miette!("mst node {cid} missing from car"))?; + let entries = decode_node(bytes) + .wrap_err_with(|| format!("failed to decode mst node {cid}"))?; + stack.extend(entries.into_iter().rev()); + } + } + } + + Ok(found) +} + +#[cfg(test)] +pub(crate) mod tests { + use super::*; + use cid::multihash::Multihash; + use jacquard_common::types::crypto::{DAG_CBOR, SHA2_256}; + use jacquard_repo::mst::TreeEntry; + use jacquard_repo::mst::util::compute_cid; + + pub(crate) fn cid(byte: u8) -> IpldCid { + let mh = Multihash::<64>::wrap(SHA2_256, &[byte; 32]).unwrap(); + IpldCid::new_v1(DAG_CBOR, mh) + } + + /// builds a node from full keys, prefix-compressing them like the wire format. + pub(crate) fn node( + entries: Vec<(&str, IpldCid, Option)>, + left: Option, + ) -> NodeData { + let mut last = ""; + let entries = entries + .into_iter() + .map(|(key, value, tree)| { + let prefix_len = common_prefix_len(last, key); + last = key; + TreeEntry { + key_suffix: key.as_bytes()[prefix_len..].to_vec().into(), + prefix_len: prefix_len as u8, + tree, + value, + } + }) + .collect(); + NodeData { left, entries } + } + + /// shared prefix in bytes, ending on a char boundary. + fn common_prefix_len(a: &str, b: &str) -> usize { + a.chars() + .zip(b.chars()) + .take_while(|(a, b)| a == b) + .map(|(a, _)| a.len_utf8()) + .sum() + } + + fn blocks(nodes: Vec<(IpldCid, NodeData)>) -> CarBlocks { + nodes + .into_iter() + .map(|(cid, node)| (cid, serde_ipld_dagcbor::to_vec(&node).unwrap().into())) + .collect() + } + + fn keys(leaves: &[(SmolStr, IpldCid)]) -> Vec<&str> { + leaves.iter().map(|(key, _)| key.as_str()).collect() + } + + #[tokio::test] + async fn leaves_match_jacquard_walk() { + use jacquard_repo::{MemoryBlockStore, Mst}; + + let mut mst = Mst::new(std::sync::Arc::new(MemoryBlockStore::new())); + for i in 0..2000 { + let collection = [ + "app.bsky.feed.like", + "app.bsky.feed.post", + "sh.tangled.repo", + ][i % 3]; + let key = format!("{collection}/{i:013}"); + mst = mst.add(&key, cid((i % 251) as u8)).await.unwrap(); + } + let (root, nodes) = mst.collect_blocks().await.unwrap(); + assert!(nodes.len() > 1, "test tree should span several nodes"); + + let ours = leaves(&nodes.into_iter().collect(), root).unwrap(); + assert_eq!(ours, mst.leaves().await.unwrap()); + } + + #[test] + fn walks_an_empty_tree() { + let root = cid(1); + assert!( + leaves(&blocks(vec![(root, node(vec![], None))]), root) + .unwrap() + .is_empty() + ); + } + + #[test] + fn rejects_a_missing_node() { + let (root, missing) = (cid(1), cid(2)); + let tree = blocks(vec![( + root, + node(vec![("a/1", cid(9), Some(missing))], None), + )]); + + let err = leaves(&tree, root).unwrap_err(); + assert!(err.to_string().contains("missing from car"), "{err}"); + } + + #[test] + fn rejects_a_subtree_reachable_twice() { + let (root, shared) = (cid(1), cid(2)); + let tree = blocks(vec![ + ( + root, + node( + vec![("a/1", cid(9), Some(shared)), ("c/1", cid(9), Some(shared))], + None, + ), + ), + (shared, node(vec![("b/1", cid(9), None)], None)), + ]); + + let err = leaves(&tree, root).unwrap_err(); + assert!( + err.to_string().contains("reachable more than once"), + "{err}" + ); + } + + #[test] + fn stops_on_a_node_that_is_its_own_left_subtree() { + // a car that is not cid-verified can claim any cid for a block + let root = cid(1); + let tree = blocks(vec![(root, node(vec![("a/1", cid(9), None)], Some(root)))]); + + let err = leaves(&tree, root).unwrap_err(); + assert!( + err.to_string().contains("reachable more than once"), + "{err}" + ); + } + + #[test] + fn rejects_keys_out_of_order() { + let root = cid(1); + let tree = blocks(vec![( + root, + node(vec![("b/1", cid(9), None), ("a/1", cid(9), None)], None), + )]); + + let err = leaves(&tree, root).unwrap_err(); + assert!(err.to_string().contains("does not sort after"), "{err}"); + } + + #[test] + fn rejects_a_duplicate_key_across_subtrees() { + let (root, right) = (cid(1), cid(2)); + let tree = blocks(vec![ + (root, node(vec![("a/1", cid(9), Some(right))], None)), + (right, node(vec![("a/1", cid(8), None)], None)), + ]); + + let err = leaves(&tree, root).unwrap_err(); + assert!(err.to_string().contains("does not sort after"), "{err}"); + } + + #[test] + fn rejects_a_prefix_that_splits_a_char() { + let root = cid(1); + let mut bad = node(vec![("é/1", cid(9), None), ("é/2", cid(9), None)], None); + bad.entries[1].prefix_len = 1; + + let err = leaves(&blocks(vec![(root, bad)]), root).unwrap_err(); + assert!( + format!("{err:?}").contains("invalid mst key prefix length 1"), + "{err:?}" + ); + } + + #[test] + fn rebuilds_prefix_compressed_keys() { + let root = cid(1); + let tree = blocks(vec![( + root, + node( + vec![ + ("app.bsky.feed.like/1", cid(9), None), + ("app.bsky.feed.post/1", cid(9), None), + ("app.bsky.feed.post/2", cid(9), None), + // both start with byte 0xc3, so they share no whole-char prefix + ("è/1", cid(9), None), + ("é/1", cid(9), None), + ], + None, + ), + )]); + + assert_eq!( + keys(&leaves(&tree, root).unwrap()), + [ + "app.bsky.feed.like/1", + "app.bsky.feed.post/1", + "app.bsky.feed.post/2", + "è/1", + "é/1" + ] + ); + } + + #[test] + fn walks_a_tree_deeper_than_the_thread_stack() { + // a recursive walk would overflow the stack long before this depth + let depth = 100_000; + let bottom = node(vec![("a/1", cid(9), None)], None); + let mut child = compute_cid(&serde_ipld_dagcbor::to_vec(&bottom).unwrap()).unwrap(); + let mut nodes = vec![(child, bottom)]; + for _ in 0..depth { + let parent = node(vec![], Some(child)); + child = compute_cid(&serde_ipld_dagcbor::to_vec(&parent).unwrap()).unwrap(); + nodes.push((child, parent)); + } + + assert_eq!(keys(&leaves(&blocks(nodes), child).unwrap()), ["a/1"]); + } +} diff --git a/src/sparse_mst.rs b/src/sparse_mst.rs index 34e6979..5d98361 100644 --- a/src/sparse_mst.rs +++ b/src/sparse_mst.rs @@ -2,11 +2,12 @@ use std::collections::BTreeMap; use bytes::Bytes; use cid::Cid as IpldCid; -use jacquard_repo::mst::NodeData; use jacquard_repo::mst::util::layer_for_key; -use miette::{IntoDiagnostic, Result, WrapErr}; +use miette::{Result, WrapErr}; use smol_str::SmolStr; +use crate::mst::{FlatEntry, decode_node}; + #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct KeyRange { start: SmolStr, @@ -85,15 +86,11 @@ fn prefix_upper_bound(prefix: &str) -> Option { } pub(crate) fn mst_node_layer(bytes: &[u8]) -> Result> { - let node: NodeData = serde_ipld_dagcbor::from_slice(bytes) - .into_diagnostic() - .wrap_err("failed to decode mst node for layer estimate")?; - decode_node_entries(&node).map(|entries| { - entries.into_iter().find_map(|entry| match entry { - FlatEntry::Leaf { key, .. } => Some(layer_for_key(key.as_str())), - FlatEntry::Tree { .. } => None, - }) - }) + let entries = decode_node(bytes).wrap_err("failed to decode mst node for layer estimate")?; + Ok(entries.into_iter().find_map(|entry| match entry { + FlatEntry::Leaf { key, .. } => Some(layer_for_key(key.as_str())), + FlatEntry::Tree { .. } => None, + })) } #[derive(Debug, Clone)] @@ -239,10 +236,8 @@ impl SparseScanner { } seen.push(bounds.clone()); - let node: NodeData = serde_ipld_dagcbor::from_slice(bytes) - .into_diagnostic() - .wrap_err_with(|| format!("failed to decode mst node {cid}"))?; - let entries = decode_node_entries(&node)?; + let entries = + decode_node(bytes).wrap_err_with(|| format!("failed to decode mst node {cid}"))?; for idx in 0..entries.len() { match &entries[idx] { @@ -273,43 +268,6 @@ struct ScanStats { node_bytes_seen: usize, } -#[derive(Debug, Clone)] -enum FlatEntry { - Tree { cid: IpldCid }, - Leaf { key: SmolStr, cid: IpldCid }, -} - -fn decode_node_entries(node: &NodeData) -> Result> { - let mut entries = Vec::new(); - if let Some(cid) = node.left { - entries.push(FlatEntry::Tree { cid }); - } - - let mut last_key = String::new(); - for entry in &node.entries { - let suffix = std::str::from_utf8(&entry.key_suffix) - .into_diagnostic() - .wrap_err("invalid utf8 in mst key suffix")?; - let prefix_len = entry.prefix_len as usize; - let prefix = last_key - .get(..prefix_len) - .ok_or_else(|| miette::miette!("invalid mst key prefix length {prefix_len}"))?; - let key = SmolStr::new(format!("{prefix}{suffix}")); - - entries.push(FlatEntry::Leaf { - key: key.clone(), - cid: entry.value, - }); - last_key = key.to_string(); - - if let Some(cid) = entry.tree { - entries.push(FlatEntry::Tree { cid }); - } - } - - Ok(entries) -} - fn previous_leaf(entries: &[FlatEntry], idx: usize) -> Option<&str> { entries[..idx].iter().rev().find_map(|entry| match entry { FlatEntry::Leaf { key, .. } => Some(key.as_str()), @@ -327,39 +285,8 @@ fn next_leaf(entries: &[FlatEntry], idx: usize) -> Option<&str> { #[cfg(test)] mod tests { use super::*; - use cid::Cid; - use cid::multihash::Multihash; - use jacquard_common::types::crypto::{DAG_CBOR, SHA2_256}; - use jacquard_repo::mst::TreeEntry; - - fn cid(byte: u8) -> IpldCid { - let hash = [byte; 32]; - let mh = Multihash::<64>::wrap(SHA2_256, &hash).unwrap(); - Cid::new_v1(DAG_CBOR, mh) - } - - fn node(entries: Vec<(&str, IpldCid, Option)>, left: Option) -> NodeData { - let mut last = String::new(); - let entries = entries - .into_iter() - .map(|(key, value, tree)| { - let prefix_len = common_prefix_len(&last, key); - let suffix = key[prefix_len..].as_bytes().to_vec(); - last = key.to_string(); - TreeEntry { - key_suffix: suffix.into(), - prefix_len: prefix_len as u8, - tree, - value, - } - }) - .collect(); - NodeData { left, entries } - } - - fn common_prefix_len(a: &str, b: &str) -> usize { - a.chars().zip(b.chars()).take_while(|(a, b)| a == b).count() - } + use crate::mst::tests::{cid, node}; + use jacquard_repo::mst::NodeData; #[test] fn builds_exact_and_wildcard_ranges() {