#[cfg(not(target_env = "msvc"))]
#[global_allocator]
static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
#[cfg(not(target_env = "msvc"))]
#[allow(non_upper_case_globals)]
#[export_name = "_rjem_malloc_conf"]
pub static malloc_conf: &[u8; 63] =
b"background_thread:true,dirty_decay_ms:2000,muzzy_decay_ms:2000\0";
#[cfg(not(target_env = "msvc"))]
fn log_allocator_tuning() {
use tikv_jemalloc_ctl::{opt, raw};
let background = opt::background_thread::read().unwrap_or(false);
let dirty_ms = unsafe { raw::read::<isize>(b"opt.dirty_decay_ms\0") }.unwrap_or(-1);
let muzzy_ms = unsafe { raw::read::<isize>(b"opt.muzzy_decay_ms\0") }.unwrap_or(-1);
if background {
tracing::info!(
"jemalloc: background_thread=true, dirty_decay_ms={}, muzzy_decay_ms={}",
dirty_ms,
muzzy_ms
);
} else {
tracing::warn!(
"jemalloc: background_thread is OFF (dirty_decay_ms={}) — pages freed by \
threads that then go idle will stay resident. The `_rjem_malloc_conf` \
symbol in main.rs is not reaching the allocator.",
dirty_ms
);
}
}
#[cfg(target_env = "msvc")]
fn log_allocator_tuning() {}
fn parse_size(s: &str) -> Result<usize, String> {
let t = s.trim();
let upper = t.to_ascii_uppercase();
let (digits, mult): (&str, usize) = if let Some(v) = upper.strip_suffix("GB") {
(&t[..v.len()], 1 << 30)
} else if let Some(v) = upper.strip_suffix("MB") {
(&t[..v.len()], 1 << 20)
} else if let Some(v) = upper.strip_suffix("KB") {
(&t[..v.len()], 1 << 10)
} else if let Some(v) = upper.strip_suffix('G') {
(&t[..v.len()], 1 << 30)
} else if let Some(v) = upper.strip_suffix('M') {
(&t[..v.len()], 1 << 20)
} else if let Some(v) = upper.strip_suffix('K') {
(&t[..v.len()], 1 << 10)
} else {
(t, 1)
};
let n: usize = digits.trim().parse().map_err(|_| {
format!("invalid size '{s}' (expected a byte count or a value like 512MB / 2GB)")
})?;
n.checked_mul(mult)
.ok_or_else(|| format!("size '{s}' overflows a usize"))
}
fn human_size(bytes: usize) -> String {
const G: usize = 1 << 30;
const M: usize = 1 << 20;
const K: usize = 1 << 10;
match bytes {
b if b >= G && b % G == 0 => format!("{}GB", b / G),
b if b >= M && b % M == 0 => format!("{}MB", b / M),
b if b >= K && b % K == 0 => format!("{}KB", b / K),
b => format!("{b}B"),
}
}
fn log_storage_profile(base: &str, p: &solidb::storage::engine::EngineProfile) {
let budget = p
.db_write_buffer_size
.map(human_size)
.unwrap_or_else(|| "unlimited".to_string());
let open_files = if p.max_open_files < 0 {
"unlimited".to_string()
} else {
p.max_open_files.to_string()
};
tracing::info!(
"Storage profile: {base} — block_cache={}, write_buffer={}/collection, \
memtable_budget={budget}, max_open_files={open_files}, \
bounded_index_cache={}, background_jobs={}",
human_size(p.block_cache_bytes),
human_size(p.write_buffer_size),
p.cache_index_and_filter_blocks,
p.max_background_jobs,
);
if p.db_write_buffer_size.is_none() {
tracing::warn!(
"Storage: no global memtable budget — total memtable RAM scales with the \
number of write-active collections ({} each) and nothing forces an early \
flush. Set --memtable-budget on instances with many collections.",
human_size(p.write_buffer_size),
);
}
if !p.cache_index_and_filter_blocks && p.max_open_files < 0 {
tracing::warn!(
"Storage: index/filter blocks are pinned per SST and the table cache is \
unlimited, so that memory grows with the dataset and is never evicted. \
Set --bounded-index-cache or --max-open-files on large datasets."
);
}
}
use clap::{Parser, Subcommand};
use solidb::server::multiplex::{ChannelListener, PeekedStream};
use solidb::{cluster::ClusterConfig, create_router, scripting::ScriptStats, StorageEngine};
use std::sync::Arc;
use sysinfo::{Pid, System};
use tokio::io::AsyncReadExt;
use tokio::sync::mpsc;
use tokio::time::Duration;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
#[derive(Parser, Debug)]
#[command(name = "solidb", version)]
#[command(about = "SolidDB - A high-performance document database", long_about = None)]
struct Args {
#[command(subcommand)]
command: Option<Command>,
#[arg(short, long, default_value_t = 6745)]
port: u16,
#[arg(long)]
host: Option<String>,
#[arg(long)]
node_id: Option<String>,
#[arg(long)]
advertise: Option<String>,
#[arg(long = "peer")]
peers: Vec<String>,
#[arg(long)]
replication_port: Option<u16>,
#[arg(long, default_value = "./data")]
data_dir: String,
#[arg(short = 'd', long)]
daemon: bool,
#[arg(long, default_value = "./solidb.pid")]
pid_file: String,
#[arg(long, default_value = "./solidb.log")]
log_file: String,
#[arg(long)]
keyfile: Option<String>,
#[arg(long)]
otlp_endpoint: Option<String>,
#[arg(long)]
no_sync_log: bool,
#[arg(long)]
dev: bool,
#[arg(long)]
no_lua: bool,
#[arg(long)]
tls_cert: Option<String>,
#[arg(long)]
tls_key: Option<String>,
#[arg(long, value_parser = parse_size, env = "SOLIDB_BLOCK_CACHE")]
block_cache: Option<usize>,
#[arg(long, value_parser = parse_size, env = "SOLIDB_MEMTABLE_BUDGET")]
memtable_budget: Option<usize>,
#[arg(long, value_parser = parse_size, env = "SOLIDB_WRITE_BUFFER_SIZE")]
write_buffer_size: Option<usize>,
#[arg(long, env = "SOLIDB_MAX_OPEN_FILES")]
max_open_files: Option<i32>,
#[arg(
long,
num_args = 0..=1,
default_missing_value = "true",
env = "SOLIDB_BOUNDED_INDEX_CACHE"
)]
bounded_index_cache: Option<bool>,
#[arg(long, env = "SOLIDB_MAX_BACKGROUND_JOBS")]
max_background_jobs: Option<i32>,
}
#[derive(Subcommand, Debug)]
enum Command {
Scripts(solidb::cli::scripts::ScriptsArgs),
Tui(solidb::cli::tui::TuiArgs),
Update,
}
fn main() -> anyhow::Result<()> {
let _ = jsonwebtoken::crypto::aws_lc::DEFAULT_PROVIDER.install_default();
let _ = dotenvy::dotenv();
let args = Args::parse();
if args.no_lua {
std::env::set_var("SOLIDB_NO_LUA", "1");
}
if let Some(command) = args.command {
return match command {
Command::Scripts(scripts_args) => solidb::cli::scripts::execute(scripts_args),
Command::Tui(tui_args) => solidb::cli::tui::execute(tui_args),
Command::Update => solidb::cli::update::execute(),
};
}
#[cfg(unix)]
if args.daemon {
use solidb::daemon::Daemonize;
use std::fs::File;
use std::path::Path;
if Path::new(&args.pid_file).exists() {
match std::fs::read_to_string(&args.pid_file) {
Ok(pid_str) => {
if let Ok(pid) = pid_str.trim().parse::<i32>() {
let mut sys = System::new_all();
sys.refresh_all();
let sys_pid = Pid::from(pid as usize);
if let Some(proc) = sys.process(sys_pid) {
let proc_name = proc.name().to_string_lossy();
if proc_name != "solidb" {
eprintln!("SECURITY ERROR: Process with PID {} is named '{}', not 'solidb'. Refusing to kill potential mismatch.", pid, proc_name);
return Ok(());
}
}
eprintln!("Found existing server with PID {}. Stopping it...", pid);
unsafe {
libc::kill(pid, libc::SIGTERM);
}
for i in 0..50 {
std::thread::sleep(std::time::Duration::from_millis(100));
let still_running = unsafe { libc::kill(pid, 0) == 0 };
if !still_running {
eprintln!("Previous server stopped successfully.");
break;
}
if i == 30 {
eprintln!("Process didn't stop gracefully, forcing shutdown...");
unsafe {
libc::kill(pid, libc::SIGKILL);
}
}
}
let _ = std::fs::remove_file(&args.pid_file);
}
}
Err(e) => {
eprintln!("Warning: Could not read PID file: {}", e);
}
}
}
let stdout = File::create(&args.log_file)?;
let stderr = File::create(&args.log_file)?;
let daemonize = Daemonize::new()
.pid_file(&args.pid_file)
.working_directory(".")
.stdout(stdout)
.stderr(stderr);
match daemonize.start() {
Ok(_) => {
}
Err(e) => {
eprintln!("Error starting daemon: {}", e);
std::process::exit(1);
}
}
}
#[cfg(not(unix))]
if args.daemon {
eprintln!("Daemon mode is only supported on Unix systems");
std::process::exit(1);
}
let runtime = tokio::runtime::Runtime::new()?;
runtime.block_on(async_main(args))
}
fn bind_host(args: &Args) -> String {
args.host
.clone()
.or_else(|| std::env::var("SOLIDB_HOST").ok())
.unwrap_or_else(|| "127.0.0.1".to_string())
}
async fn async_main(args: Args) -> anyhow::Result<()> {
let host = bind_host(&args);
tracing_subscriber::registry()
.with(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| "solidb=info,tower_http=info".into()),
)
.with(tracing_subscriber::fmt::layer())
.init();
let node_id = args
.node_id
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
let replication_port = args.replication_port.unwrap_or(args.port);
let advertise = args
.advertise
.clone()
.or_else(|| args.host.clone())
.unwrap_or_else(|| "127.0.0.1".to_string());
if let Err(reason) = check_advertise(&advertise, &args.peers) {
eprintln!("ERROR: {reason}");
eprintln!(
" Set --advertise to the address peers can reach this node on, or bind \
--host to that address."
);
std::process::exit(1);
}
let api_address = with_port(&advertise, args.port);
let repl_address = with_port(&advertise, replication_port);
let local_node = solidb::cluster::node::Node::new(
node_id.clone(),
repl_address.clone(),
api_address.clone(),
);
tracing::info!("Node ID: {}", local_node.id);
tracing::info!("Replication Address: {}", local_node.address);
tracing::info!("API Address: {}", local_node.api_address);
let cluster_config = ClusterConfig::new(
Some(node_id.clone()),
args.peers.clone(),
replication_port,
args.keyfile.clone(),
);
if !args.peers.is_empty() && cluster_config.keyfile.is_none() {
anyhow::bail!(
"Cluster peers are configured but no keyfile is available. \
Create a shared secret of 32 cryptographically random bytes, hex-encoded, \
and pass it with --keyfile (the same file on every node). \
Unix: `openssl rand -hex 32 > solidb.key`. PowerShell: \
`$b=[byte[]]::new(32);[Security.Cryptography.RandomNumberGenerator]::Fill($b);\
[BitConverter]::ToString($b).Replace('-','').ToLower() \
| Out-File -Encoding ascii solidb.key`. \
Refusing to start an unauthenticated cluster."
);
}
if args.peers.is_empty() && cluster_config.keyfile.is_none() {
tracing::warn!(
"No cluster keyfile configured: replication and cluster ports accept \
unauthenticated connections. Set --keyfile before adding peers."
);
}
log_allocator_tuning();
use solidb::storage::engine::{set_engine_profile, EngineProfile};
let (base, mut storage_profile) = if args.dev {
("dev", EngineProfile::dev())
} else {
("prod", EngineProfile::prod())
};
if let Some(bytes) = args.block_cache {
storage_profile.block_cache_bytes = bytes;
}
if let Some(bytes) = args.write_buffer_size {
storage_profile.write_buffer_size = bytes;
}
if let Some(bytes) = args.memtable_budget {
storage_profile.db_write_buffer_size = (bytes > 0).then_some(bytes);
}
if let Some(n) = args.max_open_files {
storage_profile.max_open_files = n;
}
if let Some(on) = args.bounded_index_cache {
storage_profile.cache_index_and_filter_blocks = on;
}
if let Some(n) = args.max_background_jobs {
storage_profile.max_background_jobs = n;
}
set_engine_profile(storage_profile);
log_storage_profile(base, &storage_profile);
let storage = StorageEngine::with_cluster_config(&args.data_dir, cluster_config.clone())?;
storage.initialize()?;
tracing::info!("Storage engine initialized");
let storage_for_shutdown = Arc::new(storage.clone());
let transport = Arc::new(solidb::cluster::transport::TcpTransport::new(
repl_address.clone(),
cluster_config.keyfile.clone(),
));
let cluster_state = solidb::cluster::state::ClusterState::new(node_id.clone());
let replication_log = Arc::new(
solidb::sync::log::SyncLog::new_with_options(
node_id.clone(),
&args.data_dir,
1000, args.no_sync_log,
)
.map_err(|e| anyhow::anyhow!("Failed to init replication log: {}", e))?,
);
if args.no_sync_log {
tracing::warn!(
"Sync log is DISABLED via --no-sync-log; do not use this flag on a node participating in replication"
);
}
let cluster_manager = Arc::new(solidb::cluster::manager::ClusterManager::new(
local_node.clone(),
cluster_state,
transport.clone(),
Some(replication_log.clone()),
Some(Arc::new(storage.clone())),
));
let mgr_clone2 = cluster_manager.clone();
tokio::spawn(async move {
mgr_clone2.start().await;
});
let shared_coordinator = Arc::new(solidb::sharding::coordinator::ShardCoordinator::new(
storage_for_shutdown.clone(),
Some(cluster_manager.clone()),
Some(replication_log.clone()),
));
let stats_storage = storage_for_shutdown.clone();
let stats_collector = solidb::cluster::stats::ClusterStatsCollector::new(
stats_storage,
shared_coordinator.clone(), cluster_manager.clone(),
);
tokio::spawn(async move {
stats_collector.start().await;
});
let health_config = solidb::cluster::health::HealthConfig::default();
let health_state = cluster_manager.state().clone();
let health_monitor = solidb::cluster::health::HealthMonitor::new(health_config, health_state);
tokio::spawn(async move {
health_monitor.start().await;
});
let healing_coordinator = shared_coordinator.clone();
tokio::spawn(async move {
let base = std::time::Duration::from_secs(5);
let max_backoff = std::time::Duration::from_secs(300);
let mut delay = base;
loop {
tokio::time::sleep(delay).await;
let mut failed = false;
if let Err(e) = healing_coordinator.cleanup_orphaned_shards().await {
tracing::error!("Orphaned shard cleanup failed: {}", e);
failed = true;
}
if let Err(e) = healing_coordinator.heal_shards().await {
tracing::error!("Shard healing failed: {}", e);
failed = true;
}
delay = if failed {
(delay * 2).min(max_backoff)
} else {
base
};
}
});
let blob_rebalance_config = Arc::new(solidb::sharding::RebalanceConfig::default());
let blob_worker = Arc::new(solidb::sharding::BlobRebalanceWorker::new(
storage_for_shutdown.clone(),
shared_coordinator.clone(),
Some(cluster_manager.clone()),
blob_rebalance_config,
));
let blob_worker_start = blob_worker.clone();
tokio::spawn(async move {
blob_worker_start.start().await;
});
tracing::info!("BlobRebalanceWorker started");
let worker_log = replication_log.clone();
let _worker_transport = transport.clone();
let _worker_mgr = cluster_manager.clone();
let worker_storage = Arc::new(storage.clone());
let worker_node_id = node_id.clone();
let worker_keyfile = args
.keyfile
.clone()
.unwrap_or_else(|| "solidb.key".to_string());
let worker_repl_addr = repl_address.clone();
let sync_state = Arc::new(solidb::sync::state::SyncState::new(
worker_storage.clone(),
worker_node_id.clone(),
));
let connection_pool = Arc::new(solidb::sync::transport::ConnectionPool::new(
worker_node_id.clone(),
worker_keyfile.clone(),
));
let (sync_cmd_tx, worker_cmd_rx) = solidb::sync::worker::create_command_channel();
let sync_config = solidb::sync::worker::SyncConfig::default();
let sync_worker = solidb::sync::worker::SyncWorker::new(
worker_storage,
sync_state,
connection_pool,
worker_log,
sync_config,
worker_cmd_rx,
worker_node_id,
worker_keyfile,
worker_repl_addr,
)
.with_cluster_manager(cluster_manager.clone())
.with_shard_coordinator(shared_coordinator.clone());
if !args.peers.is_empty() {
let mgr_clone3 = cluster_manager.clone();
let seeds = args.peers.clone();
let full_sync_tx = sync_cmd_tx.clone();
let startup_coordinator = shared_coordinator.clone();
tokio::spawn(async move {
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
let mut joined = false;
let mut joined_via = String::new();
for seed in seeds {
if let Err(e) = mgr_clone3.join_cluster(&seed).await {
tracing::warn!("Failed to join cluster via seed {}: {}", seed, e);
} else {
tracing::info!("Sent join request to {}", seed);
joined_via = seed.clone();
joined = true;
break; }
}
if joined {
tracing::info!("Requesting a full sync from {joined_via}");
if let Err(e) = full_sync_tx
.send(solidb::sync::worker::SyncCommand::RequestFullSync {
peer_addr: joined_via.clone(),
})
.await
{
tracing::error!(
"Could not request a full sync from {joined_via}: {e}. This node will \
only receive writes made from now on, and will not have the data that \
already existed."
);
}
tracing::info!(
"Waiting for shard tables to sync before cleaning up orphaned shards..."
);
tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;
if let Err(e) = startup_coordinator.rebalance().await {
tracing::warn!("Startup rebalance failed: {}", e);
}
match startup_coordinator.cleanup_orphaned_shards().await {
Ok(count) => {
if count > 0 {
tracing::info!(
"STARTUP: Cleaned up {} orphaned shard collections",
count
);
}
}
Err(e) => {
tracing::error!("STARTUP: Orphaned shard cleanup failed: {}", e);
}
}
match startup_coordinator.heal_shards().await {
Ok(count) => {
if count > 0 {
tracing::info!("STARTUP: Healed {} shard replicas", count);
}
}
Err(e) => {
tracing::warn!("STARTUP: Shard healing failed: {}", e);
}
}
}
});
}
let script_stats = Arc::new(ScriptStats::default());
let queue_worker = Arc::new(solidb::queue::QueueWorker::new(
Arc::new(storage.clone()),
script_stats.clone(),
));
let queue_worker_start = queue_worker.clone();
tokio::spawn(async move {
queue_worker_start.start().await;
});
let ttl_worker = Arc::new(solidb::ttl::TtlWorker::new(Arc::new(storage.clone())));
let ttl_worker_start = ttl_worker.clone();
tokio::spawn(async move {
ttl_worker_start.start().await;
});
let recovery_config = solidb::ai::RecoveryConfig::default();
let recovery_worker = Arc::new(solidb::ai::RecoveryWorker::new(
Arc::new(storage.clone()),
"_system".to_string(), recovery_config,
));
let recovery_worker_start = recovery_worker.clone();
tokio::spawn(async move {
recovery_worker_start.start().await;
});
tracing::info!("AI Recovery Worker started");
let stream_manager = Arc::new(solidb::stream::StreamManager::new(Arc::new(
storage.clone(),
)));
let http_client = Arc::new(
reqwest::Client::builder()
.pool_max_idle_per_host(10)
.pool_idle_timeout(Duration::from_secs(60))
.tcp_keepalive(Duration::from_secs(60))
.tcp_nodelay(true)
.connect_timeout(Duration::from_secs(10))
.build()
.expect("Failed to create HTTP client"),
);
tracing::info!("HTTP client with connection pooling initialized");
solidb::storage::http_client::init_http_client(http_client.as_ref().clone());
let app = create_router(
storage,
Some(cluster_manager.clone()),
Some(replication_log.clone()),
Some(shared_coordinator.clone()),
Some(queue_worker),
script_stats,
Some(stream_manager),
Some(blob_worker),
args.port,
);
let shutdown_storage = storage_for_shutdown.clone();
let tls_acceptor = match (&args.tls_cert, &args.tls_key) {
(Some(cert), Some(key)) => Some(solidb::server::tls::load_tls_acceptor(cert, key)?),
(None, None) => None,
_ => anyhow::bail!("--tls-cert and --tls-key must be provided together to enable HTTPS"),
};
if args.port == replication_port {
tracing::info!("Starting in MULTIPLEXED mode on port {}", args.port);
let addr = format!("{}:{}", host, args.port);
let listener = tokio::net::TcpListener::bind(&addr).await?;
let local_addr = listener.local_addr()?;
let (http_tx, http_rx) = mpsc::channel(8192);
let (sync_tx, sync_rx) = mpsc::channel(8192);
let channel_listener = ChannelListener::new(http_rx, local_addr);
let http_shutdown = shutdown_signal(shutdown_storage);
tokio::spawn(async move {
use axum::serve::Listener;
use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
use hyper_util::server::conn::auto::Builder as HttpConnBuilder;
use hyper_util::server::graceful::GracefulShutdown;
use hyper_util::service::TowerToHyperService;
let mut listener = channel_listener;
let graceful = GracefulShutdown::new();
tokio::pin!(http_shutdown);
let mut make_service =
app.into_make_service_with_connect_info::<std::net::SocketAddr>();
loop {
let (io, addr) = tokio::select! {
conn = listener.accept() => conn,
_ = &mut http_shutdown => {
tracing::info!("HTTP server received shutdown signal, draining connections");
break;
}
};
let mut builder = HttpConnBuilder::new(TokioExecutor::new());
builder
.http1()
.timer(TokioTimer::new())
.header_read_timeout(Duration::from_secs(30));
let tower_service = {
use tower::Service;
match make_service.call(addr).await {
Ok(svc) => svc,
Err(never) => match never {},
}
};
let service = TowerToHyperService::new(tower_service);
let conn = builder
.serve_connection_with_upgrades(TokioIo::new(io), service)
.into_owned();
let watched = graceful.watch(conn);
tokio::spawn(async move {
if let Err(e) = watched.await {
tracing::debug!("HTTP connection error: {}", e);
}
});
}
graceful.shutdown().await;
});
let sync_worker = sync_worker.with_incoming_channel(sync_rx);
tokio::spawn(async move {
sync_worker.run_background().await;
});
let driver_storage = storage_for_shutdown.clone();
let driver_tx =
solidb::driver::spawn_driver_handler(driver_storage, Some(replication_log.clone()));
tracing::info!("Native driver protocol enabled on port {}", args.port);
if tls_acceptor.is_some() {
if solidb::server::tls::tls_required() {
tracing::warn!(
"SOLIDB_TLS_REQUIRE=1: plaintext refused on port {} — native driver clients, \
sync and cluster peers cannot connect",
args.port
);
} else {
tracing::info!(
"TLS enabled on port {} in mixed mode: HTTPS clients are served over TLS, \
plaintext driver/sync/cluster connections are still accepted \
(set SOLIDB_TLS_REQUIRE=1 to refuse them)",
args.port
);
}
}
let shutdown_signal_future = async {
let ctrl_c = async {
tokio::signal::ctrl_c()
.await
.expect("failed to install Ctrl+C handler");
};
#[cfg(unix)]
let terminate = async {
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
.expect("failed to install signal handler")
.recv()
.await;
};
#[cfg(not(unix))]
let terminate = std::future::pending::<()>();
tokio::select! {
_ = ctrl_c => {},
_ = terminate => {},
}
};
tokio::pin!(shutdown_signal_future);
loop {
tokio::select! {
_ = &mut shutdown_signal_future => {
tracing::info!("Shutdown signal received in multiplexed mode, stopping...");
storage_for_shutdown.flush_all_stats();
tracing::info!("Shutdown complete");
std::process::exit(0);
}
accept_result = listener.accept() => {
let (stream, addr) = match accept_result {
Ok(conn) => conn,
Err(e) => {
tracing::error!("Accept error: {}", e);
continue;
}
};
let http_tx = http_tx.clone();
let sync_tx = sync_tx.clone();
let driver_tx = driver_tx.clone();
let connection_mgr = cluster_manager.clone();
let cluster_secret = cluster_config.keyfile.clone();
let tls = tls_acceptor.clone();
tokio::spawn(async move {
let mut stream: MuxConn = match tls {
Some(tls) => {
let mut first = [0u8; 1];
let offered_tls = match tokio::time::timeout(
std::time::Duration::from_secs(10),
stream.peek(&mut first),
)
.await
{
Ok(Ok(1)) => first[0] == solidb::server::tls::TLS_HANDSHAKE_CONTENT_TYPE,
Ok(Ok(_)) => return, Ok(Err(e)) => {
tracing::debug!(peer = %addr, "TLS sniff read failed: {}", e);
return;
}
Err(_) => {
tracing::warn!(peer = %addr, "TLS sniff timed out; dropping connection");
return;
}
};
if !offered_tls {
if solidb::server::tls::tls_required() {
tracing::warn!(peer = %addr, "plaintext connection refused (SOLIDB_TLS_REQUIRE=1)");
return;
}
Box::new(stream)
} else {
match tokio::time::timeout(
std::time::Duration::from_secs(10),
tls.accept(stream),
)
.await
{
Ok(Ok(tls_stream)) => Box::new(tls_stream),
Ok(Err(e)) => {
tracing::debug!(peer = %addr, "TLS handshake failed: {}", e);
return;
}
Err(_) => {
tracing::warn!(peer = %addr, "TLS handshake timed out; dropping connection");
return;
}
}
}
}
None => Box::new(stream),
};
let mut buf = vec![0u8; 14];
let n = match tokio::time::timeout(
std::time::Duration::from_secs(10),
stream.read(&mut buf),
)
.await
{
Ok(Ok(n)) => n,
Ok(Err(_)) => 0,
Err(_) => {
tracing::warn!(
layer = "solidb_detect",
peer = %addr,
"protocol detection read timed out; dropping connection"
);
return;
}
};
let peeked_data = buf[..n].to_vec();
if &peeked_data == b"solidb-sync-v1" {
let sync_stream: solidb::sync::transport::SyncStream = Box::new(stream);
dispatch_or_drop(&sync_tx, (sync_stream, addr.to_string()), "sync", &addr);
}
else if &peeked_data == b"solidb-drv-v1\0" {
dispatch_or_drop(&driver_tx, (stream, addr.to_string()), "driver", &addr);
}
else if peeked_data.first() == Some(&b'{') {
let peeked_stream = PeekedStream::new(stream, peeked_data.clone());
let mgr = connection_mgr.clone();
tokio::spawn(async move {
let mut buf = Vec::new();
let mut stream = tokio::io::AsyncReadExt::take(
peeked_stream,
(solidb::cluster::transport::MAX_CLUSTER_MESSAGE_SIZE + 1) as u64,
);
let read = tokio::time::timeout(
std::time::Duration::from_secs(10),
stream.read_to_end(&mut buf),
)
.await;
match read {
Ok(Ok(_)) if buf.len() <= solidb::cluster::transport::MAX_CLUSTER_MESSAGE_SIZE => {
match solidb::cluster::transport::open_cluster_message(
&buf,
cluster_secret.as_deref(),
) {
Ok(msg) => mgr.handle_message(msg).await,
Err(e) => {
tracing::warn!(
"Rejected cluster message from {}: {}",
addr,
e
);
}
}
}
Ok(Ok(_)) => {
tracing::warn!(
"Cluster message from {} exceeds size limit, dropped",
addr
);
}
_ => {
tracing::warn!(
"Cluster message read from {} failed or timed out",
addr
);
}
}
});
} else {
let peeked_stream = PeekedStream::new(stream, peeked_data.clone());
dispatch_or_drop(&http_tx, (peeked_stream, addr), "HTTP", &addr);
}
});
}
}
}
} else {
tracing::info!(
"Starting in DUAL PORT mode (API: {}, Sync: {})",
args.port,
replication_port
);
tokio::spawn(async move {
sync_worker.run().await;
});
let addr = format!("{}:{}", host, args.port);
let listener = tokio::net::TcpListener::bind(&addr).await?;
tracing::info!("Server listening on {}", addr);
match tls_acceptor {
Some(tls) => {
use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
use hyper_util::server::conn::auto::Builder as HttpConnBuilder;
use hyper_util::server::graceful::GracefulShutdown;
use hyper_util::service::TowerToHyperService;
tracing::info!("HTTPS enabled on port {}", args.port);
let shutdown = shutdown_signal(shutdown_storage);
tokio::pin!(shutdown);
let graceful = GracefulShutdown::new();
let mut make_service =
app.into_make_service_with_connect_info::<std::net::SocketAddr>();
loop {
tokio::select! {
_ = &mut shutdown => {
tracing::info!("Shutdown signal received, draining connections");
break;
}
accepted = listener.accept() => {
let (tcp, peer) = match accepted {
Ok(conn) => conn,
Err(e) => {
tracing::error!("Accept error: {}", e);
continue;
}
};
let tls = tls.clone();
let service = {
use tower::Service;
match make_service.call(peer).await {
Ok(svc) => TowerToHyperService::new(svc),
Err(never) => match never {},
}
};
let watcher = graceful.watcher();
tokio::spawn(async move {
let tls_stream = match tokio::time::timeout(
Duration::from_secs(10),
tls.accept(tcp),
)
.await
{
Ok(Ok(tls_stream)) => tls_stream,
Ok(Err(e)) => {
tracing::debug!(peer = %peer, "TLS handshake failed: {}", e);
return;
}
Err(_) => {
tracing::warn!(peer = %peer, "TLS handshake timed out; dropping connection");
return;
}
};
let mut builder = HttpConnBuilder::new(TokioExecutor::new());
builder
.http1()
.timer(TokioTimer::new())
.header_read_timeout(Duration::from_secs(30));
let conn = builder
.serve_connection_with_upgrades(
TokioIo::new(tls_stream),
service,
)
.into_owned();
if let Err(e) = watcher.watch(conn).await {
tracing::debug!("HTTPS connection error: {}", e);
}
});
}
}
}
graceful.shutdown().await;
}
None => {
axum::serve(
listener,
app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
)
.with_graceful_shutdown(shutdown_signal(shutdown_storage))
.await?;
}
}
}
Ok(())
}
type MuxConn = Box<dyn solidb::driver::handlers::DriverConnTrait>;
fn dispatch_or_drop<T>(tx: &mpsc::Sender<T>, item: T, what: &str, peer: &std::net::SocketAddr) {
match tx.try_send(item) {
Ok(()) => {}
Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
tracing::warn!(peer = %peer, "{} dispatch channel full, dropping connection", what);
}
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
tracing::error!("{} dispatch channel closed", what);
}
}
}
async fn shutdown_signal(storage: Arc<StorageEngine>) {
let ctrl_c = async {
tokio::signal::ctrl_c()
.await
.expect("failed to install Ctrl+C handler");
};
#[cfg(unix)]
let terminate = async {
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
.expect("failed to install signal handler")
.recv()
.await;
};
#[cfg(not(unix))]
let terminate = std::future::pending::<()>();
tokio::select! {
_ = ctrl_c => {},
_ = terminate => {},
}
tracing::info!("Shutdown signal received, flushing stats...");
storage.flush_all_stats();
tracing::info!("Shutdown complete");
}
fn check_advertise(advertise: &str, peers: &[String]) -> Result<(), String> {
if peers.is_empty() {
return Ok(());
}
if !is_unroutable_advertise(advertise) {
return Ok(());
}
let host = advertised_host(advertise);
let unspecified = host.is_empty()
|| host == "0.0.0.0"
|| host == "::"
|| host
.parse::<std::net::IpAddr>()
.map(|ip| ip.is_unspecified())
.unwrap_or(false);
if unspecified {
return Err(format!(
"--peer is set but this node would advertise itself as {advertise:?}, which is a \
bind address, not one a peer can dial"
));
}
let elsewhere: Vec<&String> = peers
.iter()
.filter(|peer| !is_unroutable_advertise(peer))
.collect();
if elsewhere.is_empty() {
return Ok(());
}
Err(format!(
"--peer is set but this node would advertise itself as {advertise:?}, which every peer \
reads as its own loopback — and {} is not on this host",
elsewhere[0]
))
}
fn with_port(advertise: &str, port: u16) -> String {
let host = advertise.trim();
if host.parse::<std::net::SocketAddr>().is_ok() {
return host.to_string();
}
if host.starts_with('[') && host.ends_with(']') {
return format!("{host}:{port}");
}
if host.parse::<std::net::Ipv6Addr>().is_ok() {
return format!("[{host}]:{port}");
}
if let Some((name, tail)) = host.rsplit_once(':') {
if !name.is_empty() && !name.contains(':') && tail.parse::<u16>().is_ok() {
return host.to_string();
}
}
format!("{host}:{port}")
}
fn is_unroutable_advertise(address: &str) -> bool {
let host = advertised_host(address);
if host.is_empty()
|| host == "0.0.0.0"
|| host == "::"
|| host.eq_ignore_ascii_case("localhost")
{
return true;
}
host.parse::<std::net::IpAddr>()
.map(|ip| ip.is_loopback() || ip.is_unspecified())
.unwrap_or(false)
}
fn advertised_host(address: &str) -> String {
let value = address.trim();
if let Ok(socket) = value.parse::<std::net::SocketAddr>() {
return socket.ip().to_string();
}
if let Some(inner) = value.strip_prefix('[').and_then(|v| v.strip_suffix(']')) {
return inner.to_string();
}
if let Some((name, tail)) = value.rsplit_once(':') {
if !name.is_empty() && !name.contains(':') && tail.parse::<u16>().is_ok() {
return name.to_string();
}
}
value.to_string()
}
#[cfg(test)]
mod storage_size_tests {
use super::{human_size, parse_size};
#[test]
fn suffixes_are_binary_multiples() {
assert_eq!(parse_size("512MB"), Ok(512 * 1024 * 1024));
assert_eq!(parse_size("2GB"), Ok(2 * 1024 * 1024 * 1024));
assert_eq!(parse_size("64KB"), Ok(64 * 1024));
assert_eq!(parse_size("2g"), parse_size("2GB"));
assert_eq!(parse_size("512m"), parse_size("512MB"));
assert_eq!(parse_size(" 128mb "), parse_size("128MB"));
}
#[test]
fn a_bare_number_is_a_byte_count() {
assert_eq!(parse_size("1048576"), Ok(1048576));
assert_eq!(parse_size("0"), Ok(0));
}
#[test]
fn nonsense_is_rejected_rather_than_silently_misread() {
assert!(parse_size("512QB").is_err());
assert!(parse_size("MB").is_err());
assert!(parse_size("-1").is_err());
assert!(parse_size("1.5GB").is_err());
assert!(parse_size("").is_err());
}
#[test]
fn the_startup_log_renders_what_the_flags_accept() {
for s in ["512MB", "2GB", "128KB", "8MB"] {
assert_eq!(human_size(parse_size(s).unwrap()), s);
}
assert_eq!(human_size(1_500_000), "1500000B");
}
}
#[cfg(test)]
mod advertise_tests {
use super::{advertised_host, check_advertise, is_unroutable_advertise, with_port};
#[test]
fn loopback_cannot_be_advertised_to_a_peer() {
assert!(is_unroutable_advertise("127.0.0.1"));
assert!(is_unroutable_advertise("::1"));
assert!(is_unroutable_advertise("localhost"));
}
#[test]
fn the_unspecified_address_cannot_be_advertised_either() {
assert!(is_unroutable_advertise("0.0.0.0"));
assert!(is_unroutable_advertise("::"));
assert!(is_unroutable_advertise("[::]"));
assert!(is_unroutable_advertise(""));
assert!(is_unroutable_advertise(" "));
}
#[test]
fn a_routable_address_is_accepted() {
assert!(!is_unroutable_advertise("51.15.248.118"));
assert!(!is_unroutable_advertise("10.0.0.7"));
assert!(!is_unroutable_advertise("2001:db8::1"));
}
#[test]
fn a_hostname_is_accepted_rather_than_resolved() {
assert!(!is_unroutable_advertise("db1.internal.example"));
}
#[test]
fn an_advertise_that_already_has_a_port_does_not_get_a_second_one() {
assert_eq!(with_port("10.0.0.1:6746", 6746), "10.0.0.1:6746");
assert_eq!(with_port("10.0.0.1:6746", 9999), "10.0.0.1:6746");
}
#[test]
fn a_bare_host_still_gets_the_port_appended() {
assert_eq!(with_port("10.0.0.1", 6746), "10.0.0.1:6746");
assert_eq!(with_port("db1.example.com", 6746), "db1.example.com:6746");
}
#[test]
fn a_hostname_carrying_a_port_is_recognised_even_though_no_ip_parser_accepts_it() {
assert_eq!(
with_port("db1.example.com:6746", 9999),
"db1.example.com:6746"
);
}
#[test]
fn a_bare_ipv6_address_is_bracketed_before_a_port_is_added() {
assert_eq!(with_port("2001:db8::1", 6746), "[2001:db8::1]:6746");
assert_eq!(with_port("::1", 6746), "[::1]:6746");
}
#[test]
fn a_bracketed_ipv6_address_is_handled_both_ways() {
assert_eq!(with_port("[2001:db8::1]", 6746), "[2001:db8::1]:6746");
assert_eq!(with_port("[2001:db8::1]:6746", 9999), "[2001:db8::1]:6746");
}
#[test]
fn a_port_does_not_let_loopback_past_the_guard() {
assert!(is_unroutable_advertise("127.0.0.1:6746"));
assert!(is_unroutable_advertise("[::1]:6746"));
assert!(is_unroutable_advertise("localhost:6746"));
assert!(is_unroutable_advertise("0.0.0.0:6746"));
}
#[test]
fn a_routable_address_passes_in_either_form() {
assert!(!is_unroutable_advertise("10.0.0.1"));
assert!(!is_unroutable_advertise("10.0.0.1:6746"));
assert!(!is_unroutable_advertise("db1.example.com:6746"));
assert!(!is_unroutable_advertise("[2001:db8::1]:6746"));
}
#[test]
fn the_host_is_extracted_without_brackets_or_port() {
assert_eq!(advertised_host("10.0.0.1:6746"), "10.0.0.1");
assert_eq!(advertised_host("[2001:db8::1]:6746"), "2001:db8::1");
assert_eq!(advertised_host("2001:db8::1"), "2001:db8::1");
assert_eq!(advertised_host("db1.example.com"), "db1.example.com");
}
fn peers(list: &[&str]) -> Vec<String> {
list.iter().map(|s| s.to_string()).collect()
}
#[test]
fn a_single_host_cluster_on_loopback_is_allowed() {
assert!(check_advertise("127.0.0.1:6747", &peers(&["127.0.0.1:6746"])).is_ok());
assert!(check_advertise("[::1]:6747", &peers(&["[::1]:6746"])).is_ok());
}
#[test]
fn loopback_is_refused_the_moment_one_peer_is_elsewhere() {
let error = check_advertise("127.0.0.1:6746", &peers(&["10.0.0.2:6746"])).unwrap_err();
assert!(error.contains("its own loopback"), "{error}");
assert!(error.contains("10.0.0.2:6746"), "{error}");
}
#[test]
fn a_mixed_peer_list_is_refused_rather_than_averaged() {
assert!(check_advertise(
"127.0.0.1:6746",
&peers(&["127.0.0.1:6747", "10.0.0.2:6746"])
)
.is_err());
}
#[test]
fn the_unspecified_address_is_refused_even_on_one_host() {
let error = check_advertise("0.0.0.0:6746", &peers(&["127.0.0.1:6747"])).unwrap_err();
assert!(error.contains("bind address"), "{error}");
assert!(check_advertise("::", &peers(&["127.0.0.1:6747"])).is_err());
}
#[test]
fn no_peers_means_nothing_to_check() {
assert!(check_advertise("127.0.0.1", &[]).is_ok());
assert!(check_advertise("0.0.0.0", &[]).is_ok());
}
#[test]
fn a_routable_advertise_passes_whatever_the_peers_are() {
assert!(check_advertise("10.0.0.1:6746", &peers(&["10.0.0.2:6746"])).is_ok());
assert!(check_advertise("10.0.0.1:6746", &peers(&["127.0.0.1:6747"])).is_ok());
}
}