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);
}
}

View File

@@ -0,0 +1,55 @@
[package]
name = "weed-lance-worker"
version.workspace = true
edition.workspace = true
rust-version.workspace = true
description = "SeaweedFS maintenance worker for Lance tables"
[lib]
name = "weed_lance_worker"
# The binary is not named for lance: it is the Rust side of `weed worker`, and
# the job families it registers will outgrow this crate. When a second one
# arrives the bin target moves to a crate of its own under the same name.
[[bin]]
name = "weed-worker"
path = "src/main.rs"
[dependencies]
seaweed-worker-core = { path = "../core" }
seaweed-worker-sort = { path = "../sort" }
prometheus.workspace = true
# Only the S3 backend: the other object stores lance enables by default are
# build time this worker never spends.
lance = { version = "10", default-features = false, features = ["aws"] }
lance-index = "10"
# For the manifest's own fragment type, which the sort marker summarises.
lance-table = "10"
# For the execution options the sort needs; the scanner's own stream helper
# takes the default, which has spilling off.
lance-datafusion = "10"
arrow-schema = "58"
arrow-cast = "58"
chrono = "0.4"
futures = "0.3"
anyhow.workspace = true
async-trait.workspace = true
clap = { version = "4", features = ["derive", "env"] }
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
rustls = "0.23"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
tokio.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
[dev-dependencies]
tokio = { workspace = true }
tonic.workspace = true
arrow-array = "58"
arrow-schema = "58"
arrow-cast = "58"
lance-linalg = "10"
[lints]
workspace = true

View File

@@ -0,0 +1,10 @@
//! How the worker finds Lance tables and gets at their bytes.
//!
//! It goes through the Lance namespace rather than the filer: the namespace is
//! the catalog of record, it already knows which tables are Lance, and asking it
//! to describe a table with vend_credentials is how the worker gets storage
//! credentials without holding any of its own.
pub mod namespace;
pub use namespace::{NamespaceClient, TableDescription, parse_id};

View File

@@ -0,0 +1,129 @@
use std::collections::HashMap;
use std::time::Duration;
use anyhow::{Context, Result};
use serde::Deserialize;
/// The delimiter the Lance namespace joins identifier parts with.
const DELIMITER: &str = "$";
#[derive(Debug, Deserialize)]
pub struct TableDescription {
pub location: String,
#[serde(default)]
pub storage_options: HashMap<String, String>,
#[serde(default)]
pub managed_versioning: bool,
}
#[derive(Debug, Deserialize)]
struct ListNamespacesResponse {
#[serde(default)]
namespaces: Vec<String>,
}
#[derive(Debug, Deserialize)]
struct ListTablesResponse {
#[serde(default)]
tables: Vec<String>,
}
/// A thin client for the operations this worker needs. It deliberately does not
/// wrap the whole spec: a maintenance worker lists, describes, and commits.
pub struct NamespaceClient {
base_url: String,
http: reqwest::Client,
}
/// A namespace call that has not answered by now is not going to. Without this
/// a gateway that accepts the connection and then goes quiet holds a detection
/// slot open forever, and the sweep never finishes.
const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
impl NamespaceClient {
pub fn new(base_url: impl Into<String>) -> Self {
let http = reqwest::Client::builder()
.timeout(REQUEST_TIMEOUT)
.connect_timeout(CONNECT_TIMEOUT)
.build()
// The builder only fails on a bad TLS backend, which would break
// every call anyway; a client with no timeouts is worse than a panic
// at startup, so keep the default only as a last resort.
.unwrap_or_else(|err| {
tracing::warn!("falling back to an untimed HTTP client: {err}");
reqwest::Client::new()
});
Self {
base_url: base_url.into().trim_end_matches('/').to_string(),
http,
}
}
/// Every table the namespace holds, as delimiter-joined identifiers.
pub async fn list_all_tables(&self) -> Result<Vec<String>> {
let url = format!("{}/v1/table", self.base_url);
let response: ListTablesResponse = self
.http
.get(&url)
.send()
.await
.context("list tables")?
.error_for_status()?
.json()
.await?;
Ok(response.tables)
}
/// Child namespaces of `id`; the root lists table buckets.
pub async fn list_namespaces(&self, id: &[String]) -> Result<Vec<String>> {
let url = format!("{}/v1/namespace/{}/list", self.base_url, encode_id(id));
let response: ListNamespacesResponse = self
.http
.get(&url)
.send()
.await
.context("list namespaces")?
.error_for_status()?
.json()
.await?;
Ok(response.namespaces)
}
/// Resolve a table to a location and the credentials to reach it. The
/// credentials expire, so a long compaction re-describes rather than
/// carrying one set for the whole job.
pub async fn describe_table(&self, id: &[String]) -> Result<TableDescription> {
let url = format!("{}/v1/table/{}/describe", self.base_url, encode_id(id));
let body = serde_json::json!({ "id": id, "vend_credentials": true });
let description: TableDescription = self
.http
.post(&url)
.json(&body)
.send()
.await
.context("describe table")?
.error_for_status()?
.json()
.await?;
Ok(description)
}
}
fn encode_id(id: &[String]) -> String {
if id.is_empty() {
DELIMITER.to_string()
} else {
id.join(DELIMITER)
}
}
/// Splits a delimiter-joined identifier back into parts.
pub fn parse_id(encoded: &str) -> Vec<String> {
encoded
.split(DELIMITER)
.filter(|part| !part.is_empty())
.map(|part| part.to_string())
.collect()
}

View File

@@ -0,0 +1,98 @@
//! Opening a Lance dataset with credentials the namespace vended.
//!
//! The worker holds no storage credentials of its own: it asks the namespace to
//! describe a table with `vend_credentials`, and the `storage_options` that come
//! back are handed to lance as-is. They expire, so a job that runs longer than
//! their lifetime re-describes rather than carrying one set throughout.
use std::collections::HashMap;
use anyhow::{Context, Result};
use lance::dataset::Dataset;
use lance::dataset::builder::DatasetBuilder;
use crate::catalog::{NamespaceClient, TableDescription};
/// A table the worker is about to work on.
pub struct OpenTable {
pub id: Vec<String>,
pub location: String,
pub dataset: Dataset,
}
/// What detection needs to decide whether a table is worth a job. Reading it
/// opens the dataset but touches no data files.
pub struct TableStats {
pub fragments: usize,
pub version: u64,
pub total_versions: usize,
pub rows: usize,
/// The Arrow schema as JSON, which is the only description of this table
/// anything outside the format can produce.
pub schema: Option<String>,
}
/// Storage options an operator supplies for deployments that vend none.
pub type FallbackOptions = HashMap<String, String>;
/// Resolve a table through the namespace and open it.
pub async fn open(
client: &NamespaceClient,
id: &[String],
fallback: &FallbackOptions,
) -> Result<OpenTable> {
let description = client.describe_table(id).await?;
let dataset = open_at(&description, fallback).await?;
Ok(OpenTable {
id: id.to_vec(),
location: description.location,
dataset,
})
}
/// The namespace vends object_store's own option names, so they pass straight
/// through. What it vends always wins: the fallback exists because a deployment
/// without STS vends no credentials at all, and then the worker has no other way
/// to reach the data.
async fn open_at(description: &TableDescription, fallback: &FallbackOptions) -> Result<Dataset> {
let mut options: HashMap<String, String> = fallback.clone();
options.extend(description.storage_options.clone());
DatasetBuilder::from_uri(&description.location)
.with_storage_options(options)
.load()
.await
.with_context(|| format!("open lance dataset at {}", description.location))
}
impl OpenTable {
pub async fn stats(&self) -> Result<TableStats> {
let versions = self.dataset.versions().await?;
Ok(TableStats {
fragments: self.dataset.get_fragments().len(),
version: self.dataset.version().version,
total_versions: versions.len(),
rows: self.dataset.count_rows(None).await.unwrap_or(0),
schema: schema_json(&self.dataset),
})
}
}
/// Renders the dataset's schema as JSON. Best effort: a schema that will not
/// serialise is not a reason to fail a maintenance sweep.
fn schema_json(dataset: &Dataset) -> Option<String> {
let arrow: arrow_schema::Schema = dataset.schema().into();
serde_json::to_string(
&arrow
.fields()
.iter()
.map(|f| {
serde_json::json!({
"name": f.name(),
"type": f.data_type().to_string(),
"nullable": f.is_nullable(),
})
})
.collect::<Vec<_>>(),
)
.ok()
}

View File

@@ -0,0 +1,336 @@
use std::collections::HashMap;
use anyhow::{Context, Result, anyhow};
use async_trait::async_trait;
use chrono::{Duration, Utc};
use lance::dataset::cleanup::{CleanupPolicy, cleanup_old_versions};
use seaweed_worker_core::config_form::{form, int_or, int_value, number_field};
use seaweed_worker_core::pb::{
ConfigValue, DetectionComplete, DetectionProposals, ExecuteJobRequest, JobCompleted,
JobProgressUpdate, JobProposal, JobResult, JobTypeCapability, JobTypeDescriptor,
RunDetectionRequest,
};
use seaweed_worker_core::{DetectionSender, ExecutionSender, JobHandler};
use tracing::warn;
use crate::catalog::{NamespaceClient, parse_id};
use crate::dataset;
use crate::jobs::{clamp, string_list, table_id};
pub const JOB_TYPE: &str = "lance_cleanup_versions";
const DEFAULT_RETAIN_HOURS: i64 = 168;
const DEFAULT_MIN_VERSIONS: i64 = 5;
// The ranges the form offers. Duration::hours also panics far outside this one.
const MAX_RETAIN_HOURS: i64 = 8760;
const MAX_MIN_VERSIONS: i64 = 1000;
/// Lance keeps every version until something removes it. Lance can also do this
/// itself through auto-cleanup, so this job is for deployments that would rather
/// the cluster owned the policy than each writer.
pub struct CleanupVersionsHandler {
namespace_url: String,
fallback: dataset::FallbackOptions,
metrics: Option<crate::metrics::LanceMetrics>,
}
impl CleanupVersionsHandler {
pub fn new(namespace_url: String) -> Self {
Self {
namespace_url,
fallback: dataset::FallbackOptions::new(),
metrics: None,
}
}
pub fn with_metrics(mut self, metrics: Option<crate::metrics::LanceMetrics>) -> Self {
self.metrics = metrics;
self
}
pub fn with_fallback(mut self, fallback: dataset::FallbackOptions) -> Self {
self.fallback = fallback;
self
}
fn client(&self) -> NamespaceClient {
NamespaceClient::new(self.namespace_url.clone())
}
}
/// The oldest version that may be removed while still leaving `min_versions`
/// behind, or None when the table has no more than the floor.
async fn version_floor(table: &dataset::OpenTable, min_versions: usize) -> Result<Option<u64>> {
let mut versions: Vec<u64> = table
.dataset
.versions()
.await?
.iter()
.map(|version| version.version)
.collect();
if versions.len() <= min_versions {
return Ok(None);
}
versions.sort_unstable();
Ok(Some(versions[versions.len() - min_versions]))
}
#[async_trait]
impl JobHandler for CleanupVersionsHandler {
fn capability(&self) -> JobTypeCapability {
JobTypeCapability {
job_type: JOB_TYPE.to_string(),
can_detect: true,
can_execute: true,
max_detection_concurrency: 1,
max_execution_concurrency: 1,
display_name: "Lance Version Cleanup".to_string(),
description: "Remove old Lance versions and the files only they referenced".to_string(),
weight: 10,
}
}
fn descriptor(&self) -> JobTypeDescriptor {
let mut defaults: HashMap<String, ConfigValue> = HashMap::new();
defaults.insert("retain_hours".to_string(), int_value(DEFAULT_RETAIN_HOURS));
defaults.insert(
"min_versions_to_keep".to_string(),
int_value(DEFAULT_MIN_VERSIONS),
);
JobTypeDescriptor {
job_type: JOB_TYPE.to_string(),
display_name: "Lance Version Cleanup".to_string(),
description: "Age out Lance versions the table no longer needs".to_string(),
icon: "fas fa-broom".to_string(),
descriptor_version: 1,
worker_config_form: Some(form(
"lance-cleanup-worker",
"Version cleanup",
vec![
number_field(
"retain_hours",
"Retain versions for (hours)",
"Versions younger than this are always kept",
1,
8760,
),
number_field(
"min_versions_to_keep",
"Minimum versions",
"Never leave a table with fewer versions than this, whatever their age",
1,
1000,
),
],
defaults.clone(),
)),
worker_default_values: defaults,
..Default::default()
}
}
async fn detect(
&self,
request: &RunDetectionRequest,
sender: &dyn DetectionSender,
) -> Result<()> {
let min_versions = clamp(
int_or(
&request.worker_config_values,
"min_versions_to_keep",
DEFAULT_MIN_VERSIONS,
),
1,
MAX_MIN_VERSIONS,
) as usize;
let client = self.client();
let tables = client.list_all_tables().await?;
let mut proposals = Vec::new();
for encoded in &tables {
let id = parse_id(encoded);
let table = match dataset::open(&client, &id, &self.fallback).await {
Ok(table) => table,
Err(err) => {
if let Some(counters) = &self.metrics {
counters.worker.object_skipped(JOB_TYPE, "open");
}
warn!("skipping {encoded}: {err:#}");
continue;
}
};
if let Some(counters) = &self.metrics {
counters.worker.object_seen(JOB_TYPE);
}
let stats = match table.stats().await {
Ok(stats) => stats,
Err(err) => {
if let Some(counters) = &self.metrics {
counters.worker.object_skipped(JOB_TYPE, "stats");
}
warn!("skipping {encoded}: reading its stats failed: {err:#}");
continue;
}
};
// Age is decided at execution against the retention window; a table
// at or under the floor cannot lose a version whatever its age, so
// proposing one would only produce a job with nothing to do.
if stats.total_versions <= min_versions {
continue;
}
let mut parameters: HashMap<String, ConfigValue> = HashMap::new();
parameters.insert("table_id".to_string(), string_list(&id));
proposals.push(JobProposal {
proposal_id: format!("{JOB_TYPE}:{encoded}"),
dedupe_key: format!("{JOB_TYPE}:{encoded}"),
job_type: JOB_TYPE.to_string(),
summary: format!("Clean up {encoded} ({} versions)", stats.total_versions),
detail: format!(
"{} versions retained, above the {min_versions} floor; \
those outside the retention window can go",
stats.total_versions
),
parameters,
..Default::default()
});
}
let total = proposals.len() as i32;
sender.send_proposals(DetectionProposals {
request_id: request.request_id.clone(),
job_type: JOB_TYPE.to_string(),
proposals,
has_more: false,
})?;
sender.send_complete(DetectionComplete {
request_id: request.request_id.clone(),
job_type: JOB_TYPE.to_string(),
success: true,
error_message: String::new(),
total_proposals: total,
})?;
Ok(())
}
async fn execute(
&self,
request: &ExecuteJobRequest,
sender: &dyn ExecutionSender,
) -> Result<()> {
let job = request
.job
.as_ref()
.ok_or_else(|| anyhow!("execute request carried no job"))?;
let id = table_id(&job.parameters)
.ok_or_else(|| anyhow!("job {} carried no table_id", job.job_id))?;
let retain_hours = clamp(
int_or(
&request.worker_config_values,
"retain_hours",
DEFAULT_RETAIN_HOURS,
),
0,
MAX_RETAIN_HOURS,
);
let min_versions = clamp(
int_or(
&request.worker_config_values,
"min_versions_to_keep",
DEFAULT_MIN_VERSIONS,
),
1,
MAX_MIN_VERSIONS,
) as usize;
let client = self.client();
let table = dataset::open(&client, &id, &self.fallback).await?;
let before = table.stats().await?;
// The floor is a promise about how much history survives, so it has to
// be applied here and not only when the job was proposed: by the time it
// runs, versions may have aged past the retention window, and age alone
// would take the table below what the operator asked to keep.
let Some(floor) = version_floor(&table, min_versions).await? else {
sender.send_completed(JobCompleted {
request_id: request.request_id.clone(),
job_id: job.job_id.clone(),
job_type: JOB_TYPE.to_string(),
success: true,
result: Some(JobResult {
summary: format!(
"kept all {} versions; the {min_versions} version floor leaves none to remove",
before.total_versions
),
..Default::default()
}),
..Default::default()
})?;
return Ok(());
};
sender.send_progress(JobProgressUpdate {
request_id: request.request_id.clone(),
job_id: job.job_id.clone(),
job_type: JOB_TYPE.to_string(),
progress_percent: 10.0,
stage: format!("cleaning up {} versions", before.total_versions),
..Default::default()
})?;
let policy = CleanupPolicy {
before_timestamp: Some(Utc::now() - Duration::hours(retain_hours)),
// should_clean() ANDs its clauses, so a version has to be both older
// than the window and below the floor to go.
before_version: Some(floor),
// Files this dataset cannot account for are left alone: they may
// belong to a writer that has not committed yet, and deleting them
// would corrupt a commit in flight.
delete_unverified: false,
// A tagged version is pinned on purpose, so refuse rather than
// silently dropping what someone named.
error_if_tagged_old_versions: true,
..Default::default()
};
let stats = cleanup_old_versions(&table.dataset, policy)
.await
.with_context(|| format!("clean up versions of {}", table.location))?;
if let Some(counters) = &self.metrics {
counters.versions_removed.inc_by(stats.old_versions as u64);
counters.bytes_reclaimed.inc_by(stats.bytes_removed as u64);
}
let mut output: HashMap<String, ConfigValue> = HashMap::new();
output.insert(
"old_versions_removed".to_string(),
int_value(stats.old_versions as i64),
);
output.insert(
"bytes_removed".to_string(),
int_value(stats.bytes_removed as i64),
);
output.insert(
"data_files_removed".to_string(),
int_value(stats.data_files_removed as i64),
);
sender.send_completed(JobCompleted {
request_id: request.request_id.clone(),
job_id: job.job_id.clone(),
job_type: JOB_TYPE.to_string(),
success: true,
result: Some(JobResult {
output_values: output,
summary: format!(
"removed {} versions and {} bytes",
stats.old_versions, stats.bytes_removed
),
..Default::default()
}),
..Default::default()
})?;
Ok(())
}
}

