diff --git a/docs/configuration.md b/docs/configuration.md --- a/docs/configuration.md +++ b/docs/configuration.md @@ -56,6 +56,7 @@ | variable | default | description | | :--- | :--- | :--- | | `BACKFILL_CONCURRENCY_LIMIT` | `16` (`64` full network) | maximum number of concurrent backfill tasks | +| `BACKFILL_STRATEGY` | `full` | backfill strategy: `full` keeps the existing `getRepo` path, `sparse-filter` attempts authenticated sparse collection backfill before falling back to full, `auto` uses sparse only when explicit collection filters are configured | | `REPO_FETCH_TIMEOUT` | `5min` | timeout for fetching a repository | | `VERIFY_SIGNATURES` | `full` | signature verification level: `full`, `backfill-only`, or `none` | | `PLC_URL` | `https://plc.wtf`, `https://plc.directory` (full network) | base URL(s) of the PLC directory, comma-separated | diff --git a/src/car.rs b/src/car.rs new file mode 100644 --- /dev/null +++ b/src/car.rs @@ -0,0 +1,93 @@ +use std::collections::BTreeMap; +use std::io::Cursor; + +use bytes::Bytes; +use cid::Cid as IpldCid; +use miette::{IntoDiagnostic, Result}; + +pub(crate) async fn parse_car_blocks(data: &[u8]) -> Result> { + let mut offset = 0; + let Some(header_len) = read_uvarint(data, &mut offset)? else { + return Err(miette::miette!("empty CAR file")); + }; + let header_end = offset + .checked_add(header_len) + .ok_or_else(|| miette::miette!("CAR header length overflow"))?; + if header_end > data.len() { + return Err(miette::miette!("truncated CAR header")); + } + offset = header_end; + + let mut blocks = BTreeMap::new(); + while let Some(section_len) = read_uvarint(data, &mut offset)? { + let section_end = offset + .checked_add(section_len) + .ok_or_else(|| miette::miette!("CAR block length overflow"))?; + if section_end > data.len() { + return Err(miette::miette!("truncated CAR block")); + } + + let section = &data[offset..section_end]; + offset = section_end; + + let mut cursor = Cursor::new(section); + let cid = IpldCid::read_bytes(&mut cursor).into_diagnostic()?; + let block_start = cursor.position() as usize; + if block_start >= section.len() { + return Err(miette::miette!("CAR block has no payload for {cid}")); + } + blocks.insert(cid, Bytes::copy_from_slice(§ion[block_start..])); + } + + Ok(blocks) +} + +fn read_uvarint(data: &[u8], offset: &mut usize) -> Result> { + if *offset == data.len() { + return Ok(None); + } + + let mut value = 0u64; + for shift in (0..64).step_by(7) { + if *offset >= data.len() { + return Err(miette::miette!("truncated uvarint in CAR file")); + } + + let byte = data[*offset]; + *offset += 1; + value |= u64::from(byte & 0x7f) << shift; + + if byte & 0x80 == 0 { + let len = usize::try_from(value).into_diagnostic()?; + return Ok(Some(len)); + } + } + + Err(miette::miette!("uvarint overflow in CAR file")) +} + +#[cfg(test)] +mod tests { + use super::*; + use cid::multihash::Multihash; + use jacquard_common::types::crypto::{DAG_CBOR, SHA2_256}; + + fn cid(byte: u8) -> IpldCid { + let hash = [byte; 32]; + let mh = Multihash::<64>::wrap(SHA2_256, &hash).unwrap(); + IpldCid::new_v1(DAG_CBOR, mh) + } + + #[tokio::test] + async fn parses_rootless_block_car() { + let cid = cid(1); + let mut buf = Vec::new(); + let header = iroh_car::CarHeader::new_v1(Vec::new()); + let mut writer = iroh_car::CarWriter::new(header, &mut buf); + writer.write(cid, b"block".to_vec()).await.unwrap(); + writer.finish().await.unwrap(); + + let blocks = parse_car_blocks(&buf).await.unwrap(); + assert_eq!(blocks.get(&cid).unwrap().as_ref(), b"block"); + } +} diff --git a/src/config.rs b/src/config.rs --- a/src/config.rs +++ b/src/config.rs @@ -284,6 +284,39 @@ } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BackfillStrategy { + /// always fetch full repo cars with `com.atproto.sync.getRepo`. + Full, + /// use sparse collection backfill when possible, falling back to full repo cars. + SparseFilter, + /// choose sparse collection backfill only when the configured filter supports it. + Auto, +} + +impl FromStr for BackfillStrategy { + type Err = miette::Error; + + fn from_str(s: &str) -> Result { + match s { + "full" => Ok(Self::Full), + "sparse-filter" => Ok(Self::SparseFilter), + "auto" => Ok(Self::Auto), + _ => Err(miette::miette!("invalid backfill strategy")), + } + } +} + +impl fmt::Display for BackfillStrategy { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Full => write!(f, "full"), + Self::SparseFilter => write!(f, "sparse-filter"), + Self::Auto => write!(f, "auto"), + } + } +} + #[derive(Debug, Clone)] pub struct Config { /// path to the database folder. set via `HYDRANT_DATABASE_PATH`. @@ -324,6 +357,11 @@ /// maximum number of concurrent backfill tasks. /// set via `HYDRANT_BACKFILL_CONCURRENCY_LIMIT`. pub backfill_concurrency_limit: usize, + /// backfill strategy. `full` preserves existing full-repo backfill behavior. + /// `sparse-filter` attempts authenticated sparse collection backfill first and falls back + /// to full repo backfill. `auto` uses sparse only when explicit collection filters exist. + /// set via `HYDRANT_BACKFILL_STRATEGY`. + pub backfill_strategy: BackfillStrategy, /// whether to run the network crawler. `None` defers to the default for the current mode. /// set via `HYDRANT_ENABLE_CRAWLER`. @@ -518,6 +556,7 @@ cursor_save_interval: Duration::from_secs(3), repo_fetch_timeout: Duration::from_secs(300), backfill_concurrency_limit: 16, + backfill_strategy: BackfillStrategy::Full, enable_crawler: None, crawler_max_pending_repos: 2000, crawler_resume_pending_repos: 1000, @@ -660,6 +699,7 @@ "BACKFILL_CONCURRENCY_LIMIT", defaults.backfill_concurrency_limit ); + let backfill_strategy = cfg!("BACKFILL_STRATEGY", defaults.backfill_strategy); let firehose_workers = cfg!("FIREHOSE_WORKERS", defaults.firehose_workers); let firehose_max_failures = cfg!("FIREHOSE_MAX_FAILURES", defaults.firehose_max_failures); @@ -848,6 +888,7 @@ cursor_save_interval, repo_fetch_timeout, backfill_concurrency_limit, + backfill_strategy, enable_crawler, crawler_max_pending_repos, crawler_resume_pending_repos, @@ -912,6 +953,7 @@ config_line!(f, "full network indexing", self.full_network)?; config_line!(f, "verify signatures", self.verify_signatures)?; config_line!(f, "backfill concurrency", self.backfill_concurrency_limit)?; + config_line!(f, "backfill strategy", self.backfill_strategy)?; config_line!(f, "identity cache size", self.identity_cache_size)?; config_line!( f, diff --git a/src/lib.rs b/src/lib.rs --- a/src/lib.rs +++ b/src/lib.rs @@ -40,6 +40,8 @@ #[cfg(feature = "backlinks")] pub mod backlinks; #[cfg(feature = "indexer")] +pub(crate) mod car; +#[cfg(feature = "indexer")] pub(crate) mod crawler; pub(crate) mod db; pub(crate) mod ingest; @@ -49,6 +51,8 @@ pub(crate) mod ops; pub(crate) mod patch; pub mod resolver; +#[cfg(feature = "indexer")] +pub(crate) mod sparse_mst; pub(crate) mod state; pub(crate) mod util; diff --git a/src/sparse_mst.rs b/src/sparse_mst.rs new file mode 100644 --- /dev/null +++ b/src/sparse_mst.rs @@ -0,0 +1,318 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use bytes::Bytes; +use cid::Cid as IpldCid; +use jacquard_repo::mst::NodeData; +use miette::{IntoDiagnostic, Result, WrapErr}; +use smol_str::SmolStr; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct KeyRange { + start: SmolStr, + end: Option, +} + +impl KeyRange { + pub(crate) fn contains(&self, key: &str) -> bool { + key >= self.start.as_str() && self.end.as_deref().map(|end| key < end).unwrap_or(true) + } + + fn intersects_subtree( + &self, + lower_exclusive: Option<&str>, + upper_exclusive: Option<&str>, + ) -> bool { + if let Some(upper) = upper_exclusive { + if self.start.as_str() >= upper { + return false; + } + } + + if let (Some(end), Some(lower)) = (self.end.as_deref(), lower_exclusive) { + if end <= lower { + return false; + } + } + + true + } +} + +pub(crate) fn sparse_ranges(patterns: &[SmolStr]) -> Vec { + patterns + .iter() + .filter_map(|pattern| { + let prefix = pattern + .strip_suffix(".*") + .map(|prefix| SmolStr::new(format!("{prefix}."))) + .unwrap_or_else(|| SmolStr::new(format!("{pattern}/"))); + Some(KeyRange { + end: prefix_upper_bound(prefix.as_str()), + start: prefix, + }) + }) + .collect() +} + +pub(crate) fn sparse_probe_collection(patterns: &[SmolStr]) -> Option { + patterns.iter().find_map(|pattern| { + pattern + .strip_suffix(".*") + .map(|prefix| format!("{prefix}.probe")) + .unwrap_or_else(|| pattern.to_string()) + .parse::() + .ok() + .map(|_| { + pattern + .strip_suffix(".*") + .map(|prefix| SmolStr::new(format!("{prefix}.probe"))) + .unwrap_or_else(|| pattern.clone()) + }) + }) +} + +fn prefix_upper_bound(prefix: &str) -> Option { + let mut bytes = prefix.as_bytes().to_vec(); + for i in (0..bytes.len()).rev() { + if bytes[i] < u8::MAX { + bytes[i] += 1; + bytes.truncate(i + 1); + return String::from_utf8(bytes).ok().map(SmolStr::new); + } + } + None +} + +#[derive(Debug, Clone)] +pub(crate) struct SparseScanOutput { + pub(crate) leaves: Vec<(SmolStr, IpldCid)>, + pub(crate) node_blocks_seen: usize, + pub(crate) node_bytes_seen: usize, +} + +pub(crate) struct SparseScanner { + ranges: Vec, + blocks: BTreeMap, +} + +impl SparseScanner { + pub(crate) fn new(ranges: Vec, blocks: BTreeMap) -> Self { + Self { ranges, blocks } + } + + pub(crate) fn insert_blocks(&mut self, blocks: BTreeMap) { + self.blocks.extend(blocks); + } + + pub(crate) fn scan(&self, root: IpldCid) -> Result>> { + let mut visited = BTreeSet::new(); + let mut missing = BTreeSet::new(); + let mut leaves = Vec::new(); + let mut stats = ScanStats::default(); + + self.scan_node(root, &mut visited, &mut missing, &mut leaves, &mut stats)?; + + if missing.is_empty() { + Ok(Ok(SparseScanOutput { + leaves, + node_blocks_seen: stats.node_blocks_seen, + node_bytes_seen: stats.node_bytes_seen, + })) + } else { + Ok(Err(missing.into_iter().collect())) + } + } + + pub(crate) fn take_blocks(self) -> BTreeMap { + self.blocks + } + + fn scan_node( + &self, + cid: IpldCid, + visited: &mut BTreeSet, + missing: &mut BTreeSet, + leaves: &mut Vec<(SmolStr, IpldCid)>, + stats: &mut ScanStats, + ) -> Result<()> { + if !visited.insert(cid) { + return Ok(()); + } + + let Some(bytes) = self.blocks.get(&cid) else { + missing.insert(cid); + return Ok(()); + }; + + stats.node_blocks_seen += 1; + stats.node_bytes_seen += bytes.len(); + + 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)?; + + for idx in 0..entries.len() { + match &entries[idx] { + FlatEntry::Leaf { key, cid } => { + if self.ranges.iter().any(|range| range.contains(key)) { + leaves.push((key.clone(), *cid)); + } + } + FlatEntry::Tree { cid } => { + let lower = previous_leaf(&entries, idx); + let upper = next_leaf(&entries, idx); + if self + .ranges + .iter() + .any(|range| range.intersects_subtree(lower, upper)) + { + self.scan_node(*cid, visited, missing, leaves, stats)?; + } + } + } + } + + Ok(()) + } +} + +#[derive(Default)] +struct ScanStats { + node_blocks_seen: usize, + 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()), + FlatEntry::Tree { .. } => None, + }) +} + +fn next_leaf(entries: &[FlatEntry], idx: usize) -> Option<&str> { + entries[idx + 1..].iter().find_map(|entry| match entry { + FlatEntry::Leaf { key, .. } => Some(key.as_str()), + FlatEntry::Tree { .. } => None, + }) +} + +#[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() + } + + #[test] + fn builds_exact_and_wildcard_ranges() { + let ranges = sparse_ranges(&[SmolStr::new("sh.tangled.repo"), SmolStr::new("app.bsky.*")]); + + assert!(ranges[0].contains("sh.tangled.repo/abc")); + assert!(!ranges[0].contains("sh.tangled.repo.comment/abc")); + assert!(ranges[1].contains("app.bsky.feed.post/abc")); + assert!(!ranges[1].contains("app.bskyfoo.feed.post/abc")); + assert!(!ranges[1].contains("app.csky.feed.post/abc")); + } + + #[test] + fn synthesizes_probe_collection_for_wildcard() { + assert_eq!( + sparse_probe_collection(&[SmolStr::new("sh.tangled.*")]).as_deref(), + Some("sh.tangled.probe") + ); + assert_eq!( + sparse_probe_collection(&[SmolStr::new("sh.tangled.repo")]).as_deref(), + Some("sh.tangled.repo") + ); + } + + #[test] + fn reports_missing_intersecting_subtrees_only() { + let wanted = sparse_ranges(&[SmolStr::new("sh.tangled.*")]); + let left = cid(1); + let middle = cid(2); + let right = cid(3); + let root = cid(4); + let root_node = node( + vec![ + ("app.bsky.feed.post/1", cid(5), Some(middle)), + ("zz.example.record/1", cid(6), Some(right)), + ], + Some(left), + ); + let root_bytes = serde_ipld_dagcbor::to_vec(&root_node).unwrap(); + let scanner = SparseScanner::new(wanted, BTreeMap::from([(root, root_bytes.into())])); + + let missing = scanner.scan(root).unwrap().unwrap_err(); + + assert_eq!(missing, vec![middle]); + } +} diff --git a/src/backfill/mod.rs b/src/backfill/mod.rs --- a/src/backfill/mod.rs +++ b/src/backfill/mod.rs @@ -1,24 +1,29 @@ +use crate::config::BackfillStrategy; use crate::db::types::{DbAction, DbRkey, TrimmedDid}; use crate::db::{self, CountDeltas, Db, keys, ser_repo_state}; use crate::filter::FilterMode; use crate::ops; use crate::resolver::ResolverError; +use crate::sparse_mst::{SparseScanner, sparse_probe_collection, sparse_ranges}; use crate::state::AppState; use crate::types::{Commit, GaugeState, RepoState, RepoStatus, ResyncErrorKind, ResyncState}; use fjall::Slice; +use jacquard_api::com_atproto::sync::get_blocks::{GetBlocks, GetBlocksError}; +use jacquard_api::com_atproto::sync::get_record::{GetRecord, GetRecordError}; use jacquard_api::com_atproto::sync::get_repo::{GetRepo, GetRepoError}; use jacquard_common::IntoStatic; use jacquard_common::error::{ClientError, ClientErrorKind}; -use jacquard_common::types::cid::Cid; +use jacquard_common::types::cid::{Cid as AtCid, IpldCid}; use jacquard_common::types::did::Did; +use jacquard_common::types::string::{Nsid, RecordKey}; use jacquard_common::xrpc::{XrpcError, XrpcExt}; use jacquard_repo::mst::Mst; use jacquard_repo::{BlockStore, MemoryBlockStore}; use miette::{Diagnostic, IntoDiagnostic, Result}; use reqwest::StatusCode; use smol_str::{SmolStr, ToSmolStr}; -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -43,6 +48,7 @@ http: reqwest::Client, semaphore: Arc, verify_signatures: bool, + strategy: BackfillStrategy, in_flight: Arc>>, enabled: tokio::sync::watch::Receiver, } @@ -54,6 +60,7 @@ timeout: Duration, concurrency_limit: usize, verify_signatures: bool, + strategy: BackfillStrategy, enabled: tokio::sync::watch::Receiver, ) -> Self { Self { @@ -68,6 +75,7 @@ .expect("failed to build http client"), semaphore: Arc::new(Semaphore::new(concurrency_limit)), verify_signatures, + strategy, in_flight: Arc::new(scc::HashSet::new()), enabled, } @@ -143,13 +151,15 @@ let did = did.clone(); let buffer_tx = self.buffer_tx.clone(); let verify = self.verify_signatures; + let strategy = self.strategy; let span = tracing::info_span!("backfill", did = %did); tokio::spawn( async move { let _guard = guard; let res = - did_task(&state, http, buffer_tx, &did, key, permit, verify).await; + did_task(&state, http, buffer_tx, &did, key, permit, verify, strategy) + .await; if let Err(e) = res { error!(err = %e, "process failed"); @@ -184,10 +194,11 @@ pending_key: Slice, _permit: tokio::sync::OwnedSemaphorePermit, verify_signatures: bool, + strategy: BackfillStrategy, ) -> Result<(), BackfillError> { let db = &state.db; - match process_did(state, &http, did, verify_signatures).await { + match process_did(state, &http, did, verify_signatures, strategy).await { Ok(Some(_repo_state)) => { let did_key = keys::repo_key(did); @@ -384,11 +395,478 @@ } } +#[derive(Debug)] +struct SparseBackfillSuccess { + state: RepoState<'static>, + records: usize, + node_blocks: usize, + node_bytes: usize, +} + +#[derive(Debug)] +enum SparseBackfillResult { + Imported(SparseBackfillSuccess), + Discarded, + Skipped, +} + +const SPARSE_GET_BLOCKS_CHUNK: usize = 64; +const SPARSE_MAX_SCAN_ROUNDS: usize = 256; + +async fn process_did_sparse( + app_state: &Arc, + http: &reqwest::Client, + did: &Did<'static>, + pds: &url::Url, + state: RepoState<'static>, + verify_signatures: bool, +) -> Result { + let filter = app_state.filter.load(); + let Some(probe_collection) = sparse_probe_collection(&filter.collections) else { + return Ok(SparseBackfillResult::Skipped); + }; + let ranges = sparse_ranges(&filter.collections); + if ranges.is_empty() { + return Ok(SparseBackfillResult::Skipped); + } + + let probe_collection = Nsid::new_owned(probe_collection.as_str()).into_diagnostic()?; + let probe_rkey = RecordKey::any_static("-").into_diagnostic()?; + let req = GetRecord::new() + .did(did.clone()) + .collection(probe_collection) + .rkey(probe_rkey) + .build(); + + let resp = http.xrpc(url_to_fluent_uri(pds)).send(&req).await?; + let proof = match resp.into_output() { + Ok(o) => o, + Err(XrpcError::Xrpc(GetRecordError::RecordNotFound(_))) => { + return Ok(SparseBackfillResult::Skipped); + } + Err(XrpcError::Xrpc( + GetRecordError::RepoNotFound(_) + | GetRecordError::RepoTakendown(_) + | GetRecordError::RepoSuspended(_) + | GetRecordError::RepoDeactivated(_), + )) => return Ok(SparseBackfillResult::Skipped), + Err(e) => Err(e).into_diagnostic()?, + }; + + let parsed = jacquard_repo::car::reader::parse_car_bytes(&proof.body) + .await + .into_diagnostic()?; + let root_bytes = parsed + .blocks + .get(&parsed.root) + .ok_or_else(|| miette::miette!("root block missing from sparse proof CAR"))?; + let root_commit = jacquard_repo::commit::Commit::from_cbor(root_bytes).into_diagnostic()?; + + if verify_signatures { + let pubkey = app_state.resolver.resolve_signing_key(did).await?; + root_commit + .verify(&pubkey) + .map_err(|e| miette::miette!("signature verification failed for {did}: {e}"))?; + } + + let root_cid = root_commit.data; + let root_commit = Commit::from(root_commit); + let mut scanner = SparseScanner::new(ranges, parsed.blocks); + let mut scan_rounds = 0; + let scan = loop { + match scanner.scan(root_cid)? { + Ok(scan) => break scan, + Err(missing) => { + scan_rounds += 1; + if scan_rounds > SPARSE_MAX_SCAN_ROUNDS { + return Err( + miette::miette!("sparse mst scan exceeded fetch round limit").into(), + ); + } + if missing.is_empty() { + return Ok(SparseBackfillResult::Skipped); + } + if missing.len() > SPARSE_MAX_SCAN_ROUNDS * SPARSE_GET_BLOCKS_CHUNK { + return Err( + miette::miette!("sparse mst scan exceeded missing block limit").into(), + ); + } + let blocks = fetch_blocks(http, pds, did, &missing).await?; + if missing.iter().all(|cid| !blocks.contains_key(cid)) { + return Ok(SparseBackfillResult::Skipped); + } + scanner.insert_blocks(blocks); + } + } + }; + + let mut blocks = scanner.take_blocks(); + let missing_records: Vec = scan + .leaves + .iter() + .filter_map(|(_, cid)| (!blocks.contains_key(cid)).then_some(*cid)) + .collect(); + blocks.extend(fetch_blocks(http, pds, did, &missing_records).await?); + + if let Some((_, missing)) = scan + .leaves + .iter() + .find(|(_, cid)| !blocks.contains_key(cid)) + { + return Err( + miette::miette!("sparse record block missing after getBlocks: {missing}").into(), + ); + } + + let result = + persist_sparse_backfill(app_state, did, state, root_commit, scan.leaves, blocks).await?; + + let Some((records, state)) = result else { + cleanup_discarded_repo(app_state, did).await?; + return Ok(SparseBackfillResult::Discarded); + }; + + Ok(SparseBackfillResult::Imported(SparseBackfillSuccess { + state, + records, + node_blocks: scan.node_blocks_seen, + node_bytes: scan.node_bytes_seen, + })) +} + +async fn fetch_blocks( + http: &reqwest::Client, + pds: &url::Url, + did: &Did<'static>, + cids: &[IpldCid], +) -> Result, BackfillError> { + let mut out = BTreeMap::new(); + for chunk in cids.chunks(SPARSE_GET_BLOCKS_CHUNK) { + if chunk.is_empty() { + continue; + } + + let req = GetBlocks::new() + .did(did.clone()) + .cids( + chunk + .iter() + .map(|cid| AtCid::from(cid.to_string())) + .collect::>(), + ) + .build(); + let resp = http.xrpc(url_to_fluent_uri(pds)).send(&req).await?; + let car = match resp.into_output() { + Ok(o) => o, + Err(XrpcError::Xrpc(GetBlocksError::BlockNotFound(_))) => { + return Ok(BTreeMap::new()); + } + Err(XrpcError::Xrpc( + GetBlocksError::RepoNotFound(_) + | GetBlocksError::RepoTakendown(_) + | GetBlocksError::RepoSuspended(_) + | GetBlocksError::RepoDeactivated(_), + )) => return Ok(BTreeMap::new()), + Err(e) => Err(e).into_diagnostic()?, + }; + let parsed = crate::car::parse_car_blocks(&car.body).await?; + out.extend(parsed); + } + Ok(out) +} + +async fn persist_sparse_backfill( + app_state: &Arc, + did: &Did<'static>, + mut state: RepoState<'static>, + root_commit: Commit, + leaves: Vec<(SmolStr, IpldCid)>, + blocks: BTreeMap, +) -> Result)>, BackfillError> { + let app_state = app_state.clone(); + let did = did.clone(); + tokio::task::spawn_blocking(move || { + let filter = app_state.filter.load(); + let ephemeral = app_state.ephemeral; + let only_index_links = app_state.only_index_links; + let mut count = 0; + let mut delta = 0; + let mut added_blocks = 0; + let mut collection_counts: HashMap = HashMap::new(); + let mut batch = app_state.db.inner.batch(); + + let prefix = keys::record_prefix_did(&did); + let mut existing_cids: HashMap<(SmolStr, DbRkey), SmolStr> = HashMap::new(); + + if !ephemeral { + for guard in app_state.db.records.prefix(&prefix) { + let (key, cid_bytes) = guard.into_inner().into_diagnostic()?; + let mut remaining = key[prefix.len()..].splitn(2, |b| keys::SEP.eq(b)); + let collection_raw = remaining + .next() + .ok_or_else(|| miette::miette!("invalid record key format: {key:?}"))?; + let rkey_raw = remaining + .next() + .ok_or_else(|| miette::miette!("invalid record key format: {key:?}"))?; + + let collection = std::str::from_utf8(collection_raw) + .map_err(|e| miette::miette!("invalid collection utf8: {e}"))?; + let rkey = keys::parse_rkey(rkey_raw) + .map_err(|e| miette::miette!("invalid rkey '{key:?}' for {did}: {e}"))?; + let cid = cid::Cid::read_bytes(cid_bytes.as_ref()) + .map_err(|e| miette::miette!("invalid cid '{cid_bytes:?}' for {did}: {e}"))? + .to_smolstr(); + + existing_cids.insert((collection.into(), rkey), cid); + } + } + + let mut signal_seen = filter.mode == FilterMode::Full || filter.signals.is_empty(); + + for (key, cid) in leaves { + let (collection, rkey) = ops::parse_path(&key)?; + + if !filter.matches_collection(collection) { + continue; + } + + let Some(val) = blocks.get(&cid).cloned() else { + return Err(miette::miette!("missing sparse record block {cid}")); + }; + + if !signal_seen && filter.matches_signal(collection) { + debug!(collection = %collection, "signal matched"); + signal_seen = true; + } + + let rkey = DbRkey::new(rkey); + let path = (collection.to_smolstr(), rkey.clone()); + let cid_obj = AtCid::ipld(cid); + + *collection_counts.entry(path.0.clone()).or_default() += 1; + + let existing_cid = existing_cids.remove(&path); + let action = if let Some(existing_cid) = &existing_cid { + if existing_cid == cid_obj.as_str() { + trace!(collection = %collection, rkey = %rkey, cid = %cid, "skip unchanged sparse record"); + continue; + } + DbAction::Update + } else { + DbAction::Create + }; + trace!(collection = %collection, rkey = %rkey, cid = %cid, ?action, "action sparse record"); + + let db_key = keys::record_key(&did, collection, &rkey); + let cid_raw = cid.to_bytes(); + let block_key = Slice::from(keys::block_key(collection, &cid_raw)); + if !ephemeral { + if !only_index_links { + batch.insert(&app_state.db.blocks, block_key.clone(), val.as_ref()); + } + batch.insert(&app_state.db.records, db_key, cid_raw); + #[cfg(feature = "backlinks")] + if let Ok(value) = + serde_ipld_dagcbor::from_slice::(val.as_ref()) + { + crate::backlinks::store::index_record( + &mut batch, + &app_state.db.backlinks, + did.as_str(), + collection, + &rkey.to_smolstr(), + &value, + )?; + } + } + + added_blocks += 1; + if action == DbAction::Create { + delta += 1; + } + + #[cfg(feature = "indexer_stream")] + { + let event_id = app_state.db.next_event_id.fetch_add(1, Ordering::SeqCst); + let evt = StoredEvent { + live: false, + did: TrimmedDid::from(&did), + rev: root_commit.rev, + collection: CowStr::Borrowed(collection), + rkey, + action, + data: if ephemeral { + StoredData::Block(val) + } else if only_index_links { + StoredData::Nothing + } else { + StoredData::Ptr(cid_obj.to_ipld().expect("valid cid")) + }, + }; + let bytes = rmp_serde::to_vec(&evt).into_diagnostic()?; + batch.insert(&app_state.db.events, keys::event_key(event_id), bytes); + + #[cfg(feature = "jetstream")] + { + let jetstream = crate::types::StoredJetstreamEvent::Commit { + did: TrimmedDid::from(&did).into_static(), + collection: CowStr::Borrowed(collection).into_static(), + event_id, + live: false, + }; + crate::jetstream::stage_event(&mut batch, &app_state.db, jetstream, None)?; + } + } + + count += 1; + } + + for ((collection, rkey), cid) in existing_cids { + trace!(collection = %collection, rkey = %rkey, cid = %cid, "remove sparse-stale record"); + + batch.remove( + &app_state.db.records, + keys::record_key(&did, &collection, &rkey), + ); + #[cfg(feature = "backlinks")] + crate::backlinks::store::delete_record( + &mut batch, + &app_state.db.backlinks, + did.as_str(), + &collection, + &rkey.to_smolstr(), + )?; + + #[cfg(feature = "indexer_stream")] + { + let event_id = app_state.db.next_event_id.fetch_add(1, Ordering::SeqCst); + let evt = StoredEvent { + live: false, + did: TrimmedDid::from(&did), + rev: root_commit.rev, + collection: CowStr::Borrowed(&collection), + rkey, + action: DbAction::Delete, + data: StoredData::Nothing, + }; + let bytes = rmp_serde::to_vec(&evt).into_diagnostic()?; + batch.insert(&app_state.db.events, keys::event_key(event_id), bytes); + + #[cfg(feature = "jetstream")] + { + let jetstream = crate::types::StoredJetstreamEvent::Commit { + did: TrimmedDid::from(&did).into_static(), + collection: CowStr::Borrowed(&collection).into_static(), + event_id, + live: false, + }; + crate::jetstream::stage_event(&mut batch, &app_state.db, jetstream, None)?; + } + } + + delta -= 1; + count += 1; + } + + if !signal_seen { + trace!(signals = ?filter.signals, "no signal-matching sparse records found, discarding repo"); + return Ok::<_, miette::Report>(None); + } + + state.root = Some(root_commit); + state.touch(); + + batch.insert( + &app_state.db.repos, + keys::repo_key(&did), + ser_repo_state(&state)?, + ); + + let metadata_key = keys::repo_metadata_key(&did); + let metadata_bytes = app_state + .db + .repo_metadata + .get(&metadata_key) + .into_diagnostic()? + .ok_or_else(|| miette::miette!("repo metadata not found for {}", did))?; + let mut metadata = crate::db::deser_repo_meta(&metadata_bytes)?; + metadata.tracked = true; + batch.insert( + &app_state.db.repo_metadata, + &metadata_key, + crate::db::ser_repo_meta(&metadata)?, + ); + + if !ephemeral { + db::replace_record_counts( + &mut batch, + &app_state.db, + &did, + collection_counts.iter().map(|(col, cnt)| (col.as_str(), *cnt)), + )?; + } + + let mut count_deltas = CountDeltas::default(); + if delta != 0 { + count_deltas.add("records", delta); + } + if added_blocks > 0 { + count_deltas.add("blocks", added_blocks); + } + let reservation = app_state.db.stage_count_deltas(&mut batch, &count_deltas); + batch.commit().into_diagnostic()?; + app_state.db.apply_count_deltas(&count_deltas); + drop(reservation); + + Ok::<_, miette::Report>(Some((count, state))) + }) + .await + .into_diagnostic()? + .map_err(BackfillError::from) +} + +async fn cleanup_discarded_repo( + app_state: &Arc, + did: &Did<'static>, +) -> Result<(), BackfillError> { + let metadata_key = keys::repo_metadata_key(did); + let metadata_bytes = app_state + .db + .repo_metadata + .get(&metadata_key) + .into_diagnostic()? + .ok_or_else(|| miette::miette!("repo metadata not found for {}", did))?; + let metadata = crate::db::deser_repo_meta(metadata_bytes.as_ref())?; + let did_key = keys::repo_key(did); + let backfill_pending_key = keys::pending_key(metadata.index_id); + let app_state = app_state.clone(); + + tokio::task::spawn_blocking(move || { + let mut batch = app_state.db.inner.batch(); + let mut count_deltas = CountDeltas::default(); + batch.remove(&app_state.db.repos, &did_key); + batch.remove(&app_state.db.repo_metadata, &metadata_key); + batch.remove(&app_state.db.pending, backfill_pending_key); + count_deltas.add("repos", -1); + count_deltas.add("pending", -1); + let reservation = app_state.db.stage_count_deltas(&mut batch, &count_deltas); + batch.commit().into_diagnostic().inspect(|_| { + app_state.db.apply_count_deltas(&count_deltas); + drop(reservation); + }) + }) + .await + .into_diagnostic()??; + + Ok(()) +} + async fn process_did( app_state: &Arc, http: &reqwest::Client, did: &Did<'static>, verify_signatures: bool, + strategy: BackfillStrategy, ) -> Result>, BackfillError> { debug!("starting..."); @@ -438,6 +916,54 @@ }; let _ = app_state.db.event_tx.send(ops::make_account_event(db, evt)); }; + + if strategy != BackfillStrategy::Full { + let filter = app_state.filter.load(); + let sparse_supported = !filter.collections.is_empty() + && sparse_probe_collection(&filter.collections).is_some(); + let should_try_sparse = match strategy { + BackfillStrategy::Full => false, + BackfillStrategy::SparseFilter => sparse_supported, + BackfillStrategy::Auto => sparse_supported, + }; + + if should_try_sparse { + match process_did_sparse(app_state, http, did, &pds, state.clone(), verify_signatures) + .await + { + Ok(SparseBackfillResult::Imported(sparse)) => { + #[cfg(feature = "indexer_stream")] + if sparse.state.active != previous_state.active + || sparse.state.status != previous_state.status + || previous_state.pds.is_none() + { + emit_identity(&sparse.state.status, sparse.state.active); + } + + trace!( + records = sparse.records, + node_blocks = sparse.node_blocks, + node_bytes = sparse.node_bytes, + active = sparse.state.active, + status = ?sparse.state.status, + "sparse backfill complete" + ); + return Ok(Some(previous_state)); + } + Ok(SparseBackfillResult::Discarded) => { + return Ok(None); + } + Ok(SparseBackfillResult::Skipped) => { + debug!("sparse backfill skipped, falling back to full getRepo"); + } + Err(e) => { + warn!(err = %e, "sparse backfill failed, falling back to full getRepo"); + } + } + } else if strategy == BackfillStrategy::SparseFilter { + debug!("sparse backfill requested but filter is not sparse-compatible"); + } + } // 2. fetch repo (car) let start = Instant::now(); @@ -637,7 +1163,7 @@ let rkey = DbRkey::new(rkey); let path = (collection.to_smolstr(), rkey.clone()); - let cid_obj = Cid::ipld(cid); + let cid_obj = AtCid::ipld(cid); *collection_counts.entry(path.0.clone()).or_default() += 1; diff --git a/src/bin/backfill_strategy_bench.rs b/src/bin/backfill_strategy_bench.rs new file mode 100644 --- /dev/null +++ b/src/bin/backfill_strategy_bench.rs @@ -0,0 +1,369 @@ +#[path = "../car.rs"] +mod car; +#[path = "../sparse_mst.rs"] +mod sparse_mst; + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use cid::Cid as IpldCid; +use jacquard_api::com_atproto::sync::get_blocks::GetBlocks; +use jacquard_api::com_atproto::sync::get_record::GetRecord; +use jacquard_api::com_atproto::sync::get_repo::GetRepo; +use jacquard_common::deps::bytes::Bytes; +use jacquard_common::deps::fluent_uri; +use jacquard_common::types::cid::Cid as AtCid; +use jacquard_common::types::did::Did; +use jacquard_common::types::string::{Nsid, RecordKey}; +use jacquard_common::xrpc::{XrpcError, XrpcExt}; +use jacquard_repo::{BlockStore, MemoryBlockStore, Mst}; +use miette::{IntoDiagnostic, Result, WrapErr}; +use serde::Deserialize; +use smol_str::SmolStr; +use sparse_mst::{SparseScanner, sparse_probe_collection, sparse_ranges}; +use url::Url; + +const GET_BLOCKS_CHUNK: usize = 64; + +#[derive(Debug)] +struct Args { + collection: String, + pattern: SmolStr, + index_url: Url, + plc_url: Url, + limit: usize, +} + +impl Args { + fn parse() -> Result { + let mut collection = "sh.tangled.repo".to_string(); + let mut pattern = SmolStr::new("sh.tangled.*"); + let mut index_url = Url::parse("https://lightrail.microcosm.blue").into_diagnostic()?; + let mut plc_url = Url::parse("https://plc.directory").into_diagnostic()?; + let mut limit = 5usize; + + let mut args = std::env::args().skip(1); + while let Some(arg) = args.next() { + let Some(value) = args.next() else { + return Err(miette::miette!("missing value for {arg}")); + }; + match arg.as_str() { + "--collection" => collection = value, + "--pattern" => pattern = SmolStr::new(value), + "--index" => index_url = Url::parse(&value).into_diagnostic()?, + "--plc" => plc_url = Url::parse(&value).into_diagnostic()?, + "--limit" => limit = value.parse().into_diagnostic()?, + _ => return Err(miette::miette!("unknown argument {arg}")), + } + } + + Ok(Self { + collection, + pattern, + index_url, + plc_url, + limit, + }) + } +} + +#[derive(Debug, Deserialize)] +struct ListReposByCollectionOutput { + repos: Vec, +} + +#[derive(Debug, Deserialize)] +struct RepoHit { + did: String, +} + +#[derive(Debug, Deserialize)] +struct DidDoc { + service: Vec, +} + +#[derive(Debug, Deserialize)] +struct DidService { + #[serde(rename = "type")] + kind: String, + #[serde(rename = "serviceEndpoint")] + service_endpoint: Url, +} + +#[derive(Debug)] +struct FullBench { + fetch: Duration, + parse_and_walk: Duration, + bytes: usize, + blocks: usize, + leaves: usize, + matching: usize, +} + +#[derive(Debug)] +struct SparseBench { + total: Duration, + seed_bytes: usize, + node_bytes: usize, + record_bytes: usize, + node_blocks: usize, + records: usize, +} + +#[tokio::main] +async fn main() -> Result<()> { + let args = Args::parse()?; + let http = reqwest::Client::builder() + .timeout(Duration::from_secs(120)) + .zstd(true) + .brotli(true) + .gzip(true) + .build() + .into_diagnostic()?; + + let repos = list_repos(&http, &args).await?; + let ranges = sparse_ranges(&[args.pattern.clone()]); + + println!( + "did,pds,full_fetch_ms,full_parse_walk_ms,full_bytes,full_blocks,full_leaves,full_matching,sparse_total_ms,sparse_seed_bytes,sparse_node_bytes,sparse_record_bytes,sparse_node_blocks,sparse_records" + ); + + for did in repos { + let pds = resolve_pds(&http, &args.plc_url, &did).await?; + let full = bench_full(&http, &pds, &did, &ranges).await?; + let sparse = bench_sparse(&http, &pds, &did, &[args.pattern.clone()]).await?; + println!( + "{did},{pds},{},{},{},{},{},{},{},{},{},{},{},{}", + full.fetch.as_millis(), + full.parse_and_walk.as_millis(), + full.bytes, + full.blocks, + full.leaves, + full.matching, + sparse.total.as_millis(), + sparse.seed_bytes, + sparse.node_bytes, + sparse.record_bytes, + sparse.node_blocks, + sparse.records, + ); + } + + Ok(()) +} + +async fn list_repos(http: &reqwest::Client, args: &Args) -> Result>> { + let mut url = args + .index_url + .join("/xrpc/com.atproto.sync.listReposByCollection") + .into_diagnostic()?; + url.query_pairs_mut() + .append_pair("collection", &args.collection) + .append_pair("limit", &args.limit.to_string()); + + let output = http + .get(url) + .send() + .await + .into_diagnostic()? + .error_for_status() + .into_diagnostic()? + .json::() + .await + .into_diagnostic()?; + + output + .repos + .into_iter() + .map(|repo| Did::new_owned(repo.did).into_diagnostic()) + .collect() +} + +async fn resolve_pds(http: &reqwest::Client, plc: &Url, did: &Did<'static>) -> Result { + let mut url = plc.clone(); + url.path_segments_mut() + .map_err(|_| miette::miette!("plc url cannot be a base"))? + .push(did.as_str()); + let doc = http + .get(url) + .send() + .await + .into_diagnostic()? + .error_for_status() + .into_diagnostic()? + .json::() + .await + .into_diagnostic()?; + + doc.service + .into_iter() + .find(|svc| svc.kind == "AtprotoPersonalDataServer") + .map(|svc| svc.service_endpoint) + .ok_or_else(|| miette::miette!("no pds service in did doc for {did}")) +} + +async fn bench_full( + http: &reqwest::Client, + pds: &Url, + did: &Did<'static>, + ranges: &[sparse_mst::KeyRange], +) -> Result { + let fetch_start = Instant::now(); + let req = GetRepo::new().did(did.clone()).build(); + let resp = http.xrpc(to_fluent_uri(pds)).send(&req).await?; + let car = resp + .into_output() + .map_err(|err| miette::miette!("getRepo failed for {did}: {err}"))?; + let fetch = fetch_start.elapsed(); + let bytes = car.body.len(); + + let parse_start = Instant::now(); + let parsed = jacquard_repo::car::reader::parse_car_bytes(&car.body) + .await + .into_diagnostic() + .wrap_err_with(|| format!("parse getRepo CAR for {did} ({} bytes)", car.body.len()))?; + let blocks = parsed.blocks.len(); + let store = Arc::new(MemoryBlockStore::new_from_blocks(parsed.blocks)); + let root_bytes = store + .get(&parsed.root) + .await + .into_diagnostic()? + .ok_or_else(|| miette::miette!("root block missing from getRepo car"))?; + let root_commit = jacquard_repo::commit::Commit::from_cbor(&root_bytes).into_diagnostic()?; + let mst: Mst = Mst::load(store, root_commit.data, None); + let leaves = mst.leaves().await.into_diagnostic()?; + let matching = leaves + .iter() + .filter(|(key, _)| ranges.iter().any(|range| range.contains(key))) + .count(); + let parse_and_walk = parse_start.elapsed(); + + Ok(FullBench { + fetch, + parse_and_walk, + bytes, + blocks, + leaves: leaves.len(), + matching, + }) +} + +async fn bench_sparse( + http: &reqwest::Client, + pds: &Url, + did: &Did<'static>, + patterns: &[SmolStr], +) -> Result { + let start = Instant::now(); + let ranges = sparse_ranges(patterns); + let probe_collection = sparse_probe_collection(patterns) + .ok_or_else(|| miette::miette!("no sparse-compatible probe collection"))?; + + let req = GetRecord::new() + .did(did.clone()) + .collection(Nsid::new_owned(probe_collection.as_str()).into_diagnostic()?) + .rkey(RecordKey::any_static("-").into_diagnostic()?) + .build(); + let resp = http.xrpc(to_fluent_uri(pds)).send(&req).await?; + let seed = resp + .into_output() + .map_err(|err| miette::miette!("getRecord seed failed for {did}: {err}"))?; + let seed_bytes = seed.body.len(); + let parsed = jacquard_repo::car::reader::parse_car_bytes(&seed.body) + .await + .into_diagnostic() + .wrap_err_with(|| { + format!( + "parse getRecord seed CAR for {did} ({} bytes)", + seed.body.len() + ) + })?; + let root_bytes = parsed + .blocks + .get(&parsed.root) + .ok_or_else(|| miette::miette!("root block missing from sparse seed car"))?; + let root_commit = jacquard_repo::commit::Commit::from_cbor(root_bytes).into_diagnostic()?; + let root_cid = root_commit.data; + + let mut node_fetch_bytes = 0usize; + let mut scanner = SparseScanner::new(ranges, parsed.blocks); + let scan = loop { + match scanner.scan(root_cid)? { + Ok(scan) => break scan, + Err(missing) => { + let (blocks, bytes) = fetch_blocks(http, pds, did, &missing).await?; + node_fetch_bytes += bytes; + scanner.insert_blocks(blocks); + } + } + }; + + let mut blocks = scanner.take_blocks(); + let missing_records: Vec = scan + .leaves + .iter() + .filter_map(|(_, cid)| (!blocks.contains_key(cid)).then_some(*cid)) + .collect(); + let (record_blocks, record_bytes) = fetch_blocks(http, pds, did, &missing_records).await?; + blocks.extend(record_blocks); + + Ok(SparseBench { + total: start.elapsed(), + seed_bytes, + node_bytes: scan.node_bytes_seen + node_fetch_bytes, + record_bytes, + node_blocks: scan.node_blocks_seen, + records: scan.leaves.len(), + }) +} + +async fn fetch_blocks( + http: &reqwest::Client, + pds: &Url, + did: &Did<'static>, + cids: &[IpldCid], +) -> Result<(BTreeMap, usize)> { + let mut out = BTreeMap::new(); + let mut bytes = 0usize; + for chunk in cids.chunks(GET_BLOCKS_CHUNK) { + if chunk.is_empty() { + continue; + } + + let req = GetBlocks::new() + .did(did.clone()) + .cids( + chunk + .iter() + .map(|cid| AtCid::from(cid.to_string())) + .collect::>(), + ) + .build(); + let resp = http.xrpc(to_fluent_uri(pds)).send(&req).await?; + let car = resp + .into_output() + .map_err(|err: XrpcError<_>| miette::miette!("getBlocks failed for {did}: {err}"))?; + bytes += car.body.len(); + let parsed = car::parse_car_blocks(&car.body).await.wrap_err_with(|| { + let cids = chunk + .iter() + .map(ToString::to_string) + .collect::>() + .join(","); + format!( + "parse getBlocks CAR for {did} ({} requested, {} bytes, cids={cids})", + chunk.len(), + car.body.len() + ) + })?; + out.extend(parsed); + } + Ok((out, bytes)) +} + +fn to_fluent_uri(url: &Url) -> fluent_uri::Uri { + fluent_uri::Uri::parse(url.as_str()) + .expect("validated url") + .to_owned() +} diff --git a/src/control/mod.rs b/src/control/mod.rs --- a/src/control/mod.rs +++ b/src/control/mod.rs @@ -321,6 +321,7 @@ config.verify_signatures, SignatureVerification::Full | SignatureVerification::BackfillOnly ), + config.backfill_strategy, state.backfill_enabled.subscribe(), ) .run()