edge-node: the edge worker in Rust

tools/edge-bench-rs/src/bin/edge-node.rs: leases ranges, reads the
coordinator's framed stream (one thread reads and inflates a frame
ahead, libdeflate), indexes with the shared RangeIndexer and posts
batches as gzipped JSON in order on a thread of its own (one batch may
wait). As commands/edge_node.ts: retries for network trouble and 5xx,
409 = lease lost, a batch once 500k articles or 100k files wait, a
chunk that doesn't inflate reported (fetched plain next time), the
range given back on SIGTERM/SIGINT, --max-mbit shared by the threads;
plus a heartbeat when no batch went out for a minute. One range per
thread (--threads, default one per core).

End to end against the web server (a.b.models 2017215-4017214 from
the capture, spooled): the server stored the same from both workers
(38,228 releases with every column, chunks, 481 candidates, rule
counts, 57 fragments with their JSON). Rust 9 s / 5.1 CPU-s, TypeScript
17 s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

tools/edge-bench-rs/src/bin/edge-node.rs

@@ -0,0 +1,533 @@
+//! The edge worker in Rust (REBUILD.md, part 2; commands/edge_node.ts): leases
+//! ranges from the coordinator, reads their XOVER chunks from its stream,
+//! indexes them (the same RangeIndexer as edge-bench-rs, kept identical to
+//! the TypeScript by check.ts) and posts the batches. One range per thread.
+//!
+//!   COORDINATOR_URL=… NODE_TOKEN=… edge-node [--threads N] [--batch-articles 500000] [--ranges N] [--max-mbit M]
+//!
+//! SIGTERM/SIGINT: each thread gives its range back after the current chunk
+//! (the next worker resumes at the last stored batch).
+
+use edge_bench_rs::indexer::{Batch, RangeIndexer};
+use edge_bench_rs::names::Rules;
+use edge_bench_rs::wire;
+use flate2::write::GzEncoder;
+use serde::Deserialize;
+use std::collections::HashMap;
+use std::io::{Read, Write};
+use std::sync::atomic::{AtomicBool, AtomicI64, Ordering};
+use std::sync::mpsc::{sync_channel, Receiver};
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
+
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
+/// A batch is posted once this many files (completed, or cut at an edge) wait, whatever its articles (commands/edge_node.ts).
+const MAX_PENDING_FILES: usize = 100_000;
+const RETRIES: u32 = 8;
+const IDLE: Duration = Duration::from_secs(30);
+/// Without a post for this long, the range's lease is kept alive by a heartbeat.
+const HEARTBEAT: Duration = Duration::from_secs(60);
+
+fn log(level: &str, message: &str) {
+    println!("[ {level} ] {message}");
+}
+
+#[derive(Deserialize)]
+struct LeaseJson {
+    range: RangeJson,
+    rules: RulesJson,
+}
+
+#[derive(Deserialize)]
+struct RangeJson {
+    id: i64,
+    group: String,
+    start: i64,
+    end: i64,
+    next: i64,
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct RulesJson {
+    version: i64,
+    priorities: HashMap<String, u32>,
+    #[serde(default)]
+    dropped_signatures: Vec<String>,
+    #[serde(default)]
+    dropped_tools: Option<Vec<String>>,
+    #[serde(default)]
+    per_article_max_parts: Option<u32>,
+    #[serde(default)]
+    sample_percent: Option<f64>,
+    #[serde(default)]
+    drop_executables: Option<bool>,
+}
+
+impl RulesJson {
+    fn rules(self) -> Rules {
+        let mut rules = Rules::defaults(self.priorities.into_iter().collect());
+        rules.version = self.version;
+        rules.dropped_signatures = self.dropped_signatures.into_iter().collect();
+        if let Some(tools) = self.dropped_tools {
+            rules.dropped_tools = tools.into_iter().collect();
+        }
+        if let Some(n) = self.per_article_max_parts {
+            rules.per_article_max_parts = n;
+        }
+        if let Some(p) = self.sample_percent {
+            rules.sample_percent = p;
+        }
+        if let Some(e) = self.drop_executables {
+            rules.drop_executables = e;
+        }
+        rules
+    }
+}
+
+#[derive(Deserialize)]
+struct StoredJson {
+    releases: u64,
+    candidates: u64,
+}
+
+#[derive(Debug)]
+enum ApiError {
+    /// 409: the range is no longer this node's (or the batch was stale): stop it, don't give it back.
+    LeaseLost(String),
+    Other(String),
+}
+
+impl std::fmt::Display for ApiError {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            ApiError::LeaseLost(m) | ApiError::Other(m) => f.write_str(m),
+        }
+    }
+}
+
+/// The coordinator's node API, with retries for network trouble and 5xx (as the TypeScript Api).
+struct Api {
+    base: String,
+    token: String,
+    agent: ureq::Agent,
+}
+
+impl Api {
+    fn new(base: &str, token: &str) -> Self {
+        let agent: ureq::Agent = ureq::Agent::config_builder()
+            .http_status_as_error(false)
+            .timeout_connect(Some(Duration::from_secs(30)))
+            .timeout_recv_response(Some(Duration::from_secs(600)))
+            .build()
+            .into();
+        Api { base: base.trim_end_matches('/').to_string(), token: token.to_string(), agent }
+    }
+
+    /// A request, retried; `body`: JSON bytes (gzipped when `gzip`).
+    fn request(&self, method: &str, path: &str, body: Option<&[u8]>, gzip: bool, accept: &str) -> Result<ureq::http::Response<ureq::Body>, ApiError> {
+        let url = format!("{}/api/v1/node/{path}", self.base);
+        let auth = format!("Bearer {}", self.token);
+        for attempt in 1..=RETRIES {
+            let result = match method {
+                "GET" => self.agent.get(&url).header("Authorization", &auth).header("Accept", accept).call(),
+                _ => {
+                    let mut request = self.agent.post(&url).header("Authorization", &auth).header("Accept", accept).header("Content-Type", "application/json");
+                    if gzip {
+                        request = request.header("Content-Encoding", "gzip");
+                    }
+                    request.send(body.unwrap_or(b"{}"))
+                }
+            };
+            match result {
+                Ok(mut response) => {
+                    let status = response.status().as_u16();
+                    if status == 409 {
+                        let text = response.body_mut().read_to_string().unwrap_or_default();
+                        let message = serde_json::from_str::<serde_json::Value>(&text).ok().and_then(|v| v["error"].as_str().map(str::to_string)).unwrap_or_else(|| "lease lost".into());
+                        return Err(ApiError::LeaseLost(message));
+                    }
+                    if status < 500 {
+                        return Ok(response);
+                    }
+                    if attempt == RETRIES {
+                        return Err(ApiError::Other(format!("{path}: HTTP {status}")));
+                    }
+                }
+                Err(error) => {
+                    if attempt == RETRIES {
+                        return Err(ApiError::Other(format!("{path}: {error}")));
+                    }
+                }
+            }
+            std::thread::sleep(Duration::from_millis(60_000.min(1000 * 2u64.pow(attempt))));
+        }
+        unreachable!()
+    }
+
+    fn lease(&self) -> Result<Option<LeaseJson>, ApiError> {
+        let mut response = self.request("POST", "lease", None, false, "application/json")?;
+        match response.status().as_u16() {
+            204 => Ok(None),
+            200 => {
+                let text = response.body_mut().read_to_string().map_err(|e| ApiError::Other(format!("lease: {e}")))?;
+                serde_json::from_str(&text).map(Some).map_err(|e| ApiError::Other(format!("lease: {e}")))
+            }
+            status => Err(ApiError::Other(format!("lease: HTTP {status}"))),
+        }
+    }
+
+    fn stream(&self, range: i64) -> Result<Box<dyn Read + Send>, ApiError> {
+        let response = self.request("GET", &format!("ranges/{range}/stream"), None, false, "application/octet-stream")?;
+        if response.status().as_u16() != 200 {
+            return Err(ApiError::Other(format!("stream: HTTP {}", response.status())));
+        }
+        Ok(Box::new(response.into_body().into_reader()))
+    }
+
+    /// Posts a batch (gzipped JSON): the releases and candidates stored, and the bytes posted.
+    fn batch(&self, range: i64, batch: &Batch) -> Result<(u64, u64, usize), ApiError> {
+        let json = wire::to_json(batch);
+        let mut gz = GzEncoder::new(Vec::with_capacity(json.len() / 3), flate2::Compression::new(1));
+        gz.write_all(&json).unwrap();
+        let body = gz.finish().unwrap();
+        let mut response = self.request("POST", &format!("ranges/{range}/batches"), Some(&body), true, "application/json")?;
+        if response.status().as_u16() != 200 {
+            return Err(ApiError::Other(format!("batch: HTTP {}", response.status())));
+        }
+        let text = response.body_mut().read_to_string().map_err(|e| ApiError::Other(format!("batch: {e}")))?;
+        let stored: StoredJson = serde_json::from_str(&text).map_err(|e| ApiError::Other(format!("batch: {e}")))?;
+        Ok((stored.releases, stored.candidates, body.len()))
+    }
+
+    fn call(&self, range: i64, action: &str) -> Result<(), ApiError> {
+        let response = self.request("POST", &format!("ranges/{range}/{action}"), None, false, "application/json")?;
+        match response.status().as_u16() {
+            200..=299 => Ok(()),
+            status => Err(ApiError::Other(format!("{action}: HTTP {status}"))),
+        }
+    }
+
+    fn corrupt(&self, range: i64, from: i64) -> Result<(), ApiError> {
+        let body = format!("{{\"from\":{from}}}");
+        let response = self.request("POST", &format!("ranges/{range}/corrupt"), Some(body.as_bytes()), false, "application/json")?;
+        match response.status().as_u16() {
+            200..=299 => Ok(()),
+            status => Err(ApiError::Other(format!("corrupt: HTTP {status}"))),
+        }
+    }
+}
+
+/// One frame of the range's stream: a chunk inflated (or why it didn't), or a keep-alive.
+enum Frame {
+    Chunk { from: i64, to: i64, text: Result<Vec<u8>, String> },
+    KeepAlive,
+    End,
+    Failed(String),
+}
+
+/// Reads frames (8 B from, 8 B to, 4 B length, data) and inflates them, one frame ahead of the indexer.
+fn read_frames(mut stream: Box<dyn Read + Send>, bytes_per_second: Option<f64>) -> Receiver<Frame> {
+    let (tx, rx) = sync_channel::<Frame>(2);
+    std::thread::spawn(move || {
+        let mut decompressor = libdeflater::Decompressor::new();
+        let started = Instant::now();
+        let mut read_bytes = 0f64;
+        loop {
+            let mut head = [0u8; 20];
+            match read_full(&mut stream, &mut head) {
+                Ok(0) => {
+                    let _ = tx.send(Frame::End);
+                    return;
+                }
+                Ok(20) => {}
+                Ok(n) => {
+                    let _ = tx.send(Frame::Failed(format!("stream ended inside a frame ({n} bytes)")));
+                    return;
+                }
+                Err(e) => {
+                    let _ = tx.send(Frame::Failed(format!("stream: {e}")));
+                    return;
+                }
+            }
+            let from = u64::from_be_bytes(head[0..8].try_into().unwrap()) as i64;
+            let to = u64::from_be_bytes(head[8..16].try_into().unwrap()) as i64;
+            let length = u32::from_be_bytes(head[16..20].try_into().unwrap()) as usize;
+            let mut data = vec![0u8; length];
+            match read_full(&mut stream, &mut data) {
+                Ok(n) if n == length => {}
+                Ok(n) => {
+                    let _ = tx.send(Frame::Failed(format!("stream ended inside a frame ({} of {length} bytes)", n)));
+                    return;
+                }
+                Err(e) => {
+                    let _ = tx.send(Frame::Failed(format!("stream: {e}")));
+                    return;
+                }
+            }
+            // --max-mbit: read no faster than the cap (TCP holds the coordinator back).
+            if let Some(rate) = bytes_per_second {
+                read_bytes += (20 + length) as f64;
+                let due = Duration::from_secs_f64(((read_bytes - rate) / rate).max(0.0));
+                let elapsed = started.elapsed();
+                if due > elapsed {
+                    std::thread::sleep(due - elapsed);
+                }
+            }
+            let frame = if length == 0 && from == 0 {
+                Frame::KeepAlive
+            } else {
+                Frame::Chunk { from, to, text: inflate(&data, &mut decompressor) }
+            };
+            if tx.send(frame).is_err() {
+                return;
+            }
+        }
+    });
+    rx
+}
+
+/// Reads until `buf` is full or the stream ends: the bytes read.
+fn read_full(stream: &mut dyn Read, buf: &mut [u8]) -> std::io::Result<usize> {
+    let mut at = 0;
+    while at < buf.len() {
+        match stream.read(&mut buf[at..]) {
+            Ok(0) => break,
+            Ok(n) => at += n,
+            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
+            Err(e) => return Err(e),
+        }
+    }
+    Ok(at)
+}
+
+/// A chunk as the news server sent it (zlib; gzip too), whole.
+fn inflate(raw: &[u8], d: &mut libdeflater::Decompressor) -> Result<Vec<u8>, String> {
+    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 Ok(out);
+            }
+            Err(libdeflater::DecompressionError::InsufficientSpace) => out.resize(out.len() * 2, 0),
+            Err(e) => return Err(e.to_string()),
+        }
+    }
+}
+
+struct Config {
+    batch_articles: i64,
+    bytes_per_second: Option<f64>,
+}
+
+#[derive(Default)]
+struct Totals {
+    headers: u64,
+    releases: u64,
+    candidates: u64,
+    bytes: usize,
+    post_ms: u128,
+}
+
+/// One leased range, start to finish (commands/edge_node.ts runRange).
+fn run_range(api: &Arc<Api>, lease: LeaseJson, config: &Config, stopping: &AtomicBool) {
+    let range = lease.range;
+    let started = Instant::now();
+    let mut indexer = RangeIndexer::new(&range.group, range.start, range.end, range.next, lease.rules.rules(), 5);
+    let totals = Arc::new(Mutex::new(Totals::default()));
+    let post_error: Arc<Mutex<Option<ApiError>>> = Arc::new(Mutex::new(None));
+    let last_contact = Arc::new(AtomicI64::new(0));
+    let since = Instant::now();
+    // Posts in order on their own thread; the indexer goes on meanwhile (one batch may wait).
+    let (post_tx, post_rx) = sync_channel::<Batch>(1);
+    let poster = {
+        let (api, totals, post_error, last_contact) = (api.clone(), totals.clone(), post_error.clone(), last_contact.clone());
+        std::thread::spawn(move || {
+            for batch in post_rx {
+                if post_error.lock().unwrap().is_some() {
+                    continue;
+                }
+                let at = Instant::now();
+                match api.batch(range.id, &batch) {
+                    Ok((releases, candidates, bytes)) => {
+                        let mut t = totals.lock().unwrap();
+                        t.post_ms += at.elapsed().as_millis();
+                        t.bytes += bytes;
+                        t.headers += batch.headers;
+                        t.releases += releases;
+                        t.candidates += candidates;
+                        last_contact.store(since.elapsed().as_millis() as i64, Ordering::Relaxed);
+                    }
+                    Err(e) => *post_error.lock().unwrap() = Some(e),
+                }
+            }
+        })
+    };
+    let frames = match api.stream(range.id) {
+        Ok(stream) => read_frames(stream, config.bytes_per_second),
+        Err(e) => {
+            log("error", &format!("Range {} failed: {e}; giving it back.", range.id));
+            let _ = api.call(range.id, "release");
+            return;
+        }
+    };
+    let mut batch_from = range.next;
+    let mut reached_end = false;
+    let mut corrupt: Option<(i64, i64, String)> = None;
+    let mut failure: Option<String> = None;
+    for frame in frames.iter() {
+        if stopping.load(Ordering::Relaxed) || post_error.lock().unwrap().is_some() {
+            break;
+        }
+        // Keep the lease while no batch goes out (slow ranges wait for the spool).
+        if since.elapsed().as_millis() as i64 - last_contact.load(Ordering::Relaxed) > HEARTBEAT.as_millis() as i64 {
+            last_contact.store(since.elapsed().as_millis() as i64, Ordering::Relaxed);
+            let _ = api.call(range.id, "heartbeat");
+        }
+        match frame {
+            Frame::KeepAlive => continue,
+            Frame::End => break,
+            Frame::Failed(e) => {
+                failure = Some(e);
+                break;
+            }
+            Frame::Chunk { from, to, text } => {
+                let data = match text {
+                    Ok(data) => data,
+                    Err(e) => {
+                        corrupt = Some((from, to, e));
+                        break;
+                    }
+                };
+                indexer.add_text(&String::from_utf8_lossy(&data), to + 1);
+                reached_end = to >= range.end;
+                if reached_end {
+                    indexer.finish();
+                }
+                if reached_end || to + 1 - batch_from >= config.batch_articles || indexer.pending() >= MAX_PENDING_FILES {
+                    if post_tx.send(indexer.take()).is_err() {
+                        break;
+                    }
+                    batch_from = to + 1;
+                }
+                if reached_end {
+                    break;
+                }
+            }
+        }
+    }
+    drop(post_tx);
+    let _ = poster.join();
+    let post_error = post_error.lock().unwrap().take();
+    if let Some((from, to, e)) = corrupt {
+        // The server's compressed reply cut short: reported, fetched plain next time.
+        log("warn", &format!("Range {}: chunk {from}–{to} doesn't inflate ({e}); reported, the range goes back", range.id));
+        let _ = api.corrupt(range.id, from);
+        return;
+    }
+    match post_error {
+        Some(ApiError::LeaseLost(m)) => {
+            log("warn", &format!("Range {}: {m}", range.id));
+            return;
+        }
+        Some(ApiError::Other(m)) => {
+            log("error", &format!("Range {} failed: {m}; giving it back.", range.id));
+            let _ = api.call(range.id, "release");
+            return;
+        }
+        None => {}
+    }
+    if let Some(e) = failure {
+        log("error", &format!("Range {} failed: {e}; giving it back.", range.id));
+        let _ = api.call(range.id, "release");
+        return;
+    }
+    if !reached_end {
+        log("info", &format!("Range {}: giving it back.", range.id));
+        let _ = api.call(range.id, "release");
+        return;
+    }
+    match api.call(range.id, "finish") {
+        Ok(()) => {
+            let t = totals.lock().unwrap();
+            let seconds = started.elapsed().as_secs_f64();
+            log(
+                "info",
+                &format!(
+                    "{} {}–{}: {} headers in {:.0} s ({:.0}/s), {} release chunks, {} probe candidates; posted {:.1} MB in {:.1} s.",
+                    range.group,
+                    range.start,
+                    range.end,
+                    t.headers,
+                    seconds,
+                    t.headers as f64 / seconds,
+                    t.releases,
+                    t.candidates,
+                    t.bytes as f64 / 1e6,
+                    t.post_ms as f64 / 1000.0
+                ),
+            );
+        }
+        Err(e) => log("warn", &format!("Range {}: {e}", range.id)),
+    }
+}
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    let flag = |name: &str| args.iter().position(|a| a == name).and_then(|i| args.get(i + 1)).cloned();
+    let url = std::env::var("COORDINATOR_URL").expect("COORDINATOR_URL is needed");
+    let token = std::env::var("NODE_TOKEN").expect("NODE_TOKEN is needed");
+    let threads: usize = flag("--threads")
+        .or_else(|| std::env::var("EDGE_PROCESSES").ok())
+        .map(|v| v.parse().unwrap())
+        .filter(|&n| n > 0)
+        .unwrap_or_else(|| std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1));
+    let ranges: Option<usize> = flag("--ranges").map(|v| v.parse().unwrap());
+    let max_mbit: Option<f64> = flag("--max-mbit").or_else(|| std::env::var("EDGE_MAX_MBIT").ok()).map(|v| v.parse().unwrap()).filter(|&m: &f64| m > 0.0);
+    let config = Arc::new(Config {
+        batch_articles: flag("--batch-articles").map(|v| v.parse().unwrap()).unwrap_or(500_000),
+        // The node's cap, shared by its threads.
+        bytes_per_second: max_mbit.map(|m| m * 1e6 / 8.0 / threads as f64),
+    });
+    let stopping = Arc::new(AtomicBool::new(false));
+    signal_hook::flag::register(signal_hook::consts::SIGTERM, stopping.clone()).unwrap();
+    signal_hook::flag::register(signal_hook::consts::SIGINT, stopping.clone()).unwrap();
+    let api = Arc::new(Api::new(&url, &token));
+    log("info", &format!("Starting {threads} worker threads{}.", max_mbit.map(|m| format!(" (download cap {m} Mbit/s)")).unwrap_or_default()));
+    let workers: Vec<_> = (0..threads)
+        .map(|_| {
+            let (api, config, stopping) = (api.clone(), config.clone(), stopping.clone());
+            std::thread::spawn(move || {
+                let mut done = 0;
+                while !stopping.load(Ordering::Relaxed) && ranges.is_none_or(|n| done < n) {
+                    match api.lease() {
+                        Ok(Some(lease)) => {
+                            run_range(&api, lease, &config, &stopping);
+                            done += 1;
+                        }
+                        Ok(None) if ranges.is_some() => break,
+                        Ok(None) => idle(&stopping),
+                        Err(e) => {
+                            log("warn", &format!("Lease failed: {e}"));
+                            idle(&stopping);
+                        }
+                    }
+                }
+            })
+        })
+        .collect();
+    for worker in workers {
+        let _ = worker.join();
+    }
+}
+
+fn idle(stopping: &AtomicBool) {
+    let until = Instant::now() + IDLE;
+    while Instant::now() < until && !stopping.load(Ordering::Relaxed) {
+        std::thread::sleep(Duration::from_millis(500));
+    }
+}