View File

@@ -0,0 +1,318 @@
use std::collections::HashMap;
use anyhow::{Context, Result, anyhow};
use async_trait::async_trait;
use lance::dataset::optimize::{CompactionOptions, compact_files};
use seaweed_worker_core::config_form::{form, int_or, int_value, number_field, string_value};
use seaweed_worker_core::pb::{
ConfigValue, DetectionComplete, DetectionProposals, ExecuteJobRequest, JobCompleted,
JobProgressUpdate, JobProposal, JobResult, JobTypeCapability, JobTypeDescriptor,
RunDetectionRequest, WorkerObservations,
};
use seaweed_worker_core::{DetectionSender, ExecutionSender, JobHandler};
use tracing::warn;
use crate::catalog::{NamespaceClient, parse_id};
use crate::dataset;
use crate::jobs::{FORMAT, clamp, observation, string_list, table_id};
pub const JOB_TYPE: &str = "lance_compact";
const DEFAULT_TARGET_ROWS: i64 = 1_048_576;
const DEFAULT_MIN_FRAGMENTS: i64 = 8;
// The ranges the descriptor's form offers. These are also the values that stay
// meaningful: a table needs two fragments before merging them means anything,
// and a fragment target below a thousand rows defeats the purpose of the job.
// Both are cast to usize, where a negative would arrive as an enormous number.
const TARGET_ROWS_FLOOR: i64 = 1024;
const TARGET_ROWS_CEILING: i64 = 16_777_216;
const MIN_FRAGMENTS_FLOOR: i64 = 2;
const MIN_FRAGMENTS_CEILING: i64 = 4096;
/// Lance writes one fragment per write batch, so a table fed by small appends
/// accumulates small files the same way an Iceberg table does.
pub struct CompactHandler {
namespace_url: String,
fallback: dataset::FallbackOptions,
metrics: Option<crate::metrics::LanceMetrics>,
}
impl CompactHandler {
pub fn new(namespace_url: String) -> Self {
Self {
namespace_url,
fallback: dataset::FallbackOptions::new(),
metrics: None,
}
}
pub fn with_metrics(mut self, metrics: Option<crate::metrics::LanceMetrics>) -> Self {
self.metrics = metrics;
self
}
/// Storage options to use where the namespace vends none.
pub fn with_fallback(mut self, fallback: dataset::FallbackOptions) -> Self {
self.fallback = fallback;
self
}
fn client(&self) -> NamespaceClient {
NamespaceClient::new(self.namespace_url.clone())
}
}
#[async_trait]
impl JobHandler for CompactHandler {
fn capability(&self) -> JobTypeCapability {
JobTypeCapability {
job_type: JOB_TYPE.to_string(),
can_detect: true,
can_execute: true,
max_detection_concurrency: 1,
max_execution_concurrency: 1,
display_name: "Lance Compaction".to_string(),
description: "Merge small Lance fragments into fewer, larger ones".to_string(),
weight: 20,
}
}
fn descriptor(&self) -> JobTypeDescriptor {
let mut defaults: HashMap<String, ConfigValue> = HashMap::new();
defaults.insert(
"target_rows_per_fragment".to_string(),
int_value(DEFAULT_TARGET_ROWS),
);
defaults.insert(
"min_fragments".to_string(),
int_value(DEFAULT_MIN_FRAGMENTS),
);
JobTypeDescriptor {
job_type: JOB_TYPE.to_string(),
display_name: "Lance Compaction".to_string(),
description: "Compact fragments of Lance tables".to_string(),
icon: "fas fa-compress".to_string(),
descriptor_version: 1,
worker_config_form: Some(form(
"lance-compact-worker",
"Compaction",
vec![
number_field(
"target_rows_per_fragment",
"Target rows per fragment",
"Rows to aim for when rewriting fragments",
1024,
16_777_216,
),
number_field(
"min_fragments",
"Minimum fragments",
"Leave a table alone until it has at least this many fragments",
2,
4096,
),
],
defaults.clone(),
)),
worker_default_values: defaults,
..Default::default()
}
}
/// Propose a job for every table with more fragments than the operator is
/// willing to leave alone. Opening a dataset reads its manifest, not its
/// data, so this stays cheap across a catalog.
async fn detect(
&self,
request: &RunDetectionRequest,
sender: &dyn DetectionSender,
) -> Result<()> {
let min_fragments = clamp(
int_or(
&request.worker_config_values,
"min_fragments",
DEFAULT_MIN_FRAGMENTS,
),
MIN_FRAGMENTS_FLOOR,
MIN_FRAGMENTS_CEILING,
) as usize;
let client = self.client();
let tables = client.list_all_tables().await?;
let mut proposals = Vec::new();
let mut observations = Vec::new();
for encoded in &tables {
let id = parse_id(encoded);
let table = match dataset::open(&client, &id, &self.fallback).await {
Ok(table) => table,
Err(err) => {
// A table that cannot be opened is the next run's problem,
// not a reason to abandon the whole sweep.
if let Some(counters) = &self.metrics {
counters.worker.object_skipped(JOB_TYPE, "open");
}
warn!("skipping {encoded}: {err:#}");
continue;
}
};
// One unreadable table must not end the sweep: the tables already
// read would lose their proposals, and admin would get no
// completion for this request at all.
if let Some(counters) = &self.metrics {
counters.worker.object_seen(JOB_TYPE);
}
let stats = match table.stats().await {
Ok(stats) => stats,
Err(err) => {
if let Some(counters) = &self.metrics {
counters.worker.object_skipped(JOB_TYPE, "stats");
}
warn!("skipping {encoded}: reading its stats failed: {err:#}");
continue;
}
};
// Logged because "detection proposed nothing" is otherwise
// indistinguishable from a table the worker could not read.
tracing::info!(
"compaction detection: {encoded} has {} fragments, threshold {min_fragments}",
stats.fragments
);
let mut attributes: HashMap<String, ConfigValue> = HashMap::new();
attributes.insert("fragments".to_string(), int_value(stats.fragments as i64));
attributes.insert("version".to_string(), int_value(stats.version as i64));
attributes.insert(
"versions".to_string(),
int_value(stats.total_versions as i64),
);
attributes.insert("rows".to_string(), int_value(stats.rows as i64));
if let Some(schema) = stats.schema.clone() {
attributes.insert("schema".to_string(), string_value(schema));
}
observations.push(observation(&id, FORMAT, attributes));
if stats.fragments < min_fragments {
continue;
}
let mut parameters: HashMap<String, ConfigValue> = HashMap::new();
parameters.insert("table_id".to_string(), string_list(&id));
proposals.push(JobProposal {
proposal_id: format!("{JOB_TYPE}:{encoded}"),
dedupe_key: format!("{JOB_TYPE}:{encoded}"),
job_type: JOB_TYPE.to_string(),
summary: format!("Compact {encoded} ({} fragments)", stats.fragments),
detail: format!(
"{} fragments at version {}, above the {min_fragments} the policy allows",
stats.fragments, stats.version
),
parameters,
..Default::default()
});
}
if !observations.is_empty() {
sender.send_observations(WorkerObservations {
job_type: JOB_TYPE.to_string(),
observations,
})?;
}
let total = proposals.len() as i32;
sender.send_proposals(DetectionProposals {
request_id: request.request_id.clone(),
job_type: JOB_TYPE.to_string(),
proposals,
has_more: false,
})?;
sender.send_complete(DetectionComplete {
request_id: request.request_id.clone(),
job_type: JOB_TYPE.to_string(),
success: true,
error_message: String::new(),
total_proposals: total,
})?;
Ok(())
}
async fn execute(
&self,
request: &ExecuteJobRequest,
sender: &dyn ExecutionSender,
) -> Result<()> {
let job = request
.job
.as_ref()
.ok_or_else(|| anyhow!("execute request carried no job"))?;
let id = table_id(&job.parameters)
.ok_or_else(|| anyhow!("job {} carried no table_id", job.job_id))?;
let target_rows = clamp(
int_or(
&request.worker_config_values,
"target_rows_per_fragment",
DEFAULT_TARGET_ROWS,
),
TARGET_ROWS_FLOOR,
TARGET_ROWS_CEILING,
) as usize;
let client = self.client();
// Re-resolve rather than trusting the location detection saw: the table
// may have been repointed, and the vended credentials have expired.
let mut table = dataset::open(&client, &id, &self.fallback).await?;
let before = table.stats().await?;
sender.send_progress(JobProgressUpdate {
request_id: request.request_id.clone(),
job_id: job.job_id.clone(),
job_type: JOB_TYPE.to_string(),
progress_percent: 10.0,
stage: format!("compacting {} fragments", before.fragments),
..Default::default()
})?;
let options = CompactionOptions {
target_rows_per_fragment: target_rows,
..Default::default()
};
let metrics = compact_files(&mut table.dataset, options, None)
.await
.with_context(|| format!("compact {}", table.location))?;
if let Some(counters) = &self.metrics {
counters
.fragments_removed
.inc_by(metrics.fragments_removed as u64);
}
let after = table.stats().await?;
let mut output: HashMap<String, ConfigValue> = HashMap::new();
output.insert(
"fragments_removed".to_string(),
int_value(metrics.fragments_removed as i64),
);
output.insert(
"fragments_added".to_string(),
int_value(metrics.fragments_added as i64),
);
output.insert(
"files_removed".to_string(),
int_value(metrics.files_removed as i64),
);
sender.send_completed(JobCompleted {
request_id: request.request_id.clone(),
job_id: job.job_id.clone(),
job_type: JOB_TYPE.to_string(),
success: true,
result: Some(JobResult {
output_values: output,
summary: format!("{} fragments became {}", before.fragments, after.fragments),
..Default::default()
}),
..Default::default()
})?;
Ok(())
}
}

View File

