snapshot of https://github.com/seaweedfs/seaweedfs.git (master, 2026-09-16; blobs > 131072B omitted)
Some checks failed
go: build dev binaries / cleanup (push) Successful in 4s
Codespell / Check for spelling errors (push) Failing after 3s
docker: build dev containers / build-rust-binaries (amd64, x86_64-unknown-linux-musl) (push) Failing after 3s
docker: build dev containers / build-dev-containers (push) Has been skipped
docker: build dev containers / build-rust-binaries (arm64, true, aarch64-unknown-linux-musl) (push) Failing after 3s
End to End / FUSE Mount (push) Failing after 2s
EC Integration Tests / EC Integration Tests (push) Failing after 30s
FUSE DLM Integration Tests / FUSE DLM Integration Tests (push) Failing after 3s
FUSE Volume Server Failover Tests / FUSE Volume Server Failover (push) Failing after 3s
FUSE Integration Tests / FUSE Integration Testing (push) Failing after 2s
FUSE P2P Peer Chunk Sharing Integration Tests / FUSE P2P Peer Chunk Sharing (push) Failing after 3s
go: build binary / Go Vet (push) Failing after 2s
go: build binary / Go Vet 32-bit (push) Failing after 3s
go: build binary / Build (push) Failing after 4s
go: build binary / Build cross-platform (freebsd/amd64) (push) Failing after 3s
go: build binary / Build cross-platform (windows/amd64) (push) Failing after 3s
go: build binary / Build cross-platform (darwin/arm64) (push) Failing after 3s
go: build binary / Build cross-platform (windows/arm64) (push) Failing after 3s
go: build binary / Test (push) Failing after 2s
go: build binary / Test 32-bit (push) Failing after 3s
helm: lint and test charts / lint-test (push) Failing after 3s
Java Client Integration Tests / Java Integration Tests (11) (push) Failing after 3s
Java Client Integration Tests / Java Integration Tests (17) (push) Failing after 3s
go: build dev binaries / build_dev_linux_windows (amd64, linux) (push) Failing after 2s
go: build dev binaries / build_dev_linux_windows (amd64, windows) (push) Failing after 3s
go: build dev binaries / build_dev_darwin (amd64, darwin) (push) Has been skipped
go: build dev binaries / build_dev_darwin (arm64, darwin) (push) Has been skipped
Java Client Unit Tests / Java Unit Tests (11) (push) Failing after 3s
Java Client Unit Tests / Java Unit Tests (17) (push) Failing after 2s
Java Client Unit Tests / Java Unit Tests (21) (push) Failing after 3s
Java Client Unit Tests / Java Unit Tests (8) (push) Failing after 3s
Kafka Quick Test (Load Test with Schema Registry) / Kafka Client Load Test (Quick) (push) Failing after 3s
Kafka Gateway Tests / Kafka Unit Tests (unit-tests-1) (push) Failing after 11s
Kafka Gateway Tests / Kafka Integration Tests (Critical) (integration-1) (push) Failing after 3s
Kafka Gateway Tests / Kafka End-to-End Tests (with SMQ) (e2e-1) (push) Failing after 3s
Kafka Gateway Tests / Kafka Consumer Group Tests (Highly Isolated) (consumer-group-1) (push) Failing after 2s
Kafka Gateway Tests / Kafka Client Compatibility (with SMQ) (client-compat-1) (push) Failing after 3s
Kafka Gateway Tests / Kafka SMQ Integration Tests (Full Stack) (smq-integration-1) (push) Failing after 3s
Kafka Gateway Tests / Kafka Protocol Tests (Isolated) (protocol-1) (push) Failing after 3s
KMS Tests / KMS Provider Integration Tests (push) Failing after 2s
KMS Tests / S3 KMS End-to-End Tests (push) Failing after 3s
Master Cold Start Tests / Master Cold Start Tests (push) Failing after 4s
Metadata Subscribe Integration Tests / Metadata Subscribe Integration Tests (push) Failing after 3s
Multi-Master Tests / Multi-Master Failover Tests (push) Failing after 3s
Performance / CPU and Heap Profile (push) Failing after 2s
Performance / S3 Read/Write Benchmark (rust volume) (push) Failing after 3s
pjdfstest POSIX Compliance / pjdfstest (push) Failing after 2s
Plugin Worker Integration Tests / Plugin Worker: erasure_coding (push) Failing after 3s
Plugin Worker Integration Tests / Plugin Worker: vacuum (push) Failing after 3s
Rust Volume Server Tests / Rust Integration Tests (push) Failing after 3s
Rust Volume Server Tests / Go Tests with Rust Volume (grpc - Shard 1) (push) Failing after 2s
Rust Volume Server Tests / Go Tests with Rust Volume (http - Shard 1) (push) Failing after 3s
Rust Volume Server Tests / Go Tests with Rust Volume (grpc - Shard 2) (push) Failing after 3s
Rust Volume Server Tests / Go Tests with Rust Volume (http - Shard 2) (push) Failing after 3s
Rust Volume Server Tests / Go Tests with Rust Volume (grpc - Shard 3) (push) Failing after 3s
S3 IAM Integration Tests / S3 IAM Distributed Tests (push) Failing after 2s
S3 IAM Integration Tests / S3 IAM Performance Tests (push) Failing after 3s
Performance / Throughput Benchmark (go volume) (push) Failing after 3s
Performance / Throughput Benchmark (rust volume) (push) Failing after 2s
Performance / S3 Read/Write Benchmark (go volume) (push) Failing after 3s
Plugin Worker Integration Tests / Plugin Worker: volume_balance (push) Failing after 3s
PostgreSQL Gateway Tests / PostgreSQL Basic Tests (push) Failing after 3s
Rust Volume Server Tests / Rust Unit Tests (push) Failing after 2s
Rust Volume Server Tests / Go Tests with Rust Volume (http - Shard 3) (push) Failing after 3s
Rust Plugin Worker Tests / Rust Plugin Worker Build and Unit Tests (push) Failing after 3s
rust: build dev volume server binaries / cleanup (push) Successful in 3s
S3 ETag and ACL Tests / S3 ETag + ACL Integration Tests (push) Failing after 3s
S3 IAM Integration Tests / IAM Unit Tests (push) Failing after 3s
S3 IAM Integration Tests / S3 IAM Integration Tests (advanced) (push) Failing after 3s
S3 IAM Integration Tests / S3 IAM Integration Tests (basic) (push) Failing after 2s
S3 IAM Integration Tests / S3 IAM Integration Tests (group) (push) Failing after 3s
S3 IAM Integration Tests / S3 IAM Integration Tests (policy-enforcement) (push) Failing after 2s
S3 IAM Integration Tests / S3 IAM Integration Tests (sts) (push) Failing after 4s
S3 Keycloak Integration Tests / S3 Keycloak Integration Tests (push) Failing after 3s
S3 PyArrow Parquet Tests / PyArrow Parquet Tests (Python 3.11) (push) Failing after 3s
S3 PyArrow Parquet Tests / Go Unit Tests (Implicit Directory) (push) Has been cancelled
S3 PyArrow Parquet Tests / PyArrow Parquet Tests (Python 3.12) (push) Has been cancelled
rust: build dev volume server binaries / build-rust-volume-dev-darwin (darwin-amd64, x86_64-apple-darwin) (push) Has been cancelled
rust: build dev volume server binaries / build-rust-volume-dev-darwin (darwin-arm64, aarch64-apple-darwin) (push) Has been cancelled
rust: build dev volume server binaries / build-rust-volume-dev-linux (linux-amd64, x86_64-unknown-linux-gnu) (push) Has been cancelled
S3 PyArrow Parquet Tests / PyArrow Parquet Tests (Python 3.9) (push) Has been cancelled
S3 Proxy Signature Tests / S3 Proxy Signature Verification Tests (push) Has been cancelled
S3 SDK V2 Route Disambiguation Tests / S3 SDK V2 Routing Tests (push) Has been cancelled
S3 SSE Tests / S3 SSE Integration Tests (quick) (push) Has been cancelled
S3 SSE Tests / S3 SSE Compatibility Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Metadata Persistence Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Copy Operations Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Multipart Upload Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Performance Test (push) Has been cancelled
S3 SSE Tests / S3 Volume Encryption Test (push) Has been cancelled
S3 SSE Tests / S3 SSE Integration Tests (comprehensive) (push) Has been cancelled
Spark Integration Tests / Spark Integration Tests (push) Failing after 2s
Telemetry Integration Tests / telemetry-integration-test (push) Failing after 3s
Volume Server Integration Tests / Volume Server Integration Tests (grpc - Shard 1) (push) Failing after 3s
Volume Server Integration Tests / Volume Server Integration Tests (http - Shard 1) (push) Failing after 3s
Ceph S3 tests / Basic S3 tests (KV store) (push) Has been cancelled
Ceph S3 tests / S3 Versioning & Object Lock tests (push) Has been cancelled
S3 Policy Integration Tests / S3 Policy Unit Tests (push) Failing after 3s
Ceph S3 tests / S3 CORS tests (push) Has been cancelled
Ceph S3 tests / SeaweedFS Custom S3 Copy tests (push) Has been cancelled
Ceph S3 tests / Basic S3 tests (SQL store) (push) Has been cancelled
test s3 over https using aws-cli / awscli-tests (push) Has been cancelled
TLS Rotation Integration Tests / TLS Rotation Integration Tests (push) Has been cancelled
Vacuum Integration Tests / Vacuum Integration Tests (push) Has been cancelled
S3 Policy Integration Tests / S3 Policy Variables Integration Tests (push) Failing after 3s
S3 Policy Integration Tests / S3 Policy Enforcement Integration Tests (advanced-policy) (push) Failing after 3s
S3 Policy Integration Tests / S3 Policy Enforcement Integration Tests (basic-policy) (push) Failing after 2s
S3 Policy Integration Tests / S3 Policy Enforcement Integration Tests (contextual-policy) (push) Failing after 4s
S3 Policy Integration Tests / Trusted Proxy Detection Tests (push) Failing after 2s
S3 Policy Integration Tests / S3 Policy Shell Integration Tests (push) Failing after 3s
Samba on FUSE Integration / samba-integration (push) Failing after 2s
SFTP Integration Tests / SFTP Integration Testing (push) Failing after 3s
terraform: validate and test modules / fmt, validate, plan-level tests (push) Failing after 3s
terraform: validate and test modules / local cluster smoke test (real weed) (push) Failing after 2s
Volume Server Integration Tests / Volume Server Integration Tests (grpc - Shard 3) (push) Failing after 3s
Volume Server Integration Tests / Volume Server Integration Tests (grpc - Shard 2) (push) Failing after 3s
Volume Server Integration Tests / Volume Server Integration Tests (http - Shard 2) (push) Failing after 2s
Volume Server Integration Tests / Volume Server Integration Tests (http - Shard 3) (push) Failing after 5s
mount: windows conformance / WinFsp conformance (push) Has been cancelled
mount: windows / Mount on Windows (push) Has been cancelled

