Daan MeijerandClaude Opus 5.5
c7b4ebe470
libdeflate inflates a chunk whole, in half the time of zlib-rs (1.09 -> 0.51 s per boneless range); memchr finds line ends and tabs with SIMD. The per-article timers are gone. Boneless 2026 ~700k headers/s, ath 2013 ~920k; all three captures still identical to the TypeScript. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
240 lines
9.3 KiB
Rust
240 lines
9.3 KiB
Rust
//! The edge worker's indexing in Rust, on a range captured by `edge:bench
|
||
//! --capture` (REBUILD.md, part 2: TypeScript or Rust). Same work as
|
||
//! edge:bench's replay: inflate each chunk, scan it, aggregate files, cut off
|
||
//! idle ones, and every N chunks judge and encode the releases. `--dump`
|
||
//! writes every output as JSON lines for tools/edge-bench-rs/check.ts.
|
||
//!
|
||
//! edge-bench-rs <dir> [--batch-chunks 10] [--idle 5] [--dump out.jsonl]
|
||
|
||
mod codec;
|
||
mod indexer;
|
||
mod js;
|
||
mod names;
|
||
|
||
use base64::Engine;
|
||
use indexer::{Batch, RangeIndexer};
|
||
use js::json_str;
|
||
use rustc_hash::FxHashMap;
|
||
use std::io::Write;
|
||
use std::time::Instant;
|
||
|
||
// Many small strings (message-ids, names): the system allocator was a measurable share.
|
||
#[global_allocator]
|
||
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
|
||
|
||
/// A chunk as the news server sent it (zlib; gzip also taken), whole: libdeflate wants the output
|
||
/// size, so it grows until the text fits (~3.6 times the compressed size on boneless).
|
||
fn inflate(raw: &[u8], d: &mut libdeflater::Decompressor) -> Vec<u8> {
|
||
let mut out = vec![0u8; raw.len() * 4 + 4096];
|
||
loop {
|
||
let result = if raw.starts_with(&[0x1f, 0x8b]) { d.gzip_decompress(raw, &mut out) } else { d.zlib_decompress(raw, &mut out) };
|
||
match result {
|
||
Ok(n) => {
|
||
out.truncate(n);
|
||
return out;
|
||
}
|
||
Err(libdeflater::DecompressionError::InsufficientSpace) => out.resize(out.len() * 2, 0),
|
||
Err(e) => panic!("inflate: {e}"),
|
||
}
|
||
}
|
||
}
|
||
|
||
fn json_field(json: &str, name: &str) -> String {
|
||
let at = json.find(&format!("\"{name}\":")).unwrap() + name.len() + 3;
|
||
let rest = json[at..].trim_start();
|
||
if let Some(s) = rest.strip_prefix('"') {
|
||
s[..s.find('"').unwrap()].to_string()
|
||
} else {
|
||
rest[..rest.find([',', '}']).unwrap()].trim().to_string()
|
||
}
|
||
}
|
||
|
||
fn peak_rss_mb() -> u64 {
|
||
let status = std::fs::read_to_string("/proc/self/status").unwrap_or_default();
|
||
status.lines().find(|l| l.starts_with("VmHWM:")).and_then(|l| l.split_whitespace().nth(1)).and_then(|k| k.parse::<u64>().ok()).unwrap_or(0) / 1024
|
||
}
|
||
|
||
fn opt_num(n: Option<i64>, out: &mut String) {
|
||
match n {
|
||
Some(n) => out.push_str(&n.to_string()),
|
||
None => out.push_str("null"),
|
||
}
|
||
}
|
||
|
||
fn dump_batch(batch: &Batch, out: &mut impl Write) {
|
||
let b64 = base64::engine::general_purpose::STANDARD;
|
||
let mut line = String::new();
|
||
let mut emit = |line: &mut String| {
|
||
line.push('\n');
|
||
out.write_all(line.as_bytes()).unwrap();
|
||
line.clear();
|
||
};
|
||
line.push_str(&format!("{{\"t\":\"b\",\"from\":{},\"next\":{},\"headers\":{}}}", batch.from, batch.next, batch.headers));
|
||
emit(&mut line);
|
||
let release = |r: &indexer::ReleaseChunk, kind: &str, line: &mut String| {
|
||
line.push_str(&format!("{{\"t\":\"{kind}\",\"key\":"));
|
||
json_str(&r.key, line);
|
||
line.push_str(",\"name\":");
|
||
json_str(&r.name, line);
|
||
line.push_str(&format!(",\"source\":\"{}\",\"fileTotal\":", r.source));
|
||
opt_num(r.file_total, line);
|
||
line.push_str(",\"poster\":");
|
||
json_str(&r.poster, line);
|
||
line.push_str(",\"groups\":[");
|
||
for (i, g) in r.groups.iter().enumerate() {
|
||
if i > 0 {
|
||
line.push(',');
|
||
}
|
||
json_str(g, line);
|
||
}
|
||
line.push_str("],\"postedAt\":");
|
||
opt_num(r.posted_at, line);
|
||
line.push_str(&format!(
|
||
",\"files\":{},\"segments\":{},\"bytes\":{},\"firstArticle\":{},\"lastArticle\":{},\"readable\":{},\"hasPar2\":{},\"payload\":\"{}\"",
|
||
r.files,
|
||
r.segments,
|
||
r.bytes,
|
||
r.first_article,
|
||
r.last_article,
|
||
r.readable,
|
||
r.has_par2,
|
||
b64.encode(&r.payload)
|
||
));
|
||
};
|
||
for r in &batch.releases {
|
||
release(r, "r", &mut line);
|
||
line.push('}');
|
||
emit(&mut line);
|
||
}
|
||
for (r, rule, key) in &batch.samples {
|
||
release(r, "s", &mut line);
|
||
line.push_str(&format!(",\"rule\":\"{rule}\",\"ruleKey\":"));
|
||
json_str(key, &mut line);
|
||
line.push('}');
|
||
emit(&mut line);
|
||
}
|
||
for c in &batch.candidates {
|
||
line.push_str("{\"t\":\"c\",\"releaseKey\":");
|
||
json_str(&c.release_key, &mut line);
|
||
line.push_str(&format!(",\"kind\":\"{}\",\"messageId\":", c.kind));
|
||
json_str(&c.message_id, &mut line);
|
||
line.push_str(",\"file\":");
|
||
json_str(&c.file, &mut line);
|
||
line.push('}');
|
||
emit(&mut line);
|
||
}
|
||
for (rule, key, articles, files) in &batch.counts {
|
||
line.push_str(&format!("{{\"t\":\"n\",\"rule\":\"{rule}\",\"key\":"));
|
||
json_str(key, &mut line);
|
||
line.push_str(&format!(",\"articles\":{articles},\"files\":{files}}}"));
|
||
emit(&mut line);
|
||
}
|
||
for f in &batch.fragments {
|
||
line.push_str("{\"t\":\"f\",\"fileKey\":");
|
||
json_str(&f.file_key, &mut line);
|
||
line.push_str(",\"name\":");
|
||
json_str(&f.name, &mut line);
|
||
line.push_str(&format!(",\"totalParts\":{},\"postedAt\":", f.total_parts));
|
||
opt_num(f.posted_at, &mut line);
|
||
line.push_str(",\"segments\":[");
|
||
for (i, (n, id, bytes, article)) in f.segments.iter().enumerate() {
|
||
if i > 0 {
|
||
line.push(',');
|
||
}
|
||
line.push_str(&format!("[{n},"));
|
||
json_str(id, &mut line);
|
||
line.push_str(&format!(",{bytes},{article}]"));
|
||
}
|
||
line.push_str("]}");
|
||
emit(&mut line);
|
||
}
|
||
}
|
||
|
||
fn main() {
|
||
let args: Vec<String> = std::env::args().collect();
|
||
let dir = args.get(1).expect("usage: edge-bench-rs <dir> [--batch-chunks N] [--idle N] [--dump file]").clone();
|
||
let flag = |name: &str| args.iter().position(|a| a == name).and_then(|i| args.get(i + 1)).cloned();
|
||
let batch_chunks: usize = flag("--batch-chunks").map(|v| v.parse().unwrap()).unwrap_or(10);
|
||
let idle: u64 = flag("--idle").map(|v| v.parse().unwrap()).unwrap_or(5);
|
||
let mut dump = flag("--dump").map(|p| std::io::BufWriter::new(std::fs::File::create(p).unwrap()));
|
||
|
||
let range = std::fs::read_to_string(format!("{dir}/range.json")).unwrap();
|
||
let group = json_field(&range, "group");
|
||
let start: i64 = json_field(&range, "start").parse().unwrap();
|
||
let end: i64 = json_field(&range, "end").parse().unwrap();
|
||
let mut files: Vec<String> = std::fs::read_dir(&dir).unwrap().map(|e| e.unwrap().file_name().into_string().unwrap()).filter(|f| f.ends_with(".gz")).collect();
|
||
files.sort();
|
||
|
||
let mut priorities = FxHashMap::default();
|
||
priorities.insert(group.clone(), 1);
|
||
let mut indexer = RangeIndexer::new(&group, start, end, priorities, idle);
|
||
|
||
let (mut inflate_s, mut index_s) = (0.0, 0.0);
|
||
let (mut compressed, mut text_bytes) = (0usize, 0usize);
|
||
let mut out = (0u64, 0usize, 0usize, 0usize, 0usize, 0usize, 0usize, 0usize); // headers, releases, files, segments, payload, candidates, fragments, samples
|
||
let mut keys = std::collections::HashSet::new();
|
||
let mut decompressor = libdeflater::Decompressor::new();
|
||
let wall = Instant::now();
|
||
for (i, file) in files.iter().enumerate() {
|
||
let raw = std::fs::read(format!("{dir}/{file}")).unwrap();
|
||
compressed += raw.len();
|
||
let t = Instant::now();
|
||
let data = inflate(&raw, &mut decompressor);
|
||
inflate_s += t.elapsed().as_secs_f64();
|
||
text_bytes += data.len();
|
||
let t = Instant::now();
|
||
let text = String::from_utf8_lossy(&data);
|
||
let next: i64 = file.split('-').nth(2).unwrap().trim_end_matches(".gz").parse::<i64>().unwrap() + 1;
|
||
indexer.add_text(&text, next);
|
||
let last = i + 1 == files.len();
|
||
if last {
|
||
indexer.finish();
|
||
}
|
||
if !last && (i + 1) % batch_chunks != 0 {
|
||
index_s += t.elapsed().as_secs_f64();
|
||
continue;
|
||
}
|
||
let batch = indexer.take();
|
||
index_s += t.elapsed().as_secs_f64();
|
||
out.0 += batch.headers;
|
||
out.1 += batch.releases.len();
|
||
for r in &batch.releases {
|
||
keys.insert(r.key.clone());
|
||
out.2 += r.files;
|
||
out.3 += r.segments;
|
||
out.4 += r.payload.len();
|
||
}
|
||
out.5 += batch.candidates.len();
|
||
out.6 += batch.fragments.len();
|
||
out.7 += batch.samples.len();
|
||
if let Some(d) = dump.as_mut() {
|
||
dump_batch(&batch, d);
|
||
}
|
||
}
|
||
let total = inflate_s + index_s;
|
||
#[allow(static_mut_refs)]
|
||
let p = unsafe { indexer::PROF };
|
||
eprintln!(" of which aggregate {:.2} s, evict+cutoff {:.2} s, judge+encode {:.2} s", p[3] as f64 / 1e9, p[4] as f64 / 1e9, p[5] as f64 / 1e9);
|
||
println!("{group} {start}–{end}: {} chunks, {} MB compressed, {} MB text", files.len(), compressed / 1_000_000, text_bytes / 1_000_000);
|
||
println!(
|
||
"{} headers in {:.2} s ({:.0}/s on one core; wall {:.2} s)",
|
||
out.0,
|
||
total,
|
||
out.0 as f64 / total,
|
||
wall.elapsed().as_secs_f64()
|
||
);
|
||
println!(" inflate {inflate_s:.2} s, index {index_s:.2} s; peak RSS {} MB", peak_rss_mb());
|
||
println!(
|
||
" kept: {} release chunks ({} releases), {} files, {} segments, payload {:.1} MB ({:.1} B/segment); {} candidates, {} fragments, {} samples",
|
||
out.1,
|
||
keys.len(),
|
||
out.2,
|
||
out.3,
|
||
out.4 as f64 / 1e6,
|
||
out.4 as f64 / out.3.max(1) as f64,
|
||
out.5,
|
||
out.6,
|
||
out.7
|
||
);
|
||
}
|