diff --git a/Cargo.lock b/Cargo.lock index 374d5282..4313b11d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -548,6 +548,7 @@ dependencies = [ "futures", "metrics", "metrics-exporter-prometheus", + "socket2 0.5.10", "tikv-jemallocator", "tokio", "tracing", @@ -906,7 +907,7 @@ dependencies = [ "hyper", "libc", "pin-project-lite", - "socket2", + "socket2 0.6.2", "tokio", "tower-service", "tracing", @@ -1827,6 +1828,16 @@ version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +[[package]] +name = "socket2" +version = "0.5.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e22376abed350d73dd1cd119b57ffccad95b4e585a7cda43e286245ce23c0678" +dependencies = [ + "libc", + "windows-sys 0.52.0", +] + [[package]] name = "socket2" version = "0.6.2" @@ -1996,7 +2007,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "signal-hook-registry", - "socket2", + "socket2 0.6.2", "tokio-macros", "windows-sys 0.61.2", ] diff --git a/crates/ember-core/src/lib.rs b/crates/ember-core/src/lib.rs index 60873630..22d1c05d 100644 --- a/crates/ember-core/src/lib.rs +++ b/crates/ember-core/src/lib.rs @@ -18,5 +18,5 @@ pub use keyspace::{ EvictionPolicy, IncrError, Keyspace, KeyspaceStats, ShardConfig, TtlResult, WriteError, WrongType, ZAddResult, }; -pub use shard::{ShardPersistenceConfig, ShardRequest, ShardResponse}; +pub use shard::{ShardHandle, ShardPersistenceConfig, ShardRequest, ShardResponse, spawn_shard}; pub use types::Value; diff --git a/crates/ember-server/Cargo.toml b/crates/ember-server/Cargo.toml index ad96d653..bcd5158c 100644 --- a/crates/ember-server/Cargo.toml +++ b/crates/ember-server/Cargo.toml @@ -26,6 +26,7 @@ clap = { workspace = true } metrics = { workspace = true } metrics-exporter-prometheus = { workspace = true } futures = "0.3" +socket2 = { version = "0.5", features = ["all"] } # optional: better multi-threaded allocation performance tikv-jemallocator = { version = "0.6", optional = true } diff --git a/crates/ember-server/src/main.rs b/crates/ember-server/src/main.rs index 6086b2f7..49844fc1 100644 --- a/crates/ember-server/src/main.rs +++ b/crates/ember-server/src/main.rs @@ -68,6 +68,11 @@ struct Args { /// number of shards (worker threads). defaults to available CPU cores #[arg(long)] shards: Option, + + /// use worker-per-core architecture with SO_REUSEPORT for better scaling. + /// each worker gets its own accept loop and tokio runtime. + #[arg(long)] + workers: bool, } #[tokio::main] @@ -168,18 +173,33 @@ async fn main() { enabled: args.slowlog_log_slower_than >= 0, }; - info!(shards = shard_count, "ember server starting..."); - - if let Err(e) = server::run( - addr, - shard_count, - engine_config, - None, - args.metrics_port.is_some(), - slowlog_config, - ) - .await - { + info!(shards = shard_count, workers = args.workers, "ember server starting..."); + + let result = if args.workers { + // Worker mode: multiple accept loops with SO_REUSEPORT. + server::run_with_workers( + addr, + shard_count, + engine_config, + None, + args.metrics_port.is_some(), + slowlog_config, + ) + .await + } else { + // Default mode: single accept loop in tokio multi-threaded runtime. + server::run( + addr, + shard_count, + engine_config, + None, + args.metrics_port.is_some(), + slowlog_config, + ) + .await + }; + + if let Err(e) = result { eprintln!("server error: {e}"); std::process::exit(1); } diff --git a/crates/ember-server/src/server.rs b/crates/ember-server/src/server.rs index 5b7a6141..d587b9cd 100644 --- a/crates/ember-server/src/server.rs +++ b/crates/ember-server/src/server.rs @@ -1,5 +1,9 @@ //! TCP server that accepts client connections and spawns handler tasks. //! +//! Supports two modes: +//! - Single-listener: one accept loop in the main tokio runtime (default) +//! - Worker-per-core: multiple accept loops via SO_REUSEPORT (better scaling) +//! //! Handles graceful shutdown on SIGINT/SIGTERM: stops accepting new //! connections and waits for in-flight requests to drain before exiting. @@ -158,3 +162,166 @@ pub async fn run( Ok(()) } + +/// Runs the server with multiple accept loops using SO_REUSEPORT. +/// +/// Spawns multiple accept tasks on the tokio runtime, each with its own +/// TcpListener bound to the same port via SO_REUSEPORT. The kernel +/// distributes incoming connections across accept loops. +/// +/// This reduces contention on the single accept loop while keeping +/// everything on the same tokio runtime as the shards. +pub async fn run_with_workers( + addr: SocketAddr, + shard_count: usize, + config: EngineConfig, + max_connections: Option, + metrics_enabled: bool, + slowlog_config: SlowLogConfig, +) -> Result<(), Box> { + // ensure data directory exists if persistence is configured + if let Some(ref pcfg) = config.persistence { + std::fs::create_dir_all(&pcfg.data_dir)?; + } + + let aof_enabled = config + .persistence + .as_ref() + .map(|p| p.append_only) + .unwrap_or(false); + let max_memory = config + .shard + .max_memory + .map(|per_shard| per_shard * shard_count); + + let engine = Engine::with_config(shard_count, config); + let max_conn = max_connections.unwrap_or(DEFAULT_MAX_CONNECTIONS); + + if metrics_enabled { + crate::metrics::spawn_stats_poller(engine.clone()); + } + + let ctx = Arc::new(ServerContext { + start_time: Instant::now(), + version: env!("CARGO_PKG_VERSION"), + shard_count, + max_connections: max_conn, + max_memory, + aof_enabled, + metrics_enabled, + connections_accepted: AtomicU64::new(0), + connections_active: AtomicU64::new(0), + commands_processed: AtomicU64::new(0), + }); + + let slow_log = Arc::new(SlowLog::new(slowlog_config)); + let semaphore = Arc::new(Semaphore::new(max_conn)); + + info!( + "starting {} accept loops on {addr} with SO_REUSEPORT (max {max_conn} connections)", + shard_count + ); + + // Spawn multiple accept loops, each with its own listener via SO_REUSEPORT + let mut accept_handles = Vec::with_capacity(shard_count); + for acceptor_id in 0..shard_count { + let listener = create_reuseport_listener(addr)?; + let engine = engine.clone(); + let ctx = Arc::clone(&ctx); + let slow_log = Arc::clone(&slow_log); + let semaphore = Arc::clone(&semaphore); + + let handle = tokio::spawn(async move { + run_accept_loop(acceptor_id, listener, engine, ctx, slow_log, semaphore).await + }); + accept_handles.push(handle); + } + + // Wait for shutdown signal + tokio::signal::ctrl_c().await?; + info!("shutdown signal received, stopping accept loops..."); + + // Cancel all accept loops + for handle in &accept_handles { + handle.abort(); + } + + // Wait for connection drain + info!("waiting for active connections to close..."); + let _ = semaphore.acquire_many(max_conn as u32).await; + info!("all connections drained, shutting down"); + + Ok(()) +} + +/// Creates a TcpListener with SO_REUSEPORT enabled. +fn create_reuseport_listener(addr: SocketAddr) -> Result> { + let socket = socket2::Socket::new( + socket2::Domain::for_address(addr), + socket2::Type::STREAM, + Some(socket2::Protocol::TCP), + )?; + + socket.set_reuse_port(true)?; + socket.set_reuse_address(true)?; + socket.set_nonblocking(true)?; + socket.bind(&addr.into())?; + socket.listen(1024)?; + + Ok(TcpListener::from_std(socket.into())?) +} + +/// Runs a single accept loop. +async fn run_accept_loop( + acceptor_id: usize, + listener: TcpListener, + engine: Engine, + ctx: Arc, + slow_log: Arc, + semaphore: Arc, +) { + info!(acceptor_id, "accept loop started"); + + loop { + let (stream, peer) = match listener.accept().await { + Ok(result) => result, + Err(e) => { + error!(acceptor_id, "accept error: {e}"); + continue; + } + }; + + let permit = match semaphore.clone().try_acquire_owned() { + Ok(permit) => permit, + Err(_) => { + warn!(acceptor_id, "connection limit reached, dropping {peer}"); + if ctx.metrics_enabled { + crate::metrics::on_connection_rejected(); + } + drop(stream); + continue; + } + }; + + if ctx.metrics_enabled { + crate::metrics::on_connection_accepted(); + } + ctx.connections_accepted.fetch_add(1, Ordering::Relaxed); + ctx.connections_active.fetch_add(1, Ordering::Relaxed); + + let engine = engine.clone(); + let ctx = Arc::clone(&ctx); + let slow_log = Arc::clone(&slow_log); + + tokio::spawn(async move { + if let Err(e) = connection::handle(stream, engine, &ctx, &slow_log).await { + error!("connection error from {peer}: {e}"); + } + ctx.connections_active.fetch_sub(1, Ordering::Relaxed); + if ctx.metrics_enabled { + crate::metrics::on_connection_closed(); + } + drop(permit); + }); + } +}