@@ -0,0 +1,275 @@
use std::collections::HashMap;
use anyhow::{Context, Result, anyhow};
use async_trait::async_trait;
use lance::index::DatasetIndexExt;
use lance_index::optimize::OptimizeOptions;
use seaweed_worker_core::config_form::{form, int_or, int_value, number_field};
use seaweed_worker_core::pb::{
ConfigValue, DetectionComplete, DetectionProposals, ExecuteJobRequest, JobCompleted,
JobProgressUpdate, JobProposal, JobResult, JobTypeCapability, JobTypeDescriptor,
RunDetectionRequest,
};
use seaweed_worker_core::{DetectionSender, ExecutionSender, JobHandler};
use tracing::warn;
use crate::catalog::{NamespaceClient, parse_id};
use crate::dataset::{self, OpenTable};
use crate::jobs::{clamp, string_list, table_id};
pub const JOB_TYPE: &str = "lance_optimize_indices";
const DEFAULT_MAX_UNINDEXED_ROWS: i64 = 100_000;
// Zero is a real setting - reindex as soon as any row is uncovered - so the
// floor is what stays meaningful rather than what the form offers. The ceiling
// is the form's, and the point of both is that this is cast to u64: a negative
// would arrive as an enormous budget and mean "never reindex".
const MAX_UNINDEXED_FLOOR: i64 = 0;
const MAX_UNINDEXED_CEILING: i64 = 100_000_000;
/// Rows written after an index was built are not covered by it, so a vector
/// search quietly misses them. This is the job with no Iceberg equivalent, and
/// the reason a neglected Lance table is a correctness problem rather than a
/// slow one.
pub struct OptimizeIndicesHandler {
namespace_url: String,
fallback: dataset::FallbackOptions,
metrics: Option<crate::metrics::LanceMetrics>,
}
impl OptimizeIndicesHandler {
pub fn new(namespace_url: String) -> Self {
Self {
namespace_url,
fallback: dataset::FallbackOptions::new(),
metrics: None,
}
}
pub fn with_metrics(mut self, metrics: Option<crate::metrics::LanceMetrics>) -> Self {
self.metrics = metrics;
self
}
pub fn with_fallback(mut self, fallback: dataset::FallbackOptions) -> Self {
self.fallback = fallback;
self
}
fn client(&self) -> NamespaceClient {
NamespaceClient::new(self.namespace_url.clone())
}
}
/// Rows no index covers, summed across a table's indices. A table with no
/// indices at all has nothing to optimize, which is different from a table whose
/// indices have fallen behind.
async fn unindexed_rows(table: &OpenTable) -> Result<Option<u64>> {
let indices = table.dataset.load_indices().await?;
if indices.is_empty() {
return Ok(None);
}
let mut worst = 0u64;
let mut names: Vec<String> = indices.iter().map(|index| index.name.clone()).collect();
names.sort();
names.dedup();
for name in names {
let raw = table.dataset.index_statistics(&name).await?;
let stats: serde_json::Value = serde_json::from_str(&raw)
.with_context(|| format!("parse index statistics for {name}"))?;
let unindexed = stats
.get("num_unindexed_rows")
.and_then(|value| value.as_u64())
.unwrap_or(0);
worst = worst.max(unindexed);
}
Ok(Some(worst))
}
#[async_trait]
impl JobHandler for OptimizeIndicesHandler {
fn capability(&self) -> JobTypeCapability {
JobTypeCapability {
job_type: JOB_TYPE.to_string(),
can_detect: true,
can_execute: true,
max_detection_concurrency: 1,
max_execution_concurrency: 1,
display_name: "Lance Index Optimization".to_string(),
description: "Extend indices to cover rows written since they were built".to_string(),
weight: 30,
}
}
fn descriptor(&self) -> JobTypeDescriptor {
let mut defaults: HashMap<String, ConfigValue> = HashMap::new();
defaults.insert(
"max_unindexed_rows".to_string(),
int_value(DEFAULT_MAX_UNINDEXED_ROWS),
);
JobTypeDescriptor {
job_type: JOB_TYPE.to_string(),
display_name: "Lance Index Optimization".to_string(),
description: "Keep vector and scalar indices covering the whole table".to_string(),
icon: "fas fa-magnifying-glass-chart".to_string(),
descriptor_version: 1,
worker_config_form: Some(form(
"lance-indices-worker",
"Index optimization",
vec![number_field(
"max_unindexed_rows",
"Unindexed row budget",
"Reindex once a table has more rows than this outside its indices",
1_000,
100_000_000,
)],
defaults.clone(),
)),
worker_default_values: defaults,
..Default::default()
}
}
async fn detect(
&self,
request: &RunDetectionRequest,
sender: &dyn DetectionSender,
) -> Result<()> {
let budget = clamp(
int_or(
&request.worker_config_values,
"max_unindexed_rows",
DEFAULT_MAX_UNINDEXED_ROWS,
),
MAX_UNINDEXED_FLOOR,
MAX_UNINDEXED_CEILING,
) as u64;
let client = self.client();
let tables = client.list_all_tables().await?;
let mut proposals = Vec::new();
for encoded in &tables {
let id = parse_id(encoded);
let table = match dataset::open(&client, &id, &self.fallback).await {
Ok(table) => table,
Err(err) => {
if let Some(counters) = &self.metrics {
counters.worker.object_skipped(JOB_TYPE, "open");
}
warn!("skipping {encoded}: {err:#}");
continue;
}
};
if let Some(counters) = &self.metrics {
counters.worker.object_seen(JOB_TYPE);
}
let unindexed = match unindexed_rows(&table).await {
Ok(Some(unindexed)) => unindexed,
Ok(None) => continue,
Err(err) => {
if let Some(counters) = &self.metrics {
counters.worker.object_skipped(JOB_TYPE, "index_stats");
}
warn!("skipping {encoded}: reading its index stats failed: {err:#}");
continue;
}
};
if unindexed <= budget {
continue;
}
let mut parameters: HashMap<String, ConfigValue> = HashMap::new();
parameters.insert("table_id".to_string(), string_list(&id));
proposals.push(JobProposal {
proposal_id: format!("{JOB_TYPE}:{encoded}"),
dedupe_key: format!("{JOB_TYPE}:{encoded}"),
job_type: JOB_TYPE.to_string(),
summary: format!("Reindex {encoded} ({unindexed} rows uncovered)"),
detail: format!(
"{unindexed} rows sit outside the indices, above the {budget} the policy allows; \
a search of this table misses them"
),
parameters,
..Default::default()
});
}
let total = proposals.len() as i32;
sender.send_proposals(DetectionProposals {
request_id: request.request_id.clone(),
job_type: JOB_TYPE.to_string(),
proposals,
has_more: false,
})?;
sender.send_complete(DetectionComplete {
request_id: request.request_id.clone(),
job_type: JOB_TYPE.to_string(),
success: true,
error_message: String::new(),
total_proposals: total,
})?;
Ok(())
}
async fn execute(
&self,
request: &ExecuteJobRequest,
sender: &dyn ExecutionSender,
) -> Result<()> {
let job = request
.job
.as_ref()
.ok_or_else(|| anyhow!("execute request carried no job"))?;
let id = table_id(&job.parameters)
.ok_or_else(|| anyhow!("job {} carried no table_id", job.job_id))?;
let client = self.client();
let mut table = dataset::open(&client, &id, &self.fallback).await?;
let before = unindexed_rows(&table).await?.unwrap_or(0);
sender.send_progress(JobProgressUpdate {
request_id: request.request_id.clone(),
job_id: job.job_id.clone(),
job_type: JOB_TYPE.to_string(),
progress_percent: 10.0,
stage: format!("indexing {before} uncovered rows"),
..Default::default()
})?;
// Merging every delta keeps read latency from drifting as the table is
// reindexed again and again; leaving them unmerged is how an index ends
// up fast to write and slow to search.
table
.dataset
.optimize_indices(&OptimizeOptions::default())
.await
.with_context(|| format!("optimize indices of {}", table.location))?;
let after = unindexed_rows(&table).await?.unwrap_or(0);
if let Some(counters) = &self.metrics {
counters.rows_indexed.inc_by(before.saturating_sub(after));
}
let mut output: HashMap<String, ConfigValue> = HashMap::new();
output.insert(
"unindexed_rows_before".to_string(),
int_value(before as i64),
);
output.insert("unindexed_rows_after".to_string(), int_value(after as i64));
sender.send_completed(JobCompleted {
request_id: request.request_id.clone(),
job_id: job.job_id.clone(),
job_type: JOB_TYPE.to_string(),
success: true,
result: Some(JobResult {
output_values: output,
summary: format!("{before} uncovered rows became {after}"),
..Default::default()
}),
..Default::default()
})?;
Ok(())
}
}

View File

@@ -0,0 +1,118 @@
//! One module per job type. Each declares its capability and the settings form
//! admin renders for it, then does the work.
//!
//! Detection opens each table and reads its manifest rather than its data, so a
//! sweep across a catalog stays cheap. Execution re-resolves the table instead
//! of trusting what detection saw: it may have been repointed, and the vended
//! credentials expire.
pub mod cleanup;
pub mod compact;
pub mod indices;
pub mod sort;
use std::collections::HashMap;
use std::sync::Arc;
use seaweed_worker_core::JobHandler;
use seaweed_worker_core::pb::{ConfigValue, ObjectObservation, StringList, config_value::Kind};
use crate::catalog::parse_id;
/// A table identifier travels in a proposal's parameters and comes back on the
/// job, so both sides agree on one encoding.
pub(crate) fn string_list(parts: &[String]) -> ConfigValue {
ConfigValue {
kind: Some(Kind::StringList(StringList {
values: parts.to_vec(),
})),
}
}
pub(crate) fn table_id(parameters: &HashMap<String, ConfigValue>) -> Option<Vec<String>> {
match parameters.get("table_id")?.kind.as_ref()? {
Kind::StringList(list) => Some(list.values.clone()),
Kind::StringValue(encoded) => Some(parse_id(encoded)),
_ => None,
}
}
/// Every handler this worker serves. A worker process may serve several job
/// types, which is why WorkerHello carries a list.
pub fn handlers(
namespace_url: String,
fallback: crate::dataset::FallbackOptions,
metrics: Option<crate::metrics::LanceMetrics>,
) -> Vec<Arc<dyn JobHandler>> {
vec![
Arc::new(
compact::CompactHandler::new(namespace_url.clone())
.with_fallback(fallback.clone())
.with_metrics(metrics.clone()),
),
Arc::new(
indices::OptimizeIndicesHandler::new(namespace_url.clone())
.with_fallback(fallback.clone())
.with_metrics(metrics.clone()),
),
Arc::new(
sort::SortHandler::new(namespace_url.clone())
.with_fallback(fallback.clone())
.with_metrics(metrics.clone()),
),
Arc::new(
cleanup::CleanupVersionsHandler::new(namespace_url)
.with_fallback(fallback)
.with_metrics(metrics),
),
]
}
/// Holds a configured value to the range its form offers. A value from outside
/// it is one the UI could not have produced, and every one of these is cast to
/// an unsigned type: a negative arrives as an enormous number, which silently
/// turns a threshold into "never" rather than failing loudly.
pub(crate) fn clamp(value: i64, low: i64, high: i64) -> i64 {
value.max(low).min(high)
}
/// The format the catalog records for the tables this worker maintains.
pub const FORMAT: &str = "LANCE";
/// Builds the observation a detection sweep reports for one table. Detection
/// has already opened the dataset to decide whether it needs work, so saying
/// what it saw costs nothing, and for a format the cluster cannot read this is
/// the only description of the table anything can produce.
pub(crate) fn observation(
id: &[String],
format: &str,
attributes: HashMap<String, ConfigValue>,
) -> ObjectObservation {
ObjectObservation {
object_id: id.to_vec(),
object_kind: "table".to_string(),
format: format.to_string(),
attributes,
observed_at: Some(std::time::SystemTime::now().into()),
}
}
#[cfg(test)]
mod tests {
use super::*;
// Every configured threshold is cast to an unsigned type before use. A
// negative one would arrive as an enormous number and quietly mean "never",
// which looks exactly like a worker with nothing to do.
#[test]
fn clamp_keeps_a_negative_from_wrapping() {
assert_eq!(clamp(-1, 2, 4096) as usize, 2);
assert_eq!(clamp(i64::MIN, 0, 100_000_000) as u64, 0);
}
#[test]
fn clamp_holds_the_ceiling_and_passes_the_middle() {
assert_eq!(clamp(i64::MAX, 0, 8760), 8760);
assert_eq!(clamp(168, 0, 8760), 168);
}
}

View File

