diff --git a/.planning b/.planning index ad071500..62fd2d1a 160000 --- a/.planning +++ b/.planning @@ -1 +1 @@ -Subproject commit ad0715006c0d38098d00af5f9556920667a0417d +Subproject commit 62fd2d1afddb06ca0c268770af964bbdf6b3a086 diff --git a/CHANGELOG.md b/CHANGELOG.md index 664183be..596b6130 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added — Dispatch Observability (Phase 177) + +- **`moon_dispatch_path_total{path=...}` Prometheus counter**: four-way classification of every command by shard-routing decision — `local_inline` (SIMD fast path), `local` (standard local branch), `cross_read_fast` (RwLock shared-read bypass of SPSC), `cross_spsc` (deferred cross-shard write via `PipelineBatchSlotted`). Ratio `cross_spsc / Σ` is the ground-truth signal for dispatch-layer optimization work. Zero-allocation hot-path overhead (`&'static str` labels, `#[inline]` with early-return on `!METRICS_INITIALIZED`). Verified on macOS + Linux: counter sums close exactly to driven traffic, no overcount. + +### Fixed — CI Hygiene + +- **`tests/pipeline_auto_index.rs`**: tighten outer cfg from `runtime-tokio` to `all(runtime-tokio, text-index)` so the file compiles to zero tests when text-index is disabled. Previously the file compiled but the FT.SEARCH text fast path was `#[cfg]`-ed out, causing `@name:corpus` queries to fall through to the KNN-only parser and panic with "invalid KNN query syntax". +- **4 FT unwraps**: add inline `#[allow(clippy::unwrap_used)]` with invariant justifications in `vector_search/ft_text_search.rs` (3 sites inside `apply_post_processing` where `do_summarize` / `do_highlight` implies the Option is Some) and `handler_monoio/ft.rs:165` (`is_text` was derived from `query_bytes.as_ref().map_or(false, _)`). Restores the audit-unwrap baseline to 0. + ### Changed - **`text-index` is now a default feature.** BM25 full-text search (`FT.SEARCH` BM25 mode), `FT.AGGREGATE`, and three-way RRF hybrid fusion are included in all standard builds. No longer requires `--features text-index`. To exclude it (e.g. minimal embedded builds): `--no-default-features --features runtime-monoio,jemalloc,graph`. diff --git a/src/admin/metrics_setup.rs b/src/admin/metrics_setup.rs index c411b8f2..f4a6ad15 100644 --- a/src/admin/metrics_setup.rs +++ b/src/admin/metrics_setup.rs @@ -490,6 +490,56 @@ pub fn record_spsc_drain(shard_id: usize, count: u64) { histogram!("moon_spsc_drain_batch_size", "shard" => shard).record(count as f64); } +// ── Dispatch routing counters (Phase 177, Step 6) ─────────────────────── +// Three-way split of the connection hot path so we can quantify what +// fraction of traffic is hitting the expensive cross-shard SPSC dispatch +// vs. the free local / shared-read fast paths. Ratio of these counters +// is the ground-truth signal for validating dispatch-layer optimizations +// (HotShardMessage split, outbox batching, waker relay fusion). + +/// Command executed on the connection's own shard (no cross-thread hop). +#[inline] +pub fn record_dispatch_local() { + if !METRICS_INITIALIZED.load(Ordering::Relaxed) { + return; + } + counter!("moon_dispatch_path_total", "path" => "local").increment(1); +} + +/// Command executed on a remote shard via the shared-read fast path +/// (RwLock read on the target shard's database, no SPSC message). +#[inline] +pub fn record_dispatch_cross_read_fastpath() { + if !METRICS_INITIALIZED.load(Ordering::Relaxed) { + return; + } + counter!("moon_dispatch_path_total", "path" => "cross_read_fast").increment(1); +} + +/// Command deferred to cross-shard SPSC dispatch (the slow path). +/// Recorded when a command is enqueued into a `remote_groups` bucket that +/// will be flushed as a `PipelineBatchSlotted` message. +#[inline] +pub fn record_dispatch_cross_spsc() { + if !METRICS_INITIALIZED.load(Ordering::Relaxed) { + return; + } + counter!("moon_dispatch_path_total", "path" => "cross_spsc").increment(1); +} + +/// Command handled by the inline GET/SET fast path +/// (`try_inline_dispatch_loop` in `server/conn/blocking.rs`) — the hottest +/// local branch, which bypasses the standard frame-by-frame routing and +/// therefore the three counters above. Recorded in a single batch increment +/// per dispatch loop to keep the call site out of the per-command hot path. +#[inline] +pub fn record_dispatch_local_inline(count: u64) { + if count == 0 || !METRICS_INITIALIZED.load(Ordering::Relaxed) { + return; + } + counter!("moon_dispatch_path_total", "path" => "local_inline").increment(count); +} + // ── Vector search metrics (v0.1.6) ───────────────────────────────────── /// Record a cache hit for FT.CACHESEARCH. @@ -859,3 +909,29 @@ pub fn spawn_metrics_publisher() { } }); } + +#[cfg(test)] +mod tests { + use super::*; + + // Smoke tests for the dispatch-path counters added in Phase 177, Step 6. + // A full assertion on the prometheus state would require initialising a + // recorder and a scraping harness — out of scope here. These tests just + // pin the contract that the helpers are safe to call on the hot path + // before `init_metrics` has run and therefore never panic or allocate + // unexpectedly when the exporter is disabled (admin_port = 0). + + #[test] + fn dispatch_path_counters_no_op_before_init() { + // METRICS_INITIALIZED starts false; all three helpers must early-return. + // We just assert they do not panic. The absence of a global recorder + // means counter!() would otherwise be a no-op, but the guard is what + // we actually care about: no string allocation, no label churn. + assert!(!METRICS_INITIALIZED.load(Ordering::Relaxed)); + record_dispatch_local(); + record_dispatch_cross_read_fastpath(); + record_dispatch_cross_spsc(); + record_dispatch_local_inline(0); // count == 0 must short-circuit even when init + record_dispatch_local_inline(7); + } +} diff --git a/src/command/graph/graph_read.rs b/src/command/graph/graph_read.rs index ab0d504a..ece27276 100644 --- a/src/command/graph/graph_read.rs +++ b/src/command/graph/graph_read.rs @@ -73,11 +73,11 @@ fn json_to_graph_value(v: &serde_json::Value) -> cypher::executor::Value { serde_json::Value::Array(arr) => { cypher::executor::Value::List(arr.iter().map(json_to_graph_value).collect()) } - serde_json::Value::Object(map) => { - cypher::executor::Value::Map( - map.iter().map(|(k, v)| (k.clone(), json_to_graph_value(v))).collect() - ) - } + serde_json::Value::Object(map) => cypher::executor::Value::Map( + map.iter() + .map(|(k, v)| (k.clone(), json_to_graph_value(v))) + .collect(), + ), } } diff --git a/src/command/vector_search/ft_text_search.rs b/src/command/vector_search/ft_text_search.rs index b04871e0..039b63f4 100644 --- a/src/command/vector_search/ft_text_search.rs +++ b/src/command/vector_search/ft_text_search.rs @@ -758,6 +758,8 @@ pub fn apply_post_processing( // SUMMARIZE first (extracts passage), then HIGHLIGHT (wraps terms). let processed = if do_summarize { + // Safe: `do_summarize = summarize_opts.map_or(false, ..)` so do_summarize=true ⇒ Some. + #[allow(clippy::unwrap_used)] let opts = summarize_opts.unwrap(); let summarized = summarize_field( original_text, @@ -768,6 +770,8 @@ pub fn apply_post_processing( &opts.separator, ); if do_highlight { + // Safe: `do_highlight = highlight_opts.map_or(false, ..)` so do_highlight=true ⇒ Some. + #[allow(clippy::unwrap_used)] let hopts = highlight_opts.unwrap(); highlight_field( &summarized, @@ -781,6 +785,9 @@ pub fn apply_post_processing( summarized } } else { + // Safe: reached only when do_summarize=false but the `!do_highlight && !do_summarize` + // guard above ensured at least one was true, so do_highlight=true ⇒ Some. + #[allow(clippy::unwrap_used)] let hopts = highlight_opts.unwrap(); highlight_field( original_text, diff --git a/src/persistence/wal_v3/segment.rs b/src/persistence/wal_v3/segment.rs index a6929e99..599c169b 100644 --- a/src/persistence/wal_v3/segment.rs +++ b/src/persistence/wal_v3/segment.rs @@ -131,22 +131,21 @@ impl WalWriterV3 { lsn } - /// Flush the in-memory buffer to disk and fsync. + /// Write the in-memory buffer to the OS page cache (no fsync). /// - /// After this returns, all appended records are durable on stable storage. - pub fn flush_sync(&mut self) -> std::io::Result<()> { + /// Data reaches the kernel but is NOT guaranteed durable until + /// `sync_data()` is called (typically on the 1s timer or shutdown). + fn flush_write(&mut self) -> std::io::Result<()> { if self.buf.is_empty() { return Ok(()); } - // Check if rotation is needed before writing if self.write_offset + self.buf.len() as u64 > self.segment_size { self.rotate_segment()?; } if let Some(ref mut file) = self.current_file { file.write_all(&self.buf)?; - file.sync_data()?; self.write_offset += self.buf.len() as u64; self.buf.clear(); } @@ -154,15 +153,38 @@ impl WalWriterV3 { Ok(()) } - /// Flush if buffer exceeds a threshold (matches v2 pattern). + /// Flush the in-memory buffer to disk and fsync. + /// + /// After this returns, all appended records are durable on stable storage. + pub fn flush_sync(&mut self) -> std::io::Result<()> { + self.flush_write()?; + if let Some(ref mut file) = self.current_file { + file.sync_data()?; + } + Ok(()) + } + + /// Flush if buffer exceeds a threshold — write only, no fsync. + /// + /// Matches WAL v2 pattern: frequent writes to OS page cache, + /// durable sync deferred to the 1s timer (`sync_data`). pub fn flush_if_needed(&mut self) -> std::io::Result<()> { if self.buf.len() >= 4096 { - self.flush_sync() + self.flush_write() } else { Ok(()) } } + /// Fsync without writing — call after `flush_write` / `flush_if_needed` + /// to make all previously written records durable. + pub fn sync_data(&mut self) -> std::io::Result<()> { + if let Some(ref mut file) = self.current_file { + file.sync_data()?; + } + Ok(()) + } + /// Return the current (next-to-be-assigned) LSN. #[inline] pub fn current_lsn(&self) -> u64 { diff --git a/src/server/conn/handler_monoio.rs b/src/server/conn/handler_monoio.rs deleted file mode 100644 index d5fb38df..00000000 --- a/src/server/conn/handler_monoio.rs +++ /dev/null @@ -1,4206 +0,0 @@ -// Note: some imports/variables may be conditionally used across feature flags -//! Monoio connection handler using ownership-based I/O (AsyncReadRent/AsyncWriteRent). -//! -//! Extracted from `server/connection.rs` (Plan 48-02). - -use crate::runtime::cancel::CancellationToken; -use crate::runtime::channel; -use bytes::{Bytes, BytesMut}; -use ringbuf::traits::Producer; -use std::cell::RefCell; -use std::collections::HashMap; -use std::rc::Rc; -use std::sync::Arc; - -use crate::command::connection as conn_cmd; -use crate::command::metadata; -use crate::command::mq::{ - ERR_MQ_NOT_DURABLE, ERR_MQ_UNKNOWN_SUB, parse_mq_subcommand, validate_mq_ack, - validate_mq_create, validate_mq_dlqlen, validate_mq_pop, validate_mq_publish, validate_mq_push, - validate_mq_trigger, -}; -use crate::command::temporal::{ - ERR_ENTITY_NOT_FOUND, ERR_GRAPH_NOT_FOUND, capture_wall_ms, is_temporal_invalidate, - is_temporal_snapshot_at, validate_invalidate, validate_snapshot_at, -}; -use crate::command::transaction::{ - ERR_MULTI_TXN_CONFLICT, is_txn_abort, is_txn_begin, is_txn_commit, txn_abort_validate, - txn_begin_validate, txn_commit_validate, -}; -use crate::command::workspace::{ - ERR_WS_ALREADY_BOUND, ERR_WS_NOT_FOUND, ERR_WS_UNKNOWN_SUB, parse_workspace_id_from_bytes, - parse_ws_subcommand, validate_ws_auth, validate_ws_create, validate_ws_drop, validate_ws_info, - validate_ws_list, -}; -use crate::command::{DispatchResult, dispatch, dispatch_read, is_dispatch_read_supported}; -use crate::mq::is_mq_command; -use crate::persistence::aof::{self, AofMessage}; -use crate::protocol::Frame; -use crate::pubsub::subscriber::Subscriber; -use crate::pubsub::{self}; -use crate::shard::dispatch::{ShardMessage, key_to_shard}; -use crate::shard::mesh::ChannelMesh; -use crate::storage::eviction::{try_evict_if_needed, try_evict_if_needed_async_spill}; -use crate::storage::stream::StreamId; -use crate::tracking::TrackingState; -use crate::transaction::CrossStoreTxn; -use crate::workspace::{ - WorkspaceId, is_ws_command, strip_workspace_prefix_from_response, workspace_rewrite_args, -}; - -use super::affinity::MigratedConnectionState; -use super::shared::resolve_ft_search_as_of_lsn; -use super::{ - apply_resp3_conversion, convert_blocking_to_nonblocking, execute_transaction_sharded, - extract_bytes, extract_command, extract_primary_key, handle_blocking_command_monoio, - handle_config, is_multi_key_command, propagate_subscription, try_inline_dispatch_loop, - unpropagate_subscription, -}; -use crate::framevec; -use crate::server::codec::RespCodec; -// ResponseSlotPool NOT used on monoio — its AtomicWaker doesn't cross -// monoio's single-threaded (!Send) executor boundary. Use oneshot channels. - -/// Result of `handle_connection_sharded_monoio` execution. -/// -/// Same purpose as the Tokio handler's `HandlerResult`: the generic handler cannot -/// perform FD extraction, so it returns the stream when migration is triggered. -#[cfg(feature = "runtime-monoio")] -pub enum MonoioHandlerResult { - /// Normal connection close. - Done, - /// Migration triggered: caller should extract raw FD and send via SPSC. - MigrateConnection { - state: MigratedConnectionState, - target_shard: usize, - }, -} - -/// Monoio connection handler using ownership-based I/O (AsyncReadRent/AsyncWriteRent). -/// -/// Reads RESP frames from the TCP stream, dispatches commands through the same -/// `crate::command::dispatch()` path as the tokio handler, and writes responses back. -/// -/// MVP scope: SET/GET/PING/QUIT/SELECT/DEL/COMMAND and all other commands supported -/// by `dispatch()`. Skips pub/sub, blocking, tracking, cluster, replication, and ACL -/// enforcement -- those parameters are accepted but unused for future wiring. -/// -/// Key difference from tokio path: monoio's `stream.read(buf)` takes ownership of the -/// buffer and returns `(Result, buf)`. We use a `Vec` intermediate for the -/// read since monoio's IoBufMut is implemented for Vec, then copy into BytesMut -/// for codec parsing. -#[cfg(feature = "runtime-monoio")] -#[tracing::instrument(skip_all, level = "debug")] -pub(crate) async fn handle_connection_sharded_monoio< - S: monoio::io::AsyncReadRent + monoio::io::AsyncWriteRent, ->( - mut stream: S, - peer_addr: String, - ctx: &super::core::ConnectionContext, - shutdown: CancellationToken, - client_id: u64, - can_migrate: bool, - initial_read_buf: BytesMut, - pending_wakers: Rc>>, - migrated_state: Option<&MigratedConnectionState>, -) -> (MonoioHandlerResult, Option) { - use monoio::io::AsyncWriteRentExt; - - // NOTE: do NOT call record_connection_opened() here — the caller - // (conn_accept.rs) already increments via try_accept_connection(). - - let mut read_buf = if initial_read_buf.is_empty() { - BytesMut::with_capacity(8192) - } else { - let mut buf = initial_read_buf; - buf.reserve(8192); - buf - }; - let mut write_buf = BytesMut::with_capacity(8192); - let mut codec = RespCodec::default(); - let mut conn = super::core::ConnectionState::new( - client_id, - peer_addr.clone(), - &ctx.requirepass, - ctx.shard_id, - ctx.num_shards, - can_migrate, - ctx.runtime_config.read().acllog_max_len, - migrated_state, - ); - conn.refresh_acl_cache(&ctx.acl_table); - let db_count = ctx.shard_databases.db_count(); - - // Register in global client registry for CLIENT LIST/INFO/KILL. - crate::client_registry::register( - client_id, - peer_addr.clone(), - conn.current_user.clone(), - ctx.shard_id, - ); - struct RegistryGuard(u64); - impl Drop for RegistryGuard { - fn drop(&mut self) { - crate::client_registry::deregister(self.0); - } - } - let _registry_guard = RegistryGuard(client_id); - - // Functions API registry (per-connection, lazy init) — kept as local because Rc> is !Send - let func_registry = Rc::new(RefCell::new(crate::scripting::FunctionRegistry::new())); - - // Pre-allocate read buffer outside the loop to avoid per-read heap allocation. - // Monoio's ownership I/O takes ownership and returns the buffer, so we reassign. - let mut tmp_buf = vec![0u8; 8192]; - - // Client idle timeout: 0 = disabled (read once, avoid lock on hot path) - let idle_timeout_secs = ctx.runtime_config.read().timeout; - let idle_timeout = if idle_timeout_secs > 0 { - Some(std::time::Duration::from_secs(idle_timeout_secs)) - } else { - None - }; - - // Pre-allocate batch containers outside the loop to avoid per-batch heap allocation. - // These are cleared and reused each iteration instead of being recreated. - let mut responses: Vec = Vec::with_capacity(64); - let mut remote_groups: HashMap< - usize, - Vec<(usize, std::sync::Arc, Option, Bytes)>, - > = HashMap::with_capacity(ctx.num_shards); - let mut reply_futures: Vec<(Vec<(usize, Option, Bytes)>, usize)> = - Vec::with_capacity(ctx.num_shards); - - // Pre-allocate frames Vec outside the loop; reused via .clear() each iteration. - let mut frames: Vec = Vec::with_capacity(64); - - loop { - // Check if CLIENT KILL targeted this connection - if crate::client_registry::is_killed(client_id) { - break; - } - - // Subscriber mode: bidirectional select on client commands + published messages - if conn.subscription_count > 0 { - #[allow(clippy::unwrap_used)] - // conn.pubsub_rx is always Some when conn.subscription_count > 0 - let rx = conn.pubsub_rx.as_ref().unwrap(); - let sub_tmp_buf = vec![0u8; 8192]; - monoio::select! { - read_result = stream.read(sub_tmp_buf) => { - let (result, buf) = read_result; - match result { - Ok(0) => { - // Client half-closed — break out of loop. - // Stream drop (end of function) triggers monoio's cleanup. - break; - } - Ok(n) => { - read_buf.extend_from_slice(&buf[..n]); - // Parse frames from buffer - loop { - match codec.decode_frame(&mut read_buf) { - Ok(Some(frame)) => { - if let Some((cmd, cmd_args)) = extract_command(&frame) { - match cmd { - _ if cmd.eq_ignore_ascii_case(b"SUBSCRIBE") => { - if cmd_args.is_empty() { - let err = Frame::Error(Bytes::from_static(b"ERR wrong number of arguments for 'subscribe' command")); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&err, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if wr.is_err() { return (MonoioHandlerResult::Done, None); } - continue; - } - for arg in cmd_args { - if let Some(channel) = extract_bytes(arg) { - // ACL channel permission check - let denied = { - #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable - let acl_guard = ctx.acl_table.read().unwrap(); - acl_guard.check_channel_permission(&conn.current_user, channel.as_ref()) - }; - if let Some(deny_reason) = denied { - let err = Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason))); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&err, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if wr.is_err() { return (MonoioHandlerResult::Done, None); } - continue; - } - #[allow(clippy::unwrap_used)] // conn.pubsub_tx is always Some when in subscriber mode - let sub = Subscriber::with_protocol( - conn.pubsub_tx.clone().unwrap(), - conn.subscriber_id, - conn.protocol_version >= 3, - ); - ctx.pubsub_registry.write().subscribe(channel.clone(), sub); - propagate_subscription(&ctx.all_remote_sub_maps, &channel, ctx.shard_id, ctx.num_shards, false); - conn.subscription_count += 1; - // Register pub/sub affinity for this client IP - if conn.subscription_count == 1 { - if let Ok(addr) = peer_addr.parse::() { - ctx.pubsub_affinity.write().register(addr.ip(), ctx.shard_id); - } - } - let resp = pubsub::subscribe_response(&channel, conn.subscription_count); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&resp, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if wr.is_err() { return (MonoioHandlerResult::Done, None); } - } - } - } - _ if cmd.eq_ignore_ascii_case(b"UNSUBSCRIBE") => { - if cmd_args.is_empty() { - let removed = ctx.pubsub_registry.write().unsubscribe_all(conn.subscriber_id); - for ch in &removed { - unpropagate_subscription(&ctx.all_remote_sub_maps, ch, ctx.shard_id, ctx.num_shards, false); - } - if removed.is_empty() { - conn.subscription_count = ctx.pubsub_registry.read().total_subscription_count(conn.subscriber_id); - let resp = pubsub::unsubscribe_response(&Bytes::from_static(b""), conn.subscription_count); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&resp, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if wr.is_err() { return (MonoioHandlerResult::Done, None); } - } else { - for ch in &removed { - conn.subscription_count = conn.subscription_count.saturating_sub(1); - let resp = pubsub::unsubscribe_response(ch, conn.subscription_count); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&resp, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if wr.is_err() { return (MonoioHandlerResult::Done, None); } - } - } - } else { - for arg in cmd_args { - if let Some(channel) = extract_bytes(arg) { - ctx.pubsub_registry.write().unsubscribe(channel.as_ref(), conn.subscriber_id); - unpropagate_subscription(&ctx.all_remote_sub_maps, &channel, ctx.shard_id, ctx.num_shards, false); - conn.subscription_count = conn.subscription_count.saturating_sub(1); - let resp = pubsub::unsubscribe_response(&channel, conn.subscription_count); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&resp, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if wr.is_err() { return (MonoioHandlerResult::Done, None); } - } - } - } - } - _ if cmd.eq_ignore_ascii_case(b"PSUBSCRIBE") => { - if cmd_args.is_empty() { - let err = Frame::Error(Bytes::from_static(b"ERR wrong number of arguments for 'psubscribe' command")); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&err, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if wr.is_err() { return (MonoioHandlerResult::Done, None); } - continue; - } - for arg in cmd_args { - if let Some(pattern) = extract_bytes(arg) { - // ACL channel permission check - let denied = { - #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable - let acl_guard = ctx.acl_table.read().unwrap(); - acl_guard.check_channel_permission(&conn.current_user, pattern.as_ref()) - }; - if let Some(deny_reason) = denied { - let err = Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason))); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&err, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if wr.is_err() { return (MonoioHandlerResult::Done, None); } - continue; - } - #[allow(clippy::unwrap_used)] // conn.pubsub_tx is always Some when in subscriber mode - let sub = Subscriber::with_protocol( - conn.pubsub_tx.clone().unwrap(), - conn.subscriber_id, - conn.protocol_version >= 3, - ); - ctx.pubsub_registry.write().psubscribe(pattern.clone(), sub); - propagate_subscription(&ctx.all_remote_sub_maps, &pattern, ctx.shard_id, ctx.num_shards, true); - conn.subscription_count += 1; - // Register pub/sub affinity for this client IP - if conn.subscription_count == 1 { - if let Ok(addr) = peer_addr.parse::() { - ctx.pubsub_affinity.write().register(addr.ip(), ctx.shard_id); - } - } - let resp = pubsub::psubscribe_response(&pattern, conn.subscription_count); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&resp, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if wr.is_err() { return (MonoioHandlerResult::Done, None); } - } - } - } - _ if cmd.eq_ignore_ascii_case(b"PUNSUBSCRIBE") => { - if cmd_args.is_empty() { - let removed = ctx.pubsub_registry.write().punsubscribe_all(conn.subscriber_id); - for pat in &removed { - unpropagate_subscription(&ctx.all_remote_sub_maps, pat, ctx.shard_id, ctx.num_shards, true); - } - if removed.is_empty() { - conn.subscription_count = ctx.pubsub_registry.read().total_subscription_count(conn.subscriber_id); - let resp = pubsub::punsubscribe_response(&Bytes::from_static(b""), conn.subscription_count); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&resp, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if wr.is_err() { return (MonoioHandlerResult::Done, None); } - } else { - for pat in &removed { - conn.subscription_count = conn.subscription_count.saturating_sub(1); - let resp = pubsub::punsubscribe_response(pat, conn.subscription_count); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&resp, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if wr.is_err() { return (MonoioHandlerResult::Done, None); } - } - } - } else { - for arg in cmd_args { - if let Some(pattern) = extract_bytes(arg) { - ctx.pubsub_registry.write().punsubscribe(pattern.as_ref(), conn.subscriber_id); - unpropagate_subscription(&ctx.all_remote_sub_maps, &pattern, ctx.shard_id, ctx.num_shards, true); - conn.subscription_count = conn.subscription_count.saturating_sub(1); - let resp = pubsub::punsubscribe_response(&pattern, conn.subscription_count); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&resp, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if wr.is_err() { return (MonoioHandlerResult::Done, None); } - } - } - } - } - _ if cmd.eq_ignore_ascii_case(b"PING") => { - let resp = Frame::Array(framevec![ - Frame::BulkString(Bytes::from_static(b"pong")), - Frame::BulkString(Bytes::from_static(b"")), - ]); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&resp, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if wr.is_err() { return (MonoioHandlerResult::Done, None); } - } - _ if cmd.eq_ignore_ascii_case(b"QUIT") => { - let resp = Frame::SimpleString(Bytes::from_static(b"OK")); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&resp, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - let _ = wr; // ignore write error on quit - return (MonoioHandlerResult::Done, None); // exit connection - } - _ => { - let cmd_str = String::from_utf8_lossy(cmd); - let err = Frame::Error(Bytes::from(format!( - "ERR Can't execute '{}': only (P)SUBSCRIBE / (P)UNSUBSCRIBE / PING / QUIT are allowed in this context", - cmd_str.to_lowercase() - ))); - let mut resp_buf = BytesMut::new(); - codec.encode_frame(&err, &mut resp_buf); - let data = resp_buf.freeze(); - let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if wr.is_err() { return (MonoioHandlerResult::Done, None); } - } - } - } - } - Ok(None) => break, // need more data - Err(_) => return (MonoioHandlerResult::Done, None), // parse error - } - } - } - Err(_) => break, // connection error - } - } - msg = rx.recv_async() => { - match msg { - Ok(data) => { - // Data is pre-serialized RESP bytes — write directly - let (result, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; - if result.is_err() { break; } - } - Err(_) => break, // all senders dropped - } - } - _ = shutdown.cancelled() => { break; } - } - continue; - } - - // Read data from stream using monoio ownership I/O. - // Reuse pre-allocated buffer; restore length to 8192 for the read. - if tmp_buf.len() < 8192 { - tmp_buf.resize(8192, 0); - } - if let Some(dur) = idle_timeout { - // Timeout-aware read: select between read and sleep. - // monoio::select! drops the losing future, so tmp_buf ownership transfers. - // We allocate a fresh buffer when timeout is enabled (safety feature, not hot path). - let timeout_buf = std::mem::take(&mut tmp_buf); - monoio::select! { - read_result = stream.read(timeout_buf) => { - let (result, returned_buf) = read_result; - tmp_buf = returned_buf; - match result { - Ok(0) => break, - Ok(n) => { read_buf.extend_from_slice(&tmp_buf[..n]); } - Err(_) => break, - } - } - _ = monoio::time::sleep(dur) => { - tracing::debug!("Connection {} idle timeout ({}s)", client_id, idle_timeout_secs); - break; - } - } - } else { - let (result, returned_buf) = stream.read(tmp_buf).await; - tmp_buf = returned_buf; - match result { - Ok(0) => break, - Ok(n) => { - read_buf.extend_from_slice(&tmp_buf[..n]); - } - Err(_) => break, - } - } - - // Inline dispatch: handle GET/SET directly from raw bytes without Frame - // construction or dispatch table lookup. For multi-shard, only local keys - // are inlined; remote keys fall through to normal cross-shard dispatch. - // Skip inline dispatch when not conn.authenticated — AUTH must go through normal path. - // Skip inline dispatch for workspace-bound connections: workspace key prefix - // injection happens in the normal dispatch path (line 2356+) and the inline - // path has no equivalent, so GET/SET would use the raw (un-prefixed) key, - // causing GET to return nil even after SET succeeds (LVAL-01 bug). - if conn.authenticated && conn.workspace_id.is_none() { - // Inline writes are only safe when every side-effect handled by - // the normal dispatch path is either covered by the inline path - // or provably unnecessary: - // - `cached_acl_unrestricted`: ACL check can be skipped - // - `!in_multi`: writes must be queued into the transaction - // - `!tracking_enabled`: CLIENT TRACKING invalidation required - // - `!is_replica`: replica rejects writes with READONLY - // - `spill_sender.is_none()`: tiered storage needs async spill - // eviction, not the synchronous delete path. - // - // The replica check does a non-blocking `try_read` on the shared - // `RwLock`. If the lock is momentarily held for - // write (role change in progress), fail safe by disabling inline - // writes for this batch — the normal dispatch path will do the - // full check next iteration. - let is_replica = ctx.repl_state.as_ref().is_some_and(|rs| { - rs.try_read().is_ok_and(|g| { - matches!( - g.role, - crate::replication::state::ReplicationRole::Replica { .. } - ) - }) - }); - let can_inline_writes = conn.acl_skip_allowed() - && !conn.in_multi - && !conn.tracking_state.enabled - && !is_replica - && ctx.spill_sender.is_none(); - let inlined = try_inline_dispatch_loop( - &mut read_buf, - &mut write_buf, - &ctx.shard_databases, - ctx.shard_id, - conn.selected_db, - &ctx.aof_tx, - ctx.cached_clock.ms(), - ctx.num_shards, - can_inline_writes, - &ctx.runtime_config, - ); - if inlined > 0 && read_buf.is_empty() { - // All commands were inlined -- flush write_buf and continue - if !write_buf.is_empty() { - let data = write_buf.split().freeze(); - let (result, _): (std::io::Result, bytes::Bytes) = - stream.write_all(data).await; - if result.is_err() { - break; - } - } - continue; - } - // If read_buf still has data, fall through to normal Frame parsing - // for remaining commands. Inlined responses are already in write_buf. - } - - // Parse all complete frames from the read buffer (reuse pre-allocated Vec, cap at 1024) - frames.clear(); - loop { - match codec.decode_frame(&mut read_buf) { - Ok(Some(frame)) => { - frames.push(frame); - if frames.len() >= 1024 { - break; - } - } - Ok(None) => break, - Err(_) => return (MonoioHandlerResult::Done, None), // parse error, close connection - } - } - - if frames.is_empty() { - continue; - } - - // CLIENT PAUSE: delay processing if server is paused - crate::client_pause::expire_if_needed(); - if let Some(remaining) = crate::client_pause::check_pause(true) { - monoio::time::sleep(remaining).await; - } - - // Process frames with shard routing, cross-shard dispatch, and AOF logging - // Note: do NOT clear write_buf -- it may contain responses from inline dispatch. - // The inline path appends directly; the normal path appends via encode_frame below. - // write_buf is cleared via .split().freeze() at the flush point each iteration. - let mut should_quit = false; - - // Pipeline batch optimization: reuse pre-allocated containers (clear, not re-create). - responses.clear(); - remote_groups.clear(); - // Accumulate cross-shard PUBLISH pairs per target shard for batch dispatch - let mut publish_batches: std::collections::HashMap> = - std::collections::HashMap::new(); - - // Refresh time once per batch — sub-millisecond accuracy not needed per-command. - { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - guard.refresh_now_from_cache(&ctx.cached_clock); - } - - let mut auth_delay_ms: u64 = 0; - - for frame in frames.drain(..) { - // --- AUTH gate --- - if !conn.authenticated { - match extract_command(&frame) { - Some((cmd, cmd_args)) if cmd.eq_ignore_ascii_case(b"AUTH") => { - let (response, opt_user) = conn_cmd::auth_acl(cmd_args, &ctx.acl_table); - if let Some(uname) = opt_user { - conn.authenticated = true; - conn.current_user = uname; - conn.refresh_acl_cache(&ctx.acl_table); - if let Ok(addr) = peer_addr.parse::() { - crate::auth_ratelimit::record_success(addr.ip()); - } - } else { - if let Ok(addr) = peer_addr.parse::() { - auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); - } - conn.acl_log.push(crate::acl::AclLogEntry { - reason: "auth".to_string(), - object: "AUTH".to_string(), - username: conn.current_user.clone(), - client_addr: peer_addr.clone(), - timestamp_ms: std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_millis() - as u64, - }); - } - responses.push(response); - continue; - } - Some((cmd, cmd_args)) if cmd.eq_ignore_ascii_case(b"HELLO") => { - let (response, new_proto, new_name, opt_user) = conn_cmd::hello_acl( - cmd_args, - conn.protocol_version, - client_id, - &ctx.acl_table, - &mut conn.authenticated, - ); - if !matches!(&response, Frame::Error(_)) { - conn.protocol_version = new_proto; - } - if let Some(name) = new_name { - conn.client_name = Some(name); - } - if let Some(ref uname) = opt_user { - conn.current_user = uname.clone(); - conn.refresh_acl_cache(&ctx.acl_table); - } - // HELLO AUTH rate limiting - if matches!(&response, Frame::Error(_)) { - if let Ok(addr) = peer_addr.parse::() { - auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); - } - } else if opt_user.is_some() { - if let Ok(addr) = peer_addr.parse::() { - crate::auth_ratelimit::record_success(addr.ip()); - } - } - responses.push(response); - continue; - } - Some((cmd, _)) if cmd.eq_ignore_ascii_case(b"QUIT") => { - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - should_quit = true; - break; - } - _ => { - responses.push(Frame::Error(Bytes::from_static( - b"NOAUTH Authentication required.", - ))); - continue; - } - } - } - - let (cmd, cmd_args) = match extract_command(&frame) { - Some(pair) => pair, - None => { - responses.push(Frame::Error(Bytes::from_static( - b"ERR invalid command format", - ))); - continue; - } - }; - - // --- QUIT --- - if cmd.eq_ignore_ascii_case(b"QUIT") { - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - should_quit = true; - break; - } - - // --- ASKING: set per-connection flag for next command --- - if cmd.eq_ignore_ascii_case(b"ASKING") { - conn.asking = true; - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - continue; - } - - // --- CLUSTER subcommands --- - if cmd.eq_ignore_ascii_case(b"CLUSTER") { - if let Some(ref cs) = ctx.cluster_state { - #[allow(clippy::unwrap_used)] // Fallback "127.0.0.1:6379" is a valid literal - let self_addr: std::net::SocketAddr = format!("127.0.0.1:{}", ctx.config_port) - .parse() - .unwrap_or_else(|_| "127.0.0.1:6379".parse().unwrap()); - let resp = - crate::cluster::command::handle_cluster_command(cmd_args, cs, self_addr); - responses.push(resp); - } else { - responses.push(Frame::Error(Bytes::from_static( - b"ERR This instance has cluster support disabled", - ))); - } - continue; - } - - // --- Lua scripting: EVALSHA --- - if cmd.eq_ignore_ascii_case(b"EVALSHA") { - let response = { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let db_count = ctx.shard_databases.db_count(); - let db = &mut guard; - crate::scripting::handle_evalsha( - &ctx.lua, - &ctx.script_cache, - cmd_args, - db, - ctx.shard_id, - ctx.num_shards, - conn.selected_db, - db_count, - ) - }; - responses.push(response); - continue; - } - - // --- Lua scripting: EVAL --- - if cmd.eq_ignore_ascii_case(b"EVAL") { - let response = { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let db_count = ctx.shard_databases.db_count(); - let db = &mut guard; - crate::scripting::handle_eval( - &ctx.lua, - &ctx.script_cache, - cmd_args, - db, - ctx.shard_id, - ctx.num_shards, - conn.selected_db, - db_count, - ) - }; - responses.push(response); - continue; - } - - // --- SCRIPT subcommands: LOAD, EXISTS, FLUSH --- - if cmd.eq_ignore_ascii_case(b"SCRIPT") { - let (response, fanout) = - crate::scripting::handle_script_subcommand(&ctx.script_cache, cmd_args); - if let Some((sha1, script_bytes)) = fanout { - let mut producers = ctx.dispatch_tx.borrow_mut(); - for target in 0..ctx.num_shards { - if target == ctx.shard_id { - continue; - } - let idx = ChannelMesh::target_index(ctx.shard_id, target); - let msg = ShardMessage::ScriptLoad { - sha1: sha1.clone(), - script: script_bytes.clone(), - }; - if producers[idx].try_push(msg).is_ok() { - ctx.spsc_notifiers[target].notify_one(); - } - } - drop(producers); - } - responses.push(response); - continue; - } - - // --- Cluster slot routing (pre-dispatch) --- - if crate::cluster::cluster_enabled() { - if let Some(ref cs) = ctx.cluster_state { - let was_asking = conn.asking; - conn.asking = false; - - let maybe_key = extract_primary_key(cmd, cmd_args); - if let Some(key) = maybe_key { - let slot = crate::cluster::slots::slot_for_key(key); - #[allow(clippy::unwrap_used)] - // std RwLock: poison = prior panic = unrecoverable - let route = cs.read().unwrap().route_slot(slot, was_asking); - match route { - crate::cluster::SlotRoute::Local => {} // proceed - other => { - let err_frame = other.into_error_frame(slot); - responses.push(err_frame); - continue; - } - } - - // CROSSSLOT check for multi-key commands - if is_multi_key_command(cmd, cmd_args) { - let first_slot = slot; - let mut cross_slot = false; - for arg in cmd_args.iter().skip(1) { - if let Some(k) = match arg { - Frame::BulkString(b) => Some(b.as_ref()), - _ => None, - } { - if crate::cluster::slots::slot_for_key(k) != first_slot { - cross_slot = true; - break; - } - } - } - if cross_slot { - responses.push(Frame::Error(Bytes::from_static( - b"CROSSSLOT Keys in request don't hash to the same slot", - ))); - continue; - } - } - } - } - } - - // --- AUTH (already conn.authenticated) --- - if cmd.eq_ignore_ascii_case(b"AUTH") { - let (response, opt_user) = conn_cmd::auth_acl(cmd_args, &ctx.acl_table); - if let Some(uname) = opt_user { - conn.current_user = uname; - conn.refresh_acl_cache(&ctx.acl_table); - if let Ok(addr) = peer_addr.parse::() { - crate::auth_ratelimit::record_success(addr.ip()); - } - } else if let Ok(addr) = peer_addr.parse::() { - auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); - } - responses.push(response); - continue; - } - - // --- HELLO (protocol negotiation, ACL-aware) --- - if cmd.eq_ignore_ascii_case(b"HELLO") { - let (response, new_proto, new_name, opt_user) = conn_cmd::hello_acl( - cmd_args, - conn.protocol_version, - client_id, - &ctx.acl_table, - &mut conn.authenticated, - ); - if !matches!(&response, Frame::Error(_)) { - conn.protocol_version = new_proto; - } - if let Some(name) = new_name { - conn.client_name = Some(name); - } - if let Some(ref uname) = opt_user { - conn.current_user = uname.clone(); - conn.refresh_acl_cache(&ctx.acl_table); - } - if matches!(&response, Frame::Error(_)) { - if let Ok(addr) = peer_addr.parse::() { - auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); - } - } else if opt_user.is_some() { - if let Ok(addr) = peer_addr.parse::() { - crate::auth_ratelimit::record_success(addr.ip()); - } - } - responses.push(response); - continue; - } - - // --- ACL command (intercepted at connection level) --- - if cmd.eq_ignore_ascii_case(b"ACL") { - let response = crate::command::acl::handle_acl( - cmd_args, - &ctx.acl_table, - &mut conn.acl_log, - &conn.current_user, - &peer_addr, - &ctx.runtime_config, - ); - responses.push(response); - continue; - } - - // --- CONFIG GET/SET --- - if cmd.eq_ignore_ascii_case(b"CONFIG") { - responses.push(handle_config(cmd_args, &ctx.runtime_config, &ctx.config)); - continue; - } - - // --- REPLICAOF / SLAVEOF --- - if cmd.eq_ignore_ascii_case(b"REPLICAOF") || cmd.eq_ignore_ascii_case(b"SLAVEOF") { - use crate::command::connection::{ReplicaofAction, replicaof}; - let (resp, action) = replicaof(cmd_args); - if let Some(action) = action { - if let Some(ref rs) = ctx.repl_state { - match action { - ReplicaofAction::StartReplication { host, port } => { - if let Ok(mut rs_guard) = rs.write() { - rs_guard.role = crate::replication::state::ReplicationRole::Replica { - host: host.clone(), - port, - state: crate::replication::handshake::ReplicaHandshakeState::PingPending, - }; - } - let rs_clone = Arc::clone(rs); - let cfg = crate::replication::replica::ReplicaTaskConfig { - master_host: host, - master_port: port, - repl_state: rs_clone, - num_shards: ctx.num_shards, - persistence_dir: None, - listening_port: 0, - }; - monoio::spawn(crate::replication::replica::run_replica_task(cfg)); - } - ReplicaofAction::PromoteToMaster => { - use crate::replication::state::generate_repl_id; - if let Ok(mut rs_guard) = rs.write() { - rs_guard.repl_id2 = rs_guard.repl_id.clone(); - rs_guard.repl_id = generate_repl_id(); - rs_guard.role = - crate::replication::state::ReplicationRole::Master; - } - } - ReplicaofAction::NoOp => {} - } - } - } - responses.push(resp); - continue; - } - - // --- REPLCONF --- - if cmd.eq_ignore_ascii_case(b"REPLCONF") { - let resp = crate::command::connection::replconf(cmd_args); - responses.push(resp); - continue; - } - - // --- INFO (with replication section) --- - if cmd.eq_ignore_ascii_case(b"INFO") { - let guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); - let response_text = { - let resp_frame = conn_cmd::info_readonly(&guard, cmd_args); - match resp_frame { - Frame::BulkString(b) => String::from_utf8_lossy(&b).to_string(), - _ => String::new(), - } - }; - drop(guard); - let mut response_text = response_text; - if let Some(ref rs) = ctx.repl_state { - if let Ok(rs_guard) = rs.try_read() { - response_text.push_str( - &crate::replication::handshake::build_info_replication(&rs_guard), - ); - } - } - responses.push(Frame::BulkString(Bytes::from(response_text))); - continue; - } - - // --- READONLY enforcement: reject writes on replicas --- - if let Some(ref rs) = ctx.repl_state { - if let Ok(rs_guard) = rs.try_read() { - if matches!( - rs_guard.role, - crate::replication::state::ReplicationRole::Replica { .. } - ) { - if metadata::is_write(cmd) { - responses.push(Frame::Error(Bytes::from_static( - b"READONLY You can't write against a read only replica.", - ))); - continue; - } - } - } - } - - // --- CLIENT subcommands (ID, SETNAME, GETNAME, TRACKING) --- - if cmd.eq_ignore_ascii_case(b"CLIENT") { - if let Some(sub) = cmd_args.first() { - if let Some(sub_bytes) = extract_bytes(sub) { - if sub_bytes.eq_ignore_ascii_case(b"ID") { - responses.push(conn_cmd::client_id(client_id)); - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"SETNAME") { - if cmd_args.len() != 2 { - responses.push(Frame::Error(Bytes::from_static( - b"ERR wrong number of arguments for 'CLIENT SETNAME' command", - ))); - } else { - conn.client_name = extract_bytes(&cmd_args[1]); - let name_str = conn - .client_name - .as_ref() - .map(|b| String::from_utf8_lossy(b).to_string()); - crate::client_registry::update(client_id, |e| { - e.name = name_str; - }); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"GETNAME") { - responses.push(match &conn.client_name { - Some(name) => Frame::BulkString(name.clone()), - None => Frame::Null, - }); - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"TRACKING") { - match crate::command::client::parse_tracking_args(cmd_args) { - Ok(config_parsed) => { - if config_parsed.enable { - conn.tracking_state.enabled = true; - conn.tracking_state.bcast = config_parsed.bcast; - conn.tracking_state.noloop = config_parsed.noloop; - conn.tracking_state.optin = config_parsed.optin; - conn.tracking_state.optout = config_parsed.optout; - - if conn.tracking_rx.is_none() { - let (tx, rx) = channel::mpsc_bounded::(256); - conn.tracking_state.invalidation_tx = Some(tx.clone()); - conn.tracking_rx = Some(rx); - - let mut table = ctx.tracking_table.borrow_mut(); - table.register_client(client_id, tx); - if let Some(target) = config_parsed.redirect { - table.set_redirect(client_id, target); - } - for prefix in &config_parsed.prefixes { - table.register_prefix( - client_id, - prefix.clone(), - config_parsed.noloop, - ); - } - } - responses - .push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } else { - conn.tracking_state = TrackingState::default(); - ctx.tracking_table.borrow_mut().untrack_all(client_id); - conn.tracking_rx = None; - responses - .push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - continue; - } - Err(err_frame) => { - responses.push(err_frame); - continue; - } - } - } - // Admin CLIENT subcommands (LIST, INFO, KILL, PAUSE, UNPAUSE, - // NO-EVICT, NO-TOUCH) fall through to the ACL gate below. - } - } - // Fall through — admin subcommands handled after ACL check. - } - - // --- PUBLISH: local delivery + cross-shard fan-out --- - if cmd.eq_ignore_ascii_case(b"PUBLISH") { - if cmd_args.len() != 2 { - responses.push(Frame::Error(Bytes::from_static( - b"ERR wrong number of arguments for 'publish' command", - ))); - } else { - let channel = extract_bytes(&cmd_args[0]); - let message = extract_bytes(&cmd_args[1]); - // ACL channel permission check for PUBLISH - if let Some(ref ch) = channel { - let denied = { - #[allow(clippy::unwrap_used)] - // std RwLock: poison = prior panic = unrecoverable - let acl_guard = ctx.acl_table.read().unwrap(); - acl_guard.check_channel_permission(&conn.current_user, ch.as_ref()) - }; - if let Some(deny_reason) = denied { - responses - .push(Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason)))); - continue; - } - } - match (channel, message) { - (Some(ch), Some(msg)) => { - let local_count = { ctx.pubsub_registry.write().publish(&ch, &msg) }; - // Targeted fanout: only send to shards that have subscribers - let targets = ctx.remote_subscriber_map.read().target_shards(&ch); - if targets.is_empty() { - // Fast path: no remote subscribers - responses.push(Frame::Integer(local_count)); - } else { - let remote_targets: Vec = - targets.into_iter().filter(|&t| t != ctx.shard_id).collect(); - if remote_targets.is_empty() { - responses.push(Frame::Integer(local_count)); - } else { - // Accumulate into per-shard batches for coalesced dispatch - let resp_idx = responses.len(); - responses.push(Frame::Integer(local_count)); // placeholder, updated after batch flush - for target in &remote_targets { - publish_batches.entry(*target).or_default().push(( - resp_idx, - ch.clone(), - msg.clone(), - )); - } - } - } - } - _ => responses.push(Frame::Error(Bytes::from_static( - b"ERR invalid channel or message", - ))), - } - } - continue; - } - - // --- SUBSCRIBE / PSUBSCRIBE --- - if cmd.eq_ignore_ascii_case(b"SUBSCRIBE") || cmd.eq_ignore_ascii_case(b"PSUBSCRIBE") { - let is_pattern = cmd.eq_ignore_ascii_case(b"PSUBSCRIBE"); - if cmd_args.is_empty() { - let cmd_name = if is_pattern { - "psubscribe" - } else { - "subscribe" - }; - let err = Frame::Error(Bytes::from(format!( - "ERR wrong number of arguments for '{}' command", - cmd_name - ))); - responses.push(err); - continue; - } - // Allocate pubsub channel if not yet created - if conn.pubsub_tx.is_none() { - let (tx, rx) = channel::mpsc_bounded::(256); - conn.pubsub_tx = Some(tx); - conn.pubsub_rx = Some(rx); - } - if conn.subscriber_id == 0 { - conn.subscriber_id = crate::pubsub::next_subscriber_id(); - } - // Flush accumulated responses before entering subscriber mode - for resp in &responses { - codec.encode_frame(resp, &mut write_buf); - } - for arg in cmd_args { - if let Some(ch) = extract_bytes(arg) { - // ACL channel permission check - { - #[allow(clippy::unwrap_used)] - // std RwLock: poison = prior panic = unrecoverable - let acl_guard = ctx.acl_table.read().unwrap(); - if let Some(deny_reason) = - acl_guard.check_channel_permission(&conn.current_user, ch.as_ref()) - { - drop(acl_guard); - let err = - Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason))); - codec.encode_frame(&err, &mut write_buf); - continue; - } - } - #[allow(clippy::unwrap_used)] - // conn.pubsub_tx is set to Some just above before this loop - let sub = Subscriber::with_protocol( - conn.pubsub_tx.clone().unwrap(), - conn.subscriber_id, - conn.protocol_version >= 3, - ); - if is_pattern { - ctx.pubsub_registry.write().psubscribe(ch.clone(), sub); - } else { - ctx.pubsub_registry.write().subscribe(ch.clone(), sub); - } - propagate_subscription( - &ctx.all_remote_sub_maps, - &ch, - ctx.shard_id, - ctx.num_shards, - is_pattern, - ); - conn.subscription_count += 1; - // Register pub/sub affinity for this client IP - if conn.subscription_count == 1 { - if let Ok(addr) = peer_addr.parse::() { - ctx.pubsub_affinity - .write() - .register(addr.ip(), ctx.shard_id); - } - } - let resp = if is_pattern { - pubsub::psubscribe_response(&ch, conn.subscription_count) - } else { - pubsub::subscribe_response(&ch, conn.subscription_count) - }; - codec.encode_frame(&resp, &mut write_buf); - } - } - // Flush responses and re-enter loop (next iteration enters subscriber mode) - if !write_buf.is_empty() { - let data = write_buf.split().freeze(); - let (result, _): (std::io::Result, bytes::Bytes) = - stream.write_all(data).await; - if result.is_err() { - return (MonoioHandlerResult::Done, None); - } - } - responses.clear(); - break; // break out of frame loop to re-enter main loop in subscriber mode - } - - // --- UNSUBSCRIBE / PUNSUBSCRIBE (in normal mode, no-op if not subscribed) --- - if cmd.eq_ignore_ascii_case(b"UNSUBSCRIBE") || cmd.eq_ignore_ascii_case(b"PUNSUBSCRIBE") - { - let is_pattern = cmd.eq_ignore_ascii_case(b"PUNSUBSCRIBE"); - let resp = if is_pattern { - pubsub::punsubscribe_response(&Bytes::from_static(b""), 0) - } else { - pubsub::unsubscribe_response(&Bytes::from_static(b""), 0) - }; - responses.push(resp); - continue; - } - - // --- PUBSUB introspection (zero-SPSC: direct shared-read) --- - if cmd.eq_ignore_ascii_case(b"PUBSUB") { - if cmd_args.is_empty() { - responses.push(Frame::Error(Bytes::from_static( - b"ERR wrong number of arguments for 'pubsub' command", - ))); - continue; - } - let subcmd = extract_bytes(&cmd_args[0]); - match subcmd { - Some(ref sc) if sc.eq_ignore_ascii_case(b"CHANNELS") => { - let pattern = if cmd_args.len() > 1 { - extract_bytes(&cmd_args[1]) - } else { - None - }; - let mut all_channels: std::collections::HashSet = - std::collections::HashSet::new(); - for reg in &ctx.all_pubsub_registries { - let guard = reg.read(); - all_channels.extend(guard.active_channels(pattern.as_deref())); - } - let arr: Vec = - all_channels.into_iter().map(Frame::BulkString).collect(); - responses.push(Frame::Array(arr.into())); - } - Some(ref sc) if sc.eq_ignore_ascii_case(b"NUMSUB") => { - let channels: Vec = cmd_args[1..] - .iter() - .filter_map(|a| extract_bytes(a)) - .collect(); - let mut counts: std::collections::HashMap = - std::collections::HashMap::new(); - for reg in &ctx.all_pubsub_registries { - let guard = reg.read(); - for (ch, c) in guard.numsub(&channels) { - *counts.entry(ch).or_insert(0) += c; - } - } - let mut arr = Vec::with_capacity(channels.len() * 2); - for ch in &channels { - arr.push(Frame::BulkString(ch.clone())); - arr.push(Frame::Integer(*counts.get(ch).unwrap_or(&0))); - } - responses.push(Frame::Array(arr.into())); - } - Some(ref sc) if sc.eq_ignore_ascii_case(b"NUMPAT") => { - let mut total: usize = 0; - for reg in &ctx.all_pubsub_registries { - total += reg.read().numpat(); - } - responses.push(Frame::Integer(total as i64)); - } - _ => { - responses.push(Frame::Error(Bytes::from_static( - b"ERR unknown subcommand or wrong number of arguments for 'pubsub' command", - ))); - } - } - continue; - } - - // --- BGSAVE: trigger per-shard cooperative snapshot --- - if cmd.eq_ignore_ascii_case(b"BGSAVE") { - let response = crate::command::persistence::bgsave_start_sharded( - &ctx.snapshot_trigger_tx, - ctx.num_shards, - ); - responses.push(response); - continue; - } - // SAVE -- not supported in sharded mode - if cmd.eq_ignore_ascii_case(b"SAVE") { - responses.push(Frame::Error(Bytes::from_static( - b"ERR SAVE not supported in sharded mode, use BGSAVE", - ))); - continue; - } - // LASTSAVE -- return timestamp of last successful save - if cmd.eq_ignore_ascii_case(b"LASTSAVE") { - responses.push(crate::command::persistence::handle_lastsave()); - continue; - } - // BGREWRITEAOF -- multi-part AOF rewrite - if cmd.eq_ignore_ascii_case(b"BGREWRITEAOF") { - if let Some(ref tx) = ctx.aof_tx { - responses.push(crate::command::persistence::bgrewriteaof_start_sharded( - tx, - ctx.shard_databases.clone(), - )); - } else { - responses.push(Frame::Error(Bytes::from_static(b"ERR AOF is not enabled"))); - } - continue; - } - - // === ACL permission check (NOPERM gate) === - // Exempt commands (AUTH, HELLO, QUIT, ACL) already handled above. - // Fast path: skip RwLock + HashMap probe for unrestricted users - // whose per-connection cache is still fresh (no ACL mutation has - // occurred since the cache was populated). A stale cache MUST - // NOT bypass this check — see `ConnectionState::acl_skip_allowed`. - if !conn.acl_skip_allowed() { - #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable - let acl_guard = ctx.acl_table.read().unwrap(); - if let Some(deny_reason) = - acl_guard.check_command_permission(&conn.current_user, cmd, cmd_args) - { - drop(acl_guard); - conn.acl_log.push(crate::acl::AclLogEntry { - reason: "command".to_string(), - object: String::from_utf8_lossy(cmd).to_ascii_lowercase(), - username: conn.current_user.clone(), - client_addr: peer_addr.clone(), - timestamp_ms: std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as u64, - }); - responses.push(Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason)))); - continue; - } - - // === ACL key pattern check (same lock guard) === - let is_write_for_acl = metadata::is_write(cmd); - if let Some(deny_reason) = acl_guard.check_key_permission( - &conn.current_user, - cmd, - cmd_args, - is_write_for_acl, - ) { - drop(acl_guard); - conn.acl_log.push(crate::acl::AclLogEntry { - reason: "command".to_string(), - object: String::from_utf8_lossy(cmd).to_ascii_lowercase(), - username: conn.current_user.clone(), - client_addr: peer_addr.clone(), - timestamp_ms: std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as u64, - }); - responses.push(Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason)))); - continue; - } - } // !cached_acl_unrestricted - - // --- CLIENT admin subcommands (LIST, INFO, KILL, PAUSE, UNPAUSE) --- - // Placed AFTER ACL check so restricted users cannot access admin ops. - if cmd.eq_ignore_ascii_case(b"CLIENT") { - if let Some(sub) = cmd_args.first() { - if let Some(sub_bytes) = extract_bytes(sub) { - if sub_bytes.eq_ignore_ascii_case(b"LIST") { - crate::client_registry::update(client_id, |e| { - e.db = conn.selected_db; - e.last_cmd_at = std::time::Instant::now(); - e.flags = crate::client_registry::ClientFlags { - subscriber: conn.subscription_count > 0, - in_multi: conn.in_multi, - blocked: false, - }; - }); - let list = crate::client_registry::client_list(); - responses.push(Frame::BulkString(Bytes::from(list))); - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"INFO") { - crate::client_registry::update(client_id, |e| { - e.db = conn.selected_db; - e.last_cmd_at = std::time::Instant::now(); - }); - let info = - crate::client_registry::client_info(client_id).unwrap_or_default(); - responses.push(Frame::BulkString(Bytes::from(info))); - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"KILL") { - let raw_args: Vec<&[u8]> = cmd_args[1..] - .iter() - .filter_map(|f| match f { - Frame::BulkString(b) => Some(b.as_ref()), - Frame::SimpleString(b) => Some(b.as_ref()), - _ => None, - }) - .collect(); - match crate::client_registry::parse_kill_args(&raw_args) { - Some(filter) => { - let count = crate::client_registry::kill_clients(&filter); - responses.push(Frame::Integer(count as i64)); - } - None => { - responses.push(Frame::Error(Bytes::from_static( - b"ERR syntax error. Usage: CLIENT KILL [ID id] [ADDR addr] [USER user]", - ))); - } - } - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"PAUSE") { - if cmd_args.len() < 2 { - responses.push(Frame::Error(Bytes::from_static( - b"ERR wrong number of arguments for 'CLIENT PAUSE' command", - ))); - } else { - let timeout_bytes = match &cmd_args[1] { - Frame::BulkString(b) => Some(b.as_ref()), - Frame::SimpleString(b) => Some(b.as_ref()), - _ => None, - }; - match timeout_bytes - .and_then(|b| std::str::from_utf8(b).ok()) - .and_then(|s| s.parse::().ok()) - { - Some(ms) => { - let mode = if cmd_args.len() > 2 { - match &cmd_args[2] { - Frame::BulkString(b) | Frame::SimpleString(b) - if b.eq_ignore_ascii_case(b"WRITE") => - { - crate::client_pause::PauseMode::Write - } - _ => crate::client_pause::PauseMode::All, - } - } else { - crate::client_pause::PauseMode::All - }; - crate::client_pause::pause(ms, mode); - responses - .push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - None => { - responses.push(Frame::Error(Bytes::from_static( - b"ERR timeout is not a valid integer or out of range", - ))); - } - } - } - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"UNPAUSE") { - crate::client_pause::unpause(); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"NO-EVICT") - || sub_bytes.eq_ignore_ascii_case(b"NO-TOUCH") - { - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - continue; - } - // Unknown CLIENT subcommand - responses.push(Frame::Error(Bytes::from(format!( - "ERR unknown subcommand '{}'", - String::from_utf8_lossy(&sub_bytes) - )))); - continue; - } - } - responses.push(Frame::Error(Bytes::from_static( - b"ERR wrong number of arguments for 'client' command", - ))); - continue; - } - - // --- Functions API: FUNCTION/FCALL/FCALL_RO --- - // Placed AFTER ACL check. Respects MULTI queue — if conn.in_multi, - // fall through to the MULTI queue gate instead of executing. - if !conn.in_multi { - if cmd.eq_ignore_ascii_case(b"FUNCTION") { - let response = crate::command::functions::handle_function( - &mut func_registry.borrow_mut(), - cmd_args, - ); - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"FCALL") { - let response = { - let mut guard = - ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let db_count = ctx.shard_databases.db_count(); - crate::command::functions::handle_fcall( - &func_registry.borrow(), - cmd_args, - &mut guard, - ctx.shard_id, - ctx.num_shards, - conn.selected_db, - db_count, - ) - }; - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"FCALL_RO") { - let response = { - let mut guard = - ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let db_count = ctx.shard_databases.db_count(); - crate::command::functions::handle_fcall_ro( - &func_registry.borrow(), - cmd_args, - &mut guard, - ctx.shard_id, - ctx.num_shards, - conn.selected_db, - db_count, - ) - }; - responses.push(response); - continue; - } - } - - // --- TXN.BEGIN --- - if is_txn_begin(cmd, cmd_args) { - match txn_begin_validate(conn.in_multi, conn.in_cross_txn()) { - Ok(()) => { - // Get next txn_id and snapshot_lsn from vector store's transaction manager - let mut vector_store = ctx.shard_databases.vector_store(ctx.shard_id); - let active = vector_store.txn_manager_mut().begin(); - conn.active_cross_txn = - Some(CrossStoreTxn::new(active.txn_id, active.snapshot_lsn)); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - Err(e) => responses.push(e), - } - continue; - } - - // --- TXN.COMMIT --- - if is_txn_commit(cmd, cmd_args) { - match txn_commit_validate(conn.in_cross_txn()) { - Ok(()) => { - if let Some(txn) = conn.active_cross_txn.take() { - let mut vector_store = ctx.shard_databases.vector_store(ctx.shard_id); - vector_store.txn_manager_mut().commit(txn.txn_id); - drop(vector_store); - - // Write XactCommit WAL record with committed KV state - let txn_id = txn.txn_id; - if !txn.kv_undo.is_empty() { - let db_guard = - ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); - let payload = - crate::persistence::wal_v3::record::encode_xact_commit_payload( - txn_id, - txn.kv_undo.records(), - &*db_guard, - ); - drop(db_guard); - let mut wal_buf = Vec::new(); - crate::persistence::wal_v3::record::write_wal_v3_record( - &mut wal_buf, - txn_id, - crate::persistence::wal_v3::record::WalRecordType::XactCommit, - &payload, - ); - ctx.shard_databases - .wal_append(ctx.shard_id, bytes::Bytes::from(wal_buf)); - } - - // Release KV write intents from shard side-table - ctx.shard_databases - .kv_intents(ctx.shard_id) - .release_txn(txn_id); - - // Drain deferred HNSW inserts (post-commit hook). - // The drain prevents phantom neighbors on abort. - // Actual HNSW graph insertion happens during compaction, - // not at commit time (point is already in mutable segment). - let drain_count = ctx - .shard_databases - .hnsw_queue(ctx.shard_id) - .drain_for_txn(txn_id) - .count(); - if drain_count > 0 { - tracing::debug!( - txn_id, - count = drain_count, - "Drained deferred HNSW inserts" - ); - } - - // Materialize MQ intents: enqueue deferred MQ.PUBLISH messages - if !txn.mq_intents.is_empty() { - let mut db_guard = - ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - for intent in &txn.mq_intents { - if let Ok(Some(stream)) = - db_guard.get_stream_mut(&intent.queue_key) - { - if stream.durable { - let msg_id = stream.next_auto_id(); - stream.add(msg_id, intent.fields.clone()); - } - } - } - } - - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } else { - responses - .push(Frame::Error(Bytes::from_static(b"ERR not in transaction"))); - } - } - Err(e) => responses.push(e), - } - continue; - } - - // --- TXN.ABORT --- - if is_txn_abort(cmd, cmd_args) { - match txn_abort_validate(conn.in_cross_txn()) { - Ok(()) => { - if let Some(txn) = conn.active_cross_txn.take() { - // Shared rollback (Phase 166 Plan 03): - // KV undo -> graph intents reverse -> vector - // tombstone -> side-table release. See - // src/transaction/abort.rs for lock ordering. - crate::transaction::abort::abort_cross_store_txn( - &ctx.shard_databases, - ctx.shard_id, - conn.selected_db, - txn, - ); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } else { - responses - .push(Frame::Error(Bytes::from_static(b"ERR not in transaction"))); - } - } - Err(e) => responses.push(e), - } - continue; - } - - // --- TEMPORAL.SNAPSHOT_AT --- - if is_temporal_snapshot_at(cmd) { - match validate_snapshot_at(cmd_args) { - Ok(()) => { - let wall_ms = capture_wall_ms(); - // Get current LSN from the vector store's transaction manager - let lsn = { - let vector_store = ctx.shard_databases.vector_store(ctx.shard_id); - vector_store.txn_manager().current_lsn() - }; - // Lazy-init and record the wall-clock -> LSN binding - { - let mut guard = ctx.shard_databases.temporal_registry(ctx.shard_id); - let registry = guard.get_or_insert_with(|| { - Box::new(crate::temporal::TemporalRegistry::new()) - }); - registry.record(wall_ms, lsn); - } - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - Err(e) => responses.push(e), - } - continue; - } - - // --- TEMPORAL.INVALIDATE --- - if is_temporal_invalidate(cmd) { - match validate_invalidate(cmd_args) { - Ok((entity_id, is_node, graph_name)) => { - let wall_ms = capture_wall_ms(); - #[cfg(feature = "graph")] - { - let mut gs = ctx.shard_databases.graph_store_write(ctx.shard_id); - if let Some(named_graph) = gs.get_graph_mut(&graph_name) { - let mutated = if is_node { - let node_key: crate::graph::types::NodeKey = - slotmap::KeyData::from_ffi(entity_id).into(); - if let Some(node) = named_graph.write_buf.get_node_mut(node_key) - { - node.valid_to = wall_ms; - true - } else { - false - } - } else { - let edge_key: crate::graph::types::EdgeKey = - slotmap::KeyData::from_ffi(entity_id).into(); - if let Some(edge) = named_graph.write_buf.get_edge_mut(edge_key) - { - edge.valid_to = wall_ms; - true - } else { - false - } - }; - if mutated { - // Write GraphTemporal WAL record with literal system_from - let payload = - crate::persistence::wal_v3::record::encode_graph_temporal( - entity_id, is_node, wall_ms, wall_ms, - ); - gs.wal_pending.push(payload); - let wal_records = gs.drain_wal(); - drop(gs); - for record in wal_records { - ctx.shard_databases - .wal_append(ctx.shard_id, Bytes::from(record)); - } - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } else { - drop(gs); - responses.push(Frame::Error(Bytes::from_static( - ERR_ENTITY_NOT_FOUND, - ))); - } - } else { - drop(gs); - responses - .push(Frame::Error(Bytes::from_static(ERR_GRAPH_NOT_FOUND))); - } - } - #[cfg(not(feature = "graph"))] - { - let _ = (entity_id, is_node, graph_name, wall_ms); - responses.push(Frame::Error(Bytes::from_static( - b"ERR graph feature not enabled", - ))); - } - } - Err(e) => responses.push(e), - } - continue; - } - - // --- WS.* --- - if is_ws_command(cmd) { - let sub = match parse_ws_subcommand(cmd_args) { - Ok(s) => s, - Err(e) => { - responses.push(e); - continue; - } - }; - - if sub.eq_ignore_ascii_case(b"CREATE") { - match validate_ws_create(cmd_args) { - Ok(ws_name) => { - let ws_id = WorkspaceId::new_v7(); - let created_at = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as i64; - let meta = crate::workspace::WorkspaceMetadata { - id: ws_id, - name: ws_name.clone(), - created_at, - }; - { - let mut guard = - ctx.shard_databases.workspace_registry(ctx.shard_id); - let reg = guard.get_or_insert_with(|| { - Box::new(crate::workspace::WorkspaceRegistry::new()) - }); - reg.insert(ws_id, meta); - } - // WAL: WorkspaceCreate record - let payload = crate::workspace::wal::encode_workspace_create( - ws_id.as_bytes(), - &ws_name, - ); - let mut wal_buf = Vec::new(); - crate::persistence::wal_v3::record::write_wal_v3_record( - &mut wal_buf, - 0, - crate::persistence::wal_v3::record::WalRecordType::WorkspaceCreate, - &payload, - ); - ctx.shard_databases - .wal_append(ctx.shard_id, Bytes::from(wal_buf)); - responses.push(Frame::BulkString(Bytes::from(ws_id.to_string()))); - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"DROP") { - match validate_ws_drop(cmd_args) { - Ok(ws_id_raw) => { - match parse_workspace_id_from_bytes(&ws_id_raw) { - Some(ws_id) => { - let removed = { - let mut guard = - ctx.shard_databases.workspace_registry(ctx.shard_id); - match guard.as_mut() { - Some(reg) => reg.remove(&ws_id).is_some(), - None => false, - } - }; - if removed { - // WAL: WorkspaceDrop record - let payload = crate::workspace::wal::encode_workspace_drop( - ws_id.as_bytes(), - ); - let mut wal_buf = Vec::new(); - crate::persistence::wal_v3::record::write_wal_v3_record( - &mut wal_buf, 0, - crate::persistence::wal_v3::record::WalRecordType::WorkspaceDrop, - &payload, - ); - ctx.shard_databases - .wal_append(ctx.shard_id, Bytes::from(wal_buf)); - // Best-effort cleanup: delete all KV keys with ws prefix (WS-03). - { - let prefix = format!("{{{}}}:", ws_id.as_hex()); - let mut db_guard = - ctx.shard_databases.write_db(ctx.shard_id, 0); - let keys_to_delete: Vec> = db_guard - .keys() - .filter(|k| { - k.as_bytes().starts_with(prefix.as_bytes()) - }) - .map(|k| k.as_bytes().to_vec()) - .collect(); - for key in &keys_to_delete { - db_guard.remove(key); - } - } - responses - .push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } else { - responses.push(Frame::Error(Bytes::from_static( - ERR_WS_NOT_FOUND, - ))); - } - } - None => responses.push(Frame::Error(Bytes::from_static( - crate::command::workspace::ERR_WS_INVALID_ID, - ))), - } - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"LIST") { - match validate_ws_list(cmd_args) { - Ok(()) => { - let guard = ctx.shard_databases.workspace_registry(ctx.shard_id); - let entries: Vec = match guard.as_ref() { - Some(reg) => reg - .iter() - .map(|(id, meta)| { - Frame::Array( - vec![ - Frame::BulkString(Bytes::from(id.to_string())), - Frame::BulkString(meta.name.clone()), - Frame::Integer(meta.created_at), - ] - .into(), - ) - }) - .collect(), - None => vec![], - }; - responses.push(Frame::Array(entries.into())); - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"INFO") { - match validate_ws_info(cmd_args) { - Ok(ws_id_raw) => match parse_workspace_id_from_bytes(&ws_id_raw) { - Some(ws_id) => { - let guard = ctx.shard_databases.workspace_registry(ctx.shard_id); - let found = guard.as_ref().and_then(|reg| reg.get(&ws_id)); - match found { - Some(meta) => { - responses.push(Frame::Array( - vec![ - Frame::BulkString(Bytes::from_static(b"id")), - Frame::BulkString(Bytes::from(meta.id.to_string())), - Frame::BulkString(Bytes::from_static(b"name")), - Frame::BulkString(meta.name.clone()), - Frame::BulkString(Bytes::from_static( - b"created_at", - )), - Frame::Integer(meta.created_at), - ] - .into(), - )); - } - None => responses - .push(Frame::Error(Bytes::from_static(ERR_WS_NOT_FOUND))), - } - } - None => responses.push(Frame::Error(Bytes::from_static( - crate::command::workspace::ERR_WS_INVALID_ID, - ))), - }, - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"AUTH") { - match validate_ws_auth(cmd_args) { - Ok(ws_id_raw) => { - if conn.workspace_id.is_some() { - responses - .push(Frame::Error(Bytes::from_static(ERR_WS_ALREADY_BOUND))); - } else { - match parse_workspace_id_from_bytes(&ws_id_raw) { - Some(ws_id) => { - let found = { - let guard = ctx - .shard_databases - .workspace_registry(ctx.shard_id); - guard - .as_ref() - .map_or(false, |reg| reg.get(&ws_id).is_some()) - }; - if found { - conn.workspace_id = Some(ws_id); - responses.push(Frame::SimpleString( - Bytes::from_static(b"OK"), - )); - } else { - responses.push(Frame::Error(Bytes::from_static( - ERR_WS_NOT_FOUND, - ))); - } - } - None => responses.push(Frame::Error(Bytes::from_static( - crate::command::workspace::ERR_WS_INVALID_ID, - ))), - } - } - } - Err(e) => responses.push(e), - } - continue; - } - - // Unknown WS subcommand - responses.push(Frame::Error(Bytes::from_static(ERR_WS_UNKNOWN_SUB))); - continue; - } - - // --- MQ.* --- - if is_mq_command(cmd) { - let sub = match parse_mq_subcommand(cmd_args) { - Ok(s) => s, - Err(e) => { - responses.push(e); - continue; - } - }; - - if sub.eq_ignore_ascii_case(b"CREATE") { - match validate_mq_create(cmd_args) { - Ok((queue_key, max_delivery_count, _debounce_ms)) => { - let effective_key = crate::workspace::workspace_key( - conn.workspace_id.as_ref(), - &queue_key, - ); - // Create or get Stream in db - let mut db_guard = - ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - match db_guard.get_or_create_stream(&effective_key) { - Ok(stream) => { - stream.durable = true; - stream.max_delivery_count = max_delivery_count; - // Auto-create __mq_consumers consumer group if not exists - let group_name = Bytes::from_static(b"__mq_consumers"); - let _ = stream.create_group(group_name, StreamId::ZERO); - } - Err(e) => { - responses.push(e); - continue; - } - } - drop(db_guard); - - // Store config in per-shard registry - let config = crate::mq::DurableStreamConfig::new( - effective_key.clone(), - max_delivery_count, - ); - { - let mut guard = - ctx.shard_databases.durable_queue_registry(ctx.shard_id); - let reg = guard.get_or_insert_with(|| { - Box::new(crate::mq::DurableQueueRegistry::new()) - }); - reg.insert(effective_key.clone(), config); - } - - // WAL: MqCreate record - let payload = crate::mq::wal::encode_mq_create( - &effective_key, - max_delivery_count, - ); - let mut wal_buf = Vec::new(); - crate::persistence::wal_v3::record::write_wal_v3_record( - &mut wal_buf, - 0, - crate::persistence::wal_v3::record::WalRecordType::MqCreate, - &payload, - ); - ctx.shard_databases - .wal_append(ctx.shard_id, Bytes::from(wal_buf)); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"PUSH") { - match validate_mq_push(cmd_args) { - Ok((queue_key, fields)) => { - let effective_key = crate::workspace::workspace_key( - conn.workspace_id.as_ref(), - &queue_key, - ); - let mut db_guard = - ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - match db_guard.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - if !stream.durable { - responses.push(Frame::Error(Bytes::from_static( - ERR_MQ_NOT_DURABLE, - ))); - } else { - let msg_id = stream.next_auto_id(); - let msg_id = stream.add(msg_id, fields); - drop(db_guard); - // Mark pending for any registered triggers - { - let mut trig_guard = - ctx.shard_databases.trigger_registry(ctx.shard_id); - if let Some(reg) = trig_guard.as_mut() { - let trig_key = if let Some(ws_id) = - conn.workspace_id.as_ref() - { - let ws_hex = ws_id.as_hex(); - let mut k = Vec::with_capacity( - ws_hex.len() + 1 + queue_key.len(), - ); - k.extend_from_slice(ws_hex.as_bytes()); - k.push(b':'); - k.extend_from_slice(&queue_key); - Bytes::from(k) - } else { - queue_key.clone() - }; - if let Some(trig_entry) = reg.get_mut(&trig_key) { - if trig_entry.pending_fire_ms == 0 { - let fire_at = ctx.cached_clock.ms() - + trig_entry.debounce_ms; - trig_entry.pending_fire_ms = fire_at; - } - } - } - } - responses.push(Frame::BulkString(Bytes::from(format!( - "{}-{}", - msg_id.ms, msg_id.seq - )))); - } - } - Ok(None) => { - responses - .push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); - } - Err(e) => responses.push(e), - } - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"POP") { - match validate_mq_pop(cmd_args) { - Ok((queue_key, count)) => { - let effective_key = crate::workspace::workspace_key( - conn.workspace_id.as_ref(), - &queue_key, - ); - let group_name = Bytes::from_static(b"__mq_consumers"); - let consumer_name = Bytes::from_static(b"__mq_default"); - let mut db_guard = - ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - - // Read max_delivery_count before mutating - let mdc = match db_guard.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - if !stream.durable { - responses.push(Frame::Error(Bytes::from_static( - ERR_MQ_NOT_DURABLE, - ))); - continue; - } - stream.max_delivery_count - } - Ok(None) => { - responses - .push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); - continue; - } - Err(e) => { - responses.push(e); - continue; - } - }; - - // Request more than count to account for DLQ routing - let request_count = count + (mdc as usize); - let stream = match db_guard.get_stream_mut(&effective_key) { - Ok(Some(s)) => s, - _ => { - responses - .push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); - continue; - } - }; - let claimed = match stream.read_group_new( - &group_name, - &consumer_name, - Some(request_count), - false, - ) { - Ok(entries) => entries, - Err(_) => { - responses.push(Frame::Array(vec![].into())); - continue; - } - }; - - // Filter: check PEL delivery_count for DLQ routing - let mut results = Vec::new(); - let mut dlq_entries: Vec<(StreamId, Vec<(Bytes, Bytes)>)> = Vec::new(); - let mut dlq_ack_ids: Vec = Vec::new(); - for (id, fields) in &claimed { - // Check delivery_count from PEL - let delivery_count = stream - .groups - .get(group_name.as_ref()) - .and_then(|g| g.pel.get(id)) - .map(|pe| pe.delivery_count) - .unwrap_or(1); - if mdc > 0 && delivery_count >= mdc as u64 { - // Route to DLQ - dlq_entries.push((*id, fields.clone())); - dlq_ack_ids.push(*id); - } else if results.len() < count { - results.push((*id, fields.clone())); - } - } - - // XACK DLQ entries from __mq_consumers group - if !dlq_ack_ids.is_empty() { - let _ = stream.xack(&group_name, &dlq_ack_ids); - } - - // Move DLQ entries to DLQ stream - if !dlq_entries.is_empty() { - let dlq_key = { - let mut buf = Vec::with_capacity(effective_key.len() + 8); - buf.extend_from_slice(&effective_key); - buf.extend_from_slice(b"::mq:dlq"); - Bytes::from(buf) - }; - match db_guard.get_or_create_stream(&dlq_key) { - Ok(dlq_stream) => { - for (_id, fields) in dlq_entries { - let dlq_id = dlq_stream.next_auto_id(); - dlq_stream.add(dlq_id, fields); - } - } - Err(_) => {} // DLQ creation failed -- skip silently - } - } - - // Format results as array of arrays (XREADGROUP format) - let result_frames: Vec = results - .iter() - .map(|(id, fields)| { - let mut entry_frames = Vec::with_capacity(2); - entry_frames.push(Frame::BulkString(Bytes::from(format!( - "{}-{}", - id.ms, id.seq - )))); - let field_frames: Vec = fields - .iter() - .flat_map(|(f, v)| { - vec![ - Frame::BulkString(f.clone()), - Frame::BulkString(v.clone()), - ] - }) - .collect(); - entry_frames.push(Frame::Array(field_frames.into())); - Frame::Array(entry_frames.into()) - }) - .collect(); - responses.push(Frame::Array(result_frames.into())); - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"ACK") { - match validate_mq_ack(cmd_args) { - Ok((queue_key, msg_ids)) => { - let effective_key = crate::workspace::workspace_key( - conn.workspace_id.as_ref(), - &queue_key, - ); - let ids: Vec = msg_ids - .iter() - .map(|(ms, seq)| StreamId { ms: *ms, seq: *seq }) - .collect(); - let mut db_guard = - ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - match db_guard.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - let group_name = Bytes::from_static(b"__mq_consumers"); - match stream.xack(&group_name, &ids) { - Ok(acked_count) => { - drop(db_guard); - // Emit MqAck WAL record for each acked ID - for (ms, seq) in &msg_ids { - let payload = crate::mq::wal::encode_mq_ack( - &effective_key, - *ms, - *seq, - ); - let mut wal_buf = Vec::new(); - crate::persistence::wal_v3::record::write_wal_v3_record( - &mut wal_buf, 0, - crate::persistence::wal_v3::record::WalRecordType::MqAck, - &payload, - ); - ctx.shard_databases - .wal_append(ctx.shard_id, Bytes::from(wal_buf)); - } - responses.push(Frame::Integer(acked_count as i64)); - } - Err(_) => responses.push(Frame::Integer(0)), - } - } - Ok(None) => responses.push(Frame::Integer(0)), - Err(_) => responses.push(Frame::Integer(0)), - } - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"DLQLEN") { - match validate_mq_dlqlen(cmd_args) { - Ok(queue_key) => { - let effective_key = crate::workspace::workspace_key( - conn.workspace_id.as_ref(), - &queue_key, - ); - let dlq_key = { - let mut buf = Vec::with_capacity(effective_key.len() + 8); - buf.extend_from_slice(&effective_key); - buf.extend_from_slice(b"::mq:dlq"); - Bytes::from(buf) - }; - let mut db_guard = - ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let len = match db_guard.get_stream_mut(&dlq_key) { - Ok(Some(stream)) => stream.length as i64, - _ => 0i64, - }; - responses.push(Frame::Integer(len)); - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"TRIGGER") { - match validate_mq_trigger(cmd_args) { - Ok((queue_key, callback_cmd, debounce_ms)) => { - let effective_key = crate::workspace::workspace_key( - conn.workspace_id.as_ref(), - &queue_key, - ); - let trig_key = if let Some(ws_id) = conn.workspace_id.as_ref() { - let ws_hex = ws_id.as_hex(); - let mut k = Vec::with_capacity(ws_hex.len() + 1 + queue_key.len()); - k.extend_from_slice(ws_hex.as_bytes()); - k.push(b':'); - k.extend_from_slice(&queue_key); - Bytes::from(k) - } else { - queue_key.clone() - }; - let entry = crate::mq::TriggerEntry { - queue_key: effective_key, - callback_cmd, - debounce_ms, - last_fire_ms: 0, - pending_fire_ms: 0, - }; - { - let mut guard = ctx.shard_databases.trigger_registry(ctx.shard_id); - let reg = guard.get_or_insert_with(|| { - Box::new(crate::mq::TriggerRegistry::new()) - }); - reg.register(trig_key, entry); - } - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"PUBLISH") { - match validate_mq_publish(cmd_args) { - Ok((queue_key, fields)) => { - let effective_key = crate::workspace::workspace_key( - conn.workspace_id.as_ref(), - &queue_key, - ); - if let Some(ref mut txn) = conn.active_cross_txn { - txn.record_mq(effective_key, fields); - responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); - } else { - responses.push(Frame::Error(Bytes::from_static( - b"ERR MQ PUBLISH requires an active transaction (use TXN BEGIN first)", - ))); - } - } - Err(e) => responses.push(e), - } - continue; - } - - // Unknown MQ subcommand - responses.push(Frame::Error(Bytes::from_static(ERR_MQ_UNKNOWN_SUB))); - continue; - } - - // --- MULTI --- - if cmd.eq_ignore_ascii_case(b"MULTI") { - if conn.in_cross_txn() { - responses.push(Frame::Error(Bytes::from_static(ERR_MULTI_TXN_CONFLICT))); - } else if conn.in_multi { - responses.push(Frame::Error(Bytes::from_static( - b"ERR MULTI calls can not be nested", - ))); - } else { - conn.in_multi = true; - conn.command_queue.clear(); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - continue; - } - - // --- EXEC --- - if cmd.eq_ignore_ascii_case(b"EXEC") { - if !conn.in_multi { - responses.push(Frame::Error(Bytes::from_static(b"ERR EXEC without MULTI"))); - } else { - conn.in_multi = false; - let result = execute_transaction_sharded( - &ctx.shard_databases, - ctx.shard_id, - &conn.command_queue, - conn.selected_db, - &ctx.cached_clock, - ); - conn.command_queue.clear(); - responses.push(result); - } - continue; - } - - // --- DISCARD --- - if cmd.eq_ignore_ascii_case(b"DISCARD") { - if !conn.in_multi { - responses.push(Frame::Error(Bytes::from_static( - b"ERR DISCARD without MULTI", - ))); - } else { - conn.in_multi = false; - conn.command_queue.clear(); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - continue; - } - - // --- Workspace key prefix injection --- - // MUST happen before key_to_shard() so the {ws_id} hash tag determines - // shard routing. This is the ONLY code path where workspace prefixing - // occurs (WS-07, WS-12). All subsequent dispatch uses cmd_args (shadowed). - let rewritten = conn - .workspace_id - .as_ref() - .map(|ws_id| workspace_rewrite_args(cmd, cmd_args, ws_id)); - let cmd_args: &[Frame] = rewritten.as_deref().unwrap_or(cmd_args); - - // --- BLOCKING COMMANDS --- - if cmd.eq_ignore_ascii_case(b"BLPOP") - || cmd.eq_ignore_ascii_case(b"BRPOP") - || cmd.eq_ignore_ascii_case(b"BLMOVE") - || cmd.eq_ignore_ascii_case(b"BZPOPMIN") - || cmd.eq_ignore_ascii_case(b"BZPOPMAX") - || cmd.eq_ignore_ascii_case(b"BLMPOP") - || cmd.eq_ignore_ascii_case(b"BRPOPLPUSH") - || cmd.eq_ignore_ascii_case(b"BZMPOP") - { - // Inside MULTI: queue as non-blocking variant - if conn.in_multi { - let nb_frame = convert_blocking_to_nonblocking(cmd, cmd_args); - conn.command_queue.push(nb_frame); - responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); - continue; - } - - // Flush accumulated responses before blocking - for resp in &responses { - codec.encode_frame(resp, &mut write_buf); - } - if !write_buf.is_empty() { - let data = write_buf.split().freeze(); - let (result, _): (std::io::Result, bytes::Bytes) = - stream.write_all(data).await; - if result.is_err() { - return (MonoioHandlerResult::Done, None); - } - } - - let blocking_response = handle_blocking_command_monoio( - cmd, - cmd_args, - conn.selected_db, - &ctx.shard_databases, - &ctx.blocking_registry, - ctx.shard_id, - ctx.num_shards, - &ctx.dispatch_tx, - &shutdown, - &ctx.spsc_notifiers, - ) - .await; - - // Encode blocking response directly - codec.encode_frame(&blocking_response, &mut write_buf); - responses.clear(); - break; // Blocking command ends the pipeline batch - } - - // --- MULTI queue mode: queue commands when in transaction --- - if conn.in_multi { - conn.command_queue.push(frame); - responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); - continue; - } - - // --- Cross-shard aggregation commands: KEYS, SCAN, DBSIZE --- - if ctx.num_shards > 1 { - if cmd.eq_ignore_ascii_case(b"KEYS") { - let mut response = crate::shard::coordinator::coordinate_keys( - cmd_args, - ctx.shard_id, - ctx.num_shards, - conn.selected_db, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - &ctx.cached_clock, - &(), // monoio: coordinator uses oneshot, not response_pool - ) - .await; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"SCAN") { - let mut response = crate::shard::coordinator::coordinate_scan( - cmd_args, - ctx.shard_id, - ctx.num_shards, - conn.selected_db, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - &ctx.cached_clock, - &(), // monoio: coordinator uses oneshot, not response_pool - ) - .await; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"DBSIZE") { - let response = crate::shard::coordinator::coordinate_dbsize( - ctx.shard_id, - ctx.num_shards, - conn.selected_db, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - &(), // monoio: coordinator uses oneshot, not response_pool - ) - .await; - responses.push(response); - continue; - } - - // --- Multi-key commands: MGET, MSET, DEL, UNLINK, EXISTS --- - if is_multi_key_command(cmd, cmd_args) { - let response = crate::shard::coordinator::coordinate_multi_key( - cmd, - cmd_args, - ctx.shard_id, - ctx.num_shards, - conn.selected_db, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - &ctx.cached_clock, - &(), // monoio: coordinator uses oneshot, not response_pool - ) - .await; - responses.push(response); - continue; - } - } - - // --- FT.* vector search commands --- - // Local shard: direct VectorStore access via ctx.shard_databases. - // Remote shards: SPSC dispatch. Works with any shard count (including 1). - if cmd.len() > 3 && cmd[..3].eq_ignore_ascii_case(b"FT.") { - if ctx.num_shards > 1 { - // Multi-shard: dispatch via SPSC - #[cfg(feature = "text-index")] - if cmd.eq_ignore_ascii_case(b"FT.AGGREGATE") { - // ── FT.AGGREGATE: multi-shard scatter-gather (Phase 152 Plan 03). ── - // Mirrors handler_sharded.rs:1458. scatter_text_aggregate acquires - // its own per-shard guards inside the single-shard block internally, - // so we never hold a MutexGuard across the .await below. - let parsed = - match crate::command::vector_search::ft_aggregate::parse_aggregate_args( - cmd_args, - ) { - Ok(p) => p, - Err(err_frame) => { - responses.push(err_frame); - continue; - } - }; - let response = crate::shard::scatter_aggregate::scatter_text_aggregate( - parsed.index_name, - parsed.query, - parsed.pipeline, - ctx.shard_id, - ctx.num_shards, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - ) - .await; - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { - // Check if this is a text query BEFORE trying parse_ft_search_args - // (which would return an error for non-KNN queries). - let query_bytes = cmd_args - .get(1) - .and_then(|f| crate::command::vector_search::extract_bulk(f)); - let is_text = query_bytes - .as_ref() - .map_or(false, |q| crate::command::vector_search::is_text_query(q)); - - // ── HYBRID multi-shard path (Phase 152 Plan 05, D-13) ────────── - #[cfg(feature = "text-index")] - { - match crate::command::vector_search::hybrid::parse_hybrid_modifier( - cmd_args, - ) { - Ok(Some(partial)) => { - let index_name = match cmd_args.first().and_then(|f| { - crate::command::vector_search::extract_bulk(f) - }) { - Some(b) => b, - None => { - responses.push(Frame::Error(Bytes::from_static( - b"ERR invalid index name", - ))); - continue; - } - }; - let text_query = match query_bytes.clone() { - Some(q) => q, - None => { - responses.push(Frame::Error(Bytes::from_static( - b"ERR invalid query", - ))); - continue; - } - }; - let (limit_offset, limit_count) = - crate::command::vector_search::parse_limit_clause(cmd_args); - let top_k = if limit_count == usize::MAX { - limit_offset.saturating_add(10).max(1) - } else { - limit_offset.saturating_add(limit_count).max(1) - }; - let hq = crate::command::vector_search::hybrid::HybridQuery { - index_name, - text_query, - dense_field: partial.dense_field, - dense_blob: partial.dense_blob, - sparse: partial.sparse, - weights: partial.weights, - k_per_stream: partial.k_per_stream, - top_k, - offset: limit_offset, - count: limit_count, - }; - // Phase 171 HYB-02 / SCAT-02: resolve AS_OF / TXN LSN - // ONCE on the coordinator and forward to the scatter - // helper so responders honor temporal snapshots. - let as_of_lsn = match resolve_ft_search_as_of_lsn( - cmd_args, - Some(&ctx.shard_databases), - ctx.shard_id, - conn.active_cross_txn.as_ref(), - ) { - Ok(lsn) => lsn, - Err(err_frame) => { - responses.push(err_frame); - continue; - } - }; - let response = - crate::shard::scatter_hybrid::scatter_hybrid_search( - hq, - as_of_lsn, - ctx.shard_id, - ctx.num_shards, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - ) - .await; - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response( - ws_id, - cmd, - &mut response, - ); - } - responses.push(response); - continue; - } - Ok(None) => { /* fall through */ } - Err(err_frame) => { - responses.push(err_frame); - continue; - } - } - } - - if is_text { - // ── Text FT.SEARCH: two-phase DFS scatter-gather ────────────────── - let index_name = match cmd_args - .first() - .and_then(|f| crate::command::vector_search::extract_bulk(f)) - { - Some(b) => b, - None => { - responses.push(Frame::Error(Bytes::from_static( - b"ERR invalid index name", - ))); - continue; - } - }; - let query_str = query_bytes.unwrap(); - - // B-01 SITE 2 FIX (Plan 152-06): FieldFilter short-circuit BEFORE - // the analyzer-first parse_result block. TAG queries (and Plan 07 - // NumericRange) route through the InvertedSearch fan-out — no - // analyzer touched, no field_idx resolution (the filter carries - // its own field name; search_tag resolves against tag_fields). - #[cfg(feature = "text-index")] - { - match crate::command::vector_search::pre_parse_field_filter( - query_str.as_ref(), - ) { - Ok(Some(clause)) => { - if let Some(filter) = clause.filter { - let (offset, count) = - crate::command::vector_search::parse_limit_clause( - cmd_args, - ); - let top_k = if count == usize::MAX { - 10000 - } else { - offset.saturating_add(count) - } - .max(1); - let response = - crate::shard::coordinator::scatter_text_search_filter( - index_name, - filter, - top_k, - offset, - count, - ctx.shard_id, - ctx.num_shards, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - ) - .await; - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response( - ws_id, - cmd, - &mut response, - ); - } - responses.push(response); - continue; - } - } - Ok(None) => { /* fall through to BM25 path */ } - Err(e) => { - responses.push(Frame::Error(Bytes::from(e.to_owned()))); - continue; - } - } - } - - // Parse query and resolve field_idx inside a block scope so the - // MutexGuard from text_store() is dropped BEFORE .await. - // We use the TextIndex's own field_analyzers (same pipeline used at index time). - type ParseResult = Result< - (Vec, Option), - String, - >; - let parse_result: ParseResult = { - let ts = ctx.shard_databases.text_store(ctx.shard_id); - match ts.get_index(&index_name) { - None => Err("ERR no such index".to_owned()), - Some(text_index) => { - match text_index.field_analyzers.first() { - None => Err("ERR index has no TEXT fields".to_owned()), - Some(analyzer) => { - // analyzer borrows text_index which borrows ts — all in this block. - let parsed = - crate::command::vector_search::parse_text_query( - &query_str, analyzer, - ); - match parsed { - Err(e) => Err(e.to_owned()), - Ok(clause) => { - let field_idx = match &clause.field_name { - None => Ok(None), - Some(field_name) => match text_index - .text_fields - .iter() - .position(|f| { - f.field_name - .as_ref() - .eq_ignore_ascii_case( - field_name.as_ref(), - ) - }) { - Some(idx) => Ok(Some(idx)), - None => Err(format!( - "ERR unknown field '{}'", - String::from_utf8_lossy( - field_name - ) - )), - }, - }; - field_idx.map(|idx| (clause.terms, idx)) - } - } - } - } - } - } - }; // MutexGuard dropped here - - let (query_terms, field_idx) = match parse_result { - Ok(t) => t, - Err(e) => { - responses.push(Frame::Error(Bytes::from(e))); - continue; - } - }; - - let (offset, count) = - crate::command::vector_search::parse_limit_clause(cmd_args); - let top_k = if count == usize::MAX { - 10000 - } else { - offset.saturating_add(count) - } - .max(1); - - // Parse optional HIGHLIGHT/SUMMARIZE clauses from args. - let highlight_opts = - crate::command::vector_search::parse_highlight_clause(cmd_args); - let summarize_opts = - crate::command::vector_search::parse_summarize_clause(cmd_args); - - let response = crate::shard::coordinator::scatter_text_search( - index_name, - query_terms, - field_idx, - top_k, - offset, - count, - ctx.shard_id, - ctx.num_shards, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - highlight_opts, - summarize_opts, - ) - .await; - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - - // ── Vector FT.SEARCH (KNN / SPARSE): existing path ──────────────── - // Phase 171 SCAT-01: resolve AS_OF / TXN snapshot LSN ONCE on - // the coordinator and forward through the scatter helper so - // every responder honors the same temporal snapshot. - let response = - match crate::command::vector_search::parse_ft_search_args(cmd_args) { - Ok((index_name, query_blob, k, filter, _offset, _count)) => { - if filter.is_some() { - Frame::Error(Bytes::from_static( - b"ERR FILTER not supported in multi-shard mode yet", - )) - } else { - match resolve_ft_search_as_of_lsn( - cmd_args, - Some(&ctx.shard_databases), - ctx.shard_id, - conn.active_cross_txn.as_ref(), - ) { - Err(err_frame) => err_frame, - Ok(as_of_lsn) => { - crate::shard::coordinator::scatter_vector_search_remote( - index_name, - query_blob, - k, - as_of_lsn, - ctx.shard_id, - ctx.num_shards, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - ) - .await - } - } - } - } - Err(err_frame) => err_frame, - }; - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"FT.CREATE") - || cmd.eq_ignore_ascii_case(b"FT.DROPINDEX") - { - // Broadcast to ALL shards so every shard has the index - let response = crate::shard::coordinator::broadcast_vector_command( - std::sync::Arc::new(frame), - ctx.shard_id, - ctx.num_shards, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - ) - .await; - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"FT.INFO") { - let response = { - let vs = ctx.shard_databases.vector_store(ctx.shard_id); - let ts = ctx.shard_databases.text_store(ctx.shard_id); - crate::command::vector_search::ft_info(&vs, &ts, cmd_args) - }; - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"FT._LIST") { - let response = { - let vs = ctx.shard_databases.vector_store(ctx.shard_id); - crate::command::vector_search::ft_list(&vs) - }; - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"FT.COMPACT") { - let response = { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - let mut ts = ctx.shard_databases.text_store(ctx.shard_id); - crate::command::vector_search::ft_compact(&mut vs, &mut ts, cmd_args) - }; - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"FT.CACHESEARCH") { - let response = { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - crate::command::vector_search::cache_search::ft_cachesearch( - &mut vs, cmd_args, - ) - }; - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"FT.CONFIG") { - let response = { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - let mut ts = ctx.shard_databases.text_store(ctx.shard_id); - crate::command::vector_search::ft_config(&mut vs, &mut ts, cmd_args) - }; - responses.push(response); - continue; - } - // FT.RECOMMEND, FT.NAVIGATE, FT.EXPAND need db/graph — dispatch locally - if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") - || cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") - || cmd.eq_ignore_ascii_case(b"FT.EXPAND") - { - let response = { - let sdb = &ctx.shard_databases; - let mut vs = sdb.vector_store(ctx.shard_id); - if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") { - let mut db_guard = sdb.write_db(ctx.shard_id, 0); - crate::command::vector_search::recommend::ft_recommend( - &mut vs, - cmd_args, - Some(&mut *db_guard), - ) - } else if cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") { - #[cfg(feature = "graph")] - { - let graph_guard = sdb.graph_store_read(ctx.shard_id); - crate::command::vector_search::navigate::ft_navigate( - &mut vs, - Some(&graph_guard), - cmd_args, - None, - ) - } - #[cfg(not(feature = "graph"))] - { - Frame::Error(Bytes::from_static( - b"ERR FT.NAVIGATE requires graph feature", - )) - } - } else { - // FT.EXPAND - #[cfg(feature = "graph")] - { - let graph_guard = sdb.graph_store_read(ctx.shard_id); - crate::command::vector_search::ft_expand(&graph_guard, cmd_args) - } - #[cfg(not(feature = "graph"))] - { - Frame::Error(Bytes::from_static( - b"ERR FT.EXPAND requires graph feature", - )) - } - } - }; - responses.push(response); - continue; - } - responses.push(Frame::Error(Bytes::from_static(b"ERR unknown FT command"))); - continue; - } else { - // Single-shard: no SPSC channels needed. - // Dispatch directly to shard's VectorStore via shared access. - // - // ── 154-01 single-shard FT.AGGREGATE fast path ──────────────── - // scatter_text_aggregate internally fast-paths num_shards == 1 - // to execute_local_full with locally-acquired guards dropped - // before any .await. Calling the scatter entry here (instead - // of execute_local_full directly) keeps the dispatch body - // byte-symmetric with the multi-shard site above. - #[cfg(feature = "text-index")] - if cmd.eq_ignore_ascii_case(b"FT.AGGREGATE") { - let parsed = - match crate::command::vector_search::ft_aggregate::parse_aggregate_args( - cmd_args, - ) { - Ok(p) => p, - Err(err_frame) => { - responses.push(err_frame); - continue; - } - }; - let response = crate::shard::scatter_aggregate::scatter_text_aggregate( - parsed.index_name, - parsed.query, - parsed.pipeline, - ctx.shard_id, - ctx.num_shards, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - ) - .await; - responses.push(response); - continue; - } - // - // ── 151-03 single-shard text FT.SEARCH fast path ────────────── - // Bare text queries (exact / fuzzy / prefix / field-targeted) - // are not understood by `ft_search()` (which only parses - // KNN / SPARSE / HYBRID) — they would otherwise return - // `ERR invalid KNN query syntax`. Route them directly to - // `execute_text_search_local` here, the same function the - // multi-shard path uses once its per-shard scatter has - // aggregated IDFs. We skip HYBRID (existing ft_search - // handles it) and KNN/SPARSE (is_text_query returns false). - #[cfg(feature = "text-index")] - if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { - if let Some(Frame::BulkString(query_bytes)) = cmd_args.get(1) { - match crate::command::vector_search::parse_hybrid_modifier(cmd_args) { - Ok(Some(_)) => { - // HYBRID present — defer to existing ft_search() below. - } - Err(frame_err) => { - responses.push(frame_err); - continue; - } - Ok(None) => { - if crate::command::vector_search::is_text_query( - query_bytes.as_ref(), - ) { - // Step 1: index_name from cmd_args[0]. - let index_name = match cmd_args.first() { - Some(Frame::BulkString(b)) => b.clone(), - _ => { - responses.push(Frame::Error(Bytes::from_static( - b"ERR wrong number of arguments for FT.SEARCH", - ))); - continue; - } - }; - // B-01 SITE 2 FIX (single-shard 151-03 fast path, Plan 152-06): - // FieldFilter short-circuit BEFORE the analyzer lookup and - // BEFORE the text_fields.is_empty() bail. - #[cfg(feature = "text-index")] - match crate::command::vector_search::pre_parse_field_filter( - query_bytes.as_ref(), - ) { - Ok(Some(clause)) => { - if clause.filter.is_some() { - let (offset, count) = crate::command::vector_search::parse_limit_clause(cmd_args); - let top_k = if count == usize::MAX { - 10000 - } else { - offset.saturating_add(count) - } - .max(1); - let ts_guard = ctx - .shard_databases - .text_store(ctx.shard_id); - let response = match ts_guard - .get_index(&index_name) - { - None => Frame::Error(Bytes::from_static( - b"ERR no such index", - )), - Some(text_index) => { - let results = crate::command::vector_search::ft_text_search::execute_query_on_index( - text_index, &clause, None, None, top_k, - ); - crate::command::vector_search::ft_text_search::build_text_response( - &results, offset, count, - ) - } - }; - drop(ts_guard); - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() - { - strip_workspace_prefix_from_response( - ws_id, - cmd, - &mut response, - ); - } - responses.push(response); - continue; - } - } - Ok(None) => { /* fall through */ } - Err(e) => { - responses.push(Frame::Error( - Bytes::copy_from_slice(e.as_bytes()), - )); - continue; - } - } - // Step 2: acquire ts guard (single-shard monoio). - let ts_guard = ctx.shard_databases.text_store(ctx.shard_id); - // Step 3: resolve index. - let text_index = match ts_guard.get_index(&index_name) { - Some(idx) => idx, - None => { - drop(ts_guard); - responses.push(Frame::Error(Bytes::from_static( - b"ERR no such index", - ))); - continue; - } - }; - // Step 4: ensure at least one TEXT field. - if text_index.text_fields.is_empty() { - drop(ts_guard); - responses.push(Frame::Error(Bytes::from_static( - b"ERR index has no TEXT fields", - ))); - continue; - } - // Step 5: parse the query via the index's first analyzer. - let analyzer = match text_index.field_analyzers.first() { - Some(a) => a, - None => { - drop(ts_guard); - responses.push(Frame::Error(Bytes::from_static( - b"ERR index has no TEXT fields", - ))); - continue; - } - }; - let clause = - match crate::command::vector_search::parse_text_query( - query_bytes.as_ref(), - analyzer, - ) { - Ok(c) => c, - Err(msg) => { - drop(ts_guard); - responses.push(Frame::Error( - Bytes::copy_from_slice(msg.as_bytes()), - )); - continue; - } - }; - // Step 5b: resolve field_idx. - let field_idx = match &clause.field_name { - None => None, - Some(field_name) => { - match text_index.text_fields.iter().position(|f| { - f.field_name - .as_ref() - .eq_ignore_ascii_case(field_name.as_ref()) - }) { - Some(idx) => Some(idx), - None => { - let bad_name = field_name.clone(); - drop(ts_guard); - responses.push(Frame::Error(Bytes::from( - format!( - "ERR unknown field '{}'", - String::from_utf8_lossy(&bad_name) - ), - ))); - continue; - } - } - } - }; - // Step 6: extract query_terms. - let query_terms = clause.terms; - // Step 7: LIMIT parsing + top_k cap (T-151-03-02). - let (offset, count) = - crate::command::vector_search::parse_limit_clause( - cmd_args, - ); - let top_k = if count == usize::MAX { - 10000 - } else { - offset.saturating_add(count) - } - .max(1); - // Step 8: HIGHLIGHT / SUMMARIZE options. - let highlight_opts = - crate::command::vector_search::parse_highlight_clause( - cmd_args, - ); - let summarize_opts = - crate::command::vector_search::parse_summarize_clause( - cmd_args, - ); - // Step 9: acquire DB read guard iff post-processing is needed. - let db_guard_opt = if highlight_opts.is_some() - || summarize_opts.is_some() - { - Some(ctx.shard_databases.read_db(ctx.shard_id, 0)) - } else { - None - }; - // Step 10: execute + optional post-processing. - let mut response = - crate::command::vector_search::execute_text_search_local( - &ts_guard, - &index_name, - field_idx, - &query_terms, - top_k, - offset, - count, - ); - if let Some(ref db_guard) = db_guard_opt { - let term_strings: Vec = query_terms - .iter() - .map(|qt| qt.text.clone()) - .collect(); - crate::command::vector_search::apply_post_processing( - &mut response, - &term_strings, - text_index, - db_guard, - highlight_opts.as_ref(), - summarize_opts.as_ref(), - ); - } - // Explicit drop order: inner db_guard first, outer ts_guard last. - drop(db_guard_opt); - drop(ts_guard); - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response( - ws_id, - cmd, - &mut response, - ); - } - responses.push(response); - continue; - } - } - } - } - } - let response = { - let shard_databases_ref = &ctx.shard_databases; - let mut vs = shard_databases_ref.vector_store(ctx.shard_id); - let mut ts = shard_databases_ref.text_store(ctx.shard_id); - if cmd.eq_ignore_ascii_case(b"FT.CREATE") { - crate::command::vector_search::ft_create(&mut vs, &mut ts, cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { - // Resolve AS_OF temporal clause + TXN snapshot precedence (TEMP-04, ACID-09). - let as_of_lsn = match resolve_ft_search_as_of_lsn( - cmd_args, - Some(shard_databases_ref), - ctx.shard_id, - conn.active_cross_txn.as_ref(), - ) { - Ok(lsn) => lsn, - Err(err_frame) => { - responses.push(err_frame); - continue; - } - }; - // Detect SESSION keyword to provide database access for sorted set tracking - let has_session = cmd_args.iter().any(|a| { - if let Frame::BulkString(b) = a { - b.eq_ignore_ascii_case(b"SESSION") - } else { - false - } - }); - if has_session { - let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); - crate::command::vector_search::ft_search( - &mut vs, - cmd_args, - Some(&mut *db_guard), - Some(&*ts), - as_of_lsn, - ) - } else { - crate::command::vector_search::ft_search( - &mut vs, - cmd_args, - None, - Some(&*ts), - as_of_lsn, - ) - } - } else if cmd.eq_ignore_ascii_case(b"FT.DROPINDEX") { - let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); - crate::command::vector_search::ft_dropindex( - &mut vs, - &mut ts, - Some(&mut *db_guard), - cmd_args, - ) - } else if cmd.eq_ignore_ascii_case(b"FT.INFO") { - crate::command::vector_search::ft_info(&vs, &ts, cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT._LIST") { - crate::command::vector_search::ft_list(&vs) - } else if cmd.eq_ignore_ascii_case(b"FT.COMPACT") { - crate::command::vector_search::ft_compact(&mut vs, &mut ts, cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT.CACHESEARCH") { - crate::command::vector_search::cache_search::ft_cachesearch( - &mut vs, cmd_args, - ) - } else if cmd.eq_ignore_ascii_case(b"FT.CONFIG") { - crate::command::vector_search::ft_config(&mut vs, &mut ts, cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") { - let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); - crate::command::vector_search::recommend::ft_recommend( - &mut vs, - cmd_args, - Some(&mut *db_guard), - ) - } else if cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") { - #[cfg(feature = "graph")] - { - let graph_guard = - shard_databases_ref.graph_store_read(ctx.shard_id); - crate::command::vector_search::navigate::ft_navigate( - &mut vs, - Some(&graph_guard), - cmd_args, - None, - ) - } - #[cfg(not(feature = "graph"))] - { - Frame::Error(Bytes::from_static( - b"ERR FT.NAVIGATE requires graph feature", - )) - } - } else if cmd.eq_ignore_ascii_case(b"FT.EXPAND") { - #[cfg(feature = "graph")] - { - let graph_guard = - shard_databases_ref.graph_store_read(ctx.shard_id); - crate::command::vector_search::ft_expand(&graph_guard, cmd_args) - } - #[cfg(not(feature = "graph"))] - { - Frame::Error(Bytes::from_static( - b"ERR FT.EXPAND requires graph feature", - )) - } - } else { - Frame::Error(Bytes::from_static(b"ERR unknown FT.* command")) - } - }; - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - } - - // --- GRAPH.* graph commands --- - #[cfg(feature = "graph")] - if cmd.len() > 6 && cmd[..6].eq_ignore_ascii_case(b"GRAPH.") { - let (response, wal_records, cypher_intents, cypher_undo_ops) = - if crate::command::graph::is_graph_write_cmd(cmd) - || (cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") - && crate::command::graph::is_cypher_write_query(cmd_args)) - { - let mut gs = ctx.shard_databases.graph_store_write(ctx.shard_id); - let (resp, cypher_intents, undo_ops) = if cmd - .eq_ignore_ascii_case(b"GRAPH.QUERY") - { - // Phase 167 (CYP-01/02): capture Cypher-created - // nodes/edges so TXN.ABORT can roll them back via - // CrossStoreTxn::record_graph. - crate::command::graph::graph_query_or_write(&mut gs, cmd_args) - } else { - ( - crate::command::graph::dispatch_graph_write(&mut gs, cmd, cmd_args), - Vec::new(), - Vec::new(), - ) - }; - let records = gs.drain_wal(); - (resp, records, cypher_intents, undo_ops) - } else { - let gs = ctx.shard_databases.graph_store_read(ctx.shard_id); - let resp = crate::command::graph::dispatch_graph_read(&gs, cmd, cmd_args); - (resp, Vec::new(), Vec::new(), Vec::new()) - }; - // Phase 166: record graph intent for TXN rollback. - // Captures explicit ADDNODE/ADDEDGE by response id plus - // Phase 167 Cypher CREATE/MERGE via intents returned from - // graph_query_or_write. - if let Some(txn) = conn.active_cross_txn.as_mut() { - let is_node = cmd.eq_ignore_ascii_case(b"GRAPH.ADDNODE"); - let is_edge = cmd.eq_ignore_ascii_case(b"GRAPH.ADDEDGE"); - if is_node || is_edge { - if let Frame::Integer(id) = &response { - if let Some(gname) = cmd_args.first().and_then(|f| extract_bytes(f)) { - txn.record_graph(*id as u64, is_node, gname); - } - } - } - if !cypher_intents.is_empty() { - if let Some(gname) = cmd_args.first().and_then(|f| extract_bytes(f)) { - for intent in &cypher_intents { - txn.record_graph(intent.entity_id, intent.is_node, gname.clone()); - } - } - } - // Phase 174 FIX-01: push undo ops for SET/DELETE/MERGE rollback. - for undo_op in cypher_undo_ops { - txn.record_graph_undo(undo_op); - } - } - for record in wal_records { - ctx.shard_databases - .wal_append(ctx.shard_id, bytes::Bytes::from(record)); - } - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - - // --- Routing: keyless, local, or remote --- - let target_shard = - extract_primary_key(cmd, cmd_args).map(|key| key_to_shard(key, ctx.num_shards)); - - let is_local = match target_shard { - None => true, - Some(s) if s == ctx.shard_id => true, - _ => false, - }; - - // Affinity sampling: record shard target for migration decision. - // Migration is deferred until AFTER the current batch is fully processed. - if let (Some(tracker), Some(target)) = (&mut conn.affinity_tracker, target_shard) { - if let Some(migrate_to) = tracker.record(target) { - if !conn.in_multi - && conn.subscription_count == 0 - && !conn.tracking_state.enabled - { - conn.migration_target = Some(migrate_to); - } - } - } - - // Pre-classify write commands for AOF + tracking - let is_write = if ctx.aof_tx.is_some() || conn.tracking_state.enabled { - metadata::is_write(cmd) - } else { - false - }; - - if is_local { - // LOCAL PATH: split into read/write to avoid exclusive lock on reads. - // Using read_db for local reads eliminates RwLock contention with - // cross-shard shared reads from other shard threads. - if metadata::is_write(cmd) { - // WRITE PATH: eviction + dispatch under write lock. - // When disk offload is enabled, use async spill: evicted keys - // are sent to SpillThread for background pwrite to NVMe. - let rt = ctx.runtime_config.read(); - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let evict_result = if let Some(ref sender) = ctx.spill_sender { - let mut fid = ctx.spill_file_id.get(); - let dir = ctx - .disk_offload_dir - .as_deref() - .unwrap_or(std::path::Path::new(".")); - let res = try_evict_if_needed_async_spill( - &mut guard, - &rt, - sender, - dir, - &mut fid, - conn.selected_db, - ); - ctx.spill_file_id.set(fid); - res - } else { - try_evict_if_needed(&mut guard, &rt) - }; - if let Err(oom_frame) = evict_result { - drop(guard); - drop(rt); - responses.push(oom_frame); - continue; - } - drop(rt); - - // KV undo-log capture for active cross-store transactions. - // MUST happen BEFORE dispatch() overwrites the database entry. - if let Some(ref mut txn) = conn.active_cross_txn { - if cmd.eq_ignore_ascii_case(b"DEL") || cmd.eq_ignore_ascii_case(b"UNLINK") { - // Multi-key DEL: iterate all args for undo capture - for arg in cmd_args.iter() { - if let Frame::BulkString(key_bytes) = arg { - if let Some(old_entry) = guard.get(key_bytes.as_ref()).cloned() - { - txn.kv_undo.record_delete(key_bytes.clone(), old_entry); - let lsn = txn.snapshot_lsn; - let tid = txn.txn_id; - ctx.shard_databases.kv_intents(ctx.shard_id).record_write( - key_bytes.clone(), - lsn, - tid, - ); - } - // Key not found: nothing to undo, no intent registered - } - } - } else if let Some(key) = - crate::server::conn::shared::extract_primary_key(cmd, cmd_args) - { - // SET / HSET / INCR / etc. — single primary key - let old_entry = guard.get(key.as_ref()).cloned(); - let lsn = txn.snapshot_lsn; - let tid = txn.txn_id; - match old_entry { - None => txn.kv_undo.record_insert(key.clone()), - Some(entry) => txn.kv_undo.record_update(key.clone(), entry), - } - ctx.shard_databases.kv_intents(ctx.shard_id).record_write( - key.clone(), - lsn, - tid, - ); - } - } - - let dispatch_start = std::time::Instant::now(); - let result = - dispatch(&mut guard, cmd, cmd_args, &mut conn.selected_db, db_count); - let elapsed_us = dispatch_start.elapsed().as_micros() as u64; - if let Ok(cmd_str) = std::str::from_utf8(cmd) { - crate::admin::metrics_setup::record_command(cmd_str, elapsed_us); - } - if let Frame::Array(ref args) = frame { - crate::admin::metrics_setup::global_slowlog().maybe_record( - elapsed_us, - args.as_slice(), - peer_addr.as_bytes(), - conn.client_name - .as_ref() - .map_or(b"" as &[u8], |n| n.as_ref()), - ); - } - - let response = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => { - should_quit = true; - f - } - }; - - // AOF logging for successful local writes - if !matches!(response, Frame::Error(_)) && is_write { - if let Some(ref tx) = ctx.aof_tx { - let serialized = aof::serialize_command(&frame); - let _ = tx.try_send(AofMessage::Append(serialized)); - } - } - - // Auto-index HSET into vector/text stores (if key matches index prefix) - if !matches!(response, Frame::Error(_)) && cmd.eq_ignore_ascii_case(b"HSET") { - if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { - let inserted = { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - let mut ts = ctx.shard_databases.text_store(ctx.shard_id); - crate::shard::spsc_handler::auto_index_hset_public( - &mut vs, - &mut *ts, - key.as_ref(), - cmd_args, - ) - }; - // Phase 166 (Plan 02): if inside an active cross-store - // TXN, push one VectorIntent per (index_name, key_hash) - // so TXN.ABORT (Plan 166-03) can tombstone via - // MutableSegment::mark_deleted_by_key_hash. Non-TXN - // paths discard the tuples — auto_index_hset already - // performed the append. - if let Some(txn) = conn.active_cross_txn.as_mut() { - for (index_name, key_hash) in inserted { - txn.record_vector(key_hash, index_name); - } - } - } - } - - // Post-dispatch wakeup hooks for producer commands - if !matches!(response, Frame::Error(_)) { - let needs_wake = cmd.eq_ignore_ascii_case(b"LPUSH") - || cmd.eq_ignore_ascii_case(b"RPUSH") - || cmd.eq_ignore_ascii_case(b"LMOVE") - || cmd.eq_ignore_ascii_case(b"ZADD"); - if needs_wake { - if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { - let mut reg = ctx.blocking_registry.borrow_mut(); - if cmd.eq_ignore_ascii_case(b"LPUSH") - || cmd.eq_ignore_ascii_case(b"RPUSH") - || cmd.eq_ignore_ascii_case(b"LMOVE") - { - crate::blocking::wakeup::try_wake_list_waiter( - &mut reg, - &mut guard, - conn.selected_db, - &key, - ); - } else { - crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, - &mut guard, - conn.selected_db, - &key, - ); - } - } - } - } - - drop(guard); - - // Track key on write / invalidate tracked keys - if conn.tracking_state.enabled && !matches!(response, Frame::Error(_)) { - if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { - let senders = ctx - .tracking_table - .borrow_mut() - .invalidate_key(&key, client_id); - if !senders.is_empty() { - let push = crate::tracking::invalidation::invalidation_push(&[key]); - for tx in senders { - let _ = tx.try_send(push.clone()); - } - } - } - } - let mut response = apply_resp3_conversion(cmd, response, conn.protocol_version); - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - } else { - // Snapshot visibility filter for active cross-store transactions. - // MVCC: hide keys written by uncommitted foreign transactions. - if conn.in_cross_txn() { - if let Some(ref txn) = conn.active_cross_txn { - if let Some(key) = - crate::server::conn::shared::extract_primary_key(cmd, cmd_args) - { - let snapshot_lsn = txn.snapshot_lsn; - let my_txn_id = txn.txn_id; - // Clone committed treemap to release vector_store lock - // before acquiring kv_intents lock (lock ordering). - let committed = { - let vs = ctx.shard_databases.vector_store(ctx.shard_id); - vs.txn_manager().committed_treemap().clone() - }; - let visible = { - let intents = ctx.shard_databases.kv_intents(ctx.shard_id); - intents.is_key_visible( - key.as_ref(), - snapshot_lsn, - my_txn_id, - &committed, - ) - }; - if !visible { - responses.push(Frame::Null); - continue; - } - } - } - } - - // READ PATH: shared lock — no contention with other shards' reads - let guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); - let now_ms = ctx.cached_clock.ms(); - let dispatch_start = std::time::Instant::now(); - let result = dispatch_read( - &guard, - cmd, - cmd_args, - now_ms, - &mut conn.selected_db, - db_count, - ); - let elapsed_us = dispatch_start.elapsed().as_micros() as u64; - if let Ok(cmd_str) = std::str::from_utf8(cmd) { - crate::admin::metrics_setup::record_command(cmd_str, elapsed_us); - } - if let Frame::Array(ref args) = frame { - crate::admin::metrics_setup::global_slowlog().maybe_record( - elapsed_us, - args.as_slice(), - peer_addr.as_bytes(), - conn.client_name - .as_ref() - .map_or(b"" as &[u8], |n| n.as_ref()), - ); - } - drop(guard); - - let response = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => { - should_quit = true; - f - } - }; - - // Track key on local read - if conn.tracking_state.enabled && !conn.tracking_state.bcast { - if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { - ctx.tracking_table.borrow_mut().track_key( - client_id, - &key, - conn.tracking_state.noloop, - ); - } - } - let mut response = apply_resp3_conversion(cmd, response, conn.protocol_version); - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - } // end read/write split - - // (tracking and response push handled inside read/write branches above) - } else if let Some(target) = target_shard { - // TXN cross-shard guard: cross-shard writes bypass the undo log and - // cannot be rolled back on TXN.ABORT. Return an explicit error instead - // of silently permitting writes that resist rollback. - if conn.in_cross_txn() && metadata::is_write(cmd) { - responses.push(Frame::Error(bytes::Bytes::from_static( - crate::command::transaction::ERR_TXN_CROSS_SHARD, - ))); - continue; - } - // SHARED-READ FAST PATH: cross-shard reads bypass SPSC dispatch entirely. - // By this point conn.in_multi is false (MULTI queuing happens earlier with `continue`). - // Read commands execute directly on the target shard's database via RwLock read guard, - // avoiding ~88us of two async scheduling hops through the SPSC channel. - // - // Guard: if there are already pending writes for this target shard in the - // current pipeline batch, we must NOT take the fast path -- the read would - // execute before the deferred writes, violating command ordering. Fall through - // to SPSC dispatch to preserve pipeline semantics. - if !metadata::is_write(cmd) - && !remote_groups.contains_key(&target) - && is_dispatch_read_supported(cmd) - { - let guard = ctx.shard_databases.read_db(target, conn.selected_db); - let now_ms = ctx.cached_clock.ms(); - let result = dispatch_read( - &guard, - cmd, - cmd_args, - now_ms, - &mut conn.selected_db, - db_count, - ); - drop(guard); - let response = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => { - should_quit = true; - f - } - }; - // Client tracking for cross-shard reads - if conn.tracking_state.enabled && !conn.tracking_state.bcast { - if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { - ctx.tracking_table.borrow_mut().track_key( - client_id, - &key, - conn.tracking_state.noloop, - ); - } - } - let mut response = apply_resp3_conversion(cmd, response, conn.protocol_version); - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - // Cross-shard write: deferred SPSC dispatch. - // When workspace rewriting occurred, rebuild the frame with - // prefixed args so the target shard stores the correct key. - let dispatch_frame = if rewritten.is_some() { - let mut parts = Vec::with_capacity(1 + cmd_args.len()); - parts.push(Frame::BulkString(Bytes::copy_from_slice(cmd))); - parts.extend_from_slice(cmd_args); - Frame::Array(parts.into()) - } else { - frame.clone() - }; - let resp_idx = responses.len(); - responses.push(Frame::Null); // placeholder, filled after batch dispatch - // Pre-compute AOF bytes before moving frame into Arc - let aof_bytes = if ctx.aof_tx.is_some() && metadata::is_write(cmd) { - Some(aof::serialize_command(&dispatch_frame)) - } else { - None - }; - let cmd_bytes = if let Frame::Array(ref args) = dispatch_frame { - extract_bytes(&args[0]).unwrap_or_default() - } else { - Bytes::new() - }; - remote_groups.entry(target).or_default().push(( - resp_idx, - std::sync::Arc::new(dispatch_frame), - aof_bytes, - cmd_bytes, - )); - } - } - - // Phase 2a: Flush accumulated PUBLISH batches as PubSubPublishBatch messages - if !publish_batches.is_empty() { - let mut batch_slots: Vec<( - std::sync::Arc, - Vec, - )> = Vec::new(); - { - let mut producers = ctx.dispatch_tx.borrow_mut(); - for (target, entries) in publish_batches.drain() { - let n = entries.len(); - let slot = std::sync::Arc::new( - crate::shard::dispatch::PubSubResponseSlot::with_counts(1, n), - ); - let resp_indices: Vec = entries.iter().map(|(idx, _, _)| *idx).collect(); - let pairs: Vec<(Bytes, Bytes)> = - entries.into_iter().map(|(_, ch, msg)| (ch, msg)).collect(); - - let idx = ChannelMesh::target_index(ctx.shard_id, target); - let batch_msg = ShardMessage::PubSubPublishBatch { - pairs, - slot: slot.clone(), - }; - if producers[idx].try_push(batch_msg).is_ok() { - ctx.spsc_notifiers[target].notify_one(); - } else { - slot.add(0); // push failed, mark as done - } - batch_slots.push((slot, resp_indices)); - } - } - // Resolve all batch slots - for (slot, resp_indices) in &batch_slots { - crate::shard::dispatch::PubSubResponseFuture::new(slot.clone()).await; - for (i, resp_idx) in resp_indices.iter().enumerate() { - let remote_count = slot.counts[i].load(std::sync::atomic::Ordering::Relaxed); - if remote_count > 0 { - if let Frame::Integer(ref mut total) = responses[*resp_idx] { - *total += remote_count; - } - } - } - } - } - - // Phase 2b: Dispatch all deferred remote commands as batched PipelineBatch - // messages (one per target shard), await all in parallel. - if !remote_groups.is_empty() { - reply_futures.clear(); - - let mut oneshot_futures: Vec<( - Vec<(usize, Option, Bytes)>, - channel::OneshotReceiver>, - )> = Vec::new(); - for (target, entries) in remote_groups.drain() { - let (reply_tx, reply_rx) = channel::oneshot(); - let (meta, commands): ( - Vec<(usize, Option, Bytes)>, - Vec>, - ) = entries - .into_iter() - .map(|(idx, arc_frame, aof, cmd)| ((idx, aof, cmd), arc_frame)) - .unzip(); - - let msg = ShardMessage::PipelineBatch { - db_index: conn.selected_db, - commands, - reply_tx, - }; - let target_idx = ChannelMesh::target_index(ctx.shard_id, target); - { - let mut pending = msg; - loop { - let push_result = { - let mut producers = ctx.dispatch_tx.borrow_mut(); - producers[target_idx].try_push(pending) - }; - match push_result { - Ok(()) => { - tracing::trace!( - "Shard {}: pushed PipelineBatch to shard {}, notifying", - ctx.shard_id, - target - ); - ctx.spsc_notifiers[target].notify_one(); - break; - } - Err(val) => { - pending = val; - monoio::time::sleep(std::time::Duration::from_micros(10)).await; - } - } - } - } - oneshot_futures.push((meta, reply_rx)); - } - - // Poll all shard responses via pending_wakers relay (monoio cross-thread waker fix). - // monoio's !Send executor doesn't see cross-thread Waker::wake() from flume. - // Instead, the connection task registers its waker in pending_wakers; the event - // loop drains and wakes them after every SPSC cycle (~1ms). On wake, try_recv() - // checks if the response arrived; if not, re-register and yield again. - for (meta, reply_rx) in oneshot_futures.drain(..) { - tracing::trace!( - "Shard {}: awaiting cross-shard response via pending_wakers", - ctx.shard_id - ); - let shard_responses = { - let pw = pending_wakers.clone(); - loop { - match reply_rx.try_recv() { - Ok(value) => break Ok(value), - Err(flume::TryRecvError::Disconnected) => break Err(()), - Err(flume::TryRecvError::Empty) => { - // Yield once: register waker, return Pending, then Ready on wake. - let mut yielded = false; - std::future::poll_fn(|cx| { - if yielded { - std::task::Poll::Ready(()) - } else { - yielded = true; - pw.borrow_mut().push(cx.waker().clone()); - std::task::Poll::Pending - } - }) - .await; - // After wake, loop back to try_recv - } - } - } - }; - let shard_responses = match shard_responses { - Ok(r) => r, - Err(()) => { - for (resp_idx, _, _) in &meta { - responses[*resp_idx] = Frame::Error(Bytes::from_static( - b"ERR cross-shard dispatch failed", - )); - } - continue; - } - }; - for ((resp_idx, aof_bytes, cmd_name), resp) in meta.into_iter().zip(shard_responses) - { - // AOF logging for successful remote writes - if let Some(bytes) = aof_bytes { - if !matches!(resp, Frame::Error(_)) { - if let Some(ref tx) = ctx.aof_tx { - let _ = tx.try_send(AofMessage::Append(bytes)); - } - } - } - let resp = apply_resp3_conversion(&cmd_name, resp, conn.protocol_version); - responses[resp_idx] = resp; - } - } - } - - // AUTH rate limiting: delay response to slow down brute-force attacks - if auth_delay_ms > 0 { - monoio::time::sleep(std::time::Duration::from_millis(auth_delay_ms)).await; - } - - // Serialize all responses into write_buf, then do ONE write_all syscall. - for response in &responses { - codec.encode_frame(response, &mut write_buf); - } - - // Write all responses in one batch using ownership I/O - if !write_buf.is_empty() { - let data = write_buf.split().freeze(); - let (result, _data): (std::io::Result, bytes::Bytes) = - stream.write_all(data).await; - if result.is_err() { - break; - } - } - - // Update registry with current state after each batch - crate::client_registry::update(client_id, |e| { - e.db = conn.selected_db; - e.last_cmd_at = std::time::Instant::now(); - e.flags = crate::client_registry::ClientFlags { - subscriber: conn.subscription_count > 0, - in_multi: conn.in_multi, - blocked: false, - }; - }); - - // Check if migration was triggered during frame processing. - // All responses for the current batch have been written, so the - // client sees no interruption -- TCP socket stays open. - if let Some(target_shard) = conn.migration_target { - let migrated_state = MigratedConnectionState { - selected_db: conn.selected_db, - authenticated: conn.authenticated, - client_name: conn.client_name.clone(), - protocol_version: conn.protocol_version, - current_user: conn.current_user.clone(), - flags: 0, - read_buf_remainder: read_buf.split(), - client_id, - peer_addr: peer_addr.clone(), - workspace_id: conn.workspace_id, - }; - return ( - MonoioHandlerResult::MigrateConnection { - state: migrated_state, - target_shard, - }, - Some(stream), - ); - } - - if should_quit { - break; - } - - // Check shutdown (polled after each batch -- acceptable for MVP) - if shutdown.is_cancelled() { - break; - } - - // Shrink buffers if they grew too large - if read_buf.capacity() > 65536 { - let remaining = read_buf.split(); - read_buf = BytesMut::with_capacity(8192); - if !remaining.is_empty() { - read_buf.extend_from_slice(&remaining); - } - } - if write_buf.capacity() > 65536 { - write_buf = BytesMut::with_capacity(8192); - } - if tmp_buf.capacity() > 65536 { - tmp_buf = vec![0u8; 8192]; - } - } - - // --- Graceful TCP shutdown: send FIN to client to avoid CLOSE_WAIT --- - // Uses monoio's own shutdown() which properly manages the fd through - // the runtime (unlike raw libc::shutdown which corrupts monoio state). - let _ = stream.shutdown().await; - - // Phase 166: release any leaked cross-store TXN (client disconnected mid-txn). - // Idempotent: TXN.ABORT already takes() active_cross_txn so this is a no-op if abort ran. - // Closes T-161-05 — without this, a disconnect after TXN.BEGIN + SET would leak - // kv_intents and pin the key invisible for all subsequent readers. Mirrors the - // sharded runtime block in handler_sharded.rs so both paths delegate to the same - // shared helper. FIN has already been sent; shard state is still intact. - if let Some(txn) = conn.active_cross_txn.take() { - crate::transaction::abort::abort_cross_store_txn( - &ctx.shard_databases, - ctx.shard_id, - conn.selected_db, - txn, - ); - } - - // --- Disconnect cleanup: propagate unsubscribe to all shards' remote subscriber maps --- - if conn.subscriber_id > 0 { - let removed_channels = { - ctx.pubsub_registry - .write() - .unsubscribe_all(conn.subscriber_id) - }; - let removed_patterns = { - ctx.pubsub_registry - .write() - .punsubscribe_all(conn.subscriber_id) - }; - for ch in removed_channels { - unpropagate_subscription( - &ctx.all_remote_sub_maps, - &ch, - ctx.shard_id, - ctx.num_shards, - false, - ); - } - for pat in removed_patterns { - unpropagate_subscription( - &ctx.all_remote_sub_maps, - &pat, - ctx.shard_id, - ctx.num_shards, - true, - ); - } - // Remove affinity on disconnect (no subscriptions remain) - if let Ok(addr) = peer_addr.parse::() { - ctx.pubsub_affinity.write().remove(&addr.ip()); - } - } - - // NOTE: connection close is recorded by the caller (conn_accept.rs) to - // preserve symmetry with `try_accept_connection`, which owns the - // increment. Decrementing here too produces a double-decrement on the - // AtomicU64 counter — it wraps to u64::MAX on the second subtraction - // and all subsequent `try_accept_connection` comparisons against - // `maxclients` reject new connections. - (MonoioHandlerResult::Done, None) -} diff --git a/src/server/conn/handler_monoio/dispatch.rs b/src/server/conn/handler_monoio/dispatch.rs new file mode 100644 index 00000000..8610b6a2 --- /dev/null +++ b/src/server/conn/handler_monoio/dispatch.rs @@ -0,0 +1,1061 @@ +//! Connection-level command dispatchers for the monoio handler. +//! +//! AUTH gate, HELLO, ACL, CLIENT subcommands, CONFIG, BGSAVE/SAVE/LASTSAVE/BGREWRITEAOF, +//! REPLICAOF/REPLCONF, INFO, READONLY enforcement, CLUSTER, EVAL/EVALSHA/SCRIPT, +//! FUNCTION/FCALL/FCALL_RO, ACL permission gate, cross-shard KEYS/SCAN/DBSIZE, +//! multi-key commands, and blocking commands. +//! +//! Each helper returns `true` if the command was consumed (caller should `continue`). + +use bytes::Bytes; +use ringbuf::traits::Producer; +use std::cell::RefCell; +use std::rc::Rc; +use std::sync::Arc; + +use crate::command::connection as conn_cmd; +use crate::command::metadata; +use crate::protocol::Frame; +use crate::runtime::cancel::CancellationToken; +use crate::runtime::channel; +use crate::server::conn::core::{ConnectionContext, ConnectionState}; +use crate::server::conn::util::extract_bytes; +use crate::shard::dispatch::ShardMessage; +use crate::shard::mesh::ChannelMesh; +use crate::tracking::TrackingState; +use crate::workspace::strip_workspace_prefix_from_response; + +use super::{extract_command, handle_blocking_command_monoio, handle_config, is_multi_key_command}; + +/// Result of the AUTH gate check. +pub(super) enum AuthGateResult { + /// Command consumed (AUTH, HELLO succeeded/failed). Caller should `continue`. + Consumed, + /// QUIT received while not authenticated. Caller should set should_quit and `break`. + Quit, + /// Not an AUTH/HELLO/QUIT command. Caller should push NOAUTH and `continue`. + NotAuth, + /// Already authenticated -- AUTH gate does not apply. + Authenticated, +} + +/// Check the pre-authentication gate. Returns the action the caller should take. +pub(super) fn check_auth_gate( + frame: &Frame, + conn: &mut ConnectionState, + ctx: &ConnectionContext, + peer_addr: &str, + client_id: u64, + responses: &mut Vec, + auth_delay_ms: &mut u64, +) -> AuthGateResult { + if conn.authenticated { + return AuthGateResult::Authenticated; + } + match extract_command(frame) { + Some((cmd, cmd_args)) if cmd.eq_ignore_ascii_case(b"AUTH") => { + let (response, opt_user) = conn_cmd::auth_acl(cmd_args, &ctx.acl_table); + if let Some(uname) = opt_user { + conn.authenticated = true; + conn.current_user = uname; + conn.refresh_acl_cache(&ctx.acl_table); + if let Ok(addr) = peer_addr.parse::() { + crate::auth_ratelimit::record_success(addr.ip()); + } + } else { + if let Ok(addr) = peer_addr.parse::() { + *auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); + } + conn.acl_log.push(crate::acl::AclLogEntry { + reason: "auth".to_string(), + object: "AUTH".to_string(), + username: conn.current_user.clone(), + client_addr: peer_addr.to_string(), + timestamp_ms: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64, + }); + } + responses.push(response); + AuthGateResult::Consumed + } + Some((cmd, cmd_args)) if cmd.eq_ignore_ascii_case(b"HELLO") => { + let (response, new_proto, new_name, opt_user) = conn_cmd::hello_acl( + cmd_args, + conn.protocol_version, + client_id, + &ctx.acl_table, + &mut conn.authenticated, + ); + if !matches!(&response, Frame::Error(_)) { + conn.protocol_version = new_proto; + } + if let Some(name) = new_name { + conn.client_name = Some(name); + } + if let Some(ref uname) = opt_user { + conn.current_user = uname.clone(); + conn.refresh_acl_cache(&ctx.acl_table); + } + // HELLO AUTH rate limiting + if matches!(&response, Frame::Error(_)) { + if let Ok(addr) = peer_addr.parse::() { + *auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); + } + } else if opt_user.is_some() { + if let Ok(addr) = peer_addr.parse::() { + crate::auth_ratelimit::record_success(addr.ip()); + } + } + responses.push(response); + AuthGateResult::Consumed + } + Some((cmd, _)) if cmd.eq_ignore_ascii_case(b"QUIT") => { + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + AuthGateResult::Quit + } + _ => AuthGateResult::NotAuth, + } +} + +/// Handle CLUSTER subcommands. Returns `true` if consumed. +pub(super) fn try_handle_cluster( + cmd: &[u8], + cmd_args: &[Frame], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"CLUSTER") { + return false; + } + if let Some(ref cs) = ctx.cluster_state { + #[allow(clippy::unwrap_used)] // Fallback "127.0.0.1:6379" is a valid literal + let self_addr: std::net::SocketAddr = format!("127.0.0.1:{}", ctx.config_port) + .parse() + .unwrap_or_else(|_| "127.0.0.1:6379".parse().unwrap()); + let resp = crate::cluster::command::handle_cluster_command(cmd_args, cs, self_addr); + responses.push(resp); + } else { + responses.push(Frame::Error(Bytes::from_static( + b"ERR This instance has cluster support disabled", + ))); + } + true +} + +/// Handle EVALSHA command. Returns `true` if consumed. +pub(super) fn try_handle_evalsha( + cmd: &[u8], + cmd_args: &[Frame], + conn: &ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"EVALSHA") { + return false; + } + let response = { + let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + let db_count = ctx.shard_databases.db_count(); + let db = &mut guard; + crate::scripting::handle_evalsha( + &ctx.lua, + &ctx.script_cache, + cmd_args, + db, + ctx.shard_id, + ctx.num_shards, + conn.selected_db, + db_count, + ) + }; + responses.push(response); + true +} + +/// Handle EVAL command. Returns `true` if consumed. +pub(super) fn try_handle_eval( + cmd: &[u8], + cmd_args: &[Frame], + conn: &ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"EVAL") { + return false; + } + let response = { + let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + let db_count = ctx.shard_databases.db_count(); + let db = &mut guard; + crate::scripting::handle_eval( + &ctx.lua, + &ctx.script_cache, + cmd_args, + db, + ctx.shard_id, + ctx.num_shards, + conn.selected_db, + db_count, + ) + }; + responses.push(response); + true +} + +/// Handle SCRIPT subcommands (LOAD, EXISTS, FLUSH). Returns `true` if consumed. +pub(super) fn try_handle_script( + cmd: &[u8], + cmd_args: &[Frame], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"SCRIPT") { + return false; + } + let (response, fanout) = + crate::scripting::handle_script_subcommand(&ctx.script_cache, cmd_args); + if let Some((sha1, script_bytes)) = fanout { + let mut producers = ctx.dispatch_tx.borrow_mut(); + for target in 0..ctx.num_shards { + if target == ctx.shard_id { + continue; + } + let idx = ChannelMesh::target_index(ctx.shard_id, target); + let msg = ShardMessage::ScriptLoad { + sha1: sha1.clone(), + script: script_bytes.clone(), + }; + if producers[idx].try_push(msg).is_ok() { + ctx.spsc_notifiers[target].notify_one(); + } + } + drop(producers); + } + responses.push(response); + true +} + +/// Handle cluster slot routing (pre-dispatch). +/// Returns `true` if the command was redirected (MOVED/ASK/CROSSSLOT) and should be skipped. +pub(super) fn try_handle_cluster_routing( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !crate::cluster::cluster_enabled() { + return false; + } + let Some(ref cs) = ctx.cluster_state else { + return false; + }; + let was_asking = conn.asking; + conn.asking = false; + + let maybe_key = super::extract_primary_key(cmd, cmd_args); + if let Some(key) = maybe_key { + let slot = crate::cluster::slots::slot_for_key(key); + #[allow(clippy::unwrap_used)] + // std RwLock: poison = prior panic = unrecoverable + let route = cs.read().unwrap().route_slot(slot, was_asking); + match route { + crate::cluster::SlotRoute::Local => {} // proceed + other => { + let err_frame = other.into_error_frame(slot); + responses.push(err_frame); + return true; + } + } + + // CROSSSLOT check for multi-key commands + if is_multi_key_command(cmd, cmd_args) { + let first_slot = slot; + let mut cross_slot = false; + for arg in cmd_args.iter().skip(1) { + if let Some(k) = match arg { + Frame::BulkString(b) => Some(b.as_ref()), + _ => None, + } { + if crate::cluster::slots::slot_for_key(k) != first_slot { + cross_slot = true; + break; + } + } + } + if cross_slot { + responses.push(Frame::Error(Bytes::from_static( + b"CROSSSLOT Keys in request don't hash to the same slot", + ))); + return true; + } + } + } + false +} + +/// Handle AUTH command (when already authenticated). Returns `true` if consumed. +pub(super) fn try_handle_auth( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + peer_addr: &str, + auth_delay_ms: &mut u64, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"AUTH") { + return false; + } + let (response, opt_user) = conn_cmd::auth_acl(cmd_args, &ctx.acl_table); + if let Some(uname) = opt_user { + conn.current_user = uname; + conn.refresh_acl_cache(&ctx.acl_table); + if let Ok(addr) = peer_addr.parse::() { + crate::auth_ratelimit::record_success(addr.ip()); + } + } else if let Ok(addr) = peer_addr.parse::() { + *auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); + } + responses.push(response); + true +} + +/// Handle HELLO command (protocol negotiation, ACL-aware). Returns `true` if consumed. +pub(super) fn try_handle_hello( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + client_id: u64, + peer_addr: &str, + auth_delay_ms: &mut u64, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"HELLO") { + return false; + } + let (response, new_proto, new_name, opt_user) = conn_cmd::hello_acl( + cmd_args, + conn.protocol_version, + client_id, + &ctx.acl_table, + &mut conn.authenticated, + ); + if !matches!(&response, Frame::Error(_)) { + conn.protocol_version = new_proto; + } + if let Some(name) = new_name { + conn.client_name = Some(name); + } + if let Some(ref uname) = opt_user { + conn.current_user = uname.clone(); + conn.refresh_acl_cache(&ctx.acl_table); + } + if matches!(&response, Frame::Error(_)) { + if let Ok(addr) = peer_addr.parse::() { + *auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); + } + } else if opt_user.is_some() { + if let Ok(addr) = peer_addr.parse::() { + crate::auth_ratelimit::record_success(addr.ip()); + } + } + responses.push(response); + true +} + +/// Handle ACL command (intercepted at connection level). Returns `true` if consumed. +pub(super) fn try_handle_acl( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + peer_addr: &str, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"ACL") { + return false; + } + let response = crate::command::acl::handle_acl( + cmd_args, + &ctx.acl_table, + &mut conn.acl_log, + &conn.current_user, + peer_addr, + &ctx.runtime_config, + ); + responses.push(response); + true +} + +/// Handle CONFIG GET/SET. Returns `true` if consumed. +pub(super) fn try_handle_config( + cmd: &[u8], + cmd_args: &[Frame], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"CONFIG") { + return false; + } + responses.push(handle_config(cmd_args, &ctx.runtime_config, &ctx.config)); + true +} + +/// Handle REPLICAOF / SLAVEOF. Returns `true` if consumed. +pub(super) fn try_handle_replicaof( + cmd: &[u8], + cmd_args: &[Frame], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"REPLICAOF") && !cmd.eq_ignore_ascii_case(b"SLAVEOF") { + return false; + } + use crate::command::connection::{ReplicaofAction, replicaof}; + let (resp, action) = replicaof(cmd_args); + if let Some(action) = action { + if let Some(ref rs) = ctx.repl_state { + match action { + ReplicaofAction::StartReplication { host, port } => { + if let Ok(mut rs_guard) = rs.write() { + rs_guard.role = crate::replication::state::ReplicationRole::Replica { + host: host.clone(), + port, + state: + crate::replication::handshake::ReplicaHandshakeState::PingPending, + }; + } + let rs_clone = Arc::clone(rs); + let cfg = crate::replication::replica::ReplicaTaskConfig { + master_host: host, + master_port: port, + repl_state: rs_clone, + num_shards: ctx.num_shards, + persistence_dir: None, + listening_port: 0, + }; + monoio::spawn(crate::replication::replica::run_replica_task(cfg)); + } + ReplicaofAction::PromoteToMaster => { + use crate::replication::state::generate_repl_id; + if let Ok(mut rs_guard) = rs.write() { + rs_guard.repl_id2 = rs_guard.repl_id.clone(); + rs_guard.repl_id = generate_repl_id(); + rs_guard.role = crate::replication::state::ReplicationRole::Master; + } + } + ReplicaofAction::NoOp => {} + } + } + } + responses.push(resp); + true +} + +/// Handle REPLCONF command. Returns `true` if consumed. +pub(super) fn try_handle_replconf( + cmd: &[u8], + cmd_args: &[Frame], + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"REPLCONF") { + return false; + } + responses.push(crate::command::connection::replconf(cmd_args)); + true +} + +/// Handle INFO command. Returns `true` if consumed. +pub(super) fn try_handle_info( + cmd: &[u8], + cmd_args: &[Frame], + conn: &ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"INFO") { + return false; + } + let guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); + let response_text = { + let resp_frame = conn_cmd::info_readonly(&guard, cmd_args); + match resp_frame { + Frame::BulkString(b) => String::from_utf8_lossy(&b).to_string(), + _ => String::new(), + } + }; + drop(guard); + let mut response_text = response_text; + if let Some(ref rs) = ctx.repl_state { + if let Ok(rs_guard) = rs.try_read() { + response_text.push_str(&crate::replication::handshake::build_info_replication( + &rs_guard, + )); + } + } + responses.push(Frame::BulkString(Bytes::from(response_text))); + true +} + +/// Handle READONLY enforcement: reject writes on replicas. +/// Returns `true` if the command was blocked. +pub(super) fn try_enforce_readonly( + cmd: &[u8], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if let Some(ref rs) = ctx.repl_state { + if let Ok(rs_guard) = rs.try_read() { + if matches!( + rs_guard.role, + crate::replication::state::ReplicationRole::Replica { .. } + ) { + if metadata::is_write(cmd) { + responses.push(Frame::Error(Bytes::from_static( + b"READONLY You can't write against a read only replica.", + ))); + return true; + } + } + } + } + false +} + +/// Handle CLIENT subcommands that are checked BEFORE the ACL gate: +/// ID, SETNAME, GETNAME, TRACKING. +/// Returns `true` if a subcommand was consumed (caller should `continue`). +/// Returns `false` for admin subcommands (LIST, INFO, KILL, PAUSE, UNPAUSE, NO-EVICT, NO-TOUCH) +/// which must pass through the ACL gate first. +pub(super) fn try_handle_client_early( + cmd: &[u8], + cmd_args: &[Frame], + client_id: u64, + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"CLIENT") { + return false; + } + if let Some(sub) = cmd_args.first() { + if let Some(sub_bytes) = extract_bytes(sub) { + if sub_bytes.eq_ignore_ascii_case(b"ID") { + responses.push(conn_cmd::client_id(client_id)); + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"SETNAME") { + if cmd_args.len() != 2 { + responses.push(Frame::Error(Bytes::from_static( + b"ERR wrong number of arguments for 'CLIENT SETNAME' command", + ))); + } else { + conn.client_name = extract_bytes(&cmd_args[1]); + let name_str = conn + .client_name + .as_ref() + .map(|b| String::from_utf8_lossy(b).to_string()); + crate::client_registry::update(client_id, |e| { + e.name = name_str; + }); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"GETNAME") { + responses.push(match &conn.client_name { + Some(name) => Frame::BulkString(name.clone()), + None => Frame::Null, + }); + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"TRACKING") { + match crate::command::client::parse_tracking_args(cmd_args) { + Ok(config_parsed) => { + if config_parsed.enable { + conn.tracking_state.enabled = true; + conn.tracking_state.bcast = config_parsed.bcast; + conn.tracking_state.noloop = config_parsed.noloop; + conn.tracking_state.optin = config_parsed.optin; + conn.tracking_state.optout = config_parsed.optout; + + if conn.tracking_rx.is_none() { + let (tx, rx) = channel::mpsc_bounded::(256); + conn.tracking_state.invalidation_tx = Some(tx.clone()); + conn.tracking_rx = Some(rx); + + let mut table = ctx.tracking_table.borrow_mut(); + table.register_client(client_id, tx); + if let Some(target) = config_parsed.redirect { + table.set_redirect(client_id, target); + } + for prefix in &config_parsed.prefixes { + table.register_prefix( + client_id, + prefix.clone(), + config_parsed.noloop, + ); + } + } + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } else { + conn.tracking_state = TrackingState::default(); + ctx.tracking_table.borrow_mut().untrack_all(client_id); + conn.tracking_rx = None; + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + return true; + } + Err(err_frame) => { + responses.push(err_frame); + return true; + } + } + } + // Admin CLIENT subcommands (LIST, INFO, KILL, PAUSE, UNPAUSE, + // NO-EVICT, NO-TOUCH) fall through to the ACL gate below. + } + } + // Fall through -- admin subcommands handled after ACL check. + false +} + +/// Handle CLIENT admin subcommands (LIST, INFO, KILL, PAUSE, UNPAUSE, NO-EVICT, NO-TOUCH). +/// Placed AFTER ACL check so restricted users cannot access admin ops. +/// Returns `true` if consumed. +pub(super) fn try_handle_client_admin( + cmd: &[u8], + cmd_args: &[Frame], + client_id: u64, + conn: &ConnectionState, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"CLIENT") { + return false; + } + if let Some(sub) = cmd_args.first() { + if let Some(sub_bytes) = extract_bytes(sub) { + if sub_bytes.eq_ignore_ascii_case(b"LIST") { + crate::client_registry::update(client_id, |e| { + e.db = conn.selected_db; + e.last_cmd_at = std::time::Instant::now(); + e.flags = crate::client_registry::ClientFlags { + subscriber: conn.subscription_count > 0, + in_multi: conn.in_multi, + blocked: false, + }; + }); + let list = crate::client_registry::client_list(); + responses.push(Frame::BulkString(Bytes::from(list))); + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"INFO") { + crate::client_registry::update(client_id, |e| { + e.db = conn.selected_db; + e.last_cmd_at = std::time::Instant::now(); + }); + let info = crate::client_registry::client_info(client_id).unwrap_or_default(); + responses.push(Frame::BulkString(Bytes::from(info))); + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"KILL") { + let raw_args: Vec<&[u8]> = cmd_args[1..] + .iter() + .filter_map(|f| match f { + Frame::BulkString(b) => Some(b.as_ref()), + Frame::SimpleString(b) => Some(b.as_ref()), + _ => None, + }) + .collect(); + match crate::client_registry::parse_kill_args(&raw_args) { + Some(filter) => { + let count = crate::client_registry::kill_clients(&filter); + responses.push(Frame::Integer(count as i64)); + } + None => { + responses.push(Frame::Error(Bytes::from_static( + b"ERR syntax error. Usage: CLIENT KILL [ID id] [ADDR addr] [USER user]", + ))); + } + } + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"PAUSE") { + if cmd_args.len() < 2 { + responses.push(Frame::Error(Bytes::from_static( + b"ERR wrong number of arguments for 'CLIENT PAUSE' command", + ))); + } else { + let timeout_bytes = match &cmd_args[1] { + Frame::BulkString(b) => Some(b.as_ref()), + Frame::SimpleString(b) => Some(b.as_ref()), + _ => None, + }; + match timeout_bytes + .and_then(|b| std::str::from_utf8(b).ok()) + .and_then(|s| s.parse::().ok()) + { + Some(ms) => { + let mode = if cmd_args.len() > 2 { + match &cmd_args[2] { + Frame::BulkString(b) | Frame::SimpleString(b) + if b.eq_ignore_ascii_case(b"WRITE") => + { + crate::client_pause::PauseMode::Write + } + _ => crate::client_pause::PauseMode::All, + } + } else { + crate::client_pause::PauseMode::All + }; + crate::client_pause::pause(ms, mode); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + None => { + responses.push(Frame::Error(Bytes::from_static( + b"ERR timeout is not a valid integer or out of range", + ))); + } + } + } + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"UNPAUSE") { + crate::client_pause::unpause(); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"NO-EVICT") + || sub_bytes.eq_ignore_ascii_case(b"NO-TOUCH") + { + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + return true; + } + // Unknown CLIENT subcommand + responses.push(Frame::Error(Bytes::from(format!( + "ERR unknown subcommand '{}'", + String::from_utf8_lossy(&sub_bytes) + )))); + return true; + } + } + responses.push(Frame::Error(Bytes::from_static( + b"ERR wrong number of arguments for 'client' command", + ))); + true +} + +/// Handle persistence commands (BGSAVE, SAVE, LASTSAVE, BGREWRITEAOF). +/// Returns `true` if consumed. +pub(super) fn try_handle_persistence( + cmd: &[u8], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if cmd.eq_ignore_ascii_case(b"BGSAVE") { + responses.push(crate::command::persistence::bgsave_start_sharded( + &ctx.snapshot_trigger_tx, + ctx.num_shards, + )); + return true; + } + if cmd.eq_ignore_ascii_case(b"SAVE") { + responses.push(Frame::Error(Bytes::from_static( + b"ERR SAVE not supported in sharded mode, use BGSAVE", + ))); + return true; + } + if cmd.eq_ignore_ascii_case(b"LASTSAVE") { + responses.push(crate::command::persistence::handle_lastsave()); + return true; + } + if cmd.eq_ignore_ascii_case(b"BGREWRITEAOF") { + if let Some(ref tx) = ctx.aof_tx { + responses.push(crate::command::persistence::bgrewriteaof_start_sharded( + tx, + ctx.shard_databases.clone(), + )); + } else { + responses.push(Frame::Error(Bytes::from_static(b"ERR AOF is not enabled"))); + } + return true; + } + false +} + +/// Handle ACL permission check (NOPERM gate). +/// Returns `true` if the command was denied (caller should `continue`). +pub(super) fn try_enforce_acl( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + peer_addr: &str, + responses: &mut Vec, +) -> bool { + if conn.acl_skip_allowed() { + return false; + } + #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable + let acl_guard = ctx.acl_table.read().unwrap(); + if let Some(deny_reason) = acl_guard.check_command_permission(&conn.current_user, cmd, cmd_args) + { + drop(acl_guard); + conn.acl_log.push(crate::acl::AclLogEntry { + reason: "command".to_string(), + object: String::from_utf8_lossy(cmd).to_ascii_lowercase(), + username: conn.current_user.clone(), + client_addr: peer_addr.to_string(), + timestamp_ms: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64, + }); + responses.push(Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason)))); + return true; + } + + // === ACL key pattern check (same lock guard) === + let is_write_for_acl = metadata::is_write(cmd); + if let Some(deny_reason) = + acl_guard.check_key_permission(&conn.current_user, cmd, cmd_args, is_write_for_acl) + { + drop(acl_guard); + conn.acl_log.push(crate::acl::AclLogEntry { + reason: "command".to_string(), + object: String::from_utf8_lossy(cmd).to_ascii_lowercase(), + username: conn.current_user.clone(), + client_addr: peer_addr.to_string(), + timestamp_ms: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64, + }); + responses.push(Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason)))); + return true; + } + false +} + +/// Handle FUNCTION/FCALL/FCALL_RO commands. Returns `true` if consumed. +/// Placed AFTER ACL check. Skipped when conn.in_multi (fall through to MULTI queue). +pub(super) fn try_handle_functions( + cmd: &[u8], + cmd_args: &[Frame], + conn: &ConnectionState, + ctx: &ConnectionContext, + func_registry: &Rc>, + responses: &mut Vec, +) -> bool { + if conn.in_multi { + return false; + } + if cmd.eq_ignore_ascii_case(b"FUNCTION") { + let response = + crate::command::functions::handle_function(&mut func_registry.borrow_mut(), cmd_args); + responses.push(response); + return true; + } + if cmd.eq_ignore_ascii_case(b"FCALL") { + let response = { + let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + let db_count = ctx.shard_databases.db_count(); + crate::command::functions::handle_fcall( + &func_registry.borrow(), + cmd_args, + &mut guard, + ctx.shard_id, + ctx.num_shards, + conn.selected_db, + db_count, + ) + }; + responses.push(response); + return true; + } + if cmd.eq_ignore_ascii_case(b"FCALL_RO") { + let response = { + let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + let db_count = ctx.shard_databases.db_count(); + crate::command::functions::handle_fcall_ro( + &func_registry.borrow(), + cmd_args, + &mut guard, + ctx.shard_id, + ctx.num_shards, + conn.selected_db, + db_count, + ) + }; + responses.push(response); + return true; + } + false +} + +/// Handle cross-shard aggregation commands: KEYS, SCAN, DBSIZE, and multi-key commands. +/// Returns `true` if consumed. +pub(super) async fn try_handle_cross_shard_commands( + cmd: &[u8], + cmd_args: &[Frame], + conn: &ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if ctx.num_shards <= 1 { + return false; + } + if cmd.eq_ignore_ascii_case(b"KEYS") { + let mut response = crate::shard::coordinator::coordinate_keys( + cmd_args, + ctx.shard_id, + ctx.num_shards, + conn.selected_db, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + &ctx.cached_clock, + &(), // monoio: coordinator uses oneshot, not response_pool + ) + .await; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + if cmd.eq_ignore_ascii_case(b"SCAN") { + let mut response = crate::shard::coordinator::coordinate_scan( + cmd_args, + ctx.shard_id, + ctx.num_shards, + conn.selected_db, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + &ctx.cached_clock, + &(), // monoio: coordinator uses oneshot, not response_pool + ) + .await; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + if cmd.eq_ignore_ascii_case(b"DBSIZE") { + let response = crate::shard::coordinator::coordinate_dbsize( + ctx.shard_id, + ctx.num_shards, + conn.selected_db, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + &(), // monoio: coordinator uses oneshot, not response_pool + ) + .await; + responses.push(response); + return true; + } + + // --- Multi-key commands: MGET, MSET, DEL, UNLINK, EXISTS --- + if is_multi_key_command(cmd, cmd_args) { + let response = crate::shard::coordinator::coordinate_multi_key( + cmd, + cmd_args, + ctx.shard_id, + ctx.num_shards, + conn.selected_db, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + &ctx.cached_clock, + &(), // monoio: coordinator uses oneshot, not response_pool + ) + .await; + responses.push(response); + return true; + } + false +} + +/// Result of blocking command handling. +pub(super) enum BlockingResult { + /// Not a blocking command. + NotBlocking, + /// In MULTI: queued as non-blocking variant. Caller should `continue`. + Queued, + /// Blocking command handled. Caller must `break` (ends pipeline). + Handled, + /// Write error during flush. Caller should return Done. + WriteError, +} + +/// Handle blocking commands (BLPOP, BRPOP, BLMOVE, etc.). +pub(super) async fn try_handle_blocking( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, + codec: &mut crate::server::codec::RespCodec, + write_buf: &mut bytes::BytesMut, + stream: &mut S, + shutdown: &CancellationToken, +) -> BlockingResult { + if !cmd.eq_ignore_ascii_case(b"BLPOP") + && !cmd.eq_ignore_ascii_case(b"BRPOP") + && !cmd.eq_ignore_ascii_case(b"BLMOVE") + && !cmd.eq_ignore_ascii_case(b"BZPOPMIN") + && !cmd.eq_ignore_ascii_case(b"BZPOPMAX") + && !cmd.eq_ignore_ascii_case(b"BLMPOP") + && !cmd.eq_ignore_ascii_case(b"BRPOPLPUSH") + && !cmd.eq_ignore_ascii_case(b"BZMPOP") + { + return BlockingResult::NotBlocking; + } + + // Inside MULTI: queue as non-blocking variant + if conn.in_multi { + let nb_frame = super::convert_blocking_to_nonblocking(cmd, cmd_args); + conn.command_queue.push(nb_frame); + responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); + return BlockingResult::Queued; + } + + // Flush accumulated responses before blocking + for resp in &*responses { + codec.encode_frame(resp, write_buf); + } + if !write_buf.is_empty() { + use monoio::io::AsyncWriteRentExt; + let data = write_buf.split().freeze(); + let (result, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if result.is_err() { + return BlockingResult::WriteError; + } + } + + let blocking_response = handle_blocking_command_monoio( + cmd, + cmd_args, + conn.selected_db, + &ctx.shard_databases, + &ctx.blocking_registry, + ctx.shard_id, + ctx.num_shards, + &ctx.dispatch_tx, + shutdown, + &ctx.spsc_notifiers, + ) + .await; + + // Encode blocking response directly + codec.encode_frame(&blocking_response, write_buf); + responses.clear(); + BlockingResult::Handled +} diff --git a/src/server/conn/handler_monoio/ft.rs b/src/server/conn/handler_monoio/ft.rs new file mode 100644 index 00000000..f2927dd5 --- /dev/null +++ b/src/server/conn/handler_monoio/ft.rs @@ -0,0 +1,820 @@ +//! FT.* vector/text search command handlers. +//! +//! Handles FT.CREATE, FT.DROPINDEX, FT.SEARCH (vector + text), FT.AGGREGATE, +//! FT.INFO, FT.COMPACT, FT._LIST, FT.CACHESEARCH, FT.CONFIG, FT.RECOMMEND, +//! FT.NAVIGATE, FT.EXPAND. Multi-shard scatter-gather and single-shard fast paths. + +use bytes::Bytes; + +use crate::protocol::Frame; +use crate::server::conn::core::{ConnectionContext, ConnectionState}; +use crate::server::conn::shared::resolve_ft_search_as_of_lsn; +use crate::workspace::strip_workspace_prefix_from_response; + +/// Handle FT.* commands. Returns `true` if the command was consumed. +/// +/// Caller should `continue` the frame loop when this returns `true`. +pub(super) async fn try_handle_ft_command( + cmd: &[u8], + cmd_args: &[Frame], + frame: &Frame, + conn: &ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if cmd.len() <= 3 || !cmd[..3].eq_ignore_ascii_case(b"FT.") { + return false; + } + + if ctx.num_shards > 1 { + // Multi-shard: dispatch via SPSC + #[cfg(feature = "text-index")] + if cmd.eq_ignore_ascii_case(b"FT.AGGREGATE") { + // ── FT.AGGREGATE: multi-shard scatter-gather (Phase 152 Plan 03). ── + // Mirrors handler_sharded.rs:1458. scatter_text_aggregate acquires + // its own per-shard guards inside the single-shard block internally, + // so we never hold a MutexGuard across the .await below. + let parsed = + match crate::command::vector_search::ft_aggregate::parse_aggregate_args(cmd_args) { + Ok(p) => p, + Err(err_frame) => { + responses.push(err_frame); + return true; + } + }; + let response = crate::shard::scatter_aggregate::scatter_text_aggregate( + parsed.index_name, + parsed.query, + parsed.pipeline, + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await; + responses.push(response); + return true; + } + if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { + // Check if this is a text query BEFORE trying parse_ft_search_args + // (which would return an error for non-KNN queries). + let query_bytes = cmd_args + .get(1) + .and_then(|f| crate::command::vector_search::extract_bulk(f)); + let is_text = query_bytes + .as_ref() + .map_or(false, |q| crate::command::vector_search::is_text_query(q)); + + // ── HYBRID multi-shard path (Phase 152 Plan 05, D-13) ────────── + #[cfg(feature = "text-index")] + { + match crate::command::vector_search::hybrid::parse_hybrid_modifier(cmd_args) { + Ok(Some(partial)) => { + let index_name = match cmd_args + .first() + .and_then(|f| crate::command::vector_search::extract_bulk(f)) + { + Some(b) => b, + None => { + responses.push(Frame::Error(Bytes::from_static( + b"ERR invalid index name", + ))); + return true; + } + }; + let text_query = match query_bytes.clone() { + Some(q) => q, + None => { + responses + .push(Frame::Error(Bytes::from_static(b"ERR invalid query"))); + return true; + } + }; + let (limit_offset, limit_count) = + crate::command::vector_search::parse_limit_clause(cmd_args); + let top_k = if limit_count == usize::MAX { + limit_offset.saturating_add(10).max(1) + } else { + limit_offset.saturating_add(limit_count).max(1) + }; + let hq = crate::command::vector_search::hybrid::HybridQuery { + index_name, + text_query, + dense_field: partial.dense_field, + dense_blob: partial.dense_blob, + sparse: partial.sparse, + weights: partial.weights, + k_per_stream: partial.k_per_stream, + top_k, + offset: limit_offset, + count: limit_count, + }; + // Phase 171 HYB-02 / SCAT-02: resolve AS_OF / TXN LSN + // ONCE on the coordinator and forward to the scatter + // helper so responders honor temporal snapshots. + let as_of_lsn = match resolve_ft_search_as_of_lsn( + cmd_args, + Some(&ctx.shard_databases), + ctx.shard_id, + conn.active_cross_txn.as_ref(), + ) { + Ok(lsn) => lsn, + Err(err_frame) => { + responses.push(err_frame); + return true; + } + }; + let response = crate::shard::scatter_hybrid::scatter_hybrid_search( + hq, + as_of_lsn, + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await; + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + Ok(None) => { /* fall through */ } + Err(err_frame) => { + responses.push(err_frame); + return true; + } + } + } + + if is_text { + // ── Text FT.SEARCH: two-phase DFS scatter-gather ────────────────── + let index_name = match cmd_args + .first() + .and_then(|f| crate::command::vector_search::extract_bulk(f)) + { + Some(b) => b, + None => { + responses.push(Frame::Error(Bytes::from_static(b"ERR invalid index name"))); + return true; + } + }; + // `is_text` was computed as `query_bytes.as_ref().map_or(false, ...)` (line ~65), + // so reaching `if is_text { ... }` guarantees `query_bytes.is_some()`. + #[allow(clippy::unwrap_used)] + let query_str = query_bytes.unwrap(); + + // B-01 SITE 2 FIX (Plan 152-06): FieldFilter short-circuit BEFORE + // the analyzer-first parse_result block. TAG queries (and Plan 07 + // NumericRange) route through the InvertedSearch fan-out — no + // analyzer touched, no field_idx resolution (the filter carries + // its own field name; search_tag resolves against tag_fields). + #[cfg(feature = "text-index")] + { + match crate::command::vector_search::pre_parse_field_filter(query_str.as_ref()) + { + Ok(Some(clause)) => { + if let Some(filter) = clause.filter { + let (offset, count) = + crate::command::vector_search::parse_limit_clause(cmd_args); + let top_k = if count == usize::MAX { + 10000 + } else { + offset.saturating_add(count) + } + .max(1); + let response = + crate::shard::coordinator::scatter_text_search_filter( + index_name, + filter, + top_k, + offset, + count, + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await; + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + } + Ok(None) => { /* fall through to BM25 path */ } + Err(e) => { + responses.push(Frame::Error(Bytes::from(e.to_owned()))); + return true; + } + } + } + + // Parse query and resolve field_idx inside a block scope so the + // MutexGuard from text_store() is dropped BEFORE .await. + // We use the TextIndex's own field_analyzers (same pipeline used at index time). + type ParseResult = + Result<(Vec, Option), String>; + let parse_result: ParseResult = { + let ts = ctx.shard_databases.text_store(ctx.shard_id); + match ts.get_index(&index_name) { + None => Err("ERR no such index".to_owned()), + Some(text_index) => { + match text_index.field_analyzers.first() { + None => Err("ERR index has no TEXT fields".to_owned()), + Some(analyzer) => { + // analyzer borrows text_index which borrows ts — all in this block. + let parsed = crate::command::vector_search::parse_text_query( + &query_str, analyzer, + ); + match parsed { + Err(e) => Err(e.to_owned()), + Ok(clause) => { + let field_idx = match &clause.field_name { + None => Ok(None), + Some(field_name) => match text_index + .text_fields + .iter() + .position(|f| { + f.field_name.as_ref().eq_ignore_ascii_case( + field_name.as_ref(), + ) + }) { + Some(idx) => Ok(Some(idx)), + None => Err(format!( + "ERR unknown field '{}'", + String::from_utf8_lossy(field_name) + )), + }, + }; + field_idx.map(|idx| (clause.terms, idx)) + } + } + } + } + } + } + }; // MutexGuard dropped here + + let (query_terms, field_idx) = match parse_result { + Ok(t) => t, + Err(e) => { + responses.push(Frame::Error(Bytes::from(e))); + return true; + } + }; + + let (offset, count) = crate::command::vector_search::parse_limit_clause(cmd_args); + let top_k = if count == usize::MAX { + 10000 + } else { + offset.saturating_add(count) + } + .max(1); + + // Parse optional HIGHLIGHT/SUMMARIZE clauses from args. + let highlight_opts = + crate::command::vector_search::parse_highlight_clause(cmd_args); + let summarize_opts = + crate::command::vector_search::parse_summarize_clause(cmd_args); + + let response = crate::shard::coordinator::scatter_text_search( + index_name, + query_terms, + field_idx, + top_k, + offset, + count, + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + highlight_opts, + summarize_opts, + ) + .await; + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + + // ── Vector FT.SEARCH (KNN / SPARSE): existing path ──────────────── + // Phase 171 SCAT-01: resolve AS_OF / TXN snapshot LSN ONCE on + // the coordinator and forward through the scatter helper so + // every responder honors the same temporal snapshot. + let response = match crate::command::vector_search::parse_ft_search_args(cmd_args) { + Ok((index_name, query_blob, k, filter, _offset, _count)) => { + if filter.is_some() { + Frame::Error(Bytes::from_static( + b"ERR FILTER not supported in multi-shard mode yet", + )) + } else { + match resolve_ft_search_as_of_lsn( + cmd_args, + Some(&ctx.shard_databases), + ctx.shard_id, + conn.active_cross_txn.as_ref(), + ) { + Err(err_frame) => err_frame, + Ok(as_of_lsn) => { + crate::shard::coordinator::scatter_vector_search_remote( + index_name, + query_blob, + k, + as_of_lsn, + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await + } + } + } + } + Err(err_frame) => err_frame, + }; + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + if cmd.eq_ignore_ascii_case(b"FT.CREATE") || cmd.eq_ignore_ascii_case(b"FT.DROPINDEX") { + // Broadcast to ALL shards so every shard has the index + let response = crate::shard::coordinator::broadcast_vector_command( + std::sync::Arc::new(frame.clone()), + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await; + responses.push(response); + return true; + } + if cmd.eq_ignore_ascii_case(b"FT.INFO") { + let response = { + let vs = ctx.shard_databases.vector_store(ctx.shard_id); + let ts = ctx.shard_databases.text_store(ctx.shard_id); + crate::command::vector_search::ft_info(&vs, &ts, cmd_args) + }; + responses.push(response); + return true; + } + if cmd.eq_ignore_ascii_case(b"FT._LIST") { + let response = { + let vs = ctx.shard_databases.vector_store(ctx.shard_id); + crate::command::vector_search::ft_list(&vs) + }; + responses.push(response); + return true; + } + if cmd.eq_ignore_ascii_case(b"FT.COMPACT") { + let response = { + let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); + let mut ts = ctx.shard_databases.text_store(ctx.shard_id); + crate::command::vector_search::ft_compact(&mut vs, &mut ts, cmd_args) + }; + responses.push(response); + return true; + } + if cmd.eq_ignore_ascii_case(b"FT.CACHESEARCH") { + let response = { + let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); + crate::command::vector_search::cache_search::ft_cachesearch(&mut vs, cmd_args) + }; + responses.push(response); + return true; + } + if cmd.eq_ignore_ascii_case(b"FT.CONFIG") { + let response = { + let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); + let mut ts = ctx.shard_databases.text_store(ctx.shard_id); + crate::command::vector_search::ft_config(&mut vs, &mut ts, cmd_args) + }; + responses.push(response); + return true; + } + // FT.RECOMMEND, FT.NAVIGATE, FT.EXPAND need db/graph — dispatch locally + if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") + || cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") + || cmd.eq_ignore_ascii_case(b"FT.EXPAND") + { + let response = { + let sdb = &ctx.shard_databases; + let mut vs = sdb.vector_store(ctx.shard_id); + if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") { + let mut db_guard = sdb.write_db(ctx.shard_id, 0); + crate::command::vector_search::recommend::ft_recommend( + &mut vs, + cmd_args, + Some(&mut *db_guard), + ) + } else if cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") { + #[cfg(feature = "graph")] + { + let graph_guard = sdb.graph_store_read(ctx.shard_id); + crate::command::vector_search::navigate::ft_navigate( + &mut vs, + Some(&graph_guard), + cmd_args, + None, + ) + } + #[cfg(not(feature = "graph"))] + { + Frame::Error(Bytes::from_static( + b"ERR FT.NAVIGATE requires graph feature", + )) + } + } else { + // FT.EXPAND + #[cfg(feature = "graph")] + { + let graph_guard = sdb.graph_store_read(ctx.shard_id); + crate::command::vector_search::ft_expand(&graph_guard, cmd_args) + } + #[cfg(not(feature = "graph"))] + { + Frame::Error(Bytes::from_static(b"ERR FT.EXPAND requires graph feature")) + } + } + }; + responses.push(response); + return true; + } + responses.push(Frame::Error(Bytes::from_static(b"ERR unknown FT command"))); + return true; + } else { + // Single-shard: no SPSC channels needed. + // Dispatch directly to shard's VectorStore via shared access. + // + // ── 154-01 single-shard FT.AGGREGATE fast path ──────────────── + // scatter_text_aggregate internally fast-paths num_shards == 1 + // to execute_local_full with locally-acquired guards dropped + // before any .await. Calling the scatter entry here (instead + // of execute_local_full directly) keeps the dispatch body + // byte-symmetric with the multi-shard site above. + #[cfg(feature = "text-index")] + if cmd.eq_ignore_ascii_case(b"FT.AGGREGATE") { + let parsed = + match crate::command::vector_search::ft_aggregate::parse_aggregate_args(cmd_args) { + Ok(p) => p, + Err(err_frame) => { + responses.push(err_frame); + return true; + } + }; + let response = crate::shard::scatter_aggregate::scatter_text_aggregate( + parsed.index_name, + parsed.query, + parsed.pipeline, + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await; + responses.push(response); + return true; + } + // + // ── 151-03 single-shard text FT.SEARCH fast path ────────────── + // Bare text queries (exact / fuzzy / prefix / field-targeted) + // are not understood by `ft_search()` (which only parses + // KNN / SPARSE / HYBRID) — they would otherwise return + // `ERR invalid KNN query syntax`. Route them directly to + // `execute_text_search_local` here, the same function the + // multi-shard path uses once its per-shard scatter has + // aggregated IDFs. We skip HYBRID (existing ft_search + // handles it) and KNN/SPARSE (is_text_query returns false). + #[cfg(feature = "text-index")] + if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { + if let Some(Frame::BulkString(query_bytes)) = cmd_args.get(1) { + match crate::command::vector_search::parse_hybrid_modifier(cmd_args) { + Ok(Some(_)) => { + // HYBRID present — defer to existing ft_search() below. + } + Err(frame_err) => { + responses.push(frame_err); + return true; + } + Ok(None) => { + if crate::command::vector_search::is_text_query(query_bytes.as_ref()) { + // Step 1: index_name from cmd_args[0]. + let index_name = match cmd_args.first() { + Some(Frame::BulkString(b)) => b.clone(), + _ => { + responses.push(Frame::Error(Bytes::from_static( + b"ERR wrong number of arguments for FT.SEARCH", + ))); + return true; + } + }; + // B-01 SITE 2 FIX (single-shard 151-03 fast path, Plan 152-06): + // FieldFilter short-circuit BEFORE the analyzer lookup and + // BEFORE the text_fields.is_empty() bail. + #[cfg(feature = "text-index")] + match crate::command::vector_search::pre_parse_field_filter( + query_bytes.as_ref(), + ) { + Ok(Some(clause)) => { + if clause.filter.is_some() { + let (offset, count) = + crate::command::vector_search::parse_limit_clause( + cmd_args, + ); + let top_k = if count == usize::MAX { + 10000 + } else { + offset.saturating_add(count) + } + .max(1); + let ts_guard = ctx.shard_databases.text_store(ctx.shard_id); + let response = match ts_guard.get_index(&index_name) { + None => Frame::Error(Bytes::from_static( + b"ERR no such index", + )), + Some(text_index) => { + let results = crate::command::vector_search::ft_text_search::execute_query_on_index( + text_index, &clause, None, None, top_k, + ); + crate::command::vector_search::ft_text_search::build_text_response( + &results, offset, count, + ) + } + }; + drop(ts_guard); + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response( + ws_id, + cmd, + &mut response, + ); + } + responses.push(response); + return true; + } + } + Ok(None) => { /* fall through */ } + Err(e) => { + responses + .push(Frame::Error(Bytes::copy_from_slice(e.as_bytes()))); + return true; + } + } + // Step 2: acquire ts guard (single-shard monoio). + let ts_guard = ctx.shard_databases.text_store(ctx.shard_id); + // Step 3: resolve index. + let text_index = match ts_guard.get_index(&index_name) { + Some(idx) => idx, + None => { + drop(ts_guard); + responses.push(Frame::Error(Bytes::from_static( + b"ERR no such index", + ))); + return true; + } + }; + // Step 4: ensure at least one TEXT field. + if text_index.text_fields.is_empty() { + drop(ts_guard); + responses.push(Frame::Error(Bytes::from_static( + b"ERR index has no TEXT fields", + ))); + return true; + } + // Step 5: parse the query via the index's first analyzer. + let analyzer = match text_index.field_analyzers.first() { + Some(a) => a, + None => { + drop(ts_guard); + responses.push(Frame::Error(Bytes::from_static( + b"ERR index has no TEXT fields", + ))); + return true; + } + }; + let clause = match crate::command::vector_search::parse_text_query( + query_bytes.as_ref(), + analyzer, + ) { + Ok(c) => c, + Err(msg) => { + drop(ts_guard); + responses + .push(Frame::Error(Bytes::copy_from_slice(msg.as_bytes()))); + return true; + } + }; + // Step 5b: resolve field_idx. + let field_idx = match &clause.field_name { + None => None, + Some(field_name) => { + match text_index.text_fields.iter().position(|f| { + f.field_name + .as_ref() + .eq_ignore_ascii_case(field_name.as_ref()) + }) { + Some(idx) => Some(idx), + None => { + let bad_name = field_name.clone(); + drop(ts_guard); + responses.push(Frame::Error(Bytes::from(format!( + "ERR unknown field '{}'", + String::from_utf8_lossy(&bad_name) + )))); + return true; + } + } + } + }; + // Step 6: extract query_terms. + let query_terms = clause.terms; + // Step 7: LIMIT parsing + top_k cap (T-151-03-02). + let (offset, count) = + crate::command::vector_search::parse_limit_clause(cmd_args); + let top_k = if count == usize::MAX { + 10000 + } else { + offset.saturating_add(count) + } + .max(1); + // Step 8: HIGHLIGHT / SUMMARIZE options. + let highlight_opts = + crate::command::vector_search::parse_highlight_clause(cmd_args); + let summarize_opts = + crate::command::vector_search::parse_summarize_clause(cmd_args); + // Step 9: acquire DB read guard iff post-processing is needed. + let db_guard_opt = + if highlight_opts.is_some() || summarize_opts.is_some() { + Some(ctx.shard_databases.read_db(ctx.shard_id, 0)) + } else { + None + }; + // Step 10: execute + optional post-processing. + let mut response = + crate::command::vector_search::execute_text_search_local( + &ts_guard, + &index_name, + field_idx, + &query_terms, + top_k, + offset, + count, + ); + if let Some(ref db_guard) = db_guard_opt { + let term_strings: Vec = + query_terms.iter().map(|qt| qt.text.clone()).collect(); + crate::command::vector_search::apply_post_processing( + &mut response, + &term_strings, + text_index, + db_guard, + highlight_opts.as_ref(), + summarize_opts.as_ref(), + ); + } + // Explicit drop order: inner db_guard first, outer ts_guard last. + drop(db_guard_opt); + drop(ts_guard); + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + } + } + } + } + let response = { + let shard_databases_ref = &ctx.shard_databases; + let mut vs = shard_databases_ref.vector_store(ctx.shard_id); + let mut ts = shard_databases_ref.text_store(ctx.shard_id); + if cmd.eq_ignore_ascii_case(b"FT.CREATE") { + crate::command::vector_search::ft_create(&mut vs, &mut ts, cmd_args) + } else if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { + // Resolve AS_OF temporal clause + TXN snapshot precedence (TEMP-04, ACID-09). + let as_of_lsn = match resolve_ft_search_as_of_lsn( + cmd_args, + Some(shard_databases_ref), + ctx.shard_id, + conn.active_cross_txn.as_ref(), + ) { + Ok(lsn) => lsn, + Err(err_frame) => { + responses.push(err_frame); + return true; + } + }; + // Detect SESSION keyword to provide database access for sorted set tracking + let has_session = cmd_args.iter().any(|a| { + if let Frame::BulkString(b) = a { + b.eq_ignore_ascii_case(b"SESSION") + } else { + false + } + }); + if has_session { + let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); + crate::command::vector_search::ft_search( + &mut vs, + cmd_args, + Some(&mut *db_guard), + Some(&*ts), + as_of_lsn, + ) + } else { + crate::command::vector_search::ft_search( + &mut vs, + cmd_args, + None, + Some(&*ts), + as_of_lsn, + ) + } + } else if cmd.eq_ignore_ascii_case(b"FT.DROPINDEX") { + let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); + crate::command::vector_search::ft_dropindex( + &mut vs, + &mut ts, + Some(&mut *db_guard), + cmd_args, + ) + } else if cmd.eq_ignore_ascii_case(b"FT.INFO") { + crate::command::vector_search::ft_info(&vs, &ts, cmd_args) + } else if cmd.eq_ignore_ascii_case(b"FT._LIST") { + crate::command::vector_search::ft_list(&vs) + } else if cmd.eq_ignore_ascii_case(b"FT.COMPACT") { + crate::command::vector_search::ft_compact(&mut vs, &mut ts, cmd_args) + } else if cmd.eq_ignore_ascii_case(b"FT.CACHESEARCH") { + crate::command::vector_search::cache_search::ft_cachesearch(&mut vs, cmd_args) + } else if cmd.eq_ignore_ascii_case(b"FT.CONFIG") { + crate::command::vector_search::ft_config(&mut vs, &mut ts, cmd_args) + } else if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") { + let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); + crate::command::vector_search::recommend::ft_recommend( + &mut vs, + cmd_args, + Some(&mut *db_guard), + ) + } else if cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") { + #[cfg(feature = "graph")] + { + let graph_guard = shard_databases_ref.graph_store_read(ctx.shard_id); + crate::command::vector_search::navigate::ft_navigate( + &mut vs, + Some(&graph_guard), + cmd_args, + None, + ) + } + #[cfg(not(feature = "graph"))] + { + Frame::Error(Bytes::from_static( + b"ERR FT.NAVIGATE requires graph feature", + )) + } + } else if cmd.eq_ignore_ascii_case(b"FT.EXPAND") { + #[cfg(feature = "graph")] + { + let graph_guard = shard_databases_ref.graph_store_read(ctx.shard_id); + crate::command::vector_search::ft_expand(&graph_guard, cmd_args) + } + #[cfg(not(feature = "graph"))] + { + Frame::Error(Bytes::from_static(b"ERR FT.EXPAND requires graph feature")) + } + } else { + Frame::Error(Bytes::from_static(b"ERR unknown FT.* command")) + } + }; + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + // Unreachable: all paths above return true + #[allow(unreachable_code)] + true +} diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs new file mode 100644 index 00000000..86849a78 --- /dev/null +++ b/src/server/conn/handler_monoio/mod.rs @@ -0,0 +1,1491 @@ +// Note: some imports/variables may be conditionally used across feature flags +//! Monoio connection handler using ownership-based I/O (AsyncReadRent/AsyncWriteRent). +//! +//! Extracted from `server/connection.rs` (Plan 48-02). + +mod dispatch; +mod ft; +mod pubsub; +mod read; +mod txn; +mod write; + +use crate::runtime::cancel::CancellationToken; +use crate::runtime::channel; +use bytes::{Bytes, BytesMut}; +use ringbuf::traits::Producer; +use std::cell::RefCell; +use std::collections::HashMap; +use std::rc::Rc; + +use crate::command::metadata; +use crate::command::{DispatchResult, dispatch, dispatch_read, is_dispatch_read_supported}; +use crate::persistence::aof::{self, AofMessage}; +use crate::protocol::Frame; +use crate::shard::dispatch::key_to_shard; +use crate::shard::mesh::ChannelMesh; +use crate::storage::eviction::{try_evict_if_needed, try_evict_if_needed_async_spill}; +use crate::workspace::{strip_workspace_prefix_from_response, workspace_rewrite_args}; + +use super::affinity::MigratedConnectionState; +use super::{ + apply_resp3_conversion, convert_blocking_to_nonblocking, execute_transaction_sharded, + extract_bytes, extract_command, extract_primary_key, handle_blocking_command_monoio, + handle_config, is_multi_key_command, propagate_subscription, try_inline_dispatch_loop, + unpropagate_subscription, +}; +use crate::framevec; +use crate::pubsub::subscriber::Subscriber; +use crate::server::codec::RespCodec; +use crate::shard::dispatch::ShardMessage; +// ResponseSlotPool NOT used on monoio — its AtomicWaker doesn't cross +// monoio's single-threaded (!Send) executor boundary. Use oneshot channels. + +/// Result of `handle_connection_sharded_monoio` execution. +/// +/// Same purpose as the Tokio handler's `HandlerResult`: the generic handler cannot +/// perform FD extraction, so it returns the stream when migration is triggered. +#[cfg(feature = "runtime-monoio")] +pub enum MonoioHandlerResult { + /// Normal connection close. + Done, + /// Migration triggered: caller should extract raw FD and send via SPSC. + MigrateConnection { + state: MigratedConnectionState, + target_shard: usize, + }, +} + +/// Monoio connection handler using ownership-based I/O (AsyncReadRent/AsyncWriteRent). +/// Dispatches commands through `crate::command::dispatch()` with monoio's ownership I/O model. +#[cfg(feature = "runtime-monoio")] +#[tracing::instrument(skip_all, level = "debug")] +pub(crate) async fn handle_connection_sharded_monoio< + S: monoio::io::AsyncReadRent + monoio::io::AsyncWriteRent, +>( + mut stream: S, + peer_addr: String, + ctx: &super::core::ConnectionContext, + shutdown: CancellationToken, + client_id: u64, + can_migrate: bool, + initial_read_buf: BytesMut, + pending_wakers: Rc>>, + migrated_state: Option<&MigratedConnectionState>, +) -> (MonoioHandlerResult, Option) { + use monoio::io::AsyncWriteRentExt; + + // NOTE: do NOT call record_connection_opened() here — the caller + // (conn_accept.rs) already increments via try_accept_connection(). + + let mut read_buf = if initial_read_buf.is_empty() { + BytesMut::with_capacity(8192) + } else { + let mut buf = initial_read_buf; + buf.reserve(8192); + buf + }; + let mut write_buf = BytesMut::with_capacity(8192); + let mut codec = RespCodec::default(); + let mut conn = super::core::ConnectionState::new( + client_id, + peer_addr.clone(), + &ctx.requirepass, + ctx.shard_id, + ctx.num_shards, + can_migrate, + ctx.runtime_config.read().acllog_max_len, + migrated_state, + ); + conn.refresh_acl_cache(&ctx.acl_table); + let db_count = ctx.shard_databases.db_count(); + + // Register in global client registry for CLIENT LIST/INFO/KILL. + crate::client_registry::register( + client_id, + peer_addr.clone(), + conn.current_user.clone(), + ctx.shard_id, + ); + struct RegistryGuard(u64); + impl Drop for RegistryGuard { + fn drop(&mut self) { + crate::client_registry::deregister(self.0); + } + } + let _registry_guard = RegistryGuard(client_id); + + // Functions API registry (per-connection, lazy init) — kept as local because Rc> is !Send + let func_registry = Rc::new(RefCell::new(crate::scripting::FunctionRegistry::new())); + + // Pre-allocate read buffer outside the loop to avoid per-read heap allocation. + // Monoio's ownership I/O takes ownership and returns the buffer, so we reassign. + let mut tmp_buf = vec![0u8; 8192]; + + // Client idle timeout: 0 = disabled (read once, avoid lock on hot path) + let idle_timeout_secs = ctx.runtime_config.read().timeout; + let idle_timeout = if idle_timeout_secs > 0 { + Some(std::time::Duration::from_secs(idle_timeout_secs)) + } else { + None + }; + + // Pre-allocate batch containers outside the loop to avoid per-batch heap allocation. + // These are cleared and reused each iteration instead of being recreated. + let mut responses: Vec = Vec::with_capacity(64); + let mut remote_groups: HashMap< + usize, + Vec<(usize, std::sync::Arc, Option, Bytes)>, + > = HashMap::with_capacity(ctx.num_shards); + let mut reply_futures: Vec<(Vec<(usize, Option, Bytes)>, usize)> = + Vec::with_capacity(ctx.num_shards); + + // Pre-allocate frames Vec outside the loop; reused via .clear() each iteration. + let mut frames: Vec = Vec::with_capacity(64); + + loop { + // Check if CLIENT KILL targeted this connection + if crate::client_registry::is_killed(client_id) { + break; + } + + // Subscriber mode: bidirectional select on client commands + published messages + if conn.subscription_count > 0 { + #[allow(clippy::unwrap_used)] + // conn.pubsub_rx is always Some when conn.subscription_count > 0 + let rx = conn.pubsub_rx.as_ref().unwrap(); + let sub_tmp_buf = vec![0u8; 8192]; + monoio::select! { + read_result = stream.read(sub_tmp_buf) => { + let (result, buf) = read_result; + match result { + Ok(0) => { + // Client half-closed — break out of loop. + // Stream drop (end of function) triggers monoio's cleanup. + break; + } + Ok(n) => { + read_buf.extend_from_slice(&buf[..n]); + // Parse frames from buffer + loop { + match codec.decode_frame(&mut read_buf) { + Ok(Some(frame)) => { + if let Some((cmd, cmd_args)) = extract_command(&frame) { + match cmd { + _ if cmd.eq_ignore_ascii_case(b"SUBSCRIBE") => { + if cmd_args.is_empty() { + let err = Frame::Error(Bytes::from_static(b"ERR wrong number of arguments for 'subscribe' command")); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&err, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if wr.is_err() { return (MonoioHandlerResult::Done, None); } + continue; + } + for arg in cmd_args { + if let Some(channel) = extract_bytes(arg) { + // ACL channel permission check + let denied = { + #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable + let acl_guard = ctx.acl_table.read().unwrap(); + acl_guard.check_channel_permission(&conn.current_user, channel.as_ref()) + }; + if let Some(deny_reason) = denied { + let err = Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason))); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&err, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if wr.is_err() { return (MonoioHandlerResult::Done, None); } + continue; + } + #[allow(clippy::unwrap_used)] // conn.pubsub_tx is always Some when in subscriber mode + let sub = Subscriber::with_protocol( + conn.pubsub_tx.clone().unwrap(), + conn.subscriber_id, + conn.protocol_version >= 3, + ); + ctx.pubsub_registry.write().subscribe(channel.clone(), sub); + propagate_subscription(&ctx.all_remote_sub_maps, &channel, ctx.shard_id, ctx.num_shards, false); + conn.subscription_count += 1; + // Register pub/sub affinity for this client IP + if conn.subscription_count == 1 { + if let Ok(addr) = peer_addr.parse::() { + ctx.pubsub_affinity.write().register(addr.ip(), ctx.shard_id); + } + } + let resp = crate::pubsub::subscribe_response(&channel, conn.subscription_count); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&resp, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if wr.is_err() { return (MonoioHandlerResult::Done, None); } + } + } + } + _ if cmd.eq_ignore_ascii_case(b"UNSUBSCRIBE") => { + if cmd_args.is_empty() { + let removed = ctx.pubsub_registry.write().unsubscribe_all(conn.subscriber_id); + for ch in &removed { + unpropagate_subscription(&ctx.all_remote_sub_maps, ch, ctx.shard_id, ctx.num_shards, false); + } + if removed.is_empty() { + conn.subscription_count = ctx.pubsub_registry.read().total_subscription_count(conn.subscriber_id); + let resp = crate::pubsub::unsubscribe_response(&Bytes::from_static(b""), conn.subscription_count); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&resp, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if wr.is_err() { return (MonoioHandlerResult::Done, None); } + } else { + for ch in &removed { + conn.subscription_count = conn.subscription_count.saturating_sub(1); + let resp = crate::pubsub::unsubscribe_response(ch, conn.subscription_count); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&resp, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if wr.is_err() { return (MonoioHandlerResult::Done, None); } + } + } + } else { + for arg in cmd_args { + if let Some(channel) = extract_bytes(arg) { + ctx.pubsub_registry.write().unsubscribe(channel.as_ref(), conn.subscriber_id); + unpropagate_subscription(&ctx.all_remote_sub_maps, &channel, ctx.shard_id, ctx.num_shards, false); + conn.subscription_count = conn.subscription_count.saturating_sub(1); + let resp = crate::pubsub::unsubscribe_response(&channel, conn.subscription_count); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&resp, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if wr.is_err() { return (MonoioHandlerResult::Done, None); } + } + } + } + } + _ if cmd.eq_ignore_ascii_case(b"PSUBSCRIBE") => { + if cmd_args.is_empty() { + let err = Frame::Error(Bytes::from_static(b"ERR wrong number of arguments for 'psubscribe' command")); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&err, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if wr.is_err() { return (MonoioHandlerResult::Done, None); } + continue; + } + for arg in cmd_args { + if let Some(pattern) = extract_bytes(arg) { + // ACL channel permission check + let denied = { + #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable + let acl_guard = ctx.acl_table.read().unwrap(); + acl_guard.check_channel_permission(&conn.current_user, pattern.as_ref()) + }; + if let Some(deny_reason) = denied { + let err = Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason))); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&err, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if wr.is_err() { return (MonoioHandlerResult::Done, None); } + continue; + } + #[allow(clippy::unwrap_used)] // conn.pubsub_tx is always Some when in subscriber mode + let sub = Subscriber::with_protocol( + conn.pubsub_tx.clone().unwrap(), + conn.subscriber_id, + conn.protocol_version >= 3, + ); + ctx.pubsub_registry.write().psubscribe(pattern.clone(), sub); + propagate_subscription(&ctx.all_remote_sub_maps, &pattern, ctx.shard_id, ctx.num_shards, true); + conn.subscription_count += 1; + // Register pub/sub affinity for this client IP + if conn.subscription_count == 1 { + if let Ok(addr) = peer_addr.parse::() { + ctx.pubsub_affinity.write().register(addr.ip(), ctx.shard_id); + } + } + let resp = crate::pubsub::psubscribe_response(&pattern, conn.subscription_count); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&resp, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if wr.is_err() { return (MonoioHandlerResult::Done, None); } + } + } + } + _ if cmd.eq_ignore_ascii_case(b"PUNSUBSCRIBE") => { + if cmd_args.is_empty() { + let removed = ctx.pubsub_registry.write().punsubscribe_all(conn.subscriber_id); + for pat in &removed { + unpropagate_subscription(&ctx.all_remote_sub_maps, pat, ctx.shard_id, ctx.num_shards, true); + } + if removed.is_empty() { + conn.subscription_count = ctx.pubsub_registry.read().total_subscription_count(conn.subscriber_id); + let resp = crate::pubsub::punsubscribe_response(&Bytes::from_static(b""), conn.subscription_count); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&resp, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if wr.is_err() { return (MonoioHandlerResult::Done, None); } + } else { + for pat in &removed { + conn.subscription_count = conn.subscription_count.saturating_sub(1); + let resp = crate::pubsub::punsubscribe_response(pat, conn.subscription_count); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&resp, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if wr.is_err() { return (MonoioHandlerResult::Done, None); } + } + } + } else { + for arg in cmd_args { + if let Some(pattern) = extract_bytes(arg) { + ctx.pubsub_registry.write().punsubscribe(pattern.as_ref(), conn.subscriber_id); + unpropagate_subscription(&ctx.all_remote_sub_maps, &pattern, ctx.shard_id, ctx.num_shards, true); + conn.subscription_count = conn.subscription_count.saturating_sub(1); + let resp = crate::pubsub::punsubscribe_response(&pattern, conn.subscription_count); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&resp, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if wr.is_err() { return (MonoioHandlerResult::Done, None); } + } + } + } + } + _ if cmd.eq_ignore_ascii_case(b"PING") => { + let resp = Frame::Array(framevec![ + Frame::BulkString(Bytes::from_static(b"pong")), + Frame::BulkString(Bytes::from_static(b"")), + ]); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&resp, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if wr.is_err() { return (MonoioHandlerResult::Done, None); } + } + _ if cmd.eq_ignore_ascii_case(b"QUIT") => { + let resp = Frame::SimpleString(Bytes::from_static(b"OK")); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&resp, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + let _ = wr; // ignore write error on quit + return (MonoioHandlerResult::Done, None); // exit connection + } + _ => { + let cmd_str = String::from_utf8_lossy(cmd); + let err = Frame::Error(Bytes::from(format!( + "ERR Can't execute '{}': only (P)SUBSCRIBE / (P)UNSUBSCRIBE / PING / QUIT are allowed in this context", + cmd_str.to_lowercase() + ))); + let mut resp_buf = BytesMut::new(); + codec.encode_frame(&err, &mut resp_buf); + let data = resp_buf.freeze(); + let (wr, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if wr.is_err() { return (MonoioHandlerResult::Done, None); } + } + } + } + } + Ok(None) => break, // need more data + Err(_) => return (MonoioHandlerResult::Done, None), // parse error + } + } + } + Err(_) => break, // connection error + } + } + msg = rx.recv_async() => { + match msg { + Ok(data) => { + // Data is pre-serialized RESP bytes — write directly + let (result, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if result.is_err() { break; } + } + Err(_) => break, // all senders dropped + } + } + _ = shutdown.cancelled() => { break; } + } + continue; + } + + // Read data from stream using monoio ownership I/O. + // Reuse pre-allocated buffer; restore length to 8192 for the read. + if tmp_buf.len() < 8192 { + tmp_buf.resize(8192, 0); + } + if let Some(dur) = idle_timeout { + // Timeout-aware read: select between read and sleep. + // monoio::select! drops the losing future, so tmp_buf ownership transfers. + // We allocate a fresh buffer when timeout is enabled (safety feature, not hot path). + let timeout_buf = std::mem::take(&mut tmp_buf); + monoio::select! { + read_result = stream.read(timeout_buf) => { + let (result, returned_buf) = read_result; + tmp_buf = returned_buf; + match result { + Ok(0) => break, + Ok(n) => { read_buf.extend_from_slice(&tmp_buf[..n]); } + Err(_) => break, + } + } + _ = monoio::time::sleep(dur) => { + tracing::debug!("Connection {} idle timeout ({}s)", client_id, idle_timeout_secs); + break; + } + } + } else { + let (result, returned_buf) = stream.read(tmp_buf).await; + tmp_buf = returned_buf; + match result { + Ok(0) => break, + Ok(n) => { + read_buf.extend_from_slice(&tmp_buf[..n]); + } + Err(_) => break, + } + } + + // Inline dispatch: GET/SET directly from raw bytes, skipping Frame construction. + // Skip when unauthenticated or workspace-bound (prefix injection in normal path only). + if conn.authenticated && conn.workspace_id.is_none() { + // Inline writes safe only when: ACL unrestricted, !in_multi, !tracking, + // !is_replica, no spill_sender. Replica check is non-blocking try_read. + let is_replica = ctx.repl_state.as_ref().is_some_and(|rs| { + rs.try_read().is_ok_and(|g| { + matches!( + g.role, + crate::replication::state::ReplicationRole::Replica { .. } + ) + }) + }); + let can_inline_writes = conn.acl_skip_allowed() + && !conn.in_multi + && !conn.tracking_state.enabled + && !is_replica + && ctx.spill_sender.is_none(); + let inlined = try_inline_dispatch_loop( + &mut read_buf, + &mut write_buf, + &ctx.shard_databases, + ctx.shard_id, + conn.selected_db, + &ctx.aof_tx, + ctx.cached_clock.ms(), + ctx.num_shards, + can_inline_writes, + &ctx.runtime_config, + ); + crate::admin::metrics_setup::record_dispatch_local_inline(inlined as u64); + if inlined > 0 && read_buf.is_empty() { + // All commands were inlined -- flush write_buf and continue + if !write_buf.is_empty() { + let data = write_buf.split().freeze(); + let (result, _): (std::io::Result, bytes::Bytes) = + stream.write_all(data).await; + if result.is_err() { + break; + } + } + continue; + } + // If read_buf still has data, fall through to normal Frame parsing + // for remaining commands. Inlined responses are already in write_buf. + } + + // Parse all complete frames from the read buffer (reuse pre-allocated Vec, cap at 1024) + frames.clear(); + loop { + match codec.decode_frame(&mut read_buf) { + Ok(Some(frame)) => { + frames.push(frame); + if frames.len() >= 1024 { + break; + } + } + Ok(None) => break, + Err(_) => return (MonoioHandlerResult::Done, None), // parse error, close connection + } + } + + if frames.is_empty() { + continue; + } + + // CLIENT PAUSE: delay processing if server is paused + crate::client_pause::expire_if_needed(); + if let Some(remaining) = crate::client_pause::check_pause(true) { + monoio::time::sleep(remaining).await; + } + + // Process frames (do NOT clear write_buf -- may have inline dispatch responses). + let mut should_quit = false; + responses.clear(); + remote_groups.clear(); + let mut publish_batches: std::collections::HashMap> = + std::collections::HashMap::new(); + + // Refresh time once per batch — sub-millisecond accuracy not needed per-command. + { + let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + guard.refresh_now_from_cache(&ctx.cached_clock); + } + + let mut auth_delay_ms: u64 = 0; + + for frame in frames.drain(..) { + // --- AUTH gate --- + match dispatch::check_auth_gate( + &frame, + &mut conn, + ctx, + &peer_addr, + client_id, + &mut responses, + &mut auth_delay_ms, + ) { + dispatch::AuthGateResult::Consumed => continue, + dispatch::AuthGateResult::Quit => { + should_quit = true; + break; + } + dispatch::AuthGateResult::NotAuth => { + responses.push(Frame::Error(Bytes::from_static( + b"NOAUTH Authentication required.", + ))); + continue; + } + dispatch::AuthGateResult::Authenticated => {} + } + + let (cmd, cmd_args) = match extract_command(&frame) { + Some(pair) => pair, + None => { + responses.push(Frame::Error(Bytes::from_static( + b"ERR invalid command format", + ))); + continue; + } + }; + + // --- QUIT --- + if cmd.eq_ignore_ascii_case(b"QUIT") { + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + should_quit = true; + break; + } + // --- ASKING --- + if cmd.eq_ignore_ascii_case(b"ASKING") { + conn.asking = true; + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + continue; + } + // --- Connection-level commands (dispatched to dispatch.rs) --- + if dispatch::try_handle_cluster(cmd, cmd_args, ctx, &mut responses) { + continue; + } + if dispatch::try_handle_evalsha(cmd, cmd_args, &conn, ctx, &mut responses) { + continue; + } + if dispatch::try_handle_eval(cmd, cmd_args, &conn, ctx, &mut responses) { + continue; + } + if dispatch::try_handle_script(cmd, cmd_args, ctx, &mut responses) { + continue; + } + if dispatch::try_handle_cluster_routing(cmd, cmd_args, &mut conn, ctx, &mut responses) { + continue; + } + if dispatch::try_handle_auth( + cmd, + cmd_args, + &mut conn, + ctx, + &peer_addr, + &mut auth_delay_ms, + &mut responses, + ) { + continue; + } + if dispatch::try_handle_hello( + cmd, + cmd_args, + &mut conn, + ctx, + client_id, + &peer_addr, + &mut auth_delay_ms, + &mut responses, + ) { + continue; + } + if dispatch::try_handle_acl(cmd, cmd_args, &mut conn, ctx, &peer_addr, &mut responses) { + continue; + } + if dispatch::try_handle_config(cmd, cmd_args, ctx, &mut responses) { + continue; + } + if dispatch::try_handle_replicaof(cmd, cmd_args, ctx, &mut responses) { + continue; + } + if dispatch::try_handle_replconf(cmd, cmd_args, &mut responses) { + continue; + } + if dispatch::try_handle_info(cmd, cmd_args, &conn, ctx, &mut responses) { + continue; + } + if dispatch::try_enforce_readonly(cmd, ctx, &mut responses) { + continue; + } + // CLIENT early (ID, SETNAME, GETNAME, TRACKING) -- admin subcmds fall through to ACL gate + if dispatch::try_handle_client_early( + cmd, + cmd_args, + client_id, + &mut conn, + ctx, + &mut responses, + ) { + continue; + } + // --- Pub/sub commands --- + if pubsub::try_handle_publish( + cmd, + cmd_args, + &conn, + ctx, + &mut responses, + &mut publish_batches, + ) { + continue; + } + match pubsub::try_handle_subscribe_entry( + cmd, + cmd_args, + &mut conn, + ctx, + &peer_addr, + &mut responses, + &mut codec, + &mut write_buf, + &mut stream, + ) + .await + { + pubsub::SubscribeResult::NotSubscribe => {} + pubsub::SubscribeResult::ArgError => continue, + pubsub::SubscribeResult::Subscribed => break, + pubsub::SubscribeResult::WriteError => return (MonoioHandlerResult::Done, None), + } + if pubsub::try_handle_unsubscribe(cmd, &mut responses) { + continue; + } + if pubsub::try_handle_pubsub_introspection(cmd, cmd_args, ctx, &mut responses) { + continue; + } + // --- Persistence + ACL gate + CLIENT admin + Functions --- + if dispatch::try_handle_persistence(cmd, ctx, &mut responses) { + continue; + } + if dispatch::try_enforce_acl(cmd, cmd_args, &mut conn, ctx, &peer_addr, &mut responses) + { + continue; + } + if dispatch::try_handle_client_admin(cmd, cmd_args, client_id, &conn, &mut responses) { + continue; + } + if dispatch::try_handle_functions( + cmd, + cmd_args, + &conn, + ctx, + &func_registry, + &mut responses, + ) { + continue; + } + + // --- TXN.BEGIN / TXN.COMMIT / TXN.ABORT --- + if txn::try_handle_txn_begin(cmd, cmd_args, &mut conn, ctx, &mut responses) { + continue; + } + if txn::try_handle_txn_commit(cmd, cmd_args, &mut conn, ctx, &mut responses) { + continue; + } + if txn::try_handle_txn_abort(cmd, cmd_args, &mut conn, ctx, &mut responses) { + continue; + } + + // --- TEMPORAL.SNAPSHOT_AT / TEMPORAL.INVALIDATE --- + if txn::try_handle_temporal_snapshot_at(cmd, cmd_args, ctx, &mut responses) { + continue; + } + if txn::try_handle_temporal_invalidate(cmd, cmd_args, ctx, &mut responses) { + continue; + } + + // --- WS.* --- + if write::try_handle_ws_command(cmd, cmd_args, &mut conn, ctx, &mut responses) { + continue; + } + + // --- MQ.* --- + if write::try_handle_mq_command(cmd, cmd_args, &mut conn, ctx, &mut responses) { + continue; + } + + // --- MULTI / EXEC / DISCARD --- + if write::try_handle_multi_exec(cmd, &mut conn, ctx, &mut responses) { + continue; + } + + // --- Workspace key prefix injection --- + // MUST happen before key_to_shard() so the {ws_id} hash tag determines + // shard routing. This is the ONLY code path where workspace prefixing + // occurs (WS-07, WS-12). All subsequent dispatch uses cmd_args (shadowed). + let rewritten = conn + .workspace_id + .as_ref() + .map(|ws_id| workspace_rewrite_args(cmd, cmd_args, ws_id)); + let cmd_args: &[Frame] = rewritten.as_deref().unwrap_or(cmd_args); + + // --- BLOCKING COMMANDS --- + match dispatch::try_handle_blocking( + cmd, + cmd_args, + &mut conn, + ctx, + &mut responses, + &mut codec, + &mut write_buf, + &mut stream, + &shutdown, + ) + .await + { + dispatch::BlockingResult::NotBlocking => {} + dispatch::BlockingResult::Queued => continue, + dispatch::BlockingResult::Handled => break, + dispatch::BlockingResult::WriteError => return (MonoioHandlerResult::Done, None), + } + + // --- MULTI queue mode: queue commands when in transaction --- + if conn.in_multi { + conn.command_queue.push(frame); + responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); + continue; + } + + // --- Cross-shard aggregation commands: KEYS, SCAN, DBSIZE + multi-key --- + if dispatch::try_handle_cross_shard_commands(cmd, cmd_args, &conn, ctx, &mut responses) + .await + { + continue; + } + + // --- FT.* vector search commands --- + if ft::try_handle_ft_command(cmd, cmd_args, &frame, &conn, ctx, &mut responses).await { + continue; + } + + // --- GRAPH.* graph commands --- + #[cfg(feature = "graph")] + if write::try_handle_graph_command(cmd, cmd_args, &mut conn, ctx, &mut responses) { + continue; + } + + // --- Routing: keyless, local, or remote --- + let target_shard = + extract_primary_key(cmd, cmd_args).map(|key| key_to_shard(key, ctx.num_shards)); + + let is_local = match target_shard { + None => true, + Some(s) if s == ctx.shard_id => true, + _ => false, + }; + + // Affinity sampling: record shard target for migration decision. + // Migration is deferred until AFTER the current batch is fully processed. + if let (Some(tracker), Some(target)) = (&mut conn.affinity_tracker, target_shard) { + if let Some(migrate_to) = tracker.record(target) { + if !conn.in_multi + && conn.subscription_count == 0 + && !conn.tracking_state.enabled + { + conn.migration_target = Some(migrate_to); + } + } + } + + // Pre-classify write commands for AOF + tracking + let is_write = if ctx.aof_tx.is_some() || conn.tracking_state.enabled { + metadata::is_write(cmd) + } else { + false + }; + + if is_local { + crate::admin::metrics_setup::record_dispatch_local(); + if metadata::is_write(cmd) { + // WRITE PATH: eviction + dispatch under write lock. + let rt = ctx.runtime_config.read(); + let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + let evict_result = if let Some(ref sender) = ctx.spill_sender { + let mut fid = ctx.spill_file_id.get(); + let dir = ctx + .disk_offload_dir + .as_deref() + .unwrap_or(std::path::Path::new(".")); + let res = try_evict_if_needed_async_spill( + &mut guard, + &rt, + sender, + dir, + &mut fid, + conn.selected_db, + ); + ctx.spill_file_id.set(fid); + res + } else { + try_evict_if_needed(&mut guard, &rt) + }; + if let Err(oom_frame) = evict_result { + drop(guard); + drop(rt); + responses.push(oom_frame); + continue; + } + drop(rt); + + // KV undo-log capture for active cross-store transactions. + // MUST happen BEFORE dispatch() overwrites the database entry. + if let Some(ref mut txn) = conn.active_cross_txn { + if cmd.eq_ignore_ascii_case(b"DEL") || cmd.eq_ignore_ascii_case(b"UNLINK") { + // Multi-key DEL: iterate all args for undo capture + for arg in cmd_args.iter() { + if let Frame::BulkString(key_bytes) = arg { + if let Some(old_entry) = guard.get(key_bytes.as_ref()).cloned() + { + txn.kv_undo.record_delete(key_bytes.clone(), old_entry); + let lsn = txn.snapshot_lsn; + let tid = txn.txn_id; + ctx.shard_databases.kv_intents(ctx.shard_id).record_write( + key_bytes.clone(), + lsn, + tid, + ); + } + // Key not found: nothing to undo, no intent registered + } + } + } else if let Some(key) = + crate::server::conn::shared::extract_primary_key(cmd, cmd_args) + { + // SET / HSET / INCR / etc. — single primary key + let old_entry = guard.get(key.as_ref()).cloned(); + let lsn = txn.snapshot_lsn; + let tid = txn.txn_id; + match old_entry { + None => txn.kv_undo.record_insert(key.clone()), + Some(entry) => txn.kv_undo.record_update(key.clone(), entry), + } + ctx.shard_databases.kv_intents(ctx.shard_id).record_write( + key.clone(), + lsn, + tid, + ); + } + } + + let dispatch_start = std::time::Instant::now(); + let result = + dispatch(&mut guard, cmd, cmd_args, &mut conn.selected_db, db_count); + let elapsed_us = dispatch_start.elapsed().as_micros() as u64; + if let Ok(cmd_str) = std::str::from_utf8(cmd) { + crate::admin::metrics_setup::record_command(cmd_str, elapsed_us); + } + if let Frame::Array(ref args) = frame { + crate::admin::metrics_setup::global_slowlog().maybe_record( + elapsed_us, + args.as_slice(), + peer_addr.as_bytes(), + conn.client_name + .as_ref() + .map_or(b"" as &[u8], |n| n.as_ref()), + ); + } + + let response = match result { + DispatchResult::Response(f) => f, + DispatchResult::Quit(f) => { + should_quit = true; + f + } + }; + + // AOF logging for successful local writes + if !matches!(response, Frame::Error(_)) && is_write { + if let Some(ref tx) = ctx.aof_tx { + let serialized = aof::serialize_command(&frame); + let _ = tx.try_send(AofMessage::Append(serialized)); + } + } + + // Auto-index HSET into vector/text stores (if key matches index prefix) + if !matches!(response, Frame::Error(_)) && cmd.eq_ignore_ascii_case(b"HSET") { + if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { + let inserted = { + let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); + let mut ts = ctx.shard_databases.text_store(ctx.shard_id); + crate::shard::spsc_handler::auto_index_hset_public( + &mut vs, + &mut *ts, + key.as_ref(), + cmd_args, + ) + }; + // Phase 166 (Plan 02): if inside an active cross-store + // TXN, push one VectorIntent per (index_name, key_hash) + // so TXN.ABORT (Plan 166-03) can tombstone via + // MutableSegment::mark_deleted_by_key_hash. Non-TXN + // paths discard the tuples — auto_index_hset already + // performed the append. + if let Some(txn) = conn.active_cross_txn.as_mut() { + for (index_name, key_hash) in inserted { + txn.record_vector(key_hash, index_name); + } + } + } + } + + // Post-dispatch wakeup hooks for producer commands + if !matches!(response, Frame::Error(_)) { + let needs_wake = cmd.eq_ignore_ascii_case(b"LPUSH") + || cmd.eq_ignore_ascii_case(b"RPUSH") + || cmd.eq_ignore_ascii_case(b"LMOVE") + || cmd.eq_ignore_ascii_case(b"ZADD"); + if needs_wake { + if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { + let mut reg = ctx.blocking_registry.borrow_mut(); + if cmd.eq_ignore_ascii_case(b"LPUSH") + || cmd.eq_ignore_ascii_case(b"RPUSH") + || cmd.eq_ignore_ascii_case(b"LMOVE") + { + crate::blocking::wakeup::try_wake_list_waiter( + &mut reg, + &mut guard, + conn.selected_db, + &key, + ); + } else { + crate::blocking::wakeup::try_wake_zset_waiter( + &mut reg, + &mut guard, + conn.selected_db, + &key, + ); + } + } + } + } + + drop(guard); + + // Track key on write / invalidate tracked keys + if conn.tracking_state.enabled && !matches!(response, Frame::Error(_)) { + if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { + let senders = ctx + .tracking_table + .borrow_mut() + .invalidate_key(&key, client_id); + if !senders.is_empty() { + let push = crate::tracking::invalidation::invalidation_push(&[key]); + for tx in senders { + let _ = tx.try_send(push.clone()); + } + } + } + } + let mut response = apply_resp3_conversion(cmd, response, conn.protocol_version); + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + } else { + // Snapshot visibility filter for active cross-store transactions. + // MVCC: hide keys written by uncommitted foreign transactions. + if conn.in_cross_txn() { + if let Some(ref txn) = conn.active_cross_txn { + if let Some(key) = + crate::server::conn::shared::extract_primary_key(cmd, cmd_args) + { + let snapshot_lsn = txn.snapshot_lsn; + let my_txn_id = txn.txn_id; + // Clone committed treemap to release vector_store lock + // before acquiring kv_intents lock (lock ordering). + let committed = { + let vs = ctx.shard_databases.vector_store(ctx.shard_id); + vs.txn_manager().committed_treemap().clone() + }; + let visible = { + let intents = ctx.shard_databases.kv_intents(ctx.shard_id); + intents.is_key_visible( + key.as_ref(), + snapshot_lsn, + my_txn_id, + &committed, + ) + }; + if !visible { + responses.push(Frame::Null); + continue; + } + } + } + } + + // READ PATH: shared lock — no contention with other shards' reads + let guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); + let now_ms = ctx.cached_clock.ms(); + let dispatch_start = std::time::Instant::now(); + let result = dispatch_read( + &guard, + cmd, + cmd_args, + now_ms, + &mut conn.selected_db, + db_count, + ); + let elapsed_us = dispatch_start.elapsed().as_micros() as u64; + if let Ok(cmd_str) = std::str::from_utf8(cmd) { + crate::admin::metrics_setup::record_command(cmd_str, elapsed_us); + } + if let Frame::Array(ref args) = frame { + crate::admin::metrics_setup::global_slowlog().maybe_record( + elapsed_us, + args.as_slice(), + peer_addr.as_bytes(), + conn.client_name + .as_ref() + .map_or(b"" as &[u8], |n| n.as_ref()), + ); + } + drop(guard); + + let response = match result { + DispatchResult::Response(f) => f, + DispatchResult::Quit(f) => { + should_quit = true; + f + } + }; + + // Track key on local read + if conn.tracking_state.enabled && !conn.tracking_state.bcast { + if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { + ctx.tracking_table.borrow_mut().track_key( + client_id, + &key, + conn.tracking_state.noloop, + ); + } + } + let mut response = apply_resp3_conversion(cmd, response, conn.protocol_version); + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + } // end read/write split + + // (tracking and response push handled inside read/write branches above) + } else if let Some(target) = target_shard { + // TXN cross-shard guard: reject cross-shard writes in active TXN (no undo log). + if conn.in_cross_txn() && metadata::is_write(cmd) { + responses.push(Frame::Error(bytes::Bytes::from_static( + crate::command::transaction::ERR_TXN_CROSS_SHARD, + ))); + continue; + } + // SHARED-READ FAST PATH: bypass SPSC for cross-shard reads. + // Guard: skip if pending writes exist for this target (pipeline ordering). + if !metadata::is_write(cmd) + && !remote_groups.contains_key(&target) + && is_dispatch_read_supported(cmd) + { + crate::admin::metrics_setup::record_dispatch_cross_read_fastpath(); + let guard = ctx.shard_databases.read_db(target, conn.selected_db); + let now_ms = ctx.cached_clock.ms(); + let result = dispatch_read( + &guard, + cmd, + cmd_args, + now_ms, + &mut conn.selected_db, + db_count, + ); + drop(guard); + let response = match result { + DispatchResult::Response(f) => f, + DispatchResult::Quit(f) => { + should_quit = true; + f + } + }; + // Client tracking for cross-shard reads + if conn.tracking_state.enabled && !conn.tracking_state.bcast { + if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { + ctx.tracking_table.borrow_mut().track_key( + client_id, + &key, + conn.tracking_state.noloop, + ); + } + } + let mut response = apply_resp3_conversion(cmd, response, conn.protocol_version); + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + continue; + } + // Cross-shard write: deferred SPSC dispatch. + // When workspace rewriting occurred, rebuild the frame with + // prefixed args so the target shard stores the correct key. + let dispatch_frame = if rewritten.is_some() { + let mut parts = Vec::with_capacity(1 + cmd_args.len()); + parts.push(Frame::BulkString(Bytes::copy_from_slice(cmd))); + parts.extend_from_slice(cmd_args); + Frame::Array(parts.into()) + } else { + frame.clone() + }; + let resp_idx = responses.len(); + responses.push(Frame::Null); // placeholder, filled after batch dispatch + // Pre-compute AOF bytes before moving frame into Arc + let aof_bytes = if ctx.aof_tx.is_some() && metadata::is_write(cmd) { + Some(aof::serialize_command(&dispatch_frame)) + } else { + None + }; + let cmd_bytes = if let Frame::Array(ref args) = dispatch_frame { + extract_bytes(&args[0]).unwrap_or_default() + } else { + Bytes::new() + }; + remote_groups.entry(target).or_default().push(( + resp_idx, + std::sync::Arc::new(dispatch_frame), + aof_bytes, + cmd_bytes, + )); + crate::admin::metrics_setup::record_dispatch_cross_spsc(); + } + } + + // Phase 2a: Flush accumulated PUBLISH batches as PubSubPublishBatch messages + if !publish_batches.is_empty() { + let mut batch_slots: Vec<( + std::sync::Arc, + Vec, + )> = Vec::new(); + { + let mut producers = ctx.dispatch_tx.borrow_mut(); + for (target, entries) in publish_batches.drain() { + let n = entries.len(); + let slot = std::sync::Arc::new( + crate::shard::dispatch::PubSubResponseSlot::with_counts(1, n), + ); + let resp_indices: Vec = entries.iter().map(|(idx, _, _)| *idx).collect(); + let pairs: Vec<(Bytes, Bytes)> = + entries.into_iter().map(|(_, ch, msg)| (ch, msg)).collect(); + + let idx = ChannelMesh::target_index(ctx.shard_id, target); + let batch_msg = ShardMessage::PubSubPublishBatch { + pairs, + slot: slot.clone(), + }; + if producers[idx].try_push(batch_msg).is_ok() { + ctx.spsc_notifiers[target].notify_one(); + } else { + slot.add(0); // push failed, mark as done + } + batch_slots.push((slot, resp_indices)); + } + } + // Resolve all batch slots + for (slot, resp_indices) in &batch_slots { + crate::shard::dispatch::PubSubResponseFuture::new(slot.clone()).await; + for (i, resp_idx) in resp_indices.iter().enumerate() { + let remote_count = slot.counts[i].load(std::sync::atomic::Ordering::Relaxed); + if remote_count > 0 { + if let Frame::Integer(ref mut total) = responses[*resp_idx] { + *total += remote_count; + } + } + } + } + } + + // Phase 2b: Dispatch all deferred remote commands as batched PipelineBatch + // messages (one per target shard), await all in parallel. + if !remote_groups.is_empty() { + reply_futures.clear(); + + let mut oneshot_futures: Vec<( + Vec<(usize, Option, Bytes)>, + channel::OneshotReceiver>, + )> = Vec::new(); + for (target, entries) in remote_groups.drain() { + let (reply_tx, reply_rx) = channel::oneshot(); + let (meta, commands): ( + Vec<(usize, Option, Bytes)>, + Vec>, + ) = entries + .into_iter() + .map(|(idx, arc_frame, aof, cmd)| ((idx, aof, cmd), arc_frame)) + .unzip(); + + let msg = ShardMessage::PipelineBatch { + db_index: conn.selected_db, + commands, + reply_tx, + }; + let target_idx = ChannelMesh::target_index(ctx.shard_id, target); + { + let mut pending = msg; + loop { + let push_result = { + let mut producers = ctx.dispatch_tx.borrow_mut(); + producers[target_idx].try_push(pending) + }; + match push_result { + Ok(()) => { + tracing::trace!( + "Shard {}: pushed PipelineBatch to shard {}, notifying", + ctx.shard_id, + target + ); + ctx.spsc_notifiers[target].notify_one(); + break; + } + Err(val) => { + pending = val; + monoio::time::sleep(std::time::Duration::from_micros(10)).await; + } + } + } + } + oneshot_futures.push((meta, reply_rx)); + } + + // Poll all shard responses via pending_wakers relay (monoio cross-thread waker fix). + // monoio's !Send executor doesn't see cross-thread Waker::wake() from flume. + // Instead, the connection task registers its waker in pending_wakers; the event + // loop drains and wakes them after every SPSC cycle (~1ms). On wake, try_recv() + // checks if the response arrived; if not, re-register and yield again. + for (meta, reply_rx) in oneshot_futures.drain(..) { + tracing::trace!( + "Shard {}: awaiting cross-shard response via pending_wakers", + ctx.shard_id + ); + let shard_responses = { + let pw = pending_wakers.clone(); + loop { + match reply_rx.try_recv() { + Ok(value) => break Ok(value), + Err(flume::TryRecvError::Disconnected) => break Err(()), + Err(flume::TryRecvError::Empty) => { + // Yield once: register waker, return Pending, then Ready on wake. + let mut yielded = false; + std::future::poll_fn(|cx| { + if yielded { + std::task::Poll::Ready(()) + } else { + yielded = true; + pw.borrow_mut().push(cx.waker().clone()); + std::task::Poll::Pending + } + }) + .await; + // After wake, loop back to try_recv + } + } + } + }; + let shard_responses = match shard_responses { + Ok(r) => r, + Err(()) => { + for (resp_idx, _, _) in &meta { + responses[*resp_idx] = Frame::Error(Bytes::from_static( + b"ERR cross-shard dispatch failed", + )); + } + continue; + } + }; + for ((resp_idx, aof_bytes, cmd_name), resp) in meta.into_iter().zip(shard_responses) + { + // AOF logging for successful remote writes + if let Some(bytes) = aof_bytes { + if !matches!(resp, Frame::Error(_)) { + if let Some(ref tx) = ctx.aof_tx { + let _ = tx.try_send(AofMessage::Append(bytes)); + } + } + } + let resp = apply_resp3_conversion(&cmd_name, resp, conn.protocol_version); + responses[resp_idx] = resp; + } + } + } + + // AUTH rate limiting: delay response to slow down brute-force attacks + if auth_delay_ms > 0 { + monoio::time::sleep(std::time::Duration::from_millis(auth_delay_ms)).await; + } + + // Serialize all responses into write_buf, then do ONE write_all syscall. + for response in &responses { + codec.encode_frame(response, &mut write_buf); + } + + // Write all responses in one batch using ownership I/O + if !write_buf.is_empty() { + let data = write_buf.split().freeze(); + let (result, _data): (std::io::Result, bytes::Bytes) = + stream.write_all(data).await; + if result.is_err() { + break; + } + } + + // Update registry with current state after each batch + crate::client_registry::update(client_id, |e| { + e.db = conn.selected_db; + e.last_cmd_at = std::time::Instant::now(); + e.flags = crate::client_registry::ClientFlags { + subscriber: conn.subscription_count > 0, + in_multi: conn.in_multi, + blocked: false, + }; + }); + + // Check if migration was triggered during frame processing. + // All responses for the current batch have been written, so the + // client sees no interruption -- TCP socket stays open. + if let Some(target_shard) = conn.migration_target { + let migrated_state = MigratedConnectionState { + selected_db: conn.selected_db, + authenticated: conn.authenticated, + client_name: conn.client_name.clone(), + protocol_version: conn.protocol_version, + current_user: conn.current_user.clone(), + flags: 0, + read_buf_remainder: read_buf.split(), + client_id, + peer_addr: peer_addr.clone(), + workspace_id: conn.workspace_id, + }; + return ( + MonoioHandlerResult::MigrateConnection { + state: migrated_state, + target_shard, + }, + Some(stream), + ); + } + + if should_quit { + break; + } + + // Check shutdown (polled after each batch -- acceptable for MVP) + if shutdown.is_cancelled() { + break; + } + + // Shrink buffers if they grew too large + if read_buf.capacity() > 65536 { + let remaining = read_buf.split(); + read_buf = BytesMut::with_capacity(8192); + if !remaining.is_empty() { + read_buf.extend_from_slice(&remaining); + } + } + if write_buf.capacity() > 65536 { + write_buf = BytesMut::with_capacity(8192); + } + if tmp_buf.capacity() > 65536 { + tmp_buf = vec![0u8; 8192]; + } + } + + // --- Graceful TCP shutdown: send FIN to client to avoid CLOSE_WAIT --- + // Uses monoio's own shutdown() which properly manages the fd through + // the runtime (unlike raw libc::shutdown which corrupts monoio state). + let _ = stream.shutdown().await; + + // Phase 166: release any leaked cross-store TXN (client disconnected mid-txn). + // Idempotent: TXN.ABORT already takes() active_cross_txn so this is a no-op if abort ran. + // Closes T-161-05 — without this, a disconnect after TXN.BEGIN + SET would leak + // kv_intents and pin the key invisible for all subsequent readers. Mirrors the + // sharded runtime block in handler_sharded.rs so both paths delegate to the same + // shared helper. FIN has already been sent; shard state is still intact. + if let Some(txn) = conn.active_cross_txn.take() { + crate::transaction::abort::abort_cross_store_txn( + &ctx.shard_databases, + ctx.shard_id, + conn.selected_db, + txn, + ); + } + + // --- Disconnect cleanup: propagate unsubscribe to all shards' remote subscriber maps --- + if conn.subscriber_id > 0 { + let removed_channels = { + ctx.pubsub_registry + .write() + .unsubscribe_all(conn.subscriber_id) + }; + let removed_patterns = { + ctx.pubsub_registry + .write() + .punsubscribe_all(conn.subscriber_id) + }; + for ch in removed_channels { + unpropagate_subscription( + &ctx.all_remote_sub_maps, + &ch, + ctx.shard_id, + ctx.num_shards, + false, + ); + } + for pat in removed_patterns { + unpropagate_subscription( + &ctx.all_remote_sub_maps, + &pat, + ctx.shard_id, + ctx.num_shards, + true, + ); + } + // Remove affinity on disconnect (no subscriptions remain) + if let Ok(addr) = peer_addr.parse::() { + ctx.pubsub_affinity.write().remove(&addr.ip()); + } + } + + // NOTE: connection close is recorded by the caller (conn_accept.rs) to + // preserve symmetry with `try_accept_connection`, which owns the + // increment. Decrementing here too produces a double-decrement on the + // AtomicU64 counter — it wraps to u64::MAX on the second subtraction + // and all subsequent `try_accept_connection` comparisons against + // `maxclients` reject new connections. + (MonoioHandlerResult::Done, None) +} diff --git a/src/server/conn/handler_monoio/pubsub.rs b/src/server/conn/handler_monoio/pubsub.rs new file mode 100644 index 00000000..e7ac03f2 --- /dev/null +++ b/src/server/conn/handler_monoio/pubsub.rs @@ -0,0 +1,294 @@ +//! Pub/sub command handlers: PUBLISH, UNSUBSCRIBE/PUNSUBSCRIBE (no-op in normal mode), +//! and PUBSUB introspection (CHANNELS, NUMSUB, NUMPAT). +//! +//! The subscriber-mode select loop and SUBSCRIBE/PSUBSCRIBE entry remain in mod.rs +//! because monoio's ownership I/O model (AsyncReadRent/AsyncWriteRent) requires the +//! stream to be passed by value into `monoio::select!`, which is tightly coupled to +//! the connection loop state machine. + +use bytes::Bytes; + +use crate::protocol::Frame; +use crate::server::conn::core::{ConnectionContext, ConnectionState}; +use crate::server::conn::util::extract_bytes; + +/// Handle PUBLISH command. Returns `true` if the command was consumed. +pub(super) fn try_handle_publish( + cmd: &[u8], + cmd_args: &[Frame], + conn: &ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, + publish_batches: &mut std::collections::HashMap>, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"PUBLISH") { + return false; + } + if cmd_args.len() != 2 { + responses.push(Frame::Error(Bytes::from_static( + b"ERR wrong number of arguments for 'publish' command", + ))); + return true; + } + let channel = extract_bytes(&cmd_args[0]); + let message = extract_bytes(&cmd_args[1]); + // ACL channel permission check for PUBLISH + if let Some(ref ch) = channel { + let denied = { + #[allow(clippy::unwrap_used)] + // std RwLock: poison = prior panic = unrecoverable + let acl_guard = ctx.acl_table.read().unwrap(); + acl_guard.check_channel_permission(&conn.current_user, ch.as_ref()) + }; + if let Some(deny_reason) = denied { + responses.push(Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason)))); + return true; + } + } + match (channel, message) { + (Some(ch), Some(msg)) => { + let local_count = { ctx.pubsub_registry.write().publish(&ch, &msg) }; + // Targeted fanout: only send to shards that have subscribers + let targets = ctx.remote_subscriber_map.read().target_shards(&ch); + if targets.is_empty() { + // Fast path: no remote subscribers + responses.push(Frame::Integer(local_count)); + } else { + let remote_targets: Vec = + targets.into_iter().filter(|&t| t != ctx.shard_id).collect(); + if remote_targets.is_empty() { + responses.push(Frame::Integer(local_count)); + } else { + // Accumulate into per-shard batches for coalesced dispatch + let resp_idx = responses.len(); + responses.push(Frame::Integer(local_count)); // placeholder, updated after batch flush + for target in &remote_targets { + publish_batches.entry(*target).or_default().push(( + resp_idx, + ch.clone(), + msg.clone(), + )); + } + } + } + } + _ => responses.push(Frame::Error(Bytes::from_static( + b"ERR invalid channel or message", + ))), + } + true +} + +/// Handle UNSUBSCRIBE / PUNSUBSCRIBE in normal mode (no-op, not in subscriber mode). +/// Returns `true` if the command was consumed. +pub(super) fn try_handle_unsubscribe(cmd: &[u8], responses: &mut Vec) -> bool { + if cmd.eq_ignore_ascii_case(b"UNSUBSCRIBE") { + responses.push(crate::pubsub::unsubscribe_response( + &Bytes::from_static(b""), + 0, + )); + return true; + } + if cmd.eq_ignore_ascii_case(b"PUNSUBSCRIBE") { + responses.push(crate::pubsub::punsubscribe_response( + &Bytes::from_static(b""), + 0, + )); + return true; + } + false +} + +/// Result of SUBSCRIBE/PSUBSCRIBE dispatch in normal mode. +pub(super) enum SubscribeResult { + /// Not a SUBSCRIBE/PSUBSCRIBE command. + NotSubscribe, + /// Argument validation failed; error pushed to responses. Caller should `continue`. + ArgError, + /// Subscription registered, responses encoded into write_buf. + /// Caller must flush write_buf and `break` the frame loop. + Subscribed, + /// Write error during flush. Caller should return Done. + WriteError, +} + +/// Handle SUBSCRIBE / PSUBSCRIBE entry in normal (non-subscriber) mode. +/// +/// Allocates pubsub channel if needed, registers subscriptions, encodes responses +/// into write_buf, and flushes. Returns `SubscribeResult` to tell the caller +/// whether to break or continue. +pub(super) async fn try_handle_subscribe_entry( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut super::super::core::ConnectionState, + ctx: &super::super::core::ConnectionContext, + peer_addr: &str, + responses: &mut Vec, + codec: &mut crate::server::codec::RespCodec, + write_buf: &mut bytes::BytesMut, + stream: &mut S, +) -> SubscribeResult { + if !cmd.eq_ignore_ascii_case(b"SUBSCRIBE") && !cmd.eq_ignore_ascii_case(b"PSUBSCRIBE") { + return SubscribeResult::NotSubscribe; + } + let is_pattern = cmd.eq_ignore_ascii_case(b"PSUBSCRIBE"); + if cmd_args.is_empty() { + let cmd_name = if is_pattern { + "psubscribe" + } else { + "subscribe" + }; + let err = Frame::Error(Bytes::from(format!( + "ERR wrong number of arguments for '{}' command", + cmd_name + ))); + responses.push(err); + return SubscribeResult::ArgError; + } + // Allocate pubsub channel if not yet created + if conn.pubsub_tx.is_none() { + let (tx, rx) = crate::runtime::channel::mpsc_bounded::(256); + conn.pubsub_tx = Some(tx); + conn.pubsub_rx = Some(rx); + } + if conn.subscriber_id == 0 { + conn.subscriber_id = crate::pubsub::next_subscriber_id(); + } + // Flush accumulated responses before entering subscriber mode + for resp in &*responses { + codec.encode_frame(resp, write_buf); + } + for arg in cmd_args { + if let Some(ch) = extract_bytes(arg) { + // ACL channel permission check + { + #[allow(clippy::unwrap_used)] + // std RwLock: poison = prior panic = unrecoverable + let acl_guard = ctx.acl_table.read().unwrap(); + if let Some(deny_reason) = + acl_guard.check_channel_permission(&conn.current_user, ch.as_ref()) + { + drop(acl_guard); + let err = Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason))); + codec.encode_frame(&err, write_buf); + continue; + } + } + #[allow(clippy::unwrap_used)] + // conn.pubsub_tx is set to Some just above before this loop + let sub = crate::pubsub::subscriber::Subscriber::with_protocol( + conn.pubsub_tx.clone().unwrap(), + conn.subscriber_id, + conn.protocol_version >= 3, + ); + if is_pattern { + ctx.pubsub_registry.write().psubscribe(ch.clone(), sub); + } else { + ctx.pubsub_registry.write().subscribe(ch.clone(), sub); + } + super::propagate_subscription( + &ctx.all_remote_sub_maps, + &ch, + ctx.shard_id, + ctx.num_shards, + is_pattern, + ); + conn.subscription_count += 1; + // Register pub/sub affinity for this client IP + if conn.subscription_count == 1 { + if let Ok(addr) = peer_addr.parse::() { + ctx.pubsub_affinity + .write() + .register(addr.ip(), ctx.shard_id); + } + } + let resp = if is_pattern { + crate::pubsub::psubscribe_response(&ch, conn.subscription_count) + } else { + crate::pubsub::subscribe_response(&ch, conn.subscription_count) + }; + codec.encode_frame(&resp, write_buf); + } + } + // Flush responses and re-enter loop (next iteration enters subscriber mode) + if !write_buf.is_empty() { + use monoio::io::AsyncWriteRentExt; + let data = write_buf.split().freeze(); + let (result, _): (std::io::Result, bytes::Bytes) = stream.write_all(data).await; + if result.is_err() { + return SubscribeResult::WriteError; + } + } + responses.clear(); + SubscribeResult::Subscribed +} + +/// Handle PUBSUB introspection subcommands (CHANNELS, NUMSUB, NUMPAT). +/// Returns `true` if the command was consumed. +pub(super) fn try_handle_pubsub_introspection( + cmd: &[u8], + cmd_args: &[Frame], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"PUBSUB") { + return false; + } + if cmd_args.is_empty() { + responses.push(Frame::Error(Bytes::from_static( + b"ERR wrong number of arguments for 'pubsub' command", + ))); + return true; + } + let subcmd = extract_bytes(&cmd_args[0]); + match subcmd { + Some(ref sc) if sc.eq_ignore_ascii_case(b"CHANNELS") => { + let pattern = if cmd_args.len() > 1 { + extract_bytes(&cmd_args[1]) + } else { + None + }; + let mut all_channels: std::collections::HashSet = + std::collections::HashSet::new(); + for reg in &ctx.all_pubsub_registries { + let guard = reg.read(); + all_channels.extend(guard.active_channels(pattern.as_deref())); + } + let arr: Vec = all_channels.into_iter().map(Frame::BulkString).collect(); + responses.push(Frame::Array(arr.into())); + } + Some(ref sc) if sc.eq_ignore_ascii_case(b"NUMSUB") => { + let channels: Vec = cmd_args[1..] + .iter() + .filter_map(|a| extract_bytes(a)) + .collect(); + let mut counts: std::collections::HashMap = + std::collections::HashMap::new(); + for reg in &ctx.all_pubsub_registries { + let guard = reg.read(); + for (ch, c) in guard.numsub(&channels) { + *counts.entry(ch).or_insert(0) += c; + } + } + let mut arr = Vec::with_capacity(channels.len() * 2); + for ch in &channels { + arr.push(Frame::BulkString(ch.clone())); + arr.push(Frame::Integer(*counts.get(ch).unwrap_or(&0))); + } + responses.push(Frame::Array(arr.into())); + } + Some(ref sc) if sc.eq_ignore_ascii_case(b"NUMPAT") => { + let mut total: usize = 0; + for reg in &ctx.all_pubsub_registries { + total += reg.read().numpat(); + } + responses.push(Frame::Integer(total as i64)); + } + _ => { + responses.push(Frame::Error(Bytes::from_static( + b"ERR unknown subcommand or wrong number of arguments for 'pubsub' command", + ))); + } + } + true +} diff --git a/src/server/conn/handler_monoio/read.rs b/src/server/conn/handler_monoio/read.rs new file mode 100644 index 00000000..5acb68d3 --- /dev/null +++ b/src/server/conn/handler_monoio/read.rs @@ -0,0 +1 @@ +// Placeholder -- read-only dispatch paths. diff --git a/src/server/conn/handler_monoio/txn.rs b/src/server/conn/handler_monoio/txn.rs new file mode 100644 index 00000000..0d157d38 --- /dev/null +++ b/src/server/conn/handler_monoio/txn.rs @@ -0,0 +1,258 @@ +//! TXN.BEGIN / TXN.COMMIT / TXN.ABORT + TEMPORAL.SNAPSHOT_AT / TEMPORAL.INVALIDATE handlers. +//! +//! Each helper returns `true` if the command was consumed (caller should `continue`). + +use bytes::Bytes; + +#[cfg(feature = "graph")] +use crate::command::temporal::{ERR_ENTITY_NOT_FOUND, ERR_GRAPH_NOT_FOUND}; +use crate::command::temporal::{ + capture_wall_ms, is_temporal_invalidate, is_temporal_snapshot_at, validate_invalidate, + validate_snapshot_at, +}; +use crate::command::transaction::{ + is_txn_abort, is_txn_begin, is_txn_commit, txn_abort_validate, txn_begin_validate, + txn_commit_validate, +}; +use crate::protocol::Frame; +use crate::server::conn::core::ConnectionContext; +use crate::server::conn::core::ConnectionState; +use crate::transaction::CrossStoreTxn; + +/// Handle TXN.BEGIN — returns `true` if the command was consumed. +pub(super) fn try_handle_txn_begin( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !is_txn_begin(cmd, cmd_args) { + return false; + } + match txn_begin_validate(conn.in_multi, conn.in_cross_txn()) { + Ok(()) => { + // Get next txn_id and snapshot_lsn from vector store's transaction manager + let mut vector_store = ctx.shard_databases.vector_store(ctx.shard_id); + let active = vector_store.txn_manager_mut().begin(); + conn.active_cross_txn = Some(CrossStoreTxn::new(active.txn_id, active.snapshot_lsn)); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + Err(e) => responses.push(e), + } + true +} + +/// Handle TXN.COMMIT — returns `true` if the command was consumed. +pub(super) fn try_handle_txn_commit( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !is_txn_commit(cmd, cmd_args) { + return false; + } + match txn_commit_validate(conn.in_cross_txn()) { + Ok(()) => { + if let Some(txn) = conn.active_cross_txn.take() { + let mut vector_store = ctx.shard_databases.vector_store(ctx.shard_id); + vector_store.txn_manager_mut().commit(txn.txn_id); + drop(vector_store); + + // Write XactCommit WAL record with committed KV state + let txn_id = txn.txn_id; + if !txn.kv_undo.is_empty() { + let db_guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); + let payload = crate::persistence::wal_v3::record::encode_xact_commit_payload( + txn_id, + txn.kv_undo.records(), + &*db_guard, + ); + drop(db_guard); + let mut wal_buf = Vec::new(); + crate::persistence::wal_v3::record::write_wal_v3_record( + &mut wal_buf, + txn_id, + crate::persistence::wal_v3::record::WalRecordType::XactCommit, + &payload, + ); + ctx.shard_databases + .wal_append(ctx.shard_id, bytes::Bytes::from(wal_buf)); + } + + // Release KV write intents from shard side-table + ctx.shard_databases + .kv_intents(ctx.shard_id) + .release_txn(txn_id); + + // Drain deferred HNSW inserts (post-commit hook). + // The drain prevents phantom neighbors on abort. + // Actual HNSW graph insertion happens during compaction, + // not at commit time (point is already in mutable segment). + let drain_count = ctx + .shard_databases + .hnsw_queue(ctx.shard_id) + .drain_for_txn(txn_id) + .count(); + if drain_count > 0 { + tracing::debug!(txn_id, count = drain_count, "Drained deferred HNSW inserts"); + } + + // Materialize MQ intents: enqueue deferred MQ.PUBLISH messages + if !txn.mq_intents.is_empty() { + let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + for intent in &txn.mq_intents { + if let Ok(Some(stream)) = db_guard.get_stream_mut(&intent.queue_key) { + if stream.durable { + let msg_id = stream.next_auto_id(); + stream.add(msg_id, intent.fields.clone()); + } + } + } + } + + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } else { + responses.push(Frame::Error(Bytes::from_static(b"ERR not in transaction"))); + } + } + Err(e) => responses.push(e), + } + true +} + +/// Handle TXN.ABORT ��� returns `true` if the command was consumed. +pub(super) fn try_handle_txn_abort( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !is_txn_abort(cmd, cmd_args) { + return false; + } + match txn_abort_validate(conn.in_cross_txn()) { + Ok(()) => { + if let Some(txn) = conn.active_cross_txn.take() { + // Shared rollback (Phase 166 Plan 03): + // KV undo -> graph intents reverse -> vector + // tombstone -> side-table release. See + // src/transaction/abort.rs for lock ordering. + crate::transaction::abort::abort_cross_store_txn( + &ctx.shard_databases, + ctx.shard_id, + conn.selected_db, + txn, + ); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } else { + responses.push(Frame::Error(Bytes::from_static(b"ERR not in transaction"))); + } + } + Err(e) => responses.push(e), + } + true +} + +/// Handle TEMPORAL.SNAPSHOT_AT — returns `true` if the command was consumed. +pub(super) fn try_handle_temporal_snapshot_at( + cmd: &[u8], + cmd_args: &[Frame], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !is_temporal_snapshot_at(cmd) { + return false; + } + match validate_snapshot_at(cmd_args) { + Ok(()) => { + let wall_ms = capture_wall_ms(); + let lsn = { + let vector_store = ctx.shard_databases.vector_store(ctx.shard_id); + vector_store.txn_manager().current_lsn() + }; + { + let mut guard = ctx.shard_databases.temporal_registry(ctx.shard_id); + let registry = + guard.get_or_insert_with(|| Box::new(crate::temporal::TemporalRegistry::new())); + registry.record(wall_ms, lsn); + } + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + Err(e) => responses.push(e), + } + true +} + +/// Handle TEMPORAL.INVALIDATE — returns `true` if the command was consumed. +pub(super) fn try_handle_temporal_invalidate( + cmd: &[u8], + cmd_args: &[Frame], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !is_temporal_invalidate(cmd) { + return false; + } + match validate_invalidate(cmd_args) { + Ok((entity_id, is_node, graph_name)) => { + let wall_ms = capture_wall_ms(); + #[cfg(feature = "graph")] + { + let mut gs = ctx.shard_databases.graph_store_write(ctx.shard_id); + if let Some(named_graph) = gs.get_graph_mut(&graph_name) { + let mutated = if is_node { + let node_key: crate::graph::types::NodeKey = + slotmap::KeyData::from_ffi(entity_id).into(); + if let Some(node) = named_graph.write_buf.get_node_mut(node_key) { + node.valid_to = wall_ms; + true + } else { + false + } + } else { + let edge_key: crate::graph::types::EdgeKey = + slotmap::KeyData::from_ffi(entity_id).into(); + if let Some(edge) = named_graph.write_buf.get_edge_mut(edge_key) { + edge.valid_to = wall_ms; + true + } else { + false + } + }; + if mutated { + let payload = crate::persistence::wal_v3::record::encode_graph_temporal( + entity_id, is_node, wall_ms, wall_ms, + ); + gs.wal_pending.push(payload); + let wal_records = gs.drain_wal(); + drop(gs); + for record in wal_records { + ctx.shard_databases + .wal_append(ctx.shard_id, Bytes::from(record)); + } + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } else { + drop(gs); + responses.push(Frame::Error(Bytes::from_static(ERR_ENTITY_NOT_FOUND))); + } + } else { + drop(gs); + responses.push(Frame::Error(Bytes::from_static(ERR_GRAPH_NOT_FOUND))); + } + } + #[cfg(not(feature = "graph"))] + { + let _ = (entity_id, is_node, graph_name, wall_ms, ctx); + responses.push(Frame::Error(Bytes::from_static( + b"ERR graph feature not enabled", + ))); + } + } + Err(e) => responses.push(e), + } + true +} diff --git a/src/server/conn/handler_monoio/write.rs b/src/server/conn/handler_monoio/write.rs new file mode 100644 index 00000000..8e0cdb92 --- /dev/null +++ b/src/server/conn/handler_monoio/write.rs @@ -0,0 +1,741 @@ +//! Write-path command handlers: WS.*, MQ.*, MULTI/EXEC/DISCARD, GRAPH.*. +//! +//! Each helper returns `true` if the command was consumed (caller should `continue`). + +use bytes::Bytes; + +use crate::command::mq::{ + ERR_MQ_NOT_DURABLE, ERR_MQ_UNKNOWN_SUB, parse_mq_subcommand, validate_mq_ack, + validate_mq_create, validate_mq_dlqlen, validate_mq_pop, validate_mq_publish, validate_mq_push, + validate_mq_trigger, +}; +use crate::command::transaction::ERR_MULTI_TXN_CONFLICT; +use crate::command::workspace::{ + ERR_WS_ALREADY_BOUND, ERR_WS_NOT_FOUND, ERR_WS_UNKNOWN_SUB, parse_workspace_id_from_bytes, + parse_ws_subcommand, validate_ws_auth, validate_ws_create, validate_ws_drop, validate_ws_info, + validate_ws_list, +}; +use crate::mq::is_mq_command; +use crate::protocol::Frame; +use crate::server::conn::core::{ConnectionContext, ConnectionState}; +#[cfg(feature = "graph")] +use crate::server::conn::util::extract_bytes; +use crate::storage::stream::StreamId; +#[cfg(feature = "graph")] +use crate::workspace::strip_workspace_prefix_from_response; +use crate::workspace::{WorkspaceId, is_ws_command}; + +use super::execute_transaction_sharded; + +/// Handle WS.* workspace commands. Returns `true` if consumed. +pub(super) fn try_handle_ws_command( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !is_ws_command(cmd) { + return false; + } + let sub = match parse_ws_subcommand(cmd_args) { + Ok(s) => s, + Err(e) => { + responses.push(e); + return true; + } + }; + + if sub.eq_ignore_ascii_case(b"CREATE") { + match validate_ws_create(cmd_args) { + Ok(ws_name) => { + let ws_id = WorkspaceId::new_v7(); + let created_at = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as i64; + let meta = crate::workspace::WorkspaceMetadata { + id: ws_id, + name: ws_name.clone(), + created_at, + }; + { + let mut guard = ctx.shard_databases.workspace_registry(ctx.shard_id); + let reg = guard.get_or_insert_with(|| { + Box::new(crate::workspace::WorkspaceRegistry::new()) + }); + reg.insert(ws_id, meta); + } + // WAL: WorkspaceCreate record + let payload = + crate::workspace::wal::encode_workspace_create(ws_id.as_bytes(), &ws_name); + let mut wal_buf = Vec::new(); + crate::persistence::wal_v3::record::write_wal_v3_record( + &mut wal_buf, + 0, + crate::persistence::wal_v3::record::WalRecordType::WorkspaceCreate, + &payload, + ); + ctx.shard_databases + .wal_append(ctx.shard_id, Bytes::from(wal_buf)); + responses.push(Frame::BulkString(Bytes::from(ws_id.to_string()))); + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"DROP") { + match validate_ws_drop(cmd_args) { + Ok(ws_id_raw) => { + match parse_workspace_id_from_bytes(&ws_id_raw) { + Some(ws_id) => { + let removed = { + let mut guard = ctx.shard_databases.workspace_registry(ctx.shard_id); + match guard.as_mut() { + Some(reg) => reg.remove(&ws_id).is_some(), + None => false, + } + }; + if removed { + // WAL: WorkspaceDrop record + let payload = + crate::workspace::wal::encode_workspace_drop(ws_id.as_bytes()); + let mut wal_buf = Vec::new(); + crate::persistence::wal_v3::record::write_wal_v3_record( + &mut wal_buf, + 0, + crate::persistence::wal_v3::record::WalRecordType::WorkspaceDrop, + &payload, + ); + ctx.shard_databases + .wal_append(ctx.shard_id, Bytes::from(wal_buf)); + // Best-effort cleanup: delete all KV keys with ws prefix (WS-03). + { + let prefix = format!("{{{}}}:", ws_id.as_hex()); + let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, 0); + let keys_to_delete: Vec> = db_guard + .keys() + .filter(|k| k.as_bytes().starts_with(prefix.as_bytes())) + .map(|k| k.as_bytes().to_vec()) + .collect(); + for key in &keys_to_delete { + db_guard.remove(key); + } + } + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } else { + responses.push(Frame::Error(Bytes::from_static(ERR_WS_NOT_FOUND))); + } + } + None => responses.push(Frame::Error(Bytes::from_static( + crate::command::workspace::ERR_WS_INVALID_ID, + ))), + } + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"LIST") { + match validate_ws_list(cmd_args) { + Ok(()) => { + let guard = ctx.shard_databases.workspace_registry(ctx.shard_id); + let entries: Vec = match guard.as_ref() { + Some(reg) => reg + .iter() + .map(|(id, meta)| { + Frame::Array( + vec![ + Frame::BulkString(Bytes::from(id.to_string())), + Frame::BulkString(meta.name.clone()), + Frame::Integer(meta.created_at), + ] + .into(), + ) + }) + .collect(), + None => vec![], + }; + responses.push(Frame::Array(entries.into())); + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"INFO") { + match validate_ws_info(cmd_args) { + Ok(ws_id_raw) => match parse_workspace_id_from_bytes(&ws_id_raw) { + Some(ws_id) => { + let guard = ctx.shard_databases.workspace_registry(ctx.shard_id); + let found = guard.as_ref().and_then(|reg| reg.get(&ws_id)); + match found { + Some(meta) => { + responses.push(Frame::Array( + vec![ + Frame::BulkString(Bytes::from_static(b"id")), + Frame::BulkString(Bytes::from(meta.id.to_string())), + Frame::BulkString(Bytes::from_static(b"name")), + Frame::BulkString(meta.name.clone()), + Frame::BulkString(Bytes::from_static(b"created_at")), + Frame::Integer(meta.created_at), + ] + .into(), + )); + } + None => responses.push(Frame::Error(Bytes::from_static(ERR_WS_NOT_FOUND))), + } + } + None => responses.push(Frame::Error(Bytes::from_static( + crate::command::workspace::ERR_WS_INVALID_ID, + ))), + }, + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"AUTH") { + match validate_ws_auth(cmd_args) { + Ok(ws_id_raw) => { + if conn.workspace_id.is_some() { + responses.push(Frame::Error(Bytes::from_static(ERR_WS_ALREADY_BOUND))); + } else { + match parse_workspace_id_from_bytes(&ws_id_raw) { + Some(ws_id) => { + let found = { + let guard = ctx.shard_databases.workspace_registry(ctx.shard_id); + guard + .as_ref() + .map_or(false, |reg| reg.get(&ws_id).is_some()) + }; + if found { + conn.workspace_id = Some(ws_id); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } else { + responses.push(Frame::Error(Bytes::from_static(ERR_WS_NOT_FOUND))); + } + } + None => responses.push(Frame::Error(Bytes::from_static( + crate::command::workspace::ERR_WS_INVALID_ID, + ))), + } + } + } + Err(e) => responses.push(e), + } + return true; + } + + // Unknown WS subcommand + responses.push(Frame::Error(Bytes::from_static(ERR_WS_UNKNOWN_SUB))); + true +} + +/// Handle MQ.* message queue commands. Returns `true` if consumed. +pub(super) fn try_handle_mq_command( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !is_mq_command(cmd) { + return false; + } + let sub = match parse_mq_subcommand(cmd_args) { + Ok(s) => s, + Err(e) => { + responses.push(e); + return true; + } + }; + + if sub.eq_ignore_ascii_case(b"CREATE") { + match validate_mq_create(cmd_args) { + Ok((queue_key, max_delivery_count, _debounce_ms)) => { + let effective_key = + crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); + // Create or get Stream in db + let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + match db_guard.get_or_create_stream(&effective_key) { + Ok(stream) => { + stream.durable = true; + stream.max_delivery_count = max_delivery_count; + // Auto-create __mq_consumers consumer group if not exists + let group_name = Bytes::from_static(b"__mq_consumers"); + let _ = stream.create_group(group_name, StreamId::ZERO); + } + Err(e) => { + responses.push(e); + return true; + } + } + drop(db_guard); + + // Store config in per-shard registry + let config = + crate::mq::DurableStreamConfig::new(effective_key.clone(), max_delivery_count); + { + let mut guard = ctx.shard_databases.durable_queue_registry(ctx.shard_id); + let reg = guard + .get_or_insert_with(|| Box::new(crate::mq::DurableQueueRegistry::new())); + reg.insert(effective_key.clone(), config); + } + + // WAL: MqCreate record + let payload = crate::mq::wal::encode_mq_create(&effective_key, max_delivery_count); + let mut wal_buf = Vec::new(); + crate::persistence::wal_v3::record::write_wal_v3_record( + &mut wal_buf, + 0, + crate::persistence::wal_v3::record::WalRecordType::MqCreate, + &payload, + ); + ctx.shard_databases + .wal_append(ctx.shard_id, Bytes::from(wal_buf)); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"PUSH") { + match validate_mq_push(cmd_args) { + Ok((queue_key, fields)) => { + let effective_key = + crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); + let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + match db_guard.get_stream_mut(&effective_key) { + Ok(Some(stream)) => { + if !stream.durable { + responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); + } else { + let msg_id = stream.next_auto_id(); + let msg_id = stream.add(msg_id, fields); + drop(db_guard); + // Mark pending for any registered triggers + { + let mut trig_guard = + ctx.shard_databases.trigger_registry(ctx.shard_id); + if let Some(reg) = trig_guard.as_mut() { + let trig_key = if let Some(ws_id) = conn.workspace_id.as_ref() { + let ws_hex = ws_id.as_hex(); + let mut k = + Vec::with_capacity(ws_hex.len() + 1 + queue_key.len()); + k.extend_from_slice(ws_hex.as_bytes()); + k.push(b':'); + k.extend_from_slice(&queue_key); + Bytes::from(k) + } else { + queue_key.clone() + }; + if let Some(trig_entry) = reg.get_mut(&trig_key) { + if trig_entry.pending_fire_ms == 0 { + let fire_at = + ctx.cached_clock.ms() + trig_entry.debounce_ms; + trig_entry.pending_fire_ms = fire_at; + } + } + } + } + responses.push(Frame::BulkString(Bytes::from(format!( + "{}-{}", + msg_id.ms, msg_id.seq + )))); + } + } + Ok(None) => { + responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); + } + Err(e) => responses.push(e), + } + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"POP") { + match validate_mq_pop(cmd_args) { + Ok((queue_key, count)) => { + let effective_key = + crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); + let group_name = Bytes::from_static(b"__mq_consumers"); + let consumer_name = Bytes::from_static(b"__mq_default"); + let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + + // Read max_delivery_count before mutating + let mdc = match db_guard.get_stream_mut(&effective_key) { + Ok(Some(stream)) => { + if !stream.durable { + responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); + return true; + } + stream.max_delivery_count + } + Ok(None) => { + responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); + return true; + } + Err(e) => { + responses.push(e); + return true; + } + }; + + // Request more than count to account for DLQ routing + let request_count = count + (mdc as usize); + let stream = match db_guard.get_stream_mut(&effective_key) { + Ok(Some(s)) => s, + _ => { + responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); + return true; + } + }; + let claimed = match stream.read_group_new( + &group_name, + &consumer_name, + Some(request_count), + false, + ) { + Ok(entries) => entries, + Err(_) => { + responses.push(Frame::Array(vec![].into())); + return true; + } + }; + + // Filter: check PEL delivery_count for DLQ routing + let mut results = Vec::new(); + let mut dlq_entries: Vec<(StreamId, Vec<(Bytes, Bytes)>)> = Vec::new(); + let mut dlq_ack_ids: Vec = Vec::new(); + for (id, fields) in &claimed { + // Check delivery_count from PEL + let delivery_count = stream + .groups + .get(group_name.as_ref()) + .and_then(|g| g.pel.get(id)) + .map(|pe| pe.delivery_count) + .unwrap_or(1); + if mdc > 0 && delivery_count >= mdc as u64 { + // Route to DLQ + dlq_entries.push((*id, fields.clone())); + dlq_ack_ids.push(*id); + } else if results.len() < count { + results.push((*id, fields.clone())); + } + } + + // XACK DLQ entries from __mq_consumers group + if !dlq_ack_ids.is_empty() { + let _ = stream.xack(&group_name, &dlq_ack_ids); + } + + // Move DLQ entries to DLQ stream + if !dlq_entries.is_empty() { + let dlq_key = { + let mut buf = Vec::with_capacity(effective_key.len() + 8); + buf.extend_from_slice(&effective_key); + buf.extend_from_slice(b"::mq:dlq"); + Bytes::from(buf) + }; + match db_guard.get_or_create_stream(&dlq_key) { + Ok(dlq_stream) => { + for (_id, fields) in dlq_entries { + let dlq_id = dlq_stream.next_auto_id(); + dlq_stream.add(dlq_id, fields); + } + } + Err(_) => {} // DLQ creation failed -- skip silently + } + } + + // Format results as array of arrays (XREADGROUP format) + let result_frames: Vec = results + .iter() + .map(|(id, fields)| { + let mut entry_frames = Vec::with_capacity(2); + entry_frames.push(Frame::BulkString(Bytes::from(format!( + "{}-{}", + id.ms, id.seq + )))); + let field_frames: Vec = fields + .iter() + .flat_map(|(f, v)| { + vec![Frame::BulkString(f.clone()), Frame::BulkString(v.clone())] + }) + .collect(); + entry_frames.push(Frame::Array(field_frames.into())); + Frame::Array(entry_frames.into()) + }) + .collect(); + responses.push(Frame::Array(result_frames.into())); + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"ACK") { + match validate_mq_ack(cmd_args) { + Ok((queue_key, msg_ids)) => { + let effective_key = + crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); + let ids: Vec = msg_ids + .iter() + .map(|(ms, seq)| StreamId { ms: *ms, seq: *seq }) + .collect(); + let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + match db_guard.get_stream_mut(&effective_key) { + Ok(Some(stream)) => { + let group_name = Bytes::from_static(b"__mq_consumers"); + match stream.xack(&group_name, &ids) { + Ok(acked_count) => { + drop(db_guard); + // Emit MqAck WAL record for each acked ID + for (ms, seq) in &msg_ids { + let payload = + crate::mq::wal::encode_mq_ack(&effective_key, *ms, *seq); + let mut wal_buf = Vec::new(); + crate::persistence::wal_v3::record::write_wal_v3_record( + &mut wal_buf, + 0, + crate::persistence::wal_v3::record::WalRecordType::MqAck, + &payload, + ); + ctx.shard_databases + .wal_append(ctx.shard_id, Bytes::from(wal_buf)); + } + responses.push(Frame::Integer(acked_count as i64)); + } + Err(_) => responses.push(Frame::Integer(0)), + } + } + Ok(None) => responses.push(Frame::Integer(0)), + Err(_) => responses.push(Frame::Integer(0)), + } + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"DLQLEN") { + match validate_mq_dlqlen(cmd_args) { + Ok(queue_key) => { + let effective_key = + crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); + let dlq_key = { + let mut buf = Vec::with_capacity(effective_key.len() + 8); + buf.extend_from_slice(&effective_key); + buf.extend_from_slice(b"::mq:dlq"); + Bytes::from(buf) + }; + let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + let len = match db_guard.get_stream_mut(&dlq_key) { + Ok(Some(stream)) => stream.length as i64, + _ => 0i64, + }; + responses.push(Frame::Integer(len)); + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"TRIGGER") { + match validate_mq_trigger(cmd_args) { + Ok((queue_key, callback_cmd, debounce_ms)) => { + let effective_key = + crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); + let trig_key = if let Some(ws_id) = conn.workspace_id.as_ref() { + let ws_hex = ws_id.as_hex(); + let mut k = Vec::with_capacity(ws_hex.len() + 1 + queue_key.len()); + k.extend_from_slice(ws_hex.as_bytes()); + k.push(b':'); + k.extend_from_slice(&queue_key); + Bytes::from(k) + } else { + queue_key.clone() + }; + let entry = crate::mq::TriggerEntry { + queue_key: effective_key, + callback_cmd, + debounce_ms, + last_fire_ms: 0, + pending_fire_ms: 0, + }; + { + let mut guard = ctx.shard_databases.trigger_registry(ctx.shard_id); + let reg = + guard.get_or_insert_with(|| Box::new(crate::mq::TriggerRegistry::new())); + reg.register(trig_key, entry); + } + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"PUBLISH") { + match validate_mq_publish(cmd_args) { + Ok((queue_key, fields)) => { + let effective_key = + crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); + if let Some(ref mut txn) = conn.active_cross_txn { + txn.record_mq(effective_key, fields); + responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); + } else { + responses.push(Frame::Error(Bytes::from_static( + b"ERR MQ PUBLISH requires an active transaction (use TXN BEGIN first)", + ))); + } + } + Err(e) => responses.push(e), + } + return true; + } + + // Unknown MQ subcommand + responses.push(Frame::Error(Bytes::from_static(ERR_MQ_UNKNOWN_SUB))); + true +} + +/// Handle MULTI/EXEC/DISCARD commands. Returns `true` if consumed. +pub(super) fn try_handle_multi_exec( + cmd: &[u8], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + // --- MULTI --- + if cmd.eq_ignore_ascii_case(b"MULTI") { + if conn.in_cross_txn() { + responses.push(Frame::Error(Bytes::from_static(ERR_MULTI_TXN_CONFLICT))); + } else if conn.in_multi { + responses.push(Frame::Error(Bytes::from_static( + b"ERR MULTI calls can not be nested", + ))); + } else { + conn.in_multi = true; + conn.command_queue.clear(); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + return true; + } + + // --- EXEC --- + if cmd.eq_ignore_ascii_case(b"EXEC") { + if !conn.in_multi { + responses.push(Frame::Error(Bytes::from_static(b"ERR EXEC without MULTI"))); + } else { + conn.in_multi = false; + let result = execute_transaction_sharded( + &ctx.shard_databases, + ctx.shard_id, + &conn.command_queue, + conn.selected_db, + &ctx.cached_clock, + ); + conn.command_queue.clear(); + responses.push(result); + } + return true; + } + + // --- DISCARD --- + if cmd.eq_ignore_ascii_case(b"DISCARD") { + if !conn.in_multi { + responses.push(Frame::Error(Bytes::from_static( + b"ERR DISCARD without MULTI", + ))); + } else { + conn.in_multi = false; + conn.command_queue.clear(); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + return true; + } + + false +} + +/// Handle GRAPH.* graph commands. Returns `true` if consumed. +#[cfg(feature = "graph")] +pub(super) fn try_handle_graph_command( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if cmd.len() <= 6 || !cmd[..6].eq_ignore_ascii_case(b"GRAPH.") { + return false; + } + let (response, wal_records, cypher_intents, cypher_undo_ops) = + if crate::command::graph::is_graph_write_cmd(cmd) + || (cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") + && crate::command::graph::is_cypher_write_query(cmd_args)) + { + let mut gs = ctx.shard_databases.graph_store_write(ctx.shard_id); + let (resp, cypher_intents, undo_ops) = if cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") { + // Phase 167 (CYP-01/02): capture Cypher-created + // nodes/edges so TXN.ABORT can roll them back via + // CrossStoreTxn::record_graph. + crate::command::graph::graph_query_or_write(&mut gs, cmd_args) + } else { + ( + crate::command::graph::dispatch_graph_write(&mut gs, cmd, cmd_args), + Vec::new(), + Vec::new(), + ) + }; + let records = gs.drain_wal(); + (resp, records, cypher_intents, undo_ops) + } else { + let gs = ctx.shard_databases.graph_store_read(ctx.shard_id); + let resp = crate::command::graph::dispatch_graph_read(&gs, cmd, cmd_args); + (resp, Vec::new(), Vec::new(), Vec::new()) + }; + // Phase 166: record graph intent for TXN rollback. + // Captures explicit ADDNODE/ADDEDGE by response id plus + // Phase 167 Cypher CREATE/MERGE via intents returned from + // graph_query_or_write. + if let Some(txn) = conn.active_cross_txn.as_mut() { + let is_node = cmd.eq_ignore_ascii_case(b"GRAPH.ADDNODE"); + let is_edge = cmd.eq_ignore_ascii_case(b"GRAPH.ADDEDGE"); + if is_node || is_edge { + if let Frame::Integer(id) = &response { + if let Some(gname) = cmd_args.first().and_then(|f| extract_bytes(f)) { + txn.record_graph(*id as u64, is_node, gname); + } + } + } + if !cypher_intents.is_empty() { + if let Some(gname) = cmd_args.first().and_then(|f| extract_bytes(f)) { + for intent in &cypher_intents { + txn.record_graph(intent.entity_id, intent.is_node, gname.clone()); + } + } + } + // Phase 174 FIX-01: push undo ops for SET/DELETE/MERGE rollback. + for undo_op in cypher_undo_ops { + txn.record_graph_undo(undo_op); + } + } + for record in wal_records { + ctx.shard_databases + .wal_append(ctx.shard_id, bytes::Bytes::from(record)); + } + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + true +} diff --git a/src/server/conn/handler_sharded.rs b/src/server/conn/handler_sharded.rs deleted file mode 100644 index 045c310f..00000000 --- a/src/server/conn/handler_sharded.rs +++ /dev/null @@ -1,3436 +0,0 @@ -// Note: some imports/variables may be conditionally used across feature flags -//! Sharded tokio connection handlers. -//! -//! Extracted from `server/connection.rs` (Plan 48-02). -//! Contains `handle_connection_sharded` (thin wrapper) and -//! `handle_connection_sharded_inner` (generic inner handler). - -use crate::runtime::TcpStream; -use crate::runtime::cancel::CancellationToken; -use crate::runtime::channel; -use bumpalo::Bump; -use bytes::{Bytes, BytesMut}; -use ringbuf::traits::Producer; -use std::collections::HashMap; -use std::sync::Arc; - -use crate::command::connection as conn_cmd; -use crate::command::metadata; -use crate::command::mq::{ - ERR_MQ_NOT_DURABLE, ERR_MQ_UNKNOWN_SUB, parse_mq_subcommand, validate_mq_ack, - validate_mq_create, validate_mq_dlqlen, validate_mq_pop, validate_mq_publish, validate_mq_push, - validate_mq_trigger, -}; -#[cfg(feature = "graph")] -use crate::command::temporal::{ERR_ENTITY_NOT_FOUND, ERR_GRAPH_NOT_FOUND}; -use crate::command::temporal::{ - capture_wall_ms, is_temporal_invalidate, is_temporal_snapshot_at, validate_invalidate, - validate_snapshot_at, -}; -use crate::command::transaction::{ - ERR_MULTI_TXN_CONFLICT, is_txn_abort, is_txn_begin, is_txn_commit, txn_abort_validate, - txn_begin_validate, txn_commit_validate, -}; -use crate::command::workspace::{ - ERR_WS_ALREADY_BOUND, ERR_WS_NOT_FOUND, ERR_WS_UNKNOWN_SUB, parse_workspace_id_from_bytes, - parse_ws_subcommand, validate_ws_auth, validate_ws_create, validate_ws_drop, validate_ws_info, - validate_ws_list, -}; -use crate::command::{DispatchResult, dispatch, dispatch_read}; -use crate::mq::is_mq_command; -use crate::persistence::aof::{self, AofMessage}; -use crate::protocol::Frame; -use crate::shard::dispatch::{ShardMessage, key_to_shard}; -use crate::shard::mesh::ChannelMesh; -use crate::storage::eviction::try_evict_if_needed; -use crate::storage::stream::StreamId; -use crate::tracking::TrackingState; -use crate::transaction::CrossStoreTxn; -use crate::workspace::{ - WorkspaceId, is_ws_command, strip_workspace_prefix_from_response, workspace_rewrite_args, -}; - -use super::affinity::MigratedConnectionState; -use super::shared::resolve_ft_search_as_of_lsn; -use crate::server::response_slot::ResponseSlotPool; - -/// Result of `handle_connection_sharded_inner` execution. -/// -/// The generic inner handler cannot perform FD extraction (requires concrete stream type). -/// When migration is triggered, it returns `MigrateConnection` so the concrete caller -/// can extract the raw FD and send the migration message via SPSC. -pub enum HandlerResult { - /// Normal connection close (QUIT, EOF, error, shutdown). - Done, - /// AffinityTracker detected a dominant remote shard. The caller should: - /// 1. Extract the raw FD from the concrete stream (into_std + into_raw_fd) - /// 2. Send ShardMessage::MigrateConnection via SPSC to `target_shard` - /// 3. Drop the handler (connection ownership transferred) - MigrateConnection { - state: MigratedConnectionState, - target_shard: usize, - }, -} - -use super::{ - apply_resp3_conversion, convert_blocking_to_nonblocking, execute_transaction_sharded, - extract_bytes, extract_command, extract_primary_key, handle_blocking_command, handle_config, - is_multi_key_command, propagate_subscription, unpropagate_subscription, -}; - -/// Handle a single client connection on a sharded (thread-per-core) runtime. -/// -/// Runs within a shard's single-threaded Tokio runtime. Has direct mutable access -/// to the shard's databases via `Arc` (thread-safe: parking_lot RwLock -/// single-threaded scheduling means no concurrent borrows). -/// -/// Routing logic: -/// - **Keyless commands** (PING, ECHO, SELECT, etc.): execute locally, zero overhead. -/// - **Single-key, local shard**: execute directly on borrowed database -- ZERO cross-shard overhead. -/// - **Single-key, remote shard**: dispatch via SPSC `ShardMessage::Execute`, await oneshot reply. -/// - **Multi-key commands** (MGET, MSET, multi-DEL): delegate to VLL coordinator. -/// -/// Connection-level commands (AUTH, SUBSCRIBE, MULTI/EXEC) are handled at the -/// connection level same as the non-sharded handler. -#[tracing::instrument(skip_all, level = "debug")] -pub(crate) async fn handle_connection_sharded( - mut stream: TcpStream, - ctx: &super::core::ConnectionContext, - shutdown: CancellationToken, - client_id: u64, -) { - let maxclients = ctx.runtime_config.read().maxclients; - if !crate::admin::metrics_setup::try_accept_connection(maxclients) { - use tokio::io::AsyncWriteExt; - let _ = stream - .write_all(b"-ERR max number of clients reached\r\n") - .await; - return; - } - let peer_addr = stream - .peer_addr() - .map(|a| a.to_string()) - .unwrap_or_else(|_| "unknown".to_string()); - let result = handle_connection_sharded_inner( - stream, - peer_addr, - ctx, - shutdown, - client_id, - true, // can_migrate: plain TCP supports FD extraction - BytesMut::new(), - None, // fresh connection, no migrated state - ) - .await; - - // Handle migration result: extract FD from the returned stream and send via SPSC - if let ( - HandlerResult::MigrateConnection { - state, - target_shard, - }, - Some(stream), - ) = (result.0, result.1) - { - use std::os::unix::io::IntoRawFd; - match stream.into_std() { - Ok(std_stream) => { - let raw_fd = std_stream.into_raw_fd(); - let msg = ShardMessage::MigrateConnection { fd: raw_fd, state }; - let target_idx = ChannelMesh::target_index(ctx.shard_id, target_shard); - let push_result = { - let mut producers = ctx.dispatch_tx.borrow_mut(); - producers[target_idx].try_push(msg) - }; - match push_result { - Ok(()) => { - ctx.spsc_notifiers[target_shard].notify_one(); - tracing::info!( - "Shard {}: migrated connection {} to shard {}", - ctx.shard_id, - client_id, - target_shard - ); - } - Err(returned_msg) => { - // SPSC full — retry with yield before giving up. - let mut pending = Some(returned_msg); - for _ in 0..8 { - tokio::task::yield_now().await; - #[allow(clippy::unwrap_used)] - // pending is always re-filled on retry via Err(returned_msg) - let msg = pending.take().unwrap(); - let push_result = { - let mut producers = ctx.dispatch_tx.borrow_mut(); - producers[target_idx].try_push(msg) - }; - match push_result { - Ok(()) => { - ctx.spsc_notifiers[target_shard].notify_one(); - tracing::info!( - "Shard {}: migrated connection {} to shard {} (after retry)", - ctx.shard_id, - client_id, - target_shard - ); - break; - } - Err(msg) => pending = Some(msg), - } - } - if let Some(ShardMessage::MigrateConnection { fd, .. }) = pending { - use std::os::unix::io::FromRawFd; - // SAFETY: fd is a valid, uniquely-owned file descriptor obtained - // from TcpStream::into_raw_fd() above. OwnedFd closes it on drop. - drop(unsafe { std::os::unix::io::OwnedFd::from_raw_fd(fd) }); - } - tracing::warn!( - "Shard {}: migration SPSC full, connection {} lost", - ctx.shard_id, - client_id - ); - } - } - } - Err(e) => { - tracing::warn!("Shard {}: migration into_std failed: {}", ctx.shard_id, e); - // Stream consumed by into_std attempt, connection lost either way - } - } - } else { - // Only decrement connected_clients when the connection is actually closing, - // not when migrating to another shard (the connection stays alive). - crate::admin::metrics_setup::record_connection_closed(); - } -} - -/// Generic inner handler for sharded connections (Tokio runtime). -/// -/// Works with any stream implementing `AsyncRead + AsyncWrite + Unpin`, -/// enabling both plain TCP (`TcpStream`) and TLS (`tokio_rustls::server::TlsStream`). -/// -/// Returns `(HandlerResult, Option)`: the stream is returned when migration is triggered -/// so the concrete caller can extract the raw FD. `can_migrate` controls whether the -/// AffinityTracker is active (set to `false` for TLS connections). -pub(crate) async fn handle_connection_sharded_inner< - S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, ->( - stream: S, - peer_addr: String, - ctx: &super::core::ConnectionContext, - shutdown: CancellationToken, - client_id: u64, - can_migrate: bool, - initial_read_buf: BytesMut, - migrated_state: Option<&MigratedConnectionState>, -) -> (HandlerResult, Option) { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - - // Direct buffer I/O: bypass Framed/codec for the hot path. - let mut stream = stream; - let mut read_buf = if initial_read_buf.is_empty() { - BytesMut::with_capacity(8192) - } else { - // Migration buffer: leftover bytes from the source connection. - let mut buf = initial_read_buf; - buf.reserve(8192); - buf - }; - let mut write_buf = BytesMut::with_capacity(8192); - let parse_config = crate::protocol::ParseConfig::default(); - let mut conn = super::core::ConnectionState::new( - client_id, - peer_addr.clone(), - &ctx.requirepass, - ctx.shard_id, - ctx.num_shards, - can_migrate, - ctx.runtime_config.read().acllog_max_len, - migrated_state, - ); - conn.refresh_acl_cache(&ctx.acl_table); - - // Register in global client registry for CLIENT LIST/INFO/KILL. - // RegistryGuard ensures deregister on all exit paths (including early returns). - crate::client_registry::register( - client_id, - peer_addr.clone(), - conn.current_user.clone(), - ctx.shard_id, - ); - struct RegistryGuard(u64); - impl Drop for RegistryGuard { - fn drop(&mut self) { - crate::client_registry::deregister(self.0); - } - } - let _registry_guard = RegistryGuard(client_id); - - use crate::pubsub::{self, subscriber::Subscriber}; - - // Functions API registry (per-shard, lazy init) — kept as local because Rc> is !Send - let func_registry = std::rc::Rc::new(std::cell::RefCell::new( - crate::scripting::FunctionRegistry::new(), - )); - - // Per-connection arena for batch processing temporaries. - // 4KB initial capacity, grows on demand (rarely exceeds 16KB per batch). - let mut arena = Bump::with_capacity(4096); - - // Pre-allocated response slots for zero-allocation cross-shard dispatch. - let response_pool = ResponseSlotPool::new(ctx.num_shards, ctx.shard_id); - - // Client idle timeout: 0 = disabled (read once, avoid lock on hot path) - let idle_timeout_secs = ctx.runtime_config.read().timeout; - let idle_timeout = if idle_timeout_secs > 0 { - Some(std::time::Duration::from_secs(idle_timeout_secs)) - } else { - None - }; - - let mut break_outer = false; - loop { - // Check if CLIENT KILL targeted this connection - if crate::client_registry::is_killed(client_id) { - break; - } - - // --- Subscriber mode: bidirectional select on client commands + published messages --- - if conn.subscription_count > 0 { - #[allow(clippy::unwrap_used)] - // conn.pubsub_rx is always Some when conn.subscription_count > 0 - let rx = conn.pubsub_rx.as_mut().unwrap(); - tokio::select! { - n = stream.read_buf(&mut read_buf) => { - match n { - Ok(0) | Err(_) => break, - Ok(_) => {} - } - let mut sub_break = false; - loop { - match crate::protocol::parse(&mut read_buf, &parse_config) { - Ok(Some(frame)) => { - if let Some((cmd, cmd_args)) = extract_command(&frame) { - if cmd.eq_ignore_ascii_case(b"SUBSCRIBE") { - if cmd_args.is_empty() { - let err = Frame::Error(Bytes::from_static(b"ERR wrong number of arguments for 'subscribe' command")); - write_buf.clear(); - crate::protocol::serialize(&err, &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - continue; - } - for arg in cmd_args { - if let Some(ch) = extract_bytes(arg) { - #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable - let acl_deny = { ctx.acl_table.read().unwrap().check_channel_permission(&conn.current_user, ch.as_ref()) }; - if let Some(reason) = acl_deny { - let err = Frame::Error(Bytes::from(format!("NOPERM {}", reason))); - write_buf.clear(); - crate::protocol::serialize(&err, &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - continue; - } - #[allow(clippy::unwrap_used)] // conn.pubsub_tx is always Some in subscriber mode - let sub = Subscriber::with_protocol( - conn.pubsub_tx.clone().unwrap(), - conn.subscriber_id, - conn.protocol_version >= 3, - ); - { ctx.pubsub_registry.write().subscribe(ch.clone(), sub); } - conn.subscription_count += 1; - // Register pub/sub affinity for this client IP - if conn.subscription_count == 1 { - if let Ok(addr) = peer_addr.parse::() { - ctx.pubsub_affinity.write().register(addr.ip(), ctx.shard_id); - } - } - propagate_subscription(&ctx.all_remote_sub_maps, &ch, ctx.shard_id, ctx.num_shards, false); - write_buf.clear(); - let resp = pubsub::subscribe_response(&ch, conn.subscription_count); - crate::protocol::serialize(&resp, &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - } - } - } else if cmd.eq_ignore_ascii_case(b"PSUBSCRIBE") { - if cmd_args.is_empty() { - let err = Frame::Error(Bytes::from_static(b"ERR wrong number of arguments for 'psubscribe' command")); - write_buf.clear(); - crate::protocol::serialize(&err, &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - continue; - } - for arg in cmd_args { - if let Some(pat) = extract_bytes(arg) { - #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable - let acl_deny = { ctx.acl_table.read().unwrap().check_channel_permission(&conn.current_user, pat.as_ref()) }; - if let Some(reason) = acl_deny { - let err = Frame::Error(Bytes::from(format!("NOPERM {}", reason))); - write_buf.clear(); - crate::protocol::serialize(&err, &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - continue; - } - #[allow(clippy::unwrap_used)] // conn.pubsub_tx is always Some in subscriber mode - let sub = Subscriber::with_protocol( - conn.pubsub_tx.clone().unwrap(), - conn.subscriber_id, - conn.protocol_version >= 3, - ); - { ctx.pubsub_registry.write().psubscribe(pat.clone(), sub); } - conn.subscription_count += 1; - // Register pub/sub affinity for this client IP - if conn.subscription_count == 1 { - if let Ok(addr) = peer_addr.parse::() { - ctx.pubsub_affinity.write().register(addr.ip(), ctx.shard_id); - } - } - propagate_subscription(&ctx.all_remote_sub_maps, &pat, ctx.shard_id, ctx.num_shards, true); - write_buf.clear(); - let resp = pubsub::psubscribe_response(&pat, conn.subscription_count); - crate::protocol::serialize(&resp, &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - } - } - } else if cmd.eq_ignore_ascii_case(b"UNSUBSCRIBE") { - if cmd_args.is_empty() { - let removed = { ctx.pubsub_registry.write().unsubscribe_all(conn.subscriber_id) }; - if removed.is_empty() { - conn.subscription_count = ctx.pubsub_registry.read().total_subscription_count(conn.subscriber_id); - write_buf.clear(); - crate::protocol::serialize(&pubsub::unsubscribe_response(&Bytes::from_static(b""), conn.subscription_count), &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - } else { - for ch in &removed { - conn.subscription_count = conn.subscription_count.saturating_sub(1); - unpropagate_subscription(&ctx.all_remote_sub_maps, ch, ctx.shard_id, ctx.num_shards, false); - write_buf.clear(); - crate::protocol::serialize(&pubsub::unsubscribe_response(ch, conn.subscription_count), &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - } - } - } else { - for arg in cmd_args { - if let Some(ch) = extract_bytes(arg) { - { ctx.pubsub_registry.write().unsubscribe(ch.as_ref(), conn.subscriber_id); } - conn.subscription_count = conn.subscription_count.saturating_sub(1); - unpropagate_subscription(&ctx.all_remote_sub_maps, &ch, ctx.shard_id, ctx.num_shards, false); - write_buf.clear(); - crate::protocol::serialize(&pubsub::unsubscribe_response(&ch, conn.subscription_count), &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - } - } - } - } else if cmd.eq_ignore_ascii_case(b"PUNSUBSCRIBE") { - if cmd_args.is_empty() { - let removed = { ctx.pubsub_registry.write().punsubscribe_all(conn.subscriber_id) }; - if removed.is_empty() { - conn.subscription_count = ctx.pubsub_registry.read().total_subscription_count(conn.subscriber_id); - write_buf.clear(); - crate::protocol::serialize(&pubsub::punsubscribe_response(&Bytes::from_static(b""), conn.subscription_count), &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - } else { - for pat in &removed { - conn.subscription_count = conn.subscription_count.saturating_sub(1); - unpropagate_subscription(&ctx.all_remote_sub_maps, pat, ctx.shard_id, ctx.num_shards, true); - write_buf.clear(); - crate::protocol::serialize(&pubsub::punsubscribe_response(pat, conn.subscription_count), &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - } - } - } else { - for arg in cmd_args { - if let Some(pat) = extract_bytes(arg) { - { ctx.pubsub_registry.write().punsubscribe(pat.as_ref(), conn.subscriber_id); } - conn.subscription_count = conn.subscription_count.saturating_sub(1); - unpropagate_subscription(&ctx.all_remote_sub_maps, &pat, ctx.shard_id, ctx.num_shards, true); - write_buf.clear(); - crate::protocol::serialize(&pubsub::punsubscribe_response(&pat, conn.subscription_count), &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - } - } - } - } else if cmd.eq_ignore_ascii_case(b"PING") { - write_buf.clear(); - let resp = Frame::Array(crate::framevec![ - Frame::BulkString(Bytes::from_static(b"pong")), - Frame::BulkString(Bytes::from_static(b"")), - ]); - crate::protocol::serialize(&resp, &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - } else if cmd.eq_ignore_ascii_case(b"QUIT") { - write_buf.clear(); - crate::protocol::serialize(&Frame::SimpleString(Bytes::from_static(b"OK")), &mut write_buf); - let _ = stream.write_all(&write_buf).await; - sub_break = true; - break; - } else if cmd.eq_ignore_ascii_case(b"RESET") { - { ctx.pubsub_registry.write().unsubscribe_all(conn.subscriber_id); } - { ctx.pubsub_registry.write().punsubscribe_all(conn.subscriber_id); } - conn.subscription_count = 0; - write_buf.clear(); - crate::protocol::serialize(&Frame::SimpleString(Bytes::from_static(b"RESET")), &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - } else { - let cmd_str = String::from_utf8_lossy(cmd); - let err = Frame::Error(Bytes::from(format!( - "ERR Can't execute '{}': only (P)SUBSCRIBE / (P)UNSUBSCRIBE / PING / QUIT are allowed in this context", - cmd_str.to_lowercase() - ))); - write_buf.clear(); - crate::protocol::serialize(&err, &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { sub_break = true; break; } - } - } - if conn.subscription_count == 0 { break; } - } - Ok(None) => break, - Err(_) => { return (HandlerResult::Done, None); } - } - } - if sub_break { break; } - if conn.subscription_count == 0 { continue; } - } - msg = rx.recv_async() => { - match msg { - Ok(data) => { - // Data is pre-serialized RESP bytes — write directly - if stream.write_all(&data).await.is_err() { break; } - } - Err(_) => break, - } - } - _ = shutdown.cancelled() => { - write_buf.clear(); - crate::protocol::serialize(&Frame::Error(Bytes::from_static(b"ERR server shutting down")), &mut write_buf); - let _ = stream.write_all(&write_buf).await; - break; - } - } - continue; - } - tokio::select! { - result = async { - if let Some(dur) = idle_timeout { - match tokio::time::timeout(dur, stream.read_buf(&mut read_buf)).await { - Ok(r) => r, - Err(_) => Err(std::io::Error::new(std::io::ErrorKind::TimedOut, "idle timeout")), - } - } else { - stream.read_buf(&mut read_buf).await - } - } => { - match result { - Ok(0) => break, // connection closed - Ok(_) => {} - Err(e) if e.kind() == std::io::ErrorKind::TimedOut => { - tracing::debug!("Connection {} idle timeout ({}s)", client_id, idle_timeout_secs); - break; - } - Err(_) => break, - } - - // Parse all complete frames from buffer - let mut batch: Vec = Vec::with_capacity(64); - const MAX_BATCH: usize = 1024; - loop { - match crate::protocol::parse(&mut read_buf, &parse_config) { - Ok(Some(frame)) => { - batch.push(frame); - if batch.len() >= MAX_BATCH { break; } - } - Ok(None) => break, - Err(crate::protocol::ParseError::Incomplete) => break, - Err(_) => { break_outer = true; break; } - } - } - if break_outer { break; } - if batch.is_empty() { continue; } - - // CLIENT PAUSE: delay processing if server is paused - // Check with is_write=true (conservative — pauses all batches in ALL mode) - crate::client_pause::expire_if_needed(); - if let Some(remaining) = crate::client_pause::check_pause(true) { - tokio::time::sleep(remaining).await; - } - - let mut responses: Vec = Vec::with_capacity(batch.len()); - let mut should_quit = false; - let mut remote_groups: HashMap, Option, Bytes, usize)>> = HashMap::with_capacity(ctx.num_shards); - // Accumulate cross-shard PUBLISH pairs per target shard for batch dispatch - // Key: target shard ID -> Vec of (response_index, channel, message) - let mut publish_batches: HashMap> = HashMap::new(); - - // Track if AUTH rate limiting delay is needed (applied after batch response) - let mut auth_delay_ms: u64 = 0; - - for frame in batch { - // --- AUTH gate --- - if !conn.authenticated { - match extract_command(&frame) { - Some((cmd, cmd_args)) if cmd.eq_ignore_ascii_case(b"AUTH") => { - let (response, opt_user) = conn_cmd::auth_acl(cmd_args, &ctx.acl_table); - if let Some(uname) = opt_user { - conn.authenticated = true; - conn.current_user = uname; - conn.refresh_acl_cache(&ctx.acl_table); - if let Ok(addr) = peer_addr.parse::() { - crate::auth_ratelimit::record_success(addr.ip()); - } - } else { - if let Ok(addr) = peer_addr.parse::() { - auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); - } - conn.acl_log.push(crate::acl::AclLogEntry { - reason: "auth".to_string(), - object: "AUTH".to_string(), - username: conn.current_user.clone(), - client_addr: peer_addr.clone(), - timestamp_ms: std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as u64, - }); - } - responses.push(response); - continue; - } - Some((cmd, cmd_args)) if cmd.eq_ignore_ascii_case(b"HELLO") => { - let (response, new_proto, new_name, opt_user) = conn_cmd::hello_acl( - cmd_args, - conn.protocol_version, - client_id, - &ctx.acl_table, - &mut conn.authenticated, - ); - if !matches!(&response, Frame::Error(_)) { - conn.protocol_version = new_proto; - } - if let Some(name) = new_name { - conn.client_name = Some(name); - } - if let Some(ref uname) = opt_user { - conn.current_user = uname.clone(); - conn.refresh_acl_cache(&ctx.acl_table); - } - // HELLO AUTH rate limiting (same as AUTH gate) - if matches!(&response, Frame::Error(_)) { - if let Ok(addr) = peer_addr.parse::() { - auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); - } - } else if opt_user.is_some() { - if let Ok(addr) = peer_addr.parse::() { - crate::auth_ratelimit::record_success(addr.ip()); - } - } - responses.push(response); - continue; - } - Some((cmd, _)) if cmd.eq_ignore_ascii_case(b"QUIT") => { - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - should_quit = true; - break; - } - _ => { - responses.push(Frame::Error( - Bytes::from_static(b"NOAUTH Authentication required.") - )); - continue; - } - } - } - - let (cmd, cmd_args) = match extract_command(&frame) { - Some(pair) => pair, - None => { - responses.push(Frame::Error(Bytes::from_static( - b"ERR invalid command format", - ))); - continue; - } - }; - - // --- QUIT --- - if cmd.eq_ignore_ascii_case(b"QUIT") { - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - should_quit = true; - break; - } - - // --- ASKING --- - if cmd.eq_ignore_ascii_case(b"ASKING") { - conn.asking = true; - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - continue; - } - - // --- CLUSTER subcommands --- - if cmd.eq_ignore_ascii_case(b"CLUSTER") { - if let Some(ref cs) = ctx.cluster_state { - #[allow(clippy::unwrap_used)] // Fallback "127.0.0.1:6379" is a valid literal - let self_addr: std::net::SocketAddr = - format!("127.0.0.1:{}", ctx.config_port) - .parse() - .unwrap_or_else(|_| "127.0.0.1:6379".parse().unwrap()); - let resp = crate::cluster::command::handle_cluster_command( - cmd_args, cs, self_addr, - ); - responses.push(resp); - } else { - responses.push(Frame::Error(Bytes::from_static( - b"ERR This instance has cluster support disabled", - ))); - } - continue; - } - - // --- Lua scripting: EVAL / EVALSHA --- - if cmd.eq_ignore_ascii_case(b"EVAL") || cmd.eq_ignore_ascii_case(b"EVALSHA") { - let response = { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let db_count = ctx.shard_databases.db_count(); - if cmd.eq_ignore_ascii_case(b"EVAL") { - crate::scripting::handle_eval( - &ctx.lua, &ctx.script_cache, cmd_args, &mut guard, - ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, - ) - } else { - crate::scripting::handle_evalsha( - &ctx.lua, &ctx.script_cache, cmd_args, &mut guard, - ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, - ) - } - }; - responses.push(response); - continue; - } - - // --- SCRIPT subcommands --- - if cmd.eq_ignore_ascii_case(b"SCRIPT") { - let (response, fanout) = crate::scripting::handle_script_subcommand(&ctx.script_cache, cmd_args); - if let Some((sha1, script_bytes)) = fanout { - let mut producers = ctx.dispatch_tx.borrow_mut(); - for target in 0..ctx.num_shards { - if target == ctx.shard_id { continue; } - let idx = ChannelMesh::target_index(ctx.shard_id, target); - let msg = ShardMessage::ScriptLoad { sha1: sha1.clone(), script: script_bytes.clone() }; - if producers[idx].try_push(msg).is_ok() { - ctx.spsc_notifiers[target].notify_one(); - } - } - drop(producers); - } - responses.push(response); - continue; - } - - // --- Cluster slot routing (pre-dispatch) --- - if crate::cluster::cluster_enabled() { - if let Some(ref cs) = ctx.cluster_state { - let was_asking = conn.asking; - conn.asking = false; - let maybe_key = extract_primary_key(cmd, cmd_args); - if let Some(key) = maybe_key { - let slot = crate::cluster::slots::slot_for_key(key); - #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable - let route = cs.read().unwrap().route_slot(slot, was_asking); - match route { - crate::cluster::SlotRoute::Local => {} - other => { - responses.push(other.into_error_frame(slot)); - continue; - } - } - if is_multi_key_command(cmd, cmd_args) { - let first_slot = slot; - let mut cross_slot = false; - for arg in cmd_args.iter().skip(1) { - if let Some(k) = match arg { - Frame::BulkString(b) => Some(b.as_ref()), - _ => None, - } { - if crate::cluster::slots::slot_for_key(k) != first_slot { - cross_slot = true; - break; - } - } - } - if cross_slot { - responses.push(Frame::Error(Bytes::from_static( - b"CROSSSLOT Keys in request don't hash to the same slot", - ))); - continue; - } - } - } - } - } - - // --- AUTH (already conn.authenticated) --- - if cmd.eq_ignore_ascii_case(b"AUTH") { - let (response, opt_user) = conn_cmd::auth_acl(cmd_args, &ctx.acl_table); - if let Some(uname) = opt_user { - conn.current_user = uname; - conn.refresh_acl_cache(&ctx.acl_table); - if let Ok(addr) = peer_addr.parse::() { - crate::auth_ratelimit::record_success(addr.ip()); - } - } else if let Ok(addr) = peer_addr.parse::() { - auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); - } - responses.push(response); - continue; - } - - // --- HELLO --- - if cmd.eq_ignore_ascii_case(b"HELLO") { - let (response, new_proto, new_name, opt_user) = conn_cmd::hello_acl( - cmd_args, conn.protocol_version, client_id, &ctx.acl_table, &mut conn.authenticated, - ); - if !matches!(&response, Frame::Error(_)) { conn.protocol_version = new_proto; } - if let Some(name) = new_name { conn.client_name = Some(name); } - if let Some(ref uname) = opt_user { - conn.current_user = uname.clone(); - conn.refresh_acl_cache(&ctx.acl_table); - } - if matches!(&response, Frame::Error(_)) { - if let Ok(addr) = peer_addr.parse::() { - auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); - } - } else if opt_user.is_some() { - if let Ok(addr) = peer_addr.parse::() { - crate::auth_ratelimit::record_success(addr.ip()); - } - } - responses.push(response); - continue; - } - - // --- ACL --- - if cmd.eq_ignore_ascii_case(b"ACL") { - let response = crate::command::acl::handle_acl( - cmd_args, &ctx.acl_table, &mut conn.acl_log, &conn.current_user, &peer_addr, &ctx.runtime_config, - ); - responses.push(response); - continue; - } - - // === CLIENT PAUSE check === - // Extract pause info with short lock hold, then sleep outside lock scope - let pause_wait_ms = { - let rt = ctx.runtime_config.read(); - let deadline = rt.client_pause_deadline_ms; - if deadline > 0 { - let now = crate::storage::entry::current_time_ms(); - if now < deadline { - let should_pause = if rt.client_pause_write_only { - crate::command::metadata::is_write(cmd) - } else { - true - }; - if should_pause { deadline.saturating_sub(now) } else { 0 } - } else { 0 } - } else { 0 } - }; - if pause_wait_ms > 0 { - // Poll in 50ms intervals so CLIENT UNPAUSE takes effect quickly - let mut remaining = pause_wait_ms; - while remaining > 0 { - let chunk = remaining.min(50); - #[cfg(feature = "runtime-tokio")] - { - tokio::time::sleep(std::time::Duration::from_millis(chunk)).await; - } - #[cfg(feature = "runtime-monoio")] - { - monoio::time::sleep(std::time::Duration::from_millis(chunk)).await; - } - remaining = remaining.saturating_sub(chunk); - // Re-check if UNPAUSE was called - let still_paused = { - let rt = ctx.runtime_config.read(); - rt.client_pause_deadline_ms > 0 - && crate::storage::entry::current_time_ms() < rt.client_pause_deadline_ms - }; - if !still_paused { - break; - } - } - } - - // === ACL permission check === - // Must run before any command-specific handlers (CONFIG, REPLICAOF, etc.) - // so that low-privilege users cannot reach admin commands. - // Fast path: skip RwLock + HashMap for unrestricted users - // with a fresh cache. Stale caches (after ACL SETUSER / - // DELUSER / LOAD) fall through to the full check. - if !conn.acl_skip_allowed() { - #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable - let acl_guard = ctx.acl_table.read().unwrap(); - if let Some(deny_reason) = acl_guard.check_command_permission(&conn.current_user, cmd, cmd_args) { - drop(acl_guard); - conn.acl_log.push(crate::acl::AclLogEntry { - reason: "command".to_string(), - object: String::from_utf8_lossy(cmd).to_ascii_lowercase(), - username: conn.current_user.clone(), - client_addr: peer_addr.clone(), - timestamp_ms: std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_millis() as u64, - }); - responses.push(Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason)))); - continue; - } - let is_write_for_acl = metadata::is_write(cmd); - if let Some(deny_reason) = acl_guard.check_key_permission(&conn.current_user, cmd, cmd_args, is_write_for_acl) { - drop(acl_guard); - conn.acl_log.push(crate::acl::AclLogEntry { - reason: "command".to_string(), - object: String::from_utf8_lossy(cmd).to_ascii_lowercase(), - username: conn.current_user.clone(), - client_addr: peer_addr.clone(), - timestamp_ms: std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_millis() as u64, - }); - responses.push(Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason)))); - continue; - } - } - - // --- Functions API: FUNCTION/FCALL/FCALL_RO --- - // Placed AFTER ACL check. Respects MULTI queue — if conn.in_multi, - // fall through to the MULTI queue gate instead of executing. - if !conn.in_multi { - if cmd.eq_ignore_ascii_case(b"FUNCTION") { - let response = crate::command::functions::handle_function( - &mut func_registry.borrow_mut(), cmd_args, - ); - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"FCALL") { - let response = { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let db_count = ctx.shard_databases.db_count(); - crate::command::functions::handle_fcall( - &func_registry.borrow(), cmd_args, &mut guard, - ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, - ) - }; - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"FCALL_RO") { - let response = { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let db_count = ctx.shard_databases.db_count(); - crate::command::functions::handle_fcall_ro( - &func_registry.borrow(), cmd_args, &mut guard, - ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, - ) - }; - responses.push(response); - continue; - } - } - - // --- CONFIG --- - if cmd.eq_ignore_ascii_case(b"CONFIG") { - responses.push(handle_config(cmd_args, &ctx.runtime_config, &ctx.config)); - continue; - } - - // --- SLOWLOG --- - if cmd.eq_ignore_ascii_case(b"SLOWLOG") { - let sl = crate::admin::metrics_setup::global_slowlog(); - responses.push(crate::admin::slowlog::handle_slowlog(sl, cmd_args)); - continue; - } - - // --- REPLICAOF / SLAVEOF --- - if cmd.eq_ignore_ascii_case(b"REPLICAOF") || cmd.eq_ignore_ascii_case(b"SLAVEOF") { - use crate::command::connection::{replicaof, ReplicaofAction}; - let (resp, action) = replicaof(cmd_args); - if let Some(action) = action { - if let Some(ref rs) = ctx.repl_state { - match action { - ReplicaofAction::StartReplication { host, port } => { - if let Ok(mut rs_guard) = rs.write() { - rs_guard.role = crate::replication::state::ReplicationRole::Replica { - host: host.clone(), port, - state: crate::replication::handshake::ReplicaHandshakeState::PingPending, - }; - } - let rs_clone = Arc::clone(rs); - let cfg = crate::replication::replica::ReplicaTaskConfig { - master_host: host, master_port: port, - repl_state: rs_clone, num_shards: ctx.num_shards, - persistence_dir: None, listening_port: 0, - }; - tokio::task::spawn_local(crate::replication::replica::run_replica_task(cfg)); - } - ReplicaofAction::PromoteToMaster => { - use crate::replication::state::generate_repl_id; - if let Ok(mut rs_guard) = rs.write() { - rs_guard.repl_id2 = rs_guard.repl_id.clone(); - rs_guard.repl_id = generate_repl_id(); - rs_guard.role = crate::replication::state::ReplicationRole::Master; - } - } - ReplicaofAction::NoOp => {} - } - } - } - responses.push(resp); - continue; - } - - // --- REPLCONF --- - if cmd.eq_ignore_ascii_case(b"REPLCONF") { - responses.push(crate::command::connection::replconf(cmd_args)); - continue; - } - - // --- INFO --- - if cmd.eq_ignore_ascii_case(b"INFO") { - let guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); - let response_text = { - let resp_frame = conn_cmd::info_readonly(&guard, cmd_args); - match resp_frame { - Frame::BulkString(b) => String::from_utf8_lossy(&b).to_string(), - _ => String::new(), - } - }; - drop(guard); - let mut response_text = response_text; - if let Some(ref rs) = ctx.repl_state { - if let Ok(rs_guard) = rs.try_read() { - response_text.push_str( - &crate::replication::handshake::build_info_replication(&rs_guard), - ); - } - } - responses.push(Frame::BulkString(Bytes::from(response_text))); - continue; - } - - // --- READONLY enforcement --- - if let Some(ref rs) = ctx.repl_state { - if let Ok(rs_guard) = rs.try_read() { - if matches!(rs_guard.role, crate::replication::state::ReplicationRole::Replica { .. }) { - if metadata::is_write(cmd) { - responses.push(Frame::Error(Bytes::from_static( - b"READONLY You can't write against a read only replica.", - ))); - continue; - } - } - } - } - - // --- CLIENT subcommands --- - if cmd.eq_ignore_ascii_case(b"CLIENT") { - if let Some(sub) = cmd_args.first() { - if let Some(sub_bytes) = extract_bytes(sub) { - if sub_bytes.eq_ignore_ascii_case(b"ID") { - responses.push(conn_cmd::client_id(client_id)); - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"SETNAME") { - if cmd_args.len() != 2 { - responses.push(Frame::Error(Bytes::from_static( - b"ERR wrong number of arguments for 'CLIENT SETNAME' command", - ))); - } else { - conn.client_name = extract_bytes(&cmd_args[1]); - let name_str = conn.client_name.as_ref().map(|b| String::from_utf8_lossy(b).to_string()); - crate::client_registry::update(client_id, |e| { e.name = name_str; }); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"GETNAME") { - responses.push(match &conn.client_name { - Some(name) => Frame::BulkString(name.clone()), - None => Frame::Null, - }); - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"TRACKING") { - match crate::command::client::parse_tracking_args(cmd_args) { - Ok(config_parsed) => { - if config_parsed.enable { - conn.tracking_state.enabled = true; - conn.tracking_state.bcast = config_parsed.bcast; - conn.tracking_state.noloop = config_parsed.noloop; - conn.tracking_state.optin = config_parsed.optin; - conn.tracking_state.optout = config_parsed.optout; - if conn.tracking_rx.is_none() { - let (tx, rx) = channel::mpsc_bounded::(256); - conn.tracking_state.invalidation_tx = Some(tx.clone()); - conn.tracking_rx = Some(rx); - let mut table = ctx.tracking_table.borrow_mut(); - table.register_client(client_id, tx); - if let Some(target) = config_parsed.redirect { - table.set_redirect(client_id, target); - } - for prefix in &config_parsed.prefixes { - table.register_prefix(client_id, prefix.clone(), config_parsed.noloop); - } - } - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } else { - conn.tracking_state = TrackingState::default(); - ctx.tracking_table.borrow_mut().untrack_all(client_id); - conn.tracking_rx = None; - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - continue; - } - Err(err_frame) => { responses.push(err_frame); continue; } - } - } - if sub_bytes.eq_ignore_ascii_case(b"LIST") { - // Update our own entry before listing - crate::client_registry::update(client_id, |e| { - e.db = conn.selected_db; - e.last_cmd_at = std::time::Instant::now(); - e.flags = crate::client_registry::ClientFlags { - subscriber: conn.subscription_count > 0, - in_multi: conn.in_multi, - blocked: false, - }; - }); - let list = crate::client_registry::client_list(); - responses.push(Frame::BulkString(Bytes::from(list))); - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"INFO") { - crate::client_registry::update(client_id, |e| { - e.db = conn.selected_db; - e.last_cmd_at = std::time::Instant::now(); - }); - let info = crate::client_registry::client_info(client_id) - .unwrap_or_default(); - responses.push(Frame::BulkString(Bytes::from(info))); - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"KILL") { - let raw_args: Vec<&[u8]> = cmd_args[1..].iter() - .filter_map(|f| match f { - Frame::BulkString(b) => Some(b.as_ref()), - Frame::SimpleString(b) => Some(b.as_ref()), - _ => None, - }) - .collect(); - match crate::client_registry::parse_kill_args(&raw_args) { - Some(filter) => { - let count = crate::client_registry::kill_clients(&filter); - responses.push(Frame::Integer(count as i64)); - } - None => { - responses.push(Frame::Error(Bytes::from_static( - b"ERR syntax error. Usage: CLIENT KILL [ID id] [ADDR addr] [USER user]", - ))); - } - } - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"PAUSE") { - if cmd_args.len() < 2 { - responses.push(Frame::Error(Bytes::from_static( - b"ERR wrong number of arguments for 'CLIENT PAUSE' command", - ))); - } else { - let timeout_bytes = match &cmd_args[1] { - Frame::BulkString(b) => Some(b.as_ref()), - Frame::SimpleString(b) => Some(b.as_ref()), - _ => None, - }; - match timeout_bytes.and_then(|b| std::str::from_utf8(b).ok()).and_then(|s| s.parse::().ok()) { - Some(ms) => { - let mode = if cmd_args.len() > 2 { - match &cmd_args[2] { - Frame::BulkString(b) | Frame::SimpleString(b) if b.eq_ignore_ascii_case(b"WRITE") => crate::client_pause::PauseMode::Write, - _ => crate::client_pause::PauseMode::All, - } - } else { - crate::client_pause::PauseMode::All - }; - crate::client_pause::pause(ms, mode); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - None => { - responses.push(Frame::Error(Bytes::from_static( - b"ERR timeout is not a valid integer or out of range", - ))); - } - } - } - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"UNPAUSE") { - crate::client_pause::unpause(); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - continue; - } - if sub_bytes.eq_ignore_ascii_case(b"NO-EVICT") || sub_bytes.eq_ignore_ascii_case(b"NO-TOUCH") { - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - continue; - } - responses.push(Frame::Error(Bytes::from(format!( - "ERR unknown subcommand '{}'", String::from_utf8_lossy(&sub_bytes) - )))); - continue; - } - } - responses.push(Frame::Error(Bytes::from_static( - b"ERR wrong number of arguments for 'client' command", - ))); - continue; - } - - // --- TXN.BEGIN --- - if is_txn_begin(cmd, cmd_args) { - match txn_begin_validate(conn.in_multi, conn.in_cross_txn()) { - Ok(()) => { - // Get next txn_id and snapshot_lsn from vector store's transaction manager - let mut vector_store = - ctx.shard_databases.vector_store(ctx.shard_id); - let active = vector_store.txn_manager_mut().begin(); - conn.active_cross_txn = - Some(CrossStoreTxn::new(active.txn_id, active.snapshot_lsn)); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - Err(e) => responses.push(e), - } - continue; - } - - // --- TXN.COMMIT --- - if is_txn_commit(cmd, cmd_args) { - match txn_commit_validate(conn.in_cross_txn()) { - Ok(()) => { - if let Some(txn) = conn.active_cross_txn.take() { - let mut vector_store = - ctx.shard_databases.vector_store(ctx.shard_id); - vector_store.txn_manager_mut().commit(txn.txn_id); - drop(vector_store); - - // Write XactCommit WAL record with committed KV state - let txn_id = txn.txn_id; - if !txn.kv_undo.is_empty() { - let db_guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); - let payload = crate::persistence::wal_v3::record::encode_xact_commit_payload( - txn_id, - txn.kv_undo.records(), - &*db_guard, - ); - drop(db_guard); - let mut wal_buf = Vec::new(); - crate::persistence::wal_v3::record::write_wal_v3_record( - &mut wal_buf, - txn_id, - crate::persistence::wal_v3::record::WalRecordType::XactCommit, - &payload, - ); - ctx.shard_databases.wal_append(ctx.shard_id, bytes::Bytes::from(wal_buf)); - } - - // Release KV write intents from shard side-table - ctx.shard_databases.kv_intents(ctx.shard_id).release_txn(txn_id); - - // Drain deferred HNSW inserts (post-commit hook). - // The drain prevents phantom neighbors on abort. - // Actual HNSW graph insertion happens during compaction, - // not at commit time (point is already in mutable segment). - let drain_count = ctx.shard_databases - .hnsw_queue(ctx.shard_id) - .drain_for_txn(txn_id) - .count(); - if drain_count > 0 { - tracing::debug!(txn_id, count = drain_count, "Drained deferred HNSW inserts"); - } - - // Materialize MQ intents: enqueue deferred MQ.PUBLISH messages - if !txn.mq_intents.is_empty() { - let mut db_guard = ctx.shard_databases.write_db( - ctx.shard_id, conn.selected_db, - ); - for intent in &txn.mq_intents { - if let Ok(Some(stream)) = db_guard.get_stream_mut(&intent.queue_key) { - if stream.durable { - let msg_id = stream.next_auto_id(); - stream.add(msg_id, intent.fields.clone()); - } - } - } - } - - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } else { - responses.push(Frame::Error(Bytes::from_static( - b"ERR not in transaction", - ))); - } - } - Err(e) => responses.push(e), - } - continue; - } - - // --- TXN.ABORT --- - if is_txn_abort(cmd, cmd_args) { - match txn_abort_validate(conn.in_cross_txn()) { - Ok(()) => { - if let Some(txn) = conn.active_cross_txn.take() { - // Shared rollback (Phase 166 Plan 03): - // KV undo -> graph intents reverse -> vector - // tombstone -> side-table release. See - // src/transaction/abort.rs for lock ordering. - crate::transaction::abort::abort_cross_store_txn( - &ctx.shard_databases, - ctx.shard_id, - conn.selected_db, - txn, - ); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } else { - responses.push(Frame::Error(Bytes::from_static( - b"ERR not in transaction", - ))); - } - } - Err(e) => responses.push(e), - } - continue; - } - - // --- TEMPORAL.SNAPSHOT_AT --- - if is_temporal_snapshot_at(cmd) { - match validate_snapshot_at(cmd_args) { - Ok(()) => { - let wall_ms = capture_wall_ms(); - let lsn = { - let vector_store = ctx.shard_databases.vector_store(ctx.shard_id); - vector_store.txn_manager().current_lsn() - }; - { - let mut guard = - ctx.shard_databases.temporal_registry(ctx.shard_id); - let registry = guard.get_or_insert_with(|| { - Box::new(crate::temporal::TemporalRegistry::new()) - }); - registry.record(wall_ms, lsn); - } - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - Err(e) => responses.push(e), - } - continue; - } - - // --- TEMPORAL.INVALIDATE --- - if is_temporal_invalidate(cmd) { - match validate_invalidate(cmd_args) { - Ok((entity_id, is_node, graph_name)) => { - let wall_ms = capture_wall_ms(); - #[cfg(feature = "graph")] - { - let mut gs = - ctx.shard_databases.graph_store_write(ctx.shard_id); - if let Some(named_graph) = gs.get_graph_mut(&graph_name) { - let mutated = if is_node { - let node_key: crate::graph::types::NodeKey = - slotmap::KeyData::from_ffi(entity_id).into(); - if let Some(node) = - named_graph.write_buf.get_node_mut(node_key) - { - node.valid_to = wall_ms; - true - } else { - false - } - } else { - let edge_key: crate::graph::types::EdgeKey = - slotmap::KeyData::from_ffi(entity_id).into(); - if let Some(edge) = - named_graph.write_buf.get_edge_mut(edge_key) - { - edge.valid_to = wall_ms; - true - } else { - false - } - }; - if mutated { - let payload = - crate::persistence::wal_v3::record::encode_graph_temporal( - entity_id, is_node, wall_ms, wall_ms, - ); - gs.wal_pending.push(payload); - let wal_records = gs.drain_wal(); - drop(gs); - for record in wal_records { - ctx.shard_databases - .wal_append(ctx.shard_id, Bytes::from(record)); - } - responses.push(Frame::SimpleString( - Bytes::from_static(b"OK"), - )); - } else { - drop(gs); - responses.push(Frame::Error(Bytes::from_static( - ERR_ENTITY_NOT_FOUND, - ))); - } - } else { - drop(gs); - responses.push(Frame::Error(Bytes::from_static( - ERR_GRAPH_NOT_FOUND, - ))); - } - } - #[cfg(not(feature = "graph"))] - { - let _ = (entity_id, is_node, graph_name, wall_ms); - responses.push(Frame::Error(Bytes::from_static( - b"ERR graph feature not enabled", - ))); - } - } - Err(e) => responses.push(e), - } - continue; - } - - // --- WS.* --- - if is_ws_command(cmd) { - let sub = match parse_ws_subcommand(cmd_args) { - Ok(s) => s, - Err(e) => { responses.push(e); continue; } - }; - - if sub.eq_ignore_ascii_case(b"CREATE") { - match validate_ws_create(cmd_args) { - Ok(ws_name) => { - let ws_id = WorkspaceId::new_v7(); - let created_at = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as i64; - let meta = crate::workspace::WorkspaceMetadata { - id: ws_id, - name: ws_name.clone(), - created_at, - }; - { - let mut guard = ctx.shard_databases.workspace_registry(ctx.shard_id); - let reg = guard.get_or_insert_with(|| { - Box::new(crate::workspace::WorkspaceRegistry::new()) - }); - reg.insert(ws_id, meta); - } - // WAL: WorkspaceCreate record - let payload = crate::workspace::wal::encode_workspace_create( - ws_id.as_bytes(), &ws_name, - ); - let mut wal_buf = Vec::new(); - crate::persistence::wal_v3::record::write_wal_v3_record( - &mut wal_buf, 0, - crate::persistence::wal_v3::record::WalRecordType::WorkspaceCreate, - &payload, - ); - ctx.shard_databases.wal_append(ctx.shard_id, Bytes::from(wal_buf)); - responses.push(Frame::BulkString(Bytes::from(ws_id.to_string()))); - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"DROP") { - match validate_ws_drop(cmd_args) { - Ok(ws_id_raw) => { - match parse_workspace_id_from_bytes(&ws_id_raw) { - Some(ws_id) => { - let removed = { - let mut guard = ctx.shard_databases.workspace_registry(ctx.shard_id); - match guard.as_mut() { - Some(reg) => reg.remove(&ws_id).is_some(), - None => false, - } - }; - if removed { - // WAL: WorkspaceDrop record - let payload = crate::workspace::wal::encode_workspace_drop( - ws_id.as_bytes(), - ); - let mut wal_buf = Vec::new(); - crate::persistence::wal_v3::record::write_wal_v3_record( - &mut wal_buf, 0, - crate::persistence::wal_v3::record::WalRecordType::WorkspaceDrop, - &payload, - ); - ctx.shard_databases.wal_append(ctx.shard_id, Bytes::from(wal_buf)); - // Best-effort cleanup: delete all KV keys with ws prefix (WS-03). - { - let prefix = format!("{{{}}}:", ws_id.as_hex()); - let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, 0); - let keys_to_delete: Vec> = db_guard.keys() - .filter(|k| k.as_bytes().starts_with(prefix.as_bytes())) - .map(|k| k.as_bytes().to_vec()) - .collect(); - for key in &keys_to_delete { - db_guard.remove(key); - } - } - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } else { - responses.push(Frame::Error(Bytes::from_static(ERR_WS_NOT_FOUND))); - } - } - None => responses.push(Frame::Error(Bytes::from_static( - crate::command::workspace::ERR_WS_INVALID_ID, - ))), - } - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"LIST") { - match validate_ws_list(cmd_args) { - Ok(()) => { - let guard = ctx.shard_databases.workspace_registry(ctx.shard_id); - let entries: Vec = match guard.as_ref() { - Some(reg) => reg.iter().map(|(id, meta)| { - Frame::Array(vec![ - Frame::BulkString(Bytes::from(id.to_string())), - Frame::BulkString(meta.name.clone()), - Frame::Integer(meta.created_at), - ].into()) - }).collect(), - None => vec![], - }; - responses.push(Frame::Array(entries.into())); - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"INFO") { - match validate_ws_info(cmd_args) { - Ok(ws_id_raw) => { - match parse_workspace_id_from_bytes(&ws_id_raw) { - Some(ws_id) => { - let guard = ctx.shard_databases.workspace_registry(ctx.shard_id); - let found = guard.as_ref().and_then(|reg| reg.get(&ws_id)); - match found { - Some(meta) => { - responses.push(Frame::Array(vec![ - Frame::BulkString(Bytes::from_static(b"id")), - Frame::BulkString(Bytes::from(meta.id.to_string())), - Frame::BulkString(Bytes::from_static(b"name")), - Frame::BulkString(meta.name.clone()), - Frame::BulkString(Bytes::from_static(b"created_at")), - Frame::Integer(meta.created_at), - ].into())); - } - None => responses.push(Frame::Error(Bytes::from_static(ERR_WS_NOT_FOUND))), - } - } - None => responses.push(Frame::Error(Bytes::from_static( - crate::command::workspace::ERR_WS_INVALID_ID, - ))), - } - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"AUTH") { - match validate_ws_auth(cmd_args) { - Ok(ws_id_raw) => { - if conn.workspace_id.is_some() { - responses.push(Frame::Error(Bytes::from_static(ERR_WS_ALREADY_BOUND))); - } else { - match parse_workspace_id_from_bytes(&ws_id_raw) { - Some(ws_id) => { - let found = { - let guard = ctx.shard_databases.workspace_registry(ctx.shard_id); - guard.as_ref().map_or(false, |reg| reg.get(&ws_id).is_some()) - }; - if found { - conn.workspace_id = Some(ws_id); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } else { - responses.push(Frame::Error(Bytes::from_static(ERR_WS_NOT_FOUND))); - } - } - None => responses.push(Frame::Error(Bytes::from_static( - crate::command::workspace::ERR_WS_INVALID_ID, - ))), - } - } - } - Err(e) => responses.push(e), - } - continue; - } - - // Unknown WS subcommand - responses.push(Frame::Error(Bytes::from_static(ERR_WS_UNKNOWN_SUB))); - continue; - } - - // --- MQ.* --- - if is_mq_command(cmd) { - let sub = match parse_mq_subcommand(cmd_args) { - Ok(s) => s, - Err(e) => { responses.push(e); continue; } - }; - - if sub.eq_ignore_ascii_case(b"CREATE") { - match validate_mq_create(cmd_args) { - Ok((queue_key, max_delivery_count, _debounce_ms)) => { - let effective_key = crate::workspace::workspace_key( - conn.workspace_id.as_ref(), &queue_key, - ); - let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - match db_guard.get_or_create_stream(&effective_key) { - Ok(stream) => { - stream.durable = true; - stream.max_delivery_count = max_delivery_count; - let group_name = Bytes::from_static(b"__mq_consumers"); - let _ = stream.create_group(group_name, StreamId::ZERO); - } - Err(e) => { - responses.push(e); - continue; - } - } - drop(db_guard); - - let config = crate::mq::DurableStreamConfig::new( - effective_key.clone(), max_delivery_count, - ); - { - let mut guard = ctx.shard_databases.durable_queue_registry(ctx.shard_id); - let reg = guard.get_or_insert_with(|| { - Box::new(crate::mq::DurableQueueRegistry::new()) - }); - reg.insert(effective_key.clone(), config); - } - - let payload = crate::mq::wal::encode_mq_create( - &effective_key, max_delivery_count, - ); - let mut wal_buf = Vec::new(); - crate::persistence::wal_v3::record::write_wal_v3_record( - &mut wal_buf, 0, - crate::persistence::wal_v3::record::WalRecordType::MqCreate, - &payload, - ); - ctx.shard_databases.wal_append(ctx.shard_id, Bytes::from(wal_buf)); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"PUSH") { - match validate_mq_push(cmd_args) { - Ok((queue_key, fields)) => { - let effective_key = crate::workspace::workspace_key( - conn.workspace_id.as_ref(), &queue_key, - ); - let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - match db_guard.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - if !stream.durable { - responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); - } else { - let msg_id = stream.next_auto_id(); - let msg_id = stream.add(msg_id, fields); - drop(db_guard); - { - let mut trig_guard = ctx.shard_databases.trigger_registry(ctx.shard_id); - if let Some(reg) = trig_guard.as_mut() { - let trig_key = if let Some(ws_id) = conn.workspace_id.as_ref() { - let ws_hex = ws_id.as_hex(); - let mut k = Vec::with_capacity(ws_hex.len() + 1 + queue_key.len()); - k.extend_from_slice(ws_hex.as_bytes()); - k.push(b':'); - k.extend_from_slice(&queue_key); - Bytes::from(k) - } else { - queue_key.clone() - }; - if let Some(trig_entry) = reg.get_mut(&trig_key) { - if trig_entry.pending_fire_ms == 0 { - let fire_at = ctx.cached_clock.ms() + trig_entry.debounce_ms; - trig_entry.pending_fire_ms = fire_at; - } - } - } - } - responses.push(Frame::BulkString(Bytes::from( - format!("{}-{}", msg_id.ms, msg_id.seq), - ))); - } - } - Ok(None) => { - responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); - } - Err(e) => responses.push(e), - } - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"POP") { - match validate_mq_pop(cmd_args) { - Ok((queue_key, count)) => { - let effective_key = crate::workspace::workspace_key( - conn.workspace_id.as_ref(), &queue_key, - ); - let group_name = Bytes::from_static(b"__mq_consumers"); - let consumer_name = Bytes::from_static(b"__mq_default"); - let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - - let mdc = match db_guard.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - if !stream.durable { - responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); - continue; - } - stream.max_delivery_count - } - Ok(None) => { - responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); - continue; - } - Err(e) => { - responses.push(e); - continue; - } - }; - - let request_count = count + (mdc as usize); - let stream = match db_guard.get_stream_mut(&effective_key) { - Ok(Some(s)) => s, - _ => { responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); continue; } - }; - let claimed = match stream.read_group_new( - &group_name, &consumer_name, Some(request_count), false, - ) { - Ok(entries) => entries, - Err(_) => { - responses.push(Frame::Array(vec![].into())); - continue; - } - }; - - let mut results = Vec::new(); - let mut dlq_entries: Vec<(StreamId, Vec<(Bytes, Bytes)>)> = Vec::new(); - let mut dlq_ack_ids: Vec = Vec::new(); - for (id, fields) in &claimed { - let delivery_count = stream.groups.get(group_name.as_ref()) - .and_then(|g| g.pel.get(id)) - .map(|pe| pe.delivery_count) - .unwrap_or(1); - if mdc > 0 && delivery_count >= mdc as u64 { - dlq_entries.push((*id, fields.clone())); - dlq_ack_ids.push(*id); - } else if results.len() < count { - results.push((*id, fields.clone())); - } - } - - if !dlq_ack_ids.is_empty() { - let _ = stream.xack(&group_name, &dlq_ack_ids); - } - - if !dlq_entries.is_empty() { - let dlq_key = { - let mut buf = Vec::with_capacity(effective_key.len() + 8); - buf.extend_from_slice(&effective_key); - buf.extend_from_slice(b"::mq:dlq"); - Bytes::from(buf) - }; - match db_guard.get_or_create_stream(&dlq_key) { - Ok(dlq_stream) => { - for (_id, fields) in dlq_entries { - let dlq_id = dlq_stream.next_auto_id(); - dlq_stream.add(dlq_id, fields); - } - } - Err(_) => {} - } - } - - let result_frames: Vec = results.iter().map(|(id, fields)| { - let mut entry_frames = Vec::with_capacity(2); - entry_frames.push(Frame::BulkString(Bytes::from( - format!("{}-{}", id.ms, id.seq), - ))); - let field_frames: Vec = fields.iter() - .flat_map(|(f, v)| vec![ - Frame::BulkString(f.clone()), - Frame::BulkString(v.clone()), - ]) - .collect(); - entry_frames.push(Frame::Array(field_frames.into())); - Frame::Array(entry_frames.into()) - }).collect(); - responses.push(Frame::Array(result_frames.into())); - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"ACK") { - match validate_mq_ack(cmd_args) { - Ok((queue_key, msg_ids)) => { - let effective_key = crate::workspace::workspace_key( - conn.workspace_id.as_ref(), &queue_key, - ); - let ids: Vec = msg_ids.iter() - .map(|(ms, seq)| StreamId { ms: *ms, seq: *seq }) - .collect(); - let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - match db_guard.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - let group_name = Bytes::from_static(b"__mq_consumers"); - match stream.xack(&group_name, &ids) { - Ok(acked_count) => { - drop(db_guard); - for (ms, seq) in &msg_ids { - let payload = crate::mq::wal::encode_mq_ack( - &effective_key, *ms, *seq, - ); - let mut wal_buf = Vec::new(); - crate::persistence::wal_v3::record::write_wal_v3_record( - &mut wal_buf, 0, - crate::persistence::wal_v3::record::WalRecordType::MqAck, - &payload, - ); - ctx.shard_databases.wal_append( - ctx.shard_id, Bytes::from(wal_buf), - ); - } - responses.push(Frame::Integer(acked_count as i64)); - } - Err(_) => responses.push(Frame::Integer(0)), - } - } - Ok(None) => responses.push(Frame::Integer(0)), - Err(_) => responses.push(Frame::Integer(0)), - } - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"DLQLEN") { - match validate_mq_dlqlen(cmd_args) { - Ok(queue_key) => { - let effective_key = crate::workspace::workspace_key( - conn.workspace_id.as_ref(), &queue_key, - ); - let dlq_key = { - let mut buf = Vec::with_capacity(effective_key.len() + 8); - buf.extend_from_slice(&effective_key); - buf.extend_from_slice(b"::mq:dlq"); - Bytes::from(buf) - }; - let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let len = match db_guard.get_stream_mut(&dlq_key) { - Ok(Some(stream)) => stream.length as i64, - _ => 0i64, - }; - responses.push(Frame::Integer(len)); - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"TRIGGER") { - match validate_mq_trigger(cmd_args) { - Ok((queue_key, callback_cmd, debounce_ms)) => { - let effective_key = crate::workspace::workspace_key( - conn.workspace_id.as_ref(), &queue_key, - ); - let trig_key = if let Some(ws_id) = conn.workspace_id.as_ref() { - let ws_hex = ws_id.as_hex(); - let mut k = Vec::with_capacity(ws_hex.len() + 1 + queue_key.len()); - k.extend_from_slice(ws_hex.as_bytes()); - k.push(b':'); - k.extend_from_slice(&queue_key); - Bytes::from(k) - } else { - queue_key.clone() - }; - let entry = crate::mq::TriggerEntry { - queue_key: effective_key, - callback_cmd, - debounce_ms, - last_fire_ms: 0, - pending_fire_ms: 0, - }; - { - let mut guard = ctx.shard_databases.trigger_registry(ctx.shard_id); - let reg = guard.get_or_insert_with(|| { - Box::new(crate::mq::TriggerRegistry::new()) - }); - reg.register(trig_key, entry); - } - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - Err(e) => responses.push(e), - } - continue; - } - - if sub.eq_ignore_ascii_case(b"PUBLISH") { - match validate_mq_publish(cmd_args) { - Ok((queue_key, fields)) => { - let effective_key = crate::workspace::workspace_key( - conn.workspace_id.as_ref(), &queue_key, - ); - if let Some(ref mut txn) = conn.active_cross_txn { - txn.record_mq(effective_key, fields); - responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); - } else { - responses.push(Frame::Error(Bytes::from_static( - b"ERR MQ PUBLISH requires an active transaction (use TXN BEGIN first)", - ))); - } - } - Err(e) => responses.push(e), - } - continue; - } - - // Unknown MQ subcommand - responses.push(Frame::Error(Bytes::from_static(ERR_MQ_UNKNOWN_SUB))); - continue; - } - - // --- MULTI --- - if cmd.eq_ignore_ascii_case(b"MULTI") { - if conn.in_cross_txn() { - responses.push(Frame::Error(Bytes::from_static(ERR_MULTI_TXN_CONFLICT))); - } else if conn.in_multi { - responses.push(Frame::Error(Bytes::from_static( - b"ERR MULTI calls can not be nested", - ))); - } else { - conn.in_multi = true; - conn.command_queue.clear(); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - continue; - } - - // --- EXEC --- - if cmd.eq_ignore_ascii_case(b"EXEC") { - if !conn.in_multi { - responses.push(Frame::Error(Bytes::from_static(b"ERR EXEC without MULTI"))); - } else { - conn.in_multi = false; - let result = execute_transaction_sharded(&ctx.shard_databases, ctx.shard_id, &conn.command_queue, conn.selected_db, &ctx.cached_clock); - conn.command_queue.clear(); - responses.push(result); - } - continue; - } - - // --- DISCARD --- - if cmd.eq_ignore_ascii_case(b"DISCARD") { - if !conn.in_multi { - responses.push(Frame::Error(Bytes::from_static(b"ERR DISCARD without MULTI"))); - } else { - conn.in_multi = false; - conn.command_queue.clear(); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); - } - continue; - } - - // --- Workspace key prefix injection --- - // MUST happen before key_to_shard() so the {ws_id} hash tag determines - // shard routing. This is the ONLY code path where workspace prefixing - // occurs (WS-07, WS-12). All subsequent dispatch uses cmd_args (shadowed). - let rewritten = conn - .workspace_id - .as_ref() - .map(|ws_id| workspace_rewrite_args(cmd, cmd_args, ws_id)); - let cmd_args: &[Frame] = rewritten.as_deref().unwrap_or(cmd_args); - - // --- BLOCKING COMMANDS --- - if cmd.eq_ignore_ascii_case(b"BLPOP") || cmd.eq_ignore_ascii_case(b"BRPOP") - || cmd.eq_ignore_ascii_case(b"BLMOVE") || cmd.eq_ignore_ascii_case(b"BZPOPMIN") - || cmd.eq_ignore_ascii_case(b"BZPOPMAX") - || cmd.eq_ignore_ascii_case(b"BLMPOP") || cmd.eq_ignore_ascii_case(b"BRPOPLPUSH") - || cmd.eq_ignore_ascii_case(b"BZMPOP") - { - if conn.in_multi { - let nb_frame = convert_blocking_to_nonblocking(cmd, cmd_args); - conn.command_queue.push(nb_frame); - responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); - continue; - } - write_buf.clear(); - for response in responses.iter() { - if conn.protocol_version >= 3 { - crate::protocol::serialize_resp3(response, &mut write_buf); - } else { - crate::protocol::serialize(response, &mut write_buf); - } - } - if stream.write_all(&write_buf).await.is_err() { arena.reset(); return (HandlerResult::Done, None); } - let blocking_response = handle_blocking_command( - cmd, cmd_args, conn.selected_db, &ctx.shard_databases, &ctx.blocking_registry, - ctx.shard_id, ctx.num_shards, &ctx.dispatch_tx, &shutdown, - ).await; - let blocking_response = apply_resp3_conversion(cmd, blocking_response, conn.protocol_version); - responses = Vec::with_capacity(1); - responses.push(blocking_response); - break; - } - - // --- PUBLISH --- - if cmd.eq_ignore_ascii_case(b"PUBLISH") { - if cmd_args.len() != 2 { - responses.push(Frame::Error(Bytes::from_static(b"ERR wrong number of arguments for 'publish' command"))); - } else { - let channel_arg = extract_bytes(&cmd_args[0]); - let message_arg = extract_bytes(&cmd_args[1]); - // ACL channel permission check for PUBLISH - if let Some(ref ch) = channel_arg { - #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable - let acl_guard = ctx.acl_table.read().unwrap(); - if let Some(deny_reason) = acl_guard.check_channel_permission(&conn.current_user, ch.as_ref()) { - drop(acl_guard); - responses.push(Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason)))); - continue; - } - } - match (channel_arg, message_arg) { - (Some(ch), Some(msg)) => { - let local_count = { ctx.pubsub_registry.write().publish(&ch, &msg) }; - // Targeted fanout: only send to shards that have subscribers - let targets = ctx.remote_subscriber_map.read().target_shards(&ch); - if targets.is_empty() { - // Fast path: no remote subscribers, return local count immediately - responses.push(Frame::Integer(local_count)); - } else { - // Filter to remote targets only (skip self) - let remote_targets: Vec = targets.into_iter().filter(|&t| t != ctx.shard_id).collect(); - if remote_targets.is_empty() { - responses.push(Frame::Integer(local_count)); - } else { - // Accumulate into per-shard batches for coalesced dispatch - let resp_idx = responses.len(); - responses.push(Frame::Integer(local_count)); // placeholder, updated after batch flush - for target in &remote_targets { - publish_batches.entry(*target) - .or_default() - .push((resp_idx, ch.clone(), msg.clone())); - } - } - } - } - _ => responses.push(Frame::Error(Bytes::from_static(b"ERR invalid channel or message"))), - } - } - continue; - } - - // --- SUBSCRIBE / PSUBSCRIBE --- - if cmd.eq_ignore_ascii_case(b"SUBSCRIBE") || cmd.eq_ignore_ascii_case(b"PSUBSCRIBE") { - let is_pattern = cmd.eq_ignore_ascii_case(b"PSUBSCRIBE"); - let cmd_name = if is_pattern { "psubscribe" } else { "subscribe" }; - if cmd_args.is_empty() { - responses.push(Frame::Error(Bytes::from(format!( - "ERR wrong number of arguments for '{}' command", cmd_name - )))); - continue; - } - // Allocate pubsub channel if not yet created - if conn.pubsub_tx.is_none() { - let (tx, rx) = channel::mpsc_bounded::(256); - conn.pubsub_tx = Some(tx); - conn.pubsub_rx = Some(rx); - } - if conn.subscriber_id == 0 { - conn.subscriber_id = crate::pubsub::next_subscriber_id(); - } - // Flush accumulated responses before entering subscriber mode - if !responses.is_empty() { - write_buf.clear(); - for resp in &responses { - if conn.protocol_version >= 3 { - crate::protocol::serialize_resp3(resp, &mut write_buf); - } else { - crate::protocol::serialize(resp, &mut write_buf); - } - } - if stream.write_all(&write_buf).await.is_err() { return (HandlerResult::Done, None); } - responses.clear(); - } - // Process subscribe arguments - for arg in cmd_args { - if let Some(ch) = extract_bytes(arg) { - #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable - let acl_deny = { ctx.acl_table.read().unwrap().check_channel_permission(&conn.current_user, ch.as_ref()) }; - if let Some(reason) = acl_deny { - write_buf.clear(); - let err = Frame::Error(Bytes::from(format!("NOPERM {}", reason))); - crate::protocol::serialize(&err, &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { return (HandlerResult::Done, None); } - continue; - } - #[allow(clippy::unwrap_used)] // conn.pubsub_tx is set to Some just above before this loop - let sub = Subscriber::with_protocol( - conn.pubsub_tx.clone().unwrap(), - conn.subscriber_id, - conn.protocol_version >= 3, - ); - if is_pattern { - { ctx.pubsub_registry.write().psubscribe(ch.clone(), sub); } - } else { - { ctx.pubsub_registry.write().subscribe(ch.clone(), sub); } - } - conn.subscription_count += 1; - // Register pub/sub affinity for this client IP - if conn.subscription_count == 1 { - if let Ok(addr) = peer_addr.parse::() { - ctx.pubsub_affinity.write().register(addr.ip(), ctx.shard_id); - } - } - propagate_subscription(&ctx.all_remote_sub_maps, &ch, ctx.shard_id, ctx.num_shards, is_pattern); - write_buf.clear(); - let resp = if is_pattern { - pubsub::psubscribe_response(&ch, conn.subscription_count) - } else { - pubsub::subscribe_response(&ch, conn.subscription_count) - }; - crate::protocol::serialize(&resp, &mut write_buf); - if stream.write_all(&write_buf).await.is_err() { return (HandlerResult::Done, None); } - } - } - break; // break out of frame batch loop to re-enter main loop in subscriber mode - } - // UNSUBSCRIBE/PUNSUBSCRIBE in normal mode (not subscribed) - if cmd.eq_ignore_ascii_case(b"UNSUBSCRIBE") || cmd.eq_ignore_ascii_case(b"PUNSUBSCRIBE") { - let is_pattern = cmd.eq_ignore_ascii_case(b"PUNSUBSCRIBE"); - let resp = if is_pattern { - pubsub::punsubscribe_response(&Bytes::from_static(b""), 0) - } else { - pubsub::unsubscribe_response(&Bytes::from_static(b""), 0) - }; - responses.push(resp); - continue; - } - - // --- PUBSUB introspection (zero-SPSC: direct shared-read) --- - if cmd.eq_ignore_ascii_case(b"PUBSUB") { - if cmd_args.is_empty() { - responses.push(Frame::Error(Bytes::from_static(b"ERR wrong number of arguments for 'pubsub' command"))); - continue; - } - let subcmd = extract_bytes(&cmd_args[0]); - match subcmd { - Some(ref sc) if sc.eq_ignore_ascii_case(b"CHANNELS") => { - let pattern = if cmd_args.len() > 1 { extract_bytes(&cmd_args[1]) } else { None }; - let mut all_channels: std::collections::HashSet = std::collections::HashSet::new(); - for reg in &ctx.all_pubsub_registries { - let guard = reg.read(); - all_channels.extend(guard.active_channels(pattern.as_deref())); - } - let arr: Vec = all_channels.into_iter().map(Frame::BulkString).collect(); - responses.push(Frame::Array(arr.into())); - } - Some(ref sc) if sc.eq_ignore_ascii_case(b"NUMSUB") => { - let channels: Vec = cmd_args[1..].iter().filter_map(|a| extract_bytes(a)).collect(); - let mut counts: HashMap = HashMap::new(); - for reg in &ctx.all_pubsub_registries { - let guard = reg.read(); - for (ch, c) in guard.numsub(&channels) { - *counts.entry(ch).or_insert(0) += c; - } - } - let mut arr = Vec::with_capacity(channels.len() * 2); - for ch in &channels { - arr.push(Frame::BulkString(ch.clone())); - arr.push(Frame::Integer(*counts.get(ch).unwrap_or(&0))); - } - responses.push(Frame::Array(arr.into())); - } - Some(ref sc) if sc.eq_ignore_ascii_case(b"NUMPAT") => { - let mut total: usize = 0; - for reg in &ctx.all_pubsub_registries { - total += reg.read().numpat(); - } - responses.push(Frame::Integer(total as i64)); - } - _ => { - responses.push(Frame::Error(Bytes::from_static(b"ERR unknown subcommand or wrong number of arguments for 'pubsub' command"))); - } - } - continue; - } - - // --- MULTI queue mode --- - if conn.in_multi { - conn.command_queue.push(frame); - responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); - continue; - } - - // --- BGSAVE --- - if cmd.eq_ignore_ascii_case(b"BGSAVE") { - responses.push(crate::command::persistence::bgsave_start_sharded(&ctx.snapshot_trigger_tx, ctx.num_shards)); - continue; - } - if cmd.eq_ignore_ascii_case(b"SAVE") { - responses.push(Frame::Error(Bytes::from_static(b"ERR SAVE not supported in sharded mode, use BGSAVE"))); - continue; - } - if cmd.eq_ignore_ascii_case(b"LASTSAVE") { - responses.push(crate::command::persistence::handle_lastsave()); - continue; - } - if cmd.eq_ignore_ascii_case(b"BGREWRITEAOF") { - if let Some(ref tx) = ctx.aof_tx { - responses.push(crate::command::persistence::bgrewriteaof_start_sharded( - tx, - ctx.shard_databases.clone(), - )); - } else { - responses.push(Frame::Error(Bytes::from_static( - b"ERR AOF is not enabled", - ))); - } - continue; - } - - // --- Cross-shard aggregation: KEYS, SCAN, DBSIZE --- - if cmd.eq_ignore_ascii_case(b"KEYS") { - let mut response = crate::shard::coordinator::coordinate_keys(cmd_args, ctx.shard_id, ctx.num_shards, conn.selected_db, &ctx.shard_databases, &ctx.dispatch_tx, &ctx.spsc_notifiers, &ctx.cached_clock, &()).await; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"SCAN") { - let mut response = crate::shard::coordinator::coordinate_scan(cmd_args, ctx.shard_id, ctx.num_shards, conn.selected_db, &ctx.shard_databases, &ctx.dispatch_tx, &ctx.spsc_notifiers, &ctx.cached_clock, &()).await; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"DBSIZE") { - let response = crate::shard::coordinator::coordinate_dbsize(ctx.shard_id, ctx.num_shards, conn.selected_db, &ctx.shard_databases, &ctx.dispatch_tx, &ctx.spsc_notifiers, &()).await; - responses.push(response); - continue; - } - - // --- FT.* vector search commands --- - if cmd.len() > 3 && cmd[..3].eq_ignore_ascii_case(b"FT.") { - if ctx.num_shards > 1 { - // Multi-shard: dispatch via SPSC - #[cfg(feature = "text-index")] - if cmd.eq_ignore_ascii_case(b"FT.AGGREGATE") { - // ── FT.AGGREGATE: two-phase scatter-gather per Plan 03 (D-05/D-07) ── - // scatter_text_aggregate acquires its own guards internally - // inside the single-shard block, so we never hold a MutexGuard - // across the .await below. - let parsed = match crate::command::vector_search::ft_aggregate::parse_aggregate_args(cmd_args) { - Ok(p) => p, - Err(err_frame) => { - responses.push(err_frame); - continue; - } - }; - let response = crate::shard::scatter_aggregate::scatter_text_aggregate( - parsed.index_name, - parsed.query, - parsed.pipeline, - ctx.shard_id, - ctx.num_shards, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - ).await; - responses.push(response); - continue; - } - if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { - // Check if this is a text query BEFORE trying parse_ft_search_args - // (which would return an error for non-KNN queries). - let query_bytes = cmd_args.get(1).and_then(|f| crate::command::vector_search::extract_bulk(f)); - let is_text = query_bytes.as_ref().map_or(false, |q| { - crate::command::vector_search::is_text_query(q) - }); - - // ── HYBRID multi-shard path (Phase 152 Plan 05, D-13) ────── - // If the args contain a HYBRID clause, route through - // scatter_hybrid_search (three-phase DFS → fan-out → RRF - // merge). Single-shard is handled inside the scatter entry - // point itself (fast path to execute_hybrid_search_local). - #[cfg(feature = "text-index")] - { - match crate::command::vector_search::hybrid::parse_hybrid_modifier(cmd_args) { - Ok(Some(partial)) => { - let index_name = match cmd_args.first().and_then(|f| crate::command::vector_search::extract_bulk(f)) { - Some(b) => b, - None => { - responses.push(Frame::Error(Bytes::from_static(b"ERR invalid index name"))); - continue; - } - }; - let text_query = match query_bytes.clone() { - Some(q) => q, - None => { - responses.push(Frame::Error(Bytes::from_static(b"ERR invalid query"))); - continue; - } - }; - let (limit_offset, limit_count) = crate::command::vector_search::parse_limit_clause(cmd_args); - let top_k = if limit_count == usize::MAX { - limit_offset.saturating_add(10).max(1) - } else { - limit_offset.saturating_add(limit_count).max(1) - }; - let hq = crate::command::vector_search::hybrid::HybridQuery { - index_name, - text_query, - dense_field: partial.dense_field, - dense_blob: partial.dense_blob, - sparse: partial.sparse, - weights: partial.weights, - k_per_stream: partial.k_per_stream, - top_k, - offset: limit_offset, - count: limit_count, - }; - // Phase 171 HYB-02 / SCAT-02: resolve AS_OF / - // TXN LSN ONCE on the coordinator and forward to - // every responder via the scatter helper. - let as_of_lsn = match resolve_ft_search_as_of_lsn( - cmd_args, - Some(&ctx.shard_databases), - ctx.shard_id, - conn.active_cross_txn.as_ref(), - ) { - Ok(lsn) => lsn, - Err(err_frame) => { - responses.push(err_frame); - continue; - } - }; - let response = crate::shard::scatter_hybrid::scatter_hybrid_search( - hq, - as_of_lsn, - ctx.shard_id, - ctx.num_shards, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - ).await; - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - Ok(None) => { /* fall through to non-HYBRID paths */ } - Err(err_frame) => { - responses.push(err_frame); - continue; - } - } - } - - if is_text { - // ── Text FT.SEARCH: two-phase DFS scatter-gather ───────────── - let index_name = match cmd_args.first().and_then(|f| crate::command::vector_search::extract_bulk(f)) { - Some(b) => b, - None => { - responses.push(Frame::Error(Bytes::from_static(b"ERR invalid index name"))); - continue; - } - }; - let query_str = query_bytes.unwrap(); - - // B-01 SITE 3 FIX (Plan 152-06): FieldFilter short-circuit - // BEFORE the analyzer-first parse_result block. Symmetric with - // handler_monoio. TAG queries route through the InvertedSearch - // fan-out — no analyzer touched, no field_idx resolution. - #[cfg(feature = "text-index")] - { - match crate::command::vector_search::pre_parse_field_filter(query_str.as_ref()) { - Ok(Some(clause)) => { - if let Some(filter) = clause.filter { - let (offset, count) = crate::command::vector_search::parse_limit_clause(cmd_args); - let top_k = if count == usize::MAX { 10000 } else { offset.saturating_add(count) }.max(1); - let response = crate::shard::coordinator::scatter_text_search_filter( - index_name, - filter, - top_k, - offset, - count, - ctx.shard_id, - ctx.num_shards, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - ).await; - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - } - Ok(None) => { /* fall through */ } - Err(e) => { - responses.push(Frame::Error(Bytes::from(e.to_owned()))); - continue; - } - } - } - - // Parse query and resolve field_idx inside a block scope so the - // MutexGuard from text_store() is dropped BEFORE .await. - // We use the TextIndex's own field_analyzers (same pipeline used at index time). - type ParseResult = Result<(Vec, Option), String>; - let parse_result: ParseResult = { - let ts = ctx.shard_databases.text_store(ctx.shard_id); - match ts.get_index(&index_name) { - None => Err("ERR no such index".to_owned()), - Some(text_index) => { - match text_index.field_analyzers.first() { - None => Err("ERR index has no TEXT fields".to_owned()), - Some(analyzer) => { - // analyzer borrows text_index which borrows ts — all in this block. - let parsed = crate::command::vector_search::parse_text_query(&query_str, analyzer); - match parsed { - Err(e) => Err(e.to_owned()), - Ok(clause) => { - let field_idx = match &clause.field_name { - None => Ok(None), - Some(field_name) => { - match text_index.text_fields.iter().position(|f| f.field_name.as_ref().eq_ignore_ascii_case(field_name.as_ref())) { - Some(idx) => Ok(Some(idx)), - None => Err(format!("ERR unknown field '{}'", String::from_utf8_lossy(field_name))), - } - } - }; - field_idx.map(|idx| (clause.terms, idx)) - } - } - } - } - } - } - }; // MutexGuard dropped here - - let (query_terms, field_idx) = match parse_result { - Ok(t) => t, - Err(e) => { - responses.push(Frame::Error(Bytes::from(e))); - continue; - } - }; - - let (offset, count) = crate::command::vector_search::parse_limit_clause(cmd_args); - let top_k = if count == usize::MAX { 10000 } else { offset.saturating_add(count) }.max(1); - - // Parse optional HIGHLIGHT/SUMMARIZE clauses from args. - let highlight_opts = crate::command::vector_search::parse_highlight_clause(cmd_args); - let summarize_opts = crate::command::vector_search::parse_summarize_clause(cmd_args); - - let mut response = crate::shard::coordinator::scatter_text_search( - index_name, - query_terms, - field_idx, - top_k, - offset, - count, - ctx.shard_id, - ctx.num_shards, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - highlight_opts, - summarize_opts, - ).await; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - - // ── Vector FT.SEARCH (KNN / SPARSE): existing path ─────────── - // Phase 171 SCAT-01: resolve AS_OF / TXN snapshot LSN ONCE - // on the coordinator and thread it through the scatter helper - // so every responder honors the same temporal snapshot. - let response = match crate::command::vector_search::parse_ft_search_args(cmd_args) { - Ok((index_name, query_blob, k, filter, _offset, _count)) => { - if filter.is_some() { - Frame::Error(Bytes::from_static( - b"ERR FILTER not supported in multi-shard mode yet", - )) - } else { - match resolve_ft_search_as_of_lsn( - cmd_args, - Some(&ctx.shard_databases), - ctx.shard_id, - conn.active_cross_txn.as_ref(), - ) { - Err(err_frame) => err_frame, - Ok(as_of_lsn) => { - crate::shard::coordinator::scatter_vector_search_remote( - index_name, query_blob, k, as_of_lsn, - ctx.shard_id, ctx.num_shards, - &ctx.shard_databases, - &ctx.dispatch_tx, &ctx.spsc_notifiers, - ).await - } - } - } - } - Err(err_frame) => err_frame, - }; - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - let response = crate::shard::coordinator::broadcast_vector_command( - std::sync::Arc::new(frame), - ctx.shard_id, ctx.num_shards, - &ctx.shard_databases, - &ctx.dispatch_tx, &ctx.spsc_notifiers, - ).await; - responses.push(response); - continue; - } else { - // Single-shard: no SPSC channels available. - // Dispatch directly to shard's VectorStore via shared access. - // - // ── 154-01 single-shard FT.AGGREGATE fast path ──────── - // scatter_text_aggregate internally fast-paths num_shards - // == 1 to execute_local_full. Kept byte-symmetric with - // the multi-shard arm at line 1458 so the two code paths - // share a single dispatch body shape. - #[cfg(feature = "text-index")] - if cmd.eq_ignore_ascii_case(b"FT.AGGREGATE") { - let parsed = match crate::command::vector_search::ft_aggregate::parse_aggregate_args(cmd_args) { - Ok(p) => p, - Err(err_frame) => { - responses.push(err_frame); - continue; - } - }; - let response = crate::shard::scatter_aggregate::scatter_text_aggregate( - parsed.index_name, - parsed.query, - parsed.pipeline, - ctx.shard_id, - ctx.num_shards, - &ctx.shard_databases, - &ctx.dispatch_tx, - &ctx.spsc_notifiers, - ).await; - responses.push(response); - continue; - } - // - // ── 151-03 single-shard text FT.SEARCH fast path ────── - // Parity with handler_monoio.rs. See that file's block - // for the full rationale — bare text queries bypass - // ft_search() (which only parses KNN/SPARSE/HYBRID) - // and route directly to execute_text_search_local. - #[cfg(feature = "text-index")] - if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { - if let Some(Frame::BulkString(query_bytes)) = cmd_args.get(1) { - match crate::command::vector_search::parse_hybrid_modifier(cmd_args) { - Ok(Some(_)) => { - // HYBRID present — existing ft_search() below handles it. - } - Err(frame_err) => { - responses.push(frame_err); - continue; - } - Ok(None) => { - if crate::command::vector_search::is_text_query( - query_bytes.as_ref(), - ) { - // Step 1: index_name. - let index_name = match cmd_args.first() { - Some(Frame::BulkString(b)) => b.clone(), - _ => { - responses.push(Frame::Error(Bytes::from_static( - b"ERR wrong number of arguments for FT.SEARCH", - ))); - continue; - } - }; - // B-01 SITE 3 FIX (single-shard 151-03 fast path, - // Plan 152-06): FieldFilter short-circuit BEFORE - // text_fields.is_empty() bail. - #[cfg(feature = "text-index")] - match crate::command::vector_search::pre_parse_field_filter(query_bytes.as_ref()) { - Ok(Some(clause)) => { - if clause.filter.is_some() { - let (offset, count) = crate::command::vector_search::parse_limit_clause(cmd_args); - let top_k = if count == usize::MAX { 10000 } else { offset.saturating_add(count) }.max(1); - let ts_guard = ctx.shard_databases.text_store(ctx.shard_id); - let response = match ts_guard.get_index(&index_name) { - None => Frame::Error(Bytes::from_static(b"ERR no such index")), - Some(text_index) => { - let results = crate::command::vector_search::ft_text_search::execute_query_on_index( - text_index, &clause, None, None, top_k, - ); - crate::command::vector_search::ft_text_search::build_text_response( - &results, offset, count, - ) - } - }; - drop(ts_guard); - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - } - Ok(None) => { /* fall through */ } - Err(e) => { - responses.push(Frame::Error(Bytes::copy_from_slice(e.as_bytes()))); - continue; - } - } - // Step 2: ts guard. - let ts_guard = - ctx.shard_databases.text_store(ctx.shard_id); - // Step 3: index lookup. - let text_index = match ts_guard.get_index(&index_name) { - Some(idx) => idx, - None => { - drop(ts_guard); - responses.push(Frame::Error(Bytes::from_static( - b"ERR no such index", - ))); - continue; - } - }; - // Step 4: index must have at least one TEXT field. - if text_index.text_fields.is_empty() { - drop(ts_guard); - responses.push(Frame::Error(Bytes::from_static( - b"ERR index has no TEXT fields", - ))); - continue; - } - // Step 5: parse via first analyzer. - let analyzer = match text_index.field_analyzers.first() { - Some(a) => a, - None => { - drop(ts_guard); - responses.push(Frame::Error(Bytes::from_static( - b"ERR index has no TEXT fields", - ))); - continue; - } - }; - let clause = match crate::command::vector_search::parse_text_query( - query_bytes.as_ref(), - analyzer, - ) { - Ok(c) => c, - Err(msg) => { - drop(ts_guard); - responses.push(Frame::Error( - Bytes::copy_from_slice(msg.as_bytes()), - )); - continue; - } - }; - // Step 5b: resolve field_idx. - let field_idx = match &clause.field_name { - None => None, - Some(field_name) => match text_index - .text_fields - .iter() - .position(|f| { - f.field_name.as_ref().eq_ignore_ascii_case( - field_name.as_ref(), - ) - }) { - Some(idx) => Some(idx), - None => { - let bad_name = field_name.clone(); - drop(ts_guard); - responses.push(Frame::Error(Bytes::from( - format!( - "ERR unknown field '{}'", - String::from_utf8_lossy(&bad_name) - ), - ))); - continue; - } - }, - }; - // Step 6: query_terms. - let query_terms = clause.terms; - // Step 7: LIMIT + top_k cap (T-151-03-02). - let (offset, count) = - crate::command::vector_search::parse_limit_clause( - cmd_args, - ); - let top_k = if count == usize::MAX { - 10000 - } else { - offset.saturating_add(count) - } - .max(1); - // Step 8: HIGHLIGHT / SUMMARIZE. - let highlight_opts = - crate::command::vector_search::parse_highlight_clause( - cmd_args, - ); - let summarize_opts = - crate::command::vector_search::parse_summarize_clause( - cmd_args, - ); - // Step 9: DB read guard iff post-processing needed. - let db_guard_opt = if highlight_opts.is_some() - || summarize_opts.is_some() - { - Some(ctx.shard_databases.read_db(ctx.shard_id, 0)) - } else { - None - }; - // Step 10: execute + optional post-processing. - let mut response = - crate::command::vector_search::execute_text_search_local( - &ts_guard, - &index_name, - field_idx, - &query_terms, - top_k, - offset, - count, - ); - if let Some(ref db_guard) = db_guard_opt { - let term_strings: Vec = query_terms - .iter() - .map(|qt| qt.text.clone()) - .collect(); - crate::command::vector_search::apply_post_processing( - &mut response, - &term_strings, - text_index, - db_guard, - highlight_opts.as_ref(), - summarize_opts.as_ref(), - ); - } - // Explicit drop order: db_guard first, ts_guard last. - drop(db_guard_opt); - drop(ts_guard); - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - } - } - } - } - let response = { - let shard_databases_ref = &ctx.shard_databases; - let mut vs = shard_databases_ref.vector_store(ctx.shard_id); - let mut ts = shard_databases_ref.text_store(ctx.shard_id); - if cmd.eq_ignore_ascii_case(b"FT.CREATE") { - crate::command::vector_search::ft_create(&mut vs, &mut ts, cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { - // Resolve AS_OF temporal clause + TXN snapshot precedence (TEMP-04, ACID-09). - match resolve_ft_search_as_of_lsn( - cmd_args, - Some(&ctx.shard_databases), - ctx.shard_id, - conn.active_cross_txn.as_ref(), - ) { - Err(err_frame) => err_frame, - Ok(as_of_lsn) => { - let has_session = cmd_args.iter().any(|a| { - if let Frame::BulkString(b) = a { b.eq_ignore_ascii_case(b"SESSION") } else { false } - }); - if has_session { - let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); - crate::command::vector_search::ft_search(&mut vs, cmd_args, Some(&mut *db_guard), Some(&*ts), as_of_lsn) - } else { - crate::command::vector_search::ft_search(&mut vs, cmd_args, None, Some(&*ts), as_of_lsn) - } - } - } - } else if cmd.eq_ignore_ascii_case(b"FT.DROPINDEX") { - let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); - crate::command::vector_search::ft_dropindex(&mut vs, &mut ts, Some(&mut *db_guard), cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT.INFO") { - crate::command::vector_search::ft_info(&vs, &ts, cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT._LIST") { - crate::command::vector_search::ft_list(&vs) - } else if cmd.eq_ignore_ascii_case(b"FT.COMPACT") { - crate::command::vector_search::ft_compact(&mut vs, &mut ts, cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT.CACHESEARCH") { - crate::command::vector_search::cache_search::ft_cachesearch(&mut vs, cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT.CONFIG") { - crate::command::vector_search::ft_config(&mut vs, &mut ts, cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") { - let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); - crate::command::vector_search::recommend::ft_recommend(&mut vs, cmd_args, Some(&mut *db_guard)) - } else if cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") { - #[cfg(feature = "graph")] - { - let graph_guard = shard_databases_ref.graph_store_read(ctx.shard_id); - crate::command::vector_search::navigate::ft_navigate(&mut vs, Some(&graph_guard), cmd_args, None) - } - #[cfg(not(feature = "graph"))] - { - Frame::Error(Bytes::from_static(b"ERR FT.NAVIGATE requires graph feature")) - } - } else if cmd.eq_ignore_ascii_case(b"FT.EXPAND") { - #[cfg(feature = "graph")] - { - let graph_guard = shard_databases_ref.graph_store_read(ctx.shard_id); - crate::command::vector_search::ft_expand(&graph_guard, cmd_args) - } - #[cfg(not(feature = "graph"))] - { - Frame::Error(Bytes::from_static(b"ERR FT.EXPAND requires graph feature")) - } - } else { - Frame::Error(Bytes::from_static(b"ERR unknown FT.* command")) - } - }; - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - } - - // --- GRAPH.* graph commands --- - #[cfg(feature = "graph")] - if cmd.len() > 6 && cmd[..6].eq_ignore_ascii_case(b"GRAPH.") { - let (response, wal_records, cypher_intents, cypher_undo_ops) = if crate::command::graph::is_graph_write_cmd(cmd) - || (cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") - && crate::command::graph::is_cypher_write_query(cmd_args)) - { - let mut gs = ctx.shard_databases.graph_store_write(ctx.shard_id); - let (resp, cypher_intents, undo_ops) = if cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") { - // Phase 167 (CYP-01/02): capture Cypher-created - // nodes/edges so TXN.ABORT can roll them back via - // CrossStoreTxn::record_graph. - crate::command::graph::graph_query_or_write(&mut gs, cmd_args) - } else { - ( - crate::command::graph::dispatch_graph_write(&mut gs, cmd, cmd_args), - Vec::new(), - Vec::new(), - ) - }; - let records = gs.drain_wal(); - (resp, records, cypher_intents, undo_ops) - } else { - let gs = ctx.shard_databases.graph_store_read(ctx.shard_id); - let resp = crate::command::graph::dispatch_graph_read(&gs, cmd, cmd_args); - (resp, Vec::new(), Vec::new(), Vec::new()) - }; - // Phase 166: record graph intent for TXN rollback. - // Captures explicit ADDNODE/ADDEDGE by response id plus - // Phase 167 Cypher CREATE/MERGE via intents returned from - // graph_query_or_write. - if let Some(txn) = conn.active_cross_txn.as_mut() { - let is_node = cmd.eq_ignore_ascii_case(b"GRAPH.ADDNODE"); - let is_edge = cmd.eq_ignore_ascii_case(b"GRAPH.ADDEDGE"); - if is_node || is_edge { - if let Frame::Integer(id) = &response { - if let Some(gname) = cmd_args.first().and_then(|f| extract_bytes(f)) { - txn.record_graph(*id as u64, is_node, gname); - } - } - } - if !cypher_intents.is_empty() { - if let Some(gname) = cmd_args.first().and_then(|f| extract_bytes(f)) { - for intent in &cypher_intents { - txn.record_graph(intent.entity_id, intent.is_node, gname.clone()); - } - } - } - // Phase 174 FIX-01: push undo ops for SET/DELETE/MERGE rollback. - for undo_op in cypher_undo_ops { - txn.record_graph_undo(undo_op); - } - } - for record in wal_records { - ctx.shard_databases.wal_append(ctx.shard_id, bytes::Bytes::from(record)); - } - let mut response = response; - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - - // --- Multi-key commands --- - if is_multi_key_command(cmd, cmd_args) { - let response = crate::shard::coordinator::coordinate_multi_key(cmd, cmd_args, ctx.shard_id, ctx.num_shards, conn.selected_db, &ctx.shard_databases, &ctx.dispatch_tx, &ctx.spsc_notifiers, &ctx.cached_clock, &()).await; - responses.push(response); - continue; - } - - // --- Routing: keyless, local, or remote --- - let target_shard = extract_primary_key(cmd, cmd_args).map(|key| key_to_shard(key, ctx.num_shards)); - let is_local = match target_shard { - None => true, - Some(s) if s == ctx.shard_id => true, - _ => false, - }; - - // Affinity sampling: record shard target for migration decision. - // Only sample when we have a concrete target shard (key-bearing command). - // Migration is deferred until AFTER the current batch is fully processed - // and all responses are written, ensuring no command/response desync. - if let (Some(tracker), Some(target)) = (&mut conn.affinity_tracker, target_shard) { - if let Some(migrate_to) = tracker.record(target) { - // Migration preconditions: not in MULTI, no active CLIENT TRACKING - // (tracking connections need untrack_all cleanup which doesn't transfer) - if !conn.in_multi && !conn.tracking_state.enabled { - conn.migration_target = Some(migrate_to); - } - } - } - - let is_write = if ctx.aof_tx.is_some() || conn.tracking_state.enabled { metadata::is_write(cmd) } else { false }; - let aof_bytes = if is_write && ctx.aof_tx.is_some() { Some(aof::serialize_command(&frame)) } else { None }; - - if is_local { - // LOCAL PATH: split into read/write to avoid exclusive lock on reads. - // Using read_db for local reads eliminates RwLock contention with - // cross-shard shared reads from other shard threads. - if metadata::is_write(cmd) { - // WRITE PATH: single lock acquisition for eviction + dispatch - let rt = ctx.runtime_config.read(); - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - if let Err(oom_frame) = try_evict_if_needed(&mut guard, &rt) { - drop(guard); - drop(rt); - responses.push(oom_frame); - continue; - } - drop(rt); - - // KV undo-log capture for active cross-store transactions. - // MUST happen BEFORE dispatch() overwrites the database entry. - if let Some(ref mut txn) = conn.active_cross_txn { - if cmd.eq_ignore_ascii_case(b"DEL") || cmd.eq_ignore_ascii_case(b"UNLINK") { - // Multi-key DEL: iterate all args for undo capture - for arg in cmd_args.iter() { - if let Frame::BulkString(key_bytes) = arg { - if let Some(old_entry) = guard.get(key_bytes.as_ref()).cloned() { - txn.kv_undo.record_delete(key_bytes.clone(), old_entry); - let lsn = txn.snapshot_lsn; - let tid = txn.txn_id; - ctx.shard_databases.kv_intents(ctx.shard_id) - .record_write(key_bytes.clone(), lsn, tid); - } - // Key not found: nothing to undo, no intent registered - } - } - } else if let Some(key) = crate::server::conn::shared::extract_primary_key(cmd, cmd_args) { - // SET / HSET / INCR / etc. — single primary key - let old_entry = guard.get(key.as_ref()).cloned(); - let lsn = txn.snapshot_lsn; - let tid = txn.txn_id; - match old_entry { - None => txn.kv_undo.record_insert(key.clone()), - Some(entry) => txn.kv_undo.record_update(key.clone(), entry), - } - ctx.shard_databases.kv_intents(ctx.shard_id) - .record_write(key.clone(), lsn, tid); - } - } - - let db_count = ctx.shard_databases.db_count(); - guard.refresh_now_from_cache(&ctx.cached_clock); - let dispatch_start = std::time::Instant::now(); - let result = dispatch(&mut guard, cmd, cmd_args, &mut conn.selected_db, db_count); - let elapsed_us = dispatch_start.elapsed().as_micros() as u64; - if let Ok(cmd_str) = std::str::from_utf8(cmd) { - crate::admin::metrics_setup::record_command(cmd_str, elapsed_us); - } - if let Frame::Array(ref args) = frame { - crate::admin::metrics_setup::global_slowlog().maybe_record( - elapsed_us, - args.as_slice(), - peer_addr.as_bytes(), - conn.client_name.as_ref().map_or(b"" as &[u8], |n| n.as_ref()), - ); - } - let response = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => { should_quit = true; f } - }; - if matches!(response, Frame::Error(_)) { - if let Ok(cmd_str) = std::str::from_utf8(cmd) { - crate::admin::metrics_setup::record_command_error(cmd_str); - } - } else { - let needs_wake = cmd.eq_ignore_ascii_case(b"LPUSH") || cmd.eq_ignore_ascii_case(b"RPUSH") - || cmd.eq_ignore_ascii_case(b"LMOVE") || cmd.eq_ignore_ascii_case(b"ZADD"); - if needs_wake { - if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { - let mut reg = ctx.blocking_registry.borrow_mut(); - if cmd.eq_ignore_ascii_case(b"LPUSH") || cmd.eq_ignore_ascii_case(b"RPUSH") || cmd.eq_ignore_ascii_case(b"LMOVE") { - crate::blocking::wakeup::try_wake_list_waiter(&mut reg, &mut guard, conn.selected_db, &key); - } else { - crate::blocking::wakeup::try_wake_zset_waiter(&mut reg, &mut guard, conn.selected_db, &key); - } - } - } - } - drop(guard); - // Auto-index vectors on successful HSET (local write path) - // Placed AFTER drop(guard) to avoid DB→vector_store lock order - // inversion with the shard event loop (vector_store→DB). - if !matches!(response, Frame::Error(_)) - && (cmd.eq_ignore_ascii_case(b"HSET") || cmd.eq_ignore_ascii_case(b"HMSET")) - { - if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { - // Phase 166 (Plan 02): if inside an active TXN, tag each - // inserted mutable-HNSW entry with the txn_id so that - // non-TXN readers (snapshot_lsn==0) see it as uncommitted - // and exclude it until TXN.COMMIT calls - // txn_manager.commit(txn_id) (ACID-09 isolation fix). - let active_txn_id = conn - .active_cross_txn - .as_ref() - .map(|t| t.txn_id) - .unwrap_or(0); - let inserted = { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - let mut ts = ctx.shard_databases.text_store(ctx.shard_id); - if active_txn_id != 0 { - crate::shard::spsc_handler::auto_index_hset_public_txn(&mut vs, &mut *ts, &key, cmd_args, active_txn_id) - } else { - crate::shard::spsc_handler::auto_index_hset_public(&mut vs, &mut *ts, &key, cmd_args) - } - }; - // Push one VectorIntent per (index_name, key_hash) so - // TXN.ABORT (Plan 166-03) can tombstone via - // MutableSegment::mark_deleted_by_key_hash. - if let Some(txn) = conn.active_cross_txn.as_mut() { - for (index_name, key_hash) in inserted { - txn.record_vector(key_hash, index_name); - } - } - } - } - // Auto-delete vectors on DEL/UNLINK (local write path) - // Note: HDEL removes fields, not keys — it should NOT trigger - // vector deletion unless the entire key is removed. - if !matches!(response, Frame::Error(_)) - && (cmd.eq_ignore_ascii_case(b"DEL") || cmd.eq_ignore_ascii_case(b"UNLINK")) - { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - for arg in cmd_args.iter() { - if let Some(key) = extract_bytes(arg) { - vs.mark_deleted_for_key(key.as_ref()); - } - } - } - if let Some(bytes) = aof_bytes { - if !matches!(response, Frame::Error(_)) { - if let Some(ref tx) = ctx.aof_tx { let _ = tx.try_send(AofMessage::Append(bytes)); } - } - } - if conn.tracking_state.enabled && !matches!(response, Frame::Error(_)) { - if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { - let senders = ctx.tracking_table.borrow_mut().invalidate_key(&key, client_id); - if !senders.is_empty() { - let push = crate::tracking::invalidation::invalidation_push(&[key]); - for tx in senders { let _ = tx.try_send(push.clone()); } - } - } - } - let mut response = apply_resp3_conversion(cmd, response, conn.protocol_version); - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - } else { - // Snapshot visibility filter for active cross-store transactions. - // MVCC: hide keys written by uncommitted foreign transactions. - if conn.in_cross_txn() { - if let Some(ref txn) = conn.active_cross_txn { - if let Some(key) = crate::server::conn::shared::extract_primary_key(cmd, cmd_args) { - let snapshot_lsn = txn.snapshot_lsn; - let my_txn_id = txn.txn_id; - // Clone committed treemap to release vector_store lock - // before acquiring kv_intents lock (lock ordering). - let committed = { - let vs = ctx.shard_databases.vector_store(ctx.shard_id); - vs.txn_manager().committed_treemap().clone() - }; - let visible = { - let intents = ctx.shard_databases.kv_intents(ctx.shard_id); - intents.is_key_visible(key.as_ref(), snapshot_lsn, my_txn_id, &committed) - }; - if !visible { - responses.push(Frame::Null); - continue; - } - } - } - } - - // READ PATH: shared lock — no contention with other shards' reads - let guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); - let now_ms = ctx.cached_clock.ms(); - let db_count = ctx.shard_databases.db_count(); - let dispatch_start = std::time::Instant::now(); - let result = dispatch_read(&guard, cmd, cmd_args, now_ms, &mut conn.selected_db, db_count); - let elapsed_us = dispatch_start.elapsed().as_micros() as u64; - if let Ok(cmd_str) = std::str::from_utf8(cmd) { - crate::admin::metrics_setup::record_command(cmd_str, elapsed_us); - } - if let Frame::Array(ref args) = frame { - crate::admin::metrics_setup::global_slowlog().maybe_record( - elapsed_us, - args.as_slice(), - peer_addr.as_bytes(), - conn.client_name.as_ref().map_or(b"" as &[u8], |n| n.as_ref()), - ); - } - drop(guard); - let response = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => { should_quit = true; f } - }; - if matches!(response, Frame::Error(_)) { - if let Ok(cmd_str) = std::str::from_utf8(cmd) { - crate::admin::metrics_setup::record_command_error(cmd_str); - } - } - if conn.tracking_state.enabled && !conn.tracking_state.bcast { - if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { - ctx.tracking_table.borrow_mut().track_key(client_id, &key, conn.tracking_state.noloop); - } - } - let mut response = apply_resp3_conversion(cmd, response, conn.protocol_version); - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - } - } else if let Some(target) = target_shard { - // TXN cross-shard guard: cross-shard writes bypass the undo log and - // cannot be rolled back on TXN.ABORT. Return an explicit error instead - // of silently permitting writes that resist rollback. - if conn.in_cross_txn() && metadata::is_write(cmd) { - responses.push(Frame::Error(bytes::Bytes::from_static( - crate::command::transaction::ERR_TXN_CROSS_SHARD, - ))); - continue; - } - // SHARED-READ FAST PATH: cross-shard reads bypass SPSC dispatch entirely. - // By this point conn.in_multi is false (MULTI queuing happens earlier with `continue`). - // Read commands execute directly on the target shard's database via RwLock read guard, - // avoiding ~88us of two async scheduling hops through the SPSC channel. - // - // Guard: if there are already pending writes for this target shard in the - // current pipeline batch, we must NOT take the fast path -- the read would - // execute before the deferred writes, violating command ordering. Fall through - // to SPSC dispatch to preserve pipeline semantics. - if !metadata::is_write(cmd) && !remote_groups.contains_key(&target) { - let guard = ctx.shard_databases.read_db(target, conn.selected_db); - let now_ms = ctx.cached_clock.ms(); - let db_count = ctx.shard_databases.db_count(); - let result = dispatch_read(&guard, cmd, cmd_args, now_ms, &mut conn.selected_db, db_count); - drop(guard); - let response = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => { should_quit = true; f } - }; - if matches!(response, Frame::Error(_)) { - if let Ok(cmd_str) = std::str::from_utf8(cmd) { - crate::admin::metrics_setup::record_command_error(cmd_str); - } - } - // Client tracking for cross-shard reads - if conn.tracking_state.enabled && !conn.tracking_state.bcast { - if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { - ctx.tracking_table.borrow_mut().track_key(client_id, &key, conn.tracking_state.noloop); - } - } - let mut response = apply_resp3_conversion(cmd, response, conn.protocol_version); - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - // Cross-shard write: deferred SPSC dispatch. - // When workspace rewriting occurred, rebuild the frame with - // prefixed args so the target shard stores the correct key. - let dispatch_frame = if rewritten.is_some() { - let mut parts = Vec::with_capacity(1 + cmd_args.len()); - parts.push(Frame::BulkString(Bytes::copy_from_slice(cmd))); - parts.extend_from_slice(cmd_args); - Frame::Array(parts.into()) - } else { - frame - }; - let resp_idx = responses.len(); - responses.push(Frame::Null); - let cmd_bytes = if let Frame::Array(ref args) = dispatch_frame { - extract_bytes(&args[0]).unwrap_or_default() - } else { - Bytes::new() - }; - remote_groups.entry(target).or_default().push((resp_idx, std::sync::Arc::new(dispatch_frame), aof_bytes, cmd_bytes, conn.selected_db)); - } - } - - // Phase 2: Dispatch deferred remote commands (zero-allocation via ResponseSlotPool) - if !remote_groups.is_empty() { - let mut reply_futures: Vec<(Vec<(usize, Option, Bytes)>, usize)> = Vec::with_capacity(remote_groups.len()); - for (target, entries) in remote_groups { - let slot_ptr = response_pool.slot_ptr(target); - // Use the db_index captured with the first command (all commands in a - // pipeline batch targeting the same shard share the same db_index). - let batch_db = entries.first().map(|(_, _, _, _, db)| *db).unwrap_or(conn.selected_db); - let (meta, commands): (Vec<(usize, Option, Bytes)>, Vec>) = - entries.into_iter().map(|(idx, arc_frame, aof, cmd, _db)| ((idx, aof, cmd), arc_frame)).unzip(); - let msg = ShardMessage::PipelineBatchSlotted { db_index: batch_db, commands, response_slot: crate::shard::dispatch::ResponseSlotPtr(slot_ptr) }; - let target_idx = ChannelMesh::target_index(ctx.shard_id, target); - { - let mut pending = msg; - loop { - let push_result = { let mut producers = ctx.dispatch_tx.borrow_mut(); producers[target_idx].try_push(pending) }; - match push_result { - Ok(()) => { ctx.spsc_notifiers[target].notify_one(); break; } - Err(val) => { pending = val; tokio::task::yield_now().await; } - } - } - } - reply_futures.push((meta, target)); - } - let proto_ver = conn.protocol_version; - for (meta, target) in reply_futures { - let shard_responses = response_pool.future_for(target).await; - for ((resp_idx, aof_bytes, cmd_name), resp) in meta.into_iter().zip(shard_responses) { - if let Some(bytes) = aof_bytes { - if !matches!(resp, Frame::Error(_)) { - if let Some(ref tx) = ctx.aof_tx { let _ = tx.try_send(AofMessage::Append(bytes)); } - } - } - responses[resp_idx] = apply_resp3_conversion(&cmd_name, resp, proto_ver); - } - } - } - - // Phase 3: Flush accumulated PUBLISH batches as PubSubPublishBatch messages - if !publish_batches.is_empty() { - let mut batch_slots: Vec<(std::sync::Arc, Vec)> = Vec::new(); - { - let mut producers = ctx.dispatch_tx.borrow_mut(); - for (target, entries) in publish_batches.drain() { - let n = entries.len(); - let slot = std::sync::Arc::new(crate::shard::dispatch::PubSubResponseSlot::with_counts(1, n)); - let resp_indices: Vec = entries.iter().map(|(idx, _, _)| *idx).collect(); - let pairs: Vec<(Bytes, Bytes)> = entries.into_iter().map(|(_, ch, msg)| (ch, msg)).collect(); - - let idx = ChannelMesh::target_index(ctx.shard_id, target); - let batch_msg = ShardMessage::PubSubPublishBatch { - pairs, - slot: slot.clone(), - }; - if producers[idx].try_push(batch_msg).is_ok() { - ctx.spsc_notifiers[target].notify_one(); - } else { - slot.add(0); // push failed, mark as done - } - batch_slots.push((slot, resp_indices)); - } - } - // Resolve all batch slots - for (slot, resp_indices) in &batch_slots { - crate::shard::dispatch::PubSubResponseFuture::new(slot.clone()).await; - for (i, resp_idx) in resp_indices.iter().enumerate() { - let remote_count = slot.counts[i].load(std::sync::atomic::Ordering::Relaxed); - if remote_count > 0 { - if let Frame::Integer(ref mut total) = responses[*resp_idx] { *total += remote_count; } - } - } - } - } - - arena.reset(); - - // AUTH rate limiting: delay response to slow down brute-force attacks - if auth_delay_ms > 0 { - tokio::time::sleep(std::time::Duration::from_millis(auth_delay_ms)).await; - } - - write_buf.clear(); - for response in &responses { - if conn.protocol_version >= 3 { - crate::protocol::serialize_resp3(response, &mut write_buf); - } else { - crate::protocol::serialize(response, &mut write_buf); - } - } - if stream.write_all(&write_buf).await.is_err() { - return (HandlerResult::Done, None); - } - - // Update registry with current state after each batch - crate::client_registry::update(client_id, |e| { - e.db = conn.selected_db; - e.last_cmd_at = std::time::Instant::now(); - e.flags = crate::client_registry::ClientFlags { - subscriber: conn.subscription_count > 0, - in_multi: conn.in_multi, - blocked: false, - }; - }); - - // Check if migration was triggered during frame processing. - // All responses for the current batch have been written, so the - // client sees no interruption -- TCP socket stays open. - if let Some(target_shard) = conn.migration_target { - let migrated_state = MigratedConnectionState { - selected_db: conn.selected_db, - authenticated: conn.authenticated, - client_name: conn.client_name.clone(), - protocol_version: conn.protocol_version, - current_user: conn.current_user.clone(), - flags: 0, - read_buf_remainder: read_buf.split(), - client_id, - peer_addr: peer_addr.clone(), - workspace_id: conn.workspace_id, - }; - return ( - HandlerResult::MigrateConnection { state: migrated_state, target_shard }, - Some(stream), - ); - } - - if write_buf.capacity() > 65536 { write_buf = BytesMut::with_capacity(8192); } - if read_buf.capacity() > 65536 { - let remaining = read_buf.split(); - read_buf = BytesMut::with_capacity(8192); - if !remaining.is_empty() { read_buf.extend_from_slice(&remaining); } - } - - if should_quit { break; } - } - _ = shutdown.cancelled() => { - write_buf.clear(); - let shutdown_err = Frame::Error(Bytes::from_static(b"ERR server shutting down")); - if conn.protocol_version >= 3 { - crate::protocol::serialize_resp3(&shutdown_err, &mut write_buf); - } else { - crate::protocol::serialize(&shutdown_err, &mut write_buf); - } - let _ = stream.write_all(&write_buf).await; - break; - } - } - } - - // Phase 166: release any leaked cross-store TXN (client disconnected mid-txn). - // Idempotent: TXN.ABORT already takes() active_cross_txn so this is a no-op if abort ran. - // Closes T-161-05 — without this, a disconnect after TXN.BEGIN + SET would leak - // kv_intents and pin the key invisible for all subsequent readers. - if let Some(txn) = conn.active_cross_txn.take() { - crate::transaction::abort::abort_cross_store_txn( - &ctx.shard_databases, - ctx.shard_id, - conn.selected_db, - txn, - ); - } - - // Clean up pub/sub subscriptions on disconnect - if conn.subscriber_id > 0 { - let removed_channels = { - ctx.pubsub_registry - .write() - .unsubscribe_all(conn.subscriber_id) - }; - let removed_patterns = { - ctx.pubsub_registry - .write() - .punsubscribe_all(conn.subscriber_id) - }; - for ch in removed_channels { - unpropagate_subscription( - &ctx.all_remote_sub_maps, - &ch, - ctx.shard_id, - ctx.num_shards, - false, - ); - } - for pat in removed_patterns { - unpropagate_subscription( - &ctx.all_remote_sub_maps, - &pat, - ctx.shard_id, - ctx.num_shards, - true, - ); - } - // Remove affinity on disconnect (no subscriptions remain) - if let Ok(addr) = peer_addr.parse::() { - ctx.pubsub_affinity.write().remove(&addr.ip()); - } - } - - if conn.tracking_state.enabled { - ctx.tracking_table.borrow_mut().untrack_all(client_id); - } - - (HandlerResult::Done, None) -} diff --git a/src/server/conn/handler_sharded/dispatch.rs b/src/server/conn/handler_sharded/dispatch.rs new file mode 100644 index 00000000..25ae3fe7 --- /dev/null +++ b/src/server/conn/handler_sharded/dispatch.rs @@ -0,0 +1,464 @@ +//! Connection-level command dispatchers: CLIENT subcommands, CONFIG, SLOWLOG, +//! REPLICAOF/REPLCONF, INFO, READONLY, BGSAVE/SAVE/LASTSAVE/BGREWRITEAOF, +//! cross-shard KEYS/SCAN/DBSIZE. +//! +//! Each helper returns `true` if the command was consumed (caller should `continue`). + +use bytes::Bytes; +use std::sync::Arc; + +use crate::command::connection as conn_cmd; +use crate::command::metadata; +use crate::protocol::Frame; +use crate::runtime::channel; +use crate::server::conn::core::{ConnectionContext, ConnectionState}; +use crate::server::conn::util::extract_bytes; +use crate::tracking::TrackingState; +use crate::workspace::strip_workspace_prefix_from_response; + +use super::handle_config; + +/// Handle CLIENT subcommands. Returns `true` if consumed. +pub(super) fn try_handle_client_command( + cmd: &[u8], + cmd_args: &[Frame], + client_id: u64, + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"CLIENT") { + return false; + } + if let Some(sub) = cmd_args.first() { + if let Some(sub_bytes) = extract_bytes(sub) { + if sub_bytes.eq_ignore_ascii_case(b"ID") { + responses.push(conn_cmd::client_id(client_id)); + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"SETNAME") { + if cmd_args.len() != 2 { + responses.push(Frame::Error(Bytes::from_static( + b"ERR wrong number of arguments for 'CLIENT SETNAME' command", + ))); + } else { + conn.client_name = extract_bytes(&cmd_args[1]); + let name_str = conn + .client_name + .as_ref() + .map(|b| String::from_utf8_lossy(b).to_string()); + crate::client_registry::update(client_id, |e| { + e.name = name_str; + }); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"GETNAME") { + responses.push(match &conn.client_name { + Some(name) => Frame::BulkString(name.clone()), + None => Frame::Null, + }); + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"TRACKING") { + match crate::command::client::parse_tracking_args(cmd_args) { + Ok(config_parsed) => { + if config_parsed.enable { + conn.tracking_state.enabled = true; + conn.tracking_state.bcast = config_parsed.bcast; + conn.tracking_state.noloop = config_parsed.noloop; + conn.tracking_state.optin = config_parsed.optin; + conn.tracking_state.optout = config_parsed.optout; + if conn.tracking_rx.is_none() { + let (tx, rx) = channel::mpsc_bounded::(256); + conn.tracking_state.invalidation_tx = Some(tx.clone()); + conn.tracking_rx = Some(rx); + let mut table = ctx.tracking_table.borrow_mut(); + table.register_client(client_id, tx); + if let Some(target) = config_parsed.redirect { + table.set_redirect(client_id, target); + } + for prefix in &config_parsed.prefixes { + table.register_prefix( + client_id, + prefix.clone(), + config_parsed.noloop, + ); + } + } + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } else { + conn.tracking_state = TrackingState::default(); + ctx.tracking_table.borrow_mut().untrack_all(client_id); + conn.tracking_rx = None; + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + return true; + } + Err(err_frame) => { + responses.push(err_frame); + return true; + } + } + } + if sub_bytes.eq_ignore_ascii_case(b"LIST") { + // Update our own entry before listing + crate::client_registry::update(client_id, |e| { + e.db = conn.selected_db; + e.last_cmd_at = std::time::Instant::now(); + e.flags = crate::client_registry::ClientFlags { + subscriber: conn.subscription_count > 0, + in_multi: conn.in_multi, + blocked: false, + }; + }); + let list = crate::client_registry::client_list(); + responses.push(Frame::BulkString(Bytes::from(list))); + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"INFO") { + crate::client_registry::update(client_id, |e| { + e.db = conn.selected_db; + e.last_cmd_at = std::time::Instant::now(); + }); + let info = crate::client_registry::client_info(client_id).unwrap_or_default(); + responses.push(Frame::BulkString(Bytes::from(info))); + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"KILL") { + let raw_args: Vec<&[u8]> = cmd_args[1..] + .iter() + .filter_map(|f| match f { + Frame::BulkString(b) => Some(b.as_ref()), + Frame::SimpleString(b) => Some(b.as_ref()), + _ => None, + }) + .collect(); + match crate::client_registry::parse_kill_args(&raw_args) { + Some(filter) => { + let count = crate::client_registry::kill_clients(&filter); + responses.push(Frame::Integer(count as i64)); + } + None => { + responses.push(Frame::Error(Bytes::from_static( + b"ERR syntax error. Usage: CLIENT KILL [ID id] [ADDR addr] [USER user]", + ))); + } + } + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"PAUSE") { + if cmd_args.len() < 2 { + responses.push(Frame::Error(Bytes::from_static( + b"ERR wrong number of arguments for 'CLIENT PAUSE' command", + ))); + } else { + let timeout_bytes = match &cmd_args[1] { + Frame::BulkString(b) => Some(b.as_ref()), + Frame::SimpleString(b) => Some(b.as_ref()), + _ => None, + }; + match timeout_bytes + .and_then(|b| std::str::from_utf8(b).ok()) + .and_then(|s| s.parse::().ok()) + { + Some(ms) => { + let mode = if cmd_args.len() > 2 { + match &cmd_args[2] { + Frame::BulkString(b) | Frame::SimpleString(b) + if b.eq_ignore_ascii_case(b"WRITE") => + { + crate::client_pause::PauseMode::Write + } + _ => crate::client_pause::PauseMode::All, + } + } else { + crate::client_pause::PauseMode::All + }; + crate::client_pause::pause(ms, mode); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + None => { + responses.push(Frame::Error(Bytes::from_static( + b"ERR timeout is not a valid integer or out of range", + ))); + } + } + } + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"UNPAUSE") { + crate::client_pause::unpause(); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + return true; + } + if sub_bytes.eq_ignore_ascii_case(b"NO-EVICT") + || sub_bytes.eq_ignore_ascii_case(b"NO-TOUCH") + { + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + return true; + } + responses.push(Frame::Error(Bytes::from(format!( + "ERR unknown subcommand '{}'", + String::from_utf8_lossy(&sub_bytes) + )))); + return true; + } + } + responses.push(Frame::Error(Bytes::from_static( + b"ERR wrong number of arguments for 'client' command", + ))); + true +} + +/// Handle CONFIG command. Returns `true` if consumed. +pub(super) fn try_handle_config( + cmd: &[u8], + cmd_args: &[Frame], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"CONFIG") { + return false; + } + responses.push(handle_config(cmd_args, &ctx.runtime_config, &ctx.config)); + true +} + +/// Handle REPLICAOF / SLAVEOF. Returns `true` if consumed. +pub(super) fn try_handle_replicaof( + cmd: &[u8], + cmd_args: &[Frame], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"REPLICAOF") && !cmd.eq_ignore_ascii_case(b"SLAVEOF") { + return false; + } + use crate::command::connection::{ReplicaofAction, replicaof}; + let (resp, action) = replicaof(cmd_args); + if let Some(action) = action { + if let Some(ref rs) = ctx.repl_state { + match action { + ReplicaofAction::StartReplication { host, port } => { + if let Ok(mut rs_guard) = rs.write() { + rs_guard.role = crate::replication::state::ReplicationRole::Replica { + host: host.clone(), + port, + state: + crate::replication::handshake::ReplicaHandshakeState::PingPending, + }; + } + let rs_clone = Arc::clone(rs); + let cfg = crate::replication::replica::ReplicaTaskConfig { + master_host: host, + master_port: port, + repl_state: rs_clone, + num_shards: ctx.num_shards, + persistence_dir: None, + listening_port: 0, + }; + tokio::task::spawn_local(crate::replication::replica::run_replica_task(cfg)); + } + ReplicaofAction::PromoteToMaster => { + use crate::replication::state::generate_repl_id; + if let Ok(mut rs_guard) = rs.write() { + rs_guard.repl_id2 = rs_guard.repl_id.clone(); + rs_guard.repl_id = generate_repl_id(); + rs_guard.role = crate::replication::state::ReplicationRole::Master; + } + } + ReplicaofAction::NoOp => {} + } + } + } + responses.push(resp); + true +} + +/// Handle INFO command. Returns `true` if consumed. +pub(super) fn try_handle_info( + cmd: &[u8], + cmd_args: &[Frame], + conn: &ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"INFO") { + return false; + } + let guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); + let response_text = { + let resp_frame = conn_cmd::info_readonly(&guard, cmd_args); + match resp_frame { + Frame::BulkString(b) => String::from_utf8_lossy(&b).to_string(), + _ => String::new(), + } + }; + drop(guard); + let mut response_text = response_text; + if let Some(ref rs) = ctx.repl_state { + if let Ok(rs_guard) = rs.try_read() { + response_text.push_str(&crate::replication::handshake::build_info_replication( + &rs_guard, + )); + } + } + responses.push(Frame::BulkString(Bytes::from(response_text))); + true +} + +/// Handle persistence commands (BGSAVE, SAVE, LASTSAVE, BGREWRITEAOF). +/// Returns `true` if consumed. +pub(super) fn try_handle_persistence( + cmd: &[u8], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if cmd.eq_ignore_ascii_case(b"BGSAVE") { + responses.push(crate::command::persistence::bgsave_start_sharded( + &ctx.snapshot_trigger_tx, + ctx.num_shards, + )); + return true; + } + if cmd.eq_ignore_ascii_case(b"SAVE") { + responses.push(Frame::Error(Bytes::from_static( + b"ERR SAVE not supported in sharded mode, use BGSAVE", + ))); + return true; + } + if cmd.eq_ignore_ascii_case(b"LASTSAVE") { + responses.push(crate::command::persistence::handle_lastsave()); + return true; + } + if cmd.eq_ignore_ascii_case(b"BGREWRITEAOF") { + if let Some(ref tx) = ctx.aof_tx { + responses.push(crate::command::persistence::bgrewriteaof_start_sharded( + tx, + ctx.shard_databases.clone(), + )); + } else { + responses.push(Frame::Error(Bytes::from_static(b"ERR AOF is not enabled"))); + } + return true; + } + false +} + +/// Handle cross-shard KEYS, SCAN, DBSIZE aggregation. +/// Returns `true` if consumed. +pub(super) async fn try_handle_cross_shard_scan( + cmd: &[u8], + cmd_args: &[Frame], + conn: &ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if cmd.eq_ignore_ascii_case(b"KEYS") { + let mut response = crate::shard::coordinator::coordinate_keys( + cmd_args, + ctx.shard_id, + ctx.num_shards, + conn.selected_db, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + &ctx.cached_clock, + &(), + ) + .await; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + if cmd.eq_ignore_ascii_case(b"SCAN") { + let mut response = crate::shard::coordinator::coordinate_scan( + cmd_args, + ctx.shard_id, + ctx.num_shards, + conn.selected_db, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + &ctx.cached_clock, + &(), + ) + .await; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + if cmd.eq_ignore_ascii_case(b"DBSIZE") { + let response = crate::shard::coordinator::coordinate_dbsize( + ctx.shard_id, + ctx.num_shards, + conn.selected_db, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + &(), + ) + .await; + responses.push(response); + return true; + } + false +} + +/// Handle READONLY enforcement for replicas. +/// Returns `true` if the command was blocked (caller should `continue`). +pub(super) fn try_enforce_readonly( + cmd: &[u8], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if let Some(ref rs) = ctx.repl_state { + if let Ok(rs_guard) = rs.try_read() { + if matches!( + rs_guard.role, + crate::replication::state::ReplicationRole::Replica { .. } + ) { + if metadata::is_write(cmd) { + responses.push(Frame::Error(Bytes::from_static( + b"READONLY You can't write against a read only replica.", + ))); + return true; + } + } + } + } + false +} + +/// Handle SLOWLOG command. Returns `true` if consumed. +pub(super) fn try_handle_slowlog( + cmd: &[u8], + cmd_args: &[Frame], + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"SLOWLOG") { + return false; + } + let sl = crate::admin::metrics_setup::global_slowlog(); + responses.push(crate::admin::slowlog::handle_slowlog(sl, cmd_args)); + true +} + +/// Handle REPLCONF command. Returns `true` if consumed. +pub(super) fn try_handle_replconf( + cmd: &[u8], + cmd_args: &[Frame], + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"REPLCONF") { + return false; + } + responses.push(crate::command::connection::replconf(cmd_args)); + true +} diff --git a/src/server/conn/handler_sharded/ft.rs b/src/server/conn/handler_sharded/ft.rs new file mode 100644 index 00000000..67c83155 --- /dev/null +++ b/src/server/conn/handler_sharded/ft.rs @@ -0,0 +1,705 @@ +//! FT.* vector/text search command handlers. +//! +//! Handles FT.CREATE, FT.DROPINDEX, FT.SEARCH (vector + text), FT.AGGREGATE, +//! FT.INFO, FT.COMPACT, FT._LIST, FT.CACHESEARCH, FT.CONFIG, FT.RECOMMEND, +//! FT.NAVIGATE, FT.EXPAND. Multi-shard scatter-gather and single-shard fast paths. + +use bytes::Bytes; + +use crate::protocol::Frame; +use crate::server::conn::core::{ConnectionContext, ConnectionState}; +use crate::server::conn::shared::resolve_ft_search_as_of_lsn; +use crate::workspace::strip_workspace_prefix_from_response; + +/// Handle FT.* commands. Returns `true` if the command was consumed. +/// +/// Caller should `continue` the frame loop when this returns `true`. +pub(super) async fn try_handle_ft_command( + cmd: &[u8], + cmd_args: &[Frame], + frame: &Frame, + conn: &ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if cmd.len() <= 3 || !cmd[..3].eq_ignore_ascii_case(b"FT.") { + return false; + } + + if ctx.num_shards > 1 { + // Multi-shard: dispatch via SPSC + #[cfg(feature = "text-index")] + if cmd.eq_ignore_ascii_case(b"FT.AGGREGATE") { + // -- FT.AGGREGATE: two-phase scatter-gather per Plan 03 (D-05/D-07) -- + // scatter_text_aggregate acquires its own guards internally + // inside the single-shard block, so we never hold a MutexGuard + // across the .await below. + let parsed = + match crate::command::vector_search::ft_aggregate::parse_aggregate_args(cmd_args) { + Ok(p) => p, + Err(err_frame) => { + responses.push(err_frame); + return true; + } + }; + let response = crate::shard::scatter_aggregate::scatter_text_aggregate( + parsed.index_name, + parsed.query, + parsed.pipeline, + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await; + responses.push(response); + return true; + } + if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { + // Check if this is a text query BEFORE trying parse_ft_search_args + // (which would return an error for non-KNN queries). + let query_bytes = cmd_args + .get(1) + .and_then(|f| crate::command::vector_search::extract_bulk(f)); + let is_text = query_bytes + .as_ref() + .map_or(false, |q| crate::command::vector_search::is_text_query(q)); + + // -- HYBRID multi-shard path (Phase 152 Plan 05, D-13) -- + // If the args contain a HYBRID clause, route through + // scatter_hybrid_search (three-phase DFS -> fan-out -> RRF + // merge). Single-shard is handled inside the scatter entry + // point itself (fast path to execute_hybrid_search_local). + #[cfg(feature = "text-index")] + { + match crate::command::vector_search::hybrid::parse_hybrid_modifier(cmd_args) { + Ok(Some(partial)) => { + let index_name = match cmd_args + .first() + .and_then(|f| crate::command::vector_search::extract_bulk(f)) + { + Some(b) => b, + None => { + responses.push(Frame::Error(Bytes::from_static( + b"ERR invalid index name", + ))); + return true; + } + }; + let text_query = match query_bytes.clone() { + Some(q) => q, + None => { + responses + .push(Frame::Error(Bytes::from_static(b"ERR invalid query"))); + return true; + } + }; + let (limit_offset, limit_count) = + crate::command::vector_search::parse_limit_clause(cmd_args); + let top_k = if limit_count == usize::MAX { + limit_offset.saturating_add(10).max(1) + } else { + limit_offset.saturating_add(limit_count).max(1) + }; + let hq = crate::command::vector_search::hybrid::HybridQuery { + index_name, + text_query, + dense_field: partial.dense_field, + dense_blob: partial.dense_blob, + sparse: partial.sparse, + weights: partial.weights, + k_per_stream: partial.k_per_stream, + top_k, + offset: limit_offset, + count: limit_count, + }; + // Phase 171 HYB-02 / SCAT-02: resolve AS_OF / + // TXN LSN ONCE on the coordinator and forward to + // every responder via the scatter helper. + let as_of_lsn = match resolve_ft_search_as_of_lsn( + cmd_args, + Some(&ctx.shard_databases), + ctx.shard_id, + conn.active_cross_txn.as_ref(), + ) { + Ok(lsn) => lsn, + Err(err_frame) => { + responses.push(err_frame); + return true; + } + }; + let response = crate::shard::scatter_hybrid::scatter_hybrid_search( + hq, + as_of_lsn, + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await; + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + Ok(None) => { /* fall through to non-HYBRID paths */ } + Err(err_frame) => { + responses.push(err_frame); + return true; + } + } + } + + if is_text { + // -- Text FT.SEARCH: two-phase DFS scatter-gather -- + let index_name = match cmd_args + .first() + .and_then(|f| crate::command::vector_search::extract_bulk(f)) + { + Some(b) => b, + None => { + responses.push(Frame::Error(Bytes::from_static(b"ERR invalid index name"))); + return true; + } + }; + #[allow(clippy::unwrap_used)] // query_bytes is Some when is_text is true + let query_str = query_bytes.unwrap(); + + // B-01 SITE 3 FIX (Plan 152-06): FieldFilter short-circuit + // BEFORE the analyzer-first parse_result block. Symmetric with + // handler_monoio. TAG queries route through the InvertedSearch + // fan-out -- no analyzer touched, no field_idx resolution. + #[cfg(feature = "text-index")] + { + match crate::command::vector_search::pre_parse_field_filter(query_str.as_ref()) + { + Ok(Some(clause)) => { + if let Some(filter) = clause.filter { + let (offset, count) = + crate::command::vector_search::parse_limit_clause(cmd_args); + let top_k = if count == usize::MAX { + 10000 + } else { + offset.saturating_add(count) + } + .max(1); + let response = + crate::shard::coordinator::scatter_text_search_filter( + index_name, + filter, + top_k, + offset, + count, + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await; + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + } + Ok(None) => { /* fall through */ } + Err(e) => { + responses.push(Frame::Error(Bytes::from(e.to_owned()))); + return true; + } + } + } + + // Parse query and resolve field_idx inside a block scope so the + // MutexGuard from text_store() is dropped BEFORE .await. + // We use the TextIndex's own field_analyzers (same pipeline used at index time). + type ParseResult = + Result<(Vec, Option), String>; + let parse_result: ParseResult = { + let ts = ctx.shard_databases.text_store(ctx.shard_id); + match ts.get_index(&index_name) { + None => Err("ERR no such index".to_owned()), + Some(text_index) => { + match text_index.field_analyzers.first() { + None => Err("ERR index has no TEXT fields".to_owned()), + Some(analyzer) => { + // analyzer borrows text_index which borrows ts -- all in this block. + let parsed = crate::command::vector_search::parse_text_query( + &query_str, analyzer, + ); + match parsed { + Err(e) => Err(e.to_owned()), + Ok(clause) => { + let field_idx = match &clause.field_name { + None => Ok(None), + Some(field_name) => { + match text_index.text_fields.iter().position( + |f| { + f.field_name + .as_ref() + .eq_ignore_ascii_case( + field_name.as_ref(), + ) + }, + ) { + Some(idx) => Ok(Some(idx)), + None => Err(format!( + "ERR unknown field '{}'", + String::from_utf8_lossy(field_name) + )), + } + } + }; + field_idx.map(|idx| (clause.terms, idx)) + } + } + } + } + } + } + }; // MutexGuard dropped here + + let (query_terms, field_idx) = match parse_result { + Ok(t) => t, + Err(e) => { + responses.push(Frame::Error(Bytes::from(e))); + return true; + } + }; + + let (offset, count) = crate::command::vector_search::parse_limit_clause(cmd_args); + let top_k = if count == usize::MAX { + 10000 + } else { + offset.saturating_add(count) + } + .max(1); + + // Parse optional HIGHLIGHT/SUMMARIZE clauses from args. + let highlight_opts = + crate::command::vector_search::parse_highlight_clause(cmd_args); + let summarize_opts = + crate::command::vector_search::parse_summarize_clause(cmd_args); + + let mut response = crate::shard::coordinator::scatter_text_search( + index_name, + query_terms, + field_idx, + top_k, + offset, + count, + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + highlight_opts, + summarize_opts, + ) + .await; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + + // -- Vector FT.SEARCH (KNN / SPARSE): existing path -- + // Phase 171 SCAT-01: resolve AS_OF / TXN snapshot LSN ONCE + // on the coordinator and thread it through the scatter helper + // so every responder honors the same temporal snapshot. + let response = match crate::command::vector_search::parse_ft_search_args(cmd_args) { + Ok((index_name, query_blob, k, filter, _offset, _count)) => { + if filter.is_some() { + Frame::Error(Bytes::from_static( + b"ERR FILTER not supported in multi-shard mode yet", + )) + } else { + match resolve_ft_search_as_of_lsn( + cmd_args, + Some(&ctx.shard_databases), + ctx.shard_id, + conn.active_cross_txn.as_ref(), + ) { + Err(err_frame) => err_frame, + Ok(as_of_lsn) => { + crate::shard::coordinator::scatter_vector_search_remote( + index_name, + query_blob, + k, + as_of_lsn, + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await + } + } + } + } + Err(err_frame) => err_frame, + }; + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + let response = crate::shard::coordinator::broadcast_vector_command( + std::sync::Arc::new(frame.clone()), + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await; + responses.push(response); + return true; + } + + // Single-shard: no SPSC channels available. + // Dispatch directly to shard's VectorStore via shared access. + // + // -- 154-01 single-shard FT.AGGREGATE fast path -- + // scatter_text_aggregate internally fast-paths num_shards + // == 1 to execute_local_full. + #[cfg(feature = "text-index")] + if cmd.eq_ignore_ascii_case(b"FT.AGGREGATE") { + let parsed = + match crate::command::vector_search::ft_aggregate::parse_aggregate_args(cmd_args) { + Ok(p) => p, + Err(err_frame) => { + responses.push(err_frame); + return true; + } + }; + let response = crate::shard::scatter_aggregate::scatter_text_aggregate( + parsed.index_name, + parsed.query, + parsed.pipeline, + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await; + responses.push(response); + return true; + } + // + // -- 151-03 single-shard text FT.SEARCH fast path -- + // Parity with handler_monoio.rs. Bare text queries bypass + // ft_search() (which only parses KNN/SPARSE/HYBRID) + // and route directly to execute_text_search_local. + #[cfg(feature = "text-index")] + if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { + if let Some(Frame::BulkString(query_bytes)) = cmd_args.get(1) { + match crate::command::vector_search::parse_hybrid_modifier(cmd_args) { + Ok(Some(_)) => { + // HYBRID present -- existing ft_search() below handles it. + } + Err(frame_err) => { + responses.push(frame_err); + return true; + } + Ok(None) => { + if crate::command::vector_search::is_text_query(query_bytes.as_ref()) { + // Step 1: index_name. + let index_name = match cmd_args.first() { + Some(Frame::BulkString(b)) => b.clone(), + _ => { + responses.push(Frame::Error(Bytes::from_static( + b"ERR wrong number of arguments for FT.SEARCH", + ))); + return true; + } + }; + // B-01 SITE 3 FIX (single-shard 151-03 fast path, + // Plan 152-06): FieldFilter short-circuit BEFORE + // text_fields.is_empty() bail. + #[cfg(feature = "text-index")] + match crate::command::vector_search::pre_parse_field_filter( + query_bytes.as_ref(), + ) { + Ok(Some(clause)) => { + if clause.filter.is_some() { + let (offset, count) = + crate::command::vector_search::parse_limit_clause(cmd_args); + let top_k = if count == usize::MAX { + 10000 + } else { + offset.saturating_add(count) + } + .max(1); + let ts_guard = ctx.shard_databases.text_store(ctx.shard_id); + let response = match ts_guard.get_index(&index_name) { + None => { + Frame::Error(Bytes::from_static(b"ERR no such index")) + } + Some(text_index) => { + let results = crate::command::vector_search::ft_text_search::execute_query_on_index( + text_index, &clause, None, None, top_k, + ); + crate::command::vector_search::ft_text_search::build_text_response( + &results, offset, count, + ) + } + }; + drop(ts_guard); + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response( + ws_id, + cmd, + &mut response, + ); + } + responses.push(response); + return true; + } + } + Ok(None) => { /* fall through */ } + Err(e) => { + responses.push(Frame::Error(Bytes::copy_from_slice(e.as_bytes()))); + return true; + } + } + // Step 2: ts guard. + let ts_guard = ctx.shard_databases.text_store(ctx.shard_id); + // Step 3: index lookup. + let text_index = match ts_guard.get_index(&index_name) { + Some(idx) => idx, + None => { + drop(ts_guard); + responses + .push(Frame::Error(Bytes::from_static(b"ERR no such index"))); + return true; + } + }; + // Step 4: index must have at least one TEXT field. + if text_index.text_fields.is_empty() { + drop(ts_guard); + responses.push(Frame::Error(Bytes::from_static( + b"ERR index has no TEXT fields", + ))); + return true; + } + // Step 5: parse via first analyzer. + let analyzer = match text_index.field_analyzers.first() { + Some(a) => a, + None => { + drop(ts_guard); + responses.push(Frame::Error(Bytes::from_static( + b"ERR index has no TEXT fields", + ))); + return true; + } + }; + let clause = match crate::command::vector_search::parse_text_query( + query_bytes.as_ref(), + analyzer, + ) { + Ok(c) => c, + Err(msg) => { + drop(ts_guard); + responses + .push(Frame::Error(Bytes::copy_from_slice(msg.as_bytes()))); + return true; + } + }; + // Step 5b: resolve field_idx. + let field_idx = match &clause.field_name { + None => None, + Some(field_name) => match text_index.text_fields.iter().position(|f| { + f.field_name + .as_ref() + .eq_ignore_ascii_case(field_name.as_ref()) + }) { + Some(idx) => Some(idx), + None => { + let bad_name = field_name.clone(); + drop(ts_guard); + responses.push(Frame::Error(Bytes::from(format!( + "ERR unknown field '{}'", + String::from_utf8_lossy(&bad_name) + )))); + return true; + } + }, + }; + // Step 6: query_terms. + let query_terms = clause.terms; + // Step 7: LIMIT + top_k cap (T-151-03-02). + let (offset, count) = + crate::command::vector_search::parse_limit_clause(cmd_args); + let top_k = if count == usize::MAX { + 10000 + } else { + offset.saturating_add(count) + } + .max(1); + // Step 8: HIGHLIGHT / SUMMARIZE. + let highlight_opts = + crate::command::vector_search::parse_highlight_clause(cmd_args); + let summarize_opts = + crate::command::vector_search::parse_summarize_clause(cmd_args); + // Step 9: DB read guard iff post-processing needed. + let db_guard_opt = if highlight_opts.is_some() || summarize_opts.is_some() { + Some(ctx.shard_databases.read_db(ctx.shard_id, 0)) + } else { + None + }; + // Step 10: execute + optional post-processing. + let mut response = crate::command::vector_search::execute_text_search_local( + &ts_guard, + &index_name, + field_idx, + &query_terms, + top_k, + offset, + count, + ); + if let Some(ref db_guard) = db_guard_opt { + let term_strings: Vec = + query_terms.iter().map(|qt| qt.text.clone()).collect(); + crate::command::vector_search::apply_post_processing( + &mut response, + &term_strings, + text_index, + db_guard, + highlight_opts.as_ref(), + summarize_opts.as_ref(), + ); + } + // Explicit drop order: db_guard first, ts_guard last. + drop(db_guard_opt); + drop(ts_guard); + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + return true; + } + } + } + } + } + let response = { + let shard_databases_ref = &ctx.shard_databases; + let mut vs = shard_databases_ref.vector_store(ctx.shard_id); + let mut ts = shard_databases_ref.text_store(ctx.shard_id); + if cmd.eq_ignore_ascii_case(b"FT.CREATE") { + crate::command::vector_search::ft_create(&mut vs, &mut ts, cmd_args) + } else if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { + // Resolve AS_OF temporal clause + TXN snapshot precedence (TEMP-04, ACID-09). + match resolve_ft_search_as_of_lsn( + cmd_args, + Some(&ctx.shard_databases), + ctx.shard_id, + conn.active_cross_txn.as_ref(), + ) { + Err(err_frame) => err_frame, + Ok(as_of_lsn) => { + let has_session = cmd_args.iter().any(|a| { + if let Frame::BulkString(b) = a { + b.eq_ignore_ascii_case(b"SESSION") + } else { + false + } + }); + if has_session { + let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); + crate::command::vector_search::ft_search( + &mut vs, + cmd_args, + Some(&mut *db_guard), + Some(&*ts), + as_of_lsn, + ) + } else { + crate::command::vector_search::ft_search( + &mut vs, + cmd_args, + None, + Some(&*ts), + as_of_lsn, + ) + } + } + } + } else if cmd.eq_ignore_ascii_case(b"FT.DROPINDEX") { + let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); + crate::command::vector_search::ft_dropindex( + &mut vs, + &mut ts, + Some(&mut *db_guard), + cmd_args, + ) + } else if cmd.eq_ignore_ascii_case(b"FT.INFO") { + crate::command::vector_search::ft_info(&vs, &ts, cmd_args) + } else if cmd.eq_ignore_ascii_case(b"FT._LIST") { + crate::command::vector_search::ft_list(&vs) + } else if cmd.eq_ignore_ascii_case(b"FT.COMPACT") { + crate::command::vector_search::ft_compact(&mut vs, &mut ts, cmd_args) + } else if cmd.eq_ignore_ascii_case(b"FT.CACHESEARCH") { + crate::command::vector_search::cache_search::ft_cachesearch(&mut vs, cmd_args) + } else if cmd.eq_ignore_ascii_case(b"FT.CONFIG") { + crate::command::vector_search::ft_config(&mut vs, &mut ts, cmd_args) + } else if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") { + let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); + crate::command::vector_search::recommend::ft_recommend( + &mut vs, + cmd_args, + Some(&mut *db_guard), + ) + } else if cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") { + #[cfg(feature = "graph")] + { + let graph_guard = shard_databases_ref.graph_store_read(ctx.shard_id); + crate::command::vector_search::navigate::ft_navigate( + &mut vs, + Some(&graph_guard), + cmd_args, + None, + ) + } + #[cfg(not(feature = "graph"))] + { + Frame::Error(Bytes::from_static( + b"ERR FT.NAVIGATE requires graph feature", + )) + } + } else if cmd.eq_ignore_ascii_case(b"FT.EXPAND") { + #[cfg(feature = "graph")] + { + let graph_guard = shard_databases_ref.graph_store_read(ctx.shard_id); + crate::command::vector_search::ft_expand(&graph_guard, cmd_args) + } + #[cfg(not(feature = "graph"))] + { + Frame::Error(Bytes::from_static(b"ERR FT.EXPAND requires graph feature")) + } + } else { + Frame::Error(Bytes::from_static(b"ERR unknown FT.* command")) + } + }; + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + true +} diff --git a/src/server/conn/handler_sharded/mod.rs b/src/server/conn/handler_sharded/mod.rs new file mode 100644 index 00000000..351fec5b --- /dev/null +++ b/src/server/conn/handler_sharded/mod.rs @@ -0,0 +1,1412 @@ +// Note: some imports/variables may be conditionally used across feature flags +//! Sharded tokio connection handlers. +//! +//! Extracted from `server/connection.rs` (Plan 48-02). +//! Contains `handle_connection_sharded` (thin wrapper) and +//! `handle_connection_sharded_inner` (generic inner handler). + +use crate::runtime::TcpStream; +use crate::runtime::cancel::CancellationToken; +use bumpalo::Bump; +use bytes::{Bytes, BytesMut}; +use ringbuf::traits::Producer; +use std::collections::HashMap; + +use crate::command::connection as conn_cmd; +use crate::command::metadata; +use crate::command::{DispatchResult, dispatch, dispatch_read}; +use crate::persistence::aof::{self, AofMessage}; +use crate::protocol::Frame; +use crate::shard::dispatch::{ShardMessage, key_to_shard}; +use crate::shard::mesh::ChannelMesh; +use crate::storage::eviction::try_evict_if_needed; +use crate::workspace::{strip_workspace_prefix_from_response, workspace_rewrite_args}; + +use super::affinity::MigratedConnectionState; +use crate::server::response_slot::ResponseSlotPool; + +mod dispatch; +mod ft; +mod pubsub; +mod read; +mod txn; +mod write; + +/// Result of `handle_connection_sharded_inner` execution. +/// +/// The generic inner handler cannot perform FD extraction (requires concrete stream type). +/// When migration is triggered, it returns `MigrateConnection` so the concrete caller +/// can extract the raw FD and send the migration message via SPSC. +pub enum HandlerResult { + /// Normal connection close (QUIT, EOF, error, shutdown). + Done, + /// AffinityTracker detected a dominant remote shard. The caller should: + /// 1. Extract the raw FD from the concrete stream (into_std + into_raw_fd) + /// 2. Send ShardMessage::MigrateConnection via SPSC to `target_shard` + /// 3. Drop the handler (connection ownership transferred) + MigrateConnection { + state: MigratedConnectionState, + target_shard: usize, + }, +} + +use super::{ + apply_resp3_conversion, convert_blocking_to_nonblocking, execute_transaction_sharded, + extract_bytes, extract_command, extract_primary_key, handle_blocking_command, handle_config, + is_multi_key_command, unpropagate_subscription, +}; + +/// Handle a single client connection on a sharded (thread-per-core) runtime. +/// +/// Runs within a shard's single-threaded Tokio runtime. Has direct mutable access +/// to the shard's databases via `Arc` (thread-safe: parking_lot RwLock +/// single-threaded scheduling means no concurrent borrows). +/// +/// Routing logic: +/// - **Keyless commands** (PING, ECHO, SELECT, etc.): execute locally, zero overhead. +/// - **Single-key, local shard**: execute directly on borrowed database -- ZERO cross-shard overhead. +/// - **Single-key, remote shard**: dispatch via SPSC `ShardMessage::Execute`, await oneshot reply. +/// - **Multi-key commands** (MGET, MSET, multi-DEL): delegate to VLL coordinator. +/// +/// Connection-level commands (AUTH, SUBSCRIBE, MULTI/EXEC) are handled at the +/// connection level same as the non-sharded handler. +#[tracing::instrument(skip_all, level = "debug")] +pub(crate) async fn handle_connection_sharded( + mut stream: TcpStream, + ctx: &super::core::ConnectionContext, + shutdown: CancellationToken, + client_id: u64, +) { + let maxclients = ctx.runtime_config.read().maxclients; + if !crate::admin::metrics_setup::try_accept_connection(maxclients) { + use tokio::io::AsyncWriteExt; + let _ = stream + .write_all(b"-ERR max number of clients reached\r\n") + .await; + return; + } + let peer_addr = stream + .peer_addr() + .map(|a| a.to_string()) + .unwrap_or_else(|_| "unknown".to_string()); + let result = handle_connection_sharded_inner( + stream, + peer_addr, + ctx, + shutdown, + client_id, + true, // can_migrate: plain TCP supports FD extraction + BytesMut::new(), + None, // fresh connection, no migrated state + ) + .await; + + // Handle migration result: extract FD from the returned stream and send via SPSC + if let ( + HandlerResult::MigrateConnection { + state, + target_shard, + }, + Some(stream), + ) = (result.0, result.1) + { + use std::os::unix::io::IntoRawFd; + match stream.into_std() { + Ok(std_stream) => { + let raw_fd = std_stream.into_raw_fd(); + let msg = ShardMessage::MigrateConnection { fd: raw_fd, state }; + let target_idx = ChannelMesh::target_index(ctx.shard_id, target_shard); + let push_result = { + let mut producers = ctx.dispatch_tx.borrow_mut(); + producers[target_idx].try_push(msg) + }; + match push_result { + Ok(()) => { + ctx.spsc_notifiers[target_shard].notify_one(); + tracing::info!( + "Shard {}: migrated connection {} to shard {}", + ctx.shard_id, + client_id, + target_shard + ); + } + Err(returned_msg) => { + // SPSC full — retry with yield before giving up. + let mut pending = Some(returned_msg); + for _ in 0..8 { + tokio::task::yield_now().await; + #[allow(clippy::unwrap_used)] + // pending is always re-filled on retry via Err(returned_msg) + let msg = pending.take().unwrap(); + let push_result = { + let mut producers = ctx.dispatch_tx.borrow_mut(); + producers[target_idx].try_push(msg) + }; + match push_result { + Ok(()) => { + ctx.spsc_notifiers[target_shard].notify_one(); + tracing::info!( + "Shard {}: migrated connection {} to shard {} (after retry)", + ctx.shard_id, + client_id, + target_shard + ); + break; + } + Err(msg) => pending = Some(msg), + } + } + if let Some(ShardMessage::MigrateConnection { fd, .. }) = pending { + use std::os::unix::io::FromRawFd; + // SAFETY: fd is a valid, uniquely-owned file descriptor obtained + // from TcpStream::into_raw_fd() above. OwnedFd closes it on drop. + drop(unsafe { std::os::unix::io::OwnedFd::from_raw_fd(fd) }); + } + tracing::warn!( + "Shard {}: migration SPSC full, connection {} lost", + ctx.shard_id, + client_id + ); + } + } + } + Err(e) => { + tracing::warn!("Shard {}: migration into_std failed: {}", ctx.shard_id, e); + // Stream consumed by into_std attempt, connection lost either way + } + } + } else { + // Only decrement connected_clients when the connection is actually closing, + // not when migrating to another shard (the connection stays alive). + crate::admin::metrics_setup::record_connection_closed(); + } +} + +/// Generic inner handler for sharded connections (Tokio runtime). +/// +/// Works with any stream implementing `AsyncRead + AsyncWrite + Unpin`, +/// enabling both plain TCP (`TcpStream`) and TLS (`tokio_rustls::server::TlsStream`). +/// +/// Returns `(HandlerResult, Option)`: the stream is returned when migration is triggered +/// so the concrete caller can extract the raw FD. `can_migrate` controls whether the +/// AffinityTracker is active (set to `false` for TLS connections). +pub(crate) async fn handle_connection_sharded_inner< + S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, +>( + stream: S, + peer_addr: String, + ctx: &super::core::ConnectionContext, + shutdown: CancellationToken, + client_id: u64, + can_migrate: bool, + initial_read_buf: BytesMut, + migrated_state: Option<&MigratedConnectionState>, +) -> (HandlerResult, Option) { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + // Direct buffer I/O: bypass Framed/codec for the hot path. + let mut stream = stream; + let mut read_buf = if initial_read_buf.is_empty() { + BytesMut::with_capacity(8192) + } else { + // Migration buffer: leftover bytes from the source connection. + let mut buf = initial_read_buf; + buf.reserve(8192); + buf + }; + let mut write_buf = BytesMut::with_capacity(8192); + let parse_config = crate::protocol::ParseConfig::default(); + let mut conn = super::core::ConnectionState::new( + client_id, + peer_addr.clone(), + &ctx.requirepass, + ctx.shard_id, + ctx.num_shards, + can_migrate, + ctx.runtime_config.read().acllog_max_len, + migrated_state, + ); + conn.refresh_acl_cache(&ctx.acl_table); + + // Register in global client registry for CLIENT LIST/INFO/KILL. + // RegistryGuard ensures deregister on all exit paths (including early returns). + crate::client_registry::register( + client_id, + peer_addr.clone(), + conn.current_user.clone(), + ctx.shard_id, + ); + struct RegistryGuard(u64); + impl Drop for RegistryGuard { + fn drop(&mut self) { + crate::client_registry::deregister(self.0); + } + } + let _registry_guard = RegistryGuard(client_id); + + // Functions API registry (per-shard, lazy init) — kept as local because Rc> is !Send + let func_registry = std::rc::Rc::new(std::cell::RefCell::new( + crate::scripting::FunctionRegistry::new(), + )); + + // Per-connection arena for batch processing temporaries. + // 4KB initial capacity, grows on demand (rarely exceeds 16KB per batch). + let mut arena = Bump::with_capacity(4096); + + // Pre-allocated response slots for zero-allocation cross-shard dispatch. + let response_pool = ResponseSlotPool::new(ctx.num_shards, ctx.shard_id); + + // Client idle timeout: 0 = disabled (read once, avoid lock on hot path) + let idle_timeout_secs = ctx.runtime_config.read().timeout; + let idle_timeout = if idle_timeout_secs > 0 { + Some(std::time::Duration::from_secs(idle_timeout_secs)) + } else { + None + }; + + let mut break_outer = false; + loop { + // Check if CLIENT KILL targeted this connection + if crate::client_registry::is_killed(client_id) { + break; + } + + // --- Subscriber mode: bidirectional select on client commands + published messages --- + if conn.subscription_count > 0 { + match pubsub::run_subscriber_step( + &mut stream, + &mut read_buf, + &mut write_buf, + &parse_config, + &mut conn, + ctx, + &peer_addr, + &shutdown, + ) + .await + { + pubsub::SubscriberAction::Continue => { + continue; + } + pubsub::SubscriberAction::BreakOuter => { + break; + } + pubsub::SubscriberAction::EarlyReturn => { + return (HandlerResult::Done, None); + } + } + } + tokio::select! { + result = async { + if let Some(dur) = idle_timeout { + match tokio::time::timeout(dur, stream.read_buf(&mut read_buf)).await { + Ok(r) => r, + Err(_) => Err(std::io::Error::new(std::io::ErrorKind::TimedOut, "idle timeout")), + } + } else { + stream.read_buf(&mut read_buf).await + } + } => { + match result { + Ok(0) => break, // connection closed + Ok(_) => {} + Err(e) if e.kind() == std::io::ErrorKind::TimedOut => { + tracing::debug!("Connection {} idle timeout ({}s)", client_id, idle_timeout_secs); + break; + } + Err(_) => break, + } + + // Parse all complete frames from buffer + let mut batch: Vec = Vec::with_capacity(64); + const MAX_BATCH: usize = 1024; + loop { + match crate::protocol::parse(&mut read_buf, &parse_config) { + Ok(Some(frame)) => { + batch.push(frame); + if batch.len() >= MAX_BATCH { break; } + } + Ok(None) => break, + Err(crate::protocol::ParseError::Incomplete) => break, + Err(_) => { break_outer = true; break; } + } + } + if break_outer { break; } + if batch.is_empty() { continue; } + + // CLIENT PAUSE: delay processing if server is paused + // Check with is_write=true (conservative — pauses all batches in ALL mode) + crate::client_pause::expire_if_needed(); + if let Some(remaining) = crate::client_pause::check_pause(true) { + tokio::time::sleep(remaining).await; + } + + let mut responses: Vec = Vec::with_capacity(batch.len()); + let mut should_quit = false; + let mut remote_groups: HashMap, Option, Bytes, usize)>> = HashMap::with_capacity(ctx.num_shards); + // Accumulate cross-shard PUBLISH pairs per target shard for batch dispatch + // Key: target shard ID -> Vec of (response_index, channel, message) + let mut publish_batches: HashMap> = HashMap::new(); + + // Track if AUTH rate limiting delay is needed (applied after batch response) + let mut auth_delay_ms: u64 = 0; + + for frame in batch { + // --- AUTH gate --- + if !conn.authenticated { + match extract_command(&frame) { + Some((cmd, cmd_args)) if cmd.eq_ignore_ascii_case(b"AUTH") => { + let (response, opt_user) = conn_cmd::auth_acl(cmd_args, &ctx.acl_table); + if let Some(uname) = opt_user { + conn.authenticated = true; + conn.current_user = uname; + conn.refresh_acl_cache(&ctx.acl_table); + if let Ok(addr) = peer_addr.parse::() { + crate::auth_ratelimit::record_success(addr.ip()); + } + } else { + if let Ok(addr) = peer_addr.parse::() { + auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); + } + conn.acl_log.push(crate::acl::AclLogEntry { + reason: "auth".to_string(), + object: "AUTH".to_string(), + username: conn.current_user.clone(), + client_addr: peer_addr.clone(), + timestamp_ms: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64, + }); + } + responses.push(response); + continue; + } + Some((cmd, cmd_args)) if cmd.eq_ignore_ascii_case(b"HELLO") => { + let (response, new_proto, new_name, opt_user) = conn_cmd::hello_acl( + cmd_args, + conn.protocol_version, + client_id, + &ctx.acl_table, + &mut conn.authenticated, + ); + if !matches!(&response, Frame::Error(_)) { + conn.protocol_version = new_proto; + } + if let Some(name) = new_name { + conn.client_name = Some(name); + } + if let Some(ref uname) = opt_user { + conn.current_user = uname.clone(); + conn.refresh_acl_cache(&ctx.acl_table); + } + // HELLO AUTH rate limiting (same as AUTH gate) + if matches!(&response, Frame::Error(_)) { + if let Ok(addr) = peer_addr.parse::() { + auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); + } + } else if opt_user.is_some() { + if let Ok(addr) = peer_addr.parse::() { + crate::auth_ratelimit::record_success(addr.ip()); + } + } + responses.push(response); + continue; + } + Some((cmd, _)) if cmd.eq_ignore_ascii_case(b"QUIT") => { + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + should_quit = true; + break; + } + _ => { + responses.push(Frame::Error( + Bytes::from_static(b"NOAUTH Authentication required.") + )); + continue; + } + } + } + + let (cmd, cmd_args) = match extract_command(&frame) { + Some(pair) => pair, + None => { + responses.push(Frame::Error(Bytes::from_static( + b"ERR invalid command format", + ))); + continue; + } + }; + + // --- QUIT --- + if cmd.eq_ignore_ascii_case(b"QUIT") { + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + should_quit = true; + break; + } + + // --- ASKING --- + if cmd.eq_ignore_ascii_case(b"ASKING") { + conn.asking = true; + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + continue; + } + + // --- CLUSTER subcommands --- + if cmd.eq_ignore_ascii_case(b"CLUSTER") { + if let Some(ref cs) = ctx.cluster_state { + #[allow(clippy::unwrap_used)] // Fallback "127.0.0.1:6379" is a valid literal + let self_addr: std::net::SocketAddr = + format!("127.0.0.1:{}", ctx.config_port) + .parse() + .unwrap_or_else(|_| "127.0.0.1:6379".parse().unwrap()); + let resp = crate::cluster::command::handle_cluster_command( + cmd_args, cs, self_addr, + ); + responses.push(resp); + } else { + responses.push(Frame::Error(Bytes::from_static( + b"ERR This instance has cluster support disabled", + ))); + } + continue; + } + + // --- Lua scripting: EVAL / EVALSHA --- + if cmd.eq_ignore_ascii_case(b"EVAL") || cmd.eq_ignore_ascii_case(b"EVALSHA") { + let response = { + let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + let db_count = ctx.shard_databases.db_count(); + if cmd.eq_ignore_ascii_case(b"EVAL") { + crate::scripting::handle_eval( + &ctx.lua, &ctx.script_cache, cmd_args, &mut guard, + ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, + ) + } else { + crate::scripting::handle_evalsha( + &ctx.lua, &ctx.script_cache, cmd_args, &mut guard, + ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, + ) + } + }; + responses.push(response); + continue; + } + + // --- SCRIPT subcommands --- + if cmd.eq_ignore_ascii_case(b"SCRIPT") { + let (response, fanout) = crate::scripting::handle_script_subcommand(&ctx.script_cache, cmd_args); + if let Some((sha1, script_bytes)) = fanout { + let mut producers = ctx.dispatch_tx.borrow_mut(); + for target in 0..ctx.num_shards { + if target == ctx.shard_id { continue; } + let idx = ChannelMesh::target_index(ctx.shard_id, target); + let msg = ShardMessage::ScriptLoad { sha1: sha1.clone(), script: script_bytes.clone() }; + if producers[idx].try_push(msg).is_ok() { + ctx.spsc_notifiers[target].notify_one(); + } + } + drop(producers); + } + responses.push(response); + continue; + } + + // --- Cluster slot routing (pre-dispatch) --- + if crate::cluster::cluster_enabled() { + if let Some(ref cs) = ctx.cluster_state { + let was_asking = conn.asking; + conn.asking = false; + let maybe_key = extract_primary_key(cmd, cmd_args); + if let Some(key) = maybe_key { + let slot = crate::cluster::slots::slot_for_key(key); + #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable + let route = cs.read().unwrap().route_slot(slot, was_asking); + match route { + crate::cluster::SlotRoute::Local => {} + other => { + responses.push(other.into_error_frame(slot)); + continue; + } + } + if is_multi_key_command(cmd, cmd_args) { + let first_slot = slot; + let mut cross_slot = false; + for arg in cmd_args.iter().skip(1) { + if let Some(k) = match arg { + Frame::BulkString(b) => Some(b.as_ref()), + _ => None, + } { + if crate::cluster::slots::slot_for_key(k) != first_slot { + cross_slot = true; + break; + } + } + } + if cross_slot { + responses.push(Frame::Error(Bytes::from_static( + b"CROSSSLOT Keys in request don't hash to the same slot", + ))); + continue; + } + } + } + } + } + + // --- AUTH (already conn.authenticated) --- + if cmd.eq_ignore_ascii_case(b"AUTH") { + let (response, opt_user) = conn_cmd::auth_acl(cmd_args, &ctx.acl_table); + if let Some(uname) = opt_user { + conn.current_user = uname; + conn.refresh_acl_cache(&ctx.acl_table); + if let Ok(addr) = peer_addr.parse::() { + crate::auth_ratelimit::record_success(addr.ip()); + } + } else if let Ok(addr) = peer_addr.parse::() { + auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); + } + responses.push(response); + continue; + } + + // --- HELLO --- + if cmd.eq_ignore_ascii_case(b"HELLO") { + let (response, new_proto, new_name, opt_user) = conn_cmd::hello_acl( + cmd_args, conn.protocol_version, client_id, &ctx.acl_table, &mut conn.authenticated, + ); + if !matches!(&response, Frame::Error(_)) { conn.protocol_version = new_proto; } + if let Some(name) = new_name { conn.client_name = Some(name); } + if let Some(ref uname) = opt_user { + conn.current_user = uname.clone(); + conn.refresh_acl_cache(&ctx.acl_table); + } + if matches!(&response, Frame::Error(_)) { + if let Ok(addr) = peer_addr.parse::() { + auth_delay_ms += crate::auth_ratelimit::record_failure(addr.ip()); + } + } else if opt_user.is_some() { + if let Ok(addr) = peer_addr.parse::() { + crate::auth_ratelimit::record_success(addr.ip()); + } + } + responses.push(response); + continue; + } + + // --- ACL --- + if cmd.eq_ignore_ascii_case(b"ACL") { + let response = crate::command::acl::handle_acl( + cmd_args, &ctx.acl_table, &mut conn.acl_log, &conn.current_user, &peer_addr, &ctx.runtime_config, + ); + responses.push(response); + continue; + } + + // === CLIENT PAUSE check === + // Extract pause info with short lock hold, then sleep outside lock scope + let pause_wait_ms = { + let rt = ctx.runtime_config.read(); + let deadline = rt.client_pause_deadline_ms; + if deadline > 0 { + let now = crate::storage::entry::current_time_ms(); + if now < deadline { + let should_pause = if rt.client_pause_write_only { + crate::command::metadata::is_write(cmd) + } else { + true + }; + if should_pause { deadline.saturating_sub(now) } else { 0 } + } else { 0 } + } else { 0 } + }; + if pause_wait_ms > 0 { + // Poll in 50ms intervals so CLIENT UNPAUSE takes effect quickly + let mut remaining = pause_wait_ms; + while remaining > 0 { + let chunk = remaining.min(50); + #[cfg(feature = "runtime-tokio")] + { + tokio::time::sleep(std::time::Duration::from_millis(chunk)).await; + } + #[cfg(feature = "runtime-monoio")] + { + monoio::time::sleep(std::time::Duration::from_millis(chunk)).await; + } + remaining = remaining.saturating_sub(chunk); + // Re-check if UNPAUSE was called + let still_paused = { + let rt = ctx.runtime_config.read(); + rt.client_pause_deadline_ms > 0 + && crate::storage::entry::current_time_ms() < rt.client_pause_deadline_ms + }; + if !still_paused { + break; + } + } + } + + // === ACL permission check === + // Must run before any command-specific handlers (CONFIG, REPLICAOF, etc.) + // so that low-privilege users cannot reach admin commands. + // Fast path: skip RwLock + HashMap for unrestricted users + // with a fresh cache. Stale caches (after ACL SETUSER / + // DELUSER / LOAD) fall through to the full check. + if !conn.acl_skip_allowed() { + #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable + let acl_guard = ctx.acl_table.read().unwrap(); + if let Some(deny_reason) = acl_guard.check_command_permission(&conn.current_user, cmd, cmd_args) { + drop(acl_guard); + conn.acl_log.push(crate::acl::AclLogEntry { + reason: "command".to_string(), + object: String::from_utf8_lossy(cmd).to_ascii_lowercase(), + username: conn.current_user.clone(), + client_addr: peer_addr.clone(), + timestamp_ms: std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_millis() as u64, + }); + responses.push(Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason)))); + continue; + } + let is_write_for_acl = metadata::is_write(cmd); + if let Some(deny_reason) = acl_guard.check_key_permission(&conn.current_user, cmd, cmd_args, is_write_for_acl) { + drop(acl_guard); + conn.acl_log.push(crate::acl::AclLogEntry { + reason: "command".to_string(), + object: String::from_utf8_lossy(cmd).to_ascii_lowercase(), + username: conn.current_user.clone(), + client_addr: peer_addr.clone(), + timestamp_ms: std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_millis() as u64, + }); + responses.push(Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason)))); + continue; + } + } + + // --- Functions API: FUNCTION/FCALL/FCALL_RO --- + // Placed AFTER ACL check. Respects MULTI queue — if conn.in_multi, + // fall through to the MULTI queue gate instead of executing. + if !conn.in_multi { + if cmd.eq_ignore_ascii_case(b"FUNCTION") { + let response = crate::command::functions::handle_function( + &mut func_registry.borrow_mut(), cmd_args, + ); + responses.push(response); + continue; + } + if cmd.eq_ignore_ascii_case(b"FCALL") { + let response = { + let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + let db_count = ctx.shard_databases.db_count(); + crate::command::functions::handle_fcall( + &func_registry.borrow(), cmd_args, &mut guard, + ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, + ) + }; + responses.push(response); + continue; + } + if cmd.eq_ignore_ascii_case(b"FCALL_RO") { + let response = { + let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + let db_count = ctx.shard_databases.db_count(); + crate::command::functions::handle_fcall_ro( + &func_registry.borrow(), cmd_args, &mut guard, + ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, + ) + }; + responses.push(response); + continue; + } + } + + // --- CONFIG --- + if dispatch::try_handle_config(cmd, cmd_args, ctx, &mut responses) { + continue; + } + + // --- SLOWLOG --- + if dispatch::try_handle_slowlog(cmd, cmd_args, &mut responses) { + continue; + } + + // --- REPLICAOF / SLAVEOF --- + if dispatch::try_handle_replicaof(cmd, cmd_args, ctx, &mut responses) { + continue; + } + + // --- REPLCONF --- + if dispatch::try_handle_replconf(cmd, cmd_args, &mut responses) { + continue; + } + + // --- INFO --- + if dispatch::try_handle_info(cmd, cmd_args, &conn, ctx, &mut responses) { + continue; + } + + // --- READONLY enforcement --- + if dispatch::try_enforce_readonly(cmd, ctx, &mut responses) { + continue; + } + + // --- CLIENT subcommands --- + if dispatch::try_handle_client_command(cmd, cmd_args, client_id, &mut conn, ctx, &mut responses) { + continue; + } + + // --- TXN.BEGIN / TXN.COMMIT / TXN.ABORT --- + if txn::try_handle_txn_begin(cmd, cmd_args, &mut conn, ctx, &mut responses) { + continue; + } + if txn::try_handle_txn_commit(cmd, cmd_args, &mut conn, ctx, &mut responses) { + continue; + } + if txn::try_handle_txn_abort(cmd, cmd_args, &mut conn, ctx, &mut responses) { + continue; + } + + // --- TEMPORAL.SNAPSHOT_AT / TEMPORAL.INVALIDATE --- + if txn::try_handle_temporal_snapshot_at(cmd, cmd_args, ctx, &mut responses) { + continue; + } + if txn::try_handle_temporal_invalidate(cmd, cmd_args, ctx, &mut responses) { + continue; + } + + // --- WS.* --- + if write::try_handle_ws_command(cmd, cmd_args, &mut conn, ctx, &mut responses) { + continue; + } + + // --- MQ.* --- + if write::try_handle_mq_command(cmd, cmd_args, &mut conn, ctx, &mut responses) { + continue; + } + + // --- MULTI / EXEC_CMD / DISCARD --- + if write::try_handle_multi_exec(cmd, &mut conn, ctx, &mut responses) { + continue; + } + + // --- Workspace key prefix injection --- + // MUST happen before key_to_shard() so the {ws_id} hash tag determines + // shard routing. This is the ONLY code path where workspace prefixing + // occurs (WS-07, WS-12). All subsequent dispatch uses cmd_args (shadowed). + let rewritten = conn + .workspace_id + .as_ref() + .map(|ws_id| workspace_rewrite_args(cmd, cmd_args, ws_id)); + let cmd_args: &[Frame] = rewritten.as_deref().unwrap_or(cmd_args); + + // --- BLOCKING COMMANDS --- + if cmd.eq_ignore_ascii_case(b"BLPOP") || cmd.eq_ignore_ascii_case(b"BRPOP") + || cmd.eq_ignore_ascii_case(b"BLMOVE") || cmd.eq_ignore_ascii_case(b"BZPOPMIN") + || cmd.eq_ignore_ascii_case(b"BZPOPMAX") + || cmd.eq_ignore_ascii_case(b"BLMPOP") || cmd.eq_ignore_ascii_case(b"BRPOPLPUSH") + || cmd.eq_ignore_ascii_case(b"BZMPOP") + { + if conn.in_multi { + let nb_frame = convert_blocking_to_nonblocking(cmd, cmd_args); + conn.command_queue.push(nb_frame); + responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); + continue; + } + write_buf.clear(); + for response in responses.iter() { + if conn.protocol_version >= 3 { + crate::protocol::serialize_resp3(response, &mut write_buf); + } else { + crate::protocol::serialize(response, &mut write_buf); + } + } + if stream.write_all(&write_buf).await.is_err() { arena.reset(); return (HandlerResult::Done, None); } + let blocking_response = handle_blocking_command( + cmd, cmd_args, conn.selected_db, &ctx.shard_databases, &ctx.blocking_registry, + ctx.shard_id, ctx.num_shards, &ctx.dispatch_tx, &shutdown, + ).await; + let blocking_response = apply_resp3_conversion(cmd, blocking_response, conn.protocol_version); + responses = Vec::with_capacity(1); + responses.push(blocking_response); + break; + } + + // --- PUBLISH --- + if pubsub::try_handle_publish(cmd, cmd_args, &conn, ctx, &mut responses, &mut publish_batches) { + continue; + } + + // --- SUBSCRIBE / PSUBSCRIBE --- + if let Some(action) = pubsub::try_handle_subscribe( + cmd, cmd_args, &mut stream, &mut write_buf, + &mut conn, ctx, &peer_addr, &mut responses, + ).await { + match action { + pubsub::SubscriberAction::Continue => { continue; } + pubsub::SubscriberAction::BreakOuter => { break; } + pubsub::SubscriberAction::EarlyReturn => { return (HandlerResult::Done, None); } + } + } + // UNSUBSCRIBE/PUNSUBSCRIBE in normal mode (not subscribed) + if pubsub::try_handle_unsubscribe(cmd, &mut responses) { + continue; + } + + // --- PUBSUB introspection (zero-SPSC: direct shared-read) --- + if pubsub::try_handle_pubsub_introspection(cmd, cmd_args, ctx, &mut responses) { + continue; + } + + // --- MULTI queue mode --- + if conn.in_multi { + conn.command_queue.push(frame); + responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); + continue; + } + + // --- BGSAVE / SAVE / LASTSAVE / BGREWRITEAOF --- + if dispatch::try_handle_persistence(cmd, ctx, &mut responses) { + continue; + } + + // --- Cross-shard aggregation: KEYS, SCAN, DBSIZE --- + if dispatch::try_handle_cross_shard_scan(cmd, cmd_args, &conn, ctx, &mut responses).await { + continue; + } + + // --- FT.* vector search commands --- + if ft::try_handle_ft_command(cmd, cmd_args, &frame, &conn, ctx, &mut responses).await { + continue; + } + + // --- GRAPH.* graph commands --- + #[cfg(feature = "graph")] + if write::try_handle_graph_command(cmd, cmd_args, &mut conn, ctx, &mut responses) { + continue; + } + + // --- Multi-key commands --- + if is_multi_key_command(cmd, cmd_args) { + let response = crate::shard::coordinator::coordinate_multi_key(cmd, cmd_args, ctx.shard_id, ctx.num_shards, conn.selected_db, &ctx.shard_databases, &ctx.dispatch_tx, &ctx.spsc_notifiers, &ctx.cached_clock, &()).await; + responses.push(response); + continue; + } + + // --- Routing: keyless, local, or remote --- + let target_shard = extract_primary_key(cmd, cmd_args).map(|key| key_to_shard(key, ctx.num_shards)); + let is_local = match target_shard { + None => true, + Some(s) if s == ctx.shard_id => true, + _ => false, + }; + + // Affinity sampling: record shard target for migration decision. + // Only sample when we have a concrete target shard (key-bearing command). + // Migration is deferred until AFTER the current batch is fully processed + // and all responses are written, ensuring no command/response desync. + if let (Some(tracker), Some(target)) = (&mut conn.affinity_tracker, target_shard) { + if let Some(migrate_to) = tracker.record(target) { + // Migration preconditions: not in MULTI, no active CLIENT TRACKING + // (tracking connections need untrack_all cleanup which doesn't transfer) + if !conn.in_multi && !conn.tracking_state.enabled { + conn.migration_target = Some(migrate_to); + } + } + } + + let is_write = if ctx.aof_tx.is_some() || conn.tracking_state.enabled { metadata::is_write(cmd) } else { false }; + let aof_bytes = if is_write && ctx.aof_tx.is_some() { Some(aof::serialize_command(&frame)) } else { None }; + + if is_local { + // LOCAL PATH: split into read/write to avoid exclusive lock on reads. + // Using read_db for local reads eliminates RwLock contention with + // cross-shard shared reads from other shard threads. + crate::admin::metrics_setup::record_dispatch_local(); + if metadata::is_write(cmd) { + // WRITE PATH: single lock acquisition for eviction + dispatch + let rt = ctx.runtime_config.read(); + let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + if let Err(oom_frame) = try_evict_if_needed(&mut guard, &rt) { + drop(guard); + drop(rt); + responses.push(oom_frame); + continue; + } + drop(rt); + + // KV undo-log capture for active cross-store transactions. + // MUST happen BEFORE dispatch() overwrites the database entry. + if let Some(ref mut txn) = conn.active_cross_txn { + if cmd.eq_ignore_ascii_case(b"DEL") || cmd.eq_ignore_ascii_case(b"UNLINK") { + // Multi-key DEL: iterate all args for undo capture + for arg in cmd_args.iter() { + if let Frame::BulkString(key_bytes) = arg { + if let Some(old_entry) = guard.get(key_bytes.as_ref()).cloned() { + txn.kv_undo.record_delete(key_bytes.clone(), old_entry); + let lsn = txn.snapshot_lsn; + let tid = txn.txn_id; + ctx.shard_databases.kv_intents(ctx.shard_id) + .record_write(key_bytes.clone(), lsn, tid); + } + // Key not found: nothing to undo, no intent registered + } + } + } else if let Some(key) = crate::server::conn::shared::extract_primary_key(cmd, cmd_args) { + // SET / HSET / INCR / etc. — single primary key + let old_entry = guard.get(key.as_ref()).cloned(); + let lsn = txn.snapshot_lsn; + let tid = txn.txn_id; + match old_entry { + None => txn.kv_undo.record_insert(key.clone()), + Some(entry) => txn.kv_undo.record_update(key.clone(), entry), + } + ctx.shard_databases.kv_intents(ctx.shard_id) + .record_write(key.clone(), lsn, tid); + } + } + + let db_count = ctx.shard_databases.db_count(); + guard.refresh_now_from_cache(&ctx.cached_clock); + let dispatch_start = std::time::Instant::now(); + let result = dispatch(&mut guard, cmd, cmd_args, &mut conn.selected_db, db_count); + let elapsed_us = dispatch_start.elapsed().as_micros() as u64; + if let Ok(cmd_str) = std::str::from_utf8(cmd) { + crate::admin::metrics_setup::record_command(cmd_str, elapsed_us); + } + if let Frame::Array(ref args) = frame { + crate::admin::metrics_setup::global_slowlog().maybe_record( + elapsed_us, + args.as_slice(), + peer_addr.as_bytes(), + conn.client_name.as_ref().map_or(b"" as &[u8], |n| n.as_ref()), + ); + } + let response = match result { + DispatchResult::Response(f) => f, + DispatchResult::Quit(f) => { should_quit = true; f } + }; + if matches!(response, Frame::Error(_)) { + if let Ok(cmd_str) = std::str::from_utf8(cmd) { + crate::admin::metrics_setup::record_command_error(cmd_str); + } + } else { + let needs_wake = cmd.eq_ignore_ascii_case(b"LPUSH") || cmd.eq_ignore_ascii_case(b"RPUSH") + || cmd.eq_ignore_ascii_case(b"LMOVE") || cmd.eq_ignore_ascii_case(b"ZADD"); + if needs_wake { + if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { + let mut reg = ctx.blocking_registry.borrow_mut(); + if cmd.eq_ignore_ascii_case(b"LPUSH") || cmd.eq_ignore_ascii_case(b"RPUSH") || cmd.eq_ignore_ascii_case(b"LMOVE") { + crate::blocking::wakeup::try_wake_list_waiter(&mut reg, &mut guard, conn.selected_db, &key); + } else { + crate::blocking::wakeup::try_wake_zset_waiter(&mut reg, &mut guard, conn.selected_db, &key); + } + } + } + } + drop(guard); + // Auto-index vectors on successful HSET (local write path) + // Placed AFTER drop(guard) to avoid DB→vector_store lock order + // inversion with the shard event loop (vector_store→DB). + if !matches!(response, Frame::Error(_)) + && (cmd.eq_ignore_ascii_case(b"HSET") || cmd.eq_ignore_ascii_case(b"HMSET")) + { + if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { + // Phase 166 (Plan 02): if inside an active TXN, tag each + // inserted mutable-HNSW entry with the txn_id so that + // non-TXN readers (snapshot_lsn==0) see it as uncommitted + // and exclude it until TXN.COMMIT calls + // txn_manager.commit(txn_id) (ACID-09 isolation fix). + let active_txn_id = conn + .active_cross_txn + .as_ref() + .map(|t| t.txn_id) + .unwrap_or(0); + let inserted = { + let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); + let mut ts = ctx.shard_databases.text_store(ctx.shard_id); + if active_txn_id != 0 { + crate::shard::spsc_handler::auto_index_hset_public_txn(&mut vs, &mut *ts, &key, cmd_args, active_txn_id) + } else { + crate::shard::spsc_handler::auto_index_hset_public(&mut vs, &mut *ts, &key, cmd_args) + } + }; + // Push one VectorIntent per (index_name, key_hash) so + // TXN.ABORT (Plan 166-03) can tombstone via + // MutableSegment::mark_deleted_by_key_hash. + if let Some(txn) = conn.active_cross_txn.as_mut() { + for (index_name, key_hash) in inserted { + txn.record_vector(key_hash, index_name); + } + } + } + } + // Auto-delete vectors on DEL/UNLINK (local write path) + // Note: HDEL removes fields, not keys — it should NOT trigger + // vector deletion unless the entire key is removed. + if !matches!(response, Frame::Error(_)) + && (cmd.eq_ignore_ascii_case(b"DEL") || cmd.eq_ignore_ascii_case(b"UNLINK")) + { + let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); + for arg in cmd_args.iter() { + if let Some(key) = extract_bytes(arg) { + vs.mark_deleted_for_key(key.as_ref()); + } + } + } + if let Some(bytes) = aof_bytes { + if !matches!(response, Frame::Error(_)) { + if let Some(ref tx) = ctx.aof_tx { let _ = tx.try_send(AofMessage::Append(bytes)); } + } + } + if conn.tracking_state.enabled && !matches!(response, Frame::Error(_)) { + if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { + let senders = ctx.tracking_table.borrow_mut().invalidate_key(&key, client_id); + if !senders.is_empty() { + let push = crate::tracking::invalidation::invalidation_push(&[key]); + for tx in senders { let _ = tx.try_send(push.clone()); } + } + } + } + let mut response = apply_resp3_conversion(cmd, response, conn.protocol_version); + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + } else { + // Snapshot visibility filter for active cross-store transactions. + // MVCC: hide keys written by uncommitted foreign transactions. + if conn.in_cross_txn() { + if let Some(ref txn) = conn.active_cross_txn { + if let Some(key) = crate::server::conn::shared::extract_primary_key(cmd, cmd_args) { + let snapshot_lsn = txn.snapshot_lsn; + let my_txn_id = txn.txn_id; + // Clone committed treemap to release vector_store lock + // before acquiring kv_intents lock (lock ordering). + let committed = { + let vs = ctx.shard_databases.vector_store(ctx.shard_id); + vs.txn_manager().committed_treemap().clone() + }; + let visible = { + let intents = ctx.shard_databases.kv_intents(ctx.shard_id); + intents.is_key_visible(key.as_ref(), snapshot_lsn, my_txn_id, &committed) + }; + if !visible { + responses.push(Frame::Null); + continue; + } + } + } + } + + // READ PATH: shared lock — no contention with other shards' reads + let guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); + let now_ms = ctx.cached_clock.ms(); + let db_count = ctx.shard_databases.db_count(); + let dispatch_start = std::time::Instant::now(); + let result = dispatch_read(&guard, cmd, cmd_args, now_ms, &mut conn.selected_db, db_count); + let elapsed_us = dispatch_start.elapsed().as_micros() as u64; + if let Ok(cmd_str) = std::str::from_utf8(cmd) { + crate::admin::metrics_setup::record_command(cmd_str, elapsed_us); + } + if let Frame::Array(ref args) = frame { + crate::admin::metrics_setup::global_slowlog().maybe_record( + elapsed_us, + args.as_slice(), + peer_addr.as_bytes(), + conn.client_name.as_ref().map_or(b"" as &[u8], |n| n.as_ref()), + ); + } + drop(guard); + let response = match result { + DispatchResult::Response(f) => f, + DispatchResult::Quit(f) => { should_quit = true; f } + }; + if matches!(response, Frame::Error(_)) { + if let Ok(cmd_str) = std::str::from_utf8(cmd) { + crate::admin::metrics_setup::record_command_error(cmd_str); + } + } + if conn.tracking_state.enabled && !conn.tracking_state.bcast { + if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { + ctx.tracking_table.borrow_mut().track_key(client_id, &key, conn.tracking_state.noloop); + } + } + let mut response = apply_resp3_conversion(cmd, response, conn.protocol_version); + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + } + } else if let Some(target) = target_shard { + // TXN cross-shard guard: cross-shard writes bypass the undo log and + // cannot be rolled back on TXN.ABORT. Return an explicit error instead + // of silently permitting writes that resist rollback. + if conn.in_cross_txn() && metadata::is_write(cmd) { + responses.push(Frame::Error(bytes::Bytes::from_static( + crate::command::transaction::ERR_TXN_CROSS_SHARD, + ))); + continue; + } + // SHARED-READ FAST PATH: cross-shard reads bypass SPSC dispatch entirely. + // By this point conn.in_multi is false (MULTI queuing happens earlier with `continue`). + // Read commands execute directly on the target shard's database via RwLock read guard, + // avoiding ~88us of two async scheduling hops through the SPSC channel. + // + // Guard: if there are already pending writes for this target shard in the + // current pipeline batch, we must NOT take the fast path -- the read would + // execute before the deferred writes, violating command ordering. Fall through + // to SPSC dispatch to preserve pipeline semantics. + if !metadata::is_write(cmd) && !remote_groups.contains_key(&target) { + crate::admin::metrics_setup::record_dispatch_cross_read_fastpath(); + let guard = ctx.shard_databases.read_db(target, conn.selected_db); + let now_ms = ctx.cached_clock.ms(); + let db_count = ctx.shard_databases.db_count(); + let result = dispatch_read(&guard, cmd, cmd_args, now_ms, &mut conn.selected_db, db_count); + drop(guard); + let response = match result { + DispatchResult::Response(f) => f, + DispatchResult::Quit(f) => { should_quit = true; f } + }; + if matches!(response, Frame::Error(_)) { + if let Ok(cmd_str) = std::str::from_utf8(cmd) { + crate::admin::metrics_setup::record_command_error(cmd_str); + } + } + // Client tracking for cross-shard reads + if conn.tracking_state.enabled && !conn.tracking_state.bcast { + if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { + ctx.tracking_table.borrow_mut().track_key(client_id, &key, conn.tracking_state.noloop); + } + } + let mut response = apply_resp3_conversion(cmd, response, conn.protocol_version); + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + continue; + } + // Cross-shard write: deferred SPSC dispatch. + // When workspace rewriting occurred, rebuild the frame with + // prefixed args so the target shard stores the correct key. + let dispatch_frame = if rewritten.is_some() { + let mut parts = Vec::with_capacity(1 + cmd_args.len()); + parts.push(Frame::BulkString(Bytes::copy_from_slice(cmd))); + parts.extend_from_slice(cmd_args); + Frame::Array(parts.into()) + } else { + frame + }; + let resp_idx = responses.len(); + responses.push(Frame::Null); + let cmd_bytes = if let Frame::Array(ref args) = dispatch_frame { + extract_bytes(&args[0]).unwrap_or_default() + } else { + Bytes::new() + }; + remote_groups.entry(target).or_default().push((resp_idx, std::sync::Arc::new(dispatch_frame), aof_bytes, cmd_bytes, conn.selected_db)); + crate::admin::metrics_setup::record_dispatch_cross_spsc(); + } + } + + // Phase 2: Dispatch deferred remote commands (zero-allocation via ResponseSlotPool) + if !remote_groups.is_empty() { + let mut reply_futures: Vec<(Vec<(usize, Option, Bytes)>, usize)> = Vec::with_capacity(remote_groups.len()); + for (target, entries) in remote_groups { + let slot_ptr = response_pool.slot_ptr(target); + // Use the db_index captured with the first command (all commands in a + // pipeline batch targeting the same shard share the same db_index). + let batch_db = entries.first().map(|(_, _, _, _, db)| *db).unwrap_or(conn.selected_db); + let (meta, commands): (Vec<(usize, Option, Bytes)>, Vec>) = + entries.into_iter().map(|(idx, arc_frame, aof, cmd, _db)| ((idx, aof, cmd), arc_frame)).unzip(); + let msg = ShardMessage::PipelineBatchSlotted { db_index: batch_db, commands, response_slot: crate::shard::dispatch::ResponseSlotPtr(slot_ptr) }; + let target_idx = ChannelMesh::target_index(ctx.shard_id, target); + { + let mut pending = msg; + loop { + let push_result = { let mut producers = ctx.dispatch_tx.borrow_mut(); producers[target_idx].try_push(pending) }; + match push_result { + Ok(()) => { ctx.spsc_notifiers[target].notify_one(); break; } + Err(val) => { pending = val; tokio::task::yield_now().await; } + } + } + } + reply_futures.push((meta, target)); + } + let proto_ver = conn.protocol_version; + for (meta, target) in reply_futures { + let shard_responses = response_pool.future_for(target).await; + for ((resp_idx, aof_bytes, cmd_name), resp) in meta.into_iter().zip(shard_responses) { + if let Some(bytes) = aof_bytes { + if !matches!(resp, Frame::Error(_)) { + if let Some(ref tx) = ctx.aof_tx { let _ = tx.try_send(AofMessage::Append(bytes)); } + } + } + responses[resp_idx] = apply_resp3_conversion(&cmd_name, resp, proto_ver); + } + } + } + + // Phase 3: Flush accumulated PUBLISH batches as PubSubPublishBatch messages + if !publish_batches.is_empty() { + let mut batch_slots: Vec<(std::sync::Arc, Vec)> = Vec::new(); + { + let mut producers = ctx.dispatch_tx.borrow_mut(); + for (target, entries) in publish_batches.drain() { + let n = entries.len(); + let slot = std::sync::Arc::new(crate::shard::dispatch::PubSubResponseSlot::with_counts(1, n)); + let resp_indices: Vec = entries.iter().map(|(idx, _, _)| *idx).collect(); + let pairs: Vec<(Bytes, Bytes)> = entries.into_iter().map(|(_, ch, msg)| (ch, msg)).collect(); + + let idx = ChannelMesh::target_index(ctx.shard_id, target); + let batch_msg = ShardMessage::PubSubPublishBatch { + pairs, + slot: slot.clone(), + }; + if producers[idx].try_push(batch_msg).is_ok() { + ctx.spsc_notifiers[target].notify_one(); + } else { + slot.add(0); // push failed, mark as done + } + batch_slots.push((slot, resp_indices)); + } + } + // Resolve all batch slots + for (slot, resp_indices) in &batch_slots { + crate::shard::dispatch::PubSubResponseFuture::new(slot.clone()).await; + for (i, resp_idx) in resp_indices.iter().enumerate() { + let remote_count = slot.counts[i].load(std::sync::atomic::Ordering::Relaxed); + if remote_count > 0 { + if let Frame::Integer(ref mut total) = responses[*resp_idx] { *total += remote_count; } + } + } + } + } + + arena.reset(); + + // AUTH rate limiting: delay response to slow down brute-force attacks + if auth_delay_ms > 0 { + tokio::time::sleep(std::time::Duration::from_millis(auth_delay_ms)).await; + } + + write_buf.clear(); + for response in &responses { + if conn.protocol_version >= 3 { + crate::protocol::serialize_resp3(response, &mut write_buf); + } else { + crate::protocol::serialize(response, &mut write_buf); + } + } + if stream.write_all(&write_buf).await.is_err() { + return (HandlerResult::Done, None); + } + + // Update registry with current state after each batch + crate::client_registry::update(client_id, |e| { + e.db = conn.selected_db; + e.last_cmd_at = std::time::Instant::now(); + e.flags = crate::client_registry::ClientFlags { + subscriber: conn.subscription_count > 0, + in_multi: conn.in_multi, + blocked: false, + }; + }); + + // Check if migration was triggered during frame processing. + // All responses for the current batch have been written, so the + // client sees no interruption -- TCP socket stays open. + if let Some(target_shard) = conn.migration_target { + let migrated_state = MigratedConnectionState { + selected_db: conn.selected_db, + authenticated: conn.authenticated, + client_name: conn.client_name.clone(), + protocol_version: conn.protocol_version, + current_user: conn.current_user.clone(), + flags: 0, + read_buf_remainder: read_buf.split(), + client_id, + peer_addr: peer_addr.clone(), + workspace_id: conn.workspace_id, + }; + return ( + HandlerResult::MigrateConnection { state: migrated_state, target_shard }, + Some(stream), + ); + } + + if write_buf.capacity() > 65536 { write_buf = BytesMut::with_capacity(8192); } + if read_buf.capacity() > 65536 { + let remaining = read_buf.split(); + read_buf = BytesMut::with_capacity(8192); + if !remaining.is_empty() { read_buf.extend_from_slice(&remaining); } + } + + if should_quit { break; } + } + _ = shutdown.cancelled() => { + write_buf.clear(); + let shutdown_err = Frame::Error(Bytes::from_static(b"ERR server shutting down")); + if conn.protocol_version >= 3 { + crate::protocol::serialize_resp3(&shutdown_err, &mut write_buf); + } else { + crate::protocol::serialize(&shutdown_err, &mut write_buf); + } + let _ = stream.write_all(&write_buf).await; + break; + } + } + } + + // Phase 166: release any leaked cross-store TXN (client disconnected mid-txn). + // Idempotent: TXN.ABORT already takes() active_cross_txn so this is a no-op if abort ran. + // Closes T-161-05 — without this, a disconnect after TXN.BEGIN + SET would leak + // kv_intents and pin the key invisible for all subsequent readers. + if let Some(txn) = conn.active_cross_txn.take() { + crate::transaction::abort::abort_cross_store_txn( + &ctx.shard_databases, + ctx.shard_id, + conn.selected_db, + txn, + ); + } + + // Clean up pub/sub subscriptions on disconnect + if conn.subscriber_id > 0 { + let removed_channels = { + ctx.pubsub_registry + .write() + .unsubscribe_all(conn.subscriber_id) + }; + let removed_patterns = { + ctx.pubsub_registry + .write() + .punsubscribe_all(conn.subscriber_id) + }; + for ch in removed_channels { + unpropagate_subscription( + &ctx.all_remote_sub_maps, + &ch, + ctx.shard_id, + ctx.num_shards, + false, + ); + } + for pat in removed_patterns { + unpropagate_subscription( + &ctx.all_remote_sub_maps, + &pat, + ctx.shard_id, + ctx.num_shards, + true, + ); + } + // Remove affinity on disconnect (no subscriptions remain) + if let Ok(addr) = peer_addr.parse::() { + ctx.pubsub_affinity.write().remove(&addr.ip()); + } + } + + if conn.tracking_state.enabled { + ctx.tracking_table.borrow_mut().untrack_all(client_id); + } + + (HandlerResult::Done, None) +} diff --git a/src/server/conn/handler_sharded/pubsub.rs b/src/server/conn/handler_sharded/pubsub.rs new file mode 100644 index 00000000..05f6d6c9 --- /dev/null +++ b/src/server/conn/handler_sharded/pubsub.rs @@ -0,0 +1,529 @@ +//! Pub/sub command handlers: subscriber-mode loop, SUBSCRIBE/PSUBSCRIBE, +//! PUBLISH, UNSUBSCRIBE/PUNSUBSCRIBE, and PUBSUB introspection. + +use bytes::{Bytes, BytesMut}; +use std::collections::HashMap; + +use crate::protocol::Frame; +use crate::pubsub::{self, subscriber::Subscriber}; +use crate::runtime::cancel::CancellationToken; +use crate::runtime::channel; +use crate::server::conn::core::{ConnectionContext, ConnectionState}; +use crate::server::conn::util::{ + extract_bytes, extract_command, propagate_subscription, unpropagate_subscription, +}; + +/// Signal from subscriber-mode or subscribe-entry helpers to the outer loop. +pub(super) enum SubscriberAction { + /// Normal: re-enter the main loop (connection still in subscriber mode or exited it). + Continue, + /// Break the outer connection loop (QUIT, stream error, shutdown). + BreakOuter, + /// Return `(HandlerResult::Done, None)` immediately (unrecoverable parse error). + EarlyReturn, +} + +/// Run one iteration of the subscriber-mode select loop. +/// +/// Called when `conn.subscription_count > 0`. Processes incoming client commands +/// (limited to (P)(UN)SUBSCRIBE, PING, QUIT, RESET) and published messages. +/// +/// Returns `SubscriberAction` to tell the caller what to do next. +pub(super) async fn run_subscriber_step( + stream: &mut S, + read_buf: &mut BytesMut, + write_buf: &mut BytesMut, + parse_config: &crate::protocol::ParseConfig, + conn: &mut ConnectionState, + ctx: &ConnectionContext, + peer_addr: &str, + shutdown: &CancellationToken, +) -> SubscriberAction { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + #[allow(clippy::unwrap_used)] + // conn.pubsub_rx is always Some when conn.subscription_count > 0 + let rx = conn.pubsub_rx.as_mut().unwrap(); + tokio::select! { + n = stream.read_buf(read_buf) => { + match n { + Ok(0) | Err(_) => return SubscriberAction::BreakOuter, + Ok(_) => {} + } + let mut sub_break = false; + loop { + match crate::protocol::parse(read_buf, parse_config) { + Ok(Some(frame)) => { + if let Some((cmd, cmd_args)) = extract_command(&frame) { + if cmd.eq_ignore_ascii_case(b"SUBSCRIBE") { + if cmd_args.is_empty() { + let err = Frame::Error(Bytes::from_static(b"ERR wrong number of arguments for 'subscribe' command")); + write_buf.clear(); + crate::protocol::serialize(&err, write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + continue; + } + for arg in cmd_args { + if let Some(ch) = extract_bytes(arg) { + #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable + let acl_deny = { ctx.acl_table.read().unwrap().check_channel_permission(&conn.current_user, ch.as_ref()) }; + if let Some(reason) = acl_deny { + let err = Frame::Error(Bytes::from(format!("NOPERM {}", reason))); + write_buf.clear(); + crate::protocol::serialize(&err, write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + continue; + } + #[allow(clippy::unwrap_used)] // conn.pubsub_tx is always Some in subscriber mode + let sub = Subscriber::with_protocol( + conn.pubsub_tx.clone().unwrap(), + conn.subscriber_id, + conn.protocol_version >= 3, + ); + { ctx.pubsub_registry.write().subscribe(ch.clone(), sub); } + conn.subscription_count += 1; + // Register pub/sub affinity for this client IP + if conn.subscription_count == 1 { + if let Ok(addr) = peer_addr.parse::() { + ctx.pubsub_affinity.write().register(addr.ip(), ctx.shard_id); + } + } + propagate_subscription(&ctx.all_remote_sub_maps, &ch, ctx.shard_id, ctx.num_shards, false); + write_buf.clear(); + let resp = pubsub::subscribe_response(&ch, conn.subscription_count); + crate::protocol::serialize(&resp, write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + } + } + } else if cmd.eq_ignore_ascii_case(b"PSUBSCRIBE") { + if cmd_args.is_empty() { + let err = Frame::Error(Bytes::from_static(b"ERR wrong number of arguments for 'psubscribe' command")); + write_buf.clear(); + crate::protocol::serialize(&err, write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + continue; + } + for arg in cmd_args { + if let Some(pat) = extract_bytes(arg) { + #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable + let acl_deny = { ctx.acl_table.read().unwrap().check_channel_permission(&conn.current_user, pat.as_ref()) }; + if let Some(reason) = acl_deny { + let err = Frame::Error(Bytes::from(format!("NOPERM {}", reason))); + write_buf.clear(); + crate::protocol::serialize(&err, write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + continue; + } + #[allow(clippy::unwrap_used)] // conn.pubsub_tx is always Some in subscriber mode + let sub = Subscriber::with_protocol( + conn.pubsub_tx.clone().unwrap(), + conn.subscriber_id, + conn.protocol_version >= 3, + ); + { ctx.pubsub_registry.write().psubscribe(pat.clone(), sub); } + conn.subscription_count += 1; + // Register pub/sub affinity for this client IP + if conn.subscription_count == 1 { + if let Ok(addr) = peer_addr.parse::() { + ctx.pubsub_affinity.write().register(addr.ip(), ctx.shard_id); + } + } + propagate_subscription(&ctx.all_remote_sub_maps, &pat, ctx.shard_id, ctx.num_shards, true); + write_buf.clear(); + let resp = pubsub::psubscribe_response(&pat, conn.subscription_count); + crate::protocol::serialize(&resp, write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + } + } + } else if cmd.eq_ignore_ascii_case(b"UNSUBSCRIBE") { + if cmd_args.is_empty() { + let removed = { ctx.pubsub_registry.write().unsubscribe_all(conn.subscriber_id) }; + if removed.is_empty() { + conn.subscription_count = ctx.pubsub_registry.read().total_subscription_count(conn.subscriber_id); + write_buf.clear(); + crate::protocol::serialize(&pubsub::unsubscribe_response(&Bytes::from_static(b""), conn.subscription_count), write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + } else { + for ch in &removed { + conn.subscription_count = conn.subscription_count.saturating_sub(1); + unpropagate_subscription(&ctx.all_remote_sub_maps, ch, ctx.shard_id, ctx.num_shards, false); + write_buf.clear(); + crate::protocol::serialize(&pubsub::unsubscribe_response(ch, conn.subscription_count), write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + } + } + } else { + for arg in cmd_args { + if let Some(ch) = extract_bytes(arg) { + { ctx.pubsub_registry.write().unsubscribe(ch.as_ref(), conn.subscriber_id); } + conn.subscription_count = conn.subscription_count.saturating_sub(1); + unpropagate_subscription(&ctx.all_remote_sub_maps, &ch, ctx.shard_id, ctx.num_shards, false); + write_buf.clear(); + crate::protocol::serialize(&pubsub::unsubscribe_response(&ch, conn.subscription_count), write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + } + } + } + } else if cmd.eq_ignore_ascii_case(b"PUNSUBSCRIBE") { + if cmd_args.is_empty() { + let removed = { ctx.pubsub_registry.write().punsubscribe_all(conn.subscriber_id) }; + if removed.is_empty() { + conn.subscription_count = ctx.pubsub_registry.read().total_subscription_count(conn.subscriber_id); + write_buf.clear(); + crate::protocol::serialize(&pubsub::punsubscribe_response(&Bytes::from_static(b""), conn.subscription_count), write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + } else { + for pat in &removed { + conn.subscription_count = conn.subscription_count.saturating_sub(1); + unpropagate_subscription(&ctx.all_remote_sub_maps, pat, ctx.shard_id, ctx.num_shards, true); + write_buf.clear(); + crate::protocol::serialize(&pubsub::punsubscribe_response(pat, conn.subscription_count), write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + } + } + } else { + for arg in cmd_args { + if let Some(pat) = extract_bytes(arg) { + { ctx.pubsub_registry.write().punsubscribe(pat.as_ref(), conn.subscriber_id); } + conn.subscription_count = conn.subscription_count.saturating_sub(1); + unpropagate_subscription(&ctx.all_remote_sub_maps, &pat, ctx.shard_id, ctx.num_shards, true); + write_buf.clear(); + crate::protocol::serialize(&pubsub::punsubscribe_response(&pat, conn.subscription_count), write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + } + } + } + } else if cmd.eq_ignore_ascii_case(b"PING") { + write_buf.clear(); + let resp = Frame::Array(crate::framevec![ + Frame::BulkString(Bytes::from_static(b"pong")), + Frame::BulkString(Bytes::from_static(b"")), + ]); + crate::protocol::serialize(&resp, write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + } else if cmd.eq_ignore_ascii_case(b"QUIT") { + write_buf.clear(); + crate::protocol::serialize(&Frame::SimpleString(Bytes::from_static(b"OK")), write_buf); + let _ = stream.write_all(write_buf).await; + sub_break = true; + break; + } else if cmd.eq_ignore_ascii_case(b"RESET") { + { ctx.pubsub_registry.write().unsubscribe_all(conn.subscriber_id); } + { ctx.pubsub_registry.write().punsubscribe_all(conn.subscriber_id); } + conn.subscription_count = 0; + write_buf.clear(); + crate::protocol::serialize(&Frame::SimpleString(Bytes::from_static(b"RESET")), write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + } else { + let cmd_str = String::from_utf8_lossy(cmd); + let err = Frame::Error(Bytes::from(format!( + "ERR Can't execute '{}': only (P)SUBSCRIBE / (P)UNSUBSCRIBE / PING / QUIT are allowed in this context", + cmd_str.to_lowercase() + ))); + write_buf.clear(); + crate::protocol::serialize(&err, write_buf); + if stream.write_all(write_buf).await.is_err() { sub_break = true; break; } + } + } + if conn.subscription_count == 0 { break; } + } + Ok(None) => break, + Err(_) => { return SubscriberAction::EarlyReturn; } + } + } + if sub_break { return SubscriberAction::BreakOuter; } + if conn.subscription_count == 0 { return SubscriberAction::Continue; } + } + msg = rx.recv_async() => { + match msg { + Ok(data) => { + // Data is pre-serialized RESP bytes — write directly + if stream.write_all(&data).await.is_err() { return SubscriberAction::BreakOuter; } + } + Err(_) => return SubscriberAction::BreakOuter, + } + } + _ = shutdown.cancelled() => { + write_buf.clear(); + crate::protocol::serialize(&Frame::Error(Bytes::from_static(b"ERR server shutting down")), write_buf); + let _ = stream.write_all(write_buf).await; + return SubscriberAction::BreakOuter; + } + } + SubscriberAction::Continue +} + +/// Handle SUBSCRIBE / PSUBSCRIBE in normal (non-subscriber) mode. +/// +/// Returns `Some(SubscriberAction)` if the command was consumed, `None` otherwise. +/// The caller should `break` from the frame batch loop on `BreakOuter` (to re-enter +/// subscriber mode) or `return (HandlerResult::Done, None)` on `EarlyReturn`. +pub(super) async fn try_handle_subscribe< + S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, +>( + cmd: &[u8], + cmd_args: &[Frame], + stream: &mut S, + write_buf: &mut BytesMut, + conn: &mut ConnectionState, + ctx: &ConnectionContext, + peer_addr: &str, + responses: &mut Vec, +) -> Option { + use tokio::io::AsyncWriteExt; + + if !cmd.eq_ignore_ascii_case(b"SUBSCRIBE") && !cmd.eq_ignore_ascii_case(b"PSUBSCRIBE") { + return None; + } + let is_pattern = cmd.eq_ignore_ascii_case(b"PSUBSCRIBE"); + let cmd_name = if is_pattern { + "psubscribe" + } else { + "subscribe" + }; + if cmd_args.is_empty() { + responses.push(Frame::Error(Bytes::from(format!( + "ERR wrong number of arguments for '{}' command", + cmd_name + )))); + return Some(SubscriberAction::Continue); + } + // Allocate pubsub channel if not yet created + if conn.pubsub_tx.is_none() { + let (tx, rx) = channel::mpsc_bounded::(256); + conn.pubsub_tx = Some(tx); + conn.pubsub_rx = Some(rx); + } + if conn.subscriber_id == 0 { + conn.subscriber_id = crate::pubsub::next_subscriber_id(); + } + // Flush accumulated responses before entering subscriber mode + if !responses.is_empty() { + write_buf.clear(); + for resp in responses.iter() { + if conn.protocol_version >= 3 { + crate::protocol::serialize_resp3(resp, write_buf); + } else { + crate::protocol::serialize(resp, write_buf); + } + } + if stream.write_all(write_buf).await.is_err() { + return Some(SubscriberAction::EarlyReturn); + } + responses.clear(); + } + // Process subscribe arguments + for arg in cmd_args { + if let Some(ch) = extract_bytes(arg) { + #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable + let acl_deny = { + ctx.acl_table + .read() + .unwrap() + .check_channel_permission(&conn.current_user, ch.as_ref()) + }; + if let Some(reason) = acl_deny { + write_buf.clear(); + let err = Frame::Error(Bytes::from(format!("NOPERM {}", reason))); + crate::protocol::serialize(&err, write_buf); + if stream.write_all(write_buf).await.is_err() { + return Some(SubscriberAction::EarlyReturn); + } + continue; + } + #[allow(clippy::unwrap_used)] + // conn.pubsub_tx is set to Some just above before this loop + let sub = Subscriber::with_protocol( + conn.pubsub_tx.clone().unwrap(), + conn.subscriber_id, + conn.protocol_version >= 3, + ); + if is_pattern { + ctx.pubsub_registry.write().psubscribe(ch.clone(), sub); + } else { + ctx.pubsub_registry.write().subscribe(ch.clone(), sub); + } + conn.subscription_count += 1; + // Register pub/sub affinity for this client IP + if conn.subscription_count == 1 { + if let Ok(addr) = peer_addr.parse::() { + ctx.pubsub_affinity + .write() + .register(addr.ip(), ctx.shard_id); + } + } + propagate_subscription( + &ctx.all_remote_sub_maps, + &ch, + ctx.shard_id, + ctx.num_shards, + is_pattern, + ); + write_buf.clear(); + let resp = if is_pattern { + pubsub::psubscribe_response(&ch, conn.subscription_count) + } else { + pubsub::subscribe_response(&ch, conn.subscription_count) + }; + crate::protocol::serialize(&resp, write_buf); + if stream.write_all(write_buf).await.is_err() { + return Some(SubscriberAction::EarlyReturn); + } + } + } + // break out of frame batch loop to re-enter main loop in subscriber mode + Some(SubscriberAction::BreakOuter) +} + +/// Handle UNSUBSCRIBE / PUNSUBSCRIBE in normal (non-subscriber) mode. +/// Returns `true` if the command was consumed (caller should `continue`). +pub(super) fn try_handle_unsubscribe(cmd: &[u8], responses: &mut Vec) -> bool { + if !cmd.eq_ignore_ascii_case(b"UNSUBSCRIBE") && !cmd.eq_ignore_ascii_case(b"PUNSUBSCRIBE") { + return false; + } + let is_pattern = cmd.eq_ignore_ascii_case(b"PUNSUBSCRIBE"); + let resp = if is_pattern { + pubsub::punsubscribe_response(&Bytes::from_static(b""), 0) + } else { + pubsub::unsubscribe_response(&Bytes::from_static(b""), 0) + }; + responses.push(resp); + true +} + +/// Handle PUBLISH command. +/// Returns `true` if the command was consumed (caller should `continue`). +pub(super) fn try_handle_publish( + cmd: &[u8], + cmd_args: &[Frame], + conn: &ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, + publish_batches: &mut HashMap>, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"PUBLISH") { + return false; + } + if cmd_args.len() != 2 { + responses.push(Frame::Error(Bytes::from_static( + b"ERR wrong number of arguments for 'publish' command", + ))); + } else { + let channel_arg = extract_bytes(&cmd_args[0]); + let message_arg = extract_bytes(&cmd_args[1]); + // ACL channel permission check for PUBLISH + if let Some(ref ch) = channel_arg { + #[allow(clippy::unwrap_used)] // std RwLock: poison = prior panic = unrecoverable + let acl_guard = ctx.acl_table.read().unwrap(); + if let Some(deny_reason) = + acl_guard.check_channel_permission(&conn.current_user, ch.as_ref()) + { + drop(acl_guard); + responses.push(Frame::Error(Bytes::from(format!("NOPERM {}", deny_reason)))); + return true; + } + } + match (channel_arg, message_arg) { + (Some(ch), Some(msg)) => { + let local_count = { ctx.pubsub_registry.write().publish(&ch, &msg) }; + // Targeted fanout: only send to shards that have subscribers + let targets = ctx.remote_subscriber_map.read().target_shards(&ch); + if targets.is_empty() { + // Fast path: no remote subscribers, return local count immediately + responses.push(Frame::Integer(local_count)); + } else { + // Filter to remote targets only (skip self) + let remote_targets: Vec = + targets.into_iter().filter(|&t| t != ctx.shard_id).collect(); + if remote_targets.is_empty() { + responses.push(Frame::Integer(local_count)); + } else { + // Accumulate into per-shard batches for coalesced dispatch + let resp_idx = responses.len(); + responses.push(Frame::Integer(local_count)); // placeholder, updated after batch flush + for target in &remote_targets { + publish_batches.entry(*target).or_default().push(( + resp_idx, + ch.clone(), + msg.clone(), + )); + } + } + } + } + _ => responses.push(Frame::Error(Bytes::from_static( + b"ERR invalid channel or message", + ))), + } + } + true +} + +/// Handle PUBSUB introspection commands (CHANNELS, NUMSUB, NUMPAT). +/// Returns `true` if the command was consumed (caller should `continue`). +pub(super) fn try_handle_pubsub_introspection( + cmd: &[u8], + cmd_args: &[Frame], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !cmd.eq_ignore_ascii_case(b"PUBSUB") { + return false; + } + if cmd_args.is_empty() { + responses.push(Frame::Error(Bytes::from_static( + b"ERR wrong number of arguments for 'pubsub' command", + ))); + return true; + } + let subcmd = extract_bytes(&cmd_args[0]); + match subcmd { + Some(ref sc) if sc.eq_ignore_ascii_case(b"CHANNELS") => { + let pattern = if cmd_args.len() > 1 { + extract_bytes(&cmd_args[1]) + } else { + None + }; + let mut all_channels: std::collections::HashSet = + std::collections::HashSet::new(); + for reg in &ctx.all_pubsub_registries { + let guard = reg.read(); + all_channels.extend(guard.active_channels(pattern.as_deref())); + } + let arr: Vec = all_channels.into_iter().map(Frame::BulkString).collect(); + responses.push(Frame::Array(arr.into())); + } + Some(ref sc) if sc.eq_ignore_ascii_case(b"NUMSUB") => { + let channels: Vec = cmd_args[1..] + .iter() + .filter_map(|a| extract_bytes(a)) + .collect(); + let mut counts: HashMap = HashMap::new(); + for reg in &ctx.all_pubsub_registries { + let guard = reg.read(); + for (ch, c) in guard.numsub(&channels) { + *counts.entry(ch).or_insert(0) += c; + } + } + let mut arr = Vec::with_capacity(channels.len() * 2); + for ch in &channels { + arr.push(Frame::BulkString(ch.clone())); + arr.push(Frame::Integer(*counts.get(ch).unwrap_or(&0))); + } + responses.push(Frame::Array(arr.into())); + } + Some(ref sc) if sc.eq_ignore_ascii_case(b"NUMPAT") => { + let mut total: usize = 0; + for reg in &ctx.all_pubsub_registries { + total += reg.read().numpat(); + } + responses.push(Frame::Integer(total as i64)); + } + _ => { + responses.push(Frame::Error(Bytes::from_static( + b"ERR unknown subcommand or wrong number of arguments for 'pubsub' command", + ))); + } + } + true +} diff --git a/src/server/conn/handler_sharded/read.rs b/src/server/conn/handler_sharded/read.rs new file mode 100644 index 00000000..aaf1084d --- /dev/null +++ b/src/server/conn/handler_sharded/read.rs @@ -0,0 +1 @@ +// Extracted in Task 3: read-only shard dispatch paths. diff --git a/src/server/conn/handler_sharded/txn.rs b/src/server/conn/handler_sharded/txn.rs new file mode 100644 index 00000000..0d157d38 --- /dev/null +++ b/src/server/conn/handler_sharded/txn.rs @@ -0,0 +1,258 @@ +//! TXN.BEGIN / TXN.COMMIT / TXN.ABORT + TEMPORAL.SNAPSHOT_AT / TEMPORAL.INVALIDATE handlers. +//! +//! Each helper returns `true` if the command was consumed (caller should `continue`). + +use bytes::Bytes; + +#[cfg(feature = "graph")] +use crate::command::temporal::{ERR_ENTITY_NOT_FOUND, ERR_GRAPH_NOT_FOUND}; +use crate::command::temporal::{ + capture_wall_ms, is_temporal_invalidate, is_temporal_snapshot_at, validate_invalidate, + validate_snapshot_at, +}; +use crate::command::transaction::{ + is_txn_abort, is_txn_begin, is_txn_commit, txn_abort_validate, txn_begin_validate, + txn_commit_validate, +}; +use crate::protocol::Frame; +use crate::server::conn::core::ConnectionContext; +use crate::server::conn::core::ConnectionState; +use crate::transaction::CrossStoreTxn; + +/// Handle TXN.BEGIN — returns `true` if the command was consumed. +pub(super) fn try_handle_txn_begin( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !is_txn_begin(cmd, cmd_args) { + return false; + } + match txn_begin_validate(conn.in_multi, conn.in_cross_txn()) { + Ok(()) => { + // Get next txn_id and snapshot_lsn from vector store's transaction manager + let mut vector_store = ctx.shard_databases.vector_store(ctx.shard_id); + let active = vector_store.txn_manager_mut().begin(); + conn.active_cross_txn = Some(CrossStoreTxn::new(active.txn_id, active.snapshot_lsn)); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + Err(e) => responses.push(e), + } + true +} + +/// Handle TXN.COMMIT — returns `true` if the command was consumed. +pub(super) fn try_handle_txn_commit( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !is_txn_commit(cmd, cmd_args) { + return false; + } + match txn_commit_validate(conn.in_cross_txn()) { + Ok(()) => { + if let Some(txn) = conn.active_cross_txn.take() { + let mut vector_store = ctx.shard_databases.vector_store(ctx.shard_id); + vector_store.txn_manager_mut().commit(txn.txn_id); + drop(vector_store); + + // Write XactCommit WAL record with committed KV state + let txn_id = txn.txn_id; + if !txn.kv_undo.is_empty() { + let db_guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); + let payload = crate::persistence::wal_v3::record::encode_xact_commit_payload( + txn_id, + txn.kv_undo.records(), + &*db_guard, + ); + drop(db_guard); + let mut wal_buf = Vec::new(); + crate::persistence::wal_v3::record::write_wal_v3_record( + &mut wal_buf, + txn_id, + crate::persistence::wal_v3::record::WalRecordType::XactCommit, + &payload, + ); + ctx.shard_databases + .wal_append(ctx.shard_id, bytes::Bytes::from(wal_buf)); + } + + // Release KV write intents from shard side-table + ctx.shard_databases + .kv_intents(ctx.shard_id) + .release_txn(txn_id); + + // Drain deferred HNSW inserts (post-commit hook). + // The drain prevents phantom neighbors on abort. + // Actual HNSW graph insertion happens during compaction, + // not at commit time (point is already in mutable segment). + let drain_count = ctx + .shard_databases + .hnsw_queue(ctx.shard_id) + .drain_for_txn(txn_id) + .count(); + if drain_count > 0 { + tracing::debug!(txn_id, count = drain_count, "Drained deferred HNSW inserts"); + } + + // Materialize MQ intents: enqueue deferred MQ.PUBLISH messages + if !txn.mq_intents.is_empty() { + let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + for intent in &txn.mq_intents { + if let Ok(Some(stream)) = db_guard.get_stream_mut(&intent.queue_key) { + if stream.durable { + let msg_id = stream.next_auto_id(); + stream.add(msg_id, intent.fields.clone()); + } + } + } + } + + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } else { + responses.push(Frame::Error(Bytes::from_static(b"ERR not in transaction"))); + } + } + Err(e) => responses.push(e), + } + true +} + +/// Handle TXN.ABORT ��� returns `true` if the command was consumed. +pub(super) fn try_handle_txn_abort( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !is_txn_abort(cmd, cmd_args) { + return false; + } + match txn_abort_validate(conn.in_cross_txn()) { + Ok(()) => { + if let Some(txn) = conn.active_cross_txn.take() { + // Shared rollback (Phase 166 Plan 03): + // KV undo -> graph intents reverse -> vector + // tombstone -> side-table release. See + // src/transaction/abort.rs for lock ordering. + crate::transaction::abort::abort_cross_store_txn( + &ctx.shard_databases, + ctx.shard_id, + conn.selected_db, + txn, + ); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } else { + responses.push(Frame::Error(Bytes::from_static(b"ERR not in transaction"))); + } + } + Err(e) => responses.push(e), + } + true +} + +/// Handle TEMPORAL.SNAPSHOT_AT — returns `true` if the command was consumed. +pub(super) fn try_handle_temporal_snapshot_at( + cmd: &[u8], + cmd_args: &[Frame], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !is_temporal_snapshot_at(cmd) { + return false; + } + match validate_snapshot_at(cmd_args) { + Ok(()) => { + let wall_ms = capture_wall_ms(); + let lsn = { + let vector_store = ctx.shard_databases.vector_store(ctx.shard_id); + vector_store.txn_manager().current_lsn() + }; + { + let mut guard = ctx.shard_databases.temporal_registry(ctx.shard_id); + let registry = + guard.get_or_insert_with(|| Box::new(crate::temporal::TemporalRegistry::new())); + registry.record(wall_ms, lsn); + } + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + Err(e) => responses.push(e), + } + true +} + +/// Handle TEMPORAL.INVALIDATE — returns `true` if the command was consumed. +pub(super) fn try_handle_temporal_invalidate( + cmd: &[u8], + cmd_args: &[Frame], + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !is_temporal_invalidate(cmd) { + return false; + } + match validate_invalidate(cmd_args) { + Ok((entity_id, is_node, graph_name)) => { + let wall_ms = capture_wall_ms(); + #[cfg(feature = "graph")] + { + let mut gs = ctx.shard_databases.graph_store_write(ctx.shard_id); + if let Some(named_graph) = gs.get_graph_mut(&graph_name) { + let mutated = if is_node { + let node_key: crate::graph::types::NodeKey = + slotmap::KeyData::from_ffi(entity_id).into(); + if let Some(node) = named_graph.write_buf.get_node_mut(node_key) { + node.valid_to = wall_ms; + true + } else { + false + } + } else { + let edge_key: crate::graph::types::EdgeKey = + slotmap::KeyData::from_ffi(entity_id).into(); + if let Some(edge) = named_graph.write_buf.get_edge_mut(edge_key) { + edge.valid_to = wall_ms; + true + } else { + false + } + }; + if mutated { + let payload = crate::persistence::wal_v3::record::encode_graph_temporal( + entity_id, is_node, wall_ms, wall_ms, + ); + gs.wal_pending.push(payload); + let wal_records = gs.drain_wal(); + drop(gs); + for record in wal_records { + ctx.shard_databases + .wal_append(ctx.shard_id, Bytes::from(record)); + } + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } else { + drop(gs); + responses.push(Frame::Error(Bytes::from_static(ERR_ENTITY_NOT_FOUND))); + } + } else { + drop(gs); + responses.push(Frame::Error(Bytes::from_static(ERR_GRAPH_NOT_FOUND))); + } + } + #[cfg(not(feature = "graph"))] + { + let _ = (entity_id, is_node, graph_name, wall_ms, ctx); + responses.push(Frame::Error(Bytes::from_static( + b"ERR graph feature not enabled", + ))); + } + } + Err(e) => responses.push(e), + } + true +} diff --git a/src/server/conn/handler_sharded/write.rs b/src/server/conn/handler_sharded/write.rs new file mode 100644 index 00000000..5e32ae03 --- /dev/null +++ b/src/server/conn/handler_sharded/write.rs @@ -0,0 +1,717 @@ +//! Write-path command handlers: WS.*, MQ.*, MULTI/EXEC/DISCARD, GRAPH.*. +//! +//! Each helper returns `true` if the command was consumed (caller should `continue`). + +use bytes::Bytes; + +use crate::command::mq::{ + ERR_MQ_NOT_DURABLE, ERR_MQ_UNKNOWN_SUB, parse_mq_subcommand, validate_mq_ack, + validate_mq_create, validate_mq_dlqlen, validate_mq_pop, validate_mq_publish, validate_mq_push, + validate_mq_trigger, +}; +use crate::command::transaction::ERR_MULTI_TXN_CONFLICT; +use crate::command::workspace::{ + ERR_WS_ALREADY_BOUND, ERR_WS_NOT_FOUND, ERR_WS_UNKNOWN_SUB, parse_workspace_id_from_bytes, + parse_ws_subcommand, validate_ws_auth, validate_ws_create, validate_ws_drop, validate_ws_info, + validate_ws_list, +}; +use crate::mq::is_mq_command; +use crate::protocol::Frame; +use crate::server::conn::core::{ConnectionContext, ConnectionState}; +#[cfg(feature = "graph")] +use crate::server::conn::util::extract_bytes; +use crate::storage::stream::StreamId; +#[cfg(feature = "graph")] +use crate::workspace::strip_workspace_prefix_from_response; +use crate::workspace::{WorkspaceId, is_ws_command}; + +use super::execute_transaction_sharded; + +/// Handle WS.* workspace commands. Returns `true` if consumed. +pub(super) fn try_handle_ws_command( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !is_ws_command(cmd) { + return false; + } + let sub = match parse_ws_subcommand(cmd_args) { + Ok(s) => s, + Err(e) => { + responses.push(e); + return true; + } + }; + + if sub.eq_ignore_ascii_case(b"CREATE") { + match validate_ws_create(cmd_args) { + Ok(ws_name) => { + let ws_id = WorkspaceId::new_v7(); + let created_at = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as i64; + let meta = crate::workspace::WorkspaceMetadata { + id: ws_id, + name: ws_name.clone(), + created_at, + }; + { + let mut guard = ctx.shard_databases.workspace_registry(ctx.shard_id); + let reg = guard.get_or_insert_with(|| { + Box::new(crate::workspace::WorkspaceRegistry::new()) + }); + reg.insert(ws_id, meta); + } + // WAL: WorkspaceCreate record + let payload = + crate::workspace::wal::encode_workspace_create(ws_id.as_bytes(), &ws_name); + let mut wal_buf = Vec::new(); + crate::persistence::wal_v3::record::write_wal_v3_record( + &mut wal_buf, + 0, + crate::persistence::wal_v3::record::WalRecordType::WorkspaceCreate, + &payload, + ); + ctx.shard_databases + .wal_append(ctx.shard_id, Bytes::from(wal_buf)); + responses.push(Frame::BulkString(Bytes::from(ws_id.to_string()))); + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"DROP") { + match validate_ws_drop(cmd_args) { + Ok(ws_id_raw) => match parse_workspace_id_from_bytes(&ws_id_raw) { + Some(ws_id) => { + let removed = { + let mut guard = ctx.shard_databases.workspace_registry(ctx.shard_id); + match guard.as_mut() { + Some(reg) => reg.remove(&ws_id).is_some(), + None => false, + } + }; + if removed { + // WAL: WorkspaceDrop record + let payload = + crate::workspace::wal::encode_workspace_drop(ws_id.as_bytes()); + let mut wal_buf = Vec::new(); + crate::persistence::wal_v3::record::write_wal_v3_record( + &mut wal_buf, + 0, + crate::persistence::wal_v3::record::WalRecordType::WorkspaceDrop, + &payload, + ); + ctx.shard_databases + .wal_append(ctx.shard_id, Bytes::from(wal_buf)); + // Best-effort cleanup: delete all KV keys with ws + // prefix (WS-03). + { + let prefix = format!("{{{}}}:", ws_id.as_hex()); + let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, 0); + let keys_to_delete: Vec> = db_guard + .keys() + .filter(|k| k.as_bytes().starts_with(prefix.as_bytes())) + .map(|k| k.as_bytes().to_vec()) + .collect(); + for key in &keys_to_delete { + db_guard.remove(key); + } + } + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } else { + responses.push(Frame::Error(Bytes::from_static(ERR_WS_NOT_FOUND))); + } + } + None => responses.push(Frame::Error(Bytes::from_static( + crate::command::workspace::ERR_WS_INVALID_ID, + ))), + }, + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"LIST") { + match validate_ws_list(cmd_args) { + Ok(()) => { + let guard = ctx.shard_databases.workspace_registry(ctx.shard_id); + let entries: Vec = match guard.as_ref() { + Some(reg) => reg + .iter() + .map(|(id, meta)| { + Frame::Array( + vec![ + Frame::BulkString(Bytes::from(id.to_string())), + Frame::BulkString(meta.name.clone()), + Frame::Integer(meta.created_at), + ] + .into(), + ) + }) + .collect(), + None => vec![], + }; + responses.push(Frame::Array(entries.into())); + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"INFO") { + match validate_ws_info(cmd_args) { + Ok(ws_id_raw) => match parse_workspace_id_from_bytes(&ws_id_raw) { + Some(ws_id) => { + let guard = ctx.shard_databases.workspace_registry(ctx.shard_id); + let found = guard.as_ref().and_then(|reg| reg.get(&ws_id)); + match found { + Some(meta) => { + responses.push(Frame::Array( + vec![ + Frame::BulkString(Bytes::from_static(b"id")), + Frame::BulkString(Bytes::from(meta.id.to_string())), + Frame::BulkString(Bytes::from_static(b"name")), + Frame::BulkString(meta.name.clone()), + Frame::BulkString(Bytes::from_static(b"created_at")), + Frame::Integer(meta.created_at), + ] + .into(), + )); + } + None => responses.push(Frame::Error(Bytes::from_static(ERR_WS_NOT_FOUND))), + } + } + None => responses.push(Frame::Error(Bytes::from_static( + crate::command::workspace::ERR_WS_INVALID_ID, + ))), + }, + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"AUTH") { + match validate_ws_auth(cmd_args) { + Ok(ws_id_raw) => { + if conn.workspace_id.is_some() { + responses.push(Frame::Error(Bytes::from_static(ERR_WS_ALREADY_BOUND))); + } else { + match parse_workspace_id_from_bytes(&ws_id_raw) { + Some(ws_id) => { + let found = { + let guard = ctx.shard_databases.workspace_registry(ctx.shard_id); + guard + .as_ref() + .map_or(false, |reg| reg.get(&ws_id).is_some()) + }; + if found { + conn.workspace_id = Some(ws_id); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } else { + responses.push(Frame::Error(Bytes::from_static(ERR_WS_NOT_FOUND))); + } + } + None => responses.push(Frame::Error(Bytes::from_static( + crate::command::workspace::ERR_WS_INVALID_ID, + ))), + } + } + } + Err(e) => responses.push(e), + } + return true; + } + + // Unknown WS subcommand + responses.push(Frame::Error(Bytes::from_static(ERR_WS_UNKNOWN_SUB))); + true +} + +/// Handle MQ.* message queue commands. Returns `true` if consumed. +pub(super) fn try_handle_mq_command( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if !is_mq_command(cmd) { + return false; + } + let sub = match parse_mq_subcommand(cmd_args) { + Ok(s) => s, + Err(e) => { + responses.push(e); + return true; + } + }; + + if sub.eq_ignore_ascii_case(b"CREATE") { + match validate_mq_create(cmd_args) { + Ok((queue_key, max_delivery_count, _debounce_ms)) => { + let effective_key = + crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); + let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + match db_guard.get_or_create_stream(&effective_key) { + Ok(stream) => { + stream.durable = true; + stream.max_delivery_count = max_delivery_count; + let group_name = Bytes::from_static(b"__mq_consumers"); + let _ = stream.create_group(group_name, StreamId::ZERO); + } + Err(e) => { + responses.push(e); + return true; + } + } + drop(db_guard); + + let config = + crate::mq::DurableStreamConfig::new(effective_key.clone(), max_delivery_count); + { + let mut guard = ctx.shard_databases.durable_queue_registry(ctx.shard_id); + let reg = guard + .get_or_insert_with(|| Box::new(crate::mq::DurableQueueRegistry::new())); + reg.insert(effective_key.clone(), config); + } + + let payload = crate::mq::wal::encode_mq_create(&effective_key, max_delivery_count); + let mut wal_buf = Vec::new(); + crate::persistence::wal_v3::record::write_wal_v3_record( + &mut wal_buf, + 0, + crate::persistence::wal_v3::record::WalRecordType::MqCreate, + &payload, + ); + ctx.shard_databases + .wal_append(ctx.shard_id, Bytes::from(wal_buf)); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"PUSH") { + match validate_mq_push(cmd_args) { + Ok((queue_key, fields)) => { + let effective_key = + crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); + let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + match db_guard.get_stream_mut(&effective_key) { + Ok(Some(stream)) => { + if !stream.durable { + responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); + } else { + let msg_id = stream.next_auto_id(); + let msg_id = stream.add(msg_id, fields); + drop(db_guard); + { + let mut trig_guard = + ctx.shard_databases.trigger_registry(ctx.shard_id); + if let Some(reg) = trig_guard.as_mut() { + let trig_key = if let Some(ws_id) = conn.workspace_id.as_ref() { + let ws_hex = ws_id.as_hex(); + let mut k = + Vec::with_capacity(ws_hex.len() + 1 + queue_key.len()); + k.extend_from_slice(ws_hex.as_bytes()); + k.push(b':'); + k.extend_from_slice(&queue_key); + Bytes::from(k) + } else { + queue_key.clone() + }; + if let Some(trig_entry) = reg.get_mut(&trig_key) { + if trig_entry.pending_fire_ms == 0 { + let fire_at = + ctx.cached_clock.ms() + trig_entry.debounce_ms; + trig_entry.pending_fire_ms = fire_at; + } + } + } + } + responses.push(Frame::BulkString(Bytes::from(format!( + "{}-{}", + msg_id.ms, msg_id.seq + )))); + } + } + Ok(None) => { + responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); + } + Err(e) => responses.push(e), + } + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"POP") { + match validate_mq_pop(cmd_args) { + Ok((queue_key, count)) => { + let effective_key = + crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); + let group_name = Bytes::from_static(b"__mq_consumers"); + let consumer_name = Bytes::from_static(b"__mq_default"); + let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + + let mdc = match db_guard.get_stream_mut(&effective_key) { + Ok(Some(stream)) => { + if !stream.durable { + responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); + return true; + } + stream.max_delivery_count + } + Ok(None) => { + responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); + return true; + } + Err(e) => { + responses.push(e); + return true; + } + }; + + let request_count = count + (mdc as usize); + let stream = match db_guard.get_stream_mut(&effective_key) { + Ok(Some(s)) => s, + _ => { + responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); + return true; + } + }; + let claimed = match stream.read_group_new( + &group_name, + &consumer_name, + Some(request_count), + false, + ) { + Ok(entries) => entries, + Err(_) => { + responses.push(Frame::Array(vec![].into())); + return true; + } + }; + + let mut results = Vec::new(); + let mut dlq_entries: Vec<(StreamId, Vec<(Bytes, Bytes)>)> = Vec::new(); + let mut dlq_ack_ids: Vec = Vec::new(); + for (id, fields) in &claimed { + let delivery_count = stream + .groups + .get(group_name.as_ref()) + .and_then(|g| g.pel.get(id)) + .map(|pe| pe.delivery_count) + .unwrap_or(1); + if mdc > 0 && delivery_count >= mdc as u64 { + dlq_entries.push((*id, fields.clone())); + dlq_ack_ids.push(*id); + } else if results.len() < count { + results.push((*id, fields.clone())); + } + } + + if !dlq_ack_ids.is_empty() { + let _ = stream.xack(&group_name, &dlq_ack_ids); + } + + if !dlq_entries.is_empty() { + let dlq_key = { + let mut buf = Vec::with_capacity(effective_key.len() + 8); + buf.extend_from_slice(&effective_key); + buf.extend_from_slice(b"::mq:dlq"); + Bytes::from(buf) + }; + if let Ok(dlq_stream) = db_guard.get_or_create_stream(&dlq_key) { + for (_id, fields) in dlq_entries { + let dlq_id = dlq_stream.next_auto_id(); + dlq_stream.add(dlq_id, fields); + } + } + } + + let result_frames: Vec = results + .iter() + .map(|(id, fields)| { + let mut entry_frames = Vec::with_capacity(2); + entry_frames.push(Frame::BulkString(Bytes::from(format!( + "{}-{}", + id.ms, id.seq + )))); + let field_frames: Vec = fields + .iter() + .flat_map(|(f, v)| { + vec![Frame::BulkString(f.clone()), Frame::BulkString(v.clone())] + }) + .collect(); + entry_frames.push(Frame::Array(field_frames.into())); + Frame::Array(entry_frames.into()) + }) + .collect(); + responses.push(Frame::Array(result_frames.into())); + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"ACK") { + match validate_mq_ack(cmd_args) { + Ok((queue_key, msg_ids)) => { + let effective_key = + crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); + let ids: Vec = msg_ids + .iter() + .map(|(ms, seq)| StreamId { ms: *ms, seq: *seq }) + .collect(); + let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + match db_guard.get_stream_mut(&effective_key) { + Ok(Some(stream)) => { + let group_name = Bytes::from_static(b"__mq_consumers"); + match stream.xack(&group_name, &ids) { + Ok(acked_count) => { + drop(db_guard); + for (ms, seq) in &msg_ids { + let payload = + crate::mq::wal::encode_mq_ack(&effective_key, *ms, *seq); + let mut wal_buf = Vec::new(); + crate::persistence::wal_v3::record::write_wal_v3_record( + &mut wal_buf, + 0, + crate::persistence::wal_v3::record::WalRecordType::MqAck, + &payload, + ); + ctx.shard_databases + .wal_append(ctx.shard_id, Bytes::from(wal_buf)); + } + responses.push(Frame::Integer(acked_count as i64)); + } + Err(_) => responses.push(Frame::Integer(0)), + } + } + Ok(None) => responses.push(Frame::Integer(0)), + Err(_) => responses.push(Frame::Integer(0)), + } + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"DLQLEN") { + match validate_mq_dlqlen(cmd_args) { + Ok(queue_key) => { + let effective_key = + crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); + let dlq_key = { + let mut buf = Vec::with_capacity(effective_key.len() + 8); + buf.extend_from_slice(&effective_key); + buf.extend_from_slice(b"::mq:dlq"); + Bytes::from(buf) + }; + let mut db_guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + let len = match db_guard.get_stream_mut(&dlq_key) { + Ok(Some(stream)) => stream.length as i64, + _ => 0i64, + }; + responses.push(Frame::Integer(len)); + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"TRIGGER") { + match validate_mq_trigger(cmd_args) { + Ok((queue_key, callback_cmd, debounce_ms)) => { + let effective_key = + crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); + let trig_key = if let Some(ws_id) = conn.workspace_id.as_ref() { + let ws_hex = ws_id.as_hex(); + let mut k = Vec::with_capacity(ws_hex.len() + 1 + queue_key.len()); + k.extend_from_slice(ws_hex.as_bytes()); + k.push(b':'); + k.extend_from_slice(&queue_key); + Bytes::from(k) + } else { + queue_key.clone() + }; + let entry = crate::mq::TriggerEntry { + queue_key: effective_key, + callback_cmd, + debounce_ms, + last_fire_ms: 0, + pending_fire_ms: 0, + }; + { + let mut guard = ctx.shard_databases.trigger_registry(ctx.shard_id); + let reg = + guard.get_or_insert_with(|| Box::new(crate::mq::TriggerRegistry::new())); + reg.register(trig_key, entry); + } + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + Err(e) => responses.push(e), + } + return true; + } + + if sub.eq_ignore_ascii_case(b"PUBLISH") { + match validate_mq_publish(cmd_args) { + Ok((queue_key, fields)) => { + let effective_key = + crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); + if let Some(ref mut txn) = conn.active_cross_txn { + txn.record_mq(effective_key, fields); + responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); + } else { + responses.push(Frame::Error(Bytes::from_static( + b"ERR MQ PUBLISH requires an active transaction (use TXN BEGIN first)", + ))); + } + } + Err(e) => responses.push(e), + } + return true; + } + + // Unknown MQ subcommand + responses.push(Frame::Error(Bytes::from_static(ERR_MQ_UNKNOWN_SUB))); + true +} + +/// Handle MULTI/EXEC/DISCARD commands. Returns `true` if consumed. +pub(super) fn try_handle_multi_exec( + cmd: &[u8], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + // --- MULTI --- + if cmd.eq_ignore_ascii_case(b"MULTI") { + if conn.in_cross_txn() { + responses.push(Frame::Error(Bytes::from_static(ERR_MULTI_TXN_CONFLICT))); + } else if conn.in_multi { + responses.push(Frame::Error(Bytes::from_static( + b"ERR MULTI calls can not be nested", + ))); + } else { + conn.in_multi = true; + conn.command_queue.clear(); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + return true; + } + + // --- EXEC --- + if cmd.eq_ignore_ascii_case(b"EXEC") { + if !conn.in_multi { + responses.push(Frame::Error(Bytes::from_static(b"ERR EXEC without MULTI"))); + } else { + conn.in_multi = false; + let result = execute_transaction_sharded( + &ctx.shard_databases, + ctx.shard_id, + &conn.command_queue, + conn.selected_db, + &ctx.cached_clock, + ); + conn.command_queue.clear(); + responses.push(result); + } + return true; + } + + // --- DISCARD --- + if cmd.eq_ignore_ascii_case(b"DISCARD") { + if !conn.in_multi { + responses.push(Frame::Error(Bytes::from_static( + b"ERR DISCARD without MULTI", + ))); + } else { + conn.in_multi = false; + conn.command_queue.clear(); + responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); + } + return true; + } + + false +} + +/// Handle GRAPH.* graph commands. Returns `true` if consumed. +#[cfg(feature = "graph")] +pub(super) fn try_handle_graph_command( + cmd: &[u8], + cmd_args: &[Frame], + conn: &mut ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { + if cmd.len() <= 6 || !cmd[..6].eq_ignore_ascii_case(b"GRAPH.") { + return false; + } + let (response, wal_records, cypher_intents, cypher_undo_ops) = + if crate::command::graph::is_graph_write_cmd(cmd) + || (cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") + && crate::command::graph::is_cypher_write_query(cmd_args)) + { + let mut gs = ctx.shard_databases.graph_store_write(ctx.shard_id); + let (resp, cypher_intents, undo_ops) = if cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") { + crate::command::graph::graph_query_or_write(&mut gs, cmd_args) + } else { + ( + crate::command::graph::dispatch_graph_write(&mut gs, cmd, cmd_args), + Vec::new(), + Vec::new(), + ) + }; + let records = gs.drain_wal(); + (resp, records, cypher_intents, undo_ops) + } else { + let gs = ctx.shard_databases.graph_store_read(ctx.shard_id); + let resp = crate::command::graph::dispatch_graph_read(&gs, cmd, cmd_args); + (resp, Vec::new(), Vec::new(), Vec::new()) + }; + // Phase 166: record graph intent for TXN rollback. + if let Some(txn) = conn.active_cross_txn.as_mut() { + let is_node = cmd.eq_ignore_ascii_case(b"GRAPH.ADDNODE"); + let is_edge = cmd.eq_ignore_ascii_case(b"GRAPH.ADDEDGE"); + if is_node || is_edge { + if let Frame::Integer(id) = &response { + if let Some(gname) = cmd_args.first().and_then(|f| extract_bytes(f)) { + txn.record_graph(*id as u64, is_node, gname); + } + } + } + if !cypher_intents.is_empty() { + if let Some(gname) = cmd_args.first().and_then(|f| extract_bytes(f)) { + for intent in &cypher_intents { + txn.record_graph(intent.entity_id, intent.is_node, gname.clone()); + } + } + } + // Phase 174 FIX-01: push undo ops for SET/DELETE/MERGE rollback. + for undo_op in cypher_undo_ops { + txn.record_graph_undo(undo_op); + } + } + for record in wal_records { + ctx.shard_databases + .wal_append(ctx.shard_id, bytes::Bytes::from(record)); + } + let mut response = response; + if let Some(ws_id) = conn.workspace_id.as_ref() { + strip_workspace_prefix_from_response(ws_id, cmd, &mut response); + } + responses.push(response); + true +} diff --git a/src/server/conn/mod.rs b/src/server/conn/mod.rs index aa17fe66..0357cabd 100644 --- a/src/server/conn/mod.rs +++ b/src/server/conn/mod.rs @@ -32,6 +32,7 @@ pub(crate) use shared::{SharedDatabases, execute_transaction}; pub(crate) use shared::{ execute_transaction_sharded, extract_primary_key, handle_config, is_multi_key_command, }; +#[allow(unused_imports)] pub(crate) use util::{ apply_resp3_conversion, extract_bytes, extract_command, propagate_subscription, restore_migrated_state, unpropagate_subscription, diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index 1a0b1935..1c51e7c7 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -371,19 +371,31 @@ impl super::Shard { let mut snapshot_state: Option = None; let mut snapshot_reply_tx: Option>> = None; - // Per-shard WAL writer (created only when persistence is actually enabled). + // Per-shard WAL v2 writer — created only when persistence is enabled AND + // disk-offload is disabled. When disk-offload is on, WAL v3 supersedes v2 + // (stronger guarantees: per-record LSN + CRC32C), so v2 is skipped entirely + // to avoid double-write overhead. let appendonly_enabled = runtime_config.read().appendonly != "no"; let mut wal_writer: Option = match (&persistence_dir, appendonly_enabled) { - (Some(dir), true) => match WalWriter::new(shard_id, std::path::Path::new(dir)) { - Ok(w) => { - info!("Shard {}: WAL writer initialized", shard_id); - Some(w) - } - Err(e) => { - tracing::warn!("Shard {}: WAL init failed: {}", shard_id, e); - None + (Some(dir), true) if !server_config.disk_offload_enabled() => { + match WalWriter::new(shard_id, std::path::Path::new(dir)) { + Ok(w) => { + info!("Shard {}: WAL writer initialized", shard_id); + Some(w) + } + Err(e) => { + tracing::warn!("Shard {}: WAL init failed: {}", shard_id, e); + None + } } - }, + } + (Some(_), true) => { + info!( + "Shard {}: WAL v2 skipped (disk-offload active, WAL v3 supersedes)", + shard_id + ); + None + } (Some(_), false) => { info!("Shard {}: WAL skipped (appendonly=no)", shard_id); None diff --git a/src/shard/spsc_handler.rs b/src/shard/spsc_handler.rs index cc7ae4c1..35389e66 100644 --- a/src/shard/spsc_handler.rs +++ b/src/shard/spsc_handler.rs @@ -1846,16 +1846,15 @@ pub(crate) fn wal_append_and_fanout( repl_state: &Option>>, shard_id: usize, ) { - // 1a. WAL v2 append (disk durability, legacy path) - if let Some(w) = wal_writer { - w.append(data); - } - // 1b. WAL v3 append (disk-offload mode: per-record LSN, CRC32C) + // WAL v3 supersedes v2 — skip v2 append when v3 is active to avoid + // double-write overhead (2 write syscalls per SPSC drain batch). if let Some(w3) = wal_v3_writer { w3.append( crate::persistence::wal_v3::record::WalRecordType::Command, data, ); + } else if let Some(w) = wal_writer { + w.append(data); } // 2. Replication backlog (in-memory circular buffer for partial resync) if let Some(backlog) = repl_backlog { diff --git a/src/shard/timers.rs b/src/shard/timers.rs index e1030038..7e3e557d 100644 --- a/src/shard/timers.rs +++ b/src/shard/timers.rs @@ -140,7 +140,9 @@ pub(crate) fn fire_pending_mq_triggers( /// WAL v3 fsync on 1-second interval (mirrors v2 everysec pattern). /// -/// Calls `flush_sync()` which writes buffered data and fsyncs the segment file. +/// Flush any buffered WAL v3 data and fsync to stable storage. +/// +/// Called on the 1s timer. Writes remaining buffer contents then fsyncs. /// Only active when disk-offload is enabled and WalWriterV3 was successfully initialized. pub(crate) fn sync_wal_v3(wal_v3: &mut Option) { if let Some(wal) = wal_v3 { diff --git a/tests/pipeline_auto_index.rs b/tests/pipeline_auto_index.rs index 7374d37b..66d09aaf 100644 --- a/tests/pipeline_auto_index.rs +++ b/tests/pipeline_auto_index.rs @@ -21,7 +21,7 @@ //! cargo test --no-default-features \ //! --features runtime-tokio,jemalloc,text-index,graph \ //! --test pipeline_auto_index -#![cfg(feature = "runtime-tokio")] +#![cfg(all(feature = "runtime-tokio", feature = "text-index"))] use moon::config::ServerConfig; use moon::runtime::cancel::CancellationToken;