Files
usenet-indexer/edge/src/bin/edge-node.rs
T
Daan MeijerandClaude Opus 5.5 3d24db5565
Test and deploy / test (push) Successful in 50s
Test and deploy / edge (push) Successful in 1m55s
Test and deploy / deploy (push) Successful in 1m47s
Rules drop posters at the header pass: the CelebMovie family (rule set version 3)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-16 11:36:46 +02:00

568 lines
22 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! 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, 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::indexer::{Batch, RangeIndexer};
use edge::names::Rules;
use edge::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>,
#[serde(default)]
dropped_posters: Option<Vec<String>>,
}
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;
}
if let Some(posters) = self.dropped_posters {
rules.dropped_posters = posters;
}
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),
/// 413: the batch was over the server's body limit; the range is tried again in smaller batches.
TooLarge(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::TooLarge(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() == 413 {
return Err(ApiError::TooLarge(format!("batch: HTTP {} ({} MB gzipped, {} MB of JSON)", response.status(), body.len() / 1_000_000, json.len() / 1_000_000)));
}
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>,
/// Ranges whose batch was too large (413), and how much smaller their batches are now (halved each time).
smaller: Mutex<HashMap<i64, i64>>,
}
impl Config {
/// Articles and pending files a batch of this range is posted at.
fn limits(&self, range: i64) -> (i64, usize) {
let divisor = self.smaller.lock().unwrap().get(&range).copied().unwrap_or(1);
(self.batch_articles / divisor, MAX_PENDING_FILES / divisor as usize)
}
}
#[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 (batch_articles, max_pending) = config.limits(range.id);
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 >= batch_articles || indexer.pending() >= max_pending {
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::TooLarge(m)) => {
// Batches already stored stay; from the resume point on, half the size (a batch is at least a chunk).
let divisor = {
let mut smaller = config.smaller.lock().unwrap();
let d = smaller.entry(range.id).or_insert(1);
*d = (*d * 2).min(64);
*d
};
log("warn", &format!("Range {}: {m}; giving it back, its batches 1/{divisor} the size from now on.", range.id));
let _ = api.call(range.id, "release");
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),
smaller: Mutex::new(HashMap::new()),
});
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));
}
}