@@ -0,0 +1,488 @@
use std::collections::HashMap;
use std::sync::Arc;
use anyhow::{Context, Result, anyhow};
use async_trait::async_trait;
use lance::dataset::scanner::ColumnOrdering;
use lance::dataset::transaction::Operation;
use lance::dataset::write::{CommitBuilder, InsertBuilder};
use lance::dataset::{WriteDestination, WriteMode, WriteParams};
use lance::index::DatasetIndexExt;
use lance_datafusion::exec::{LanceExecutionOptions, execute_plan};
use seaweed_worker_core::config_form::{form, int_or, int_value, string_or, string_value};
use seaweed_worker_core::pb::{
ConfigValue, DetectionComplete, DetectionProposals, ExecuteJobRequest, JobCompleted,
JobProgressUpdate, JobProposal, JobResult, JobTypeCapability, JobTypeDescriptor,
RunDetectionRequest, WorkerObservations,
};
use seaweed_worker_core::{DetectionSender, ExecutionSender, JobHandler};
use seaweed_worker_sort::{
CONFIG_MAX_ROWS_PER_FILE, CONFIG_MEMORY_BUDGET_MB, CONFIG_MIN_UNSORTED_ROWS,
CONFIG_SORT_FIELDS, DECLARED_FIELDS_KEY, FragmentSummary, SortSpec, SortState, config_fields,
resolve, verdict,
};
use tracing::warn;
use crate::catalog::{NamespaceClient, parse_id};
use crate::dataset;
use crate::jobs::{FORMAT, clamp, observation, string_list, table_id};
pub const JOB_TYPE: &str = "lance_sort";
const DEFAULT_MIN_UNSORTED_ROWS: i64 = 1_048_576;
const DEFAULT_MEMORY_BUDGET_MB: i64 = 512;
const DEFAULT_MAX_ROWS_PER_FILE: i64 = 1_048_576;
// The ranges the shared form offers, repeated here because a value from outside
// one is a value the UI could not have produced and every one of these is cast
// to an unsigned type on the way in.
const MIN_UNSORTED_ROWS_FLOOR: i64 = 1;
const MIN_UNSORTED_ROWS_CEILING: i64 = 1_000_000_000;
const MEMORY_BUDGET_FLOOR: i64 = 64;
const MEMORY_BUDGET_CEILING: i64 = 1_048_576;
const MAX_ROWS_PER_FILE_FLOOR: i64 = 1024;
const MAX_ROWS_PER_FILE_CEILING: i64 = 16_777_216;
const BYTES_PER_MB: u64 = 1024 * 1024;
/// Rewrites a table so its rows are stored in the order its fields declare,
/// which is what lets a range scan read a few files instead of all of them.
///
/// Lance appends fragments in write order and has no notion of a sorted table,
/// so nothing but a rewrite establishes that order, and nothing but another
/// rewrite restores it once rows have been appended.
pub struct SortHandler {
namespace_url: String,
fallback: dataset::FallbackOptions,
metrics: Option<crate::metrics::LanceMetrics>,
}
impl SortHandler {
pub fn new(namespace_url: String) -> Self {
Self {
namespace_url,
fallback: dataset::FallbackOptions::new(),
metrics: None,
}
}
pub fn with_metrics(mut self, metrics: Option<crate::metrics::LanceMetrics>) -> Self {
self.metrics = metrics;
self
}
/// Storage options to use where the namespace vends none.
pub fn with_fallback(mut self, fallback: dataset::FallbackOptions) -> Self {
self.fallback = fallback;
self
}
fn client(&self) -> NamespaceClient {
NamespaceClient::new(self.namespace_url.clone())
}
}
/// The fragments a dataset holds, in fragment order, as the sort crate reads
/// them: the data files each one names, and the live rows it holds.
///
/// Appends only add fragments after the existing ones and lance hands out
/// fragment ids monotonically even across an overwrite, so ordering by id makes
/// "the fragments the sort wrote are still the first ones" a prefix test.
fn fragment_summaries(dataset: &lance::Dataset) -> Vec<FragmentSummary> {
let mut fragments = dataset.get_fragments();
fragments.sort_by_key(|fragment| fragment.id());
fragments
.iter()
.map(|fragment| summarize(fragment.metadata()))
.collect()
}
/// The manifest's own account of one fragment. `num_rows` is physical rows less
/// deletions, and None where the manifest does not record them.
fn summarize(fragment: &lance_table::format::Fragment) -> FragmentSummary {
FragmentSummary {
files: fragment
.files
.iter()
.map(|file| file.path.clone())
.collect(),
rows: fragment.num_rows().map(|rows| rows as u64),
}
}
/// The scan ordering for a spec. Lance validates the column names against the
/// schema when the ordering is set, so a spec naming a column the table does
/// not have fails before anything is written.
fn orderings(spec: &SortSpec) -> Vec<ColumnOrdering> {
spec.fields
.iter()
.map(|field| ColumnOrdering {
ascending: !field.descending,
nulls_first: field.nulls_first,
column_name: field.path.clone(),
})
.collect()
}
#[async_trait]
impl JobHandler for SortHandler {
fn capability(&self) -> JobTypeCapability {
JobTypeCapability {
job_type: JOB_TYPE.to_string(),
can_detect: true,
can_execute: true,
max_detection_concurrency: 1,
max_execution_concurrency: 1,
display_name: "Lance Sort".to_string(),
description: "Rewrite Lance tables in the order their fields declare. The rewrite replaces every fragment, so indices do not survive it and want rebuilding afterwards.".to_string(),
weight: 40,
}
}
fn descriptor(&self) -> JobTypeDescriptor {
let mut defaults: HashMap<String, ConfigValue> = HashMap::new();
defaults.insert(CONFIG_SORT_FIELDS.to_string(), string_value(""));
defaults.insert(
CONFIG_MIN_UNSORTED_ROWS.to_string(),
int_value(DEFAULT_MIN_UNSORTED_ROWS),
);
defaults.insert(
CONFIG_MEMORY_BUDGET_MB.to_string(),
int_value(DEFAULT_MEMORY_BUDGET_MB),
);
defaults.insert(
CONFIG_MAX_ROWS_PER_FILE.to_string(),
int_value(DEFAULT_MAX_ROWS_PER_FILE),
);
JobTypeDescriptor {
job_type: JOB_TYPE.to_string(),
display_name: "Lance Sort".to_string(),
description: "Sort Lance tables by their declared fields".to_string(),
icon: "fas fa-sort".to_string(),
descriptor_version: 1,
worker_config_form: Some(form(
"lance-sort-worker",
"Sort",
config_fields(),
defaults.clone(),
)),
worker_default_values: defaults,
..Default::default()
}
}
/// Propose a job for every table whose declared order the data no longer
/// follows. Opening a dataset reads its manifest and its configuration, not
/// its data, so this stays cheap across a catalog.
async fn detect(
&self,
request: &RunDetectionRequest,
sender: &dyn DetectionSender,
) -> Result<()> {
let configured = string_or(&request.worker_config_values, CONFIG_SORT_FIELDS, "");
let min_unsorted_rows = clamp(
int_or(
&request.worker_config_values,
CONFIG_MIN_UNSORTED_ROWS,
DEFAULT_MIN_UNSORTED_ROWS,
),
MIN_UNSORTED_ROWS_FLOOR,
MIN_UNSORTED_ROWS_CEILING,
) as u64;
let client = self.client();
let tables = client.list_all_tables().await?;
let mut proposals = Vec::new();
let mut observations = Vec::new();
for encoded in &tables {
let id = parse_id(encoded);
let table = match dataset::open(&client, &id, &self.fallback).await {
Ok(table) => table,
Err(err) => {
// One unreadable table must not end the sweep: the tables
// already read would lose their proposals.
if let Some(counters) = &self.metrics {
counters.worker.object_skipped(JOB_TYPE, "open");
}
warn!("skipping {encoded}: {err:#}");
continue;
}
};
if let Some(counters) = &self.metrics {
counters.worker.object_seen(JOB_TYPE);
}
let stats = match table.stats().await {
Ok(stats) => stats,
Err(err) => {
if let Some(counters) = &self.metrics {
counters.worker.object_skipped(JOB_TYPE, "stats");
}
warn!("skipping {encoded}: reading its stats failed: {err:#}");
continue;
}
};
let config = table.dataset.config().clone();
let spec = match resolve(
config.get(DECLARED_FIELDS_KEY).map(String::as_str),
&configured,
) {
Ok(Some(spec)) => spec,
Ok(None) => {
// Not every table is one an operator wants sorted, and a
// table with no order named anywhere is not a failure.
if let Some(counters) = &self.metrics {
counters.worker.object_skipped(JOB_TYPE, "no_sort_order");
}
continue;
}
Err(err) => {
if let Some(counters) = &self.metrics {
counters.worker.object_skipped(JOB_TYPE, "sort_order");
}
warn!("skipping {encoded}: {err:#}");
continue;
}
};
let state = SortState::from_config(&config);
let decision = verdict(
&spec,
&state,
&fragment_summaries(&table.dataset),
min_unsorted_rows,
);
let mut attributes: HashMap<String, ConfigValue> = HashMap::new();
attributes.insert("rows".to_string(), int_value(stats.rows as i64));
attributes.insert("version".to_string(), int_value(stats.version as i64));
attributes.insert("fragments".to_string(), int_value(stats.fragments as i64));
attributes.insert("sort_fields".to_string(), string_value(spec.to_string()));
attributes.insert("sort_state".to_string(), string_value(decision.reason()));
if let Some(schema) = stats.schema.clone() {
attributes.insert("schema".to_string(), string_value(schema));
}
observations.push(observation(&id, FORMAT, attributes));
// Logged because "detection proposed nothing" is otherwise
// indistinguishable from a table the worker could not read.
tracing::info!(
"sort detection: {encoded} has {} rows in {spec} order: {}",
stats.rows,
decision.reason()
);
if !decision.needs_sort() {
continue;
}
let mut parameters: HashMap<String, ConfigValue> = HashMap::new();
parameters.insert("table_id".to_string(), string_list(&id));
proposals.push(JobProposal {
proposal_id: format!("{JOB_TYPE}:{encoded}"),
dedupe_key: format!("{JOB_TYPE}:{encoded}"),
job_type: JOB_TYPE.to_string(),
summary: format!("Sort {encoded} by {spec}"),
detail: format!("{} rows, {}", stats.rows, decision.reason()),
parameters,
..Default::default()
});
}
if !observations.is_empty() {
sender.send_observations(WorkerObservations {
job_type: JOB_TYPE.to_string(),
observations,
})?;
}
let total = proposals.len() as i32;
sender.send_proposals(DetectionProposals {
request_id: request.request_id.clone(),
job_type: JOB_TYPE.to_string(),
proposals,
has_more: false,
})?;
sender.send_complete(DetectionComplete {
request_id: request.request_id.clone(),
job_type: JOB_TYPE.to_string(),
success: true,
error_message: String::new(),
total_proposals: total,
})?;
Ok(())
}
async fn execute(
&self,
request: &ExecuteJobRequest,
sender: &dyn ExecutionSender,
) -> Result<()> {
let job = request
.job
.as_ref()
.ok_or_else(|| anyhow!("execute request carried no job"))?;
let id = table_id(&job.parameters)
.ok_or_else(|| anyhow!("job {} carried no table_id", job.job_id))?;
let configured = string_or(&request.worker_config_values, CONFIG_SORT_FIELDS, "");
let memory_budget_mb = clamp(
int_or(
&request.worker_config_values,
CONFIG_MEMORY_BUDGET_MB,
DEFAULT_MEMORY_BUDGET_MB,
),
MEMORY_BUDGET_FLOOR,
MEMORY_BUDGET_CEILING,
) as u64;
let max_rows_per_file = clamp(
int_or(
&request.worker_config_values,
CONFIG_MAX_ROWS_PER_FILE,
DEFAULT_MAX_ROWS_PER_FILE,
),
MAX_ROWS_PER_FILE_FLOOR,
MAX_ROWS_PER_FILE_CEILING,
) as usize;
let client = self.client();
// Re-resolve rather than trusting what detection saw: the table may have
// been repointed, its declared order changed, and the vended credentials
// have expired.
let table = dataset::open(&client, &id, &self.fallback).await?;
let before = table.stats().await?;
let config = table.dataset.config().clone();
let spec = resolve(
config.get(DECLARED_FIELDS_KEY).map(String::as_str),
&configured,
)?
.ok_or_else(|| anyhow!("neither the table nor this worker names an order to sort by"))?;
sender.send_progress(JobProgressUpdate {
request_id: request.request_id.clone(),
job_id: job.job_id.clone(),
job_type: JOB_TYPE.to_string(),
progress_percent: 10.0,
stage: format!("sorting {} rows by {spec}", before.rows),
..Default::default()
})?;
let mut scanner = table.dataset.scan();
scanner
.order_by(Some(orderings(&spec)))
.with_context(|| format!("order the scan of {} by {spec}", table.location))?;
let plan = scanner
.create_plan()
.await
.context("plan the ordered scan")?;
// Spilling is off in LanceExecutionOptions::default(), and the scanner's
// own stream helper takes that default, so the options are built here
// instead: a sort that cannot spill is one bounded by memory, which is
// the whole thing this job exists to avoid.
let stream = execute_plan(
plan,
LanceExecutionOptions {
use_spilling: true,
mem_pool_size: Some(memory_budget_mb * BYTES_PER_MB),
..Default::default()
},
)
.context("run the ordered scan")?;
let params = WriteParams {
mode: WriteMode::Overwrite,
max_rows_per_file,
..Default::default()
};
let destination = Arc::new(table.dataset.clone());
let mut transaction = InsertBuilder::new(WriteDestination::Dataset(destination.clone()))
.with_params(&params)
.execute_uncommitted_stream(stream)
.await
.context("write the sorted fragments")?;
// The marker rides in the same commit as the data it describes:
// Operation::Overwrite is the one place lance takes configuration
// values alongside fragments, and a marker written as a second commit
// would be lost if the worker died in between, re-sorting the whole
// table on the next sweep.
match &mut transaction.operation {
Operation::Overwrite {
config_upsert_values,
fragments,
..
} => {
// The files these fragments name were written a moment ago and
// keep their names whatever version this commit becomes, which
// is what lets the marker travel inside the commit it describes.
let written: Vec<FragmentSummary> = fragments.iter().map(summarize).collect();
*config_upsert_values =
Some(SortState::record(&spec, &written, before.rows as u64));
}
other => {
return Err(anyhow!(
"a sorted rewrite produced {other} instead of an overwrite"
));
}
}
let sorted = CommitBuilder::new(WriteDestination::Dataset(destination))
.execute(transaction)
.await
.context("commit the sorted rewrite")?;
if let Some(counters) = &self.metrics {
counters.rows_sorted.inc_by(before.rows as u64);
}
// An overwrite carries fragments, schema and configuration — not index
// metadata — so a table that had indices no longer does. Saying so is
// the difference between a slow table and a mystery.
let dropped_indices = match table.dataset.load_indices().await {
Ok(indices) => indices.len(),
Err(err) => {
warn!("could not read the indices of {}: {err:#}", table.location);
0
}
};
if dropped_indices > 0 {
warn!(
"{} lost {dropped_indices} index/indices to the sorted rewrite; run the index job to rebuild them",
table.location
);
}
let fragments_after = sorted.get_fragments().len();
let mut output: HashMap<String, ConfigValue> = HashMap::new();
output.insert("rows_sorted".to_string(), int_value(before.rows as i64));
output.insert("sort_fields".to_string(), string_value(spec.to_string()));
output.insert(
"fragments_before".to_string(),
int_value(before.fragments as i64),
);
output.insert(
"fragments_after".to_string(),
int_value(fragments_after as i64),
);
output.insert(
"indices_dropped".to_string(),
int_value(dropped_indices as i64),
);
sender.send_completed(JobCompleted {
request_id: request.request_id.clone(),
job_id: job.job_id.clone(),
job_type: JOB_TYPE.to_string(),
success: true,
result: Some(JobResult {
output_values: output,
summary: format!("{} rows rewritten in {spec} order", before.rows),
..Default::default()
}),
..Default::default()
})?;
Ok(())
}
}

View File

@@ -0,0 +1,15 @@
//! Maintenance jobs for Lance tables.
//!
//! A Lance dataset needs three things done to it over time: its fragments
//! compacted, its indices extended to cover rows written after they were built,
//! and its old versions removed. None can run in the Go worker, because all
//! three read and rewrite Lance files. This crate is the worker that can.
pub mod catalog;
pub mod dataset;
pub mod jobs;
pub mod metrics;
pub mod preview;
pub mod tls;
pub use jobs::handlers;

View File

@@ -0,0 +1,193 @@
use std::time::Duration;
use anyhow::{Context, Result};
use clap::Parser;
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use seaweed_worker_core::{Metrics, Registry, TlsOptions, WorkerOptions};
use weed_lance_worker::handlers;
use weed_lance_worker::metrics::LanceMetrics;
/// Mirrors `weed worker`'s flags, because this is the same contract from another
/// language and an operator should not have to learn a second set of names.
#[derive(Parser, Debug)]
#[command(name = "weed-worker", about = "SeaweedFS maintenance worker")]
struct Args {
/// Admin server gRPC address.
#[arg(long, default_value = "localhost:23646", env = "WEED_ADMIN")]
admin: String,
/// Worker identity reported to admin.
#[arg(long, default_value = "lance-worker", env = "WEED_WORKER_ID")]
id: String,
/// Lance namespace the worker lists tables from.
#[arg(
long,
default_value = "http://localhost:9101",
env = "WEED_LANCE_NAMESPACE"
)]
namespace: String,
#[arg(long, default_value = "10", env = "WEED_HEARTBEAT_SECONDS")]
heartbeat_seconds: u64,
#[arg(long, default_value = "1")]
max_concurrency: i32,
/// Storage credentials to use where the namespace vends none. A gateway
/// with STS configured vends its own and these are ignored.
#[arg(long, env = "WEED_S3_ACCESS_KEY")]
access_key: Option<String>,
#[arg(long, env = "WEED_S3_SECRET_KEY")]
secret_key: Option<String>,
/// mTLS for the admin stream, the same certificates the Go worker reads
/// from the [grpc.worker] section of security.toml. All three together, or
/// none, in which case the stream is plaintext.
#[arg(long, env = "WEED_GRPC_CA")]
tls_ca: Option<String>,
#[arg(long, env = "WEED_GRPC_CLIENT_CERT")]
tls_cert: Option<String>,
#[arg(long, env = "WEED_GRPC_CLIENT_KEY")]
tls_key: Option<String>,
/// Name to verify admin's certificate against, when it is not the address
/// this worker dials.
#[arg(long)]
tls_server_name: Option<String>,
/// Serve /health, /ready and /metrics on this port, the way
/// `weed worker -metricsPort` does. Zero, the default, serves nothing.
#[arg(long, default_value = "0", env = "WEED_METRICS_PORT")]
metrics_port: u16,
/// Address the metrics server binds. Loopback by default, since the
/// endpoint is unauthenticated.
#[arg(long, default_value = "127.0.0.1", env = "WEED_METRICS_IP")]
metrics_ip: String,
}
impl Args {
/// The TLS configuration, or an error when the three certificate paths do
/// not arrive together: a CA on its own would silently give one-way TLS,
/// which the cluster's mutual setup refuses anyway.
fn tls(&self) -> Result<Option<TlsOptions>> {
match (
self.tls_ca.clone(),
self.tls_cert.clone(),
self.tls_key.clone(),
) {
(None, None, None) => Ok(None),
(Some(ca_path), Some(client_cert_path), Some(client_key_path)) => {
Ok(Some(TlsOptions {
ca_path,
client_cert_path,
client_key_path,
server_name: self.tls_server_name.clone(),
}))
}
_ => Err(anyhow::anyhow!(
"--tls-ca, --tls-cert and --tls-key must be given together"
)),
}
}
}
/// 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<SocketAddr> {
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<()> {
weed_lance_worker::tls::install_default_crypto_provider();
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
)
.init();
let args = Args::parse();
let tls = args.tls()?;
let options = WorkerOptions {
admin_address: args.admin,
worker_id: args.id,
heartbeat_interval: Duration::from_secs(args.heartbeat_seconds),
max_detection_concurrency: args.max_concurrency,
max_execution_concurrency: args.max_concurrency,
tls,
..Default::default()
};
let mut fallback = weed_lance_worker::dataset::FallbackOptions::new();
if let (Some(access), Some(secret)) = (args.access_key, args.secret_key) {
fallback.insert("aws_access_key_id".to_string(), access);
fallback.insert("aws_secret_access_key".to_string(), secret);
}
let metrics = Metrics::new(&options.worker_id, &options.worker_version)?;
let lance_metrics = LanceMetrics::new(&metrics)?;
if args.metrics_port > 0 {
seaweed_worker_core::metrics::spawn(
metrics.clone(),
metrics_address(&args.metrics_ip, args.metrics_port)?,
);
}
let mut registry = Registry::new().with_preview(Arc::new(
weed_lance_worker::preview::LancePreview::new(args.namespace.clone(), fallback.clone()),
));
for handler in handlers(args.namespace, fallback, Some(lance_metrics)) {
registry = registry.register(handler);
}
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());
}
}

View File

@@ -0,0 +1,50 @@
//! What the Lance jobs reclaimed, in numbers a scrape can add up.
//!
//! The generic worker metrics say a job ran and how long it took. These say what
//! it did: a compaction sweep that runs every minute and removes nothing is a
//! different thing from one that never runs, and only these tell them apart.
use anyhow::Result;
use prometheus::IntCounter;
use seaweed_worker_core::Metrics;
/// Counters the Lance jobs raise, registered on the worker's own registry so
/// they are served from the same endpoint.
#[derive(Clone)]
pub struct LanceMetrics {
/// The worker's own metrics, so a job can also record what it read and what
/// it had to skip - the pair that tells "nothing needed doing" apart from
/// "nothing could be read".
pub worker: Metrics,
pub fragments_removed: IntCounter,
pub rows_indexed: IntCounter,
pub rows_sorted: IntCounter,
pub versions_removed: IntCounter,
pub bytes_reclaimed: IntCounter,
}
impl LanceMetrics {
pub fn new(metrics: &Metrics) -> Result<Self> {
Ok(Self {
worker: metrics.clone(),
fragments_removed: metrics.counter(
"lance_fragments_removed_total",
"Fragments merged away by compaction.",
)?,
rows_indexed: metrics.counter(
"lance_rows_indexed_total",
"Rows brought under an index that did not cover them.",
)?,
rows_sorted: metrics.counter(
"lance_rows_sorted_total",
"Rows rewritten into their table's declared order.",
)?,
versions_removed: metrics.counter(
"lance_versions_removed_total",
"Dataset versions removed by cleanup.",
)?,
bytes_reclaimed: metrics
.counter("lance_bytes_reclaimed_total", "Bytes freed by cleanup.")?,
})
}
}

View File