This commit is contained in:
upstream-mirror
2026-09-16 03:23:42 +05:00
commit 89f5be005d
3997 changed files with 961970 additions and 0 deletions

View File

@@ -0,0 +1,32 @@
[package]
name = "seaweed-worker-core"
version.workspace = true
edition.workspace = true
rust-version.workspace = true
description = "SeaweedFS plugin.proto worker contract"
[lib]
name = "seaweed_worker_core"
[dependencies]
anyhow.workspace = true
async-trait.workspace = true
prost.workspace = true
prost-types.workspace = true
tokio.workspace = true
tokio-stream.workspace = true
tonic.workspace = true
tonic-prost.workspace = true
axum.workspace = true
prometheus.workspace = true
tracing.workspace = true
[build-dependencies]
tonic-prost-build.workspace = true
# Ships protoc with the build so neither CI nor a developer needs a system
# install, and so the version is pinned rather than whatever the platform's
# package manager happens to carry. The same crate seaweed-volume uses.
protoc-bin-vendored = "3"
[lints]
workspace = true

View File

@@ -0,0 +1,25 @@
fn main() -> Result<(), Box<dyn std::error::Error>> {
// Use the protoc that ships with protoc-bin-vendored rather than a system
// one, so the build needs no package manager and always sees the same
// version. An explicit PROTOC still wins, for packagers supplying their own
// and for the lance crates, whose own build scripts read the same variable.
if std::env::var_os("PROTOC").is_none() {
// SAFETY: a build script's main runs single-threaded before anything
// else in this process, so no other thread can be reading the
// environment concurrently.
unsafe {
std::env::set_var("PROTOC", protoc_bin_vendored::protoc_bin_path()?);
}
}
// Compiled straight out of the Go tree, the way seaweed-volume already reads
// filer.proto, so the contract cannot drift from a vendored copy.
tonic_prost_build::configure()
// The server half is only for tests, which stand up a fake admin.
.build_server(true)
.build_client(true)
.protoc_arg("--experimental_allow_proto3_optional")
.compile_protos(&["../../../weed/pb/plugin.proto"], &["../../../weed/pb/"])?;
println!("cargo:rerun-if-changed=../../../weed/pb/plugin.proto");
Ok(())
}

View File

