#[path = "../car.rs"]
#[allow(dead_code)]
mod car;
#[path = "../mst.rs"]
#[allow(dead_code)]
mod mst;
#[path = "../sparse_mst.rs"]
#[allow(dead_code)]
mod sparse_mst;
use std::io::Write;
use std::time::{Duration, Instant};
use cid::Cid as IpldCid;
use futures::{StreamExt, stream};
use jacquard_common::types::did::Did;
use jacquard_common::types::string::{Nsid, Rkey};
use miette::{IntoDiagnostic, Result, WrapErr};
use serde::Deserialize;
use smol_str::SmolStr;
use sparse_mst::{SparseScanner, mst_node_layer, sparse_probe_collection, sparse_ranges};
use url::Url;
const GET_BLOCKS_CHUNK: usize = 64;
const GET_BLOCKS_PARALLELISM: usize = 4;
#[derive(Debug)]
struct Args {
collection: String,
pattern: SmolStr,
index_url: Url,
plc_url: Url,
limit: usize,
probe_collection: Option<SmolStr>,
probe_rkey: SmolStr,
}
impl Args {
fn parse() -> Result<Self> {
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 probe_collection = None;
let mut probe_rkey = SmolStr::new("-");
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()?,
"--probe-collection" => probe_collection = Some(SmolStr::new(value)),
"--probe-rkey" => probe_rkey = SmolStr::new(value),
_ => return Err(miette::miette!("unknown argument {arg}")),
}
}
Ok(Self {
collection,
pattern,
index_url,
plc_url,
limit,
probe_collection,
probe_rkey,
})
}
}
#[derive(Debug, Deserialize)]
struct ListReposByCollectionOutput {
repos: Vec<RepoHit>,
}
#[derive(Debug, Deserialize)]
struct RepoHit {
did: String,
}
#[derive(Debug, Deserialize)]
struct DidDoc {
service: Vec<DidService>,
}
#[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,
requests: usize,
auto_requests: usize,
root_layer: Option<usize>,
seed_bytes: usize,
node_bytes: usize,
record_bytes: usize,
node_blocks: usize,
records: usize,
}
#[derive(Debug, Default)]
struct BenchTotals {
repos: usize,
full_ms: u128,
full_bytes: usize,
sparse_ms: u128,
sparse_bytes: usize,
}
impl BenchTotals {
fn add(&mut self, full: &FullBench, sparse: &SparseBench) {
self.repos += 1;
self.full_ms += full.fetch.as_millis() + full.parse_and_walk.as_millis();
self.full_bytes += full.bytes;
self.sparse_ms += sparse.total.as_millis();
self.sparse_bytes += sparse.seed_bytes + sparse.node_bytes + sparse.record_bytes;
}
}
#[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_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();
for did in repos {
let pds = match resolve_pds(&http, &args.plc_url, &did).await {
Ok(pds) => pds,
Err(err) => {
eprintln!("skipping {did}: resolve failed: {err:?}");
continue;
}
};
let full = match bench_full(&http, &pds, &did, &ranges).await {
Ok(full) => full,
Err(err) => {
eprintln!("skipping {did}: full bench failed: {err:?}");
continue;
}
};
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:?}");
continue;
}
};
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.requests,
sparse.auto_requests,
sparse
.root_layer
.map(|layer| layer.to_string())
.unwrap_or_default(),
sparse.seed_bytes,
sparse.node_bytes,
sparse.record_bytes,
sparse.node_blocks,
sparse.records,
);
totals.add(&full, &sparse);
}
std::io::stdout().flush().into_diagnostic()?;
eprintln!(
"summary: repos={}, full_ms={}, full_bytes={}, sparse_ms={}, sparse_bytes={}",
totals.repos, totals.full_ms, totals.full_bytes, totals.sparse_ms, totals.sparse_bytes
);
Ok(())
}
async fn list_repos(http: &reqwest::Client, args: &Args) -> Result<Vec<Did>> {
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::<ListReposByCollectionOutput>()
.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) -> Result<Url> {
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::<DidDoc>()
.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,
ranges: &[sparse_mst::KeyRange],
) -> Result<FullBench> {
let fetch_start = Instant::now();
let mut url = pds.join("xrpc/com.atproto.sync.getRepo").into_diagnostic()?;
url.query_pairs_mut().append_pair("did", did.as_str());
let car_bytes = http
.get(url)
.header(reqwest::header::ACCEPT, "application/vnd.ipld.car")
.send()
.await
.into_diagnostic()?
.error_for_status()
.map_err(|err| miette::miette!("getRepo failed for {did}: {err}"))?
.bytes()
.await
.into_diagnostic()?;
let fetch = fetch_start.elapsed();
let bytes = car_bytes.len();
let parse_start = Instant::now();
let parsed = car::parse_car(car_bytes, false)
.wrap_err_with(|| format!("parse getRepo CAR for {did} ({bytes} bytes)"))?;
let blocks = parsed.blocks.map().len();
let root_bytes = parsed
.blocks
.get(&parsed.root)
.ok_or_else(|| miette::miette!("root block missing from getRepo car"))?;
let root_commit = mst::AtpCommit::from_cbor(root_bytes)?;
let leaves = mst::leaves(parsed.blocks.map(), root_commit.data)?;
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,
patterns: &[SmolStr],
probe_collection: Option<&str>,
probe_rkey: &str,
) -> Result<SparseBench> {
let start = Instant::now();
let ranges = sparse_ranges(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 probe_collection: Nsid = Nsid::new_owned(probe_collection.as_str()).into_diagnostic()?;
let probe_rkey: Rkey = Rkey::new_owned(probe_rkey).into_diagnostic()?;
let mut url = pds.join("xrpc/com.atproto.sync.getRecord").into_diagnostic()?;
url.query_pairs_mut()
.append_pair("did", did.as_str())
.append_pair("collection", probe_collection.as_str())
.append_pair("rkey", probe_rkey.as_ref());
let seed_bytes_buf = http
.get(url)
.header(reqwest::header::ACCEPT, "application/vnd.ipld.car")
.send()
.await
.into_diagnostic()?
.error_for_status()
.map_err(|err| miette::miette!("getRecord seed failed for {did}: {err}"))?
.bytes()
.await
.into_diagnostic()?;
let seed_bytes = seed_bytes_buf.len();
let parsed = car::parse_commit_car(seed_bytes_buf, false)
.wrap_err_with(|| format!("parse getRecord seed CAR for {did} ({seed_bytes} bytes)"))?;
let root_bytes = parsed
.blocks
.get(&parsed.root)
.ok_or_else(|| miette::miette!("root block missing from sparse seed car"))?;
let root_commit = mst::AtpCommit::from_cbor(root_bytes)?;
let root_cid = root_commit.data;
let root_layer = parsed
.blocks
.get(&root_cid)
.map(|bytes| mst_node_layer(bytes))
.transpose()?
.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, chunk_requests) =
fetch_blocks(http, pds, did, &missing).await?;
node_fetch_bytes += bytes;
requests += chunk_requests;
scanner.insert_blocks(blocks);
}
}
};
let mut blocks = scanner.take_blocks();
let missing_records: Vec<IpldCid> = scan
.leaves
.iter()
.filter_map(|(_, cid)| blocks.get(cid).is_none().then_some(*cid))
.collect();
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,
record_bytes,
node_blocks: scan.node_blocks_seen,
records: scan.leaves.len(),
})
}
async fn fetch_blocks(
http: &reqwest::Client,
pds: &Url,
did: &Did,
cids: &[IpldCid],
) -> Result<(car::Blocks<car::CommitBlocks>, usize, usize)> {
let mut out = car::Blocks::default();
let mut bytes = 0usize;
let fetches = cids
.chunks(GET_BLOCKS_CHUNK)
.filter(|chunk| !chunk.is_empty())
.map(|chunk| fetch_block_chunk(http.clone(), pds.clone(), did.clone(), chunk.to_vec()))
.collect::<Vec<_>>();
let requests = fetches.len();
let fetches = stream::iter(fetches).buffer_unordered(GET_BLOCKS_PARALLELISM);
futures::pin_mut!(fetches);
while let Some(chunk) = fetches.next().await {
let (blocks, chunk_bytes) = chunk?;
bytes += chunk_bytes;
out.extend(blocks);
}
Ok((out, bytes, requests))
}
async fn fetch_block_chunk(
http: reqwest::Client,
pds: Url,
did: Did,
cids: Vec<IpldCid>,
) -> Result<(car::Blocks<car::CommitBlocks>, usize)> {
let mut url = pds.join("xrpc/com.atproto.sync.getBlocks").into_diagnostic()?;
{
let mut query = url.query_pairs_mut();
query.append_pair("did", did.as_str());
for cid in &cids {
query.append_pair("cids", &cid.to_string());
}
}
let car_bytes = http
.get(url)
.header(reqwest::header::ACCEPT, "application/vnd.ipld.car")
.send()
.await
.into_diagnostic()?
.error_for_status()
.map_err(|err| miette::miette!("getBlocks failed for {did}: {err}"))?
.bytes()
.await
.into_diagnostic()?;
let bytes = car_bytes.len();
let parsed = car::parse_car_blocks(car_bytes, false).wrap_err_with(|| {
let cids = cids
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
.join(",");
format!(
"parse getBlocks CAR for {did} ({} requested, {} bytes, cids={cids})",
cids.len(),
bytes
)
})?;
Ok((parsed, bytes))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn benchmark_api_json_fixtures_decode() {
let list: ListReposByCollectionOutput = serde_json::from_value(serde_json::json!({
"repos": [{"did": "did:web:example.com"}],
}))
.unwrap();
assert_eq!(list.repos.len(), 1);
assert_eq!(list.repos[0].did, "did:web:example.com");
let doc: DidDoc = serde_json::from_value(serde_json::json!({
"service": [{
"type": "AtprotoPersonalDataServer",
"serviceEndpoint": "https://pds.example",
}],
}))
.unwrap();
assert_eq!(doc.service.len(), 1);
assert_eq!(doc.service[0].kind, "AtprotoPersonalDataServer");
assert_eq!(
doc.service[0].service_endpoint.as_str(),
"https://pds.example/"
);
}
}