@@ -0,0 +1,86 @@
//! Sample rows of a Lance table, for a page that cannot read the format.
//!
//! Admin renders an Iceberg table by reading its Parquet files directly. There
//! is no Go Lance reader, so for Lance the worker does the reading and hands
//! back text.
use anyhow::{Context, Result};
use arrow_cast::display::{ArrayFormatter, FormatOptions};
use async_trait::async_trait;
use futures::TryStreamExt;
use seaweed_worker_core::{Preview, PreviewProvider};
use crate::catalog::NamespaceClient;
use crate::dataset::{self, FallbackOptions};
use crate::jobs::FORMAT;
pub struct LancePreview {
namespace_url: String,
fallback: FallbackOptions,
}
impl LancePreview {
pub fn new(namespace_url: String, fallback: FallbackOptions) -> Self {
Self {
namespace_url,
fallback,
}
}
}
#[async_trait]
impl PreviewProvider for LancePreview {
fn format(&self) -> &str {
FORMAT
}
async fn preview(&self, object_id: &[String], row_limit: usize) -> Result<Preview> {
let client = NamespaceClient::new(self.namespace_url.clone());
let table = dataset::open(&client, object_id, &self.fallback).await?;
let total_rows = table.dataset.count_rows(None).await.unwrap_or(0) as i64;
let mut scanner = table.dataset.scan();
scanner.limit(Some(row_limit as i64), None)?;
let batches: Vec<_> = scanner
.try_into_stream()
.await?
.try_collect()
.await
.with_context(|| format!("read rows from {}", table.location))?;
let columns = table
.dataset
.schema()
.fields
.iter()
.map(|field| field.name.clone())
.collect();
let mut rows = Vec::new();
for batch in &batches {
// ArrayFormatter renders each Arrow type the way its own tooling
// does, so a vector column reads as a vector rather than as bytes.
let formatters = batch
.columns()
.iter()
.map(|array| ArrayFormatter::try_new(array.as_ref(), &FormatOptions::default()))
.collect::<Result<Vec<_>, _>>()?;
for index in 0..batch.num_rows() {
rows.push(
formatters
.iter()
.map(|formatter| formatter.value(index).to_string())
.collect(),
);
}
}
rows.truncate(row_limit);
Ok(Preview {
columns,
rows,
total_rows,
})
}
}

View File

@@ -0,0 +1,35 @@
use rustls::crypto::aws_lc_rs;
// aws-lc-rs and ring both get linked transitively (lance's aws backend pulls
// aws-lc-rs, reqwest's rustls-tls pulls ring), so rustls can't auto-select a
// provider and tonic's client TLS panics on first use. Pin the default to
// aws-lc-rs, matching the Rust volume server. Idempotent.
pub fn install_default_crypto_provider() {
let _ = aws_lc_rs::default_provider().install_default();
}
#[cfg(test)]
mod tests {
use tonic::transport::{Certificate, Channel, ClientTlsConfig, Identity};
use super::install_default_crypto_provider;
const TEST_CERT_PEM: &str = "-----BEGIN CERTIFICATE-----\nMIIBPDCB76ADAgECAhRuRPQgeAu43BT/M7EfAWSdapVdYDAFBgMrZXAwFDESMBAG\nA1UEAwwJbG9jYWxob3N0MB4XDTI2MDcwNTE2MTUyOVoXDTM2MDcwMjE2MTUyOVow\nFDESMBAGA1UEAwwJbG9jYWxob3N0MCowBQYDK2VwAyEAr/3bNIFI+8V32oCiY6y+\nXRFmZpdNQ2g//VtRkT+nQg+jUzBRMB0GA1UdDgQWBBTsy9tLf1zPiXCQfgci6zNi\ndEzRSjAfBgNVHSMEGDAWgBTsy9tLf1zPiXCQfgci6zNidEzRSjAPBgNVHRMBAf8E\nBTADAQH/MAUGAytlcANBAIvsdw0IbvOBBkb9cd7BfMJfIP9pQQrAL03pCRWJFnFh\nSysaLVgFXI4T078IiaM874oO+iB+5vNbWEpc7CkGow4=\n-----END CERTIFICATE-----\n";
const TEST_KEY_PEM: &str = "-----BEGIN PRIVATE KEY-----\nMC4CAQAwBQYDK2VwBCIEIHbyn71Kk+Y7KT3sBctit7uZpErpoH6qDbFj6P8qGaZH\n-----END PRIVATE KEY-----\n";
// Without install_default_crypto_provider this panics in the lance crate,
// where both aws-lc-rs and ring are linked. Mirrors the volume server's
// test_build_grpc_endpoint_with_tls_resolves_crypto_provider.
#[tokio::test]
async fn tls_channel_builds_after_crypto_provider_installed() {
install_default_crypto_provider();
let config = ClientTlsConfig::new()
.ca_certificate(Certificate::from_pem(TEST_CERT_PEM))
.identity(Identity::from_pem(TEST_CERT_PEM, TEST_KEY_PEM));
let endpoint = Channel::from_shared("https://127.0.0.1:9")
.unwrap()
.tls_config(config)
.unwrap();
assert_eq!(endpoint.uri().scheme_str(), Some("https"));
}
}

View File

@@ -0,0 +1,97 @@
//! What every integration test in this crate needs to drive a handler: a sender
//! that keeps what the handler sent, and the two settings that make a handler
//! usable against a test gateway.
//!
//! Each test binary compiles this module on its own and uses part of it.
#![allow(dead_code)]
use std::collections::HashMap;
use std::sync::Mutex;
use anyhow::Result;
use seaweed_worker_core::pb::{
ActivityEvent, ConfigValue, DetectionComplete, DetectionProposals, JobCompleted,
JobProgressUpdate, JobProposal, ObjectObservation, WorkerObservations, config_value::Kind,
};
use seaweed_worker_core::{DetectionSender, ExecutionSender};
/// Keeps everything a handler reports, so a test can assert on it afterwards.
#[derive(Default)]
pub struct Recorder {
pub proposals: Mutex<Vec<JobProposal>>,
pub observations: Mutex<Vec<ObjectObservation>>,
pub completed: Mutex<Vec<JobCompleted>>,
}
impl DetectionSender for Recorder {
fn send_proposals(&self, proposals: DetectionProposals) -> Result<()> {
self.proposals.lock().unwrap().extend(proposals.proposals);
Ok(())
}
fn send_complete(&self, _complete: DetectionComplete) -> Result<()> {
Ok(())
}
fn send_activity(&self, _activity: ActivityEvent) -> Result<()> {
Ok(())
}
fn send_observations(&self, observations: WorkerObservations) -> Result<()> {
self.observations
.lock()
.unwrap()
.extend(observations.observations);
Ok(())
}
}
impl ExecutionSender for Recorder {
fn send_progress(&self, _progress: JobProgressUpdate) -> Result<()> {
Ok(())
}
fn send_completed(&self, completed: JobCompleted) -> Result<()> {
self.completed.lock().unwrap().push(completed);
Ok(())
}
}
/// The gateway to run against, or nothing, in which case the caller skips:
/// these handlers rewrite real files, and there is nothing to learn from them
/// against a fake.
pub fn namespace_url() -> Option<String> {
std::env::var("WEED_LANCE_NAMESPACE")
.ok()
.filter(|s| !s.is_empty())
}
pub fn int_config(name: &str, value: i64) -> HashMap<String, ConfigValue> {
let mut values = HashMap::new();
values.insert(
name.to_string(),
ConfigValue {
kind: Some(Kind::Int64Value(value)),
},
);
values
}
/// What to reach the data with where the namespace vends nothing of its own -
/// a gateway without STS vends no credentials, and one bound to a wildcard
/// address vends no endpoint either. Taken from the environment so a harness
/// can point these tests at a gateway that does check what it is given.
pub fn fallback() -> weed_lance_worker::dataset::FallbackOptions {
let mut options = weed_lance_worker::dataset::FallbackOptions::new();
let from_env =
|name: &str, default: &str| std::env::var(name).unwrap_or_else(|_| default.to_string());
options.insert(
"aws_access_key_id".to_string(),
from_env("AWS_ACCESS_KEY_ID", "any"),
);
options.insert(
"aws_secret_access_key".to_string(),
from_env("AWS_SECRET_ACCESS_KEY", "any"),
);
if let Ok(endpoint) = std::env::var("AWS_ENDPOINT_URL") {
options.insert("aws_endpoint".to_string(), endpoint);
options.insert("allow_http".to_string(), "true".to_string());
}
options
}

View File