@@ -0,0 +1,73 @@
//! SeaweedFS addresses the way the Go tree does.
//!
//! An operator gives a worker the admin's HTTP address, and the gRPC port is
//! derived from it rather than asked for separately. Dialling the HTTP port
//! instead fails as "frame with invalid size", which reads like a protocol bug
//! rather than a wrong port, so getting this right is worth its own module.
//! Mirrors pb.ServerToGrpcAddress in weed/pb/grpc_client_server.go.
const GRPC_PORT_OFFSET: u16 = 10000;
/// Converts `host:port` to the gRPC address, and accepts the explicit
/// `host:port.grpcPort` form the Go side also understands.
pub fn server_to_grpc_address(server: &str) -> Option<String> {
let (host, port_part) = server.rsplit_once(':')?;
// "port.grpcPort" states the gRPC port outright.
if let Some((_, grpc_port)) = port_part.split_once('.')
&& let Ok(port) = grpc_port.parse::<u16>()
{
return Some(join_host_port(host, port));
}
let port: u16 = port_part.parse().ok()?;
Some(join_host_port(host, port.checked_add(GRPC_PORT_OFFSET)?))
}
fn join_host_port(host: &str, port: u16) -> String {
// An IPv6 literal has to keep its brackets or the port reads as part of it.
if host.contains(':') && !host.starts_with('[') {
format!("[{host}]:{port}")
} else {
format!("{host}:{port}")
}
}
#[cfg(test)]
mod tests {
use super::server_to_grpc_address;
#[test]
fn derives_the_grpc_port() {
assert_eq!(
server_to_grpc_address("localhost:23646").as_deref(),
Some("localhost:33646")
);
assert_eq!(
server_to_grpc_address("127.0.0.1:9333").as_deref(),
Some("127.0.0.1:19333")
);
}
#[test]
fn honours_an_explicit_grpc_port() {
assert_eq!(
server_to_grpc_address("localhost:23646.33999").as_deref(),
Some("localhost:33999")
);
}
#[test]
fn brackets_ipv6_literals() {
assert_eq!(
server_to_grpc_address("::1:23646").as_deref(),
Some("[::1]:33646")
);
}
#[test]
fn rejects_what_it_cannot_parse() {
assert!(server_to_grpc_address("localhost").is_none());
assert!(server_to_grpc_address("localhost:notaport").is_none());
}
}

View File

@@ -0,0 +1,49 @@
use std::time::Duration;
/// How one worker process connects and how much work it will take on.
#[derive(Clone, Debug)]
pub struct WorkerOptions {
/// Admin gRPC address, e.g. "localhost:23646".
pub admin_address: String,
pub worker_id: String,
pub worker_version: String,
/// Advertised address; empty when the worker takes no inbound connections.
pub worker_address: String,
pub heartbeat_interval: Duration,
pub reconnect_delay: Duration,
pub max_detection_concurrency: i32,
pub max_execution_concurrency: i32,
/// mTLS for the control stream, mirroring the Go worker's `[grpc.worker]`
/// section of security.toml. None means plaintext, which is the default the
/// Go worker also takes when no certificates are configured.
pub tls: Option<TlsOptions>,
}
/// Certificates for the control stream. All three are required together: the
/// cluster's gRPC TLS is mutual, so a CA without a client identity gets refused
/// by admin rather than falling back to one-way TLS.
#[derive(Clone, Debug)]
pub struct TlsOptions {
pub ca_path: String,
pub client_cert_path: String,
pub client_key_path: String,
/// Name to verify the server certificate against, when the address a worker
/// dials is not the name the certificate carries.
pub server_name: Option<String>,
}
impl Default for WorkerOptions {
fn default() -> Self {
Self {
admin_address: "localhost:23646".to_string(),
worker_id: String::new(),
worker_version: env!("CARGO_PKG_VERSION").to_string(),
worker_address: String::new(),
heartbeat_interval: Duration::from_secs(10),
reconnect_delay: Duration::from_secs(5),
max_detection_concurrency: 1,
max_execution_concurrency: 1,
tls: None,
}
}
}

View File

@@ -0,0 +1,106 @@
//! Builders for the config forms a worker returns in its JobTypeDescriptor.
//!
//! Admin renders these into the job's settings page, so a worker written in any
//! language gets a UI without touching Go or templ. That only holds if the
//! field types and defaults are right, which is why they are built here rather
//! than spelled out at each call site.
use std::collections::HashMap;
use crate::pb::{
ConfigField, ConfigFieldType, ConfigForm, ConfigSection, ConfigValue, config_value::Kind,
};
pub fn int_value(value: i64) -> ConfigValue {
ConfigValue {
kind: Some(Kind::Int64Value(value)),
}
}
pub fn bool_value(value: bool) -> ConfigValue {
ConfigValue {
kind: Some(Kind::BoolValue(value)),
}
}
pub fn string_value(value: impl Into<String>) -> ConfigValue {
ConfigValue {
kind: Some(Kind::StringValue(value.into())),
}
}
/// Reads an integer a request carried, falling back when admin sent nothing.
pub fn int_or(values: &HashMap<String, ConfigValue>, name: &str, fallback: i64) -> i64 {
match values.get(name).and_then(|v| v.kind.as_ref()) {
Some(Kind::Int64Value(value)) => *value,
Some(Kind::DoubleValue(value)) => *value as i64,
_ => fallback,
}
}
pub fn bool_or(values: &HashMap<String, ConfigValue>, name: &str, fallback: bool) -> bool {
match values.get(name).and_then(|v| v.kind.as_ref()) {
Some(Kind::BoolValue(value)) => *value,
_ => fallback,
}
}
pub fn string_or(values: &HashMap<String, ConfigValue>, name: &str, fallback: &str) -> String {
match values.get(name).and_then(|v| v.kind.as_ref()) {
Some(Kind::StringValue(value)) if !value.is_empty() => value.clone(),
_ => fallback.to_string(),
}
}
pub fn number_field(name: &str, label: &str, description: &str, min: i64, max: i64) -> ConfigField {
ConfigField {
name: name.to_string(),
label: label.to_string(),
description: description.to_string(),
field_type: ConfigFieldType::Int64 as i32,
min_value: Some(int_value(min)),
max_value: Some(int_value(max)),
..Default::default()
}
}
pub fn text_field(name: &str, label: &str, description: &str, placeholder: &str) -> ConfigField {
ConfigField {
name: name.to_string(),
label: label.to_string(),
description: description.to_string(),
field_type: ConfigFieldType::String as i32,
placeholder: placeholder.to_string(),
..Default::default()
}
}
pub fn bool_field(name: &str, label: &str, description: &str) -> ConfigField {
ConfigField {
name: name.to_string(),
label: label.to_string(),
description: description.to_string(),
field_type: ConfigFieldType::Bool as i32,
..Default::default()
}
}
pub fn form(
form_id: &str,
title: &str,
fields: Vec<ConfigField>,
defaults: HashMap<String, ConfigValue>,
) -> ConfigForm {
ConfigForm {
form_id: form_id.to_string(),
title: title.to_string(),
description: String::new(),
sections: vec![ConfigSection {
section_id: format!("{form_id}-main"),
title: title.to_string(),
description: String::new(),
fields,
}],
default_values: defaults,
}
}

View File

@@ -0,0 +1,29 @@
//! The SeaweedFS plugin worker contract, in Rust.
//!
//! `weed/pb/plugin.proto` is a language-agnostic gRPC stream: a worker connects
//! out to admin, announces the job types it can detect and execute, and then
//! answers requests on that one stream. `weed worker -admin=host:23646` is the
//! Go implementation of the same contract from outside the admin process; this
//! is the Rust one, and it knows nothing about any particular job.
pub mod address;
pub mod config;
pub mod config_form;
pub mod metrics;
pub mod registry;
pub mod senders;
pub mod stream;
/// Generated plugin.proto types.
pub mod pb {
// prost gives every oneof its own enum; the variant sizes are the
// messages' own, not a choice made here.
#![allow(clippy::large_enum_variant)]
tonic::include_proto!("plugin");
}
pub use config::{TlsOptions, WorkerOptions};
pub use metrics::Metrics;
pub use registry::{JobHandler, Preview, PreviewProvider, Registry};
pub use senders::{DetectionSender, ExecutionSender};
pub use stream::run;

