mkv

Unnamed repository; edit this file 'description' to name the repository.
Log | Files | Refs | README

main.rs (2652B)

      1 use std::time::Duration;
      2 
      3 use clap::{Parser, Subcommand};
      4 
      5 #[derive(Parser)]
      6 #[command(name = "mkv", about = "Distributed key-value store")]
      7 struct Cli {
      8     #[arg(short, long, env = "MKV_DB", default_value = "/tmp/mkv/index.db")]
      9     db: String,
     10 
     11     #[arg(
     12         short,
     13         long,
     14         env = "MKV_VOLUMES",
     15         required = true,
     16         value_delimiter = ','
     17     )]
     18     volumes: Vec<String>,
     19 
     20     #[arg(short, long, env = "MKV_REPLICAS", default_value_t = 2)]
     21     replicas: usize,
     22 
     23     /// Timeout for volume health checks (in milliseconds)
     24     #[arg(long, env = "MKV_VOLTIMEOUT", default_value_t = 1000)]
     25     voltimeout: u64,
     26 
     27     #[command(subcommand)]
     28     command: Commands,
     29 }
     30 
     31 #[derive(Subcommand)]
     32 enum Commands {
     33     /// Start the index server
     34     Serve {
     35         #[arg(short, long, env = "MKV_PORT", default_value_t = 3000)]
     36         port: u16,
     37     },
     38     /// Rebuild SQLite index from volume servers
     39     Rebuild,
     40     /// Rebalance data after adding/removing volumes
     41     Rebalance {
     42         #[arg(long)]
     43         dry_run: bool,
     44     },
     45 }
     46 
     47 async fn shutdown_signal() {
     48     let ctrl_c = tokio::signal::ctrl_c();
     49     #[cfg(unix)]
     50     {
     51         let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
     52             .expect("failed to install SIGTERM handler");
     53         tokio::select! {
     54             _ = ctrl_c => tracing::info!("Received SIGINT, shutting down..."),
     55             _ = sigterm.recv() => tracing::info!("Received SIGTERM, shutting down..."),
     56         }
     57     }
     58     #[cfg(not(unix))]
     59     {
     60         ctrl_c.await.expect("failed to listen for ctrl-c");
     61         tracing::info!("Received ctrl-c, shutting down...");
     62     }
     63 }
     64 
     65 #[tokio::main]
     66 async fn main() {
     67     tracing_subscriber::fmt::init();
     68     let cli = Cli::parse();
     69 
     70     let args = mkv::Args {
     71         db_path: cli.db,
     72         volumes: cli.volumes,
     73         replicas: cli.replicas,
     74         voltimeout: Duration::from_millis(cli.voltimeout),
     75     };
     76 
     77     match cli.command {
     78         Commands::Serve { port } => {
     79             let app = mkv::build_app(&args);
     80             let addr = format!("0.0.0.0:{port}");
     81             tracing::info!("Listening on {addr}");
     82             let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
     83             axum::serve(listener, app)
     84                 .with_graceful_shutdown(shutdown_signal())
     85                 .await
     86                 .unwrap();
     87             tracing::info!("Shutdown complete");
     88         }
     89         Commands::Rebuild => {
     90             mkv::rebuild::run(&args).await;
     91         }
     92         Commands::Rebalance { dry_run } => {
     93             mkv::rebalance::run(&args, dry_run).await;
     94         }
     95     }
     96 }