seaweedfs/seaweed-worker/crates/lance/src/main.rs

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:12535", 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,

    /// Comma-separated master addresses: the worker asks them for the filers
    /// it keeps tables through.
    #[arg(long, alias = "master", env = "WEED_MASTERS")]
    masters: String,
    /// The filer group whose filers the worker uses.
    #[arg(
        long,
        alias = "filer.group",
        env = "WEED_FILER_GROUP",
        default_value = ""
    )]
    filer_group: String,
    /// Encrypt every chunk written for the volume servers, as `weed s3
    /// -encryptVolumeData` does.
    #[arg(
        long = "encryptVolumeData",
        alias = "encrypt-volume-data",
        env = "WEED_ENCRYPT_VOLUME_DATA"
    )]
    encrypt_volume_data: bool,
    /// The cluster's security.toml; by default the file weed processes find
    /// in -config_dir, ., ~/.seaweedfs and /etc/seaweedfs.
    #[arg(long, env = "WEED_WORKER_SECURITY_TOML")]
    security_toml: Option<std::path::PathBuf>,

    /// 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>,

    /// The cluster's `jwt.filer_signing.key`. The Lance namespace knows the
    /// worker by a token derived from it, and admins that authenticate the
    /// worker stream with a signing key require it. When not given, read from
    /// where the Go worker reads it: WEED_JWT_FILER_SIGNING_KEY, then
    /// security.toml.
    #[arg(long, env = "WEED_FILER_SIGNING_KEY")]
    filer_signing_key: Option<String>,

    /// Lifetime of the token minted from the signing key, mirroring
    /// `jwt.filer_signing.expires_after_seconds`. Zero mints a token with no
    /// expiry. When not given, the expiry configured beside a discovered key,
    /// else zero.
    #[arg(long)]
    filer_signing_key_expires_after_seconds: Option<u64>,

    /// 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()?;
    // One security.toml for everything: the admin stream's and the catalog's
    // signing key, and the cluster connection.
    let security =
        weed_lance_worker::security::Security::load(args.security_toml.as_deref(), None)?;
    let filer_signing = match args.filer_signing_key.filter(|key| !key.is_empty()) {
        Some(key) => Some((key, 0)),
        None => security
            .get("jwt.filer_signing.key")
            .filter(|key| !key.is_empty())
            .map(|key| {
                if let Some(path) = security.path() {
                    tracing::info!(source = %path.display(), "using jwt.filer_signing.key");
                }
                let expires = security
                    .get("jwt.filer_signing.expires_after_seconds")
                    .and_then(|v| v.parse().ok())
                    .unwrap_or(0);
                (key, expires)
            }),
    };
    let (filer_signing_key, configured_expiry) = filer_signing.unzip();
    let signing_key = weed_lance_worker::catalog::SigningKey::new(filer_signing_key.clone());
    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,
        signing_key: filer_signing_key,
        signing_key_expires_after_seconds: args
            .filer_signing_key_expires_after_seconds
            .or(configured_expiry)
            .unwrap_or(0),
        ..Default::default()
    };

    let cluster =
        weed_lance_worker::cluster::Cluster::new(&args.masters, &args.filer_group, &security)?;
    let storage = weed_lance_worker::dataset::Storage(Some(
        cluster.storage(&security, args.encrypt_volume_data).await?,
    ));

    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(), storage.clone())
            .with_signing_key(signing_key.clone()),
    ));
    for handler in handlers(args.namespace, storage, signing_key, 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());
    }
}