View File

@@ -0,0 +1,404 @@
//! What a worker will say about itself when nobody is watching the logs.
//!
//! The Go plugin worker serves `/health`, `/ready` and `/metrics` on an optional
//! port (`weed worker -metricsPort`); this is the same contract, so one scrape
//! config covers workers in either language. Names follow the Go side's
//! convention, `SeaweedFS_<subsystem>_<name>`, under a `worker` subsystem that
//! nothing else uses.
//!
//! Nothing here is Lance-specific. A worker for another format registers its own
//! collectors on the same registry and gets the same endpoint.
use std::net::SocketAddr;
use std::sync::Arc;
use anyhow::{Context, Result};
use prometheus::{
Encoder, HistogramOpts, HistogramVec, IntCounter, IntCounterVec, IntGauge, IntGaugeVec, Opts,
Registry, TextEncoder,
};
use std::sync::atomic::{AtomicBool, Ordering};
use tracing::{info, warn};
const NAMESPACE: &str = "SeaweedFS";
const SUBSYSTEM: &str = "worker";
/// The metrics one worker process publishes, and the registry they live on.
///
/// Cloneable because the stream loop, the job handlers and the HTTP server all
/// hold it; everything inside is already shared.
#[derive(Clone)]
pub struct Metrics {
registry: Registry,
connected: IntGauge,
connects: IntCounterVec,
slots_used: IntGaugeVec,
slots_total: IntGaugeVec,
detections: IntCounterVec,
detection_seconds: HistogramVec,
proposals: IntCounterVec,
objects_seen: IntCounterVec,
objects_skipped: IntCounterVec,
jobs: IntCounterVec,
job_seconds: HistogramVec,
previews: IntCounterVec,
ready: Arc<AtomicBool>,
}
impl Metrics {
pub fn new(worker_id: &str, worker_version: &str) -> Result<Self> {
let registry = Registry::new();
// Build info as a constant 1, the way the Go side does it, so a scrape
// can tell which worker and which build answered.
let build = IntGaugeVec::new(
Opts::new("build_info", "Worker build information.")
.namespace(NAMESPACE)
.subsystem(SUBSYSTEM),
&["worker_id", "version"],
)?;
build.with_label_values(&[worker_id, worker_version]).set(1);
registry.register(Box::new(build))?;
let connected = IntGauge::with_opts(
Opts::new("connected", "1 while the admin control stream is up.")
.namespace(NAMESPACE)
.subsystem(SUBSYSTEM),
)?;
let connects = IntCounterVec::new(
Opts::new(
"stream_events_total",
"Control stream lifecycle events by outcome (connected, closed, failed, shutdown).",
)
.namespace(NAMESPACE)
.subsystem(SUBSYSTEM),
&["event"],
)?;
let slots_used = IntGaugeVec::new(
Opts::new("slots_used", "Slots currently held, by lane.")
.namespace(NAMESPACE)
.subsystem(SUBSYSTEM),
&["lane"],
)?;
let slots_total = IntGaugeVec::new(
Opts::new("slots_total", "Slots this worker advertises, by lane.")
.namespace(NAMESPACE)
.subsystem(SUBSYSTEM),
&["lane"],
)?;
let detections = IntCounterVec::new(
Opts::new(
"detections_total",
"Detection sweeps by job type and outcome.",
)
.namespace(NAMESPACE)
.subsystem(SUBSYSTEM),
&["job_type", "result"],
)?;
let detection_seconds = HistogramVec::new(
HistogramOpts::new("detection_seconds", "How long a detection sweep took.")
.namespace(NAMESPACE)
.subsystem(SUBSYSTEM)
.buckets(prometheus::exponential_buckets(0.01, 2.0, 14)?),
&["job_type"],
)?;
let proposals = IntCounterVec::new(
Opts::new("proposals_total", "Jobs proposed to admin, by job type.")
.namespace(NAMESPACE)
.subsystem(SUBSYSTEM),
&["job_type"],
)?;
let objects_seen = IntCounterVec::new(
Opts::new(
"objects_seen_total",
"Objects a detection sweep read, by job type.",
)
.namespace(NAMESPACE)
.subsystem(SUBSYSTEM),
&["job_type"],
)?;
let objects_skipped = IntCounterVec::new(
Opts::new(
"objects_skipped_total",
"Objects a sweep could not read, by job type and reason. A sweep that \
proposes nothing looks the same as one that could read nothing; this is \
what tells them apart.",
)
.namespace(NAMESPACE)
.subsystem(SUBSYSTEM),
&["job_type", "reason"],
)?;
let jobs = IntCounterVec::new(
Opts::new("jobs_total", "Jobs executed by job type and outcome.")
.namespace(NAMESPACE)
.subsystem(SUBSYSTEM),
&["job_type", "result"],
)?;
let job_seconds = HistogramVec::new(
HistogramOpts::new("job_seconds", "How long a job took to run.")
.namespace(NAMESPACE)
.subsystem(SUBSYSTEM)
.buckets(prometheus::exponential_buckets(0.05, 2.0, 16)?),
&["job_type"],
)?;
let previews = IntCounterVec::new(
Opts::new(
"previews_total",
"Object previews served to admin, by outcome.",
)
.namespace(NAMESPACE)
.subsystem(SUBSYSTEM),
&["result"],
)?;
for collector in [
Box::new(connected.clone()) as Box<dyn prometheus::core::Collector>,
Box::new(connects.clone()),
Box::new(slots_used.clone()),
Box::new(slots_total.clone()),
Box::new(detections.clone()),
Box::new(detection_seconds.clone()),
Box::new(proposals.clone()),
Box::new(objects_seen.clone()),
Box::new(objects_skipped.clone()),
Box::new(jobs.clone()),
Box::new(job_seconds.clone()),
Box::new(previews.clone()),
] {
registry.register(collector)?;
}
Ok(Self {
registry,
connected,
connects,
slots_used,
slots_total,
detections,
detection_seconds,
proposals,
objects_seen,
objects_skipped,
jobs,
job_seconds,
previews,
ready: Arc::new(AtomicBool::new(false)),
})
}
/// The registry, so a worker can add collectors of its own.
pub fn registry(&self) -> &Registry {
&self.registry
}
pub fn stream_connected(&self) {
self.connected.set(1);
self.ready.store(true, Ordering::Relaxed);
self.connects.with_label_values(&["connected"]).inc();
}
/// `event` is why the stream ended: closed, failed, or shutdown.
pub fn stream_ended(&self, event: &str) {
self.connected.set(0);
self.ready.store(false, Ordering::Relaxed);
self.connects.with_label_values(&[event]).inc();
}
pub fn set_slots(&self, lane: &str, used: i64, total: i64) {
self.slots_used.with_label_values(&[lane]).set(used);
self.slots_total.with_label_values(&[lane]).set(total);
}
pub fn detection_finished(&self, job_type: &str, result: &str, seconds: f64, proposals: usize) {
self.detections
.with_label_values(&[job_type, result])
.inc_by(1);
self.detection_seconds
.with_label_values(&[job_type])
.observe(seconds);
self.proposals
.with_label_values(&[job_type])
.inc_by(proposals as u64);
}
pub fn object_seen(&self, job_type: &str) {
self.objects_seen.with_label_values(&[job_type]).inc();
}
pub fn object_skipped(&self, job_type: &str, reason: &str) {
self.objects_skipped
.with_label_values(&[job_type, reason])
.inc();
}
pub fn job_finished(&self, job_type: &str, result: &str, seconds: f64) {
self.jobs.with_label_values(&[job_type, result]).inc();
self.job_seconds
.with_label_values(&[job_type])
.observe(seconds);
}
pub fn preview_finished(&self, result: &str) {
self.previews.with_label_values(&[result]).inc();
}
/// A counter this worker's own jobs can raise, e.g. fragments removed.
/// Registered lazily so a format's numbers live beside the generic ones
/// without core having to know what they are.
pub fn counter(&self, name: &str, help: &str) -> Result<IntCounter> {
let counter = IntCounter::with_opts(
Opts::new(name, help)
.namespace(NAMESPACE)
.subsystem(SUBSYSTEM),
)?;
self.registry.register(Box::new(counter.clone()))?;
Ok(counter)
}
fn gather(&self) -> Result<String> {
let mut buffer = Vec::new();
TextEncoder::new().encode(&self.registry.gather(), &mut buffer)?;
Ok(String::from_utf8(buffer)?)
}
fn is_ready(&self) -> bool {
self.ready.load(Ordering::Relaxed)
}
}
/// Serves the metrics endpoints until the process stops. Failing to bind is
/// logged rather than fatal: a worker that cannot publish metrics should still
/// do its work.
pub async fn serve(metrics: Metrics, addr: SocketAddr) -> Result<()> {
use axum::Router;
use axum::extract::State;
use axum::http::StatusCode;
use axum::routing::get;
let app = Router::new()
.route("/health", get(|| async { StatusCode::OK }))
.route(
"/ready",
get(|State(metrics): State<Metrics>| async move {
if metrics.is_ready() {
StatusCode::OK
} else {
StatusCode::SERVICE_UNAVAILABLE
}
}),
)
.route(
"/metrics",
get(|State(metrics): State<Metrics>| async move {
match metrics.gather() {
Ok(body) => (StatusCode::OK, body),
Err(err) => (StatusCode::INTERNAL_SERVER_ERROR, format!("{err:#}")),
}
}),
)
.with_state(metrics);
let listener = tokio::net::TcpListener::bind(addr)
.await
.with_context(|| format!("bind the metrics server to {addr}"))?;
info!("serving worker metrics on http://{addr}/metrics");
axum::serve(listener, app)
.await
.context("metrics server stopped")
}
/// Starts the metrics server in the background, warning rather than failing when
/// it cannot start.
pub fn spawn(metrics: Metrics, addr: SocketAddr) {
tokio::spawn(async move {
if let Err(err) = serve(metrics, addr).await {
warn!("worker metrics server: {err:#}");
}
});
}
#[cfg(test)]
mod tests {
use super::*;
fn metrics() -> Metrics {
Metrics::new("worker-1", "0.1.0").expect("build metrics")
}
// A scrape has to be able to tell these apart: a worker that swept and found
// nothing to do, and one that could not read anything it swept.
#[test]
fn a_sweep_reports_what_it_read_and_what_it_could_not() {
let metrics = metrics();
metrics.object_seen("lance_compact");
metrics.object_seen("lance_compact");
metrics.object_skipped("lance_compact", "open");
metrics.detection_finished("lance_compact", "ok", 0.5, 0);
let text = metrics.gather().expect("gather");
assert!(text.contains("SeaweedFS_worker_objects_seen_total{job_type=\"lance_compact\"} 2"));
assert!(text.contains(
"SeaweedFS_worker_objects_skipped_total{job_type=\"lance_compact\",reason=\"open\"} 1"
));
assert!(text.contains(
"SeaweedFS_worker_detections_total{job_type=\"lance_compact\",result=\"ok\"} 1"
));
assert!(text.contains("SeaweedFS_worker_proposals_total{job_type=\"lance_compact\"} 0"));
}
// /ready follows the stream, because a worker with no admin behind it is
// running but not doing anything.
#[test]
fn readiness_follows_the_stream() {
let metrics = metrics();
assert!(!metrics.is_ready(), "not ready before the stream is up");
metrics.stream_connected();
assert!(metrics.is_ready());
assert!(
metrics
.gather()
.unwrap()
.contains("SeaweedFS_worker_connected 1")
);
metrics.stream_ended("closed");
assert!(!metrics.is_ready());
assert!(
metrics
.gather()
.unwrap()
.contains("SeaweedFS_worker_connected 0")
);
assert!(
metrics
.gather()
.unwrap()
.contains("SeaweedFS_worker_stream_events_total{event=\"closed\"} 1")
);
}
#[test]
fn build_info_names_the_worker() {
let text = metrics().gather().expect("gather");
assert!(
text.contains(
"SeaweedFS_worker_build_info{version=\"0.1.0\",worker_id=\"worker-1\"} 1"
)
);
}
// A format's own numbers land on the same registry, so one endpoint serves
// both and a second worker implementation needs no new plumbing.
#[test]
fn a_worker_can_add_counters_of_its_own() {
let metrics = metrics();
let counter = metrics
.counter("lance_fragments_removed_total", "Fragments merged away.")
.expect("register");
counter.inc_by(16);
let text = metrics.gather().expect("gather");
assert!(text.contains("SeaweedFS_worker_lance_fragments_removed_total 16"));
}
}

