perf: add SO_REUSEPORT multiple accept loops (--workers flag)

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;

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 }

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<usize>,
+
+    /// 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);
     }

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<usize>,
+    metrics_enabled: bool,
+    slowlog_config: SlowLogConfig,
+) -> Result<(), Box<dyn std::error::Error>> {
+    // 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<TcpListener, Box<dyn std::error::Error>> {
+    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<ServerContext>,
+    slow_log: Arc<SlowLog>,
+    semaphore: Arc<Semaphore>,
+) {
+    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);
+        });
+    }
+}