@@ -0,0 +1,540 @@
//! Drives the compaction handler against a live namespace.
//!
//! Skipped unless WEED_LANCE_NAMESPACE names one, the way the Go integration
//! tests skip without Docker: compaction rewrites real files, and there is
//! nothing to learn from it against a fake.
use std::collections::HashMap;
use anyhow::Result;
use seaweed_worker_core::pb::{
ExecuteJobRequest, JobSpec, RunDetectionRequest, config_value::Kind,
};
use seaweed_worker_core::{JobHandler, PreviewProvider};
use weed_lance_worker::catalog::NamespaceClient;
use weed_lance_worker::jobs::cleanup::CleanupVersionsHandler;
use weed_lance_worker::jobs::compact::{CompactHandler, JOB_TYPE};
use weed_lance_worker::jobs::indices::OptimizeIndicesHandler;
use weed_lance_worker::preview::LancePreview;
mod common;
use common::{Recorder, fallback, int_config, namespace_url};
/// These tests drive one live gateway and one shared catalog: `list_all_tables`
/// sweeps everything, so a table another test is writing shows up in this test's
/// detection. Rust runs a binary's tests concurrently, so take a lock.
static GATEWAY: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
/// Declares a table through the namespace and writes `fragments` one-row
/// appends into it, so a test brings its own state instead of depending on
/// whatever a previous run left behind.
async fn seed_fragmented_table(url: &str, name: &str, fragments: usize) -> Result<String> {
seed_table(url, name, fragments, 1, false).await
}
/// Writes `batches` appends of `rows_each` into a freshly declared table, and
/// optionally builds a vector index after the first batch so the later ones are
/// rows no index covers.
async fn seed_table(
url: &str,
name: &str,
batches: usize,
rows_each: usize,
with_index: bool,
) -> Result<String> {
use arrow_array::{
FixedSizeListArray, Float32Array, Int64Array, RecordBatch, RecordBatchIterator,
};
use arrow_schema::{DataType, Field, Schema};
use lance::dataset::{Dataset, WriteMode, WriteParams};
use lance::io::{ObjectStoreParams, StorageOptionsAccessor};
use std::sync::Arc;
// Declaring is the namespace's job, not the worker's, so the test asks for
// it directly rather than widening the client the worker uses. The bucket and
// namespace come first: a table cannot be declared under a parent that does
// not exist, and a test that assumes one is a test that only passes twice.
let http = reqwest::Client::new();
for parent in ["vec", "vec$ml"] {
http.post(format!("{url}/v1/namespace/{parent}/create"))
.json(&serde_json::json!({"mode": "EXIST_OK"}))
.send()
.await?
.error_for_status()?;
}
let encoded = format!("vec$ml${name}");
http.post(format!("{url}/v1/table/{encoded}/declare"))
.json(&serde_json::json!({}))
.send()
.await?
.error_for_status()?;
let client = NamespaceClient::new(url.to_string());
let id = vec!["vec".to_string(), "ml".to_string(), name.to_string()];
let description = client.describe_table(&id).await?;
let mut options = description.storage_options.clone();
options.extend(fallback());
const DIM: i32 = 16;
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new(
"vec",
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), DIM),
false,
),
]));
for i in 0..batches {
let ids: Vec<i64> = (0..rows_each).map(|r| (i * rows_each + r) as i64).collect();
let values: Vec<f32> = ids
.iter()
.flat_map(|id| (0..DIM).map(move |d| (*id as f32) + d as f32))
.collect();
let vectors = FixedSizeListArray::new(
Arc::new(Field::new("item", DataType::Float32, true)),
DIM,
Arc::new(Float32Array::from(values)),
None,
);
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int64Array::from(ids)), Arc::new(vectors)],
)?;
let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
let params = WriteParams {
mode: if i == 0 {
WriteMode::Overwrite
} else {
WriteMode::Append
},
store_params: Some(ObjectStoreParams {
storage_options_accessor: Some(std::sync::Arc::new(
StorageOptionsAccessor::with_static_options(options.clone()),
)),
..Default::default()
}),
..Default::default()
};
let dataset = Dataset::write(reader, description.location.as_str(), Some(params)).await?;
// The index is built after the first batch, so everything appended
// afterwards is a row it does not cover.
if with_index && i == 0 {
use lance::index::DatasetIndexExt;
use lance::index::vector::VectorIndexParams;
use lance_index::IndexType;
use lance_index::vector::{ivf::IvfBuildParams, pq::PQBuildParams};
let mut dataset = dataset;
let params = VectorIndexParams::with_ivf_pq_params(
lance_linalg::distance::MetricType::L2,
IvfBuildParams::new(1),
PQBuildParams::new(4, 8),
);
dataset
.create_index(&["vec"], IndexType::Vector, None, &params, true)
.await?;
}
}
Ok(encoded)
}
/// A table with more fragments than the policy allows is proposed, and running
/// the proposal leaves it with fewer than it started with.
#[tokio::test]
async fn compacts_a_fragmented_table() {
let _gateway = GATEWAY.lock().await;
let Some(url) = namespace_url() else {
eprintln!("WEED_LANCE_NAMESPACE is unset, skipping");
return;
};
// Seeded here rather than by a script, so the test is repeatable: a previous
// run compacts the table it depended on.
let encoded = seed_fragmented_table(&url, "compactme", 12)
.await
.expect("seed a fragmented table");
let handler = CompactHandler::new(url).with_fallback(fallback());
let recorder = Recorder::default();
let request = RunDetectionRequest {
request_id: "detect-1".to_string(),
job_type: JOB_TYPE.to_string(),
worker_config_values: int_config("min_fragments", 4),
..Default::default()
};
handler
.detect(&request, &recorder)
.await
.expect("detection failed");
let proposals = recorder.proposals.lock().unwrap().clone();
assert!(
!proposals.is_empty(),
"expected a proposal for the fragmented table"
);
let proposal = proposals
.iter()
.find(|p| p.summary.contains(encoded.as_str()))
.expect("no proposal for the seeded table");
// Detection opened the dataset to decide, so it reports what it saw. This is
// the only description of a Lance table anything outside the format can give.
let observations = recorder.observations.lock().unwrap().clone();
let observed = observations
.iter()
.find(|o| o.object_id.last().map(String::as_str) == Some("compactme"))
.expect("detection reported no observation for the seeded table");
assert_eq!(observed.format, "LANCE");
for attribute in ["fragments", "rows", "versions", "schema"] {
assert!(
observed.attributes.contains_key(attribute),
"observation is missing {attribute}: {:?}",
observed.attributes.keys().collect::<Vec<_>>()
);
}
let execute = ExecuteJobRequest {
request_id: "execute-1".to_string(),
job: Some(JobSpec {
job_id: "job-1".to_string(),
job_type: JOB_TYPE.to_string(),
parameters: proposal.parameters.clone(),
..Default::default()
}),
worker_config_values: int_config("target_rows_per_fragment", 1_048_576),
..Default::default()
};
handler
.execute(&execute, &recorder)
.await
.expect("execution failed");
let completed = recorder.completed.lock().unwrap().clone();
let result = completed.first().expect("no completion reported");
assert!(
result.success,
"compaction reported failure: {}",
result.error_message
);
let summary = result
.result
.as_ref()
.map(|r| r.summary.clone())
.unwrap_or_default();
assert!(
summary.contains("became"),
"completion carried no fragment counts: {summary}"
);
eprintln!("compaction result: {summary}");
}
/// A table with more versions than the floor is proposed, and running the job
/// reports what it removed. The compaction test above leaves one behind.
#[tokio::test]
async fn cleans_up_old_versions() {
let _gateway = GATEWAY.lock().await;
let Some(url) = namespace_url() else {
eprintln!("WEED_LANCE_NAMESPACE is unset, skipping");
return;
};
let encoded = seed_table(&url, "cleanme", 6, 4, false)
.await
.expect("seed a table with versions to clean");
let handler = CleanupVersionsHandler::new(url).with_fallback(fallback());
let recorder = Recorder::default();
let request = RunDetectionRequest {
request_id: "detect-cleanup".to_string(),
job_type: "lance_cleanup_versions".to_string(),
worker_config_values: int_config("min_versions_to_keep", 2),
..Default::default()
};
handler
.detect(&request, &recorder)
.await
.expect("detection failed");
let proposals = recorder.proposals.lock().unwrap().clone();
let proposal = proposals
.iter()
.find(|p| p.summary.contains(encoded.as_str()))
.cloned()
.expect("no cleanup proposal for the seeded table");
// Retain nothing, so every version outside the current one is fair game and
// the job has something to report rather than a no-op.
let execute = ExecuteJobRequest {
request_id: "execute-cleanup".to_string(),
job: Some(JobSpec {
job_id: "job-cleanup".to_string(),
job_type: "lance_cleanup_versions".to_string(),
parameters: proposal.parameters.clone(),
..Default::default()
}),
worker_config_values: int_config("retain_hours", 0),
..Default::default()
};
handler
.execute(&execute, &recorder)
.await
.expect("cleanup failed");
let completed = recorder.completed.lock().unwrap().clone();
let result = completed.last().expect("no completion reported");
assert!(
result.success,
"cleanup reported failure: {}",
result.error_message
);
eprintln!(
"cleanup result: {}",
result.result.as_ref().unwrap().summary
);
}
/// A table with no indices has nothing to optimize, so detection proposes
/// nothing rather than queueing work that would do nothing.
#[tokio::test]
async fn skips_tables_without_indices() {
let _gateway = GATEWAY.lock().await;
let Some(url) = namespace_url() else {
eprintln!("WEED_LANCE_NAMESPACE is unset, skipping");
return;
};
let encoded = seed_table(&url, "noindex", 2, 8, false)
.await
.expect("seed a table without an index");
let handler = OptimizeIndicesHandler::new(url).with_fallback(fallback());
let recorder = Recorder::default();
let request = RunDetectionRequest {
request_id: "detect-indices".to_string(),
job_type: "lance_optimize_indices".to_string(),
worker_config_values: int_config("max_unindexed_rows", 1),
..Default::default()
};
handler
.detect(&request, &recorder)
.await
.expect("detection failed");
// Judge this table only: the catalog holds every other test's tables too,
// and an indexed one with uncovered rows is supposed to be proposed.
assert!(
!recorder
.proposals
.lock()
.unwrap()
.iter()
.any(|p| p.summary.contains(encoded.as_str())),
"a table with no indices must not be proposed for reindexing"
);
}
/// The job with no Iceberg equivalent: rows appended after an index was built
/// are invisible to a search of it until this runs. Needs a table with an index
/// and rows outside it, which `indexed.py` in the scratchpad seeds.
#[tokio::test]
async fn reindexes_rows_an_index_does_not_cover() {
let _gateway = GATEWAY.lock().await;
let Some(url) = namespace_url() else {
eprintln!("WEED_LANCE_NAMESPACE is unset, skipping");
return;
};
let encoded = seed_table(&url, "reindexme", 2, 512, true)
.await
.expect("seed an indexed table with uncovered rows");
let handler = OptimizeIndicesHandler::new(url).with_fallback(fallback());
let recorder = Recorder::default();
let request = RunDetectionRequest {
request_id: "detect-reindex".to_string(),
job_type: "lance_optimize_indices".to_string(),
worker_config_values: int_config("max_unindexed_rows", 100),
..Default::default()
};
handler
.detect(&request, &recorder)
.await
.expect("detection failed");
let proposals = recorder.proposals.lock().unwrap().clone();
let proposal = proposals
.iter()
.find(|p| p.summary.contains(encoded.as_str()))
.cloned()
.expect("the seeded indexed table was not proposed");
let execute = ExecuteJobRequest {
request_id: "execute-reindex".to_string(),
job: Some(JobSpec {
job_id: "job-reindex".to_string(),
job_type: "lance_optimize_indices".to_string(),
parameters: proposal.parameters.clone(),
..Default::default()
}),
..Default::default()
};
handler
.execute(&execute, &recorder)
.await
.expect("reindex failed");
let completed = recorder.completed.lock().unwrap().clone();
let result = completed.last().expect("no completion reported");
assert!(
result.success,
"reindex reported failure: {}",
result.error_message
);
let output = &result.result.as_ref().unwrap().output_values;
let after = match output
.get("unindexed_rows_after")
.and_then(|v| v.kind.as_ref())
{
Some(Kind::Int64Value(value)) => *value,
other => panic!("no unindexed_rows_after in {other:?}"),
};
assert_eq!(
after, 0,
"rows are still outside the index after optimizing"
);
eprintln!(
"reindex result: {}",
result.result.as_ref().unwrap().summary
);
}
/// The UI's whole reason for asking a worker: admin cannot read a Lance table,
/// so the rows have to come back already rendered.
#[tokio::test]
async fn previews_rows_of_a_table() {
let _gateway = GATEWAY.lock().await;
let Some(url) = namespace_url() else {
eprintln!("set WEED_LANCE_NAMESPACE to run this test");
return;
};
seed_table(&url, "previewme", 2, 3, false)
.await
.expect("seed a table to preview");
let provider = LancePreview::new(url, fallback());
let id = vec!["vec".to_string(), "ml".to_string(), "previewme".to_string()];
let preview = provider.preview(&id, 4).await.expect("preview the table");
assert_eq!(preview.columns, vec!["id".to_string(), "vec".to_string()]);
assert_eq!(preview.total_rows, 6, "total is the table, not the sample");
assert_eq!(preview.rows.len(), 4, "row_limit bounds the sample");
assert!(
preview.rows[0][1].starts_with('['),
"a vector column should render as a list, got {:?}",
preview.rows[0][1]
);
}
/// The claim that removed managed versioning is not "one writer wins the
/// conditional PUT" - that is only the mechanism. It is that concurrent writers
/// lose nothing: the loser sees the conflict, rebases, and commits again. Eight
/// writers appending at once must leave all eight batches in the table.
#[tokio::test]
async fn concurrent_writers_keep_every_commit() {
let _gateway = GATEWAY.lock().await;
let Some(url) = namespace_url() else {
eprintln!("set WEED_LANCE_NAMESPACE to run this test");
return;
};
const WRITERS: i64 = 8;
const ROWS_EACH: i64 = 4;
seed_table(&url, "racers", 1, ROWS_EACH as usize, false)
.await
.expect("seed the table the writers will append to");
let client = NamespaceClient::new(url.clone());
let id = vec!["vec".to_string(), "ml".to_string(), "racers".to_string()];
let description = client.describe_table(&id).await.expect("describe");
let mut options = description.storage_options.clone();
options.extend(fallback());
let writes = (0..WRITERS).map(|writer| {
let location = description.location.clone();
let options = options.clone();
tokio::spawn(
async move { append_rows(&location, &options, writer * 1000, ROWS_EACH).await },
)
});
for (writer, handle) in writes.enumerate() {
handle
.await
.expect("writer panicked")
.unwrap_or_else(|err| panic!("writer {writer} failed to commit: {err:#}"));
}
let table = weed_lance_worker::dataset::open(&client, &id, &fallback())
.await
.expect("reopen the table");
let rows = table.dataset.count_rows(None).await.expect("count rows");
let expected = (ROWS_EACH + WRITERS * ROWS_EACH) as usize;
assert_eq!(
rows, expected,
"concurrent commits lost data: {rows} rows, want {expected}"
);
}
/// Appends one batch to an existing dataset, the way an independent writer would.
async fn append_rows(
location: &str,
options: &HashMap<String, String>,
first_id: i64,
rows: i64,
) -> Result<()> {
use arrow_array::{
FixedSizeListArray, Float32Array, Int64Array, RecordBatch, RecordBatchIterator,
};
use arrow_schema::{DataType, Field, Schema};
use lance::dataset::{Dataset, WriteMode, WriteParams};
use lance::io::{ObjectStoreParams, StorageOptionsAccessor};
use std::sync::Arc;
const DIM: i32 = 16;
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new(
"vec",
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), DIM),
false,
),
]));
let ids: Vec<i64> = (0..rows).map(|r| first_id + r).collect();
let values: Vec<f32> = ids
.iter()
.flat_map(|id| (0..DIM).map(move |d| (*id as f32) + d as f32))
.collect();
let vectors = FixedSizeListArray::new(
Arc::new(Field::new("item", DataType::Float32, true)),
DIM,
Arc::new(Float32Array::from(values)),
None,
);
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int64Array::from(ids)), Arc::new(vectors)],
)?;
let params = WriteParams {
mode: WriteMode::Append,
store_params: Some(ObjectStoreParams {
storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options(
options.clone(),
))),
..Default::default()
}),
..Default::default()
};
Dataset::write(
RecordBatchIterator::new(vec![Ok(batch)], schema.clone()),
location,
Some(params),
)
.await?;
Ok(())
}

View File

@@ -0,0 +1,123 @@
//! Maintains one table someone else wrote, and nothing else.
//!
//! The Go suite in test/s3tables/lifecycle drives a table through its whole
//! life - declared through the namespace, written by a Lance client, maintained,
//! read again, dropped - and this is the maintenance step of it, so that step
//! goes through the handlers a deployed worker runs rather than through the two
//! lance calls they wrap. The table to work on comes in as WEED_LANCE_TABLE
//! because the table is the Go test's, seeded and checked there.
//!
//! Skipped, like the rest of this crate's integration tests, unless
//! WEED_LANCE_NAMESPACE names a live gateway.
use std::collections::HashMap;
use seaweed_worker_core::JobHandler;
use seaweed_worker_core::pb::{ConfigValue, ExecuteJobRequest, JobSpec, RunDetectionRequest};
use weed_lance_worker::jobs::cleanup::{self, CleanupVersionsHandler};
use weed_lance_worker::jobs::compact::{CompactHandler, JOB_TYPE as COMPACT_JOB_TYPE};
mod common;
use common::{Recorder, fallback, int_config, namespace_url};
fn table() -> Option<String> {
std::env::var("WEED_LANCE_TABLE")
.ok()
.filter(|s| !s.is_empty())
}
/// Compacts the named table's fragments, then drops the versions compaction
/// left behind. Both go through the handler's own detect-then-execute path: a
/// proposal the worker would not have made is not one worth running.
#[tokio::test]
async fn maintains_the_named_table() {
let (Some(url), Some(name)) = (namespace_url(), table()) else {
eprintln!("WEED_LANCE_NAMESPACE or WEED_LANCE_TABLE is unset, skipping");
return;
};
// A table written as a fragment per append is the case this exists for, so
// anything above one fragment is worth merging here.
let mut config = int_config("min_fragments", 2);
config.extend(int_config("target_rows_per_fragment", 1_048_576));
let compact = CompactHandler::new(url.clone()).with_fallback(fallback());
run(&compact, COMPACT_JOB_TYPE, &name, config, true).await;
// Keep the current version and retain nothing else, so the versions
// compaction superseded are actually removed rather than counted.
let mut config = int_config("min_versions_to_keep", 1);
config.extend(int_config("retain_hours", 0));
let cleanup = CleanupVersionsHandler::new(url).with_fallback(fallback());
run(&cleanup, cleanup::JOB_TYPE, &name, config, true).await;
}
/// Detects, finds the proposal for `name`, and executes it. The config carries
/// every key both halves read; each ignores what it does not know.
async fn run<H: JobHandler>(
handler: &H,
job_type: &str,
name: &str,
config: HashMap<String, ConfigValue>,
required: bool,
) {
let recorder = Recorder::default();
handler
.detect(
&RunDetectionRequest {
request_id: format!("detect-{job_type}"),
job_type: job_type.to_string(),
worker_config_values: config.clone(),
..Default::default()
},
&recorder,
)
.await
.unwrap_or_else(|err| panic!("{job_type} detection failed: {err}"));
let proposals = recorder.proposals.lock().unwrap().clone();
let Some(proposal) = proposals.iter().find(|p| p.summary.contains(name)) else {
assert!(
!required,
"{job_type} proposed nothing for {name}, out of {} proposals",
proposals.len()
);
eprintln!("{job_type}: nothing to do for {name}");
return;
};
handler
.execute(
&ExecuteJobRequest {
request_id: format!("execute-{job_type}"),
job: Some(JobSpec {
job_id: format!("job-{job_type}"),
job_type: job_type.to_string(),
parameters: proposal.parameters.clone(),
..Default::default()
}),
worker_config_values: config,
..Default::default()
},
&recorder,
)
.await
.unwrap_or_else(|err| panic!("{job_type} execution failed: {err}"));
let completed = recorder.completed.lock().unwrap().clone();
let result = completed
.last()
.unwrap_or_else(|| panic!("{job_type} reported no completion"));
assert!(
result.success,
"{job_type} reported failure: {}",
result.error_message
);
eprintln!(
"{job_type}: {}",
result
.result
.as_ref()
.map(|r| r.summary.clone())
.unwrap_or_default()
);
}

View File

