usenet-indexer/tools/edge-bench-rs/src/main.rs

//! 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] [--wire batches.jsonl]

use base64::Engine;
use edge_bench_rs::indexer::{self, Batch, RangeIndexer};
use edge_bench_rs::js::json_str;
use edge_bench_rs::{names, wire};
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}\",\"upload\":{},\"key\":", r.upload));
        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(",\"postedLast\":");
        opt_num(r.posted_last, 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(&format!("{{\"t\":\"c\",\"upload\":{},\"releaseKey\":", c.upload));
        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()));
    // Each batch as the worker posts it (wire.rs), one JSON per line: check.ts compares it with toWire.
    let mut wire_out = flag("--wire").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, start, names::Rules::defaults(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);
        }
        if let Some(w) = wire_out.as_mut() {
            w.write_all(&wire::to_json(&batch)).unwrap();
            w.write_all(b"\n").unwrap();
        }
    }
    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
    );
}