View File

@@ -0,0 +1,89 @@
use std::collections::HashMap;
use std::sync::Arc;
use anyhow::Result;
use async_trait::async_trait;
use crate::pb::{ExecuteJobRequest, JobTypeCapability, JobTypeDescriptor, RunDetectionRequest};
use crate::senders::{DetectionSender, ExecutionSender};
/// One job type, worker side. Mirrors the Go JobHandler interface in
/// weed/plugin/worker/worker.go so the two stay readable against each other.
#[async_trait]
pub trait JobHandler: Send + Sync {
fn capability(&self) -> JobTypeCapability;
/// The descriptor admin renders as this job's settings page.
fn descriptor(&self) -> JobTypeDescriptor;
async fn detect(
&self,
request: &RunDetectionRequest,
sender: &dyn DetectionSender,
) -> Result<()>;
async fn execute(
&self,
request: &ExecuteJobRequest,
sender: &dyn ExecutionSender,
) -> Result<()>;
}
/// Sample rows of one object, already rendered as text. The worker formats
/// them because it is the only side that knows the object's types.
pub struct Preview {
pub columns: Vec<String>,
pub rows: Vec<Vec<String>>,
/// Rows in the object, which is not the number sampled.
pub total_rows: i64,
}
/// Reads sample rows of a format admin cannot parse itself.
///
/// This is deliberately not a JobHandler: a preview is answered while someone
/// waits on a page, so it neither schedules nor reports progress.
#[async_trait]
pub trait PreviewProvider: Send + Sync {
/// The format this provider reads, matched case-insensitively against what
/// the catalog recorded.
fn format(&self) -> &str;
async fn preview(&self, object_id: &[String], row_limit: usize) -> Result<Preview>;
}
/// The handlers one worker process serves. A process may serve several job
/// types, which is why WorkerHello carries a list of capabilities.
#[derive(Default, Clone)]
pub struct Registry {
handlers: HashMap<String, Arc<dyn JobHandler>>,
previews: HashMap<String, Arc<dyn PreviewProvider>>,
}
impl Registry {
pub fn new() -> Self {
Self::default()
}
pub fn register(mut self, handler: Arc<dyn JobHandler>) -> Self {
self.handlers.insert(handler.capability().job_type, handler);
self
}
pub fn with_preview(mut self, provider: Arc<dyn PreviewProvider>) -> Self {
self.previews
.insert(provider.format().to_ascii_uppercase(), provider);
self
}
pub fn get(&self, job_type: &str) -> Option<Arc<dyn JobHandler>> {
self.handlers.get(job_type).cloned()
}
pub fn preview_provider(&self, format: &str) -> Option<Arc<dyn PreviewProvider>> {
self.previews.get(&format.to_ascii_uppercase()).cloned()
}
pub fn capabilities(&self) -> Vec<JobTypeCapability> {
self.handlers.values().map(|h| h.capability()).collect()
}
pub fn is_empty(&self) -> bool {
self.handlers.is_empty()
}
}

