Daan MeijerandClaude Fable 5.1
91984d3625
a.b.storage's stored 2016 ranges are not contiguous (its 2016 window lies in two article regions); the batch after a gap kept the previous range's end as its start and the harness found no leased range for it. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
386 lines
17 KiB
Rust
386 lines
17 KiB
Rust
//! The Rust indexer over stored chunks (REBUILD.md, part 2 and "Voyager 3"): a range captured by `edge:bench
|
||
//! --capture` or stored in the corpus, or a group's corpus ranges in order through one window (`--group`, the
|
||
//! replay harness behind `replay:corpus`). Same work as the edge-node worker: inflate each chunk, scan it,
|
||
//! aggregate files, finish releases through the window, and every N chunks judge and encode them. `--dump`
|
||
//! writes every output as JSON lines; `--wire` the batches as the worker posts them (`-`: stdout).
|
||
//!
|
||
//! edge-bench <range dir> | --group <group dir> [--from article] [--to article] [--limit ranges] [--rules rules.json]
|
||
//! [--batch-chunks 10] [--idle 5] [--no-finish] [--dump out.jsonl] [--wire batches.jsonl|-] [--releases out.jsonl]
|
||
|
||
use base64::Engine;
|
||
use edge::indexer::{self, Batch, RangeIndexer};
|
||
use edge::js::json_str;
|
||
use edge::{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"),
|
||
}
|
||
}
|
||
|
||
/// The --releases dump: per release chunk its key, name, rule, counter total, poster, time, counts and first subjects.
|
||
fn dump_releases(batch: &Batch, group: &str, out: &mut impl Write) {
|
||
let mut line = String::new();
|
||
for r in &batch.releases {
|
||
line.push_str("{\"group\":");
|
||
json_str(group, &mut line);
|
||
line.push_str(",\"key\":");
|
||
json_str(&r.key, &mut line);
|
||
line.push_str(",\"name\":");
|
||
json_str(&r.name, &mut line);
|
||
line.push_str(&format!(",\"source\":\"{}\",\"fileTotal\":", r.source));
|
||
opt_num(r.file_total, &mut line);
|
||
line.push_str(",\"poster\":");
|
||
json_str(&r.poster, &mut line);
|
||
line.push_str(",\"postedAt\":");
|
||
opt_num(r.posted_at, &mut line);
|
||
line.push_str(",\"postedLast\":");
|
||
opt_num(r.posted_last, &mut line);
|
||
line.push_str(&format!(",\"files\":{},\"segments\":{},\"bytes\":{},\"firstArticle\":{},\"hasPar2\":{},\"subjects\":[", r.files, r.segments, r.bytes, r.first_article, r.has_par2));
|
||
for (i, s) in r.subjects.iter().enumerate() {
|
||
if i > 0 {
|
||
line.push(',');
|
||
}
|
||
json_str(s, &mut line);
|
||
}
|
||
line.push_str("]}\n");
|
||
out.write_all(line.as_bytes()).unwrap();
|
||
line.clear();
|
||
}
|
||
}
|
||
|
||
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();
|
||
};
|
||
let opt = |t: Option<i64>| t.map_or("null".to_string(), |t| t.to_string());
|
||
line.push_str(&format!(
|
||
"{{\"t\":\"b\",\"from\":{},\"next\":{},\"headers\":{},\"postedLow\":{},\"postedHigh\":{}}}",
|
||
batch.from, batch.next, batch.headers, opt(batch.posted_low), opt(batch.posted_high)
|
||
));
|
||
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\":{},\"finished\":{},\"edge\":{},\"partial\":{},\"payload\":\"{}\"",
|
||
r.files,
|
||
r.segments,
|
||
r.bytes,
|
||
r.first_article,
|
||
r.last_article,
|
||
r.readable,
|
||
r.has_par2,
|
||
r.finished,
|
||
r.edge,
|
||
r.partial,
|
||
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);
|
||
}
|
||
}
|
||
|
||
/// A stored range: its directory and article span (range.json).
|
||
struct Range {
|
||
dir: String,
|
||
start: i64,
|
||
end: i64,
|
||
}
|
||
|
||
/// The range in a directory with a range.json, and its group.
|
||
fn range_in(dir: &str) -> Option<(String, Range)> {
|
||
let json = std::fs::read_to_string(format!("{dir}/range.json")).ok()?;
|
||
Some((json_field(&json, "group"), Range { dir: dir.to_string(), start: json_field(&json, "start").parse().ok()?, end: json_field(&json, "end").parse().ok()? }))
|
||
}
|
||
|
||
/// A group directory's complete ranges in article order (the corpus layout: `<group>/<start>-<end>/`).
|
||
fn ranges_in(group_dir: &str) -> (String, Vec<Range>) {
|
||
let mut group = String::new();
|
||
let mut ranges: Vec<Range> = Vec::new();
|
||
for entry in std::fs::read_dir(group_dir).unwrap_or_else(|e| panic!("{group_dir}: {e}")) {
|
||
let path = entry.unwrap().path();
|
||
if !path.is_dir() {
|
||
continue;
|
||
}
|
||
if let Some((g, range)) = range_in(path.to_str().unwrap()) {
|
||
group = g;
|
||
ranges.push(range);
|
||
}
|
||
}
|
||
ranges.sort_by_key(|r| r.start);
|
||
(group, ranges)
|
||
}
|
||
|
||
fn date(ms: Option<i64>) -> String {
|
||
match ms.and_then(chrono::DateTime::from_timestamp_millis) {
|
||
Some(t) => t.format("%Y-%m-%d %H:%M").to_string(),
|
||
None => "-".to_string(),
|
||
}
|
||
}
|
||
|
||
/// Files waiting in finished releases after which a batch goes out whatever its chunks (edge-node's limit).
|
||
const MAX_PENDING_FILES: usize = 100_000;
|
||
|
||
fn main() {
|
||
let args: Vec<String> = std::env::args().collect();
|
||
let usage = "usage: edge-bench <range dir> | --group <group dir> [--from article] [--to article] [--limit ranges] [--rules rules.json] [--batch-chunks N] [--idle N] [--no-finish] [--dump file] [--wire file|-] [--releases file]";
|
||
let flag = |name: &str| args.iter().position(|a| a == name).and_then(|i| args.get(i + 1)).cloned();
|
||
let has = |name: &str| args.iter().any(|a| a == name);
|
||
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: a file, or stdout (`-`) for replay:corpus.
|
||
let wire_path = flag("--wire");
|
||
let wire_stdout = wire_path.as_deref() == Some("-");
|
||
let mut wire_out: Option<Box<dyn Write>> = match wire_path.as_deref() {
|
||
Some("-") => Some(Box::new(std::io::BufWriter::with_capacity(1 << 20, std::io::stdout().lock()))),
|
||
Some(p) => Some(Box::new(std::io::BufWriter::new(std::fs::File::create(p).unwrap()))),
|
||
None => None,
|
||
};
|
||
// One line per release chunk without its payload, with its first subjects (deploy/corpus_sweep.py).
|
||
let mut releases_out = flag("--releases").map(|p| std::io::BufWriter::new(std::fs::File::create(p).unwrap()));
|
||
|
||
// One range, or a group's ranges in order through one window (the window carries over from range to range;
|
||
// the end is only finished after the last one, unless --no-finish: a replay to be continued).
|
||
let (group, ranges) = match flag("--group") {
|
||
Some(dir) => {
|
||
let (group, mut ranges) = ranges_in(&dir);
|
||
// --from: ranges ending before it are skipped (a replay continued); --to: ranges starting after it (one
|
||
// window of the group's history, the corpus holding others); --limit: this many at most.
|
||
if let Some(from) = flag("--from").map(|v| v.parse::<i64>().unwrap()) {
|
||
ranges.retain(|r| r.end >= from);
|
||
}
|
||
if let Some(to) = flag("--to").map(|v| v.parse::<i64>().unwrap()) {
|
||
ranges.retain(|r| r.start <= to);
|
||
}
|
||
if let Some(limit) = flag("--limit").map(|v| v.parse::<usize>().unwrap()) {
|
||
ranges.truncate(limit);
|
||
}
|
||
(group, ranges)
|
||
}
|
||
None => {
|
||
let positional = (1..args.len()).find(|&i| !args[i].starts_with("--") && !args[i - 1].starts_with("--")).map(|i| args[i].clone());
|
||
let dir = positional.expect(usage);
|
||
let (group, range) = range_in(&dir).unwrap_or_else(|| panic!("{dir}: no range.json"));
|
||
(group, vec![range])
|
||
}
|
||
};
|
||
if ranges.is_empty() {
|
||
eprintln!("no complete range found");
|
||
std::process::exit(2);
|
||
}
|
||
let finish = !has("--no-finish");
|
||
let rules = match flag("--rules") {
|
||
Some(path) => {
|
||
let json = std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("{path}: {e}"));
|
||
serde_json::from_str::<edge::lease::RulesJson>(&json).unwrap_or_else(|e| panic!("{path}: {e}")).rules()
|
||
}
|
||
None => {
|
||
let mut priorities = FxHashMap::default();
|
||
priorities.insert(group.clone(), 1);
|
||
names::Rules::defaults(priorities)
|
||
}
|
||
};
|
||
// The report goes to stderr when the wire takes stdout.
|
||
let say = |line: String| {
|
||
if wire_stdout {
|
||
eprintln!("{line}");
|
||
} else {
|
||
println!("{line}");
|
||
}
|
||
};
|
||
|
||
let (start, end) = (ranges[0].start, ranges[ranges.len() - 1].end);
|
||
let mut indexer = RangeIndexer::new(&group, start, end, start, rules, idle);
|
||
|
||
let (mut inflate_s, mut index_s) = (0.0, 0.0);
|
||
let (mut compressed, mut text_bytes, mut chunks) = (0usize, 0usize, 0usize);
|
||
let mut out = (0u64, 0usize, 0usize, 0usize, 0usize, 0usize, 0usize); // headers, releases, files, segments, payload, candidates, samples
|
||
let mut finished = 0usize;
|
||
let mut keys = std::collections::HashSet::new();
|
||
let mut decompressor = libdeflater::Decompressor::new();
|
||
let wall = Instant::now();
|
||
let mut since_batch = 0usize;
|
||
for (ri, range) in ranges.iter().enumerate() {
|
||
let range_wall = Instant::now();
|
||
let headers_before = out.0;
|
||
// After a gap between stored ranges the batch starts at this range (every range ends with a batch taken).
|
||
if ri > 0 {
|
||
indexer.skip_to(range.start);
|
||
}
|
||
let mut files: Vec<String> = std::fs::read_dir(&range.dir).unwrap().map(|e| e.unwrap().file_name().into_string().unwrap()).filter(|f| f.ends_with(".gz")).collect();
|
||
files.sort();
|
||
for (i, file) in files.iter().enumerate() {
|
||
let raw = std::fs::read(format!("{}/{file}", range.dir)).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();
|
||
chunks += 1;
|
||
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);
|
||
since_batch += 1;
|
||
let range_end = i + 1 == files.len();
|
||
let last = range_end && ri + 1 == ranges.len();
|
||
if last && finish {
|
||
indexer.finish();
|
||
}
|
||
if !range_end && since_batch < batch_chunks && indexer.pending() < MAX_PENDING_FILES {
|
||
index_s += t.elapsed().as_secs_f64();
|
||
continue;
|
||
}
|
||
let batch = indexer.take();
|
||
since_batch = 0;
|
||
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();
|
||
finished += r.finished as usize;
|
||
}
|
||
out.5 += batch.candidates.len();
|
||
out.6 += 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();
|
||
w.flush().unwrap();
|
||
}
|
||
if let Some(w) = releases_out.as_mut() {
|
||
dump_releases(&batch, &group, w);
|
||
}
|
||
}
|
||
if ranges.len() > 1 {
|
||
eprintln!(
|
||
"{} {}–{}: {} headers in {:.1} s; frontier {}; {} releases out ({} finished), {} releases and {} files open; RSS {} MB",
|
||
group,
|
||
range.start,
|
||
range.end,
|
||
out.0 - headers_before,
|
||
range_wall.elapsed().as_secs_f64(),
|
||
date(indexer.frontier()),
|
||
out.1,
|
||
finished,
|
||
indexer.open_releases(),
|
||
indexer.open_files(),
|
||
peak_rss_mb()
|
||
);
|
||
}
|
||
}
|
||
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);
|
||
say(format!("{group} {start}–{end}: {} ranges, {chunks} chunks, {} MB compressed, {} MB text", ranges.len(), compressed / 1_000_000, text_bytes / 1_000_000));
|
||
say(format!("{} 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()));
|
||
say(format!(" inflate {inflate_s:.2} s, index {index_s:.2} s; peak RSS {} MB", peak_rss_mb()));
|
||
say(format!(
|
||
" kept: {} release chunks ({} releases, {} finished), {} files, {} segments, payload {:.1} MB ({:.1} B/segment); {} candidates, {} samples",
|
||
out.1,
|
||
keys.len(),
|
||
finished,
|
||
out.2,
|
||
out.3,
|
||
out.4 as f64 / 1e6,
|
||
out.4 as f64 / out.3.max(1) as f64,
|
||
out.5,
|
||
out.6
|
||
));
|
||
}
|