diff --git a/Cargo.lock b/Cargo.lock index 374d5282..031363ac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -460,6 +460,20 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +[[package]] +name = "dashmap" +version = "6.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + [[package]] name = "derive_more" version = "1.0.0" @@ -567,9 +581,11 @@ version = "0.2.2" dependencies = [ "bytes", "criterion", + "dashmap", "ember-persistence", "ember-protocol", "ordered-float", + "parking_lot", "rand 0.9.2", "tempfile", "thiserror 2.0.18", @@ -780,6 +796,12 @@ dependencies = [ "ahash 0.7.8", ] +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + [[package]] name = "hashbrown" version = "0.15.5" diff --git a/crates/ember-core/Cargo.toml b/crates/ember-core/Cargo.toml index 387fa522..3c990b39 100644 --- a/crates/ember-core/Cargo.toml +++ b/crates/ember-core/Cargo.toml @@ -21,6 +21,8 @@ tokio = { workspace = true } tracing = { workspace = true } rand = { workspace = true } ordered-float = { workspace = true } +dashmap = "6" +parking_lot = "0.12" [dev-dependencies] tempfile = "3" @@ -34,3 +36,7 @@ harness = false [[bench]] name = "engine" harness = false + +[[bench]] +name = "concurrent" +harness = false diff --git a/crates/ember-core/benches/concurrent.rs b/crates/ember-core/benches/concurrent.rs new file mode 100644 index 00000000..ac470529 --- /dev/null +++ b/crates/ember-core/benches/concurrent.rs @@ -0,0 +1,108 @@ +//! Benchmark comparing sharded engine vs concurrent keyspace. + +use bytes::Bytes; +use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput}; +use ember_core::concurrent::ConcurrentKeyspace; +use ember_core::keyspace::EvictionPolicy; +use std::sync::Arc; + +fn bench_concurrent_set(c: &mut Criterion) { + let ks = Arc::new(ConcurrentKeyspace::new(None, EvictionPolicy::NoEviction)); + + let mut group = c.benchmark_group("concurrent"); + group.throughput(Throughput::Elements(1)); + + group.bench_function("set", |b| { + let mut i = 0u64; + b.iter(|| { + let key = format!("key:{}", i); + i = i.wrapping_add(1); + ks.set(key, Bytes::from_static(b"value"), None); + black_box(()) + }) + }); + + // Pre-populate for get benchmark + for i in 0..10000 { + ks.set(format!("key:{}", i), Bytes::from_static(b"value"), None); + } + + group.bench_function("get_existing", |b| { + let mut i = 0u64; + b.iter(|| { + let key = format!("key:{}", i % 10000); + i = i.wrapping_add(1); + black_box(ks.get(&key)) + }) + }); + + group.bench_function("get_missing", |b| { + let mut i = 0u64; + b.iter(|| { + let key = format!("missing:{}", i); + i = i.wrapping_add(1); + black_box(ks.get(&key)) + }) + }); + + group.finish(); +} + +fn bench_concurrent_multithread(c: &mut Criterion) { + use std::thread; + + let ks = Arc::new(ConcurrentKeyspace::new(None, EvictionPolicy::NoEviction)); + + // Pre-populate + for i in 0..100000 { + ks.set(format!("key:{}", i), Bytes::from_static(b"value"), None); + } + + let mut group = c.benchmark_group("concurrent_mt"); + group.throughput(Throughput::Elements(8)); // 8 threads + + group.bench_function("get_8_threads", |b| { + b.iter(|| { + let handles: Vec<_> = (0..8) + .map(|t| { + let ks = Arc::clone(&ks); + thread::spawn(move || { + for i in 0..1000 { + let key = format!("key:{}", (t * 1000 + i) % 100000); + black_box(ks.get(&key)); + } + }) + }) + .collect(); + + for h in handles { + h.join().unwrap(); + } + }) + }); + + group.bench_function("set_8_threads", |b| { + b.iter(|| { + let handles: Vec<_> = (0..8) + .map(|t| { + let ks = Arc::clone(&ks); + thread::spawn(move || { + for i in 0..1000 { + let key = format!("thread:{}:key:{}", t, i); + ks.set(key, Bytes::from_static(b"value"), None); + } + }) + }) + .collect(); + + for h in handles { + h.join().unwrap(); + } + }) + }); + + group.finish(); +} + +criterion_group!(benches, bench_concurrent_set, bench_concurrent_multithread); +criterion_main!(benches); diff --git a/crates/ember-core/benches/keyspace.rs b/crates/ember-core/benches/keyspace.rs index 75e9032b..bceef5a0 100644 --- a/crates/ember-core/benches/keyspace.rs +++ b/crates/ember-core/benches/keyspace.rs @@ -93,7 +93,7 @@ fn bench_mixed(c: &mut Criterion) { let mut i = 0u64; b.iter(|| { let key = format!("key:{}", i % KEY_COUNT as u64); - if i % 2 == 0 { + if i.is_multiple_of(2) { let _ = black_box(ks.get(&key)); } else { black_box(ks.set(key, value.clone(), None)); diff --git a/crates/ember-core/src/concurrent.rs b/crates/ember-core/src/concurrent.rs new file mode 100644 index 00000000..677ac99c --- /dev/null +++ b/crates/ember-core/src/concurrent.rs @@ -0,0 +1,299 @@ +//! Concurrent keyspace using DashMap for lock-free multi-threaded access. +//! +//! This is an alternative to the sharded architecture that eliminates channel +//! overhead by allowing direct access from multiple connection handlers. + +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; + +use bytes::Bytes; +use dashmap::DashMap; + +use crate::keyspace::{EvictionPolicy, TtlResult}; + +/// An entry in the concurrent keyspace. +#[derive(Debug, Clone)] +struct Entry { + value: Bytes, + expires_at: Option, + size: usize, +} + +impl Entry { + fn is_expired(&self) -> bool { + self.expires_at + .map(|t| Instant::now() >= t) + .unwrap_or(false) + } +} + +/// A concurrent keyspace backed by DashMap. +/// +/// Provides thread-safe access to key-value data without channel overhead. +/// All operations are lock-free for non-conflicting keys. +#[derive(Debug)] +pub struct ConcurrentKeyspace { + data: DashMap, + memory_used: AtomicUsize, + max_memory: Option, + eviction_policy: EvictionPolicy, + ops_count: AtomicU64, +} + +impl ConcurrentKeyspace { + /// Creates a new concurrent keyspace with optional memory limit. + pub fn new(max_memory: Option, eviction_policy: EvictionPolicy) -> Self { + Self { + data: DashMap::new(), + memory_used: AtomicUsize::new(0), + max_memory, + eviction_policy, + ops_count: AtomicU64::new(0), + } + } + + /// Gets a value by key, returning None if not found or expired. + pub fn get(&self, key: &str) -> Option { + self.ops_count.fetch_add(1, Ordering::Relaxed); + + let entry = self.data.get(key)?; + + if entry.is_expired() { + drop(entry); + // Remove expired entry + if let Some((_, removed)) = self.data.remove(key) { + self.memory_used.fetch_sub(removed.size, Ordering::Relaxed); + } + return None; + } + + Some(entry.value.clone()) + } + + /// Sets a key-value pair with optional TTL. + pub fn set(&self, key: String, value: Bytes, ttl: Option) -> bool { + self.ops_count.fetch_add(1, Ordering::Relaxed); + + let entry_size = key.len() + value.len() + 64; // rough overhead estimate + let expires_at = ttl.map(|d| Instant::now() + d); + + // Check memory limit + if let Some(max) = self.max_memory { + let current = self.memory_used.load(Ordering::Relaxed); + if current + entry_size > max { + if self.eviction_policy == EvictionPolicy::NoEviction { + return false; + } + // Simple eviction: remove some entries + self.evict_entries(entry_size); + } + } + + let entry = Entry { + value, + expires_at, + size: entry_size, + }; + + // Update memory tracking + if let Some(old) = self.data.insert(key, entry) { + // Replace: adjust memory + let diff = entry_size as isize - old.size as isize; + if diff > 0 { + self.memory_used.fetch_add(diff as usize, Ordering::Relaxed); + } else { + self.memory_used + .fetch_sub((-diff) as usize, Ordering::Relaxed); + } + } else { + self.memory_used.fetch_add(entry_size, Ordering::Relaxed); + } + + true + } + + /// Deletes a key, returning true if it existed. + pub fn del(&self, key: &str) -> bool { + self.ops_count.fetch_add(1, Ordering::Relaxed); + + if let Some((_, removed)) = self.data.remove(key) { + self.memory_used.fetch_sub(removed.size, Ordering::Relaxed); + true + } else { + false + } + } + + /// Checks if a key exists (and is not expired). + pub fn exists(&self, key: &str) -> bool { + self.get(key).is_some() + } + + /// Returns the TTL of a key. + pub fn ttl(&self, key: &str) -> TtlResult { + match self.data.get(key) { + None => TtlResult::NotFound, + Some(entry) => { + if entry.is_expired() { + TtlResult::NotFound + } else { + match entry.expires_at { + None => TtlResult::NoExpiry, + Some(t) => { + let remaining = t.saturating_duration_since(Instant::now()); + TtlResult::Seconds(remaining.as_secs()) + } + } + } + } + } + } + + /// Sets expiration on a key. + pub fn expire(&self, key: &str, seconds: u64) -> bool { + self.ops_count.fetch_add(1, Ordering::Relaxed); + + if let Some(mut entry) = self.data.get_mut(key) { + if entry.is_expired() { + return false; + } + entry.expires_at = Some(Instant::now() + Duration::from_secs(seconds)); + true + } else { + false + } + } + + /// Returns the number of keys. + pub fn len(&self) -> usize { + self.data.len() + } + + /// Returns true if the keyspace is empty. + pub fn is_empty(&self) -> bool { + self.data.is_empty() + } + + /// Returns memory usage in bytes. + pub fn memory_used(&self) -> usize { + self.memory_used.load(Ordering::Relaxed) + } + + /// Returns the operation count. + pub fn ops_count(&self) -> u64 { + self.ops_count.load(Ordering::Relaxed) + } + + /// Clears all keys. + pub fn clear(&self) { + self.data.clear(); + self.memory_used.store(0, Ordering::Relaxed); + } + + /// Simple eviction: remove approximately `needed` bytes worth of entries. + fn evict_entries(&self, needed: usize) { + let mut freed = 0usize; + let mut keys_to_remove = Vec::new(); + + // Collect keys to remove (can't remove while iterating) + for entry in self.data.iter() { + if freed >= needed { + break; + } + keys_to_remove.push(entry.key().clone()); + freed += entry.value().size; + } + + // Remove collected keys + for key in keys_to_remove { + if let Some((_, removed)) = self.data.remove(&key) { + self.memory_used.fetch_sub(removed.size, Ordering::Relaxed); + } + } + } +} + +impl Default for ConcurrentKeyspace { + fn default() -> Self { + Self::new(None, EvictionPolicy::NoEviction) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn set_and_get() { + let ks = ConcurrentKeyspace::default(); + assert!(ks.set("key".into(), Bytes::from("value"), None)); + assert_eq!(ks.get("key"), Some(Bytes::from("value"))); + } + + #[test] + fn get_missing() { + let ks = ConcurrentKeyspace::default(); + assert_eq!(ks.get("missing"), None); + } + + #[test] + fn del_existing() { + let ks = ConcurrentKeyspace::default(); + ks.set("key".into(), Bytes::from("value"), None); + assert!(ks.del("key")); + assert_eq!(ks.get("key"), None); + } + + #[test] + fn del_missing() { + let ks = ConcurrentKeyspace::default(); + assert!(!ks.del("missing")); + } + + #[test] + fn exists_check() { + let ks = ConcurrentKeyspace::default(); + ks.set("key".into(), Bytes::from("value"), None); + assert!(ks.exists("key")); + assert!(!ks.exists("missing")); + } + + #[test] + fn ttl_expires() { + let ks = ConcurrentKeyspace::default(); + ks.set( + "key".into(), + Bytes::from("value"), + Some(Duration::from_millis(10)), + ); + assert!(matches!(ks.ttl("key"), TtlResult::Seconds(_))); + std::thread::sleep(Duration::from_millis(20)); + assert_eq!(ks.get("key"), None); + } + + #[test] + fn concurrent_access() { + use std::sync::Arc; + use std::thread; + + let ks = Arc::new(ConcurrentKeyspace::default()); + let mut handles = vec![]; + + // Spawn multiple threads doing concurrent sets + for i in 0..8 { + let ks = Arc::clone(&ks); + handles.push(thread::spawn(move || { + for j in 0..1000 { + let key = format!("key-{}-{}", i, j); + ks.set(key, Bytes::from("value"), None); + } + })); + } + + for h in handles { + h.join().unwrap(); + } + + assert_eq!(ks.len(), 8000); + } +} diff --git a/crates/ember-core/src/lib.rs b/crates/ember-core/src/lib.rs index 60873630..d44fd95f 100644 --- a/crates/ember-core/src/lib.rs +++ b/crates/ember-core/src/lib.rs @@ -4,6 +4,7 @@ //! Designed around a thread-per-core, shared-nothing architecture //! where each shard independently manages a partition of keys. +pub mod concurrent; pub mod engine; pub mod error; pub mod expiry; @@ -12,6 +13,7 @@ pub mod memory; pub mod shard; pub mod types; +pub use concurrent::ConcurrentKeyspace; pub use engine::{Engine, EngineConfig}; pub use error::ShardError; pub use keyspace::{ diff --git a/crates/ember-server/src/concurrent_handler.rs b/crates/ember-server/src/concurrent_handler.rs new file mode 100644 index 00000000..67ab34c0 --- /dev/null +++ b/crates/ember-server/src/concurrent_handler.rs @@ -0,0 +1,200 @@ +//! Concurrent handler that bypasses shard channels for GET/SET. +//! +//! Uses DashMap-backed ConcurrentKeyspace for lock-free multi-threaded access. +//! Falls back to sharded engine for complex commands. + +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use bytes::BytesMut; +use ember_core::{ConcurrentKeyspace, Engine, TtlResult}; +use ember_protocol::{parse_frame, Command, Frame, SetExpire}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpStream; + +use crate::server::ServerContext; +use crate::slowlog::SlowLog; + +const BUF_CAPACITY: usize = 4096; +const MAX_BUF_SIZE: usize = 64 * 1024 * 1024; +const IDLE_TIMEOUT: Duration = Duration::from_secs(300); + +/// Handles a connection using the concurrent keyspace for GET/SET. +pub async fn handle( + mut stream: TcpStream, + keyspace: Arc, + engine: Engine, // fallback for complex commands + ctx: &Arc, + slow_log: &Arc, +) -> Result<(), Box> { + stream.set_nodelay(true)?; + + let mut buf = BytesMut::with_capacity(BUF_CAPACITY); + let mut out = BytesMut::with_capacity(BUF_CAPACITY); + + loop { + if buf.len() > MAX_BUF_SIZE { + let msg = "ERR max buffer size exceeded, closing connection"; + let mut err_buf = BytesMut::new(); + Frame::Error(msg.into()).serialize(&mut err_buf); + let _ = stream.write_all(&err_buf).await; + return Ok(()); + } + + match tokio::time::timeout(IDLE_TIMEOUT, stream.read_buf(&mut buf)).await { + Ok(Ok(0)) => return Ok(()), + Ok(Ok(_)) => {} + Ok(Err(e)) => return Err(e.into()), + Err(_) => return Ok(()), + } + + out.clear(); + loop { + match parse_frame(&buf) { + Ok(Some((frame, consumed))) => { + let _ = buf.split_to(consumed); + let response = process(frame, &keyspace, &engine, ctx, slow_log).await; + response.serialize(&mut out); + } + Ok(None) => break, + Err(e) => { + let msg = format!("ERR protocol error: {e}"); + Frame::Error(msg).serialize(&mut out); + stream.write_all(&out).await?; + return Ok(()); + } + } + } + + if !out.is_empty() { + stream.write_all(&out).await?; + } + } +} + +async fn process( + frame: Frame, + keyspace: &Arc, + engine: &Engine, + ctx: &Arc, + slow_log: &Arc, +) -> Frame { + match Command::from_frame(frame) { + Ok(cmd) => { + let cmd_name = cmd.command_name(); + let needs_timing = ctx.metrics_enabled || slow_log.is_enabled(); + let start = if needs_timing { + Some(Instant::now()) + } else { + None + }; + + let response = execute_concurrent(cmd, keyspace, engine).await; + ctx.commands_processed.fetch_add(1, Ordering::Relaxed); + + if let Some(start) = start { + let elapsed = start.elapsed(); + slow_log.maybe_record(elapsed, cmd_name); + if ctx.metrics_enabled { + let is_error = matches!(&response, Frame::Error(_)); + crate::metrics::record_command(cmd_name, elapsed, is_error); + } + } + + response + } + Err(e) => Frame::Error(format!("ERR {e}")), + } +} + +/// Execute commands using concurrent keyspace for GET/SET, fallback for others. +async fn execute_concurrent( + cmd: Command, + keyspace: &Arc, + _engine: &Engine, +) -> Frame { + match cmd { + // Hot path: direct access without channels + Command::Get { key } => match keyspace.get(&key) { + Some(data) => Frame::Bulk(data), + None => Frame::Null, + }, + + Command::Set { + key, + value, + expire, + nx, + xx, + } => { + // Handle NX/XX flags + let exists = keyspace.exists(&key); + if nx && exists { + return Frame::Null; + } + if xx && !exists { + return Frame::Null; + } + + let ttl = expire.map(|e| match e { + SetExpire::Ex(secs) => Duration::from_secs(secs), + SetExpire::Px(millis) => Duration::from_millis(millis), + }); + + if keyspace.set(key, value, ttl) { + Frame::Simple("OK".into()) + } else { + Frame::Error("OOM command not allowed when used memory > 'maxmemory'".into()) + } + } + + Command::Del { keys } => { + let mut count = 0i64; + for key in keys { + if keyspace.del(&key) { + count += 1; + } + } + Frame::Integer(count) + } + + Command::Exists { keys } => { + let mut count = 0i64; + for key in keys { + if keyspace.exists(&key) { + count += 1; + } + } + Frame::Integer(count) + } + + Command::Expire { key, seconds } => { + let result = keyspace.expire(&key, seconds); + Frame::Integer(if result { 1 } else { 0 }) + } + + Command::Ttl { key } => match keyspace.ttl(&key) { + TtlResult::Seconds(s) => Frame::Integer(s as i64), + TtlResult::NoExpiry => Frame::Integer(-1), + TtlResult::NotFound => Frame::Integer(-2), + TtlResult::Milliseconds(ms) => Frame::Integer((ms / 1000) as i64), + }, + + Command::Ping(None) => Frame::Simple("PONG".into()), + Command::Ping(Some(msg)) => Frame::Bulk(msg), + Command::Echo(msg) => Frame::Bulk(msg), + + Command::DbSize => Frame::Integer(keyspace.len() as i64), + + Command::FlushDb => { + keyspace.clear(); + Frame::Simple("OK".into()) + } + + // For unsupported commands, return an error + Command::Unknown(name) => Frame::Error(format!("ERR unknown command '{name}'")), + + _ => Frame::Error("ERR command not supported in concurrent mode".into()), + } +} diff --git a/crates/ember-server/src/main.rs b/crates/ember-server/src/main.rs index 6086b2f7..ae456ceb 100644 --- a/crates/ember-server/src/main.rs +++ b/crates/ember-server/src/main.rs @@ -4,6 +4,7 @@ #[global_allocator] static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; +mod concurrent_handler; mod config; mod connection; mod metrics; @@ -68,6 +69,11 @@ struct Args { /// number of shards (worker threads). defaults to available CPU cores #[arg(long)] shards: Option, + + /// use concurrent keyspace (DashMap) instead of sharded channels. + /// experimental: bypasses channel overhead for GET/SET commands. + #[arg(long)] + concurrent: bool, } #[tokio::main] @@ -168,18 +174,37 @@ 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, + concurrent = args.concurrent, + "ember server starting..." + ); + + let result = if args.concurrent { + server::run_concurrent( + addr, + shard_count, + engine_config, + max_memory, + eviction_policy, + None, + args.metrics_port.is_some(), + slowlog_config, + ) + .await + } else { + 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..72c17b2d 100644 --- a/crates/ember-server/src/server.rs +++ b/crates/ember-server/src/server.rs @@ -8,7 +8,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::Instant; -use ember_core::{Engine, EngineConfig}; +use ember_core::{ConcurrentKeyspace, Engine, EngineConfig, EvictionPolicy}; use tokio::net::TcpListener; use tokio::sync::Semaphore; use tracing::{error, info, warn}; @@ -158,3 +158,117 @@ pub async fn run( Ok(()) } + +/// Runs the server with a concurrent keyspace (DashMap-backed). +/// +/// This mode bypasses shard channels for GET/SET operations, accessing +/// the keyspace directly from connection handlers. Falls back to the +/// sharded engine for complex commands. +#[allow(clippy::too_many_arguments)] +pub async fn run_concurrent( + addr: SocketAddr, + shard_count: usize, + config: EngineConfig, + max_memory: Option, + eviction_policy: EvictionPolicy, + max_connections: Option, + metrics_enabled: bool, + slowlog_config: SlowLogConfig, +) -> Result<(), Box> { + let aof_enabled = config + .persistence + .as_ref() + .map(|p| p.append_only) + .unwrap_or(false); + + // Create the concurrent keyspace + let keyspace = Arc::new(ConcurrentKeyspace::new(max_memory, eviction_policy)); + + // Also create the sharded engine for fallback on complex commands + let engine = Engine::with_config(shard_count, config); + + if metrics_enabled { + crate::metrics::spawn_stats_poller(engine.clone()); + } + + let listener = TcpListener::bind(addr).await?; + let max_conn = max_connections.unwrap_or(DEFAULT_MAX_CONNECTIONS); + let semaphore = Arc::new(Semaphore::new(max_conn)); + + 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)); + + info!("listening on {addr} with concurrent keyspace (max {max_conn} connections)"); + + let shutdown = tokio::signal::ctrl_c(); + tokio::pin!(shutdown); + + loop { + tokio::select! { + biased; + + _ = &mut shutdown => { + info!("shutdown signal received, draining connections..."); + break; + } + + result = listener.accept() => { + let (stream, peer) = result?; + + let permit = match semaphore.clone().try_acquire_owned() { + Ok(permit) => permit, + Err(_) => { + warn!("connection limit reached, dropping connection from {peer}"); + if metrics_enabled { + crate::metrics::on_connection_rejected(); + } + drop(stream); + continue; + } + }; + + if metrics_enabled { + crate::metrics::on_connection_accepted(); + } + ctx.connections_accepted.fetch_add(1, Ordering::Relaxed); + ctx.connections_active.fetch_add(1, Ordering::Relaxed); + + let keyspace = Arc::clone(&keyspace); + let engine = engine.clone(); + let ctx = Arc::clone(&ctx); + let slow_log = Arc::clone(&slow_log); + + tokio::spawn(async move { + if let Err(e) = crate::concurrent_handler::handle( + stream, keyspace, 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); + }); + } + } + } + + info!("waiting for active connections to close..."); + let _ = semaphore.acquire_many(max_conn as u32).await; + info!("all connections drained, shutting down"); + + Ok(()) +}