seaweedfs/seaweed-vector/crates/server/src/main.rs

use std::path::PathBuf;

use anyhow::{Context, Result};
use clap::Parser;
use tokio_stream::wrappers::TcpListenerStream;
use tonic::transport::Server;
use weed_lance_worker::catalog::{NamespaceClient, SigningKey};
use weed_lance_worker::dataset::FallbackOptions;
use weed_vector_server::pb::vector::vector_service_server::VectorServiceServer;
use weed_vector_server::read_cache::ReadCache;
use weed_vector_server::security::Security;
use weed_vector_server::service::VectorServer;
use weed_vector_server::tls::{self, ServerTls};

/// The VectorService the S3 gateway fans vector calls out to.
#[derive(Parser)]
struct Args {
    /// Lance Namespace base URL: the S3 gateway's Lance port (-s3.port.lance).
    #[arg(
        long,
        env = "WEED_VECTOR_NAMESPACE",
        default_value = "http://localhost:9101"
    )]
    namespace: String,
    /// gRPC listen address.
    #[arg(long, env = "WEED_VECTOR_LISTEN", default_value = "127.0.0.1:8090")]
    listen: String,
    /// object_store option fallback for deployments that vend no credentials:
    /// key=value, repeatable.
    #[arg(long = "storage-option", alias = "storage.option", value_parser = parse_option)]
    storage_options: Vec<(String, String)>,
    /// Comma-separated master addresses to register this server with; the
    /// gateway discovers it via ListClusterNodes(client_type=vector).
    #[arg(long, alias = "master", env = "WEED_VECTOR_MASTERS")]
    masters: Option<String>,
    /// Memory for caching reads of write-once table files, in MiB; 0 disables.
    #[arg(long, env = "WEED_VECTOR_READ_CACHE_MB", default_value_t = 512)]
    read_cache_mb: u64,
    /// Partitions each query probes in a table's vector index. Default: a
    /// tenth of the index's partitions, at least 8; more under a filter.
    #[arg(long, env = "WEED_VECTOR_NPROBES")]
    nprobes: Option<std::num::NonZeroUsize>,
    /// The `jwt.filer_signing.key` of the cluster. The Lance namespace vends
    /// credentials that write index data only to a caller holding it.
    #[arg(long, env = "WEED_FILER_SIGNING_KEY")]
    filer_signing_key: Option<String>,
    /// The cluster's security.toml, for gRPC TLS: this server's identity is
    /// [grpc.vector], the masters are dialed with [grpc.client]. By default
    /// the file weed processes find in -config_dir, ., ~/.seaweedfs and
    /// /etc/seaweedfs.
    #[arg(long, env = "WEED_VECTOR_SECURITY_TOML")]
    security_toml: Option<PathBuf>,
    /// The directory weed's -config_dir names, searched for security.toml
    /// first.
    #[arg(long, alias = "config_dir", env = "WEED_VECTOR_CONFIG_DIR")]
    config_dir: Option<PathBuf>,
}

fn parse_option(s: &str) -> Result<(String, String), String> {
    s.split_once('=')
        .map(|(k, v)| (k.to_string(), v.to_string()))
        .ok_or_else(|| "expected key=value".to_string())
}

// weed tools take Go-style single-dash long flags; normalize_args converts
// any -flag to --flag — the same normalization weed-volume's config.rs does —
// so -master and --masters parse the same. The masters field carries the
// singular alias clap-side.
fn normalize_args(args: Vec<String>) -> Vec<String> {
    let mut args = args;
    let mut i = 1; // skip args[0] (binary name)
    while i < args.len() {
        let arg = &args[i];
        if arg == "--" {
            break;
        }
        if arg.starts_with("--") || !arg.starts_with('-') {
            i += 1;
            continue;
        }
        let without_dash = &arg[1..];
        if without_dash.len() > 1 && !without_dash.starts_with(|c: char| c.is_ascii_digit()) {
            args[i] = format!("-{}", arg);
        }
        i += 1;
    }
    args
}

