diff --git a/crates/ember-server/src/pubsub.rs b/crates/ember-server/src/pubsub.rs index 362b9e3b..b589f5c9 100644 --- a/crates/ember-server/src/pubsub.rs +++ b/crates/ember-server/src/pubsub.rs @@ -54,12 +54,7 @@ impl PubSubManager { /// Subscribe to an exact channel. Returns a receiver for messages /// on that channel. pub fn subscribe(&self, channel: &str) -> broadcast::Receiver { - let entry = self.channels.entry(channel.to_string()).or_insert_with(|| { - let (tx, _) = broadcast::channel(CHANNEL_CAPACITY); - tx - }); - self.subscription_count.fetch_add(1, Ordering::Relaxed); - entry.subscribe() + self.subscribe_to(&self.channels, channel) } /// Unsubscribe from an exact channel. Returns true if the channel @@ -68,23 +63,28 @@ impl PubSubManager { /// Note: the actual receiver is dropped by the caller. This just /// cleans up empty channels and adjusts the subscription count. pub fn unsubscribe(&self, channel: &str) -> bool { - if let Some(entry) = self.channels.get(channel) { - self.subscription_count.fetch_sub(1, Ordering::Relaxed); - // if no receivers left, remove the channel entirely - if entry.receiver_count() <= 1 { - drop(entry); - self.channels.remove(channel); - } - true - } else { - false - } + self.unsubscribe_from(&self.channels, channel) } /// Subscribe to a glob pattern. Returns a receiver for messages /// matching the pattern. pub fn psubscribe(&self, pattern: &str) -> broadcast::Receiver { - let entry = self.patterns.entry(pattern.to_string()).or_insert_with(|| { + self.subscribe_to(&self.patterns, pattern) + } + + /// Unsubscribe from a pattern. Returns true if the pattern existed + /// in the registry. + pub fn punsubscribe(&self, pattern: &str) -> bool { + self.unsubscribe_from(&self.patterns, pattern) + } + + /// Subscribes to a key in the given map (channels or patterns). + fn subscribe_to( + &self, + map: &DashMap>, + key: &str, + ) -> broadcast::Receiver { + let entry = map.entry(key.to_string()).or_insert_with(|| { let (tx, _) = broadcast::channel(CHANNEL_CAPACITY); tx }); @@ -92,14 +92,18 @@ impl PubSubManager { entry.subscribe() } - /// Unsubscribe from a pattern. Returns true if the pattern existed - /// in the registry. - pub fn punsubscribe(&self, pattern: &str) -> bool { - if let Some(entry) = self.patterns.get(pattern) { + /// Unsubscribes from a key in the given map. Returns true if the + /// key existed. Removes the entry when no receivers remain. + fn unsubscribe_from( + &self, + map: &DashMap>, + key: &str, + ) -> bool { + if let Some(entry) = map.get(key) { self.subscription_count.fetch_sub(1, Ordering::Relaxed); if entry.receiver_count() <= 1 { drop(entry); - self.patterns.remove(pattern); + map.remove(key); } true } else { diff --git a/crates/ember-server/src/server.rs b/crates/ember-server/src/server.rs index e677172c..62624e2b 100644 --- a/crates/ember-server/src/server.rs +++ b/crates/ember-server/src/server.rs @@ -10,8 +10,8 @@ use std::time::{Duration, Instant}; use ember_core::{ConcurrentKeyspace, Engine, EngineConfig, EvictionPolicy}; use tokio::io::AsyncWriteExt; -use tokio::net::TcpListener; -use tokio::sync::Semaphore; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tokio_rustls::TlsAcceptor; use tracing::{error, info, warn}; @@ -159,37 +159,16 @@ pub async fn run( result = listener.accept() => { let (stream, peer) = result?; - // disable Nagle's algorithm — cache servers need low-latency writes. - // done here before passing to handler so TLS streams also benefit. - if let Err(e) = stream.set_nodelay(true) { - warn!("failed to set TCP_NODELAY: {e}"); - } - - // protected mode: reject non-loopback connections when no - // password is set and the server is bound to a public address if is_protected_mode_violation(&ctx, &peer) { reject_protected_mode(stream).await; continue; } - 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; - } + let permit = match accept_connection(&stream, peer, &ctx, &semaphore) { + Some(p) => p, + None => 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 engine = engine.clone(); let ctx = Arc::clone(&ctx); let slow_log = Arc::clone(&slow_log); @@ -199,11 +178,7 @@ pub async fn run( if let Err(e) = connection::handle(stream, engine, &ctx, &slow_log, &pubsub).await { error!("connection error from {peer}: {e}"); } - ctx.connections_active.fetch_sub(1, Ordering::Relaxed); - if ctx.metrics_enabled { - crate::metrics::on_connection_closed(); - } - // permit is dropped here, releasing the slot + on_connection_done(&ctx); drop(permit); }); } @@ -212,81 +187,112 @@ pub async fn run( result = tls_accept() => { let (stream, peer, acceptor) = result?; - if let Err(e) = stream.set_nodelay(true) { - warn!("failed to set TCP_NODELAY: {e}"); - } - - // protected mode check on TLS connections too if is_protected_mode_violation(&ctx, &peer) { reject_protected_mode(stream).await; continue; } - let permit = match semaphore.clone().try_acquire_owned() { - Ok(permit) => permit, - Err(_) => { - warn!("connection limit reached, dropping TLS connection from {peer}"); - if metrics_enabled { - crate::metrics::on_connection_rejected(); - } - drop(stream); - continue; - } + let permit = match accept_connection(&stream, peer, &ctx, &semaphore) { + Some(p) => p, + None => 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 engine = engine.clone(); let ctx = Arc::clone(&ctx); let slow_log = Arc::clone(&slow_log); let pubsub = Arc::clone(&pubsub); tokio::spawn(async move { - // perform TLS handshake with timeout to prevent slowloris - let handshake = tokio::time::timeout( - Duration::from_secs(10), - acceptor.accept(stream), - ); - match handshake.await { - Ok(Ok(tls_stream)) => { - if let Err(e) = connection::handle(tls_stream, engine, &ctx, &slow_log, &pubsub).await { - error!("TLS connection error from {peer}: {e}"); - } - } - Ok(Err(e)) => { - warn!("TLS handshake failed from {peer}: {e}"); + if let Some(tls_stream) = tls_handshake(acceptor, stream, peer).await { + if let Err(e) = connection::handle(tls_stream, engine, &ctx, &slow_log, &pubsub).await { + error!("TLS connection error from {peer}: {e}"); } - Err(_) => { - warn!("TLS handshake timed out from {peer}"); - } - } - ctx.connections_active.fetch_sub(1, Ordering::Relaxed); - if ctx.metrics_enabled { - crate::metrics::on_connection_closed(); } + on_connection_done(&ctx); drop(permit); }); } } } - // wait for all connection handlers to finish, with a timeout + drain_connections(&semaphore, max_conn).await; + Ok(()) +} + +/// Sets nodelay, acquires a semaphore permit, and updates connection metrics. +/// Returns `None` if the connection limit was reached — the caller should +/// `continue`. Protected mode must be checked by the caller before this. +fn accept_connection( + stream: &TcpStream, + peer: SocketAddr, + ctx: &Arc, + semaphore: &Arc, +) -> Option { + if let Err(e) = stream.set_nodelay(true) { + warn!("failed to set TCP_NODELAY: {e}"); + } + + let permit = match semaphore.clone().try_acquire_owned() { + Ok(permit) => permit, + Err(_) => { + warn!("connection limit reached, dropping connection from {peer}"); + if ctx.metrics_enabled { + crate::metrics::on_connection_rejected(); + } + return None; + } + }; + + 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); + + Some(permit) +} + +/// Performs TLS handshake with a 10-second timeout to prevent slowloris attacks. +/// Returns `None` on handshake failure or timeout. +async fn tls_handshake( + acceptor: TlsAcceptor, + stream: TcpStream, + peer: SocketAddr, +) -> Option> { + let handshake = tokio::time::timeout(Duration::from_secs(10), acceptor.accept(stream)); + match handshake.await { + Ok(Ok(tls_stream)) => Some(tls_stream), + Ok(Err(e)) => { + warn!("TLS handshake failed from {peer}: {e}"); + None + } + Err(_) => { + warn!("TLS handshake timed out from {peer}"); + None + } + } +} + +/// Decrements active connection count and records the close metric. +fn on_connection_done(ctx: &ServerContext) { + ctx.connections_active.fetch_sub(1, Ordering::Relaxed); + if ctx.metrics_enabled { + crate::metrics::on_connection_closed(); + } +} + +/// Waits for all connections to drain with a 30-second timeout. +async fn drain_connections(semaphore: &Arc, max_conn: usize) { info!("waiting for active connections to close..."); let drain = semaphore.acquire_many(max_conn as u32); - match tokio::time::timeout(std::time::Duration::from_secs(30), drain).await { + match tokio::time::timeout(Duration::from_secs(30), drain).await { Ok(_) => info!("all connections drained, shutting down"), Err(_) => warn!("shutdown timeout after 30s, forcing exit"), } - - Ok(()) } /// Sends the protected mode rejection message and closes the connection. -async fn reject_protected_mode(mut stream: tokio::net::TcpStream) { +async fn reject_protected_mode(mut stream: TcpStream) { let msg = "-DENIED Ember is running in protected mode \ because no password is set. In this mode \ connections are only accepted from the loopback \ @@ -395,33 +401,16 @@ pub async fn run_concurrent( result = listener.accept() => { let (stream, peer) = result?; - if let Err(e) = stream.set_nodelay(true) { - warn!("failed to set TCP_NODELAY: {e}"); - } - if is_protected_mode_violation(&ctx, &peer) { reject_protected_mode(stream).await; continue; } - 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; - } + let permit = match accept_connection(&stream, peer, &ctx, &semaphore) { + Some(p) => p, + None => 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); @@ -434,10 +423,7 @@ pub async fn run_concurrent( ).await { error!("connection error from {peer}: {e}"); } - ctx.connections_active.fetch_sub(1, Ordering::Relaxed); - if ctx.metrics_enabled { - crate::metrics::on_connection_closed(); - } + on_connection_done(&ctx); drop(permit); }); } @@ -446,33 +432,16 @@ pub async fn run_concurrent( result = tls_accept() => { let (stream, peer, acceptor) = result?; - if let Err(e) = stream.set_nodelay(true) { - warn!("failed to set TCP_NODELAY: {e}"); - } - if is_protected_mode_violation(&ctx, &peer) { reject_protected_mode(stream).await; continue; } - let permit = match semaphore.clone().try_acquire_owned() { - Ok(permit) => permit, - Err(_) => { - warn!("connection limit reached, dropping TLS connection from {peer}"); - if metrics_enabled { - crate::metrics::on_connection_rejected(); - } - drop(stream); - continue; - } + let permit = match accept_connection(&stream, peer, &ctx, &semaphore) { + Some(p) => p, + None => 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); @@ -480,42 +449,21 @@ pub async fn run_concurrent( let pubsub = Arc::clone(&pubsub); tokio::spawn(async move { - let handshake = tokio::time::timeout( - Duration::from_secs(10), - acceptor.accept(stream), - ); - match handshake.await { - Ok(Ok(tls_stream)) => { - if let Err(e) = crate::concurrent_handler::handle( - tls_stream, keyspace, engine, &ctx, &slow_log, &pubsub - ).await { - error!("TLS connection error from {peer}: {e}"); - } - } - Ok(Err(e)) => { - warn!("TLS handshake failed from {peer}: {e}"); - } - Err(_) => { - warn!("TLS handshake timed out from {peer}"); + if let Some(tls_stream) = tls_handshake(acceptor, stream, peer).await { + if let Err(e) = crate::concurrent_handler::handle( + tls_stream, keyspace, engine, &ctx, &slow_log, &pubsub + ).await { + error!("TLS connection error from {peer}: {e}"); } } - ctx.connections_active.fetch_sub(1, Ordering::Relaxed); - if ctx.metrics_enabled { - crate::metrics::on_connection_closed(); - } + on_connection_done(&ctx); drop(permit); }); } } } - info!("waiting for active connections to close..."); - let drain = semaphore.acquire_many(max_conn as u32); - match tokio::time::timeout(std::time::Duration::from_secs(30), drain).await { - Ok(_) => info!("all connections drained, shutting down"), - Err(_) => warn!("shutdown timeout after 30s, forcing exit"), - } - + drain_connections(&semaphore, max_conn).await; Ok(()) } diff --git a/crates/ember-server/src/slowlog.rs b/crates/ember-server/src/slowlog.rs index da3ccc2b..72195362 100644 --- a/crates/ember-server/src/slowlog.rs +++ b/crates/ember-server/src/slowlog.rs @@ -86,14 +86,7 @@ impl SlowLog { return; } - let mut inner = match self.inner.lock() { - Ok(guard) => guard, - Err(poisoned) => { - let mut guard = poisoned.into_inner(); - guard.entries.clear(); - guard - } - }; + let mut inner = self.lock_or_clear(); let id = inner.next_id; inner.next_id += 1; @@ -114,10 +107,7 @@ impl SlowLog { /// /// If `count` is `None`, returns all entries. pub fn get(&self, count: Option) -> Vec { - let inner = match self.inner.lock() { - Ok(guard) => guard, - Err(poisoned) => poisoned.into_inner(), - }; + let inner = self.lock_inner(); let n = count .unwrap_or(inner.entries.len()) .min(inner.entries.len()); @@ -126,19 +116,34 @@ impl SlowLog { /// Returns the number of entries currently in the log. pub fn len(&self) -> usize { - match self.inner.lock() { - Ok(guard) => guard.entries.len(), - Err(poisoned) => poisoned.into_inner().entries.len(), - } + self.lock_inner().entries.len() } /// Clears all entries from the log. pub fn reset(&self) { - let mut inner = match self.inner.lock() { + self.lock_inner().entries.clear(); + } + + /// Acquires the lock, clearing entries on poison. Used by writes + /// where we'd rather start fresh than propagate a panic. + fn lock_or_clear(&self) -> std::sync::MutexGuard<'_, SlowLogInner> { + match self.inner.lock() { + Ok(guard) => guard, + Err(poisoned) => { + let mut guard = poisoned.into_inner(); + guard.entries.clear(); + guard + } + } + } + + /// Acquires the lock, recovering silently on poison. Used by reads + /// and non-critical writes (reset, len). + fn lock_inner(&self) -> std::sync::MutexGuard<'_, SlowLogInner> { + match self.inner.lock() { Ok(guard) => guard, Err(poisoned) => poisoned.into_inner(), - }; - inner.entries.clear(); + } } }