View File

@@ -0,0 +1,203 @@
use anyhow::Result;
use tokio::sync::mpsc;
use crate::pb::{
ActivityEvent, DetectionComplete, DetectionProposals, JobCompleted, JobProgressUpdate,
WorkerObservations, WorkerToAdminMessage, worker_to_admin_message::Body,
};
/// Replies to one detection request.
pub trait DetectionSender: Send + Sync {
fn send_proposals(&self, proposals: DetectionProposals) -> Result<()>;
fn send_complete(&self, complete: DetectionComplete) -> Result<()>;
fn send_activity(&self, activity: ActivityEvent) -> Result<()>;
/// Reports what the worker saw while deciding. Admin caches the last one
/// per object and serves it back for display; nothing is scheduled from it.
fn send_observations(&self, observations: WorkerObservations) -> Result<()>;
}
/// Replies to one execution request.
pub trait ExecutionSender: Send + Sync {
fn send_progress(&self, progress: JobProgressUpdate) -> Result<()>;
fn send_completed(&self, completed: JobCompleted) -> Result<()>;
}
/// Both senders write to the single outbound stream, so they share one channel.
#[derive(Clone)]
pub struct StreamSender {
worker_id: String,
tx: mpsc::UnboundedSender<WorkerToAdminMessage>,
}
impl StreamSender {
pub fn new(worker_id: String, tx: mpsc::UnboundedSender<WorkerToAdminMessage>) -> Self {
Self { worker_id, tx }
}
pub fn send(&self, body: Body) -> Result<()> {
self.tx.send(WorkerToAdminMessage {
worker_id: self.worker_id.clone(),
sent_at: Some(std::time::SystemTime::now().into()),
body: Some(body),
})?;
Ok(())
}
}
impl DetectionSender for StreamSender {
fn send_proposals(&self, proposals: DetectionProposals) -> Result<()> {
self.send(Body::DetectionProposals(proposals))
}
fn send_complete(&self, complete: DetectionComplete) -> Result<()> {
self.send(Body::DetectionComplete(complete))
}
fn send_activity(&self, _activity: ActivityEvent) -> Result<()> {
// Activity rides inside progress and completion messages rather than
// being a body of its own, so there is nothing to send on its own here.
Ok(())
}
fn send_observations(&self, observations: WorkerObservations) -> Result<()> {
self.send(Body::Observations(observations))
}
}
impl ExecutionSender for StreamSender {
fn send_progress(&self, progress: JobProgressUpdate) -> Result<()> {
self.send(Body::JobProgressUpdate(progress))
}
fn send_completed(&self, completed: JobCompleted) -> Result<()> {
self.send(Body::JobCompleted(completed))
}
}
/// Wraps a sender so what passes through it is counted. The handlers report
/// their results to admin and nowhere else, so this is where a scrape can learn
/// what happened without every handler having to know about metrics.
pub struct MeteredSender<'a> {
inner: &'a StreamSender,
proposals: std::sync::atomic::AtomicUsize,
failed: std::sync::atomic::AtomicBool,
}
impl<'a> MeteredSender<'a> {
pub fn new(inner: &'a StreamSender) -> Self {
Self {
inner,
proposals: std::sync::atomic::AtomicUsize::new(0),
failed: std::sync::atomic::AtomicBool::new(false),
}
}
/// How many proposals went out, for the detection counter.
pub fn proposals(&self) -> usize {
self.proposals.load(std::sync::atomic::Ordering::Relaxed)
}
/// Whether the handler reported a failure of its own. A handler that fails
/// by returning an error is counted by the caller; this catches the one that
/// reports failure and returns Ok.
pub fn reported_failure(&self) -> bool {
self.failed.load(std::sync::atomic::Ordering::Relaxed)
}
}
impl DetectionSender for MeteredSender<'_> {
fn send_proposals(&self, proposals: DetectionProposals) -> Result<()> {
// 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<()> {
if !complete.success {
self.failed
.store(true, std::sync::atomic::Ordering::Relaxed);
}
self.inner.send_complete(complete)
}
fn send_activity(&self, activity: ActivityEvent) -> Result<()> {
self.inner.send_activity(activity)
}
fn send_observations(&self, observations: WorkerObservations) -> Result<()> {
self.inner.send_observations(observations)
}
}
impl ExecutionSender for MeteredSender<'_> {
fn send_progress(&self, progress: JobProgressUpdate) -> Result<()> {
self.inner.send_progress(progress)
}
fn send_completed(&self, completed: JobCompleted) -> Result<()> {
if !completed.success {
self.failed
.store(true, std::sync::atomic::Ordering::Relaxed);
}
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());
}
}

View File

