tidaldb/tidal-server/src/main.rs

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

use clap::{Args, Parser, Subcommand};
use tidal_server::{
    cluster::{ClusterState, build_cluster_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 single-process cluster.
    ///
    /// Cluster mode replicates between regions over the real tidal-net gRPC
    /// transport (loopback), but every region runs inside THIS one process —
    /// there is no host/process isolation, so it is not production HA. It
    /// refuses to start without this flag (or the
    /// `TIDAL_ALLOW_EXPERIMENTAL_CLUSTER=1` env var).
    #[arg(long)]
    experimental_cluster: bool,
}

#[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, &args.listen, api_key).await
}

async fn run_cluster(args: ClusterArgs) -> Result<()> {
    // Honest gate: cluster mode replicates between regions over the real
    // tidal-net gRPC transport, but every region runs inside this single
    // 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)?;

    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_cluster(state, &args.listen, api_key).await
}

async fn serve_cluster(state: ClusterState, addr: &str, api_key: Option<Arc<str>>) -> 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_cluster_router(state, api_key))
        .with_graceful_shutdown(cluster_shutdown_signal(shutdown_state.clone()))
        .await?;

    // axum::serve has returned, so the router (and every `Arc<ClusterState>` it
    // held) is dropped. We should now be the sole owner; reclaim ownership and
    // shut every node down deterministically (checkpoint + WAL fsync + thread
    // join) before the process exits. If an Arc unexpectedly lingers, the
    // `Drop` backstop on `ClusterState` still runs the same shutdown.
    match Arc::try_unwrap(shutdown_state) {
        Ok(mut owned) => owned.shutdown(),
        Err(arc) => {
            tracing::warn!(
                strong = Arc::strong_count(&arc),
                "cluster state still shared after serve returned; relying on Drop for shutdown"
            );
        }
    }
    Ok(())
}

async fn cluster_shutdown_signal(state: Arc<ClusterState>) {
    #[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}");
                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 => {}
    }

    state.set_shutting_down();
    tracing::info!("shutdown signal received, 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
        }
    }
}

async fn serve(state: ServerState, addr: &str, api_key: Option<Arc<str>>) -> 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<ServerState>` it
    // held) is dropped. Deterministically reclaim sole ownership and run the
    // final durable shutdown HERE, exactly like `serve_cluster` — instead of
    // letting an implicit drop fire the `TidalDb::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) => {
            // Dropping the sole-owner `ServerState` here 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 now, not deferred.
            drop(owned);
            tracing::info!("standalone shutdown: database closed (checkpoint + WAL fsync)");
        }
        Err(arc) => {
            tracing::warn!(
                strong = Arc::strong_count(&arc),
                "server state still shared after serve returned; relying on Drop for final flush"
            );
        }
    }
    Ok(())
}

async fn shutdown_signal(state: Arc<ServerState>) {
    // 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!("readiness flipped to not-ready, draining in-flight requests");
}