From a8a139c6d92f9ac95518739b97dc240aee951f83 Mon Sep 17 00:00:00 2001 From: Daan Meijer Date: Fri, 9 Oct 2026 20:42:09 +0200 Subject: [PATCH] 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 --- tools/edge-bench-rs/src/bin/edge-node.rs | 533 +++++++++++++++++++++++ 1 file changed, 533 insertions(+) create mode 100644 tools/edge-bench-rs/src/bin/edge-node.rs diff --git a/tools/edge-bench-rs/src/bin/edge-node.rs b/tools/edge-bench-rs/src/bin/edge-node.rs new file mode 100644 index 0000000..0a3b970 --- /dev/null +++ b/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, + #[serde(default)] + dropped_signatures: Vec, + #[serde(default)] + dropped_tools: Option>, + #[serde(default)] + per_article_max_parts: Option, + #[serde(default)] + sample_percent: Option, + #[serde(default)] + drop_executables: Option, +} + +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, 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::(&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, 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, 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, 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, bytes_per_second: Option) -> Receiver { + let (tx, rx) = sync_channel::(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 { + 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, 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, +} + +#[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, 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>> = 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::(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 = 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 = 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 = flag("--ranges").map(|v| v.parse().unwrap()); + let max_mbit: Option = 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)); + } +}