@@ -0,0 +1,504 @@
//! Drives the sort handler against a live namespace.
//!
//! Skipped unless WEED_LANCE_NAMESPACE names one, like the other integration
//! tests here: the job rewrites every fragment of a real table and commits, and
//! the commit is the half worth testing.
use anyhow::Result;
use seaweed_worker_core::JobHandler;
use seaweed_worker_core::pb::{
ConfigValue, ExecuteJobRequest, JobSpec, RunDetectionRequest, config_value::Kind,
};
use weed_lance_worker::catalog::NamespaceClient;
use weed_lance_worker::jobs::sort::{JOB_TYPE, SortHandler};
mod common;
use common::{Recorder, fallback, namespace_url};
/// One live gateway and one shared catalog, and `list_all_tables` sweeps
/// everything, so these tests take a lock the way the compaction ones do.
static GATEWAY: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
fn config(values: &[(&str, ConfigValue)]) -> std::collections::HashMap<String, ConfigValue> {
values
.iter()
.map(|(name, value)| (name.to_string(), value.clone()))
.collect()
}
fn int(value: i64) -> ConfigValue {
ConfigValue {
kind: Some(Kind::Int64Value(value)),
}
}
fn text(value: &str) -> ConfigValue {
ConfigValue {
kind: Some(Kind::StringValue(value.to_string())),
}
}
/// Declares a table and writes `batches` appends whose ids descend, so the rows
/// land in an order no sort would produce. `declared` is written into the
/// dataset's own configuration, which is where a table says how it wants to be
/// sorted.
async fn seed_unsorted_table(
url: &str,
name: &str,
batches: i64,
rows_each: i64,
declared: Option<&str>,
) -> Result<String> {
use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator};
use arrow_schema::{DataType, Field, Schema};
use lance::dataset::{Dataset, WriteMode, WriteParams};
use lance::io::{ObjectStoreParams, StorageOptionsAccessor};
use std::sync::Arc;
let http = reqwest::Client::new();
for parent in ["vec", "vec$ml"] {
http.post(format!("{url}/v1/namespace/{parent}/create"))
.json(&serde_json::json!({"mode": "EXIST_OK"}))
.send()
.await?
.error_for_status()?;
}
let encoded = format!("vec$ml${name}");
http.post(format!("{url}/v1/table/{encoded}/declare"))
.json(&serde_json::json!({}))
.send()
.await?
.error_for_status()?;
let client = NamespaceClient::new(url.to_string());
let id = vec!["vec".to_string(), "ml".to_string(), name.to_string()];
let description = client.describe_table(&id).await?;
let mut options = description.storage_options.clone();
options.extend(fallback());
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
let total = batches * rows_each;
let mut dataset = None;
for batch_index in 0..batches {
// Ids descend across and within the batches, so the physical order is
// the reverse of the order the table declares.
let ids: Vec<i64> = (0..rows_each)
.map(|row| total - 1 - (batch_index * rows_each + row))
.collect();
let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(ids))])?;
let params = WriteParams {
mode: if batch_index == 0 {
WriteMode::Overwrite
} else {
WriteMode::Append
},
store_params: Some(ObjectStoreParams {
storage_options_accessor: Some(Arc::new(
StorageOptionsAccessor::with_static_options(options.clone()),
)),
..Default::default()
}),
..Default::default()
};
dataset = Some(
Dataset::write(
RecordBatchIterator::new(vec![Ok(batch)], schema.clone()),
description.location.as_str(),
Some(params),
)
.await?,
);
}
if let Some(declared) = declared {
let mut dataset = dataset.expect("at least one batch was written");
dataset
.update_config([("seaweedfs.sort.fields", declared)])
.await?;
}
Ok(encoded)
}
/// Reads the ids back in the order they are stored, which is what a sort is
/// supposed to change.
async fn stored_ids(url: &str, name: &str) -> Result<Vec<i64>> {
use arrow_array::Int64Array;
use futures::TryStreamExt;
let client = NamespaceClient::new(url.to_string());
let id = vec!["vec".to_string(), "ml".to_string(), name.to_string()];
let table = weed_lance_worker::dataset::open(&client, &id, &fallback()).await?;
let stream = table.dataset.scan().try_into_stream().await?;
let batches: Vec<_> = stream.try_collect().await?;
let mut ids = Vec::new();
for batch in batches {
let column = batch
.column_by_name("id")
.expect("the seeded table has an id column")
.as_any()
.downcast_ref::<Int64Array>()
.expect("id is an Int64 column")
.clone();
ids.extend(column.values().iter().copied());
}
Ok(ids)
}
/// The whole loop: a table declaring an order its rows do not follow is
/// proposed, sorting it stores the rows in that order, and the marker the same
/// commit carries is what stops the next sweep proposing it again.
#[tokio::test]
async fn sorts_a_table_and_stops_proposing_it() {
let _gateway = GATEWAY.lock().await;
let Some(url) = namespace_url() else {
eprintln!("WEED_LANCE_NAMESPACE is unset, skipping");
return;
};
let encoded = seed_unsorted_table(&url, "sortme", 4, 250, Some("id asc"))
.await
.expect("seed an unsorted table");
let before = stored_ids(&url, "sortme").await.expect("read the ids back");
assert!(
before.windows(2).any(|pair| pair[0] > pair[1]),
"the seeded table was already sorted, so the test proves nothing"
);
let handler = SortHandler::new(url.clone()).with_fallback(fallback());
let recorder = Recorder::default();
let detection = RunDetectionRequest {
request_id: "detect-sort".to_string(),
job_type: JOB_TYPE.to_string(),
worker_config_values: config(&[("min_unsorted_rows", int(1))]),
..Default::default()
};
handler
.detect(&detection, &recorder)
.await
.expect("detection failed");
let proposal = recorder
.proposals
.lock()
.unwrap()
.iter()
.find(|p| p.summary.contains(encoded.as_str()))
.cloned()
.expect("the unsorted table was not proposed");
assert!(
proposal.summary.contains("id asc"),
"the proposal does not say what order it would use: {}",
proposal.summary
);
let execute = ExecuteJobRequest {
request_id: "execute-sort".to_string(),
job: Some(JobSpec {
job_id: "job-sort".to_string(),
job_type: JOB_TYPE.to_string(),
parameters: proposal.parameters.clone(),
..Default::default()
}),
worker_config_values: config(&[
("memory_budget_mb", int(64)),
("max_rows_per_file", int(1024)),
]),
..Default::default()
};
handler
.execute(&execute, &recorder)
.await
.expect("sorting failed");
let completed = recorder.completed.lock().unwrap().clone();
let result = completed.last().expect("no completion reported");
assert!(
result.success,
"the sort reported failure: {}",
result.error_message
);
let after = stored_ids(&url, "sortme")
.await
.expect("read the sorted ids back");
assert_eq!(after.len(), before.len(), "the sort lost or invented rows");
assert!(
after.windows(2).all(|pair| pair[0] <= pair[1]),
"the rows are not stored in ascending id order"
);
// The marker rode in the same commit as the data, so a second sweep finds
// the table up to date without anything else having written it.
let second = Recorder::default();
handler
.detect(&detection, &second)
.await
.expect("second detection failed");
assert!(
!second
.proposals
.lock()
.unwrap()
.iter()
.any(|p| p.summary.contains(encoded.as_str())),
"a table that was just sorted must not be proposed again"
);
eprintln!("sort result: {}", result.result.as_ref().unwrap().summary);
}
/// A table that declares no order, with a worker configuring none either, is
/// not something to rewrite on a guess.
#[tokio::test]
async fn leaves_a_table_that_declares_no_order_alone() {
let _gateway = GATEWAY.lock().await;
let Some(url) = namespace_url() else {
eprintln!("WEED_LANCE_NAMESPACE is unset, skipping");
return;
};
let encoded = seed_unsorted_table(&url, "unsorted", 2, 16, None)
.await
.expect("seed a table with no declared order");
let handler = SortHandler::new(url.clone()).with_fallback(fallback());
let recorder = Recorder::default();
handler
.detect(
&RunDetectionRequest {
request_id: "detect-none".to_string(),
job_type: JOB_TYPE.to_string(),
worker_config_values: config(&[
("min_unsorted_rows", int(1)),
("sort_fields", text("")),
]),
..Default::default()
},
&recorder,
)
.await
.expect("detection failed");
assert!(
!recorder
.proposals
.lock()
.unwrap()
.iter()
.any(|p| p.summary.contains(encoded.as_str())),
"a table with no order named anywhere must not be proposed"
);
}
/// The case a row-count threshold cannot see: the table replaced by a write of
/// its own, growing by far less than `min_unsorted_rows`. Every row moved, so
/// the table is unsorted, and detection has to say so however small the delta.
#[tokio::test]
async fn proposes_a_replacement_that_barely_changes_the_row_count() {
let _gateway = GATEWAY.lock().await;
let Some(url) = namespace_url() else {
eprintln!("WEED_LANCE_NAMESPACE is unset, skipping");
return;
};
let encoded = seed_unsorted_table(&url, "replaceme", 4, 250, Some("id asc"))
.await
.expect("seed an unsorted table");
let handler = SortHandler::new(url.clone()).with_fallback(fallback());
// A threshold no realistic append would cross, so nothing but the identity
// of the data can be what proposes this table.
let detection = RunDetectionRequest {
request_id: "detect-replace".to_string(),
job_type: JOB_TYPE.to_string(),
worker_config_values: config(&[("min_unsorted_rows", int(1_000_000))]),
..Default::default()
};
let first = Recorder::default();
handler.detect(&detection, &first).await.expect("detection");
let proposal = first
.proposals
.lock()
.unwrap()
.iter()
.find(|p| p.summary.contains(encoded.as_str()))
.cloned()
.expect("a never-sorted table was not proposed");
handler
.execute(
&ExecuteJobRequest {
request_id: "execute-replace".to_string(),
job: Some(JobSpec {
job_id: "job-replace".to_string(),
job_type: JOB_TYPE.to_string(),
parameters: proposal.parameters.clone(),
..Default::default()
}),
worker_config_values: config(&[("max_rows_per_file", int(1024))]),
..Default::default()
},
&first,
)
.await
.expect("sorting failed");
// Sorted, and left alone by the next sweep.
let after_sort = Recorder::default();
handler
.detect(&detection, &after_sort)
.await
.expect("detection");
assert!(
!after_sort
.proposals
.lock()
.unwrap()
.iter()
.any(|p| p.summary.contains(encoded.as_str())),
"a table that was just sorted must not be proposed again"
);
// Now replace the data: 1002 rows where there were 1000, shuffled again.
seed_unsorted_table(&url, "replaceme", 3, 334, Some("id asc"))
.await
.expect("replace the table's data");
let after_replace = Recorder::default();
handler
.detect(&detection, &after_replace)
.await
.expect("detection");
assert!(
after_replace
.proposals
.lock()
.unwrap()
.iter()
.any(|p| p.summary.contains(encoded.as_str())),
"a replaced table must be proposed even when the row count barely moved"
);
}
/// Deletions must not hide appended rows. This is the arithmetic against a real
/// manifest rather than a hand-built one: the sorted fragments keep their files
/// and lose most of their rows, a smaller batch is appended after them, and the
/// table has to be proposed *because rows were appended* — not because it looks
/// replaced, and not left alone because it now holds fewer rows than when it
/// was sorted.
#[tokio::test]
async fn deletions_do_not_hide_appended_rows() {
let _gateway = GATEWAY.lock().await;
let Some(url) = namespace_url() else {
eprintln!("WEED_LANCE_NAMESPACE is unset, skipping");
return;
};
let encoded = seed_unsorted_table(&url, "deleteme", 4, 250, Some("id asc"))
.await
.expect("seed an unsorted table");
let handler = SortHandler::new(url.clone()).with_fallback(fallback());
let detection = RunDetectionRequest {
request_id: "detect-delete".to_string(),
job_type: JOB_TYPE.to_string(),
worker_config_values: config(&[("min_unsorted_rows", int(100))]),
..Default::default()
};
let recorder = Recorder::default();
handler
.detect(&detection, &recorder)
.await
.expect("detection");
let proposal = recorder
.proposals
.lock()
.unwrap()
.iter()
.find(|p| p.summary.contains(encoded.as_str()))
.cloned()
.expect("a never-sorted table was not proposed");
handler
.execute(
&ExecuteJobRequest {
request_id: "execute-delete".to_string(),
job: Some(JobSpec {
job_id: "job-delete".to_string(),
job_type: JOB_TYPE.to_string(),
parameters: proposal.parameters.clone(),
..Default::default()
}),
worker_config_values: config(&[("max_rows_per_file", int(1024))]),
..Default::default()
},
&recorder,
)
.await
.expect("sorting failed");
// Delete 900 of the 1000 sorted rows, then append 200 unsorted ones. The
// table is smaller than when it was sorted, so net growth reads as zero.
let client = NamespaceClient::new(url.clone());
let id = vec!["vec".to_string(), "ml".to_string(), "deleteme".to_string()];
let mut table = weed_lance_worker::dataset::open(&client, &id, &fallback())
.await
.expect("reopen the sorted table");
table
.dataset
.delete("id < 900")
.await
.expect("delete most of the sorted rows");
append_ids(&url, "deleteme", 5_000, 200)
.await
.expect("append rows after the sorted ones");
let after = Recorder::default();
handler.detect(&detection, &after).await.expect("detection");
let proposal = after
.proposals
.lock()
.unwrap()
.iter()
.find(|p| p.summary.contains(encoded.as_str()))
.cloned()
.expect("a table with 200 unsorted rows appended was not proposed");
assert!(
proposal.detail.contains("200 rows appended"),
"expected the appended rows to be counted, got {:?}",
proposal.detail
);
}
/// Appends one batch of ascending ids to an existing table.
async fn append_ids(url: &str, name: &str, first_id: i64, rows: i64) -> Result<()> {
use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator};
use arrow_schema::{DataType, Field, Schema};
use lance::dataset::{Dataset, WriteMode, WriteParams};
use lance::io::{ObjectStoreParams, StorageOptionsAccessor};
use std::sync::Arc;
let client = NamespaceClient::new(url.to_string());
let id = vec!["vec".to_string(), "ml".to_string(), name.to_string()];
let description = client.describe_table(&id).await?;
let mut options = description.storage_options.clone();
options.extend(fallback());
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
let ids: Vec<i64> = (0..rows).map(|row| first_id + row).collect();
let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(ids))])?;
let params = WriteParams {
mode: WriteMode::Append,
store_params: Some(ObjectStoreParams {
storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options(
options,
))),
..Default::default()
}),
..Default::default()
};
Dataset::write(
RecordBatchIterator::new(vec![Ok(batch)], schema.clone()),
description.location.as_str(),
Some(params),
)
.await?;
Ok(())
}

View File

@@ -0,0 +1,16 @@
[package]
name = "seaweed-worker-sort"
version.workspace = true
edition.workspace = true
rust-version.workspace = true
description = "The sort specification shared by SeaweedFS sorting jobs"
[lib]
name = "seaweed_worker_sort"
[dependencies]
seaweed-worker-core = { path = "../core" }
anyhow.workspace = true
[lints]
workspace = true

View File

