diff --git a/CHANGELOG.md b/CHANGELOG.md index c9be9e31f..db7b97d10 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,137 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 fresh connection on mid-startup connection resets (CI flake: per-shard SO_REUSEPORT listeners accept-then-reset during init). +### Fixed — OOM eviction bypass closure (PR #TBD) + +- **Cross-shard SPSC write legs now enforce `--maxmemory`.** `Execute`, + `MultiExecute`, `PipelineBatch` and their `*Slotted` variants + (`src/shard/spsc_handler.rs`) executed writes against the TARGET shard's + `Database` directly, bypassing the connection handlers' eviction gate + entirely — a scatter-gather write (e.g. a pipeline of individually-routed + `SET`s hashing to a remote shard) could grow that shard's memory past + `maxmemory` without limit. A new `spsc_eviction_gate` helper mirrors + `run_write_eviction_gate` (`handler_monoio/mod.rs`) exactly — same elastic + per-shard budget (GAP-1), same spill-vs-plain branching — threaded through + `drain_spsc_shared`/`handle_shard_message_shared` and gated by a + once-per-drain-cycle `evict_active` snapshot (perf parity with + `batch_eviction_active`: zero extra lock acquires when `maxmemory` is + unset). **Perf follow-up (same PR):** that snapshot, and the Lua bridge's + gate below, each used to answer "is `maxmemory` nonzero?" with either a + `runtime_config.read()` per drain cycle or a generation/`Cell` snapshot. + Both now read a single process-global `AtomicU64` + (`storage::eviction::{publish_maxmemory, maxmemory_is_set}`), published at + every production write site of `RuntimeConfig.maxmemory` (`CONFIG SET + maxmemory`, server startup) — one Relaxed load, no lock, no per-script + Cell bookkeeping. +- **Lua `redis.call`/`redis.pcall` writes now enforce `--maxmemory`.** + EVAL/EVALSHA carry no WRITE command flag, so the dispatch-level OOM check + never saw them, and the bridge (`src/scripting/bridge.rs`) never ran the + check before executing a write inside a script — a tight `redis.call('SET', + ...)` loop could write arbitrarily far past the cap. A new + `LuaEvictionCtx`, captured by value into the `redis.call`/`redis.pcall` + closures at VM-setup time (`setup_lua_vm`, once per shard — zero new + `unsafe`, zero per-call allocation), runs the same gate before a WRITE + `redis.call` executes; `redis.call` raises the OOM error as a Lua runtime + error (script aborts, matching Redis semantics), `redis.pcall` returns it + as an `{err = ...}` table. Read-only scripts are unaffected. +- **FCALL-internal `redis.call`/`redis.pcall` writes now enforce + `--maxmemory` too.** `FUNCTION LOAD` (`src/scripting/functions.rs`) + creates its own per-library sandboxed Lua VM, separate from the shard's + shared EVAL/EVALSHA VM — that per-library VM registered `redis.call`/ + `redis.pcall` with `LuaEvictionCtx::disabled()` unconditionally, so a + `FUNCTION` whose body wrote in a loop could grow memory past the cap + independent of the EVAL/EVALSHA fix above. `FunctionRegistry::new` now + takes a `LuaEvictionCtx`, stored on the registry and cloned into every + library's VM at `create_library` time; both production construction sites + (`handler_monoio/mod.rs`, `handler_sharded/mod.rs`) build a real ctx from + the same shard handles `setup_lua_vm` already uses for the shared VM + (`ConnectionContext::{shard_databases, runtime_config, shard_id, + spill_sender, spill_file_id, disk_offload_dir}`); only test-only callers + pass `LuaEvictionCtx::disabled()`. +- **Cross-shard `MOVE`/`COPY ... DB n` two-database intercept now runs on + EVERY `ShardMessage` SPSC arm, not just the plain `Execute` arm.** The + intercept (special-cased ahead of the generic single-db write path because + both commands need two `&mut Database` borrows at once) previously existed + only in `Execute`. Every other arm — including `PipelineBatchSlotted`, the + one real client traffic (`handler_monoio`/`handler_sharded`'s pipelined + remote-write path) actually uses — fell through to the generic single-db + `key_extra::copy`, which parses and silently ignores the `DB` clause: a + remote `COPY src dst DB n` performed a same-db copy instead (data + corruption, confirmed pre-fix via manual probe: 4/20 remote `COPY`s landed + in the wrong db), and a remote `MOVE` returned a loud-but-wrong "MOVE + requires handler-level dispatch (cross-db operation)" error instead of + moving the key. The intercept is now extracted into a shared helper + (`src/shard/spsc_two_db::try_two_db_intercept`) that `Execute`, + `MultiExecute`, `PipelineBatch`, `ExecuteSlotted`, `MultiExecuteSlotted` + and `PipelineBatchSlotted` all call verbatim; cross-db `COPY` still runs + `spsc_eviction_gate` against the destination db before `copy_core` (`COPY` + duplicates the value, so it needs the gate; same-shard `MOVE` is net-zero + and stays ungated, same rationale as `DEL`). +- **Found, NOT fixed (out of scope): cross-shard `COPY`'s destination key is + only reliably retrievable when it hash-tag-colocates with the source key.** + Moon shards purely by key hash, with no db-awareness. A `COPY src dst DB n` + executes on whichever shard owns `src` (that's the routing key for the + whole command) and writes `dst` into that SAME shard's db space — it does + not (and cannot, without a real cross-shard two-phase protocol) relocate + the write to whichever shard `dst`'s OWN hash would naturally pick. A later + plain `GET dst` routes by `hash(dst)`, so it only reliably finds the value + when `src` and `dst` share a hash tag (or `--shards 1`). `MOVE` is exempt + (the key name is unchanged, so `hash(key)` is identical before and after). + This is a pre-existing, inherent property of the per-key-hash sharding + model applied to `COPY` with a renamed destination — the plain `Execute` + arm this fix mirrors has the identical property — not a regression from + extending the intercept to every arm. Fixing it is a materially larger + change (routing the destination write to `dst`'s own natural shard) than + "apply the existing intercept to more arms" and is out of scope here. +- **Found, NOT fixed (out of scope): Moon's eviction gate is per-DATABASE, + not per-shard-aggregate-across-databases.** `try_evict_if_needed_budget` + (and every SPSC write arm, including the new cross-db `COPY` gate above) + measures `db.estimated_memory()` for the ONE database being written to, + not the sum across all 16 logical dbs a shard holds — consistent + throughout Moon's write path, but a deviation from real Redis's + instance-wide `maxmemory`. A cross-db `COPY` into a mostly-empty + destination db can undercount total shard pressure. Uniform, pre-existing + behavior; fixing it is a cross-cutting change to the eviction model, not a + `COPY`/`MOVE`-specific one — out of scope here. +- New wire-level regression suite `tests/spsc_two_db.rs` (3 cases, both + runtimes) exercises the one arm real client traffic uses + (`PipelineBatchSlotted`): pipelined cross-shard `COPY ... DB n` (RED + before the fix — `MOVE`'s loud wrong error and `COPY`'s silent same-db + copy — both confirmed via a git-stash re-run against the pre-fix source), + pipelined cross-shard `MOVE`, and a same-db `COPY` control (routes through + the coordinator's multi-key scatter path, unaffected by this fix, kept as + a negative control). The `COPY` case uses `{i}`-hash-tagged src/dst key + pairs so the destination key is independently `GET`-able post-fix — see + the "destination hash-tag" caveat above; an untagged version of the same + test only succeeds for the ~1-in-4 keys where `hash(dst)` collides with + `hash(src)` by chance. +- New wire-level regression suite `tests/oom_bypass_closure.rs` (8 cases, + both runtimes): direct-SET control (A), cross-shard pipeline (B), Lua EVAL + loop (C), read-only EVAL under pressure not blocked (D), cross-shard COPY + under pressure (E), `CONFIG SET maxmemory` atomic publish/un-publish (F), + FCALL-internal write loop (G), FCALL control with no maxmemory (H). + B, C, E and G RED before their respective fixes. D locks the write-only + scope of the Lua gate. Case F (added with the atomic-snapshot perf + follow-up) starts the server with `--maxmemory 0 --disk-offload disable` + (both required so the atomic starts genuinely unset — omitting + `--maxmemory` trips the auto-guardrail's nonzero default, and + disk-offload's spill-sender presence ORs into the same gate) and drives a + cross-shard pipeline before and after `CONFIG SET maxmemory`/`CONFIG SET + maxmemory 0`; RED before the CONFIG SET call site published the atomic + (775/3000 OOM — the local-only share) GREEN after (3000/3000). Case G + (FCALL) RED before the FunctionRegistry fix — a 2000-iteration + `redis.call('SET', ...)` loop inside an FCALL'd function completed with + `'DONE'` instead of an OOM error against a 2MB-capped `noeviction` server; + GREEN after. Case H locks the write-only/OOM-only scope of the FCALL gate, + mirroring case D for EVAL. Case E was revised for the cross-shard SPSC + intercept fix above: it now pre-loads db 1 (the `COPY` destination) with + the same ballast as db 0 before the OOM phase, since Moon's per-database + eviction gate (see the "found, not fixed" note above) needs the + destination db under its OWN pressure to trip — the original version + passed only as a side effect of the pre-fix same-db-copy bug (0/300 OOM + with the destination-db gate disabled and routing intact, confirmed via a + targeted stash of just that gate block; 76/300 fully fixed). + ### Added — int8 symmetric ADC for SQ8 vector search (task #13) - New per-candidate integer dot-product path for SQ8 asymmetric distance diff --git a/src/command/config.rs b/src/command/config.rs index a6d78e8ba..b08f223fc 100644 --- a/src/command/config.rs +++ b/src/command/config.rs @@ -102,6 +102,10 @@ pub fn config_set(runtime_config: &mut RuntimeConfig, args: &[Frame]) -> Frame { runtime_config.maxmemory = v; // Keep the inline write path's lock-free pre-gate in sync. crate::storage::eviction::publish_maxmemory_hints(&*runtime_config); + // Keep the SPSC-drain / Lua-bridge "is maxmemory set?" atomic + // in sync (Gap C) — a missed publish here silently bypasses + // the eviction gate on both paths. + crate::storage::eviction::publish_maxmemory(v as u64); } Err(_) => { return Frame::Error(Bytes::from(format!( diff --git a/src/config.rs b/src/config.rs index 1e3b12e64..f93e60e20 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1024,10 +1024,26 @@ fn detect_memory_limit_bytes() -> Option { } } -/// Non-Linux: no portable, dependency-free memory-limit probe. The guardrail -/// is skipped (operator sets `--maxmemory` explicitly on dev hosts). Production -/// targets Linux per the platform policy. -#[cfg(not(target_os = "linux"))] +/// macOS (first-class target): probe physical RAM via `sysctl -n hw.memsize`. +/// A spawned command instead of `sysctlbyname` FFI keeps this free of new +/// unsafe blocks; it runs once at startup so the fork cost is irrelevant. +/// No container limit concept applies on macOS, so host RAM is the limit. +#[cfg(target_os = "macos")] +fn detect_memory_limit_bytes() -> Option { + let out = std::process::Command::new("sysctl") + .args(["-n", "hw.memsize"]) + .output() + .ok()?; + if !out.status.success() { + return None; + } + std::str::from_utf8(&out.stdout).ok()?.trim().parse().ok() +} + +/// Other platforms: no portable, dependency-free memory-limit probe. The +/// guardrail is skipped (operator sets `--maxmemory` explicitly). Production +/// targets Linux/macOS per the platform policy. +#[cfg(not(any(target_os = "linux", target_os = "macos")))] fn detect_memory_limit_bytes() -> Option { None } @@ -1539,6 +1555,19 @@ mod tests { assert_eq!(config.to_runtime_config().maxmemory, 0); } + #[test] + #[cfg(target_os = "macos")] + fn macos_detects_memory_limit() { + // macOS is a first-class target; booting UNLIMITED + noeviction by + // default (guardrail Skipped) was RSS/CPU/OOM review item 4 — the + // hw.memsize probe must feed the same 80% guardrail as Linux. + let detected = detect_memory_limit_bytes(); + assert!( + detected.is_some_and(|b| b > 1 << 30), + "hw.memsize probe failed: {detected:?}" + ); + } + #[test] fn parse_meminfo_extracts_memtotal_bytes() { let sample = "MemTotal: 16384256 kB\nMemFree: 1000 kB\n"; diff --git a/src/main.rs b/src/main.rs index 12f8779a3..e4d9aa2b4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -873,6 +873,10 @@ fn main() -> anyhow::Result<()> { // per-shard hint divides by it) — the inline write path's eviction // pre-gate reads these instead of taking the runtime-config lock. moon::storage::eviction::publish_maxmemory_hints(&runtime_config_shared.read()); + // Publish the SPSC-drain / Lua-bridge "is maxmemory set?" atomic (Gap C). + // Missing this at startup would silently bypass the eviction gate for + // any server launched with --maxmemory (until the first CONFIG SET). + moon::storage::eviction::publish_maxmemory(runtime_config_shared.read().maxmemory as u64); moon::config::log_maxmemory_sharding(runtime_config_shared.read().maxmemory, num_shards); let server_config_shared: std::sync::Arc = { std::sync::Arc::new(config.clone()) }; diff --git a/src/scripting/bridge.rs b/src/scripting/bridge.rs index c8b0a2831..76088588c 100644 --- a/src/scripting/bridge.rs +++ b/src/scripting/bridge.rs @@ -7,11 +7,115 @@ //! 3. The pointer is cleared immediately after script execution use std::cell::Cell; +use std::path::PathBuf; +use std::rc::Rc; +use std::sync::Arc; use mlua::prelude::*; +use crate::config::RuntimeConfig; use crate::protocol::Frame; +use crate::shard::shared_databases::ShardDatabases; use crate::storage::engine::StorageEngine; +use crate::storage::eviction::{ + try_evict_if_needed_async_spill_budget, try_evict_if_needed_budget, +}; +use crate::storage::tiered::spill_thread::SpillRequest; + +/// Shard context needed to enforce `--maxmemory` eviction before a Lua +/// `redis.call`/`redis.pcall` WRITE actually mutates the database. +/// +/// # Why closure-capture instead of a thread-local +/// +/// The bridge already uses a thread-local raw pointer (`CURRENT_DB`) for the +/// `Database` itself, because that pointer's target changes on every script +/// invocation. This context is different: it is the same for every script run +/// on a given shard for the shard's entire lifetime (the shard's +/// `ShardDatabases`/`RuntimeConfig`/spill handles never change identity). +/// `redis.call`/`redis.pcall` are Lua closures created exactly once per shard +/// by [`crate::scripting::setup_lua_vm`], at a point where the caller already +/// owns cloneable handles to all of this — so it is captured directly into +/// the `move` closure. This adds zero new `unsafe` code (the existing +/// `CURRENT_DB` unsafe deref is untouched) and zero per-call allocation. The +/// common case (`maxmemory` unset, no spill) is decided by a single Relaxed +/// load of the process-global [`crate::storage::eviction::maxmemory_is_set`] +/// atomic, so a tight `redis.call('SET', ...)` loop never takes the +/// `RuntimeConfig` lock at all in that case. +#[derive(Clone)] +pub struct LuaEvictionCtx(Option); + +#[derive(Clone)] +struct LuaEvictionInner { + shard_databases: Arc, + runtime_config: Arc>, + shard_id: usize, + spill_sender: Option>, + spill_file_id: Rc>, + disk_offload_dir: Option, +} + +impl LuaEvictionCtx { + /// No-op gate. Used by unit tests (no real shard context available). + /// Production call sites (Lua EVAL/EVALSHA and Lua FUNCTION/FCALL) must + /// build a real ctx via [`LuaEvictionCtx::new`] — see + /// `src/shard/conn_accept.rs` and `src/scripting/functions.rs`. + pub fn disabled() -> Self { + LuaEvictionCtx(None) + } + + /// Real gate, built from the shard's own handles at VM-setup time. + #[allow(clippy::too_many_arguments)] + pub fn new( + shard_databases: Arc, + runtime_config: Arc>, + shard_id: usize, + spill_sender: Option>, + spill_file_id: Rc>, + disk_offload_dir: Option, + ) -> Self { + LuaEvictionCtx(Some(LuaEvictionInner { + shard_databases, + runtime_config, + shard_id, + spill_sender, + spill_file_id, + disk_offload_dir, + })) + } + + /// Run the same eviction/OOM gate the connection handlers use + /// (`run_write_eviction_gate` in `handler_monoio/mod.rs`), against `db` + /// (the shard's `Database`, already borrowed by the caller via the + /// `CURRENT_DB` thread-local). Returns the standard OOM `Frame::Error` + /// on failure; `Ok(())` if within budget, eviction succeeded, or + /// `maxmemory` is unset. + fn gate(&self, db: &mut crate::storage::Database, db_index: usize) -> Result<(), Frame> { + let Some(inner) = self.0.as_ref() else { + return Ok(()); + }; + // Lock-free fast path: a script issuing thousands of writes checks + // the process-global atomic (Gap C), not the RuntimeConfig lock. + if inner.spill_sender.is_none() && !crate::storage::eviction::maxmemory_is_set() { + return Ok(()); + } + let rt = inner.runtime_config.read(); + let budget = inner.shard_databases.elastic_budget(inner.shard_id); + if let Some(sender) = &inner.spill_sender { + let mut fid = inner.spill_file_id.get(); + let dir = inner + .disk_offload_dir + .as_deref() + .unwrap_or(std::path::Path::new(".")); + let res = try_evict_if_needed_async_spill_budget( + db, &rt, sender, dir, &mut fid, db_index, budget, + ); + inner.spill_file_id.set(fid); + res + } else { + try_evict_if_needed_budget(db, &rt, budget) + } + } +} thread_local! { /// Raw pointer to the current shard's Database during script execution. @@ -59,7 +163,11 @@ pub fn script_had_write() -> bool { /// /// If `propagate_errors` is true (redis.call), Frame::Error results are raised as Lua errors. /// If false (redis.pcall), errors are returned as {err = "..."} tables. -pub fn make_redis_call_fn(lua: &Lua, propagate_errors: bool) -> mlua::Result { +pub fn make_redis_call_fn( + lua: &Lua, + propagate_errors: bool, + eviction_ctx: LuaEvictionCtx, +) -> mlua::Result { lua.create_function(move |lua, args: LuaMultiValue| { // Convert all Lua arguments to Frames let frames: Vec = args @@ -110,6 +218,14 @@ pub fn make_redis_call_fn(lua: &Lua, propagate_errors: bool) -> mlua::Result, /// Reverse index: function_name -> library_name (for fast FCALL lookup). func_to_lib: HashMap, + /// OOM/eviction gate for `redis.call`/`redis.pcall` writes issued from + /// FCALL'd functions (Gap B). Each library gets its own sandboxed Lua VM + /// (`create_library` below), so this must be captured here and cloned + /// into every VM at creation time rather than wired once — mirrors how + /// `setup_lua_vm` captures the shard's `LuaEvictionCtx` for the shared + /// EVAL/EVALSHA VM. Production callers (`handler_monoio`/`handler_sharded`) + /// must pass a real ctx built from the shard's handles; only test-only + /// callers may pass `LuaEvictionCtx::disabled()`. + eviction_ctx: crate::scripting::bridge::LuaEvictionCtx, } impl FunctionRegistry { - pub fn new() -> Self { + pub fn new(eviction_ctx: crate::scripting::bridge::LuaEvictionCtx) -> Self { FunctionRegistry { libraries: HashMap::new(), func_to_lib: HashMap::new(), + eviction_ctx, } } @@ -283,7 +293,10 @@ impl FunctionRegistry { let lua = Rc::new(mlua::Lua::new()); crate::scripting::sandbox::setup_sandbox(&lua) .map_err(|e| LoadError::LuaError(e.to_string()))?; - crate::scripting::sandbox::register_redis_api(&lua) + // FCALL-internal writes run the same OOM/eviction gate as + // EVAL/EVALSHA (Gap B) — this per-library VM gets its own clone of + // the registry's `eviction_ctx`. + crate::scripting::sandbox::register_redis_api(&lua, self.eviction_ctx.clone()) .map_err(|e| LoadError::LuaError(e.to_string()))?; // Create a table to store registered functions @@ -526,7 +539,7 @@ mod tests { #[test] fn test_load_and_lookup() { - let mut reg = FunctionRegistry::new(); + let mut reg = FunctionRegistry::new(crate::scripting::bridge::LuaEvictionCtx::disabled()); let body = b"#!lua name=mylib\nredis.register_function('hello', function() return 'world' end)"; let name = reg.load(body, false).unwrap(); @@ -539,7 +552,7 @@ mod tests { #[test] fn test_load_duplicate_without_replace() { - let mut reg = FunctionRegistry::new(); + let mut reg = FunctionRegistry::new(crate::scripting::bridge::LuaEvictionCtx::disabled()); let body = b"#!lua name=mylib\nredis.register_function('hello', function() return 'world' end)"; reg.load(body, false).unwrap(); @@ -551,7 +564,7 @@ mod tests { #[test] fn test_load_replace() { - let mut reg = FunctionRegistry::new(); + let mut reg = FunctionRegistry::new(crate::scripting::bridge::LuaEvictionCtx::disabled()); let body1 = b"#!lua name=mylib\nredis.register_function('hello', function() return 'world' end)"; let body2 = @@ -563,7 +576,7 @@ mod tests { #[test] fn test_delete() { - let mut reg = FunctionRegistry::new(); + let mut reg = FunctionRegistry::new(crate::scripting::bridge::LuaEvictionCtx::disabled()); let body = b"#!lua name=mylib\nredis.register_function('hello', function() return 'world' end)"; reg.load(body, false).unwrap(); @@ -573,7 +586,7 @@ mod tests { #[test] fn test_flush() { - let mut reg = FunctionRegistry::new(); + let mut reg = FunctionRegistry::new(crate::scripting::bridge::LuaEvictionCtx::disabled()); let body = b"#!lua name=mylib\nredis.register_function('hello', function() return 'world' end)"; reg.load(body, false).unwrap(); @@ -583,7 +596,7 @@ mod tests { #[test] fn test_list() { - let mut reg = FunctionRegistry::new(); + let mut reg = FunctionRegistry::new(crate::scripting::bridge::LuaEvictionCtx::disabled()); let body = b"#!lua name=mylib\nredis.register_function('hello', function() return 'world' end)"; reg.load(body, false).unwrap(); @@ -594,7 +607,7 @@ mod tests { #[test] fn test_table_form_registration() { - let mut reg = FunctionRegistry::new(); + let mut reg = FunctionRegistry::new(crate::scripting::bridge::LuaEvictionCtx::disabled()); let body = b"#!lua name=mylib\nredis.register_function{function_name='hello', callback=function() return 'world' end, description='test func'}"; let name = reg.load(body, false).unwrap(); assert_eq!(name, Bytes::from_static(b"mylib")); @@ -604,7 +617,7 @@ mod tests { #[test] fn test_call_function() { - let mut reg = FunctionRegistry::new(); + let mut reg = FunctionRegistry::new(crate::scripting::bridge::LuaEvictionCtx::disabled()); let body = b"#!lua name=mylib\nredis.register_function('hello', function() return 'world' end)"; reg.load(body, false).unwrap(); @@ -616,7 +629,7 @@ mod tests { #[test] fn test_call_function_not_found() { - let reg = FunctionRegistry::new(); + let reg = FunctionRegistry::new(crate::scripting::bridge::LuaEvictionCtx::disabled()); let mut db = Database::new(); let result = reg.call_function(b"nonexistent", vec![], vec![], &mut db, 0, 1, false); assert!(matches!(result, Frame::Error(_))); diff --git a/src/scripting/mod.rs b/src/scripting/mod.rs index d1111ee16..cffef3009 100644 --- a/src/scripting/mod.rs +++ b/src/scripting/mod.rs @@ -18,10 +18,15 @@ use crate::storage::Database; /// Create and return a fully sandboxed Lua 5.4 VM with redis.* API registered. /// Must be called on the shard thread (Lua is !Send). -pub fn setup_lua_vm() -> mlua::Result> { +/// +/// `eviction_ctx` is captured into the `redis.call`/`redis.pcall` closures for +/// the lifetime of this VM (M3 OOM-bypass fix) — pass +/// [`bridge::LuaEvictionCtx::disabled()`] when no shard context applies +/// (tests). +pub fn setup_lua_vm(eviction_ctx: bridge::LuaEvictionCtx) -> mlua::Result> { let lua = Rc::new(Lua::new()); sandbox::setup_sandbox(&lua)?; - sandbox::register_redis_api(&lua)?; + sandbox::register_redis_api(&lua, eviction_ctx)?; Ok(lua) } @@ -363,7 +368,7 @@ mod tests { #[test] fn test_setup_lua_vm() { - let lua = setup_lua_vm().unwrap(); + let lua = setup_lua_vm(bridge::LuaEvictionCtx::disabled()).unwrap(); // Should have redis table let redis: LuaValue = lua.globals().get("redis").unwrap(); assert!(matches!(redis, LuaValue::Table(_))); @@ -375,7 +380,7 @@ mod tests { #[test] fn test_run_script_simple() { - let lua = setup_lua_vm().unwrap(); + let lua = setup_lua_vm(bridge::LuaEvictionCtx::disabled()).unwrap(); let mut db = Database::new(); let result = run_script(&lua, b"return 42", vec![], vec![], &mut db, 0, 1); @@ -384,7 +389,7 @@ mod tests { #[test] fn test_run_script_keys_argv() { - let lua = setup_lua_vm().unwrap(); + let lua = setup_lua_vm(bridge::LuaEvictionCtx::disabled()).unwrap(); let mut db = Database::new(); let result = run_script( @@ -401,7 +406,7 @@ mod tests { #[test] fn test_run_script_with_redis_call() { - let lua = setup_lua_vm().unwrap(); + let lua = setup_lua_vm(bridge::LuaEvictionCtx::disabled()).unwrap(); let mut db = Database::new(); // SET and GET via redis.call @@ -419,7 +424,7 @@ mod tests { #[test] fn test_run_script_redis_pcall_catches_error() { - let lua = setup_lua_vm().unwrap(); + let lua = setup_lua_vm(bridge::LuaEvictionCtx::disabled()).unwrap(); let mut db = Database::new(); // pcall should catch errors as table @@ -438,7 +443,7 @@ mod tests { #[test] fn test_run_script_type_conversions() { - let lua = setup_lua_vm().unwrap(); + let lua = setup_lua_vm(bridge::LuaEvictionCtx::disabled()).unwrap(); let mut db = Database::new(); // Return string @@ -521,7 +526,7 @@ mod tests { #[test] fn test_handle_eval_basic() { - let lua = setup_lua_vm().unwrap(); + let lua = setup_lua_vm(bridge::LuaEvictionCtx::disabled()).unwrap(); let cache = Rc::new(RefCell::new(ScriptCache::new())); let mut db = Database::new(); @@ -536,7 +541,7 @@ mod tests { #[test] fn test_handle_evalsha_noscript() { - let lua = setup_lua_vm().unwrap(); + let lua = setup_lua_vm(bridge::LuaEvictionCtx::disabled()).unwrap(); let cache = Rc::new(RefCell::new(ScriptCache::new())); let mut db = Database::new(); @@ -556,7 +561,7 @@ mod tests { #[test] fn test_handle_evalsha_after_eval() { - let lua = setup_lua_vm().unwrap(); + let lua = setup_lua_vm(bridge::LuaEvictionCtx::disabled()).unwrap(); let cache = Rc::new(RefCell::new(ScriptCache::new())); let mut db = Database::new(); diff --git a/src/scripting/sandbox.rs b/src/scripting/sandbox.rs index c32be6e7d..975364728 100644 --- a/src/scripting/sandbox.rs +++ b/src/scripting/sandbox.rs @@ -64,19 +64,28 @@ pub fn remove_timeout_hook(lua: &Lua) { /// /// Registers: redis.call, redis.pcall, redis.log, redis.error_reply, /// redis.status_reply, redis.sha1hex, and LOG_* level constants. -pub fn register_redis_api(lua: &Lua) -> mlua::Result<()> { +/// +/// `eviction_ctx` is threaded into the `redis.call`/`redis.pcall` closures so +/// every WRITE invoked from a script runs the same `--maxmemory` gate the +/// connection handlers use (M3 fix). Pass +/// [`crate::scripting::bridge::LuaEvictionCtx::disabled()`] when no shard +/// context is available (unit tests, FCALL's pre-existing documented gap). +pub fn register_redis_api( + lua: &Lua, + eviction_ctx: crate::scripting::bridge::LuaEvictionCtx, +) -> mlua::Result<()> { let redis_table = lua.create_table()?; // redis.call -- propagate errors as Lua errors redis_table.set( "call", - crate::scripting::bridge::make_redis_call_fn(lua, true)?, + crate::scripting::bridge::make_redis_call_fn(lua, true, eviction_ctx.clone())?, )?; // redis.pcall -- catch errors as {err = string} table redis_table.set( "pcall", - crate::scripting::bridge::make_redis_call_fn(lua, false)?, + crate::scripting::bridge::make_redis_call_fn(lua, false, eviction_ctx)?, )?; // redis.log(level, message) diff --git a/src/server/conn/core.rs b/src/server/conn/core.rs index dc2e2365d..385dbb643 100644 --- a/src/server/conn/core.rs +++ b/src/server/conn/core.rs @@ -84,11 +84,11 @@ pub(crate) struct ConnectionContext { /// pub/sub `SUBSCRIBE` and by per-connection `AffinityTracker` migration /// decisions when a connection converges ≥10/16 ops on a remote shard. pub pubsub_affinity: Arc>, - #[allow(dead_code)] // Only used by monoio handler (tiered storage) + // Used by the monoio handler's tiered-storage eviction path AND by both + // handlers' `FunctionRegistry::new` construction site (Gap B) to build + // the real `LuaEvictionCtx` for FCALL-internal writes. pub spill_sender: Option>, - #[allow(dead_code)] // Only used by monoio handler (tiered storage) pub spill_file_id: Rc>, - #[allow(dead_code)] // Only used by monoio handler (tiered storage) pub disk_offload_dir: Option, } diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index dc85513fb..719657981 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -168,8 +168,20 @@ pub(crate) async fn handle_connection_sharded_monoio< } 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())); + // Functions API registry (per-connection, lazy init) — kept as local because Rc> is !Send. + // Real eviction ctx (Gap B): FCALL-internal `redis.call` writes must run + // the same OOM gate as EVAL/EVALSHA, same handles `LuaEvictionCtx::new` + // uses elsewhere on this shard (conn_accept.rs's `setup_lua_vm` call). + let func_registry = Rc::new(RefCell::new(crate::scripting::FunctionRegistry::new( + crate::scripting::bridge::LuaEvictionCtx::new( + ctx.shard_databases.clone(), + ctx.runtime_config.clone(), + ctx.shard_id, + ctx.spill_sender.clone(), + ctx.spill_file_id.clone(), + ctx.disk_offload_dir.clone(), + ), + ))); // 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. diff --git a/src/server/conn/handler_sharded/mod.rs b/src/server/conn/handler_sharded/mod.rs index 894b5089b..b36744988 100644 --- a/src/server/conn/handler_sharded/mod.rs +++ b/src/server/conn/handler_sharded/mod.rs @@ -276,9 +276,19 @@ pub(crate) async fn handle_connection_sharded_inner< } let _registry_guard = RegistryGuard(client_id); - // Functions API registry (per-shard, lazy init) — kept as local because Rc> is !Send + // Functions API registry (per-shard, lazy init) — kept as local because Rc> is !Send. + // Real eviction ctx (Gap B): FCALL-internal `redis.call` writes must run + // the same OOM gate as EVAL/EVALSHA, same handles `LuaEvictionCtx::new` + // uses elsewhere on this shard (conn_accept.rs's `setup_lua_vm` call). let func_registry = std::rc::Rc::new(std::cell::RefCell::new( - crate::scripting::FunctionRegistry::new(), + crate::scripting::FunctionRegistry::new(crate::scripting::bridge::LuaEvictionCtx::new( + ctx.shard_databases.clone(), + ctx.runtime_config.clone(), + ctx.shard_id, + ctx.spill_sender.clone(), + ctx.spill_file_id.clone(), + ctx.disk_offload_dir.clone(), + )), )); // Per-connection arena for batch processing temporaries. diff --git a/src/shard/conn_accept.rs b/src/shard/conn_accept.rs index 434968b4f..950291cd3 100644 --- a/src/shard/conn_accept.rs +++ b/src/shard/conn_accept.rs @@ -160,8 +160,21 @@ pub(crate) fn spawn_tokio_connection( let lua = { let mut lua_opt = lua_rc.borrow_mut(); if lua_opt.is_none() { - *lua_opt = - Some(crate::scripting::setup_lua_vm().expect("Lua VM initialization failed")); + // M3 fix: bake this shard's OOM eviction context into the + // redis.call/pcall closures at VM-setup time (once per shard, + // shared by every connection landing here — shard_id/spill + // handles/runtime_config are identical across connections). + let eviction_ctx = crate::scripting::bridge::LuaEvictionCtx::new( + sdbs.clone(), + runtime_config.clone(), + shard_id, + spill_sender.clone(), + spill_file_id.clone(), + disk_offload_dir.clone(), + ); + *lua_opt = Some( + crate::scripting::setup_lua_vm(eviction_ctx).expect("Lua VM initialization failed"), + ); } lua_opt.as_ref().unwrap().clone() }; @@ -364,8 +377,18 @@ pub(crate) fn spawn_migrated_tokio_connection( let lua = { let mut lua_opt = lua_rc.borrow_mut(); if lua_opt.is_none() { + // M3 fix: see spawn_tokio_connection for rationale. + let eviction_ctx = crate::scripting::bridge::LuaEvictionCtx::new( + sdbs.clone(), + runtime_config.clone(), + shard_id, + spill_sender.clone(), + spill_file_id.clone(), + disk_offload_dir.clone(), + ); *lua_opt = Some( - crate::scripting::setup_lua_vm().expect("Lua VM initialization failed"), + crate::scripting::setup_lua_vm(eviction_ctx) + .expect("Lua VM initialization failed"), ); } lua_opt.as_ref().unwrap().clone() @@ -513,8 +536,18 @@ pub(crate) fn spawn_monoio_connection( let lua = { let mut lua_opt = lua_rc.borrow_mut(); if lua_opt.is_none() { + // M3 fix: see spawn_tokio_connection for rationale. + let eviction_ctx = crate::scripting::bridge::LuaEvictionCtx::new( + sdbs.clone(), + runtime_config.clone(), + shard_id, + spill_tx.clone(), + spill_fid.clone(), + do_dir.clone(), + ); *lua_opt = Some( - crate::scripting::setup_lua_vm().expect("Lua VM initialization failed"), + crate::scripting::setup_lua_vm(eviction_ctx) + .expect("Lua VM initialization failed"), ); } lua_opt.as_ref().unwrap().clone() @@ -840,8 +873,18 @@ pub(crate) fn spawn_migrated_monoio_connection( let lua = { let mut lua_opt = lua_rc.borrow_mut(); if lua_opt.is_none() { + // M3 fix: see spawn_tokio_connection for rationale. + let eviction_ctx = crate::scripting::bridge::LuaEvictionCtx::new( + sdbs.clone(), + runtime_config.clone(), + shard_id, + spill_sender.clone(), + spill_file_id.clone(), + disk_offload_dir.clone(), + ); *lua_opt = Some( - crate::scripting::setup_lua_vm().expect("Lua VM initialization failed"), + crate::scripting::setup_lua_vm(eviction_ctx) + .expect("Lua VM initialization failed"), ); } lua_opt.as_ref().unwrap().clone() diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index 11cb3c8f7..4734c2a33 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -1210,6 +1210,10 @@ impl super::Shard { !appendonly_enabled || !cdc_registry.is_empty() } }, + &runtime_config, + spill_sender.as_ref(), + &spill_file_id, + disk_offload_dir.as_deref(), ); if hit_cap { // M3: capped drain may have left a tail — re-arm immediately @@ -1312,6 +1316,10 @@ impl super::Shard { !appendonly_enabled || !cdc_registry.is_empty() } }, + &runtime_config, + spill_sender.as_ref(), + &spill_file_id, + disk_offload_dir.as_deref(), ); if hit_cap { // M3: capped drain may have left a tail — re-arm immediately @@ -1890,6 +1898,10 @@ impl super::Shard { !appendonly_enabled || !cdc_registry.is_empty() } }, + &runtime_config, + spill_sender.as_ref(), + &spill_file_id, + disk_offload_dir.as_deref(), ); if hit_cap { // M3: the drain stopped at its per-cycle cap (or a snapshot diff --git a/src/shard/mod.rs b/src/shard/mod.rs index 38113144c..6904cda9a 100644 --- a/src/shard/mod.rs +++ b/src/shard/mod.rs @@ -23,6 +23,10 @@ pub mod segment_stall; pub mod shared_databases; pub mod slice; pub mod spsc_handler; +/// Shared MOVE/COPY-DB two-database intercept for every `ShardMessage` SPSC +/// arm (Gap A). Split out of `spsc_handler.rs` per the repo's file-size +/// convention rather than growing that file further. +pub(crate) mod spsc_two_db; pub mod timers; pub mod uring_handler; @@ -367,6 +371,10 @@ mod tests { ), None, // aof_pool — None in tests true, // wal_kv_log — legacy behavior in tests + &std::sync::Arc::new(parking_lot::RwLock::new(RuntimeConfig::default())), // M2 fix: no shard context needed (maxmemory unset) + None, + &Rc::new(std::cell::Cell::new(1u64)), + None, ); // Subscriber now receives pre-serialized RESP bytes @@ -430,6 +438,10 @@ mod tests { ), None, // aof_pool — None in tests true, // wal_kv_log — legacy behavior in tests + &std::sync::Arc::new(parking_lot::RwLock::new(RuntimeConfig::default())), // M2 fix: no shard context needed (maxmemory unset) + None, + &Rc::new(std::cell::Cell::new(1u64)), + None, ); } diff --git a/src/shard/spsc_handler.rs b/src/shard/spsc_handler.rs index 4f8a227ae..d96b7f1a7 100644 --- a/src/shard/spsc_handler.rs +++ b/src/shard/spsc_handler.rs @@ -3,7 +3,7 @@ //! Extracted from shard/mod.rs to reduce file size. These are synchronous //! functions called from the event loop's select! arms. -use std::cell::RefCell; +use std::cell::{Cell, RefCell}; use std::rc::Rc; use std::sync::Arc; @@ -14,6 +14,7 @@ use tracing::info; use crate::blocking::BlockingRegistry; use crate::command::metadata; use crate::command::{DispatchResult, dispatch as cmd_dispatch}; +use crate::config::RuntimeConfig; use crate::persistence::aof; use crate::persistence::snapshot::SnapshotState; use crate::persistence::wal::WalWriter; @@ -23,6 +24,10 @@ use crate::replication::backlog::ReplicationBacklog; use crate::runtime::channel; use crate::storage::Database; use crate::storage::entry::CachedClock; +use crate::storage::eviction::{ + try_evict_if_needed_async_spill_budget, try_evict_if_needed_budget, +}; +use crate::storage::tiered::spill_thread::SpillRequest; use crate::command::vector_search; use crate::vector::store::VectorStore; @@ -30,6 +35,42 @@ use crate::vector::store::VectorStore; use super::dispatch::ShardMessage; use super::shared_databases::ShardDatabases; +/// SPSC-side write-path OOM/eviction gate (M2 fix). +/// +/// Mirrors `run_write_eviction_gate` in `src/server/conn/handler_monoio/mod.rs` +/// exactly (same elastic-budget lookup, same spill-vs-plain branching, same +/// OOM `Frame::Error` on failure) — this file is runtime-agnostic (shared by +/// both `runtime-monoio` and `runtime-tokio`, so it cannot call that +/// `#[cfg(feature = "runtime-monoio")]`-gated function directly. Cross-shard +/// SPSC legs (`Execute`/`MultiExecute`/`PipelineBatch` + their `*Slotted` +/// variants) execute writes against the TARGET shard's `&mut Database` +/// directly, bypassing the connection handlers' write path entirely — without +/// this gate a scatter-gather write could grow a remote shard's memory past +/// `maxmemory` without limit. +pub(super) fn spsc_eviction_gate( + db: &mut Database, + db_idx: usize, + shard_databases: &Arc, + shard_id: usize, + runtime_config: &Arc>, + spill_sender: Option<&flume::Sender>, + spill_file_id: &Rc>, + disk_offload_dir: Option<&std::path::Path>, +) -> Result<(), crate::protocol::Frame> { + let rt = runtime_config.read(); + let budget = shard_databases.elastic_budget(shard_id); + if let Some(sender) = spill_sender { + let mut fid = spill_file_id.get(); + let dir = disk_offload_dir.unwrap_or(std::path::Path::new(".")); + let res = + try_evict_if_needed_async_spill_budget(db, &rt, sender, dir, &mut fid, db_idx, budget); + spill_file_id.set(fid); + res + } else { + try_evict_if_needed_budget(db, &rt, budget) + } +} + /// Drain all SPSC consumer channels, processing cross-shard messages. /// /// SnapshotBegin messages are collected into `pending_snapshot` for deferred handling @@ -81,10 +122,26 @@ pub(crate) fn drain_spsc_shared( // drain cycle (`--wal-kv-log`; auto = false when the AOF is the recovery // authority and no CDC subscriber is attached — see wal_append_and_fanout). wal_kv_log: bool, + // M2 fix: OOM/eviction context for cross-shard write legs — mirrors what + // `ConnectionContext` carries for the local write path. + runtime_config: &Arc>, + spill_sender: Option<&flume::Sender>, + spill_file_id: &Rc>, + disk_offload_dir: Option<&std::path::Path>, ) -> bool { const MAX_DRAIN_PER_CYCLE: usize = 256; let mut drained = 0; + // Batch-level eviction gate (perf parity with handler_monoio's + // `batch_eviction_active`): snapshot "is maxmemory set?" once per drain + // cycle from the process-global atomic (Gap C) instead of taking + // `runtime_config.read()` per drain cycle. When neither maxmemory nor + // disk-offload is configured — the common non-memory-bound path — every + // write arm below skips the eviction call (and any lock acquire) + // entirely. `runtime_config` is still threaded through for the actual + // eviction pass (`spsc_eviction_gate`) when this is true. + let evict_active = spill_sender.is_some() || crate::storage::eviction::maxmemory_is_set(); + // Collect all messages first, then batch Execute/PipelineBatch under single borrow. // // Scratch buffers are thread-local (one shard per OS thread) so this @@ -202,6 +259,11 @@ pub(crate) fn drain_spsc_shared( autovacuum_daemon, aof_pool, // FIX-W1-2: thread AOF pool through SPSC drain wal_kv_log, + evict_active, + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, ); } } @@ -234,6 +296,11 @@ pub(crate) fn drain_spsc_shared( autovacuum_daemon, aof_pool, // FIX-W1-2: thread AOF pool through SPSC drain wal_kv_log, + evict_active, + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, ); } @@ -288,6 +355,14 @@ pub(crate) fn handle_shard_message_shared( // Whether KV command records should be logged to the per-shard WAL // (`--wal-kv-log`; see wal_append_and_fanout). wal_kv_log: bool, + // M2 fix: precomputed `maxmemory != 0 || disk-offload configured` (see + // `drain_spsc_shared`) — skips the eviction gate call (and its lock + // acquire) entirely on the common non-memory-bound path. + evict_active: bool, + runtime_config: &Arc>, + spill_sender: Option<&flume::Sender>, + spill_file_id: &Rc>, + disk_offload_dir: Option<&std::path::Path>, ) { match msg { ShardMessage::Execute { @@ -474,87 +549,28 @@ pub(crate) fn handle_shard_message_shared( // T2.2 MOVE — atomically moves a key between two dbs on the same shard. // T2.3 COPY DB n — copies a key to a different db on the same shard. // Both require two databases simultaneously; intercept before cmd_dispatch. - if cmd.eq_ignore_ascii_case(b"MOVE") { - use crate::command::keyspace::move_cmd as ksmv; - let mut response = match ksmv::parse_move_args(args, db_count) { - Err(e) => e, - Ok((_key, dst_db)) if dst_db == db_idx => { - crate::protocol::Frame::Integer(0) - } - Ok((key, dst_db)) => { - // SPSC runs single-threaded per shard; no concurrent MOVE can - // deadlock. slice path uses split_at_mut (no locking needed). - // Refresh expiry clock on BOTH databases before the move so - // an expired source key behaves as "not found" and an expired - // destination key doesn't shadow the insert (mirrors the - // single-DB write path at line 583). - crate::shard::slice::with_shard(|s| { - ksmv::with_two_slice_dbs( - &mut s.databases, - db_idx, - dst_db, - |src, dst| { - src.refresh_now_from_cache(cached_clock); - dst.refresh_now_from_cache(cached_clock); - ksmv::move_core(src, dst, &key) - }, - ) - }) - } - }; - if matches!(response, crate::protocol::Frame::Integer(1)) { - let serialized = aof::serialize_command(&command); - let mut aof_budget = crate::persistence::aof::AOF_SPSC_BACKPRESSURE_BOUND; - if !wal_append_and_fanout( - &serialized, - wal_writer, - wal_v3_writer, - repl_backlog, - replica_txs, - repl_state, + // Gap A: shared with every other ShardMessage arm via + // spsc_two_db::try_two_db_intercept — SPSC runs single-threaded + // per shard, so the slice path (split_at_mut, no locking) is safe. + if cmd.eq_ignore_ascii_case(b"MOVE") || cmd.eq_ignore_ascii_case(b"COPY") { + let intercepted = crate::shard::slice::with_shard(|s| { + crate::shard::spsc_two_db::try_two_db_intercept( + cmd, + args, + &mut s.databases, + db_idx, + db_count, + cached_clock, + evict_active, + shard_databases, shard_id, - aof_pool, // FIX-W1-2 - wal_kv_log, - &mut aof_budget, - ) { - response = crate::protocol::Frame::Error(bytes::Bytes::from_static( - AOF_APPEND_LOST_ERR, - )); - } - } - let _ = reply_tx.send(response); - return; - } - - if cmd.eq_ignore_ascii_case(b"COPY") { - use crate::command::keyspace::move_cmd as ksmv; - if let Some(copy_result) = ksmv::parse_copy_db_args(args, db_idx, db_count) { - let mut response = match copy_result { - Err(e) => e, - Ok(ca) => { - // Refresh expiry clock on BOTH dbs to mirror the - // single-DB write path: expired src/dst keys must - // resolve correctly before copy_core inspects them. - crate::shard::slice::with_shard(|s| { - ksmv::with_two_slice_dbs( - &mut s.databases, - db_idx, - ca.dst_db, - |src, dst| { - src.refresh_now_from_cache(cached_clock); - dst.refresh_now_from_cache(cached_clock); - ksmv::copy_core( - src, - dst, - &ca.src_key, - &ca.dst_key, - ca.replace, - ) - }, - ) - }) - } - }; + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, + ) + }); + if let Some(mut response) = intercepted { if matches!(response, crate::protocol::Frame::Integer(1)) { let serialized = aof::serialize_command(&command); let mut aof_budget = @@ -579,120 +595,151 @@ pub(crate) fn handle_shard_message_shared( let _ = reply_tx.send(response); return; } - // No DB clause or same-db: fall through to cmd_dispatch → key_extra::copy + // COPY with no DB clause or same-db: fall through to + // cmd_dispatch → key_extra::copy. (MOVE always returns + // Some(..) from the helper — parse_move_args either + // succeeds or produces an error frame — so this fall- + // through is COPY-only in practice.) } // COW intercept: capture old value before write if snapshot is active let is_write = metadata::is_write(cmd); + // M2 fix: OOM/eviction gate for the target shard, run BEFORE + // COW intercept/dispatch (same order as handler_monoio's + // write path) — a remote-shard write leg must not be able to + // grow memory past `maxmemory` just because it arrived via + // SPSC instead of a local connection. + let mut oom_frame: Option = None; // write_db and text_store (HSET auto-index) accessed in one with_shard // closure to avoid re-entrant borrow (multi-resource arm). let frame = { if is_write { crate::shard::slice::with_shard_db(db_idx, |db| { + if evict_active { + if let Err(oom) = spsc_eviction_gate( + db, + db_idx, + shard_databases, + shard_id, + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, + ) { + oom_frame = Some(oom); + return; + } + } cow_intercept(snapshot_state, db, db_idx, &command); }); } - crate::shard::slice::with_shard(|s| { - let db = &mut s.databases[db_idx]; - db.refresh_now_from_cache(cached_clock); - let mut selected = db_idx; - let result = cmd_dispatch(db, cmd, args, &mut selected, db_count); - let frame = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => f, - }; + if let Some(oom) = oom_frame.take() { + oom + } else { + crate::shard::slice::with_shard(|s| { + let db = &mut s.databases[db_idx]; + db.refresh_now_from_cache(cached_clock); + let mut selected = db_idx; + let result = cmd_dispatch(db, cmd, args, &mut selected, db_count); + let frame = match result { + DispatchResult::Response(f) => f, + DispatchResult::Quit(f) => f, + }; - // WAL append + replication fan-out for successful write commands - let mut aof_ok = true; - if is_write && !matches!(frame, crate::protocol::Frame::Error(_)) { - let serialized = aof::serialize_command(&command); - let mut aof_budget = - crate::persistence::aof::AOF_SPSC_BACKPRESSURE_BOUND; - aof_ok = wal_append_and_fanout( - &serialized, - wal_writer, - wal_v3_writer, - repl_backlog, - replica_txs, - repl_state, - shard_id, - aof_pool, // FIX-W1-2 - wal_kv_log, - &mut aof_budget, - ); - } + // WAL append + replication fan-out for successful write commands + let mut aof_ok = true; + if is_write && !matches!(frame, crate::protocol::Frame::Error(_)) { + let serialized = aof::serialize_command(&command); + let mut aof_budget = + crate::persistence::aof::AOF_SPSC_BACKPRESSURE_BOUND; + aof_ok = wal_append_and_fanout( + &serialized, + wal_writer, + wal_v3_writer, + repl_backlog, + replica_txs, + repl_state, + shard_id, + aof_pool, // FIX-W1-2 + wal_kv_log, + &mut aof_budget, + ); + } - // Post-dispatch wakeup hooks for producer commands (cross-shard blocking) - if !matches!(frame, crate::protocol::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") - || cmd.eq_ignore_ascii_case(b"XADD"); - if needs_wake { - let wake_key = if cmd.eq_ignore_ascii_case(b"LMOVE") { - args.get(1) - .and_then(|f| crate::server::connection::extract_bytes(f)) - } else { - args.first() - .and_then(|f| crate::server::connection::extract_bytes(f)) - }; - if let Some(key) = wake_key { - let mut reg = 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, db, db_idx, &key, - ); - } else if cmd.eq_ignore_ascii_case(b"ZADD") { - crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, db, db_idx, &key, - ); + // Post-dispatch wakeup hooks for producer commands (cross-shard blocking) + if !matches!(frame, crate::protocol::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") + || cmd.eq_ignore_ascii_case(b"XADD"); + if needs_wake { + let wake_key = if cmd.eq_ignore_ascii_case(b"LMOVE") { + args.get(1).and_then(|f| { + crate::server::connection::extract_bytes(f) + }) } else { - crate::blocking::wakeup::try_wake_stream_waiter( - &mut reg, db, db_idx, &key, - ); + args.first().and_then(|f| { + crate::server::connection::extract_bytes(f) + }) + }; + if let Some(key) = wake_key { + let mut reg = 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, db, db_idx, &key, + ); + } else if cmd.eq_ignore_ascii_case(b"ZADD") { + crate::blocking::wakeup::try_wake_zset_waiter( + &mut reg, db, db_idx, &key, + ); + } else { + crate::blocking::wakeup::try_wake_stream_waiter( + &mut reg, db, db_idx, &key, + ); + } } } } - } - // Auto-index: if HSET succeeded and key matches a vector index prefix, - // extract the vector field and append to mutable segment. - // vector_store and text_store accessed here (same with_shard closure). - if cmd.eq_ignore_ascii_case(b"HSET") - && !matches!(frame, crate::protocol::Frame::Error(_)) - { - if let Some(crate::protocol::Frame::BulkString(key_bytes)) = - args.first() + // Auto-index: if HSET succeeded and key matches a vector index prefix, + // extract the vector field and append to mutable segment. + // vector_store and text_store accessed here (same with_shard closure). + if cmd.eq_ignore_ascii_case(b"HSET") + && !matches!(frame, crate::protocol::Frame::Error(_)) { - // Plan 166-01: return value (index_name, key_hash) - // tuples will be consumed by Plan 166-02 to record - // VectorIntents on the active CrossStoreTxn. Discarded - // here because this path is not txn-aware yet. - let _ = auto_index_hset( - &mut s.vector_store, - &mut s.text_store, - key_bytes, - args, - 0, - ); + if let Some(crate::protocol::Frame::BulkString(key_bytes)) = + args.first() + { + // Plan 166-01: return value (index_name, key_hash) + // tuples will be consumed by Plan 166-02 to record + // VectorIntents on the active CrossStoreTxn. Discarded + // here because this path is not txn-aware yet. + let _ = auto_index_hset( + &mut s.vector_store, + &mut s.text_store, + key_bytes, + args, + 0, + ); + } } - } - // Fail-loud: the mutation is applied (wake/auto-index above - // ran on real state), but the client must not see success - // for a write whose AOF record was dropped. - if aof_ok { - frame - } else { - crate::protocol::Frame::Error(bytes::Bytes::from_static( - AOF_APPEND_LOST_ERR, - )) - } - }) + // Fail-loud: the mutation is applied (wake/auto-index above + // ran on real state), but the client must not see success + // for a write whose AOF record was dropped. + if aof_ok { + frame + } else { + crate::protocol::Frame::Error(bytes::Bytes::from_static( + AOF_APPEND_LOST_ERR, + )) + } + }) + } }; // Auto-delete is a vector_store-only operation; runs outside the gate. @@ -721,8 +768,8 @@ pub(crate) fn handle_shard_message_shared( let mut results = Vec::with_capacity(commands.len()); let db_count = shard_databases.db_count(); let db_idx = db_index.min(db_count.saturating_sub(1)); - crate::shard::slice::with_shard_db(db_idx, |guard| { - guard.refresh_now_from_cache(cached_clock); + crate::shard::slice::with_shard(|s| { + s.databases[db_idx].refresh_now_from_cache(cached_clock); // ONE backpressure budget for the whole batch: under sustained // AOF backpressure the shard thread stalls at most BOUND total, // not BOUND × batch-len (review finding, PR #211). @@ -738,8 +785,83 @@ pub(crate) fn handle_shard_message_shared( } }; + // Gap A: MOVE / COPY-DB two-db intercept, mirroring the + // plain Execute arm — needs two &mut Database borrows at + // once, so it must run before `guard` narrows to a single + // db below. + if cmd.eq_ignore_ascii_case(b"MOVE") || cmd.eq_ignore_ascii_case(b"COPY") { + if let Some(response) = crate::shard::spsc_two_db::try_two_db_intercept( + cmd, + args, + &mut s.databases, + db_idx, + db_count, + cached_clock, + evict_active, + shard_databases, + shard_id, + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, + ) { + let mut aof_ok = true; + if matches!(response, crate::protocol::Frame::Integer(1)) + && wal_fanout_has_work( + wal_writer, + wal_v3_writer, + replica_txs, + aof_pool, + wal_kv_log, + ) + { + let serialized = aof::serialize_command(cmd_frame); + aof_ok = wal_append_and_fanout( + &serialized, + wal_writer, + wal_v3_writer, + repl_backlog, + replica_txs, + repl_state, + shard_id, + aof_pool, // FIX-W1-2 + wal_kv_log, + &mut aof_budget, + ); + } + results.push(if aof_ok { + response + } else { + crate::protocol::Frame::Error(bytes::Bytes::from_static( + AOF_APPEND_LOST_ERR, + )) + }); + continue; + } + // COPY with no DB clause or same-db: fall through to + // the generic single-db write path below. + } + + let guard = &mut s.databases[db_idx]; let is_write = metadata::is_write(cmd); if is_write { + // M2 fix: same gate as the Execute arm, applied per + // command in this batch's shared `guard` borrow. + if evict_active { + if let Err(oom) = spsc_eviction_gate( + guard, + db_idx, + shard_databases, + shard_id, + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, + ) { + results.push(oom); + continue; + } + } cow_intercept(snapshot_state, guard, db_idx, cmd_frame); } @@ -834,8 +956,11 @@ pub(crate) fn handle_shard_message_shared( // write_db and text_store (HSET auto-index) accessed in one with_shard // closure to avoid re-entrant borrow (multi-resource arm). crate::shard::slice::with_shard(|s| { - let guard = &mut s.databases[db_idx]; - guard.refresh_now_from_cache(cached_clock); + // One-time refresh via a scoped temporary borrow — `guard` + // itself moves INSIDE the loop below (Gap A) so the MOVE/ + // COPY-DB branch can borrow `&mut s.databases` (both src and + // dst) for the same command. + s.databases[db_idx].refresh_now_from_cache(cached_clock); // ONE backpressure budget for the whole batch: under sustained // AOF backpressure the shard thread stalls at most BOUND total, // not BOUND × batch-len (review finding, PR #211). @@ -851,8 +976,84 @@ pub(crate) fn handle_shard_message_shared( } }; + // Gap A: MOVE / COPY-DB two-db intercept, mirroring the + // plain Execute arm — needs two &mut Database borrows at + // once. Returns BEFORE the auto-index/wake hooks below + // too, mirroring the Execute arm's pre-existing behavior + // (not new for this fix — see the Gap A commit body). + if cmd.eq_ignore_ascii_case(b"MOVE") || cmd.eq_ignore_ascii_case(b"COPY") { + if let Some(response) = crate::shard::spsc_two_db::try_two_db_intercept( + cmd, + args, + &mut s.databases, + db_idx, + db_count, + cached_clock, + evict_active, + shard_databases, + shard_id, + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, + ) { + let mut aof_ok = true; + if matches!(response, crate::protocol::Frame::Integer(1)) + && wal_fanout_has_work( + wal_writer, + wal_v3_writer, + replica_txs, + aof_pool, + wal_kv_log, + ) + { + let serialized = aof::serialize_command(cmd_frame); + aof_ok = wal_append_and_fanout( + &serialized, + wal_writer, + wal_v3_writer, + repl_backlog, + replica_txs, + repl_state, + shard_id, + aof_pool, // FIX-C4-FOLD + wal_kv_log, + &mut aof_budget, + ); + } + results.push(if aof_ok { + response + } else { + crate::protocol::Frame::Error(bytes::Bytes::from_static( + AOF_APPEND_LOST_ERR, + )) + }); + continue; + } + // COPY with no DB clause or same-db: fall through to + // the generic single-db write path below. + } + + let guard = &mut s.databases[db_idx]; let is_write = metadata::is_write(cmd); if is_write { + // M2 fix: same gate as the Execute arm, applied per + // command in this batch's shared `guard` borrow. + if evict_active { + if let Err(oom) = spsc_eviction_gate( + guard, + db_idx, + shard_databases, + shard_id, + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, + ) { + results.push(oom); + continue; + } + } cow_intercept(snapshot_state, guard, db_idx, cmd_frame); } @@ -1008,89 +1209,162 @@ pub(crate) fn handle_shard_message_shared( } }; + // Gap A: MOVE / COPY-DB two-db intercept, mirroring the plain + // Execute arm — needs two &mut Database borrows at once. + if cmd.eq_ignore_ascii_case(b"MOVE") || cmd.eq_ignore_ascii_case(b"COPY") { + let intercepted = crate::shard::slice::with_shard(|s| { + crate::shard::spsc_two_db::try_two_db_intercept( + cmd, + args, + &mut s.databases, + db_idx, + db_count, + cached_clock, + evict_active, + shard_databases, + shard_id, + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, + ) + }); + if let Some(mut response) = intercepted { + if matches!(response, crate::protocol::Frame::Integer(1)) { + let serialized = aof::serialize_command(&command); + let mut aof_budget = crate::persistence::aof::AOF_SPSC_BACKPRESSURE_BOUND; + if !wal_append_and_fanout( + &serialized, + wal_writer, + wal_v3_writer, + repl_backlog, + replica_txs, + repl_state, + shard_id, + aof_pool, // FIX-W1-2 + wal_kv_log, + &mut aof_budget, + ) { + response = crate::protocol::Frame::Error(bytes::Bytes::from_static( + AOF_APPEND_LOST_ERR, + )); + } + } + // Arc-owned slot: deref is safe, refcount keeps it alive. + let slot = &*response_slot.0; + slot.fill(vec![response]); + return; + } + // COPY with no DB clause or same-db: fall through to the + // generic single-db write path below. + } + { let is_write = metadata::is_write(cmd); + // M2 fix: see the Execute arm for rationale/ordering. + let mut oom_frame: Option = None; let frame = { if is_write { crate::shard::slice::with_shard_db(db_idx, |db| { + if evict_active { + if let Err(oom) = spsc_eviction_gate( + db, + db_idx, + shard_databases, + shard_id, + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, + ) { + oom_frame = Some(oom); + return; + } + } cow_intercept(snapshot_state, db, db_idx, &command); }); } - crate::shard::slice::with_shard(|s| { - let db = &mut s.databases[db_idx]; - db.refresh_now_from_cache(cached_clock); - let mut selected = db_idx; - let result = cmd_dispatch(db, cmd, args, &mut selected, db_count); - let frame = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => f, - }; + if let Some(oom) = oom_frame.take() { + oom + } else { + crate::shard::slice::with_shard(|s| { + let db = &mut s.databases[db_idx]; + db.refresh_now_from_cache(cached_clock); + let mut selected = db_idx; + let result = cmd_dispatch(db, cmd, args, &mut selected, db_count); + let frame = match result { + DispatchResult::Response(f) => f, + DispatchResult::Quit(f) => f, + }; - let mut aof_ok = true; - if is_write && !matches!(frame, crate::protocol::Frame::Error(_)) { - let serialized = aof::serialize_command(&command); - let mut aof_budget = - crate::persistence::aof::AOF_SPSC_BACKPRESSURE_BOUND; - aof_ok = wal_append_and_fanout( - &serialized, - wal_writer, - wal_v3_writer, - repl_backlog, - replica_txs, - repl_state, - shard_id, - aof_pool, // FIX-W1-2 - wal_kv_log, - &mut aof_budget, - ); - } + let mut aof_ok = true; + if is_write && !matches!(frame, crate::protocol::Frame::Error(_)) { + let serialized = aof::serialize_command(&command); + let mut aof_budget = + crate::persistence::aof::AOF_SPSC_BACKPRESSURE_BOUND; + aof_ok = wal_append_and_fanout( + &serialized, + wal_writer, + wal_v3_writer, + repl_backlog, + replica_txs, + repl_state, + shard_id, + aof_pool, // FIX-W1-2 + wal_kv_log, + &mut aof_budget, + ); + } - if !matches!(frame, crate::protocol::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") - || cmd.eq_ignore_ascii_case(b"XADD"); - if needs_wake { - let wake_key = if cmd.eq_ignore_ascii_case(b"LMOVE") { - args.get(1) - .and_then(|f| crate::server::connection::extract_bytes(f)) - } else { - args.first() - .and_then(|f| crate::server::connection::extract_bytes(f)) - }; - if let Some(key) = wake_key { - let mut reg = 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, db, db_idx, &key, - ); - } else if cmd.eq_ignore_ascii_case(b"ZADD") { - crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, db, db_idx, &key, - ); + if !matches!(frame, crate::protocol::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") + || cmd.eq_ignore_ascii_case(b"XADD"); + if needs_wake { + let wake_key = if cmd.eq_ignore_ascii_case(b"LMOVE") { + args.get(1).and_then(|f| { + crate::server::connection::extract_bytes(f) + }) } else { - crate::blocking::wakeup::try_wake_stream_waiter( - &mut reg, db, db_idx, &key, - ); + args.first().and_then(|f| { + crate::server::connection::extract_bytes(f) + }) + }; + if let Some(key) = wake_key { + let mut reg = 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, db, db_idx, &key, + ); + } else if cmd.eq_ignore_ascii_case(b"ZADD") { + crate::blocking::wakeup::try_wake_zset_waiter( + &mut reg, db, db_idx, &key, + ); + } else { + crate::blocking::wakeup::try_wake_stream_waiter( + &mut reg, db, db_idx, &key, + ); + } } } } - } - // Fail-loud: mutation applied, but the client must not - // see success for a write whose AOF record was dropped. - if aof_ok { - frame - } else { - crate::protocol::Frame::Error(bytes::Bytes::from_static( - AOF_APPEND_LOST_ERR, - )) - } - }) + // Fail-loud: mutation applied, but the client must not + // see success for a write whose AOF record was dropped. + if aof_ok { + frame + } else { + crate::protocol::Frame::Error(bytes::Bytes::from_static( + AOF_APPEND_LOST_ERR, + )) + } + }) + } }; // Arc-owned slot: deref is safe, refcount keeps it alive. let slot = &*response_slot.0; @@ -1105,8 +1379,8 @@ pub(crate) fn handle_shard_message_shared( let mut results = Vec::with_capacity(commands.len()); let db_count = shard_databases.db_count(); let db_idx = db_index.min(db_count.saturating_sub(1)); - crate::shard::slice::with_shard_db(db_idx, |guard| { - guard.refresh_now_from_cache(cached_clock); + crate::shard::slice::with_shard(|s| { + s.databases[db_idx].refresh_now_from_cache(cached_clock); // ONE backpressure budget for the whole batch: under sustained // AOF backpressure the shard thread stalls at most BOUND total, // not BOUND × batch-len (review finding, PR #211). @@ -1122,8 +1396,83 @@ pub(crate) fn handle_shard_message_shared( } }; + // Gap A: MOVE / COPY-DB two-db intercept, mirroring the + // plain Execute arm — needs two &mut Database borrows at + // once, so it must run before `guard` narrows to a single + // db below. + if cmd.eq_ignore_ascii_case(b"MOVE") || cmd.eq_ignore_ascii_case(b"COPY") { + if let Some(response) = crate::shard::spsc_two_db::try_two_db_intercept( + cmd, + args, + &mut s.databases, + db_idx, + db_count, + cached_clock, + evict_active, + shard_databases, + shard_id, + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, + ) { + let mut aof_ok = true; + if matches!(response, crate::protocol::Frame::Integer(1)) + && wal_fanout_has_work( + wal_writer, + wal_v3_writer, + replica_txs, + aof_pool, + wal_kv_log, + ) + { + let serialized = aof::serialize_command(cmd_frame); + aof_ok = wal_append_and_fanout( + &serialized, + wal_writer, + wal_v3_writer, + repl_backlog, + replica_txs, + repl_state, + shard_id, + aof_pool, // FIX-W1-2 + wal_kv_log, + &mut aof_budget, + ); + } + results.push(if aof_ok { + response + } else { + crate::protocol::Frame::Error(bytes::Bytes::from_static( + AOF_APPEND_LOST_ERR, + )) + }); + continue; + } + // COPY with no DB clause or same-db: fall through to + // the generic single-db write path below. + } + + let guard = &mut s.databases[db_idx]; let is_write = metadata::is_write(cmd); if is_write { + // M2 fix: same gate as the Execute arm, applied per + // command in this batch's shared `guard` borrow. + if evict_active { + if let Err(oom) = spsc_eviction_gate( + guard, + db_idx, + shard_databases, + shard_id, + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, + ) { + results.push(oom); + continue; + } + } cow_intercept(snapshot_state, guard, db_idx, cmd_frame); } @@ -1219,8 +1568,11 @@ pub(crate) fn handle_shard_message_shared( let db_idx = db_index.min(db_count.saturating_sub(1)); // write_db and text_store (HSET auto-index) in one with_shard closure. crate::shard::slice::with_shard(|s| { - let guard = &mut s.databases[db_idx]; - guard.refresh_now_from_cache(cached_clock); + // One-time refresh via a scoped temporary borrow — `guard` + // itself moves INSIDE the loop below (Gap A) so the MOVE/ + // COPY-DB branch can borrow `&mut s.databases` (both src and + // dst) for the same command. + s.databases[db_idx].refresh_now_from_cache(cached_clock); // ONE backpressure budget for the whole batch: under sustained // AOF backpressure the shard thread stalls at most BOUND total, // not BOUND × batch-len (review finding, PR #211). @@ -1236,8 +1588,84 @@ pub(crate) fn handle_shard_message_shared( } }; + // Gap A: MOVE / COPY-DB two-db intercept, mirroring the + // plain Execute arm — needs two &mut Database borrows at + // once. Returns BEFORE the auto-index/wake hooks below + // too, mirroring the Execute arm's pre-existing behavior + // (not new for this fix — see the Gap A commit body). + if cmd.eq_ignore_ascii_case(b"MOVE") || cmd.eq_ignore_ascii_case(b"COPY") { + if let Some(response) = crate::shard::spsc_two_db::try_two_db_intercept( + cmd, + args, + &mut s.databases, + db_idx, + db_count, + cached_clock, + evict_active, + shard_databases, + shard_id, + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, + ) { + let mut aof_ok = true; + if matches!(response, crate::protocol::Frame::Integer(1)) + && wal_fanout_has_work( + wal_writer, + wal_v3_writer, + replica_txs, + aof_pool, + wal_kv_log, + ) + { + let serialized = aof::serialize_command(cmd_frame); + aof_ok = wal_append_and_fanout( + &serialized, + wal_writer, + wal_v3_writer, + repl_backlog, + replica_txs, + repl_state, + shard_id, + aof_pool, // FIX-C4-FOLD + wal_kv_log, + &mut aof_budget, + ); + } + results.push(if aof_ok { + response + } else { + crate::protocol::Frame::Error(bytes::Bytes::from_static( + AOF_APPEND_LOST_ERR, + )) + }); + continue; + } + // COPY with no DB clause or same-db: fall through to + // the generic single-db write path below. + } + + let guard = &mut s.databases[db_idx]; let is_write = metadata::is_write(cmd); if is_write { + // M2 fix: same gate as the Execute arm, applied per + // command in this batch's shared `guard` borrow. + if evict_active { + if let Err(oom) = spsc_eviction_gate( + guard, + db_idx, + shard_databases, + shard_id, + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, + ) { + results.push(oom); + continue; + } + } cow_intercept(snapshot_state, guard, db_idx, cmd_frame); } @@ -3153,6 +3581,10 @@ mod drain_cap_tests { let mut cdc = Vec::new(); let mut manifest = None; let mut autovacuum = crate::shard::autovacuum::AutovacuumDaemon::new(Default::default()); + // M2 fix: no shard context needed for this drain-cap test — maxmemory + // unset (evict_active == false) is the fast, no-op path. + let rtcfg = Arc::new(parking_lot::RwLock::new(RuntimeConfig::default())); + let spill_fid = Rc::new(Cell::new(1u64)); // BlockCancel messages don't touch ShardSlice, so no init_shard needed. // First cycle: 300 queued > 256 cap -> drains exactly 256, reports tail. @@ -3180,6 +3612,10 @@ mod drain_cap_tests { &mut autovacuum, None, true, // wal_kv_log + &rtcfg, + None, + &spill_fid, + None, ); assert!( hit_cap, @@ -3211,6 +3647,10 @@ mod drain_cap_tests { &mut autovacuum, None, true, // wal_kv_log + &rtcfg, + None, + &spill_fid, + None, ); assert!( !hit_cap2, diff --git a/src/shard/spsc_two_db.rs b/src/shard/spsc_two_db.rs new file mode 100644 index 000000000..5e5ec3746 --- /dev/null +++ b/src/shard/spsc_two_db.rs @@ -0,0 +1,123 @@ +//! Shared MOVE / `COPY ... DB n` two-database intercept for every +//! `ShardMessage` SPSC arm (Gap A). +//! +//! `MOVE` and `COPY ... DB n` each need two `&mut Database` borrows at once +//! (source + destination) — something the generic per-db write path +//! (`cmd_dispatch` / `key_extra::copy`) cannot provide, since it only ever +//! sees one `&mut Database`. The plain `Execute` arm in `spsc_handler.rs` +//! special-cased both commands ahead of the generic path; this module +//! extracts that logic into one helper so every other `ShardMessage` arm +//! (`MultiExecute`, `PipelineBatch`, `ExecuteSlotted`, `MultiExecuteSlotted`, +//! `PipelineBatchSlotted`) reuses it verbatim instead of silently falling +//! through to the single-db path — which, before this fix, meant COPY with +//! a `DB` clause silently performed a same-db copy (wrong-db data +//! corruption) and MOVE returned a loud-but-wrong "cross-db not supported" +//! error on every arm except the plain `Execute` one. + +use std::cell::Cell; +use std::rc::Rc; +use std::sync::Arc; + +use crate::command::keyspace::move_cmd as ksmv; +use crate::config::RuntimeConfig; +use crate::protocol::Frame; +use crate::shard::shared_databases::ShardDatabases; +use crate::storage::Database; +use crate::storage::entry::CachedClock; +use crate::storage::tiered::spill_thread::SpillRequest; + +/// Attempt the MOVE/`COPY ... DB n` two-database intercept for one command. +/// +/// Returns `Some(response)` when `cmd` is `MOVE`, or `COPY` with a `DB` +/// clause targeting a database other than `db_idx` — the caller MUST treat +/// this as the command's full, final response: no generic dispatch, no COW +/// intercept, no auto-index hooks (mirrors the plain `Execute` arm's +/// pre-existing behavior, which returns before reaching any of those for +/// both commands — see the COW note in the Gap A commit body for why this +/// is intentional, not an oversight). +/// +/// Returns `None` for anything else — not MOVE/COPY, or a same-db `COPY` +/// with no `DB` clause — so the caller falls through to the generic +/// single-db write path (`cmd_dispatch` → `key_extra::copy` for same-db +/// COPY). +/// +/// Persistence (WAL/AOF) is deliberately NOT done here: the caller routes +/// the returned response through each arm's own per-command persistence +/// block, gated on `matches!(response, Frame::Integer(1))` — the same +/// condition the plain `Execute` arm already uses. This is STRICTER than +/// the batch arms' generic `!Error` persistence condition: a same-db +/// `MOVE`, or a `COPY` of a missing source key, returns `Integer(0)` and +/// must NOT persist. +#[allow(clippy::too_many_arguments)] +pub(crate) fn try_two_db_intercept( + cmd: &[u8], + args: &[Frame], + databases: &mut [Database], + db_idx: usize, + db_count: usize, + cached_clock: &CachedClock, + evict_active: bool, + shard_databases: &Arc, + shard_id: usize, + runtime_config: &Arc>, + spill_sender: Option<&flume::Sender>, + spill_file_id: &Rc>, + disk_offload_dir: Option<&std::path::Path>, +) -> Option { + if cmd.eq_ignore_ascii_case(b"MOVE") { + let response = match ksmv::parse_move_args(args, db_count) { + Err(e) => e, + Ok((_key, dst_db)) if dst_db == db_idx => Frame::Integer(0), + Ok((key, dst_db)) => { + // Refresh expiry clock on BOTH databases before the move so + // an expired source key behaves as "not found" and an + // expired destination key doesn't shadow the insert. + ksmv::with_two_slice_dbs(databases, db_idx, dst_db, |src, dst| { + src.refresh_now_from_cache(cached_clock); + dst.refresh_now_from_cache(cached_clock); + ksmv::move_core(src, dst, &key) + }) + } + }; + return Some(response); + } + + if cmd.eq_ignore_ascii_case(b"COPY") { + // `?` here returns `None` from THIS function (not just the match) — + // exactly the desired "no DB clause / same-db: fall through to + // cmd_dispatch" behavior `parse_copy_db_args` documents. + let copy_result = ksmv::parse_copy_db_args(args, db_idx, db_count)?; + let response = match copy_result { + Err(e) => e, + Ok(ca) => ksmv::with_two_slice_dbs(databases, db_idx, ca.dst_db, |src, dst| { + // Refresh expiry clock on BOTH dbs to mirror the single-db + // write path: expired src/dst keys must resolve correctly + // before copy_core inspects them. + src.refresh_now_from_cache(cached_clock); + dst.refresh_now_from_cache(cached_clock); + // Unlike MOVE (net-zero — the key leaves src as it lands in + // dst), cross-db COPY duplicates the value, so it must run + // the same eviction gate as any other write. Gate on the + // DESTINATION db (the one that grows) before copy_core. + if evict_active { + if let Err(oom) = crate::shard::spsc_handler::spsc_eviction_gate( + dst, + ca.dst_db, + shard_databases, + shard_id, + runtime_config, + spill_sender, + spill_file_id, + disk_offload_dir, + ) { + return oom; + } + } + ksmv::copy_core(src, dst, &ca.src_key, &ca.dst_key, ca.replace) + }), + }; + return Some(response); + } + + None +} diff --git a/src/storage/eviction.rs b/src/storage/eviction.rs index 72f0d0eb5..0e872edfd 100644 --- a/src/storage/eviction.rs +++ b/src/storage/eviction.rs @@ -33,6 +33,34 @@ static MAXMEMORY_HINT: std::sync::atomic::AtomicUsize = static MAXMEMORY_PER_SHARD_HINT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(usize::MAX); +/// Process-global "is `maxmemory` nonzero?" flag, in bytes (0 = unset / +/// unlimited). This answers a single boolean question — distinct from +/// [`MAXMEMORY_HINT`] above, which additionally encodes the exact budget for +/// the inline fast path's skip-eviction arithmetic. Two call sites used to +/// answer this same question with a lock read per drain cycle +/// (`spsc_handler.rs`'s `drain_spsc_shared`) or a generation/`Cell` snapshot +/// (`LuaEvictionCtx::gate`); both now read this atomic instead. +static MAXMEMORY_GLOBAL: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + +/// Publish the current `maxmemory` byte count. MUST be called at every +/// production write site of `RuntimeConfig.maxmemory`: server startup +/// (after `ServerConfig` resolves the guardrail-capped value into +/// `RuntimeConfig`) and `CONFIG SET maxmemory`. A missed publish is a silent +/// eviction-gate bypass — `tests/oom_bypass_closure.rs` exercises both call +/// sites. +#[inline] +pub fn publish_maxmemory(bytes: u64) { + MAXMEMORY_GLOBAL.store(bytes, std::sync::atomic::Ordering::Relaxed); +} + +/// `true` iff `maxmemory` is currently nonzero (i.e. a limit is configured). +/// Lock-free: a single Relaxed atomic load, safe to call on every SPSC drain +/// cycle or Lua `redis.call`/`redis.pcall` invocation. +#[inline] +pub fn maxmemory_is_set() -> bool { + MAXMEMORY_GLOBAL.load(std::sync::atomic::Ordering::Relaxed) != 0 +} + /// Publish the maxmemory hints. MUST be called wherever `maxmemory` (or the /// shard count it is divided by) changes: server startup after the resolved /// `num_shards` is written, and `CONFIG SET maxmemory`. A missed publish is @@ -877,6 +905,25 @@ mod tests { // compute_elastic_budget (GAP-1) // ----------------------------------------------------------------- + // ----------------------------------------------------------------- + // MAXMEMORY_GLOBAL atomic (Gap C) + // ----------------------------------------------------------------- + + #[test] + fn maxmemory_publish_and_is_set_roundtrip() { + // Unset (0) => maxmemory_is_set() is false. + publish_maxmemory(0); + assert!(!maxmemory_is_set()); + // Any nonzero byte count => true. + publish_maxmemory(1); + assert!(maxmemory_is_set()); + publish_maxmemory(64 * 1024 * 1024); + assert!(maxmemory_is_set()); + // Un-publish (CONFIG SET maxmemory 0) flips it back off. + publish_maxmemory(0); + assert!(!maxmemory_is_set()); + } + #[test] fn elastic_budget_balanced_load_keeps_base() { // Every shard at base — no surplus to borrow. diff --git a/tests/oom_bypass_closure.rs b/tests/oom_bypass_closure.rs new file mode 100644 index 000000000..d88edaee3 --- /dev/null +++ b/tests/oom_bypass_closure.rs @@ -0,0 +1,856 @@ +//! `--maxmemory` + eviction is enforced ONLY in the four connection handlers +//! (see `run_write_eviction_gate` in `src/server/conn/handler_monoio/mod.rs`). +//! Two write paths never ran that check, so memory could be driven past +//! `maxmemory` without limit: +//! +//! 1. Cross-shard SPSC write legs (`src/shard/spsc_handler.rs`): `Execute`, +//! `PipelineBatch`, `MultiExecute` (+ their `*Slotted` variants) executed +//! write commands against the TARGET shard's `&mut Database` directly, +//! bypassing the connection handler's eviction gate entirely. +//! 2. Lua `redis.call` writes (`src/scripting/bridge.rs`): EVAL/EVALSHA carry +//! no WRITE command flag (`src/command/metadata.rs`), so the dispatch-level +//! OOM check never saw them, and the bridge itself never ran the check +//! before executing a write inside a script. +//! +//! Wire-level on purpose: store-level tests cannot catch dispatch wiring — +//! see `tests/vector_del_unindex.rs` for the same rationale. +//! +//! Case B deliberately does NOT use MSET: `coordinate_mset`'s cross-shard +//! scatter path (`src/shard/coordinator.rs`) discards every remote leg's +//! reply and hardcodes `+OK` regardless of leg outcome — a separate, +//! pre-existing coordinator bug, independent of the SPSC eviction gate this +//! suite targets. A pipeline of individually-routed `SET` commands exercises +//! the exact same SPSC arms (`ExecuteSlotted`/`PipelineBatchSlotted`) while +//! faithfully propagating each remote leg's real reply back to the client, +//! cleanly isolating the SPSC-side bypass from that separate bug. +//! +//! Run alone with: +//! MOON_BIN=$PWD/target/release/moon cargo test --test oom_bypass_closure + +#![allow(clippy::unwrap_used)] + +use std::io::{BufReader, Read, Write}; +use std::net::{TcpStream, ToSocketAddrs}; +use std::process::{Child, Command}; +use std::time::{Duration, Instant}; + +// --------------------------------------------------------------------------- +// Binary resolution + server spawn (pattern: tests/vector_del_unindex.rs) +// --------------------------------------------------------------------------- + +fn find_moon_binary() -> std::path::PathBuf { + if let Ok(bin) = std::env::var("MOON_BIN") { + let p = std::path::PathBuf::from(bin); + if p.exists() { + return p; + } + } + let manifest = env!("CARGO_MANIFEST_DIR"); + let release = std::path::PathBuf::from(format!("{manifest}/target/release/moon")); + if release.exists() { + return release; + } + let debug = std::path::PathBuf::from(format!("{manifest}/target/debug/moon")); + if debug.exists() { + return debug; + } + panic!("No moon binary found. Build first or set MOON_BIN=/path/to/moon."); +} + +/// Ports below 20000 collide with other services in CI/dev; pick a free one +/// above that floor instead of a fixed low port. +fn free_port() -> u16 { + loop { + let l = std::net::TcpListener::bind("127.0.0.1:0").expect("bind :0"); + let p = l.local_addr().expect("local_addr").port(); + drop(l); + if p >= 20000 { + return p; + } + } +} + +/// `tempfile::tempdir()` defaults to `$TMPDIR`, which on macOS lives on the +/// root volume group — observed at ~95% full in dev environments, well past +/// Moon's 5%-free diskfull write-pause guard (`MOONERR diskfull`). Root the +/// test's scratch dirs under the repo's own volume instead, which has ample +/// headroom (see gotcha_vm_diskfull_shared_volume in project memory for the +/// VM-side analog of this trap). +fn test_tmpdir() -> tempfile::TempDir { + let base = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("target/oom-test-tmp"); + std::fs::create_dir_all(&base).expect("create oom-test-tmp base dir"); + tempfile::Builder::new() + .prefix("oom-bypass-") + .tempdir_in(&base) + .expect("tempdir_in target/oom-test-tmp") +} + +struct ServerGuard(Child); + +impl Drop for ServerGuard { + fn drop(&mut self) { + // kill() sends SIGKILL on all platforms via std::process::Child; + // belt-and-suspenders backstop against a leaked busy-poller (see + // gotcha_leaked_moon_busypoller_contaminates_xshard in project memory). + let _ = self.0.kill(); + let _ = self.0.wait(); + } +} + +/// Spawn moon with a tiny `--maxmemory` + `noeviction` policy so writes OOM +/// deterministically once the budget is exceeded. +fn spawn_moon_oom( + port: u16, + dir: &std::path::Path, + shards: u32, + maxmemory_bytes: u64, +) -> ServerGuard { + let child = Command::new(find_moon_binary()) + .args([ + "--port", + &port.to_string(), + "--dir", + &dir.to_string_lossy(), + "--shards", + &shards.to_string(), + "--appendonly", + "no", + "--maxmemory", + &maxmemory_bytes.to_string(), + "--maxmemory-policy", + "noeviction", + ]) + .stdout(std::fs::File::create(dir.join("moon.stdout.log")).expect("stdout log")) + .stderr(std::fs::File::create(dir.join("moon.stderr.log")).expect("stderr log")) + .spawn() + .expect("spawn moon"); + ServerGuard(child) +} + +// --------------------------------------------------------------------------- +// Minimal RESP client (binary-safe args, full-frame parser) +// --------------------------------------------------------------------------- + +#[derive(Debug, PartialEq)] +enum V { + Simple(String), + Err(String), + Int(i64), + Bulk(Vec), + Arr(Vec), + Null, +} + +impl V { + fn is_oom_error(&self) -> bool { + matches!(self, V::Err(msg) if msg.to_uppercase().contains("OOM")) + } +} + +struct Client { + reader: BufReader, + writer: TcpStream, +} + +impl Client { + fn connect(port: u16) -> Self { + let addr = format!("127.0.0.1:{port}") + .to_socket_addrs() + .unwrap() + .next() + .unwrap(); + let start = Instant::now(); + let stream = loop { + match TcpStream::connect_timeout(&addr, Duration::from_millis(200)) { + Ok(s) => break s, + Err(_) if start.elapsed() < Duration::from_secs(30) => { + std::thread::sleep(Duration::from_millis(50)); + } + Err(e) => panic!("server never accepted on port {port}: {e}"), + } + }; + stream + .set_read_timeout(Some(Duration::from_secs(30))) + .unwrap(); + let writer = stream.try_clone().unwrap(); + Client { + reader: BufReader::new(stream), + writer, + } + } + + fn encode(args: &[&[u8]]) -> Vec { + let mut out = format!("*{}\r\n", args.len()).into_bytes(); + for a in args { + out.extend_from_slice(format!("${}\r\n", a.len()).as_bytes()); + out.extend_from_slice(a); + out.extend_from_slice(b"\r\n"); + } + out + } + + fn read_line(&mut self) -> String { + let mut line = Vec::new(); + let mut b = [0u8; 1]; + loop { + self.reader.read_exact(&mut b).expect("read byte"); + if b[0] == b'\n' { + break; + } + if b[0] != b'\r' { + line.push(b[0]); + } + } + String::from_utf8_lossy(&line).into_owned() + } + + fn parse(&mut self) -> V { + let line = self.read_line(); + let (t, rest) = line.split_at(1); + match t { + "+" => V::Simple(rest.to_string()), + "-" => V::Err(rest.to_string()), + ":" => V::Int(rest.parse().expect("int")), + "$" => { + let n: i64 = rest.parse().expect("bulk len"); + if n < 0 { + return V::Null; + } + let mut buf = vec![0u8; n as usize + 2]; + self.reader.read_exact(&mut buf).expect("bulk body"); + buf.truncate(n as usize); + V::Bulk(buf) + } + "*" => { + let n: i64 = rest.parse().expect("arr len"); + if n < 0 { + return V::Null; + } + V::Arr((0..n).map(|_| self.parse()).collect()) + } + other => panic!("unexpected RESP type {other:?} (line {line:?})"), + } + } + + fn cmd(&mut self, args: &[&[u8]]) -> V { + self.writer.write_all(&Self::encode(args)).expect("send"); + self.parse() + } + + /// Send all commands in ONE write (a wire pipeline), then read all replies. + fn pipeline(&mut self, cmds: &[Vec>]) -> Vec { + let mut buf = Vec::new(); + for c in cmds { + let refs: Vec<&[u8]> = c.iter().map(|a| a.as_slice()).collect(); + buf.extend_from_slice(&Self::encode(&refs)); + } + self.writer.write_all(&buf).expect("send pipeline"); + cmds.iter().map(|_| self.parse()).collect() + } + + /// Fallible PING for the readiness probe: a connection accepted while the + /// server is still bringing up its per-shard SO_REUSEPORT listeners can be + /// RESET mid-read — that must retry with a fresh connection, not panic. + fn try_ping(&mut self) -> std::io::Result { + self.writer.write_all(b"*1\r\n$4\r\nPING\r\n")?; + let mut buf = [0u8; 7]; + self.reader.read_exact(&mut buf)?; + Ok(&buf == b"+PONG\r\n") + } +} + +fn wait_ready(port: u16) -> Client { + let start = Instant::now(); + loop { + let mut c = Client::connect(port); + // Any I/O error (or a non-PONG answer, which would desync the framing) + // drops this connection and probes again on a new one. + if let Ok(true) = c.try_ping() { + return c; + } + assert!( + start.elapsed() < Duration::from_secs(30), + "server never answered PING on port {port}" + ); + std::thread::sleep(Duration::from_millis(100)); + } +} + +fn blob(size: usize, fill: u8) -> Vec { + vec![fill; size] +} + +// --------------------------------------------------------------------------- +// Case A: direct SET control. Sanity check that the pre-existing local write +// path enforces maxmemory (proves the test's OOM plumbing/assertions work +// before we lean on them for the cross-shard/Lua cases below). +// --------------------------------------------------------------------------- + +#[test] +fn test_case_a_direct_set_control_oom() { + let dir = test_tmpdir(); + let port = free_port(); + const MAXMEMORY: u64 = 2 * 1024 * 1024; // 2MB + let _guard = spawn_moon_oom(port, dir.path(), 1, MAXMEMORY); + let mut c = wait_ready(port); + + let value = blob(64 * 1024, b'a'); // 64KB per key + const MAX_ITERS: usize = 200; // ~32 expected before OOM; generous margin. + + let mut oomed = false; + for i in 0..MAX_ITERS { + let key = format!("a:{i}"); + let r = c.cmd(&[b"SET", key.as_bytes(), &value]); + if r.is_oom_error() { + oomed = true; + break; + } + assert_eq!( + r, + V::Simple("OK".into()), + "SET a:{i} unexpected reply {r:?}" + ); + } + assert!( + oomed, + "control case: direct SET never hit OOM within {MAX_ITERS} iterations \ + at maxmemory={MAXMEMORY} — the handler-level eviction gate itself is \ + broken, which would invalidate cases B/C below." + ); +} + +// --------------------------------------------------------------------------- +// Case B: cross-shard SPSC write legs. One wire pipeline of many +// individually-keyed SET commands (no hash tags) at shards=4 — the majority +// route to a shard other than the one this connection is pinned to, via the +// SPSC ExecuteSlotted/PipelineBatchSlotted arms in spsc_handler.rs. +// +// Before the M2 fix, those arms executed writes against the target shard's +// Database with NO eviction check at all, so remote legs kept returning +OK +// far past the point where a same-shard write would already OOM. This +// asserts the OOM error rate over a large oversubscribed pipeline, which is +// robust to elastic-budget redistribution (GAP-1) without pinning an exact +// byte threshold. +// --------------------------------------------------------------------------- + +#[test] +fn test_case_b_cross_shard_pipeline_oom() { + let dir = test_tmpdir(); + let port = free_port(); + const MAXMEMORY: u64 = 2 * 1024 * 1024; // 2MB whole-instance cap + let _guard = spawn_moon_oom(port, dir.path(), 4, MAXMEMORY); + let mut c = wait_ready(port); + + const N: usize = 3000; + const VALUE_SIZE: usize = 4096; // 4KB — total attempted ~12MB vs 2MB cap. + let value = blob(VALUE_SIZE, b'b'); + + let cmds: Vec>> = (0..N) + .map(|i| { + vec![ + b"SET".to_vec(), + format!("b:{i}").into_bytes(), + value.clone(), + ] + }) + .collect(); + let replies = c.pipeline(&cmds); + assert_eq!(replies.len(), N); + + let oom_count = replies.iter().filter(|r| r.is_oom_error()).count(); + let ok_count = replies + .iter() + .filter(|r| matches!(r, V::Simple(s) if s == "OK")) + .count(); + let other_count = N - oom_count - ok_count; + + // With eviction enforced across ALL 4 shards, each shard accepts writes + // until its share of the 2MB cap is used, then OOMs the rest — but the + // exact OK count is timing-sensitive: the GAP-1 elastic budget + // redistributes per-shard caps on a 100ms snapshot tick, so early writes + // can land before budgets tighten (observed ~500-1100 OKs of 3000 across + // machines). Without the fix, only the ~1/4 of keys landing on THIS + // connection's own local shard can ever OOM (the pre-existing, + // already-correct local path) — the other ~3/4 (remote legs) return +OK + // unconditionally, bounding the pre-fix OOM count at roughly N/4 (~750). + // N/2 sits with wide margin above that ceiling and below the post-fix + // floor (~1900+ observed), so the assertion discriminates the bypass + // without being sensitive to eviction-tick timing. + assert!( + oom_count >= N / 2, + "cross-shard SPSC bypass: expected the large majority of {N} \ + oversubscribed writes to OOM once every shard's budget is \ + exhausted, got only {oom_count} OOM / {ok_count} OK / \ + {other_count} other — remote-shard legs are bypassing the \ + eviction gate" + ); +} + +// --------------------------------------------------------------------------- +// Case C: Lua redis.call writes. EVAL/EVALSHA carry no WRITE command flag, +// so the dispatch-level check never saw them; the bridge itself never ran +// the eviction check before executing a write inside a script. A tight loop +// of redis.call('SET', ...) must be denied with an OOM error once the +// script's writes exceed maxmemory, not silently write past the cap. +// --------------------------------------------------------------------------- + +#[test] +fn test_case_c_lua_redis_call_write_oom() { + let dir = test_tmpdir(); + let port = free_port(); + const MAXMEMORY: u64 = 2 * 1024 * 1024; // 2MB + let _guard = spawn_moon_oom(port, dir.path(), 1, MAXMEMORY); + let mut c = wait_ready(port); + + // 2000 * 8KB = ~16MB attempted vs a 2MB cap — comfortably oversubscribed. + let value = blob(8 * 1024, b'c'); + let script = b"for i=1,2000 do redis.call('SET', KEYS[1]..i, ARGV[1]) end return 'DONE'"; + + let r = c.cmd(&[b"EVAL", script, b"1", b"lk:", &value]); + + match &r { + V::Err(msg) => { + assert!( + msg.to_uppercase().contains("OOM"), + "EVAL errored but not with OOM: {msg:?}" + ); + } + other => panic!( + "Lua redis.call bypass: EVAL completed as {other:?} instead of \ + hitting OOM — writes inside a script are not gated by \ + --maxmemory" + ), + } +} + +// --------------------------------------------------------------------------- +// Case D: read-only Lua script under memory pressure must NOT be blocked. +// The M2 fix gates redis.call inside `if cmd_is_write { ... }` in +// scripting/bridge.rs — this locks that invariant against a future +// over-broad gate (e.g. one that checks eviction unconditionally for every +// redis.call regardless of the inner command's write flag). +// --------------------------------------------------------------------------- + +#[test] +fn test_case_d_readonly_eval_not_blocked_at_oom() { + let dir = test_tmpdir(); + let port = free_port(); + const MAXMEMORY: u64 = 2 * 1024 * 1024; // 2MB + let _guard = spawn_moon_oom(port, dir.path(), 1, MAXMEMORY); + let mut c = wait_ready(port); + + // Drive the instance past its cap first (reuses case A's control setup). + let value = blob(64 * 1024, b'd'); + let mut set_up_oom = false; + for i in 0..200 { + let key = format!("d:{i}"); + let r = c.cmd(&[b"SET", key.as_bytes(), &value]); + if r.is_oom_error() { + set_up_oom = true; + break; + } + } + assert!( + set_up_oom, + "setup: instance never reached OOM to test against" + ); + + // A key written before the cap was hit must still be readable via a + // read-only script — the eviction gate must not fire for reads. + let script = b"return redis.call('GET', KEYS[1])"; + let r = c.cmd(&[b"EVAL", script, b"1", b"d:0"]); + assert_eq!( + r, + V::Bulk(value), + "read-only EVAL was blocked (or returned the wrong value) while the \ + instance is over maxmemory — the eviction gate must only apply to \ + writes inside a script, not reads" + ); +} + +// --------------------------------------------------------------------------- +// Case E: cross-db `COPY ... DB n` under memory pressure. This locks in the +// destination-db eviction gate in `src/shard/spsc_two_db.rs`'s +// `try_two_db_intercept` (Gap A), which every `ShardMessage` SPSC arm now +// calls (previously only the plain `Execute` arm had a two-db intercept at +// all; the other arms — including `PipelineBatchSlotted`, the one real +// client traffic uses — silently ignored the `DB` clause and performed a +// same-db copy, a data-correctness bug, not an eviction one; see the CHANGELOG +// and the Gap A commit body). Moon's eviction gate is per-DATABASE (whichever +// db a write targets — `db.estimated_memory()` alone, not a shard-wide +// aggregate across all 16 dbs; see `try_evict_if_needed_budget`'s doc +// comment), consistently across every write path including this one — so +// db 1 (the COPY destination) needs its OWN pre-existing memory pressure to +// demonstrate the gate; ballast keys below give it that. +// +// Revision note: this case originally pre-loaded ONLY db 0, then asserted +// OOM on cross-db COPYs. That passed only because of the (now-fixed) Gap A +// bug: pre-fix, COPY's `DB 1` clause was silently ignored and the copy +// landed in db 0 (already loaded from phase 1), so the single-db gate +// tripped as a side effect of the routing bug, not because the destination +// db was actually under pressure. Post-fix, COPY correctly writes into db 1, +// which started empty, so the exact same setup produced 0/300 OOM — a +// regression alarm that traced back to a stale test premise, not a +// production bug (confirmed: temporarily stashing only the `if evict_active +// { spsc_eviction_gate(dst, ...) }` block inside `try_two_db_intercept`, +// while keeping the destination-db routing fix, reproduces RED here — see +// the Gap A commit body). Phase 1b below gives db 1 the same starting +// footprint as db 0 so the destination-db gate has something to trip on. +// --------------------------------------------------------------------------- + +#[test] +fn test_case_e_cross_db_copy_oom() { + let dir = test_tmpdir(); + let port = free_port(); + const MAXMEMORY: u64 = 2 * 1024 * 1024; // 2MB whole-instance cap + let _guard = spawn_moon_oom(port, dir.path(), 4, MAXMEMORY); + let mut c = wait_ready(port); + + const N: usize = 300; + const VALUE_SIZE: usize = 4096; // 4KB — N*VALUE_SIZE ~= 1.2MB, well under the 2MB cap. + let value = blob(VALUE_SIZE, b'e'); + + // Phase 1: SET N keys in db 0 (default), spread across all 4 shards. + // Comfortably under budget — every SET must succeed. + let set_cmds: Vec>> = (0..N) + .map(|i| { + vec![ + b"SET".to_vec(), + format!("e:{i}").into_bytes(), + value.clone(), + ] + }) + .collect(); + let set_replies = c.pipeline(&set_cmds); + let set_ok = set_replies + .iter() + .filter(|r| matches!(r, V::Simple(s) if s == "OK")) + .count(); + assert_eq!( + set_ok, N, + "setup: not all {N} initial SETs succeeded under budget (got {set_ok} OK) \ + — the phase-1 sizing assumption is wrong, adjust before trusting phase 2" + ); + + // Phase 1b: ballast — SET the SAME N keys (same names, so they hash to + // the SAME shards as their phase-1 counterparts) into db 1 too, so the + // COPY destination starts with the same per-shard memory footprint as + // the source db. Required because Moon's eviction gate is per-DATABASE + // (see the module doc above) — without this, phase 2's copies land in a + // near-empty db 1 and never cross budget, regardless of whether the + // destination-db gate is wired correctly. + assert_eq!( + c.cmd(&[b"SELECT", b"1"]), + V::Simple("OK".into()), + "SELECT 1 (ballast) failed" + ); + let ballast_replies = c.pipeline(&set_cmds); + let ballast_ok = ballast_replies + .iter() + .filter(|r| matches!(r, V::Simple(s) if s == "OK")) + .count(); + assert_eq!( + ballast_ok, N, + "ballast: not all {N} db-1 SETs succeeded under budget (got {ballast_ok} OK) \ + — the phase-1b sizing assumption is wrong, adjust before trusting phase 2" + ); + assert_eq!( + c.cmd(&[b"SELECT", b"0"]), + V::Simple("OK".into()), + "SELECT 0 (back for phase 2) failed" + ); + + // Phase 2: COPY each key to a DISTINCT dst key (src != dst — required, + // Redis rejects same-key COPY regardless of db) — doubles the value's + // footprint, pushing per-shard usage past the cap. + let copy_cmds: Vec>> = (0..N) + .map(|i| { + let src = format!("e:{i}").into_bytes(); + let dst = format!("e:{i}:c").into_bytes(); + vec![b"COPY".to_vec(), src, dst, b"DB".to_vec(), b"1".to_vec()] + }) + .collect(); + let copy_replies = c.pipeline(©_cmds); + assert_eq!(copy_replies.len(), N); + + let oom_count = copy_replies.iter().filter(|r| r.is_oom_error()).count(); + let ok_count = copy_replies + .iter() + .filter(|r| matches!(r, V::Int(1))) + .count(); + let other_count = N - oom_count - ok_count; + + // Threshold picked from observed data, same methodology as case B: + // gate-block-disabled (routing intact, see the Gap A commit body) run = + // 0/300 OOM; fully fixed = 76/300. N/10 (30) sits safely above the RED + // floor and safely below the GREEN observed count, robust to + // elastic-budget (GAP-1) variance across CI machines. + assert!( + oom_count >= N / 10, + "cross-shard COPY bypass: expected a substantial share of {N} COPYs \ + routed via cross-shard SPSC to OOM once their shard's budget is \ + exhausted, got only {oom_count} OOM / {ok_count} OK / \ + {other_count} other — COPY is not hitting the generic eviction \ + gate on the cross-shard write path" + ); +} + +// --------------------------------------------------------------------------- +// Case F: CONFIG SET maxmemory publishes/un-publishes the process-global +// atomic (Gap C — `crate::storage::eviction::{publish_maxmemory, +// maxmemory_is_set}`) that the cross-shard SPSC drain (`spsc_handler.rs`) +// and the Lua eviction bridge (`scripting/bridge.rs`) now read instead of a +// per-drain-cycle `RuntimeConfig` lock read / per-script generation Cell. +// Server starts WITHOUT --maxmemory (atomic unset at startup, exercising the +// startup-publish call site's absence-of-effect) so this test isolates the +// CONFIG SET write site specifically. Runs at shards=4 with the same +// cross-shard oversubscribed-pipeline shape as case B, so the OOM assertion +// specifically proves the SPSC arms observe the freshly-published atomic — +// not just the pre-existing local/inline fast path. +// --------------------------------------------------------------------------- + +fn spawn_moon_no_maxmemory(port: u16, dir: &std::path::Path, shards: u32) -> ServerGuard { + let child = Command::new(find_moon_binary()) + .args([ + "--port", + &port.to_string(), + "--dir", + &dir.to_string_lossy(), + "--shards", + &shards.to_string(), + "--appendonly", + "no", + // Disk offload defaults to `enable`, which gives every shard's + // SPSC drain a `spill_sender.is_some() == true` regardless of + // maxmemory — that ORs into `evict_active` and would mask + // whether the CONFIG SET call site actually published the + // atomic under test. Disable it so `evict_active` here is driven + // solely by `maxmemory_is_set()`. + "--disk-offload", + "disable", + // Explicit `0`, NOT omitted: omitting --maxmemory triggers the + // config auto-guardrail (config.rs ~1044-1101), which caps + // maxmemory at a nonzero fraction of detected system RAM — that + // nonzero value would make the STARTUP publish_maxmemory call + // (main.rs) publish `maxmemory_is_set() == true` immediately, + // making phase 2's OOM assertion pass regardless of whether the + // CONFIG SET call site publishes anything at all. `0` is the + // explicit, Redis-compatible "unlimited" escape hatch — it + // starts the atomic definitively unset. + "--maxmemory", + "0", + ]) + .stdout(std::fs::File::create(dir.join("moon.stdout.log")).expect("stdout log")) + .stderr(std::fs::File::create(dir.join("moon.stderr.log")).expect("stderr log")) + .spawn() + .expect("spawn moon"); + ServerGuard(child) +} + +#[test] +fn test_case_f_config_set_maxmemory_publishes_atomic() { + let dir = test_tmpdir(); + let port = free_port(); + let _guard = spawn_moon_no_maxmemory(port, dir.path(), 4); + let mut c = wait_ready(port); + + const N: usize = 3000; + const VALUE_SIZE: usize = 4096; // 4KB, same shape as case B. + let value = blob(VALUE_SIZE, b'f'); + + // Phase 1: with no maxmemory ever configured (atomic unset at startup — + // no --maxmemory flag was passed), a large cross-shard pipeline must all + // succeed. Establishes the baseline before enabling the cap. + let baseline_cmds: Vec>> = (0..N) + .map(|i| { + vec![ + b"SET".to_vec(), + format!("f0:{i}").into_bytes(), + value.clone(), + ] + }) + .collect(); + let baseline_replies = c.pipeline(&baseline_cmds); + let baseline_ok = baseline_replies + .iter() + .filter(|r| matches!(r, V::Simple(s) if s == "OK")) + .count(); + assert_eq!( + baseline_ok, N, + "setup: writes should all succeed before maxmemory is ever configured" + ); + + // Phase 2: CONFIG SET maxmemory-policy + a tiny maxmemory (no restart). + // If `publish_maxmemory` at the CONFIG SET call site + // (src/command/config.rs) is missing or wrong, `maxmemory_is_set()` + // stays false and the SPSC drain's per-cycle `evict_active` snapshot + // never turns on — remote legs would keep returning +OK regardless of + // the now-configured cap. + let r = c.cmd(&[b"CONFIG", b"SET", b"maxmemory-policy", b"noeviction"]); + assert_eq!( + r, + V::Simple("OK".into()), + "CONFIG SET maxmemory-policy failed: {r:?}" + ); + const MAXMEMORY: &[u8] = b"2097152"; // 2MB whole-instance cap, same as cases B/E. + let r = c.cmd(&[b"CONFIG", b"SET", b"maxmemory", MAXMEMORY]); + assert_eq!( + r, + V::Simple("OK".into()), + "CONFIG SET maxmemory failed: {r:?}" + ); + + let capped_cmds: Vec>> = (0..N) + .map(|i| { + vec![ + b"SET".to_vec(), + format!("f1:{i}").into_bytes(), + value.clone(), + ] + }) + .collect(); + let capped_replies = c.pipeline(&capped_cmds); + let oom_count = capped_replies.iter().filter(|r| r.is_oom_error()).count(); + let ok_count = capped_replies + .iter() + .filter(|r| matches!(r, V::Simple(s) if s == "OK")) + .count(); + assert!( + oom_count >= N / 2, + "CONFIG SET maxmemory publish: expected the large majority of {N} \ + oversubscribed cross-shard writes to OOM after CONFIG SET maxmemory \ + {MAXMEMORY:?}, got only {oom_count} OOM / {ok_count} OK — the \ + process-global maxmemory_is_set() atomic was not published at the \ + CONFIG SET call site (Gap C)" + ); + + // Phase 3: CONFIG SET maxmemory 0 (un-publish) — writes must succeed + // again, proving the atomic flips back off and doesn't latch "active" + // forever. + let r = c.cmd(&[b"CONFIG", b"SET", b"maxmemory", b"0"]); + assert_eq!( + r, + V::Simple("OK".into()), + "CONFIG SET maxmemory 0 failed: {r:?}" + ); + + let unpublish_cmds: Vec>> = (0..N) + .map(|i| { + vec![ + b"SET".to_vec(), + format!("f2:{i}").into_bytes(), + value.clone(), + ] + }) + .collect(); + let unpublish_replies = c.pipeline(&unpublish_cmds); + let unpublish_ok = unpublish_replies + .iter() + .filter(|r| matches!(r, V::Simple(s) if s == "OK")) + .count(); + assert_eq!( + unpublish_ok, N, + "CONFIG SET maxmemory 0 un-publish: expected all {N} writes to \ + succeed once maxmemory is cleared, got only {unpublish_ok} OK — \ + the maxmemory_is_set() atomic did not flip back off" + ); +} + +// --------------------------------------------------------------------------- +// Case G: FCALL-internal writes (Lua FUNCTION library). `FUNCTION LOAD` +// creates its own per-library sandboxed Lua VM (src/scripting/functions.rs), +// separate from the shard's shared EVAL/EVALSHA VM. That per-library VM +// used to register redis.call/pcall with `LuaEvictionCtx::disabled()` +// unconditionally, regardless of the EVAL/EVALSHA fix in case C — a +// FUNCTION whose body writes in a loop could grow memory past `maxmemory` +// without limit. +// --------------------------------------------------------------------------- + +#[test] +fn test_case_g_fcall_internal_write_oom() { + let dir = test_tmpdir(); + let port = free_port(); + const MAXMEMORY: u64 = 2 * 1024 * 1024; // 2MB + let _guard = spawn_moon_oom(port, dir.path(), 1, MAXMEMORY); + let mut c = wait_ready(port); + + // Registered functions are called with zero Lua args (`call_function` in + // src/scripting/functions.rs invokes `registered.call(())`) — like a + // plain EVAL body, they read `KEYS`/`ARGV` as globals, not parameters. + let lib_body: &[u8] = b"#!lua name=oomlib\n\ + local function write_loop()\n\ + local val = ARGV[1]\n\ + for i = 1, 2000 do\n\ + redis.call('SET', 'gk:' .. i, val)\n\ + end\n\ + return 'DONE'\n\ + end\n\ + redis.register_function('write_loop', write_loop)"; + + let load_reply = c.cmd(&[b"FUNCTION", b"LOAD", lib_body]); + assert_eq!( + load_reply, + V::Bulk(b"oomlib".to_vec()), + "FUNCTION LOAD failed: {load_reply:?}" + ); + + // 2000 * 8KB = ~16MB attempted vs a 2MB cap — comfortably oversubscribed + // (same shape as case C's EVAL loop). + let value = blob(8 * 1024, b'g'); + let r = c.cmd(&[b"FCALL", b"write_loop", b"0", &value]); + + match &r { + V::Err(msg) => { + assert!( + msg.to_uppercase().contains("OOM"), + "FCALL errored but not with OOM: {msg:?}" + ); + } + other => panic!( + "FCALL-internal write bypass: FCALL completed as {other:?} \ + instead of hitting OOM — writes inside a FUNCTION are not \ + gated by --maxmemory" + ), + } +} + +// --------------------------------------------------------------------------- +// Case H: control — FCALL with no maxmemory configured must succeed. Locks +// the write-only/OOM-only scope of the FCALL gate against a future +// over-broad check (mirrors case D for EVAL). +// --------------------------------------------------------------------------- + +#[test] +fn test_case_h_fcall_no_maxmemory_succeeds() { + let dir = test_tmpdir(); + let port = free_port(); + // Reuse spawn_moon_no_maxmemory (Gap C's case F helper) so this + // genuinely has no cap — spawn_moon_oom always sets one. + let _guard = spawn_moon_no_maxmemory(port, dir.path(), 1); + let mut c = wait_ready(port); + + let lib_body: &[u8] = b"#!lua name=oklib\n\ + local function write_loop()\n\ + local val = ARGV[1]\n\ + for i = 1, 50 do\n\ + redis.call('SET', 'hk:' .. i, val)\n\ + end\n\ + return 'DONE'\n\ + end\n\ + redis.register_function('write_loop', write_loop)"; + + let load_reply = c.cmd(&[b"FUNCTION", b"LOAD", lib_body]); + assert_eq!( + load_reply, + V::Bulk(b"oklib".to_vec()), + "FUNCTION LOAD failed: {load_reply:?}" + ); + + let value = blob(1024, b'h'); + let r = c.cmd(&[b"FCALL", b"write_loop", b"0", &value]); + assert_eq!( + r, + V::Bulk(b"DONE".to_vec()), + "FCALL without maxmemory configured should succeed, got {r:?}" + ); +} diff --git a/tests/spsc_two_db.rs b/tests/spsc_two_db.rs new file mode 100644 index 000000000..bf50683b5 --- /dev/null +++ b/tests/spsc_two_db.rs @@ -0,0 +1,520 @@ +//! Gap A: cross-shard `MOVE`/`COPY ... DB n` two-database intercept must run +//! on EVERY `ShardMessage` SPSC arm, not just the plain `Execute` arm. +//! +//! Before this fix, the two-database intercept +//! (`src/shard/spsc_handler.rs`'s `MOVE`/`COPY` special-casing, now +//! extracted into `src/shard/spsc_two_db::try_two_db_intercept`) existed +//! ONLY in the plain `Execute` arm. Client traffic's actual cross-shard +//! write path — both single-command and pipelined — funnels through +//! `PipelineBatchSlotted` (a batch-of-one for the single-command case), so a +//! remote `COPY src dst DB n` fell through to the generic single-db +//! `key_extra::copy`, which parses (and silently ignores) the `DB` clause +//! and performs a same-db copy instead (confirmed pre-fix via manual probe +//! in `tests/oom_bypass_closure.rs`'s case E: 4/20 remote `COPY`s landed in +//! the wrong db). `MOVE` returned a loud-but-wrong "cross-db not supported" +//! error on the same arm. +//! +//! `Execute`/`MultiExecute`/`ExecuteSlotted`/`MultiExecuteSlotted` are NOT +//! reachable from ordinary client traffic today (`Execute` is the +//! admin-console path; `MultiExecute`/`MultiExecuteSlotted` are the VLL +//! coordinator's multi-key scatter path for MGET/MSET/multi-DEL, which +//! never carries MOVE/COPY; `ExecuteSlotted`/`MultiExecuteSlotted` have no +//! live constructor at all). The two wire-level cases below exercise the +//! one arm real traffic actually uses (`PipelineBatchSlotted`), which calls +//! the shared `try_two_db_intercept` helper (`src/shard/spsc_two_db.rs`) +//! verbatim — the same helper every other arm calls — so this coverage +//! transitively validates the logic the dead/admin-only arms share, though +//! it does not independently exercise those arms' own call sites. No direct +//! unit test of `try_two_db_intercept` exists yet (a documented gap, not a +//! claim of coverage that isn't there). +//! +//! Run alone with: +//! MOON_BIN=$PWD/target/release/moon cargo test --test spsc_two_db + +#![allow(clippy::unwrap_used)] + +use std::io::{BufReader, Read, Write}; +use std::net::{TcpStream, ToSocketAddrs}; +use std::process::{Child, Command}; +use std::time::{Duration, Instant}; + +// --------------------------------------------------------------------------- +// Binary resolution + server spawn (pattern: tests/oom_bypass_closure.rs) +// --------------------------------------------------------------------------- + +fn find_moon_binary() -> std::path::PathBuf { + if let Ok(bin) = std::env::var("MOON_BIN") { + let p = std::path::PathBuf::from(bin); + if p.exists() { + return p; + } + } + let manifest = env!("CARGO_MANIFEST_DIR"); + let release = std::path::PathBuf::from(format!("{manifest}/target/release/moon")); + if release.exists() { + return release; + } + let debug = std::path::PathBuf::from(format!("{manifest}/target/debug/moon")); + if debug.exists() { + return debug; + } + panic!("No moon binary found. Build first or set MOON_BIN=/path/to/moon."); +} + +fn free_port() -> u16 { + loop { + let l = std::net::TcpListener::bind("127.0.0.1:0").expect("bind :0"); + let p = l.local_addr().expect("local_addr").port(); + drop(l); + if p >= 20000 { + return p; + } + } +} + +/// Rooted under the repo's own volume, not system `$TMPDIR` (diskfull-guard +/// trap — see gotcha_vm_diskfull_shared_volume in project memory). +fn test_tmpdir() -> tempfile::TempDir { + let base = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("target/spsc-2db-tmp"); + std::fs::create_dir_all(&base).expect("create spsc-2db-tmp base dir"); + tempfile::Builder::new() + .prefix("spsc-2db-") + .tempdir_in(&base) + .expect("tempdir_in target/spsc-2db-tmp") +} + +struct ServerGuard(Child); + +impl Drop for ServerGuard { + fn drop(&mut self) { + // Belt-and-suspenders backstop against a leaked busy-poller (see + // gotcha_leaked_moon_busypoller_contaminates_xshard in project memory). + let _ = self.0.kill(); + let _ = self.0.wait(); + } +} + +fn spawn_moon(port: u16, dir: &std::path::Path, shards: u32) -> ServerGuard { + let child = Command::new(find_moon_binary()) + .args([ + "--port", + &port.to_string(), + "--dir", + &dir.to_string_lossy(), + "--shards", + &shards.to_string(), + "--appendonly", + "no", + ]) + .stdout(std::fs::File::create(dir.join("moon.stdout.log")).expect("stdout log")) + .stderr(std::fs::File::create(dir.join("moon.stderr.log")).expect("stderr log")) + .spawn() + .expect("spawn moon"); + ServerGuard(child) +} + +// --------------------------------------------------------------------------- +// Minimal RESP client (binary-safe args, full-frame parser) +// --------------------------------------------------------------------------- + +#[derive(Debug, PartialEq)] +enum V { + Simple(String), + Err(String), + Int(i64), + Bulk(Vec), + Arr(Vec), + Null, +} + +struct Client { + reader: BufReader, + writer: TcpStream, +} + +impl Client { + fn connect(port: u16) -> Self { + let addr = format!("127.0.0.1:{port}") + .to_socket_addrs() + .unwrap() + .next() + .unwrap(); + let start = Instant::now(); + let stream = loop { + match TcpStream::connect_timeout(&addr, Duration::from_millis(200)) { + Ok(s) => break s, + Err(_) if start.elapsed() < Duration::from_secs(30) => { + std::thread::sleep(Duration::from_millis(50)); + } + Err(e) => panic!("server never accepted on port {port}: {e}"), + } + }; + stream + .set_read_timeout(Some(Duration::from_secs(30))) + .unwrap(); + let writer = stream.try_clone().unwrap(); + Client { + reader: BufReader::new(stream), + writer, + } + } + + fn encode(args: &[&[u8]]) -> Vec { + let mut out = format!("*{}\r\n", args.len()).into_bytes(); + for a in args { + out.extend_from_slice(format!("${}\r\n", a.len()).as_bytes()); + out.extend_from_slice(a); + out.extend_from_slice(b"\r\n"); + } + out + } + + fn read_line(&mut self) -> String { + let mut line = Vec::new(); + let mut b = [0u8; 1]; + loop { + self.reader.read_exact(&mut b).expect("read byte"); + if b[0] == b'\n' { + break; + } + if b[0] != b'\r' { + line.push(b[0]); + } + } + String::from_utf8_lossy(&line).into_owned() + } + + fn parse(&mut self) -> V { + let line = self.read_line(); + let (t, rest) = line.split_at(1); + match t { + "+" => V::Simple(rest.to_string()), + "-" => V::Err(rest.to_string()), + ":" => V::Int(rest.parse().expect("int")), + "$" => { + let n: i64 = rest.parse().expect("bulk len"); + if n < 0 { + return V::Null; + } + let mut buf = vec![0u8; n as usize + 2]; + self.reader.read_exact(&mut buf).expect("bulk body"); + buf.truncate(n as usize); + V::Bulk(buf) + } + "*" => { + let n: i64 = rest.parse().expect("arr len"); + if n < 0 { + return V::Null; + } + V::Arr((0..n).map(|_| self.parse()).collect()) + } + other => panic!("unexpected RESP type {other:?} (line {line:?})"), + } + } + + fn cmd(&mut self, args: &[&[u8]]) -> V { + self.writer.write_all(&Self::encode(args)).expect("send"); + self.parse() + } + + /// Send all commands in ONE write (a wire pipeline), then read all replies. + fn pipeline(&mut self, cmds: &[Vec>]) -> Vec { + let mut buf = Vec::new(); + for c in cmds { + let refs: Vec<&[u8]> = c.iter().map(|a| a.as_slice()).collect(); + buf.extend_from_slice(&Self::encode(&refs)); + } + self.writer.write_all(&buf).expect("send pipeline"); + cmds.iter().map(|_| self.parse()).collect() + } + + fn try_ping(&mut self) -> std::io::Result { + self.writer.write_all(b"*1\r\n$4\r\nPING\r\n")?; + let mut buf = [0u8; 7]; + self.reader.read_exact(&mut buf)?; + Ok(&buf == b"+PONG\r\n") + } +} + +fn wait_ready(port: u16) -> Client { + let start = Instant::now(); + loop { + let mut c = Client::connect(port); + if let Ok(true) = c.try_ping() { + return c; + } + assert!( + start.elapsed() < Duration::from_secs(30), + "server never answered PING on port {port}" + ); + std::thread::sleep(Duration::from_millis(100)); + } +} + +fn blob(size: usize, fill: u8) -> Vec { + vec![fill; size] +} + +// --------------------------------------------------------------------------- +// Case 1: pipelined cross-shard `COPY src dst DB 1`. THE RED CASE — before +// the fix, `PipelineBatchSlotted` (the arm real client traffic uses for +// remote writes) fell through to the generic single-db `key_extra::copy`, +// so `dst` landed in db 0 instead of db 1. +// --------------------------------------------------------------------------- + +#[test] +fn test_pipelined_cross_shard_copy_db_n() { + let dir = test_tmpdir(); + let port = free_port(); + let _guard = spawn_moon(port, dir.path(), 4); + let mut c = wait_ready(port); + + const N: usize = 20; + let values: Vec> = (0..N).map(|i| blob(32, i as u8)).collect(); + + // Keys use a `{i}` hash tag so src and dst hash to the SAME shard for a + // given pair `i`, while still spreading ~N pairs across all 4 shards + // (tag `i` varies, so different pairs land on different shards — this + // still routes through PipelineBatchSlotted for most pairs). + // + // This is NOT optional plumbing: COPY's destination key is a genuinely + // DIFFERENT string than the source key. Moon shards purely by key hash + // (no db-awareness), and the two-db intercept executes on whichever + // shard OWNS THE SOURCE KEY (that's the routing key for the whole + // command) — it does not (and structurally cannot, without a real + // cross-shard two-phase protocol, which is out of Gap A's scope; see + // the CHANGELOG "found, not fixed" note) relocate the destination write + // to whichever shard the destination key's OWN hash would naturally + // pick. So a later plain `GET dst` — which routes by hash(dst) — only + // reliably finds the value if src and dst are hash-tag-colocated. This + // is a pre-existing, inherent property of Moon's per-key-hash sharding + // model applied to MOVE/COPY (the plain `Execute` arm this fix mirrors + // has the identical property) — not a Gap A regression. Verified: the + // original (untagged) form of this test intermittently succeeds for + // ONLY the ~1-in-4 keys where hash(dst) happens to collide with + // hash(src) by chance, and passes reliably at `--shards 1` where the + // question is moot. + let src_key = |i: usize| format!("{{{i}}}:src"); + let dst_key = |i: usize| format!("{{{i}}}:copy"); + + // Setup: SET N keys in db 0. + let set_cmds: Vec>> = (0..N) + .map(|i| vec![b"SET".to_vec(), src_key(i).into_bytes(), values[i].clone()]) + .collect(); + let set_replies = c.pipeline(&set_cmds); + for (i, r) in set_replies.iter().enumerate() { + assert_eq!( + *r, + V::Simple("OK".into()), + "setup SET {} failed: {r:?}", + src_key(i) + ); + } + + // Pipelined COPY ... DB 1 for every key (>= 2 commands in one wire + // write, so this routes through PipelineBatch/PipelineBatchSlotted, not + // a lone ExecuteSlotted). + let copy_cmds: Vec>> = (0..N) + .map(|i| { + vec![ + b"COPY".to_vec(), + src_key(i).into_bytes(), + dst_key(i).into_bytes(), + b"DB".to_vec(), + b"1".to_vec(), + ] + }) + .collect(); + let copy_replies = c.pipeline(©_cmds); + for (i, r) in copy_replies.iter().enumerate() { + assert_eq!( + *r, + V::Int(1), + "COPY {} -> {} DB 1 did not report success: {r:?}", + src_key(i), + dst_key(i) + ); + } + + // Verify per key: the copy landed in db 1, NOT db 0. + for i in 0..N { + let dst = dst_key(i); + + assert_eq!( + c.cmd(&[b"SELECT", b"1"]), + V::Simple("OK".into()), + "SELECT 1 failed" + ); + let in_db1 = c.cmd(&[b"GET", dst.as_bytes()]); + assert_eq!( + in_db1, + V::Bulk(values[i].clone()), + "{dst} missing/wrong value in db 1 (target db) — the two-db \ + intercept did not run on the SPSC arm this pipeline used" + ); + + assert_eq!( + c.cmd(&[b"SELECT", b"0"]), + V::Simple("OK".into()), + "SELECT 0 failed" + ); + let in_db0 = c.cmd(&[b"GET", dst.as_bytes()]); + assert_eq!( + in_db0, + V::Null, + "{dst} leaked into db 0 — COPY ... DB 1 silently performed \ + a same-db copy instead of a cross-db one (the pre-fix bug)" + ); + + // Source key is untouched by COPY (still in db 0). + let src_still_present = c.cmd(&[b"GET", src_key(i).as_bytes()]); + assert_eq!( + src_still_present, + V::Bulk(values[i].clone()), + "COPY must not remove the source key {} from db 0", + src_key(i) + ); + } +} + +// --------------------------------------------------------------------------- +// Case 2: pipelined cross-shard `MOVE key 1`. Pre-fix, this arm returned the +// defensive "MOVE requires handler-level dispatch (cross-db operation)" +// error instead of moving the key (loud, but wrong). +// --------------------------------------------------------------------------- + +#[test] +fn test_pipelined_cross_shard_move() { + let dir = test_tmpdir(); + let port = free_port(); + let _guard = spawn_moon(port, dir.path(), 4); + let mut c = wait_ready(port); + + const N: usize = 20; + let values: Vec> = (0..N) + .map(|i| blob(32, (i as u8).wrapping_add(1))) + .collect(); + + let set_cmds: Vec>> = (0..N) + .map(|i| { + vec![ + b"SET".to_vec(), + format!("mv:{i}").into_bytes(), + values[i].clone(), + ] + }) + .collect(); + let set_replies = c.pipeline(&set_cmds); + for (i, r) in set_replies.iter().enumerate() { + assert_eq!(*r, V::Simple("OK".into()), "setup SET mv:{i} failed: {r:?}"); + } + + let move_cmds: Vec>> = (0..N) + .map(|i| { + vec![ + b"MOVE".to_vec(), + format!("mv:{i}").into_bytes(), + b"1".to_vec(), + ] + }) + .collect(); + let move_replies = c.pipeline(&move_cmds); + for (i, r) in move_replies.iter().enumerate() { + assert_eq!( + *r, + V::Int(1), + "MOVE mv:{i} 1 did not report success: {r:?} — pipelined \ + cross-shard MOVE is not intercepted on this SPSC arm" + ); + } + + for i in 0..N { + let key = format!("mv:{i}"); + + assert_eq!( + c.cmd(&[b"SELECT", b"0"]), + V::Simple("OK".into()), + "SELECT 0 failed" + ); + let in_db0 = c.cmd(&[b"GET", key.as_bytes()]); + assert_eq!( + in_db0, + V::Null, + "mv:{i} still present in db 0 after MOVE — key was not removed \ + from the source db" + ); + + assert_eq!( + c.cmd(&[b"SELECT", b"1"]), + V::Simple("OK".into()), + "SELECT 1 failed" + ); + let in_db1 = c.cmd(&[b"GET", key.as_bytes()]); + assert_eq!( + in_db1, + V::Bulk(values[i].clone()), + "mv:{i} missing/wrong value in db 1 after MOVE" + ); + } +} + +// --------------------------------------------------------------------------- +// Case 3: control — same-db `COPY` (no `DB` clause) still works pipelined +// cross-shard. Locks the "no DB clause / same-db falls through to the +// generic path" branch of `try_two_db_intercept` against a regression. +// --------------------------------------------------------------------------- + +#[test] +fn test_pipelined_same_db_copy_control() { + let dir = test_tmpdir(); + let port = free_port(); + let _guard = spawn_moon(port, dir.path(), 4); + let mut c = wait_ready(port); + + const N: usize = 20; + let values: Vec> = (0..N) + .map(|i| blob(32, (i as u8).wrapping_add(2))) + .collect(); + + let set_cmds: Vec>> = (0..N) + .map(|i| { + vec![ + b"SET".to_vec(), + format!("sd:{i}").into_bytes(), + values[i].clone(), + ] + }) + .collect(); + let set_replies = c.pipeline(&set_cmds); + for (i, r) in set_replies.iter().enumerate() { + assert_eq!(*r, V::Simple("OK".into()), "setup SET sd:{i} failed: {r:?}"); + } + + let copy_cmds: Vec>> = (0..N) + .map(|i| { + vec![ + b"COPY".to_vec(), + format!("sd:{i}").into_bytes(), + format!("sd:{i}:copy").into_bytes(), + ] + }) + .collect(); + let copy_replies = c.pipeline(©_cmds); + for (i, r) in copy_replies.iter().enumerate() { + assert_eq!( + *r, + V::Int(1), + "same-db COPY sd:{i} -> sd:{i}:copy did not report success: {r:?}" + ); + } + + for i in 0..N { + let dst = format!("sd:{i}:copy"); + let got = c.cmd(&[b"GET", dst.as_bytes()]); + assert_eq!( + got, + V::Bulk(values[i].clone()), + "same-db COPY control regressed: sd:{i}:copy missing/wrong value" + ); + } +}