tidaldb/tidal-server/src/main.rs

use std::{net::SocketAddr, path::PathBuf, sync::Arc};

use clap::{Args, Parser, Subcommand};
use tidal_server::{
    cluster::{
        ClusterMode, ClusterState, RegionClusterState, build_cluster_router, build_region_router,
        ensure_experimental_enabled, load_topology,
    },
    config::{CONFIG_DIR_SCHEMA_FILE, CONFIG_DIR_TOPOLOGY_FILE, load_schema, resolve_config_path},
    error::{Result, ServerError},
    router::build_router,
    state::ServerState,
};
use tidaldb::TidalDb;

#[derive(Parser)]
#[command(version, about = "HTTP wrapper for tidalDB")]
struct Cli {
    #[command(subcommand)]
    mode: Command,
}

#[derive(Subcommand)]
enum Command {
    #[command(about = "Run a single-node server wrapping one tidalDB instance")]
    Standalone(StandaloneArgs),
    #[command(about = "Run a multi-region cluster behind a single HTTP surface")]
    Cluster(ClusterArgs),
}

#[derive(Args)]
struct ClusterArgs {
    #[arg(long, default_value = "127.0.0.1:9500", env = "PORT", value_parser = parse_listen_addr)]
    listen: String,
    #[arg(long)]
    schema: Option<PathBuf>,
    #[arg(long)]
    topology: Option<PathBuf>,
    /// Directory holding `default-schema.yaml` / `default-cluster.yaml`, used
    /// when `--schema` / `--topology` are not given. Wired from the container's
    /// `TIDAL_CONFIG` env var.
    #[arg(long, env = "TIDAL_CONFIG")]
    config_dir: Option<PathBuf>,
    /// Opt in to the EXPERIMENTAL cluster mode (single- or multi-process).
    ///
    /// Single-process cluster mode replicates between regions over the real
    /// tidal-net gRPC transport (loopback), but every region runs inside THIS
    /// one process — no host/process isolation. Multi-process mode (`--region`)
    /// gives real process isolation but still lacks quorum-ack writes and
    /// automatic failure detection. Either way it refuses to start without this
    /// flag (or the `TIDAL_ALLOW_EXPERIMENTAL_CLUSTER=1` env var).
    #[arg(long)]
    experimental_cluster: bool,
    /// Run ONLY this region in this process (multi-process cluster mode).
    ///
    /// Peers are reached via the topology's per-region `grpc_addr`/`http_addr`.
    /// Omitted ⇒ the existing single-process cluster runs unchanged.
    #[arg(long, env = "TIDAL_REGION")]
    region: Option<String>,
    /// Data directory for this region's `TidalDb` (multi-process mode).
    /// Omitted ⇒ ephemeral.
    #[arg(long)]
    data_dir: Option<PathBuf>,
}

#[derive(Args)]
struct StandaloneArgs {
    #[arg(long, default_value = "127.0.0.1:9400", env = "PORT", value_parser = parse_listen_addr)]
    listen: String,
    #[arg(long)]
    schema: Option<PathBuf>,
    /// Directory holding `default-schema.yaml`, used when `--schema` is not
    /// given. Wired from the container's `TIDAL_CONFIG` env var.
    #[arg(long, env = "TIDAL_CONFIG")]
    config_dir: Option<PathBuf>,
    #[arg(long)]
    data_dir: Option<PathBuf>,
    #[arg(
        long,
        help = "Bind address for Prometheus /metrics endpoint (e.g. 127.0.0.1:9091)"
    )]
    metrics: Option<String>,
}

#[tokio::main]
async fn main() {
    if let Err(err) = run().await {
        eprintln!("error: {err}");
        std::process::exit(1);
    }
}

async fn run() -> Result<()> {
    let cli = Cli::parse();
    init_tracing();

    match cli.mode {
        Command::Standalone(args) => run_standalone(args).await,
        Command::Cluster(args) => run_cluster(args).await,
    }
}