@@ -0,0 +1,491 @@
use std::sync::Arc;
use std::time::{Duration, Instant};
use anyhow::{Context, Result, anyhow};
use tokio::sync::{Semaphore, mpsc};
use tokio_stream::wrappers::UnboundedReceiverStream;
use tonic::transport::{Certificate, Channel, ClientTlsConfig, Identity};
use tracing::{info, warn};
use crate::config::WorkerOptions;
use crate::metrics::Metrics;
use crate::pb::{
ConfigSchemaResponse, ExecuteJobRequest, JobCompleted, ObjectPreviewResponse, PreviewRow,
RequestObjectPreview, RunDetectionRequest, RunningWork, WorkerHeartbeat, WorkerHello,
admin_to_worker_message::Body as AdminBody,
plugin_control_service_client::PluginControlServiceClient,
worker_to_admin_message::Body as WorkerBody,
};
use crate::registry::Registry;
use crate::senders::{MeteredSender, StreamSender};
/// The protocol version this worker speaks, sent in WorkerHello.
const PROTOCOL_VERSION: &str = "1";
/// Connect to admin and serve the registry until the context is cancelled,
/// reconnecting on failure. The stream is the only channel: everything admin
/// asks for and everything the worker reports flows through it.
pub async fn run(options: WorkerOptions, registry: Registry) -> Result<()> {
let metrics = Metrics::new(&options.worker_id, &options.worker_version)?;
run_with_metrics(options, registry, metrics).await
}
/// Runs with metrics a caller has already made, so it can serve them and add
/// collectors of its own before the stream starts.
pub async fn run_with_metrics(
options: WorkerOptions,
registry: Registry,
metrics: Metrics,
) -> Result<()> {
if registry.is_empty() {
return Err(anyhow!("no job handlers registered"));
}
if options.max_detection_concurrency < 1 || options.max_execution_concurrency < 1 {
return Err(anyhow!(
"concurrency limits must be at least 1, got detection={} execution={}",
options.max_detection_concurrency,
options.max_execution_concurrency
));
}
let slots = Slots::new(&options);
metrics.set_slots("detection", 0, slots.detection_total as i64);
metrics.set_slots("execution", 0, slots.execution_total as i64);
loop {
match serve_once(&options, &registry, &slots, &metrics).await {
Err(err) => {
metrics.stream_ended("failed");
warn!("worker stream ended: {err:#}")
}
// Admin asked this worker to stop, so stop. Reconnecting here would
// make shutdown impossible: the worker would log back in.
Ok(Outcome::ShutdownRequested) => {
metrics.stream_ended("shutdown");
return Ok(());
}
// Admin closing a healthy stream is not an error, but reconnecting
// in silence hides the reason - two workers sharing an id evict
// each other and produce nothing but a login every few seconds.
Ok(Outcome::StreamClosed) => {
metrics.stream_ended("closed");
warn!(
"admin closed the stream; reconnecting in {:?}. If this repeats, check for \
another worker using the id {}",
options.reconnect_delay, options.worker_id
)
}
}
tokio::time::sleep(options.reconnect_delay).await;
}
}
/// Why a stream ended. Only one of these means "do not come back".
enum Outcome {
StreamClosed,
ShutdownRequested,
}
/// The capacity this worker advertises in WorkerHello. Admin schedules against
/// those numbers, so the worker has to actually hold to them - and the heartbeat
/// has to report what is in use, or admin is scheduling blind.
#[derive(Clone)]
struct Slots {
detection: Arc<Semaphore>,
execution: Arc<Semaphore>,
detection_total: i32,
execution_total: i32,
}
impl Slots {
fn new(options: &WorkerOptions) -> Self {
Self {
detection: Arc::new(Semaphore::new(options.max_detection_concurrency as usize)),
execution: Arc::new(Semaphore::new(options.max_execution_concurrency as usize)),
detection_total: options.max_detection_concurrency,
execution_total: options.max_execution_concurrency,
}
}
fn detection_used(&self) -> i32 {
self.detection_total - self.detection.available_permits() as i32
}
fn execution_used(&self) -> i32 {
self.execution_total - self.execution.available_permits() as i32
}
}
/// Dials admin, over mTLS when certificates are configured. Plaintext is the
/// default and is fine over loopback; anything else carries preview rows and
/// execution commands in the clear, and a cluster with grpc TLS on refuses the
/// connection anyway.
async fn connect(options: &WorkerOptions, grpc_address: &str) -> Result<Channel> {
let Some(tls) = options.tls.as_ref() else {
return Ok(Channel::from_shared(format!("http://{grpc_address}"))?
.connect_timeout(Duration::from_secs(10))
.connect()
.await?);
};
let ca = tokio::fs::read(&tls.ca_path)
.await
.with_context(|| format!("read CA certificate {}", tls.ca_path))?;
let cert = tokio::fs::read(&tls.client_cert_path)
.await
.with_context(|| format!("read client certificate {}", tls.client_cert_path))?;
let key = tokio::fs::read(&tls.client_key_path)
.await
.with_context(|| format!("read client key {}", tls.client_key_path))?;
let mut config = ClientTlsConfig::new()
.ca_certificate(Certificate::from_pem(ca))
.identity(Identity::from_pem(cert, key));
if let Some(server_name) = tls.server_name.as_ref() {
config = config.domain_name(server_name.clone());
}
Ok(Channel::from_shared(format!("https://{grpc_address}"))?
.tls_config(config)?
.connect_timeout(Duration::from_secs(10))
.connect()
.await?)
}
async fn serve_once(
options: &WorkerOptions,
registry: &Registry,
slots: &Slots,
metrics: &Metrics,
) -> Result<Outcome> {
// Operators give the admin's HTTP address; the gRPC port is derived, the
// same way the Go worker does it.
let grpc_address = crate::address::server_to_grpc_address(&options.admin_address)
.ok_or_else(|| anyhow!("cannot parse admin address {}", options.admin_address))?;
let channel = connect(options, &grpc_address).await?;
let mut client = PluginControlServiceClient::new(channel);
let (tx, rx) = mpsc::unbounded_channel();
let sender = StreamSender::new(options.worker_id.clone(), tx);
sender.send(WorkerBody::Hello(WorkerHello {
worker_id: options.worker_id.clone(),
worker_instance_id: options.worker_id.clone(),
address: options.worker_address.clone(),
worker_version: options.worker_version.clone(),
protocol_version: PROTOCOL_VERSION.to_string(),
capabilities: registry.capabilities(),
metadata: Default::default(),
}))?;
let mut inbound = client
.worker_stream(UnboundedReceiverStream::new(rx))
.await?
.into_inner();
let heartbeat = spawn_heartbeat(
sender.clone(),
options.clone(),
slots.clone(),
metrics.clone(),
);
while let Some(message) = inbound.message().await? {
let request_id = message.request_id.clone();
match message.body {
Some(AdminBody::Hello(hello)) => {
if !hello.accepted {
return Err(anyhow!("admin rejected this worker: {}", hello.message));
}
metrics.stream_connected();
info!(
"connected to admin at {} ({})",
options.admin_address, grpc_address
);
}
Some(AdminBody::RequestConfigSchema(request)) => {
let response = match registry.get(&request.job_type) {
Some(handler) => ConfigSchemaResponse {
request_id: request_id.clone(),
job_type: request.job_type.clone(),
success: true,
error_message: String::new(),
job_type_descriptor: Some(handler.descriptor()),
},
None => ConfigSchemaResponse {
request_id: request_id.clone(),
job_type: request.job_type.clone(),
success: false,
error_message: format!("unknown job type {}", request.job_type),
job_type_descriptor: None,
},
};
sender.send(WorkerBody::ConfigSchemaResponse(response))?;
}
Some(AdminBody::RequestObjectPreview(request)) => {
spawn_preview(
registry.clone(),
sender.clone(),
metrics.clone(),
request_id.clone(),
request,
);
}
Some(AdminBody::RunDetectionRequest(request)) => {
spawn_detection(
registry.clone(),
sender.clone(),
slots.clone(),
metrics.clone(),
request,
);
}
Some(AdminBody::ExecuteJobRequest(request)) => {
spawn_execution(
registry.clone(),
sender.clone(),
slots.clone(),
metrics.clone(),
request,
);
}
Some(AdminBody::CancelRequest(request)) => {
// Cancellation needs a per-request handle to be honoured; until
// then say so rather than silently continuing to run the job.
warn!(
"cancel requested for {} ({}) but is not implemented",
request.target_id, request.reason
);
}
Some(AdminBody::Shutdown(shutdown)) => {
info!("admin asked this worker to stop: {}", shutdown.reason);
heartbeat.abort();
return Ok(Outcome::ShutdownRequested);
}
None => {}
}
}
heartbeat.abort();
Ok(Outcome::StreamClosed)
}
/// Answers one preview request off the stream loop. Reading rows takes as long
/// as it takes, and the stream has heartbeats to keep up meanwhile.
fn spawn_preview(
registry: Registry,
sender: StreamSender,
metrics: Metrics,
request_id: String,
request: RequestObjectPreview,
) {
tokio::spawn(async move {
let limit = request.row_limit.max(1) as usize;
let response = match registry.preview_provider(&request.format) {
None => ObjectPreviewResponse {
request_id,
success: false,
error_message: format!("this worker does not read {} objects", request.format),
..Default::default()
},
Some(provider) => match provider.preview(&request.object_id, limit).await {
Ok(preview) => ObjectPreviewResponse {
request_id,
success: true,
error_message: String::new(),
columns: preview.columns,
rows: preview
.rows
.into_iter()
.map(|values| PreviewRow { values })
.collect(),
total_rows: preview.total_rows,
},
Err(err) => ObjectPreviewResponse {
request_id,
success: false,
error_message: format!("{err:#}"),
..Default::default()
},
},
};
metrics.preview_finished(if response.success { "ok" } else { "failed" });
let _ = sender.send(WorkerBody::ObjectPreviewResponse(response));
});
}
fn spawn_heartbeat(
sender: StreamSender,
options: WorkerOptions,
slots: Slots,
metrics: Metrics,
) -> tokio::task::JoinHandle<()> {
// The handle has to be the heartbeat's own, or aborting it aborts nothing
// and every reconnect leaves another ticker running.
tokio::spawn(async move {
let mut ticker = tokio::time::interval(options.heartbeat_interval);
loop {
ticker.tick().await;
// The heartbeat already computes this for admin; publish the same
// numbers so a scrape and the admin UI cannot disagree.
metrics.set_slots(
"detection",
slots.detection_used() as i64,
slots.detection_total as i64,
);
metrics.set_slots(
"execution",
slots.execution_used() as i64,
slots.execution_total as i64,
);
let beat = WorkerHeartbeat {
worker_id: options.worker_id.clone(),
running_work: Vec::<RunningWork>::new(),
detection_slots_used: slots.detection_used(),
detection_slots_total: slots.detection_total,
execution_slots_used: slots.execution_used(),
execution_slots_total: slots.execution_total,
queued_jobs_by_type: Default::default(),
metadata: Default::default(),
};
if sender.send(WorkerBody::Heartbeat(beat)).is_err() {
return;
}
}
})
}
fn spawn_detection(
registry: Registry,
sender: StreamSender,
slots: Slots,
metrics: Metrics,
request: RunDetectionRequest,
) {
tokio::spawn(async move {
let Some(handler) = registry.get(&request.job_type) else {
return;
};
// 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);
let started = Instant::now();
let outcome = handler.detect(&request, &metered).await;
let result = if let Err(err) = &outcome {
warn!("detection for {} failed: {err:#}", request.job_type);
let _ = sender.send(WorkerBody::DetectionComplete(
crate::pb::DetectionComplete {
request_id: request.request_id.clone(),
job_type: request.job_type.clone(),
success: false,
error_message: format!("{err:#}"),
total_proposals: 0,
},
));
"failed"
} else if metered.reported_failure() {
"failed"
} else {
"ok"
};
metrics.detection_finished(
&request.job_type,
result,
started.elapsed().as_secs_f64(),
metered.proposals(),
);
});
}
fn spawn_execution(
registry: Registry,
sender: StreamSender,
slots: Slots,
metrics: Metrics,
request: ExecuteJobRequest,
) {
tokio::spawn(async move {
let _permit = slots.execution.acquire().await;
let job_type = request
.job
.as_ref()
.map(|job| job.job_type.clone())
.unwrap_or_default();
let job_id = request
.job
.as_ref()
.map(|job| job.job_id.clone())
.unwrap_or_default();
let Some(handler) = registry.get(&job_type) else {
return;
};
let metered = MeteredSender::new(&sender);
let started = Instant::now();
let outcome = handler.execute(&request, &metered).await;
let result = if let Err(err) = &outcome {
warn!("job {job_id} failed: {err:#}");
let _ = sender.send(WorkerBody::JobCompleted(JobCompleted {
request_id: request.request_id.clone(),
job_id,
job_type: job_type.clone(),
success: false,
error_message: format!("{err:#}"),
..Default::default()
}));
"failed"
} else if metered.reported_failure() {
"failed"
} else {
"ok"
};
metrics.job_finished(&job_type, result, started.elapsed().as_secs_f64());
});
}
#[cfg(test)]
mod tests {
use super::*;
fn options(detection: i32, execution: i32) -> WorkerOptions {
WorkerOptions {
max_detection_concurrency: detection,
max_execution_concurrency: execution,
..Default::default()
}
}
// The heartbeat is admin's only view of how busy this worker is; before the
// permits existed it reported zero however much was running.
#[tokio::test]
async fn slots_report_what_is_held() {
let slots = Slots::new(&options(2, 3));
assert_eq!(slots.detection_used(), 0);
assert_eq!(slots.execution_used(), 0);
let held = slots.detection.acquire().await.unwrap();
assert_eq!(slots.detection_used(), 1);
assert_eq!(slots.execution_used(), 0, "the two do not share capacity");
drop(held);
assert_eq!(slots.detection_used(), 0);
}
// A limit of one means the second request waits, rather than running anyway
// as it did when every request simply spawned a task.
#[tokio::test]
async fn a_full_lane_makes_the_next_request_wait() {
let slots = Slots::new(&options(1, 1));
let held = slots.execution.clone().acquire_owned().await.unwrap();
assert_eq!(slots.execution_used(), 1);
let waiter = tokio::spawn({
let execution = slots.execution.clone();
async move { execution.acquire_owned().await.unwrap() }
});
tokio::task::yield_now().await;
assert!(!waiter.is_finished(), "the second request must not start");
drop(held);
let _second = waiter.await.expect("the waiter should be handed the slot");
assert_eq!(slots.execution_used(), 1);
}
}