diff --git a/src/backfill/sparse.rs b/src/backfill/sparse.rs index 6c91914..2971da1 100644 --- a/src/backfill/sparse.rs +++ b/src/backfill/sparse.rs @@ -3,7 +3,7 @@ use crate::backfill::error::BackfillError; use crate::config::{BackfillStrategy, RateTier}; use crate::db::types::{DbAction, DbRkey, TrimmedDid}; use crate::db::{self, CountDeltas, keys, ser_repo_state}; -use crate::filter::FilterMode; +use crate::filter::{FilterConfig, FilterMode}; use crate::ops; use crate::sparse_mst::{SparseScanner, mst_node_layer, sparse_probe_collection, sparse_ranges}; use crate::state::AppState; @@ -51,9 +51,26 @@ pub(crate) enum SparseBackfillResult { } const SPARSE_GET_BLOCKS_CHUNK: usize = 100; -const SPARSE_GET_BLOCKS_PARALLELISM: usize = 1; +const SPARSE_GET_BLOCKS_PARALLELISM: usize = 4; const SPARSE_MAX_SCAN_ROUNDS: usize = 256; const SPARSE_AUTO_FULL_MAX_ROOT_LAYER: usize = 2; +fn sparse_probe(filter: &FilterConfig) -> Option<(SmolStr, &'static str)> { + filter + .collections + .iter() + .find(|collection| !collection.ends_with(".*") && collection.ends_with(".profile")) + .or_else(|| { + filter + .signals + .iter() + .find(|signal| signal.ends_with(".profile") && filter.matches_collection(signal)) + }) + .cloned() + .map(|collection| (collection, "self")) + .or_else(|| { + sparse_probe_collection(&filter.collections).map(|collection| (collection, "-")) + }) +} pub(crate) async fn process_did_sparse( app_state: &Arc, @@ -65,7 +82,7 @@ pub(crate) async fn process_did_sparse( strategy: BackfillStrategy, ) -> Result { let filter = app_state.filter.load(); - let Some(probe_collection) = sparse_probe_collection(&filter.collections) else { + let Some((probe_collection, probe_rkey)) = sparse_probe(&filter) else { return Ok(SparseBackfillResult::Skipped); }; let ranges = sparse_ranges(&filter.collections); @@ -74,7 +91,7 @@ pub(crate) async fn process_did_sparse( } let probe_collection = Nsid::new_owned(probe_collection.as_str()).into_diagnostic()?; - let probe_rkey = RecordKey::any_static("-").into_diagnostic()?; + let probe_rkey = RecordKey::any_static(probe_rkey).into_diagnostic()?; let req = GetRecord::new() .did(did.clone()) .collection(probe_collection) @@ -138,20 +155,17 @@ pub(crate) async fn process_did_sparse( } let root_cid = root_commit.data; - if strategy == BackfillStrategy::Auto { - if let Some(root_bytes) = parsed.blocks.get(&root_cid) { - if let Some(root_layer) = mst_node_layer(root_bytes)? { - if root_layer <= SPARSE_AUTO_FULL_MAX_ROOT_LAYER { - debug!( - root_layer, - max_sparse_layer = SPARSE_AUTO_FULL_MAX_ROOT_LAYER, - "sparse auto selected full getRepo for small repo" - ); - return Ok(SparseBackfillResult::Skipped); - } - } - } - } + let auto_full_layer = if strategy == BackfillStrategy::Auto { + parsed + .blocks + .get(&root_cid) + .map(|root_bytes| mst_node_layer(root_bytes)) + .transpose()? + .flatten() + .filter(|layer| *layer <= SPARSE_AUTO_FULL_MAX_ROOT_LAYER) + } else { + None + }; let root_commit = Commit::from(root_commit); let mut scanner = SparseScanner::new(ranges, parsed.blocks); @@ -160,6 +174,14 @@ pub(crate) async fn process_did_sparse( match scanner.scan(root_cid)? { Ok(scan) => break scan, Err(missing) => { + if let Some(root_layer) = auto_full_layer { + debug!( + root_layer, + max_sparse_layer = SPARSE_AUTO_FULL_MAX_ROOT_LAYER, + "sparse auto selected full getRepo because the probe omitted MST nodes" + ); + return Ok(SparseBackfillResult::Skipped); + } scan_rounds += 1; if scan_rounds > SPARSE_MAX_SCAN_ROUNDS { return Err( @@ -193,6 +215,16 @@ pub(crate) async fn process_did_sparse( .iter() .filter_map(|cid| (!scanned_blocks.contains_key(cid)).then_some(*cid)) .collect::>(); + if !missing_records.is_empty() + && let Some(root_layer) = auto_full_layer + { + debug!( + root_layer, + max_sparse_layer = SPARSE_AUTO_FULL_MAX_ROOT_LAYER, + "sparse auto selected full getRepo because the probe omitted record blocks" + ); + return Ok(SparseBackfillResult::Skipped); + } let mut record_blocks = record_cids .iter() .filter_map(|cid| scanned_blocks.remove(cid).map(|bytes| (*cid, bytes))) @@ -244,17 +276,7 @@ async fn fetch_blocks( let fetches = cids .chunks(SPARSE_GET_BLOCKS_CHUNK) - .filter(|chunk| !chunk.is_empty()) - .map(|chunk| { - fetch_block_chunk( - http.clone(), - pds.clone(), - did.clone(), - chunk.to_vec(), - throttle, - tier, - ) - }) + .map(|chunk| fetch_block_chunk(http, pds, did, chunk, throttle, tier)) .collect::>(); let fetches = stream::iter(fetches).buffer_unordered(SPARSE_GET_BLOCKS_PARALLELISM); futures::pin_mut!(fetches); @@ -267,10 +289,10 @@ async fn fetch_blocks( } async fn fetch_block_chunk( - http: ThrottledHttpClient, - pds: url::Url, - did: Did<'static>, - cids: Vec, + http: &ThrottledHttpClient, + pds: &url::Url, + did: &Did<'_>, + cids: &[IpldCid], throttle: &ThrottleHandle, tier: &RateTier, ) -> Result, BackfillError> { @@ -280,7 +302,7 @@ async fn fetch_block_chunk( { let mut query = url.query_pairs_mut(); query.append_pair("did", did.as_str()); - for cid in &cids { + for cid in cids { query.append_pair("cids", &cid.to_string()); } } @@ -320,9 +342,7 @@ async fn fetch_block_chunk( .bytes() .await .map_err(|e| BackfillError::Transport(e.to_string().into()))?; - crate::car::parse_car_blocks(&body) - .await - .map_err(BackfillError::from) + crate::car::parse_car_blocks(body).map_err(BackfillError::from) } else { let retry_after = if status == StatusCode::TOO_MANY_REQUESTS { crate::util::parse_retry_after(&resp) @@ -614,3 +634,131 @@ async fn persist_sparse_backfill( .into_diagnostic()? .map_err(BackfillError::from) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::filter::FilterConfig; + use axum::{Router, extract::State, response::IntoResponse, routing::get}; + use cid::Cid; + use cid::multihash::Multihash; + use jacquard_common::types::crypto::{DAG_CBOR, SHA2_256}; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + use tokio::sync::Barrier; + + #[derive(Clone)] + struct FetchState { + started: Arc, + active: Arc, + max_active: Arc, + barrier: Arc, + car: bytes::Bytes, + } + + async fn delayed_car(State(state): State) -> impl IntoResponse { + let ordinal = state.started.fetch_add(1, Ordering::SeqCst) + 1; + let active = state.active.fetch_add(1, Ordering::SeqCst) + 1; + state.max_active.fetch_max(active, Ordering::SeqCst); + if ordinal <= 2 { + state.barrier.wait().await; + } + state.active.fetch_sub(1, Ordering::SeqCst); + ( + [( + reqwest::header::CONTENT_TYPE.as_str(), + "application/vnd.ipld.car", + )], + state.car, + ) + } + + fn cid(byte: u8) -> IpldCid { + let hash = [byte; 32]; + let multihash = Multihash::<64>::wrap(SHA2_256, &hash).unwrap(); + Cid::new_v1(DAG_CBOR, multihash) + } + + fn filter(collections: &[&str], signals: &[&str]) -> FilterConfig { + let mut filter = FilterConfig::new(FilterMode::Filter); + filter.collections = collections.iter().map(SmolStr::new).collect(); + filter.signals = signals.iter().map(SmolStr::new).collect(); + filter + } + + #[test] + fn sparse_probe_prefers_matching_profile_signal() { + let filter = filter( + &["sh.tangled.*"], + &["sh.tangled.actor.profile", "app.bsky.actor.profile"], + ); + + assert_eq!( + sparse_probe(&filter), + Some((SmolStr::new("sh.tangled.actor.profile"), "self")) + ); + } + + #[test] + fn sparse_probe_ignores_profile_signal_outside_filter() { + let filter = filter(&["sh.tangled.*"], &["app.bsky.actor.profile"]); + + assert_eq!( + sparse_probe(&filter), + Some((SmolStr::new("sh.tangled.probe"), "-")) + ); + } + + #[test] + fn sparse_probe_uses_self_for_exact_profile_collection() { + let filter = filter(&["app.bsky.actor.profile"], &[]); + + assert_eq!( + sparse_probe(&filter), + Some((SmolStr::new("app.bsky.actor.profile"), "self")) + ); + } + + #[tokio::test] + async fn fetches_independent_get_blocks_chunks_concurrently() { + let mut car = Vec::new(); + let mut writer = + iroh_car::CarWriter::new(iroh_car::CarHeader::new_v1(Vec::new()), &mut car); + let response_cid = cid(u8::MAX); + writer.write(response_cid, b"block".to_vec()).await.unwrap(); + writer.finish().await.unwrap(); + let state = FetchState { + started: Arc::new(AtomicUsize::new(0)), + active: Arc::new(AtomicUsize::new(0)), + max_active: Arc::new(AtomicUsize::new(0)), + barrier: Arc::new(Barrier::new(2)), + car: car.into(), + }; + let app = Router::new() + .route("/xrpc/com.atproto.sync.getBlocks", get(delayed_car)) + .with_state(state.clone()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let pds = url::Url::parse(&format!("http://{}/", listener.local_addr().unwrap())).unwrap(); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + + let throttler = crate::util::throttle::Throttler::new(1, 10); + let http = ThrottledHttpClient::new(vec![reqwest::Client::new()], throttler.clone()); + let throttle = throttler.get_handle(&pds).await; + let did = Did::new_static("did:plc:aaaaaaaaaaaaaaaaaaaaaaaa").unwrap(); + let cids = (0..=200).map(|byte| cid(byte as u8)).collect::>(); + + let blocks = tokio::time::timeout( + Duration::from_secs(5), + fetch_blocks(&http, &pds, &did, &cids, &throttle, &RateTier::trusted()), + ) + .await + .expect("getBlocks chunks were fetched serially") + .unwrap(); + + assert_eq!(blocks.get(&response_cid).unwrap().as_ref(), b"block"); + assert!( + state.max_active.load(Ordering::SeqCst) > 1, + "expected more than one getBlocks chunk in flight" + ); + } +} diff --git a/src/bin/backfill_strategy_bench.rs b/src/bin/backfill_strategy_bench.rs index b0516f1..8d35d79 100644 --- a/src/bin/backfill_strategy_bench.rs +++ b/src/bin/backfill_strategy_bench.rs @@ -36,6 +36,8 @@ struct Args { index_url: Url, plc_url: Url, limit: usize, + probe_collection: Option, + probe_rkey: SmolStr, } impl Args { @@ -45,6 +47,8 @@ impl Args { 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 probe_collection = None; + let mut probe_rkey = SmolStr::new("-"); let mut args = std::env::args().skip(1); while let Some(arg) = args.next() { @@ -57,6 +61,8 @@ impl Args { "--index" => index_url = Url::parse(&value).into_diagnostic()?, "--plc" => plc_url = Url::parse(&value).into_diagnostic()?, "--limit" => limit = value.parse().into_diagnostic()?, + "--probe-collection" => probe_collection = Some(SmolStr::new(value)), + "--probe-rkey" => probe_rkey = SmolStr::new(value), _ => return Err(miette::miette!("unknown argument {arg}")), } } @@ -67,6 +73,8 @@ impl Args { index_url, plc_url, limit, + probe_collection, + probe_rkey, }) } } @@ -107,6 +115,8 @@ struct FullBench { #[derive(Debug)] struct SparseBench { total: Duration, + requests: usize, + auto_requests: usize, root_layer: Option, seed_bytes: usize, node_bytes: usize, @@ -149,7 +159,7 @@ async fn main() -> Result<()> { 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_root_layer,sparse_seed_bytes,sparse_node_bytes,sparse_record_bytes,sparse_node_blocks,sparse_records" + "did,pds,full_fetch_ms,full_parse_walk_ms,full_bytes,full_blocks,full_leaves,full_matching,sparse_total_ms,sparse_requests,auto_requests,sparse_root_layer,sparse_seed_bytes,sparse_node_bytes,sparse_record_bytes,sparse_node_blocks,sparse_records" ); let mut totals = BenchTotals::default(); @@ -168,7 +178,16 @@ async fn main() -> Result<()> { continue; } }; - let sparse = match bench_sparse(&http, &pds, &did, &[args.pattern.clone()]).await { + let sparse = match bench_sparse( + &http, + &pds, + &did, + &[args.pattern.clone()], + args.probe_collection.as_deref(), + args.probe_rkey.as_str(), + ) + .await + { Ok(sparse) => sparse, Err(err) => { eprintln!("skipping {did}: sparse bench failed: {err:?}"); @@ -176,7 +195,7 @@ async fn main() -> Result<()> { } }; println!( - "{did},{pds},{},{},{},{},{},{},{},{},{},{},{},{},{}", + "{did},{pds},{},{},{},{},{},{},{},{},{},{},{},{},{},{},{}", full.fetch.as_millis(), full.parse_and_walk.as_millis(), full.bytes, @@ -184,6 +203,8 @@ async fn main() -> Result<()> { full.leaves, full.matching, sparse.total.as_millis(), + sparse.requests, + sparse.auto_requests, sparse .root_layer .map(|layer| layer.to_string()) @@ -307,16 +328,20 @@ async fn bench_sparse( pds: &Url, did: &Did<'static>, patterns: &[SmolStr], + probe_collection: Option<&str>, + probe_rkey: &str, ) -> Result { let start = Instant::now(); let ranges = sparse_ranges(patterns); - let probe_collection = sparse_probe_collection(patterns) + let probe_collection = probe_collection + .map(SmolStr::new) + .or_else(|| 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()?) + .rkey(RecordKey::any(probe_rkey).into_diagnostic()?) .build(); let resp = http.xrpc(to_fluent_uri(pds)).send(&req).await?; let seed = resp @@ -346,13 +371,16 @@ async fn bench_sparse( .flatten(); let mut node_fetch_bytes = 0usize; + let mut requests = 1usize; 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?; + let (blocks, bytes, chunk_requests) = + fetch_blocks(http, pds, did, &missing).await?; node_fetch_bytes += bytes; + requests += chunk_requests; scanner.insert_blocks(blocks); } } @@ -364,11 +392,21 @@ async fn bench_sparse( .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?; + let (record_blocks, record_bytes, chunk_requests) = + fetch_blocks(http, pds, did, &missing_records).await?; + requests += chunk_requests; blocks.extend(record_blocks); + let auto_requests = if requests > 1 && root_layer.is_some_and(|layer| layer <= 2) { + 2 + } else { + requests + }; + Ok(SparseBench { total: start.elapsed(), + requests, + auto_requests, root_layer, seed_bytes, node_bytes: scan.node_bytes_seen + node_fetch_bytes, @@ -383,7 +421,7 @@ async fn fetch_blocks( pds: &Url, did: &Did<'static>, cids: &[IpldCid], -) -> Result<(BTreeMap, usize)> { +) -> Result<(BTreeMap, usize, usize)> { let mut out = BTreeMap::new(); let mut bytes = 0usize; @@ -392,6 +430,7 @@ async fn fetch_blocks( .filter(|chunk| !chunk.is_empty()) .map(|chunk| fetch_block_chunk(http.clone(), pds.clone(), did.clone(), chunk.to_vec())) .collect::>(); + let requests = fetches.len(); let fetches = stream::iter(fetches).buffer_unordered(GET_BLOCKS_PARALLELISM); futures::pin_mut!(fetches); @@ -401,7 +440,7 @@ async fn fetch_blocks( out.extend(blocks); } - Ok((out, bytes)) + Ok((out, bytes, requests)) } async fn fetch_block_chunk( @@ -423,7 +462,7 @@ async fn fetch_block_chunk( .into_output() .map_err(|err: XrpcError<_>| miette::miette!("getBlocks failed for {did}: {err}"))?; let bytes = car.body.len(); - let parsed = car::parse_car_blocks(&car.body).await.wrap_err_with(|| { + let parsed = car::parse_car_blocks(car.body).wrap_err_with(|| { let cids = cids .iter() .map(ToString::to_string) diff --git a/src/car.rs b/src/car.rs index a2b04f0..26db1c0 100644 --- a/src/car.rs +++ b/src/car.rs @@ -5,9 +5,9 @@ use bytes::Bytes; use cid::Cid as IpldCid; use miette::{IntoDiagnostic, Result}; -pub(crate) async fn parse_car_blocks(data: &[u8]) -> Result> { +pub(crate) fn parse_car_blocks(data: Bytes) -> Result> { let mut offset = 0; - let Some(header_len) = read_uvarint(data, &mut offset)? else { + let Some(header_len) = read_uvarint(&data, &mut offset)? else { return Err(miette::miette!("empty CAR file")); }; let header_end = offset @@ -19,7 +19,7 @@ pub(crate) async fn parse_car_blocks(data: &[u8]) -> Result Result= section.len() { return Err(miette::miette!("CAR block has no payload for {cid}")); } - blocks.insert(cid, Bytes::copy_from_slice(§ion[block_start..])); + blocks.insert(cid, section.slice(block_start..)); } Ok(blocks) @@ -87,7 +87,7 @@ mod tests { writer.write(cid, b"block".to_vec()).await.unwrap(); writer.finish().await.unwrap(); - let blocks = parse_car_blocks(&buf).await.unwrap(); + let blocks = parse_car_blocks(buf.into()).unwrap(); assert_eq!(blocks.get(&cid).unwrap().as_ref(), b"block"); } }