84 lines
2.3 KiB
Rust
84 lines
2.3 KiB
Rust
use clap::{Parser, Subcommand};
|
|
|
|
#[derive(Parser)]
|
|
#[command(name = "mkv", about = "Distributed key-value store")]
|
|
struct Cli {
|
|
#[arg(short, long, default_value = "/tmp/mkv/index.db")]
|
|
db: String,
|
|
|
|
#[arg(short, long, required = true, value_delimiter = ',')]
|
|
volumes: Vec<String>,
|
|
|
|
#[arg(short, long, default_value_t = 2)]
|
|
replicas: usize,
|
|
|
|
#[command(subcommand)]
|
|
command: Commands,
|
|
}
|
|
|
|
#[derive(Subcommand)]
|
|
enum Commands {
|
|
/// Start the index server
|
|
Serve {
|
|
#[arg(short, long, default_value_t = 3000)]
|
|
port: u16,
|
|
},
|
|
/// Rebuild SQLite index from volume servers
|
|
Rebuild,
|
|
/// Rebalance data after adding/removing volumes
|
|
Rebalance {
|
|
#[arg(long)]
|
|
dry_run: bool,
|
|
},
|
|
}
|
|
|
|
async fn shutdown_signal() {
|
|
let ctrl_c = tokio::signal::ctrl_c();
|
|
#[cfg(unix)]
|
|
{
|
|
let mut sigterm =
|
|
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
|
|
.expect("failed to install SIGTERM handler");
|
|
tokio::select! {
|
|
_ = ctrl_c => tracing::info!("Received SIGINT, shutting down..."),
|
|
_ = sigterm.recv() => tracing::info!("Received SIGTERM, shutting down..."),
|
|
}
|
|
}
|
|
#[cfg(not(unix))]
|
|
{
|
|
ctrl_c.await.expect("failed to listen for ctrl-c");
|
|
tracing::info!("Received ctrl-c, shutting down...");
|
|
}
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
tracing_subscriber::fmt::init();
|
|
let cli = Cli::parse();
|
|
|
|
let args = mkv::Args {
|
|
db_path: cli.db,
|
|
volumes: cli.volumes,
|
|
replicas: cli.replicas,
|
|
};
|
|
|
|
match cli.command {
|
|
Commands::Serve { port } => {
|
|
let app = mkv::build_app(&args);
|
|
let addr = format!("0.0.0.0:{port}");
|
|
tracing::info!("Listening on {addr}");
|
|
let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
|
|
axum::serve(listener, app)
|
|
.with_graceful_shutdown(shutdown_signal())
|
|
.await
|
|
.unwrap();
|
|
tracing::info!("Shutdown complete");
|
|
}
|
|
Commands::Rebuild => {
|
|
mkv::rebuild::run(&args).await;
|
|
}
|
|
Commands::Rebalance { dry_run } => {
|
|
mkv::rebalance::run(&args, dry_run).await;
|
|
}
|
|
}
|
|
}
|