/// Parse a listen address from either a full `host:port` string or a bare port number.
/// When `PORT=8080` is set, clap passes `"8080"` — this normalises it to `0.0.0.0:8080`.
fn parse_listen_addr(s: &str) -> std::result::Result<String, String> {
    if s.contains(':') {
        s.parse::<SocketAddr>()
            .map(|_| s.to_string())
            .map_err(|e| format!("invalid address '{s}': {e}"))
    } else {
        s.parse::<u16>()
            .map(|port| format!("0.0.0.0:{port}"))
            .map_err(|_| format!("expected host:port or port number, got '{s}'"))
    }
}

fn init_tracing() {
    let env_filter = std::env::var("TIDAL_SERVER_LOG").unwrap_or_else(|_| "info".into());
    let _ = tracing_subscriber::fmt()
        .with_env_filter(env_filter)
        .try_init();
}

async fn run_standalone(args: StandaloneArgs) -> Result<()> {
    let schema_path = resolve_config_path(
        args.schema.as_deref(),
        args.config_dir.as_deref(),
        CONFIG_DIR_SCHEMA_FILE,
    )?;
    let (schema, profiles) = load_schema(schema_path.as_deref())?;

    let mut builder = TidalDb::builder()
        .with_schema(schema.clone())
        .with_profiles(profiles);
    if let Some(dir) = args.data_dir {
        builder = builder.with_data_dir(dir);
    } else {
        builder = builder.ephemeral();
    }
    if let Some(ref addr) = args.metrics {
        builder = builder.enable_metrics(addr);
    }

    let db = builder.open()?;
    if let Some(addr) = db.metrics_addr() {
        tracing::info!("metrics endpoint listening on http://{addr}/metrics");
    }
    let state = ServerState::new(db);

    let api_key = read_api_key();
    serve_state(state, &args.listen, api_key, build_router).await
}

async fn run_cluster(args: ClusterArgs) -> Result<()> {
    // `--region` selects true multi-process mode (this process owns ONE region
    // and peers with siblings over real gRPC); without it the existing
    // single-process cluster runs byte-for-byte unchanged.
    match args.region.clone() {
        Some(region) => run_region_cluster(args, region).await,
        None => run_single_process_cluster(args).await,
    }
}

async fn run_single_process_cluster(args: ClusterArgs) -> Result<()> {
    // Honest gate: single-process cluster mode replicates between regions over
    // the real tidal-net gRPC transport, but every region runs inside this one
    // process (no host/process isolation), so it is not production HA. Refuse to
    // start (and emit a loud WARN when permitted) so no operator mistakes it for
    // production HA.
    ensure_experimental_enabled(args.experimental_cluster, ClusterMode::SingleProcess)?;

    let schema_path = resolve_config_path(
        args.schema.as_deref(),
        args.config_dir.as_deref(),
        CONFIG_DIR_SCHEMA_FILE,
    )?;
    let topology_path = resolve_config_path(
        args.topology.as_deref(),
        args.config_dir.as_deref(),
        CONFIG_DIR_TOPOLOGY_FILE,
    )?;

    let (schema, profiles) = load_schema(schema_path.as_deref())?;
    let topology = load_topology(topology_path.as_deref())?;

    tracing::info!(
        regions = topology.regions.len(),
        leader = %topology.leader,
        "building cluster"
    );

    // `ClusterState::new` starts the follower gRPC servers via
    // `GrpcTransport::new`, which blocks on its own tokio runtime and must not
    // run inside this reactor. Build it on a dedicated thread; the constructed
    // state (and its embedded runtimes) then lives for the server's lifetime.
    let state = std::thread::Builder::new()
        .name("cluster-build".into())
        .spawn(move || ClusterState::new(&topology, schema, profiles))
        .map_err(|e| ServerError::Cluster(format!("spawn cluster builder thread: {e}")))?
        .join()
        .map_err(|_| ServerError::Cluster("cluster builder thread panicked".into()))??;

    let api_key = read_api_key();
    serve_state(state, &args.listen, api_key, build_cluster_router).await
}