/// A wildcard listen address is not reachable for peers; derive the outbound
/// IP toward the first master instead — the address the master sees anyway.
fn advertised_addr(listen: &str, masters: &[String]) -> String {
    let Some((host, _)) = listen.rsplit_once(':') else {
        return listen.to_string();
    };
    if !matches!(host, "0.0.0.0" | "::" | "" | "[::]") {
        return listen.to_string();
    }
    let port = listen.rsplit_once(':').map(|(_, p)| p).unwrap_or_default();
    if let Some(master) = masters.first() {
        // Masters may arrive as host:http.grpc or scheme://host:http —
        // the UDP probe only needs a dialable host:port.
        let probe = master.rsplit_once("://").map(|(_, r)| r).unwrap_or(master);
        let probe = match probe.rsplit_once(':') {
            Some((h, p)) => match p.split_once('.') {
                Some((_, g)) => format!("{h}:{g}"),
                None => probe.to_string(),
            },
            None => probe.to_string(),
        };
        if let Ok(sock) = std::net::UdpSocket::bind((host, 0u16)) {
            if sock.connect(&probe).is_ok() {
                if let Ok(local) = sock.local_addr() {
                    return format!("{}:{}", local.ip(), port);
                }
            }
        }
    }
    listen.to_string()
}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt::init();
    let args = Args::parse_from(normalize_args(std::env::args().collect()));
    weed_lance_worker::tls::install_default_crypto_provider();
    let security = Security::load(args.security_toml.as_deref(), args.config_dir.as_deref())?;
    if let Some(path) = security.path() {
        tracing::info!("security settings from {}", path.display());
    }
    let master_tls = match &args.masters {
        Some(_) => security.grpc_client_tls("grpc.client", "grpc.master")?,
        None => None,
    };
    let server_tls = security
        .grpc_server_tls("grpc.vector")?
        .map(ServerTls::load)
        .transpose()?;
    // Bind before registering, so the master never lists an engine that is
    // about to exit because its port is taken.
    let listener = tokio::net::TcpListener::bind(&args.listen)
        .await
        .with_context(|| format!("listen on {}", args.listen))?;
    if let Some(masters) = &args.masters {
        let masters: Vec<String> = masters.split(',').map(|s| s.trim().to_string()).collect();
        for master in &masters {
            weed_vector_server::master_client::grpc_addr(master)?;
        }
        let advertise = advertised_addr(&args.listen, &masters);
        tokio::spawn(weed_vector_server::master_client::keep_connected(
            masters,
            advertise,
            master_tls.clone(),
        ));
    }
    let fallback: FallbackOptions = args.storage_options.into_iter().collect();
    let signing_key = SigningKey::new(args.filer_signing_key);
    let client = NamespaceClient::new(args.namespace.clone())
        .with_fallback(&fallback)
        .with_signing_key(signing_key.as_ref());
    let mut service = VectorServer::new(client, fallback)
        .with_nprobes(args.nprobes.map(std::num::NonZeroUsize::get));
    if args.read_cache_mb > 0 {
        service = service.with_read_cache(ReadCache::new(args.read_cache_mb << 20));
    }
    tracing::info!(
        namespace = %args.namespace,
        listen = %args.listen,
        "weed-vector; serving gRPC {}, dialing masters {}",
        match &server_tls {
            Some(tls) if tls.files().allowed_common_names.is_empty()
                && tls.files().allowed_wildcard_domain.is_empty() =>
                "over TLS with client certificates required",
            Some(_) => "over TLS with allowed client certificates required",
            None => "in plaintext",
        },
        if master_tls.is_some() { "over TLS" } else { "in plaintext" },
    );
    let router = Server::builder()
        // Match the gateway's 1 GiB cap — put batches over the tonic 4 MiB
        // default otherwise get rejected before they reach the handler.
        .add_service(
            VectorServiceServer::new(service)
                .max_decoding_message_size(1 << 30)
                .max_encoding_message_size(1 << 30),
        );
    match server_tls {
        Some(tls) => {
            tokio::spawn(tls.clone().refresh(tls::refresh_interval()));
            router.serve_with_incoming(tls.incoming(listener)).await?;
        }
        None => {
            router
                .serve_with_incoming(TcpListenerStream::new(listener))
                .await?
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn nprobes_must_be_positive() {
        assert!(Args::try_parse_from(["weed-vector", "--nprobes", "0"]).is_err());
        let args = Args::try_parse_from(["weed-vector", "--nprobes", "12"]).unwrap();
        assert_eq!(args.nprobes.map(std::num::NonZeroUsize::get), Some(12));
    }
}