@@ -0,0 +1,612 @@
//! The sort specification the sorting jobs share.
//!
//! A worker sorts a table by the fields the table itself declares, and by the
//! fields its own configuration names for a table that declares none. Both
//! halves of that rule, and the spelling of a field list, live here rather than
//! in each job: a Lance job and an Iceberg job that disagreed about what
//! "id desc nulls-first" means would be two features wearing one name.
//!
//! The spelling is deliberately the one `weed/worker/tasks/iceberg` already
//! writes into an Iceberg snapshot's `sort-fields` summary, so the same order
//! reads the same way whichever format holds the table and whichever language
//! sorted it.
use std::collections::{HashMap, HashSet};
use std::fmt;
use anyhow::{Context, Result, bail};
use seaweed_worker_core::config_form::{number_field, text_field};
use seaweed_worker_core::pb::ConfigField;
/// Worker configuration keys. Shared so the Iceberg and Lance forms offer an
/// operator the same names for the same settings.
pub const CONFIG_SORT_FIELDS: &str = "sort_fields";
pub const CONFIG_MIN_UNSORTED_ROWS: &str = "min_unsorted_rows";
pub const CONFIG_MEMORY_BUDGET_MB: &str = "memory_budget_mb";
pub const CONFIG_MAX_ROWS_PER_FILE: &str = "max_rows_per_file";
/// The order an operator declares on a table. The worker only ever reads this
/// one.
pub const DECLARED_FIELDS_KEY: &str = "seaweedfs.sort.fields";
/// What the worker recorded about the sort it last performed. Kept apart from
/// the declaration above so that "the operator asked for this order" and "the
/// worker achieved this order" stay distinguishable — comparing them is how a
/// changed order is noticed.
pub const SORTED_FIELDS_KEY: &str = "seaweedfs.sort.sorted_fields";
pub const SORTED_ROWS_KEY: &str = "seaweedfs.sort.sorted_rows";
/// How many fragments the sort wrote, and a digest of their data file names.
/// Together they identify the data the sort produced — see [`verdict`] for why
/// a version number cannot.
pub const SORTED_FRAGMENTS_KEY: &str = "seaweedfs.sort.sorted_fragments";
pub const SORTED_DIGEST_KEY: &str = "seaweedfs.sort.sorted_digest";
/// One field of a sort order.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SortField {
pub path: String,
pub descending: bool,
pub nulls_first: bool,
}
/// A whole sort order, in the order the fields are compared.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SortSpec {
pub fields: Vec<SortField>,
}
impl SortSpec {
/// Reads "user_id asc nulls-last, ts desc" into a spec. An empty string is
/// not an error: it is how an operator says nothing, and the caller decides
/// what that means.
///
/// Direction defaults to ascending and null order to Iceberg's default for
/// the direction — nulls first ascending, nulls last descending — because
/// this spelling is shared with tables that already have that rule.
pub fn parse(text: &str) -> Result<Option<Self>> {
let mut fields = Vec::new();
let mut seen = HashSet::new();
for entry in text.split(',') {
let entry = entry.trim();
if entry.is_empty() {
continue;
}
let field =
parse_field(entry).with_context(|| format!("read the sort field {entry:?}"))?;
// Compared exactly, not case-folded: Arrow schemas are
// case-sensitive, so `id` and `ID` can be two real columns and
// folding them together would reject a valid order. The Iceberg
// job, whose format resolves names case-insensitively, does its own
// check against the table schema.
if !seen.insert(field.path.clone()) {
bail!("the sort field {:?} is named twice", field.path);
}
fields.push(field);
}
if fields.is_empty() {
return Ok(None);
}
Ok(Some(Self { fields }))
}
}
/// Renders a spec back, always spelling out direction and null order so that a
/// recorded order round-trips to the same thing it was parsed from.
impl fmt::Display for SortSpec {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
for (i, field) in self.fields.iter().enumerate() {
if i > 0 {
write!(f, ", ")?;
}
write!(
f,
"{} {} {}",
field.path,
if field.descending { "desc" } else { "asc" },
if field.nulls_first {
"nulls-first"
} else {
"nulls-last"
}
)?;
}
Ok(())
}
}
fn parse_field(entry: &str) -> Result<SortField> {
let mut tokens = entry.split_whitespace();
let path = tokens
.next()
.expect("split_whitespace yields one token for a non-empty entry")
.to_string();
let mut descending = None;
let mut nulls_first = None;
for token in tokens {
// Underscores are accepted because that is how the same words are
// spelled in configuration keys, and refusing them would only teach an
// operator that the two spellings are different things.
match token.to_ascii_lowercase().replace('_', "-").as_str() {
"asc" | "ascending" if descending.is_none() => descending = Some(false),
"desc" | "descending" if descending.is_none() => descending = Some(true),
"nulls-first" if nulls_first.is_none() => nulls_first = Some(true),
"nulls-last" if nulls_first.is_none() => nulls_first = Some(false),
other => bail!("{other:?} is not a direction or a null order"),
}
}
let descending = descending.unwrap_or(false);
Ok(SortField {
path,
descending,
nulls_first: nulls_first.unwrap_or(!descending),
})
}
/// The order to sort a table by: what the table declares wins, and the worker's
/// configuration is the fallback for a table that declares nothing. Neither one
/// is not an error — it means this is not a table the operator wants sorted.
///
/// A declaration that cannot be read is an error rather than a reason to fall
/// back: sorting by the worker's default order instead of the one the table
/// asked for would silently rewrite the table the wrong way.
pub fn resolve(declared: Option<&str>, configured: &str) -> Result<Option<SortSpec>> {
if let Some(declared) = declared
&& let Some(spec) = SortSpec::parse(declared).context("read the table's declared order")?
{
return Ok(Some(spec));
}
SortSpec::parse(configured).context("read the configured sort order")
}
/// One fragment as detection sees it.
///
/// `rows` is the live row count the manifest records — physical rows less the
/// deletions — and is None when the manifest does not say, which older
/// fragments do not.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FragmentSummary {
pub files: Vec<String>,
pub rows: Option<u64>,
}
/// Every data file of these fragments, in order.
pub fn data_files(fragments: &[FragmentSummary]) -> Vec<&str> {
fragments
.iter()
.flat_map(|fragment| fragment.files.iter().map(String::as_str))
.collect()
}
/// A digest of data file names, in the order the fragments hold them.
///
/// FNV-1a rather than the standard library's hasher, whose output is explicitly
/// not stable across releases: a marker that hashed differently after a
/// toolchain upgrade would re-sort every table once, silently.
pub fn digest<S: AsRef<str>>(files: &[S]) -> String {
let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
for file in files {
for byte in file.as_ref().as_bytes().iter().chain(std::iter::once(&0u8)) {
hash ^= *byte as u64;
hash = hash.wrapping_mul(0x1000_0000_01b3);
}
}
format!("{hash:016x}")
}
/// What the worker recorded on a table the last time it sorted it.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct SortState {
pub fields: Option<String>,
pub rows: Option<u64>,
/// The fragments the sort wrote: how many, and a digest of their data file
/// names. Data file names are chosen before the commit, which is what makes
/// them usable in a marker the same commit carries.
pub fragments: Option<usize>,
pub digest: Option<String>,
}
impl SortState {
/// Reads the marker out of a table's key/value configuration. A key that
/// will not parse is treated as absent, which asks for a sort rather than
/// skipping the table: the cost of a needless sort is time, and the cost of
/// skipping is a table that silently never gets sorted again.
pub fn from_config(config: &HashMap<String, String>) -> Self {
Self {
fields: config.get(SORTED_FIELDS_KEY).cloned(),
rows: config.get(SORTED_ROWS_KEY).and_then(|v| v.parse().ok()),
fragments: config
.get(SORTED_FRAGMENTS_KEY)
.and_then(|v| v.parse().ok()),
digest: config.get(SORTED_DIGEST_KEY).cloned(),
}
}
/// The marker to write for a sort that just finished.
pub fn record(
spec: &SortSpec,
fragments: &[FragmentSummary],
rows: u64,
) -> HashMap<String, String> {
HashMap::from([
(SORTED_FIELDS_KEY.to_string(), spec.to_string()),
(SORTED_ROWS_KEY.to_string(), rows.to_string()),
(
SORTED_FRAGMENTS_KEY.to_string(),
fragments.len().to_string(),
),
(
SORTED_DIGEST_KEY.to_string(),
digest(&data_files(fragments)),
),
])
}
}
/// Why a table does or does not need sorting.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SortVerdict {
NeverSorted,
OrderChanged,
RowsAppended(u64),
/// The table was committed to since the sort, and the row count did not
/// grow — the data is not the data that was sorted.
Rewritten,
UpToDate,
}
impl SortVerdict {
pub fn needs_sort(&self) -> bool {
!matches!(self, Self::UpToDate)
}
pub fn reason(&self) -> String {
match self {
Self::NeverSorted => "never sorted".to_string(),
Self::OrderChanged => "the declared order changed since the last sort".to_string(),
Self::RowsAppended(rows) => format!("{rows} rows appended since the last sort"),
Self::Rewritten => "the data was replaced since the last sort".to_string(),
Self::UpToDate => "sorted".to_string(),
}
}
}
/// Decides whether a table is worth sorting.
///
/// The question detection actually has to answer is "is this still the data the
/// sort wrote?", and neither the row count nor the dataset version can answer
/// it. A rewrite that leaves the row count where it was is invisible to the
/// first, and the second cannot be recorded at all: the commit carrying the
/// marker is the commit that creates the version, and a commit that rebases
/// past a conflict lands on a different number again.
///
/// Data file names can answer it. They are chosen before the commit, so the
/// marker the same commit carries can name them, and they are stable no matter
/// which version the commit ends up as. So:
///
/// - the same files → the table is exactly what the sort left;
/// - the sorted files still there, with more after them → rows were appended,
/// and only then is `min_unsorted_rows` consulted, which is the churn that
/// threshold exists to damp;
/// - anything else → the data was replaced, whatever the row count says.
///
/// Appends only ever add fragments after the existing ones, and lance hands out
/// fragment ids monotonically even across an overwrite, so "the sorted files
/// are still there" is a prefix test on the files in fragment order.
///
/// Deletes leave the data files alone and write a deletion file beside them, so
/// they read as unchanged here — which is right: deleting rows does not unsort
/// the ones that remain.
pub fn verdict(
spec: &SortSpec,
state: &SortState,
fragments: &[FragmentSummary],
min_unsorted_rows: u64,
) -> SortVerdict {
let (Some(recorded_digest), Some(recorded_fragments)) =
(state.digest.as_deref(), state.fragments)
else {
return SortVerdict::NeverSorted;
};
if state.fields.as_deref() != Some(spec.to_string().as_str()) {
return SortVerdict::OrderChanged;
}
if digest(&data_files(fragments)) == recorded_digest {
return SortVerdict::UpToDate;
}
// Not the same files. Only a prefix match means the sorted data survived
// and the rest was appended after it.
if recorded_fragments > fragments.len()
|| digest(&data_files(&fragments[..recorded_fragments])) != recorded_digest
{
return SortVerdict::Rewritten;
}
// How many rows are in the fragments appended after the sorted ones.
//
// Not the table's net growth against the recorded row count: rows deleted
// from the sorted fragments hide appended rows one for one, and enough of
// them hide every appended row, leaving a table that reads as sorted with
// an unsorted tail. Counting the appended fragments themselves is the same
// arithmetic the threshold was always meant to do.
let mut appended = 0u64;
for fragment in &fragments[recorded_fragments..] {
// A fragment whose length the manifest does not record cannot be
// counted, and a table that cannot be judged is one to sort rather than
// one to leave alone forever.
let Some(rows) = fragment.rows else {
return SortVerdict::Rewritten;
};
appended += rows;
}
if appended >= min_unsorted_rows.max(1) {
return SortVerdict::RowsAppended(appended);
}
SortVerdict::UpToDate
}
/// The settings every sorting job offers, so admin renders one form for the
/// same behaviour whichever format the job sorts.
pub fn config_fields() -> Vec<ConfigField> {
vec![
text_field(
CONFIG_SORT_FIELDS,
"Sort fields",
"Order for tables that declare none of their own, as \"user_id asc, ts desc nulls-last\". Empty sorts only the tables that declare an order.",
"user_id asc, ts desc",
),
number_field(
CONFIG_MIN_UNSORTED_ROWS,
"Minimum unsorted rows",
"Rows appended since the last sort before a table is worth sorting again",
1,
1_000_000_000,
),
number_field(
CONFIG_MEMORY_BUDGET_MB,
"Memory budget (MB)",
"Memory the sort may hold before it spills runs to disk. A table larger than this sorts more slowly rather than failing.",
64,
1_048_576,
),
number_field(
CONFIG_MAX_ROWS_PER_FILE,
"Maximum rows per file",
"Rows to write into each output file of a sorted rewrite",
1024,
16_777_216,
),
]
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_direction_and_null_order() {
let spec = SortSpec::parse("user_id, ts desc, name asc nulls-last")
.unwrap()
.unwrap();
assert_eq!(
spec.fields,
vec![
// Ascending defaults to nulls first, descending to nulls last,
// which is the rule the Iceberg tables sharing this spelling use.
SortField {
path: "user_id".into(),
descending: false,
nulls_first: true
},
SortField {
path: "ts".into(),
descending: true,
nulls_first: false
},
SortField {
path: "name".into(),
descending: false,
nulls_first: false
},
]
);
}
#[test]
fn renders_back_to_something_that_parses_the_same() {
let spec = SortSpec::parse("a, b desc nulls_first").unwrap().unwrap();
let rendered = spec.to_string();
assert_eq!(rendered, "a asc nulls-first, b desc nulls-first");
assert_eq!(SortSpec::parse(&rendered).unwrap().unwrap(), spec);
}
#[test]
fn empty_is_no_spec_rather_than_an_error() {
assert!(SortSpec::parse("").unwrap().is_none());
assert!(SortSpec::parse(" , ").unwrap().is_none());
}
#[test]
fn rejects_nonsense_and_repeats() {
assert!(SortSpec::parse("a sideways").is_err());
assert!(SortSpec::parse("a asc desc").is_err());
assert!(SortSpec::parse("a, a desc").is_err());
}
#[test]
fn the_table_wins_and_configuration_fills_in() {
let configured = "b desc";
let from_table = resolve(Some("a asc"), configured).unwrap().unwrap();
assert_eq!(from_table.fields[0].path, "a");
let from_config = resolve(None, configured).unwrap().unwrap();
assert_eq!(from_config.fields[0].path, "b");
let from_config_again = resolve(Some(" "), configured).unwrap().unwrap();
assert_eq!(from_config_again.fields[0].path, "b");
assert!(resolve(None, "").unwrap().is_none());
}
// Falling back would sort the table by an order nobody asked for.
#[test]
fn a_broken_declaration_is_an_error_not_a_fallback() {
assert!(resolve(Some("a sideways"), "b desc").is_err());
}
/// One fragment per name, one file each, with a known row count.
fn frags(names: &[(&str, u64)]) -> Vec<FragmentSummary> {
names
.iter()
.map(|(name, rows)| FragmentSummary {
files: vec![name.to_string()],
rows: Some(*rows),
})
.collect()
}
#[test]
fn verdicts_follow_the_marker() {
let spec = SortSpec::parse("a asc").unwrap().unwrap();
let sorted = frags(&[("a.lance", 500), ("b.lance", 500)]);
let fresh = SortState::default();
assert_eq!(
verdict(&spec, &fresh, &sorted, 100),
SortVerdict::NeverSorted
);
let state = SortState {
fields: Some(spec.to_string()),
rows: Some(1_000),
fragments: Some(sorted.len()),
digest: Some(digest(&data_files(&sorted))),
};
// The same files, whatever version the commit ended up as.
assert_eq!(verdict(&spec, &state, &sorted, 100), SortVerdict::UpToDate);
// Appended: the sorted fragments are still the first ones.
let appended_50 = frags(&[("a.lance", 500), ("b.lance", 500), ("c.lance", 50)]);
assert_eq!(
verdict(&spec, &state, &appended_50, 100),
SortVerdict::UpToDate
);
let appended_100 = frags(&[("a.lance", 500), ("b.lance", 500), ("c.lance", 100)]);
assert_eq!(
verdict(&spec, &state, &appended_100, 100),
SortVerdict::RowsAppended(100)
);
// Replaced: different files, and the row count is no defence — this is
// the case a row-count threshold lets through, whatever the delta.
let replaced = frags(&[("x.lance", 500), ("y.lance", 501)]);
assert_eq!(
verdict(&spec, &state, &replaced, 100),
SortVerdict::Rewritten
);
// Compacted into fewer files is a replacement too.
assert_eq!(
verdict(&spec, &state, &frags(&[("merged.lance", 1_000)]), 100),
SortVerdict::Rewritten
);
// Deleting rows rewrites no data file, so the table still reads as
// sorted — which it is.
let after_deletes = frags(&[("a.lance", 500), ("b.lance", 400)]);
assert_eq!(
verdict(&spec, &state, &after_deletes, 100),
SortVerdict::UpToDate
);
let other = SortSpec::parse("b desc").unwrap().unwrap();
assert_eq!(
verdict(&other, &state, &sorted, 100),
SortVerdict::OrderChanged
);
}
/// Deletions must not hide appended rows. Net growth against the recorded
/// row count would say this table shrank, and it would stay unsorted with
/// an unsorted tail for as long as the deletes kept pace with the appends.
#[test]
fn deletions_do_not_mask_appended_rows() {
let spec = SortSpec::parse("a asc").unwrap().unwrap();
let sorted = frags(&[("a.lance", 500), ("b.lance", 500)]);
let state = SortState {
fields: Some(spec.to_string()),
rows: Some(1_000),
fragments: Some(sorted.len()),
digest: Some(digest(&data_files(&sorted))),
};
// 800 rows deleted from the sorted fragments, 300 appended after them:
// the table is smaller than it was, and 300 rows of it are unsorted.
let mixed = frags(&[("a.lance", 100), ("b.lance", 100), ("c.lance", 300)]);
assert_eq!(
verdict(&spec, &state, &mixed, 100),
SortVerdict::RowsAppended(300)
);
// The threshold still damps a small append beside heavy deletion.
let small = frags(&[("a.lance", 100), ("b.lance", 100), ("c.lance", 20)]);
assert_eq!(verdict(&spec, &state, &small, 100), SortVerdict::UpToDate);
// And net growth is still detected where there are no deletions.
let grown = frags(&[("a.lance", 500), ("b.lance", 500), ("c.lance", 700)]);
assert_eq!(
verdict(&spec, &state, &grown, 100),
SortVerdict::RowsAppended(700)
);
}
// An appended fragment whose length the manifest does not record cannot be
// measured against the threshold, and a table that cannot be judged is one
// to sort rather than one to leave alone forever.
#[test]
fn an_uncountable_appended_fragment_is_stale() {
let spec = SortSpec::parse("a asc").unwrap().unwrap();
let sorted = frags(&[("a.lance", 500)]);
let state = SortState {
fields: Some(spec.to_string()),
rows: Some(500),
fragments: Some(1),
digest: Some(digest(&data_files(&sorted))),
};
let mut appended = sorted.clone();
appended.push(FragmentSummary {
files: vec!["b.lance".to_string()],
rows: None,
});
assert_eq!(
verdict(&spec, &state, &appended, 100),
SortVerdict::Rewritten
);
// The untouched table is still up to date: the files answer that
// without needing to count anything.
assert_eq!(verdict(&spec, &state, &sorted, 100), SortVerdict::UpToDate);
}
// Distinct columns whose names differ only in case are two columns in an
// Arrow schema, not one named twice.
#[test]
fn case_distinguishes_two_fields() {
let spec = SortSpec::parse("id asc, ID desc").unwrap().unwrap();
assert_eq!(spec.fields.len(), 2);
assert!(SortSpec::parse("id asc, id desc").is_err());
}
#[test]
fn the_marker_round_trips_through_a_config_map() {
let spec = SortSpec::parse("a asc").unwrap().unwrap();
let written = frags(&[("one.lance", 20), ("two.lance", 22)]);
let recorded = SortState::record(&spec, &written, 42);
let state = SortState::from_config(&recorded);
assert_eq!(state.rows, Some(42));
assert_eq!(state.fragments, Some(2));
assert_eq!(verdict(&spec, &state, &written, 1), SortVerdict::UpToDate);
}
}