//! 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]
);
}