129 lines
3.5 KiB
Rust
129 lines
3.5 KiB
Rust
mod config;
|
|
mod db;
|
|
mod error;
|
|
mod hasher;
|
|
mod health;
|
|
mod server;
|
|
mod volume;
|
|
|
|
use clap::{Parser, Subcommand};
|
|
use std::collections::HashSet;
|
|
use std::path::PathBuf;
|
|
use std::sync::Arc;
|
|
use tokio::sync::RwLock;
|
|
|
|
#[derive(Parser)]
|
|
#[command(name = "mkv", about = "Distributed key-value store")]
|
|
struct Cli {
|
|
#[arg(short, long, default_value = "config.toml")]
|
|
config: PathBuf,
|
|
|
|
#[command(subcommand)]
|
|
command: Commands,
|
|
}
|
|
|
|
#[derive(Subcommand)]
|
|
enum Commands {
|
|
/// Start the index server
|
|
Serve,
|
|
/// Rebuild SQLite index from volume servers
|
|
Rebuild,
|
|
/// Rebalance data after adding/removing volumes
|
|
Rebalance {
|
|
#[arg(long)]
|
|
dry_run: bool,
|
|
},
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
tracing_subscriber::fmt::init();
|
|
|
|
let cli = Cli::parse();
|
|
let config = config::Config::load(&cli.config).unwrap_or_else(|e| {
|
|
eprintln!("Failed to load config: {e}");
|
|
std::process::exit(1);
|
|
});
|
|
|
|
match cli.command {
|
|
Commands::Serve => serve(config).await,
|
|
Commands::Rebuild => {
|
|
eprintln!("rebuild not yet implemented");
|
|
std::process::exit(1);
|
|
}
|
|
Commands::Rebalance { dry_run: _ } => {
|
|
eprintln!("rebalance not yet implemented");
|
|
std::process::exit(1);
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn serve(config: config::Config) {
|
|
let db_path = &config.database.path;
|
|
|
|
// Ensure parent directory exists
|
|
if let Some(parent) = std::path::Path::new(db_path).parent() {
|
|
std::fs::create_dir_all(parent).unwrap_or_else(|e| {
|
|
eprintln!("Failed to create database directory: {e}");
|
|
std::process::exit(1);
|
|
});
|
|
}
|
|
|
|
let (writer, ready_rx) = db::spawn_writer(db_path.to_string());
|
|
ready_rx.await.expect("writer failed to initialize");
|
|
|
|
let num_readers = std::thread::available_parallelism()
|
|
.map(|n| n.get())
|
|
.unwrap_or(4);
|
|
let reads = db::ReadPool::new(db_path, num_readers);
|
|
|
|
let volume_urls = config.volume_urls();
|
|
let ring = Arc::new(RwLock::new(hasher::Ring::new(
|
|
&volume_urls,
|
|
config.server.virtual_nodes,
|
|
)));
|
|
|
|
let volume_client = volume::VolumeClient::new();
|
|
|
|
// Start with all volumes assumed healthy, health checker will update
|
|
let healthy_volumes: health::HealthyVolumes =
|
|
Arc::new(RwLock::new(HashSet::from_iter(volume_urls.clone())));
|
|
|
|
health::spawn_health_checker(
|
|
volume_client.clone(),
|
|
volume_urls.clone(),
|
|
healthy_volumes.clone(),
|
|
);
|
|
|
|
let port = config.server.port;
|
|
|
|
let state = server::AppState {
|
|
writer,
|
|
reads,
|
|
ring,
|
|
volume_client,
|
|
healthy_volumes,
|
|
config: Arc::new(config),
|
|
};
|
|
|
|
tracing::info!("Starting mkv server on port {port}");
|
|
tracing::info!(" Readers: {num_readers}");
|
|
tracing::info!(" Volumes: {volume_urls:?}");
|
|
tracing::info!(" Replication factor: {}", state.config.server.replication_factor);
|
|
|
|
let app = axum::Router::new()
|
|
.route("/", axum::routing::get(server::list_keys))
|
|
.route(
|
|
"/{*key}",
|
|
axum::routing::get(server::get_key)
|
|
.put(server::put_key)
|
|
.delete(server::delete_key)
|
|
.head(server::head_key),
|
|
)
|
|
.with_state(state);
|
|
|
|
let addr = format!("0.0.0.0:{port}");
|
|
let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
|
|
tracing::info!("Listening on {addr}");
|
|
axum::serve(listener, app).await.unwrap();
|
|
}
|