diff --git a/seaweed-worker/crates/core/src/senders.rs b/seaweed-worker/crates/core/src/senders.rs index b8a3136fb..39c5bff7e 100644 --- a/seaweed-worker/crates/core/src/senders.rs +++ b/seaweed-worker/crates/core/src/senders.rs @@ -79,18 +79,14 @@ impl ExecutionSender for StreamSender { /// what happened without every handler having to know about metrics. pub struct MeteredSender<'a> { inner: &'a StreamSender, - metrics: crate::metrics::Metrics, - job_type: String, proposals: std::sync::atomic::AtomicUsize, failed: std::sync::atomic::AtomicBool, } impl<'a> MeteredSender<'a> { - pub fn new(inner: &'a StreamSender, metrics: crate::metrics::Metrics, job_type: &str) -> Self { + pub fn new(inner: &'a StreamSender) -> Self { Self { inner, - metrics, - job_type: job_type.to_string(), proposals: std::sync::atomic::AtomicUsize::new(0), failed: std::sync::atomic::AtomicBool::new(false), } @@ -111,11 +107,13 @@ impl<'a> MeteredSender<'a> { impl DetectionSender for MeteredSender<'_> { fn send_proposals(&self, proposals: DetectionProposals) -> Result<()> { - self.proposals.fetch_add( - proposals.proposals.len(), - std::sync::atomic::Ordering::Relaxed, - ); - self.inner.send_proposals(proposals) + // Counted after the send, not before: a stream that closed mid-sweep + // would otherwise leave proposals_total claiming work admin never saw. + let count = proposals.proposals.len(); + self.inner.send_proposals(proposals)?; + self.proposals + .fetch_add(count, std::sync::atomic::Ordering::Relaxed); + Ok(()) } fn send_complete(&self, complete: DetectionComplete) -> Result<()> { @@ -145,8 +143,61 @@ impl ExecutionSender for MeteredSender<'_> { self.failed .store(true, std::sync::atomic::Ordering::Relaxed); } - let _ = &self.job_type; - let _ = &self.metrics; self.inner.send_completed(completed) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::pb::JobProposal; + + fn proposals(n: usize) -> DetectionProposals { + DetectionProposals { + proposals: (0..n).map(|_| JobProposal::default()).collect(), + ..Default::default() + } + } + + #[test] + fn proposals_are_counted_once_they_are_sent() { + let (tx, mut rx) = mpsc::unbounded_channel(); + let stream = StreamSender::new("worker-1".to_string(), tx); + let metered = MeteredSender::new(&stream); + + metered.send_proposals(proposals(3)).expect("send"); + assert_eq!(metered.proposals(), 3); + assert!(rx.try_recv().is_ok(), "the proposals reached the stream"); + } + + // Admin never saw these, so counting them would report work that was not + // handed over. + #[test] + fn proposals_are_not_counted_when_the_stream_is_gone() { + let (tx, rx) = mpsc::unbounded_channel(); + let stream = StreamSender::new("worker-1".to_string(), tx); + let metered = MeteredSender::new(&stream); + drop(rx); + + assert!(metered.send_proposals(proposals(3)).is_err()); + assert_eq!(metered.proposals(), 0); + } + + // A handler can report failure and still return Ok; the outcome has to come + // from what it said, not only from what it returned. + #[test] + fn a_reported_failure_is_remembered() { + let (tx, _rx) = mpsc::unbounded_channel(); + let stream = StreamSender::new("worker-1".to_string(), tx); + let metered = MeteredSender::new(&stream); + assert!(!metered.reported_failure()); + + metered + .send_complete(DetectionComplete { + success: false, + ..Default::default() + }) + .expect("send"); + assert!(metered.reported_failure()); + } +} diff --git a/seaweed-worker/crates/core/src/stream.rs b/seaweed-worker/crates/core/src/stream.rs index 7dae72e38..9dc7770f2 100644 --- a/seaweed-worker/crates/core/src/stream.rs +++ b/seaweed-worker/crates/core/src/stream.rs @@ -367,7 +367,7 @@ fn spawn_detection( // Held until the sweep finishes, so the worker keeps to the capacity it // advertised and the heartbeat reports the truth while it works. let _permit = slots.detection.acquire().await; - let metered = MeteredSender::new(&sender, metrics.clone(), &request.job_type); + let metered = MeteredSender::new(&sender); let started = Instant::now(); let outcome = handler.detect(&request, &metered).await; let result = if let Err(err) = &outcome { @@ -418,7 +418,7 @@ fn spawn_execution( let Some(handler) = registry.get(&job_type) else { return; }; - let metered = MeteredSender::new(&sender, metrics.clone(), &job_type); + let metered = MeteredSender::new(&sender); let started = Instant::now(); let outcome = handler.execute(&request, &metered).await; let result = if let Err(err) = &outcome { diff --git a/seaweed-worker/crates/lance/src/main.rs b/seaweed-worker/crates/lance/src/main.rs index c71c253c7..41bb1109b 100644 --- a/seaweed-worker/crates/lance/src/main.rs +++ b/seaweed-worker/crates/lance/src/main.rs @@ -2,7 +2,7 @@ use std::time::Duration; use anyhow::{Context, Result}; use clap::Parser; -use std::net::SocketAddr; +use std::net::{IpAddr, SocketAddr}; use std::sync::Arc; use seaweed_worker_core::{Metrics, Registry, TlsOptions, WorkerOptions}; @@ -101,6 +101,19 @@ impl Args { } } +/// Builds the metrics bind address. Parsing the host separately is what makes an +/// IPv6 literal work: "::1" and 9327 joined with a colon is not an address, and +/// formatting them that way turns `--metrics-ip ::1` into a startup failure. +/// Brackets are accepted too, since that is how the same address is written in a +/// URL and an operator will reasonably try it. +fn metrics_address(ip: &str, port: u16) -> Result { + let host = ip.trim().trim_start_matches('[').trim_end_matches(']'); + let parsed: IpAddr = host + .parse() + .with_context(|| format!("parse the metrics address {ip}"))?; + Ok(SocketAddr::new(parsed, port)) +} + #[tokio::main] async fn main() -> Result<()> { tracing_subscriber::fmt() @@ -131,15 +144,10 @@ async fn main() -> Result<()> { let metrics = Metrics::new(&options.worker_id, &options.worker_version)?; let lance_metrics = LanceMetrics::new(&metrics)?; if args.metrics_port > 0 { - let addr: SocketAddr = format!("{}:{}", args.metrics_ip, args.metrics_port) - .parse() - .with_context(|| { - format!( - "parse the metrics address {}:{}", - args.metrics_ip, args.metrics_port - ) - })?; - seaweed_worker_core::metrics::spawn(metrics.clone(), addr); + seaweed_worker_core::metrics::spawn( + metrics.clone(), + metrics_address(&args.metrics_ip, args.metrics_port)?, + ); } let mut registry = Registry::new().with_preview(Arc::new( @@ -151,3 +159,36 @@ async fn main() -> Result<()> { seaweed_worker_core::stream::run_with_metrics(options, registry, metrics).await } + +#[cfg(test)] +mod tests { + use super::metrics_address; + + #[test] + fn metrics_address_takes_ipv4_ipv6_and_brackets() { + assert_eq!( + metrics_address("127.0.0.1", 9327).unwrap().to_string(), + "127.0.0.1:9327" + ); + // Joining these with a colon gives "::1:9327", which is not an address. + assert_eq!( + metrics_address("::1", 9327).unwrap().to_string(), + "[::1]:9327" + ); + assert_eq!( + metrics_address("[::1]", 9327).unwrap().to_string(), + "[::1]:9327" + ); + assert_eq!( + metrics_address("0.0.0.0", 9327).unwrap().to_string(), + "0.0.0.0:9327" + ); + } + + #[test] + fn metrics_address_rejects_a_hostname() { + // Binding takes an address, not a name; saying so beats a confusing + // failure inside the server. + assert!(metrics_address("localhost", 9327).is_err()); + } +}