async fn run_region_cluster(args: ClusterArgs, region: String) -> Result<()> {
    // Multi-process mode is experimental too: real process isolation, but no
    // quorum-ack writes or automatic failure detection yet. The gate's WARN text
    // differentiates the two modes honestly.
    ensure_experimental_enabled(args.experimental_cluster, ClusterMode::MultiProcess)?;

    let schema_path = resolve_config_path(
        args.schema.as_deref(),
        args.config_dir.as_deref(),
        CONFIG_DIR_SCHEMA_FILE,
    )?;
    let topology_path = resolve_config_path(
        args.topology.as_deref(),
        args.config_dir.as_deref(),
        CONFIG_DIR_TOPOLOGY_FILE,
    )?;

    let (schema, profiles) = load_schema(schema_path.as_deref())?;
    let topology = load_topology(topology_path.as_deref())?;

    // The operator/test surface for clock-skew injection: parse i64 strictly so
    // an invalid value is a hard startup error, not a silent 0.
    let hlc_offset_ms = match std::env::var("TIDAL_HLC_SKEW_MS") {
        Ok(raw) => raw.trim().parse::<i64>().map_err(|e| {
            ServerError::SchemaConfig(format!("invalid TIDAL_HLC_SKEW_MS '{raw}': {e}"))
        })?,
        Err(_) => 0,
    };

    tracing::info!(
        region = %region,
        leader = %topology.leader,
        hlc_offset_ms,
        "building region cluster node (multi-process)"
    );

    let data_dir = args.data_dir.clone();
    // Read the bearer key up front: the region node needs its OWN copy to
    // authenticate inter-sibling POSTs (the heal item/embedding backfill). The
    // internal marker is a trust signal, not an auth bypass, so an unauthenticated
    // backfill 401s when TIDAL_API_KEY is set. Same key serves the router below.
    let api_key = read_api_key();
    let node_key = api_key.clone();
    // `RegionClusterState::new` builds the GrpcTransport via `GrpcTransport::new`,
    // which blocks on its own tokio runtime — must run off this reactor. Build it
    // on a dedicated thread, exactly like the single-process path.
    let state = std::thread::Builder::new()
        .name("region-build".into())
        .spawn(move || {
            RegionClusterState::new(
                &topology,
                &region,
                schema,
                profiles,
                data_dir,
                hlc_offset_ms,
                node_key,
            )
        })
        .map_err(|e| ServerError::Cluster(format!("spawn region builder thread: {e}")))?
        .join()
        .map_err(|_| ServerError::Cluster("region builder thread panicked".into()))??;

    serve_state(state, &args.listen, api_key, build_region_router).await
}

/// The serve/shutdown contract every server mode (standalone, single-process
/// cluster, multi-process region) satisfies, so the socket-bind → serve →
/// drain → deterministic-shutdown plumbing is written ONCE in [`serve_state`].
trait ServeState: Send + Sync + 'static {
    /// Human label for the post-serve log lines.
    const WHAT: &'static str;
    /// Flip `/health` to not-ready BEFORE axum starts draining.
    fn set_shutting_down(&self);
    /// Final deterministic shutdown (checkpoint + WAL fsync + thread join),
    /// run once `serve_state` has reclaimed sole ownership.
    fn shutdown_owned(self);
}

impl ServeState for ServerState {
    const WHAT: &'static str = "standalone server";
    fn set_shutting_down(&self) {
        // `Self::` resolves to the INHERENT method (inherent impls outrank
        // trait impls in associated-fn resolution), not this trait fn.
        Self::set_shutting_down(self);
    }
    fn shutdown_owned(self) {
        // Dropping the sole-owner `ServerState` drops its `Arc<TidalDb>`; that
        // runs `TidalDb::Drop`, which attempts every final durable flush and
        // logs the first failure at error level (SHUTDOWN-2). The drop is
        // deterministic here, not deferred.
        drop(self);
        tracing::info!("standalone shutdown: database closed (checkpoint + WAL fsync)");
    }
}

impl ServeState for ClusterState {
    const WHAT: &'static str = "cluster";
    fn set_shutting_down(&self) {
        Self::set_shutting_down(self); // the inherent method, as above
    }
    fn shutdown_owned(mut self) {
        self.shutdown();
    }
}

