Something went wrong. Try again.
very fast at protocol indexer with flexible filtering, xrpc queries, cursor-backed event stream, and more, built on fjall
rust fjall at-protocol atproto indexer
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187//! cargo run --release --example chunk_load -- http://127.0.0.1:3000 11 stream|chunks [in_flight]
use std::time::{Duration, Instant};
use futures::StreamExt;use tokio_websockets::ClientBuilder;
#[derive(Default, Debug)]struct Run { took: Duration, chunk_bytes: u64, chunk_events: u64, ws_bytes: u64, ws_events: u64, error: Option<String>,}
/// wire format is in docs/api/stream.md#chunk-format#[derive(serde::Deserialize)]#[allow(dead_code)]struct ChunkEvent<'a> { cid: Option<cid::Cid>, did: &'a str, rev: &'a str, live: bool, rkey: &'a str, action: &'a str, #[serde(default, borrow)] record: Option<&'a [u8]>, collection: &'a str,}
fn varint(star: &[u8], at: &mut usize) -> usize { let (mut n, mut shift) = (0, 0); loop { let byte = star[*at]; *at += 1; n |= usize::from(byte & 0x7f) << shift; if byte & 0x80 == 0 { return n; } shift += 7; }}
fn chunk_events(body: &[u8]) -> u64 { let star = zstd::decode_all(body).unwrap(); assert_eq!(star[..4], [0x2a, 0x6c, 0x01, 0x00]); let mut at = 4; at += varint(&star, &mut at); let mut events = 0; while at < star.len() { at += varint(&star, &mut at); let len = varint(&star, &mut at); let _: ChunkEvent = serde_ipld_dagcbor::from_slice(&star[at..at + len]).unwrap(); at += len; events += 1; } events}
async fn subscribe(base: String, mode: String, in_flight: usize, head: u64) -> Run { let mut run = Run::default(); let opened = Instant::now(); let mut cursor = 0; let mut stream = String::new(); if mode == "chunks" { let http = reqwest::Client::new(); let manifest: serde_json::Value = http .get(format!("{base}/stream/manifest")) .send() .await .unwrap() .json() .await .unwrap(); stream = manifest["stream"].as_str().unwrap().to_owned(); let shas: Vec<String> = manifest["chunks"] .as_array() .unwrap() .iter() .map(|c| c["sha256"].as_str().unwrap().to_owned()) .collect(); let mut bodies = futures::stream::iter(shas) .map(|sha| { let http = http.clone(); let base = base.clone(); async move { http.get(format!("{base}/stream/chunk/{sha}")) .send() .await .unwrap() .bytes() .await .unwrap() } }) .buffered(in_flight); while let Some(body) = bodies.next().await { run.chunk_bytes += body.len() as u64; let events = tokio::task::spawn_blocking(move || chunk_events(&body)) .await .unwrap(); run.chunk_events += events; } cursor = manifest["through"] .as_u64() .map_or(0, |through| through + 1); } if cursor <= head { let ws_base = base.replacen("http", "ws", 1); let mut uri = format!("{ws_base}/stream?cursor={cursor}"); if !stream.is_empty() { uri += &format!("&stream={stream}"); } let (mut ws, _) = match ClientBuilder::from_uri(uri.parse().unwrap()) .connect() .await { Ok(c) => c, Err(e) => { run.error = Some(e.to_string()); return run; } }; while let Some(msg) = ws.next().await { let Ok(msg) = msg else { break }; let Some(text) = msg.as_text() else { continue }; run.ws_events += 1; run.ws_bytes += text.len() as u64; let frame: serde_json::Value = serde_json::from_str(text).unwrap(); if frame["type"] == "error" { run.error = Some(text.to_owned()); break; } if frame["id"].as_u64().is_some_and(|id| id >= head) { break; } } } run.took = opened.elapsed(); run}
#[tokio::main]async fn main() { let mut args = std::env::args().skip(1); let base = args.next().unwrap(); let n: usize = args.next().unwrap().parse().unwrap(); let mode = args.next().unwrap(); let in_flight: usize = args.next().map_or(2, |n| n.parse().unwrap()); let head: serde_json::Value = reqwest::get(format!("{base}/stream/head")) .await .unwrap() .json() .await .unwrap(); let head = head["id"].as_u64().unwrap(); let started = Instant::now(); let runs = futures::future::join_all( (0..n).map(|_| tokio::spawn(subscribe(base.clone(), mode.clone(), in_flight, head))), ) .await; let wall = started.elapsed(); let mut took = Vec::new(); for run in runs { let run = run.unwrap(); println!( "{{\"took_s\":{:.3},\"chunk_bytes\":{},\"chunk_events\":{},\"ws_bytes\":{},\"ws_events\":{},\"error\":{}}}", run.took.as_secs_f64(), run.chunk_bytes, run.chunk_events, run.ws_bytes, run.ws_events, serde_json::to_string(&run.error).unwrap() ); took.push(run.took.as_secs_f64()); } took.sort_by(f64::total_cmp); println!( "{{\"total\":true,\"mode\":\"{mode}\",\"in_flight\":{in_flight},\"n\":{n},\"head\":{head},\"wall_s\":{:.3},\"median_s\":{:.3},\"max_s\":{:.3}}}", wall.as_secs_f64(), took[took.len() / 2], took[took.len() - 1] );}