impl ServeState for RegionClusterState {
    const WHAT: &'static str = "region";
    fn set_shutting_down(&self) {
        Self::set_shutting_down(self); // the inherent method, as above
    }
    fn shutdown_owned(mut self) {
        self.shutdown();
    }
}

// The `state` binding deliberately lives until after `axum::serve` returns so
// the post-serve `Arc::try_unwrap` can reclaim sole ownership and run the
// deterministic shutdown (checkpoint + WAL fsync + receiver join). Dropping it
// "early" (clippy's suggestion) would drop the node BEFORE the server stops
// accepting requests — the exact ordering bug this ordering exists to avoid.
#[allow(clippy::significant_drop_tightening)]
async fn serve_state<S: ServeState>(
    state: S,
    addr: &str,
    api_key: Option<Arc<str>>,
    build_router: impl FnOnce(Arc<S>, Option<Arc<str>>) -> axum::Router,
) -> Result<()> {
    let socket: SocketAddr = addr
        .parse()
        .map_err(|e| ServerError::BadRequest(format!("invalid addr: {e}")))?;

    let listener = tokio::net::TcpListener::bind(socket).await?;
    let actual = listener.local_addr()?;
    tracing::info!("listening on http://{actual}");

    let state = Arc::new(state);
    let shutdown_state = state.clone();

    axum::serve(listener, build_router(state, api_key))
        .with_graceful_shutdown(shutdown_signal(shutdown_state.clone()))
        .await?;

    // axum::serve has returned, so the router (and every `Arc<S>` it held) is
    // dropped. We should now be the sole owner; reclaim ownership and run the
    // final durable shutdown HERE — deterministically, before the process
    // exits — instead of letting an implicit drop fire the `Drop` backstop at
    // an unspecified point. If an Arc unexpectedly lingers we cannot observe
    // the final flush from this stack frame, so we log that we are falling
    // back to `Drop` (which still runs the same shutdown and logs any flush
    // failure at error level) rather than silently exiting 0 on a
    // possibly-failed flush.
    match Arc::try_unwrap(shutdown_state) {
        Ok(owned) => owned.shutdown_owned(),
        Err(arc) => {
            tracing::warn!(
                strong = Arc::strong_count(&arc),
                state = S::WHAT,
                "state still shared after serve returned; relying on Drop for shutdown"
            );
        }
    }
    Ok(())
}

async fn shutdown_signal<S: ServeState>(state: Arc<S>) {
    // SIGTERM is Unix-only; on other platforms we fall back to ctrl-c alone.
    #[cfg(unix)]
    let sigterm = async {
        use tokio::signal::unix::{SignalKind, signal};
        match signal(SignalKind::terminate()) {
            Ok(mut stream) => {
                stream.recv().await;
            }
            Err(err) => {
                tracing::warn!("SIGTERM handler registration failed: {err}");
                // If we cannot register SIGTERM, park this branch so the
                // select falls back to ctrl_c exclusively.
                std::future::pending::<()>().await;
            }
        }
    };
    #[cfg(not(unix))]
    let sigterm = std::future::pending::<()>();

    tokio::select! {
        result = tokio::signal::ctrl_c() => {
            if let Err(err) = result {
                tracing::warn!("ctrl-c handler error: {err}");
            }
        }
        () = sigterm => {}
    }

    // Flip readiness BEFORE axum starts draining.
    state.set_shutting_down();
    tracing::info!("shutdown signal received; readiness not-ready, draining in-flight requests");
}

/// Read the API key from the environment.
///
/// If `TIDAL_API_KEY` is not set, all requests are accepted without
/// authentication. This is appropriate for local development but should
/// never be used in production. A startup warning is emitted.
fn read_api_key() -> Option<Arc<str>> {
    match std::env::var("TIDAL_API_KEY") {
        Ok(key) if !key.is_empty() => Some(Arc::from(key.as_str())),
        _ => {
            tracing::warn!(
                "TIDAL_API_KEY is not set — all endpoints are unauthenticated. \
                 Set this variable before exposing the server to any network."
            );
            None
        }
    }
}