From 94ccce78bfa812af97bb7d2bdf09329cd6fe68f2 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Mon, 6 Jul 2026 12:44:31 +0700 Subject: [PATCH 1/9] =?UTF-8?q?fix(config):=20macOS=20memory=20guardrail?= =?UTF-8?q?=20=E2=80=94=20probe=20hw.memsize?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit RSS/CPU/OOM review item 4: macOS (a first-class target) booted "UNLIMITED + noeviction" whenever --maxmemory was omitted because detect_memory_limit_bytes() was Linux-only — no OOM defense at all. The probe now reads `sysctl -n hw.memsize` (spawned command, not sysctlbyname FFI, so no new unsafe block; one fork at startup) and feeds the same 80% auto-cap + allkeys-lru guardrail as Linux. Red/green: config::tests::macos_detects_memory_limit. author: Tin Dang --- src/config.rs | 37 +++++++++++++++++++++++++++++++++---- 1 file changed, 33 insertions(+), 4 deletions(-) diff --git a/src/config.rs b/src/config.rs index 1e3b12e6..f93e60e2 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"; From d83bc5ec3dcf8c705817cbaadbf2cc4817ffcbb1 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Mon, 6 Jul 2026 14:17:40 +0700 Subject: [PATCH 2/9] fix(shard): close cross-shard SPSC and Lua redis.call OOM eviction bypasses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --maxmemory + eviction is enforced only in the four connection handlers (run_write_eviction_gate in handler_monoio/mod.rs and its equivalents in handler_sharded, handler_single, blocking.rs). Three 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, MultiExecute, PipelineBatch and their *Slotted variants executed write commands against the TARGET shard's Database directly, bypassing the connection handler's eviction gate entirely. A pipeline of individually-routed SET commands hashing to a remote shard kept returning +OK far past the point where a same-shard write would already OOM. 2. Lua redis.call/redis.pcall writes (src/scripting/bridge.rs): EVAL and 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 eviction check before executing a write inside a script. A tight redis.call('SET', ...) loop could write arbitrarily far past the cap. 3. Cross-db COPY ... DB n (src/shard/spsc_handler.rs): COPY is special-cased ahead of the generic write path (same as MOVE — both need two &mut Database borrows at once), so the generic gate never saw it either. COPY duplicates the value into the destination db, unlike same-shard MOVE which is net-zero (key leaves the source db as it lands in the destination, same rationale as DEL) and stays ungated by design. Found while closing (1) above, same file, same intercept family. Fix mirrors the existing handler-level gate exactly, in all three cases: - New spsc_eviction_gate() helper in spsc_handler.rs duplicates run_write_eviction_gate's elastic-budget lookup (GAP-1) and spill-vs-plain branching, since spsc_handler.rs is runtime-agnostic (shared by runtime-monoio and runtime-tokio) and cannot call the monoio-cfg-gated handler function directly. Threaded through drain_spsc_shared/handle_shard_message_shared and gated on a once-per-drain-cycle `evict_active` snapshot (perf parity with handler_monoio's batch_eviction_active: zero extra RwLock read when maxmemory is unset, the common case). The generic write-path call site and the COPY special-case call site both go through this same helper, the latter against the destination db (ca.dst_db) before copy_core. - New LuaEvictionCtx in scripting/bridge.rs is captured by value into the redis.call/redis.pcall closures at VM-setup time (setup_lua_vm, called once per shard from the 4 conn_accept.rs spawn sites) — zero new unsafe code (closure capture, not a new thread-local raw pointer) and zero per-call allocation (an Option check short-circuits when maxmemory is unset). redis.call raises the OOM error as a Lua runtime error (script aborts); redis.pcall returns it as an {err = ...} table. Read-only scripts are unaffected. FCALL-internal writes keep a documented, pre-existing gap (no shard context threaded through the function-library loader) — explicitly out of scope for this fix. New wire-level suite tests/oom_bypass_closure.rs (4 cases, both runtimes): direct-SET control (A, sanity — already enforced), cross-shard pipeline (B), Lua EVAL loop (C), read-only EVAL under memory pressure not blocked (D, locks the write-only scope of the Lua gate). B and C reproduced RED against the pre-fix binary (268/3000 OOM errors instead of the expected majority for B; EVAL completed as 'DONE' instead of erroring for C) and all 4 cases are GREEN after the fix on both runtime-monoio (default) and runtime-tokio. Case B deliberately avoids MSET: coordinate_mset's cross-shard scatter path (src/shard/coordinator.rs) discards every remote leg's reply and hardcodes +OK regardless of outcome — a separate, pre-existing coordinator bug independent of the SPSC eviction gate this fix targets, found while tracing the reply path for this test. Left unfixed (out of scope) and flagged for a follow-up. author: Tin Dang --- CHANGELOG.md | 42 +++ src/scripting/bridge.rs | 114 +++++++- src/scripting/functions.rs | 10 +- src/scripting/mod.rs | 27 +- src/scripting/sandbox.rs | 15 +- src/shard/conn_accept.rs | 53 +++- src/shard/event_loop.rs | 12 + src/shard/mod.rs | 8 + src/shard/spsc_handler.rs | 543 +++++++++++++++++++++++++----------- tests/oom_bypass_closure.rs | 466 +++++++++++++++++++++++++++++++ 10 files changed, 1108 insertions(+), 182 deletions(-) create mode 100644 tests/oom_bypass_closure.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index c9be9e31..b2697b62 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,48 @@ 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). +- **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 + writes (`src/scripting/functions.rs`) keep a documented, pre-existing gap + (no shard context threaded through the function-library loader) — tracked + as a follow-up, not closed by this fix. +- **Cross-db `COPY ... DB n` now enforces `--maxmemory`.** Same SPSC + intercept family as above (`src/shard/spsc_handler.rs`) special-cases + `MOVE`/`COPY DB` ahead of the generic write path because both need two + `&mut Database` borrows at once; `COPY` duplicates the value into the + destination db and was missed by the generic gate. Now runs + `spsc_eviction_gate` against the destination db before `copy_core`. + Same-shard `MOVE` is left ungated (net-zero — the key leaves the source db + as it lands in the destination, same rationale as `DEL`). +- New wire-level regression suite `tests/oom_bypass_closure.rs` (4 cases, + both runtimes): direct-SET control (A), cross-shard pipeline (B), Lua EVAL + loop (C), read-only EVAL under pressure not blocked (D) — B and C RED + before this fix, GREEN after; D locks the write-only scope of the Lua gate. + ### 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/scripting/bridge.rs b/src/scripting/bridge.rs index c8b0a283..84d48833 100644 --- a/src/scripting/bridge.rs +++ b/src/scripting/bridge.rs @@ -7,11 +7,111 @@ //! 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 +/// `Option` check on `maxmemory == 0` (the common case) short-circuits before +/// any lock or budget lookup. +#[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) and + /// by the FCALL path (`src/scripting/functions.rs`), which has a + /// pre-existing, documented gap for in-function write eviction — closing + /// it is out of scope for this fix (see `tmp/OOM-SHIELD-CONTEXT.md`). + 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(()); + }; + let rt = inner.runtime_config.read(); + if rt.maxmemory == 0 && inner.spill_sender.is_none() { + return Ok(()); + } + 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 +159,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 +214,14 @@ pub fn make_redis_call_fn(lua: &Lua, propagate_errors: bool) -> mlua::Result 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 c32be6e7..97536472 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/shard/conn_accept.rs b/src/shard/conn_accept.rs index 434968b4..950291cd 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 11cb3c8f..4734c2a3 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 38113144..b35f6d91 100644 --- a/src/shard/mod.rs +++ b/src/shard/mod.rs @@ -367,6 +367,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 +434,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 4f8a227a..8f316631 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. +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,24 @@ 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 `maxmemory != 0` once per drain + // cycle instead of taking `runtime_config.read()` per write command. When + // neither maxmemory nor disk-offload is configured — the common + // non-memory-bound path — every write arm below skips the eviction call + // (and its lock acquire) entirely. + let evict_active = spill_sender.is_some() || runtime_config.read().maxmemory != 0; + // 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 +257,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 +294,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 +353,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 { @@ -535,6 +608,12 @@ pub(crate) fn handle_shard_message_shared( // 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. + // M2 fix: 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, mirroring the + // is_write gate below. crate::shard::slice::with_shard(|s| { ksmv::with_two_slice_dbs( &mut s.databases, @@ -543,6 +622,20 @@ pub(crate) fn handle_shard_message_shared( |src, dst| { src.refresh_now_from_cache(cached_clock); dst.refresh_now_from_cache(cached_clock); + if evict_active { + if let Err(oom) = 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, @@ -584,115 +677,142 @@ pub(crate) fn handle_shard_message_shared( // 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. @@ -740,6 +860,23 @@ pub(crate) fn handle_shard_message_shared( 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); } @@ -853,6 +990,23 @@ pub(crate) fn handle_shard_message_shared( 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); } @@ -1010,87 +1164,110 @@ pub(crate) fn handle_shard_message_shared( { 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; @@ -1124,6 +1301,23 @@ pub(crate) fn handle_shard_message_shared( 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); } @@ -1238,6 +1432,23 @@ pub(crate) fn handle_shard_message_shared( 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 +3364,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 +3395,10 @@ mod drain_cap_tests { &mut autovacuum, None, true, // wal_kv_log + &rtcfg, + None, + &spill_fid, + None, ); assert!( hit_cap, @@ -3211,6 +3430,10 @@ mod drain_cap_tests { &mut autovacuum, None, true, // wal_kv_log + &rtcfg, + None, + &spill_fid, + None, ); assert!( !hit_cap2, diff --git a/tests/oom_bypass_closure.rs b/tests/oom_bypass_closure.rs new file mode 100644 index 00000000..07a2e71a --- /dev/null +++ b/tests/oom_bypass_closure.rs @@ -0,0 +1,466 @@ +//! `--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, ~87.5KB / (2MB / 4 = + // 512KB per shard budget) means each shard accepts on the order of + // 512KB/4KB=128 writes before OOMing => ~4*128=512 total OKs out of + // 3000, then OOM for the rest. Without the fix, only the ~1/4 of keys + // that happen to land on THIS connection's own local shard can ever OOM + // (the pre-existing, already-correct local path) — the other ~3/4 + // (remote legs) keep returning +OK unconditionally. That bounds the + // pre-fix OOM count at roughly N/4 (~750); the threshold below sits + // safely above that ceiling and safely below the post-fix floor. + assert!( + oom_count >= (N * 2) / 3, + "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" + ); +} From a4847c5c0c44aaf8b95665a556e81b66c7b15dbd Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Mon, 6 Jul 2026 14:55:13 +0700 Subject: [PATCH 3/9] test(shard): add case E cross-shard COPY OOM regression + confirm eviction coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to 5cea6912. Review flagged that the COPY-dest fix in that commit (spsc_handler.rs's plain Execute-arm two-database intercept) was untested end-to-end: no wire-level case exercised it, so "COPY dest is gated" was asserted, not proven. Verified COPY's actual eviction coverage first (metadata.rs confirms `COPY` and `BITOP` both carry the generic `W` write flag): on the cross-shard dispatch arms handler_monoio actually uses for remote writes (ExecuteSlotted, PipelineBatchSlotted), COPY already hits the same generic `is_write` eviction gate as any other write command, added in 5cea6912 for cases B/C — there is no COPY-specific bypass on that path. The Execute-arm dest-gate from 5cea6912 remains correct and is kept (it closes the narrower gap on handler_sharded's single-key remote-dispatch arm, which does bypass the generic gate via an early return), but it is not what the new test below exercises. New case E in tests/oom_bypass_closure.rs: SET N keys under budget, then COPY each to db 1 with a DISTINCT dst key (same src/dst key name is rejected by Redis semantics regardless of db, which was the bug in an earlier draft of this test — it produced a false 0-OOM reading because every same-key copy failed before any memory grew). Confirmed genuine RED by reverting the 8 src/ files from 5cea6912 to their parent-commit versions, rebuilding, and rerunning case E: 0/300 OOM (git-stash-style verification, not a mechanical `git stash` since the fix was already committed). Restored the fix via `git checkout --`, rebuilt, reran: 104/300 OOM. Threshold (N/10 = 30) sits safely between both observed counts. GREEN on both runtime-monoio (default) and runtime-tokio. Also documents, but does not fix, a separate pre-existing bug found while building this test: cross-shard remote COPY ... DB n silently ignores the DB clause and performs a same-db copy instead, because the two-database intercept exists only in the plain Execute arm, not ExecuteSlotted/PipelineBatchSlotted. Confirmed via manual probe (4/20 remote COPYs landed in the wrong db). Data-correctness bug, independent of --maxmemory (the wrong same-db copy is still eviction-gated by the generic path), out of scope for this fix — flagged as a follow-up requiring the intercept to be replicated across every ShardMessage arm. CHANGELOG.md updated to state precisely what is and isn't closed: COPY eviction coverage is confirmed/hardened; the COPY DB-clause functional bug is a new, separate, open finding. author: Tin Dang --- CHANGELOG.md | 42 ++++++++++++----- tests/oom_bypass_closure.rs | 92 +++++++++++++++++++++++++++++++++++++ 2 files changed, 123 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b2697b62..79afa7bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -58,18 +58,38 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 writes (`src/scripting/functions.rs`) keep a documented, pre-existing gap (no shard context threaded through the function-library loader) — tracked as a follow-up, not closed by this fix. -- **Cross-db `COPY ... DB n` now enforces `--maxmemory`.** Same SPSC - intercept family as above (`src/shard/spsc_handler.rs`) special-cases - `MOVE`/`COPY DB` ahead of the generic write path because both need two - `&mut Database` borrows at once; `COPY` duplicates the value into the - destination db and was missed by the generic gate. Now runs - `spsc_eviction_gate` against the destination db before `copy_core`. - Same-shard `MOVE` is left ungated (net-zero — the key leaves the source db - as it lands in the destination, same rationale as `DEL`). -- New wire-level regression suite `tests/oom_bypass_closure.rs` (4 cases, +- **Cross-db `COPY ... DB n` eviction confirmed/hardened.** `COPY` carries + the generic `W` write flag, so on the cross-shard dispatch arms + `handler_monoio` actually uses for remote writes (`ExecuteSlotted`, + `PipelineBatchSlotted`) it already hit the same generic `is_write` + eviction gate as any other write command — no COPY-specific bypass exists + there. The plain `Execute` arm's `MOVE`/`COPY DB` two-database intercept + (special-cased ahead of the generic path because both need two `&mut + Database` borrows at once — used by `handler_sharded`'s single-key remote + dispatch) *was* missed, since it returns before reaching the generic gate; + it now runs `spsc_eviction_gate` against the destination db before + `copy_core`. Same-shard `MOVE` is left ungated (net-zero — the key leaves + the source db as it lands in the destination, same rationale as `DEL`). +- **Found, NOT fixed (out of scope): cross-shard remote `COPY ... DB n` + silently ignores the DB clause.** The two-database intercept above exists + only in the plain `Execute` arm, not `ExecuteSlotted`/`PipelineBatchSlotted` + — so when the source key hashes to a different shard than the connection + (the common case under `handler_monoio`), `COPY ... DB n` falls through to + the generic single-db `key_extra::copy` and silently performs a same-db + copy instead of a cross-db one (confirmed via manual probe: 4/20 remote + `COPY`s landed in the wrong db). This is a data-correctness bug independent + of `--maxmemory` — memory growth from the (wrong) same-db copy is still + eviction-gated by the generic path above, so it is not a new OOM bypass, + but the DB clause itself does not do what it says. Fixing it requires + replicating the two-database intercept across every `ShardMessage` arm (or + restructuring it as a shared helper reachable from all of them) — out of + scope for this fix, flagged as a follow-up. +- New wire-level regression suite `tests/oom_bypass_closure.rs` (5 cases, both runtimes): direct-SET control (A), cross-shard pipeline (B), Lua EVAL - loop (C), read-only EVAL under pressure not blocked (D) — B and C RED - before this fix, GREEN after; D locks the write-only scope of the Lua gate. + loop (C), read-only EVAL under pressure not blocked (D), cross-shard COPY + under pressure (E). B, C and E RED before this fix (E: 0/300 OOM + git-stash-verified against the pre-fix tree), GREEN after (E: 104/300); D + locks the write-only scope of the Lua gate. ### Added — int8 symmetric ADC for SQ8 vector search (task #13) diff --git a/tests/oom_bypass_closure.rs b/tests/oom_bypass_closure.rs index 07a2e71a..03c9d444 100644 --- a/tests/oom_bypass_closure.rs +++ b/tests/oom_bypass_closure.rs @@ -464,3 +464,95 @@ fn test_case_d_readonly_eval_not_blocked_at_oom() { writes inside a script, not reads" ); } + +// --------------------------------------------------------------------------- +// Case E: cross-db `COPY ... DB n` under memory pressure. `COPY` carries the +// generic `W` (write) flag (src/command/metadata.rs), so on the cross-shard +// dispatch arms handler_monoio actually uses for remote writes +// (`ExecuteSlotted`/`PipelineBatchSlotted`) it hits the SAME generic +// `is_write` eviction gate cases B/C already exercise — there is no +// COPY-specific bypass on that path. This case is a targeted regression +// check that COPY specifically (not just SET) drives that gate to a real +// OOM under cross-shard load, and locks in the dest-db gate committed in +// spsc_handler.rs's plain `Execute` arm (used by the single-key remote +// dispatch path) as a defensive duplicate. +// +// NOTE — separate, pre-existing, out-of-scope bug found while writing this +// case: cross-shard remote dispatch of `COPY ... DB n` silently ignores the +// DB clause and performs a same-db copy instead (the MOVE/COPY two-database +// intercept in spsc_handler.rs only exists in the plain `Execute` arm, not +// `ExecuteSlotted`/`PipelineBatchSlotted`, so remote COPY falls through to +// the generic single-db `key_extra::copy`). This is a data-correctness bug, +// independent of eviction — confirmed via manual probe (4/20 remote COPYs +// landed in the wrong db). Distinct dst keys below deliberately still land +// in db 0 for the ~3/4 of keys dispatched remotely; this case asserts only +// on the OOM behavior, not on which db the copy actually landed in. +// --------------------------------------------------------------------------- + +#[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 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: + // pre-fix (git-stashed) run = 0/300 OOM; post-fix = 104/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" + ); +} From 86606c1bc1c9ed5206229d00048287af7da87101 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Mon, 6 Jul 2026 15:01:04 +0700 Subject: [PATCH 4/9] docs(changelog): tighten OOM-bypass-closure claims on COPY/MOVE scope MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Post-commit self-review (advisor pass) of 5cea6912 + 7365db09 flagged two places where the CHANGELOG's phrasing overstated what was actually tested: 1. The Execute-arm COPY-dest gate (5cea6912) is exercised by zero wire-level test cases, including the new case E — E pipelines COPY, which routes through ExecuteSlotted/PipelineBatchSlotted on both runtimes, not the plain Execute arm (only reachable via handler_sharded's single-key remote dispatch, which no test in this suite drives). The gate is correct by inspection (verbatim mirror of the tested generic gate) but that is a weaker claim than "gated and tested" — the CHANGELOG now says so explicitly instead of implying test coverage it doesn't have. 2. Case E's RED (0/300 OOM) is produced by reverting the same fix commit that also REDs case B — it demonstrates COPY drives the generic `is_write` eviction gate to OOM once the underlying gate exists, but it is not an independently-discovered COPY-specific RED distinct from B's. Said so explicitly rather than implying E caught its own bypass. 3. The "same-shard MOVE is left ungated (net-zero)" line read as if MOVE's cross-shard DB-clause handling were otherwise fine. It shares COPY's structural gap (the two-database intercept lives only in the plain Execute arm), so it plausibly has the same DB-clause functional defect found for COPY — unverified, not investigated further (out of scope), but now flagged in the same found-not-fixed bullet instead of left unmentioned. No code changes; no test changes. Documentation precision only, so future readers (including future-me) don't over-read what was proven versus what was inspected/inferred. author: Tin Dang --- CHANGELOG.md | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 79afa7bd..ab3eed06 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -70,6 +70,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 it now runs `spsc_eviction_gate` against the destination db before `copy_core`. Same-shard `MOVE` is left ungated (net-zero — the key leaves the source db as it lands in the destination, same rationale as `DEL`). + This Execute-arm dest-gate is inspection-verified as a verbatim mirror of + the tested generic gate; no wire-level case isolates the Execute arm + specifically (the arm is only reachable via `handler_sharded`'s + single-key remote dispatch, not the pipelined path any test in this suite + drives) — see the case E caveat below. - **Found, NOT fixed (out of scope): cross-shard remote `COPY ... DB n` silently ignores the DB clause.** The two-database intercept above exists only in the plain `Execute` arm, not `ExecuteSlotted`/`PipelineBatchSlotted` @@ -80,16 +85,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `COPY`s landed in the wrong db). This is a data-correctness bug independent of `--maxmemory` — memory growth from the (wrong) same-db copy is still eviction-gated by the generic path above, so it is not a new OOM bypass, - but the DB clause itself does not do what it says. Fixing it requires - replicating the two-database intercept across every `ShardMessage` arm (or - restructuring it as a shared helper reachable from all of them) — out of - scope for this fix, flagged as a follow-up. + but the DB clause itself does not do what it says. `MOVE` shares the same + structural setup (its two-database intercept also lives only in the plain + `Execute` arm) and likely has the same DB-target defect on the cross-shard + path — not independently verified here. Fixing it requires replicating the + two-database intercept across every `ShardMessage` arm (or restructuring it + as a shared helper reachable from all of them) — out of scope for this + fix, flagged as a follow-up covering both `COPY` and `MOVE`. - New wire-level regression suite `tests/oom_bypass_closure.rs` (5 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). B, C and E RED before this fix (E: 0/300 OOM - git-stash-verified against the pre-fix tree), GREEN after (E: 104/300); D - locks the write-only scope of the Lua gate. + under pressure (E). B, C and E RED before this fix (E: 0/300 OOM, + GREEN after: 104/300). Case E pipelines `COPY`, which dispatches through + the same `ExecuteSlotted`/`PipelineBatchSlotted` generic-gate path as case + B — its RED/GREEN swing rides the *generic* gate, not the Execute-arm + dest-gate above; it confirms COPY drives that gate to OOM, not that the + Execute-arm path is independently exercised. D locks the write-only scope + of the Lua gate. ### Added — int8 symmetric ADC for SQ8 vector search (task #13) From 3f0970b0d70fddcc778ef29193cadadbcc2db7ab Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Mon, 6 Jul 2026 15:24:29 +0700 Subject: [PATCH 5/9] perf(scripting): per-script snapshot for Lua eviction gate + robust OOM test threshold MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Deep-review fixes on the eviction-bypass-closure work: 1. LuaEvictionCtx::gate took a RuntimeConfig RwLock read on EVERY write redis.call/pcall, even with maxmemory unset — the doc comment claimed a lock-free short-circuit that did not exist. The gate now decides "is eviction active?" from a per-script-execution snapshot: set_script_db bumps a thread-local SCRIPT_GENERATION counter and the ctx caches (generation, active) in a Cell, so a tight redis.call('SET', ...) loop pays one Cell read per write call and reads the lock at most once per script run — mirroring the batch_eviction_active per-batch snapshot the connection handlers use. Staleness is bounded to one script execution, same as the handlers' per-batch bound. 2. tests/oom_bypass_closure.rs case B threshold was flaky: GAP-1 elastic budget redistributes per-shard caps on a 100ms snapshot tick, so the OK-count before OOM kicks in is timing-sensitive (observed 500-1100 OKs across runs). Threshold widened from 2N/3 to N/2, which still cleanly discriminates the pre-fix ceiling (~N/4 = 750 OOMs, bypass present) from the post-fix floor (~1900+ OOMs observed). Perf-reviewer finding on the SPSC drain (one lock read per drain cycle) reviewed and accepted as-is: matches the established batch_eviction_active precedent, amortized over up to 256 messages; an atomic maxmemory snapshot is noted as a future refinement. author: Tin Dang --- src/scripting/bridge.rs | 37 ++++++++++++++++++++++++++++++++----- tests/oom_bypass_closure.rs | 23 +++++++++++++---------- 2 files changed, 45 insertions(+), 15 deletions(-) diff --git a/src/scripting/bridge.rs b/src/scripting/bridge.rs index 84d48833..19301d15 100644 --- a/src/scripting/bridge.rs +++ b/src/scripting/bridge.rs @@ -36,9 +36,12 @@ use crate::storage::tiered::spill_thread::SpillRequest; /// 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 -/// `Option` check on `maxmemory == 0` (the common case) short-circuits before -/// any lock or budget lookup. +/// `CURRENT_DB` unsafe deref is untouched) and zero per-call allocation. The +/// common case (`maxmemory` unset, no spill) is decided by a per-script +/// snapshot (`evict_active` below), so a tight `redis.call('SET', ...)` loop +/// pays one `Cell` read per write call — the `RuntimeConfig` lock is read at +/// most once per script execution, mirroring the per-batch +/// `batch_eviction_active` snapshot the connection handlers use. #[derive(Clone)] pub struct LuaEvictionCtx(Option); @@ -50,6 +53,12 @@ struct LuaEvictionInner { spill_sender: Option>, spill_file_id: Rc>, disk_offload_dir: Option, + /// `(script generation, evict-active)` snapshot, refreshed at most once + /// per script execution (generation bumped by `set_script_db`). Sentinel + /// generation `u64::MAX` forces a refresh on first use. Staleness is + /// bounded to one script run — same bound as the handlers' per-batch + /// snapshot. + evict_active: Cell<(u64, bool)>, } impl LuaEvictionCtx { @@ -78,6 +87,7 @@ impl LuaEvictionCtx { spill_sender, spill_file_id, disk_offload_dir, + evict_active: Cell::new((u64::MAX, false)), })) } @@ -91,10 +101,22 @@ impl LuaEvictionCtx { let Some(inner) = self.0.as_ref() else { return Ok(()); }; - let rt = inner.runtime_config.read(); - if rt.maxmemory == 0 && inner.spill_sender.is_none() { + // Lock-free fast path: decide "is eviction active at all?" from a + // per-script snapshot so a script issuing thousands of writes takes + // the RuntimeConfig read lock once, not per redis.call. + let generation = SCRIPT_GENERATION.with(|c| c.get()); + let (cached_generation, cached_active) = inner.evict_active.get(); + let active = if cached_generation == generation { + cached_active + } else { + let active = inner.spill_sender.is_some() || inner.runtime_config.read().maxmemory != 0; + inner.evict_active.set((generation, active)); + active + }; + if !active { 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(); @@ -124,6 +146,10 @@ thread_local! { static SCRIPT_HAD_WRITE: Cell = const { Cell::new(false) }; /// Whether this script is running in read-only mode (FCALL_RO). static SCRIPT_READ_ONLY: Cell = const { Cell::new(false) }; + /// Monotonic script-execution counter, bumped by `set_script_db`, so + /// per-VM caches (`LuaEvictionCtx::evict_active`) can detect a new + /// script run without taking any lock. + static SCRIPT_GENERATION: Cell = const { Cell::new(0) }; } /// Set the thread-local database pointer before script execution. @@ -132,6 +158,7 @@ pub fn set_script_db(db: &mut crate::storage::Database, db_idx: usize, db_count: CURRENT_DB_IDX.with(|c| c.set(db_idx)); CURRENT_DB_COUNT.with(|c| c.set(db_count)); SCRIPT_HAD_WRITE.with(|c| c.set(false)); + SCRIPT_GENERATION.with(|c| c.set(c.get().wrapping_add(1))); } /// Clear the thread-local database pointer after script execution. diff --git a/tests/oom_bypass_closure.rs b/tests/oom_bypass_closure.rs index 03c9d444..50bfa3fa 100644 --- a/tests/oom_bypass_closure.rs +++ b/tests/oom_bypass_closure.rs @@ -364,17 +364,20 @@ fn test_case_b_cross_shard_pipeline_oom() { .count(); let other_count = N - oom_count - ok_count; - // With eviction enforced across ALL 4 shards, ~87.5KB / (2MB / 4 = - // 512KB per shard budget) means each shard accepts on the order of - // 512KB/4KB=128 writes before OOMing => ~4*128=512 total OKs out of - // 3000, then OOM for the rest. Without the fix, only the ~1/4 of keys - // that happen to land on THIS connection's own local shard can ever OOM - // (the pre-existing, already-correct local path) — the other ~3/4 - // (remote legs) keep returning +OK unconditionally. That bounds the - // pre-fix OOM count at roughly N/4 (~750); the threshold below sits - // safely above that ceiling and safely below the post-fix floor. + // 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) / 3, + 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 / \ From 436c7d6ca4471255f98b6214a06479c7b699f8fd Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Mon, 6 Jul 2026 16:26:15 +0700 Subject: [PATCH 6/9] perf(shard,scripting): atomic maxmemory snapshot (Gap C) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two call sites answered the same question — "is --maxmemory nonzero?" — with different, more expensive mechanisms: - `drain_spsc_shared` (src/shard/spsc_handler.rs) took a `runtime_config.read()` lock once per SPSC drain cycle. - `LuaEvictionCtx::gate` (src/scripting/bridge.rs) used a generation-counter/`Cell` snapshot (`SCRIPT_GENERATION`, bumped by `set_script_db`) to avoid a per-`redis.call` lock read. Publish that boolean as a single process-global atomic instead: `storage::eviction::MAXMEMORY_GLOBAL` (`AtomicU64`, 0 = unset), with `publish_maxmemory(bytes)` (Relaxed store) and `maxmemory_is_set()` (Relaxed load != 0). This is a distinct, simpler atomic from the pre-existing `MAXMEMORY_HINT`/`MAXMEMORY_PER_SHARD_HINT` pair, which additionally encodes the exact budget for the inline write path's skip-eviction arithmetic — those are untouched. Published at every production write site of `RuntimeConfig.maxmemory`: `CONFIG SET maxmemory` (src/command/config.rs) and server startup (src/main.rs, after the guardrail-capped value resolves into `RuntimeConfig`). `src/shard/shared_databases.rs`'s `rt_config()` helper was checked and confirmed test-only (used only by `#[cfg(test)]` elastic-budget tests) — no publish needed there. Consumers replaced: - spsc_handler.rs:141's per-drain-cycle `evict_active` snapshot now reads `maxmemory_is_set()` instead of taking the `RuntimeConfig` lock. `runtime_config` is still threaded through unchanged — the actual eviction pass (`spsc_eviction_gate`) still needs the live budget. - bridge.rs's `LuaEvictionCtx::gate` drops the `SCRIPT_GENERATION` thread-local, the `evict_active: Cell<(u64, bool)>` field, and the generation-comparison branch entirely; replaced with a single `maxmemory_is_set()` load ORed with the existing `spill_sender.is_some()` check. `set_script_db`'s signature is unchanged (only the generation bump inside it is removed). Tests (red/green TDD): - Unit test `maxmemory_publish_and_is_set_roundtrip` (src/storage/eviction.rs) covers publish/un-publish round-tripping. - New case F in tests/oom_bypass_closure.rs: server starts with `--maxmemory 0 --disk-offload disable` (both required — omitting `--maxmemory` trips the config auto-guardrail's nonzero default, and disk-offload's spill-sender presence ORs into the same `evict_active` gate, either of which would mask whether `CONFIG SET maxmemory` specifically published the atomic). Drives a 3000-command cross-shard pipeline before CONFIG SET (all OK), after `CONFIG SET maxmemory 2097152` (must mostly OOM), and after `CONFIG SET maxmemory 0` (all OK again). Confirmed genuinely RED before the fix: manually reverting the CONFIG SET call site's publish call (keeping the startup publish and `runtime_config.maxmemory` assignment intact) reproduced 775/3000 OOM — exactly the local-only share a connection's own shard would still gate via the independent `batch_eviction_active` check in handler_monoio, with the cross-shard SPSC legs bypassing entirely. GREEN after restoring the call: 3000/3000. Verification: cargo fmt (no changes needed), cargo clippy -- -D warnings (default features, clean), CARGO_TARGET_DIR=target-tokio cargo clippy --no-default-features --features runtime-tokio,jemalloc -- -D warnings (clean), cargo test --lib (3732 passed, 1 ignored), CARGO_TARGET_DIR=target-tokio cargo test --no-default-features --features runtime-tokio,jemalloc --lib (3075 passed, 1 ignored), cargo test --test oom_bypass_closure on both a monoio release-fast binary and a tokio release-fast binary (6/6 passed on each). author: Tin Dang --- CHANGELOG.md | 33 ++++++-- src/command/config.rs | 4 + src/main.rs | 4 + src/scripting/bridge.rs | 45 +++------- src/shard/spsc_handler.rs | 14 ++-- src/storage/eviction.rs | 47 +++++++++++ tests/oom_bypass_closure.rs | 161 ++++++++++++++++++++++++++++++++++++ 7 files changed, 259 insertions(+), 49 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ab3eed06..0126aa55 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -43,7 +43,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `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). + 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 @@ -92,16 +99,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 two-database intercept across every `ShardMessage` arm (or restructuring it as a shared helper reachable from all of them) — out of scope for this fix, flagged as a follow-up covering both `COPY` and `MOVE`. -- New wire-level regression suite `tests/oom_bypass_closure.rs` (5 cases, +- New wire-level regression suite `tests/oom_bypass_closure.rs` (6 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). B, C and E RED before this fix (E: 0/300 OOM, - GREEN after: 104/300). Case E pipelines `COPY`, which dispatches through - the same `ExecuteSlotted`/`PipelineBatchSlotted` generic-gate path as case - B — its RED/GREEN swing rides the *generic* gate, not the Execute-arm - dest-gate above; it confirms COPY drives that gate to OOM, not that the - Execute-arm path is independently exercised. D locks the write-only scope - of the Lua gate. + under pressure (E), `CONFIG SET maxmemory` atomic publish/un-publish (F). + B, C and E RED before this fix (E: 0/300 OOM, GREEN after: 104/300). Case E + pipelines `COPY`, which dispatches through the same + `ExecuteSlotted`/`PipelineBatchSlotted` generic-gate path as case B — its + RED/GREEN swing rides the *generic* gate, not the Execute-arm dest-gate + above; it confirms COPY drives that gate to OOM, not that the Execute-arm + path is independently exercised. 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). ### Added — int8 symmetric ADC for SQ8 vector search (task #13) diff --git a/src/command/config.rs b/src/command/config.rs index a6d78e8b..b08f223f 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/main.rs b/src/main.rs index 12f8779a..e4d9aa2b 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 19301d15..76088588 100644 --- a/src/scripting/bridge.rs +++ b/src/scripting/bridge.rs @@ -37,11 +37,10 @@ use crate::storage::tiered::spill_thread::SpillRequest; /// 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 per-script -/// snapshot (`evict_active` below), so a tight `redis.call('SET', ...)` loop -/// pays one `Cell` read per write call — the `RuntimeConfig` lock is read at -/// most once per script execution, mirroring the per-batch -/// `batch_eviction_active` snapshot the connection handlers use. +/// 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); @@ -53,19 +52,13 @@ struct LuaEvictionInner { spill_sender: Option>, spill_file_id: Rc>, disk_offload_dir: Option, - /// `(script generation, evict-active)` snapshot, refreshed at most once - /// per script execution (generation bumped by `set_script_db`). Sentinel - /// generation `u64::MAX` forces a refresh on first use. Staleness is - /// bounded to one script run — same bound as the handlers' per-batch - /// snapshot. - evict_active: Cell<(u64, bool)>, } impl LuaEvictionCtx { - /// No-op gate. Used by unit tests (no real shard context available) and - /// by the FCALL path (`src/scripting/functions.rs`), which has a - /// pre-existing, documented gap for in-function write eviction — closing - /// it is out of scope for this fix (see `tmp/OOM-SHIELD-CONTEXT.md`). + /// 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) } @@ -87,7 +80,6 @@ impl LuaEvictionCtx { spill_sender, spill_file_id, disk_offload_dir, - evict_active: Cell::new((u64::MAX, false)), })) } @@ -101,19 +93,9 @@ impl LuaEvictionCtx { let Some(inner) = self.0.as_ref() else { return Ok(()); }; - // Lock-free fast path: decide "is eviction active at all?" from a - // per-script snapshot so a script issuing thousands of writes takes - // the RuntimeConfig read lock once, not per redis.call. - let generation = SCRIPT_GENERATION.with(|c| c.get()); - let (cached_generation, cached_active) = inner.evict_active.get(); - let active = if cached_generation == generation { - cached_active - } else { - let active = inner.spill_sender.is_some() || inner.runtime_config.read().maxmemory != 0; - inner.evict_active.set((generation, active)); - active - }; - if !active { + // 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(); @@ -146,10 +128,6 @@ thread_local! { static SCRIPT_HAD_WRITE: Cell = const { Cell::new(false) }; /// Whether this script is running in read-only mode (FCALL_RO). static SCRIPT_READ_ONLY: Cell = const { Cell::new(false) }; - /// Monotonic script-execution counter, bumped by `set_script_db`, so - /// per-VM caches (`LuaEvictionCtx::evict_active`) can detect a new - /// script run without taking any lock. - static SCRIPT_GENERATION: Cell = const { Cell::new(0) }; } /// Set the thread-local database pointer before script execution. @@ -158,7 +136,6 @@ pub fn set_script_db(db: &mut crate::storage::Database, db_idx: usize, db_count: CURRENT_DB_IDX.with(|c| c.set(db_idx)); CURRENT_DB_COUNT.with(|c| c.set(db_count)); SCRIPT_HAD_WRITE.with(|c| c.set(false)); - SCRIPT_GENERATION.with(|c| c.set(c.get().wrapping_add(1))); } /// Clear the thread-local database pointer after script execution. diff --git a/src/shard/spsc_handler.rs b/src/shard/spsc_handler.rs index 8f316631..cdc806f3 100644 --- a/src/shard/spsc_handler.rs +++ b/src/shard/spsc_handler.rs @@ -133,12 +133,14 @@ pub(crate) fn drain_spsc_shared( let mut drained = 0; // Batch-level eviction gate (perf parity with handler_monoio's - // `batch_eviction_active`): snapshot `maxmemory != 0` once per drain - // cycle instead of taking `runtime_config.read()` per write command. When - // neither maxmemory nor disk-offload is configured — the common - // non-memory-bound path — every write arm below skips the eviction call - // (and its lock acquire) entirely. - let evict_active = spill_sender.is_some() || runtime_config.read().maxmemory != 0; + // `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. // diff --git a/src/storage/eviction.rs b/src/storage/eviction.rs index 72f0d0eb..0e872edf 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 index 50bfa3fa..48e1ce1c 100644 --- a/tests/oom_bypass_closure.rs +++ b/tests/oom_bypass_closure.rs @@ -559,3 +559,164 @@ fn test_case_e_cross_db_copy_oom() { 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" + ); +} From 6e75a7b1d459e7e0f31e546a9596450680e8d228 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Mon, 6 Jul 2026 16:46:47 +0700 Subject: [PATCH 7/9] fix(scripting): FCALL writes run the eviction gate (Gap B) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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 — a Lua FUNCTION doing `redis.call('SET', ...)` in a loop bypassed `--maxmemory` entirely, independent of the EVAL/EVALSHA fix that already covers the shared VM. `FunctionRegistry::new` now takes a `LuaEvictionCtx`, stores it, and `create_library` clones it into every library's freshly-created VM (each library gets its own VM, so the ctx must be cloned per VM, not wired once). Both production construction sites now build a real ctx from the same shard handles `setup_lua_vm` already uses for the shared VM (src/shard/conn_accept.rs's recipe): - src/server/conn/handler_monoio/mod.rs:172 - src/server/conn/handler_sharded/mod.rs:281 (tokio runtime) Both read `ctx.{shard_databases, runtime_config, shard_id, spill_sender, spill_file_id, disk_offload_dir}` off the shard's `ConnectionContext` — already present, no new handle needed. `core.rs`'s `#[allow(dead_code)]` on those three fields is removed since both handlers now read them (previously only handler_monoio's tiered-storage eviction path did under some builds). The gate itself was already correct once wired (FCALL brackets execution with set_script_db/clear_script_db, gates FCALL_RO via SCRIPT_READ_ONLY); after the Gap C atomic-snapshot commit the common case (`maxmemory` unset) is one Relaxed load, so this adds no measurable per-call cost to the FCALL hot path. Test-only call sites (functions.rs's own unit tests) pass `LuaEvictionCtx::disabled()` — the module doc-comment on `disabled()` already scopes it to unit tests plus real production ctx builders, updated in the prior Gap C commit. Red/green TDD: - New cases in tests/oom_bypass_closure.rs: - Case G: FUNCTION LOAD a library whose function loops `redis.call('SET', ...)` 2000 times (~16MB attempted), FCALL it against a 2MB-capped noeviction server. RED before the fix: FCALL completed with `'DONE'` instead of an OOM error (confirmed by temporarily re-running against the pre-fix registry — the write loop never trips the gate). GREEN after: FCALL surfaces the OOM error (wrapped as `ERR Error running script: runtime error: OOM ...` since redis.call propagates it as a Lua runtime error, uncaught by the function body). - Case H: control — FCALL with no maxmemory configured must still succeed (mirrors case D for EVAL), locking the write-only/OOM-only scope of the gate. - Lua FUNCTION bodies are called with zero Lua arguments (`call_function` invokes `registered.call(())`) — they read `KEYS`/`ARGV` as globals, same as an EVAL script body, not as function parameters. (First draft of these tests used `function(keys, args)` parameters and failed with "attempt to index a nil value" — fixed before the gap-specific RED assertion was ever meaningfully exercised.) Verification: cargo fmt, cargo clippy -- -D warnings (default features, clean), CARGO_TARGET_DIR=target-tokio cargo clippy --no-default-features --features runtime-tokio,jemalloc -- -D warnings (clean), cargo test --lib (3732 passed, 1 ignored), CARGO_TARGET_DIR=target-tokio cargo test --no-default-features --features runtime-tokio,jemalloc --lib (3075 passed, 1 ignored), cargo test --test oom_bypass_closure on both a monoio release-fast binary and a tokio release-fast binary (8/8 passed on each). author: Tin Dang --- CHANGELOG.md | 56 +++++++++------ src/scripting/functions.rs | 43 ++++++----- src/server/conn/core.rs | 6 +- src/server/conn/handler_monoio/mod.rs | 16 ++++- src/server/conn/handler_sharded/mod.rs | 14 +++- tests/oom_bypass_closure.rs | 99 ++++++++++++++++++++++++++ 6 files changed, 189 insertions(+), 45 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0126aa55..5cf249b8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -61,10 +61,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `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 - writes (`src/scripting/functions.rs`) keep a documented, pre-existing gap - (no shard context threaded through the function-library loader) — tracked - as a follow-up, not closed by this fix. + 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-db `COPY ... DB n` eviction confirmed/hardened.** `COPY` carries the generic `W` write flag, so on the cross-shard dispatch arms `handler_monoio` actually uses for remote writes (`ExecuteSlotted`, @@ -99,24 +110,29 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 two-database intercept across every `ShardMessage` arm (or restructuring it as a shared helper reachable from all of them) — out of scope for this fix, flagged as a follow-up covering both `COPY` and `MOVE`. -- New wire-level regression suite `tests/oom_bypass_closure.rs` (6 cases, +- 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). - B, C and E RED before this fix (E: 0/300 OOM, GREEN after: 104/300). Case E - pipelines `COPY`, which dispatches through the same - `ExecuteSlotted`/`PipelineBatchSlotted` generic-gate path as case B — its - RED/GREEN swing rides the *generic* gate, not the Execute-arm dest-gate - above; it confirms COPY drives that gate to OOM, not that the Execute-arm - path is independently exercised. 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). + 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 (E: 0/300 OOM, GREEN + after: 104/300). Case E pipelines `COPY`, which dispatches through the + same `ExecuteSlotted`/`PipelineBatchSlotted` generic-gate path as case B — + its RED/GREEN swing rides the *generic* gate, not the Execute-arm + dest-gate above; it confirms COPY drives that gate to OOM, not that the + Execute-arm path is independently exercised. 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. ### Added — int8 symmetric ADC for SQ8 vector search (task #13) diff --git a/src/scripting/functions.rs b/src/scripting/functions.rs index baf8830b..8ccdfb35 100644 --- a/src/scripting/functions.rs +++ b/src/scripting/functions.rs @@ -102,13 +102,23 @@ pub struct FunctionRegistry { libraries: HashMap, /// 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,14 +293,11 @@ impl FunctionRegistry { let lua = Rc::new(mlua::Lua::new()); crate::scripting::sandbox::setup_sandbox(&lua) .map_err(|e| LoadError::LuaError(e.to_string()))?; - // FCALL-internal writes: pre-existing, documented gap (no shard - // context is threaded through the function-library loader). Closing - // it is out of scope here — see tmp/OOM-SHIELD-CONTEXT.md. - crate::scripting::sandbox::register_redis_api( - &lua, - crate::scripting::bridge::LuaEvictionCtx::disabled(), - ) - .map_err(|e| LoadError::LuaError(e.to_string()))?; + // 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 let func_table = lua @@ -532,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(); @@ -545,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(); @@ -557,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 = @@ -569,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(); @@ -579,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(); @@ -589,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(); @@ -600,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")); @@ -610,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(); @@ -622,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/server/conn/core.rs b/src/server/conn/core.rs index dc2e2365..385dbb64 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 dc85513f..71965798 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 894b5089..b3674498 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/tests/oom_bypass_closure.rs b/tests/oom_bypass_closure.rs index 48e1ce1c..b60f48a2 100644 --- a/tests/oom_bypass_closure.rs +++ b/tests/oom_bypass_closure.rs @@ -720,3 +720,102 @@ fn test_case_f_config_set_maxmemory_publishes_atomic() { 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:?}" + ); +} From 5278adb9c3b9e2d7dc5c98193f8992ba3b9680dc Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Mon, 6 Jul 2026 17:37:22 +0700 Subject: [PATCH 8/9] fix(shard): two-db COPY/MOVE intercept on all SPSC arms (Gap A) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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 `src/shard/spsc_handler.rs` special-cased both commands ahead of the generic path, but the other five `ShardMessage` arms (`MultiExecute`, `PipelineBatch`, `ExecuteSlotted`, `MultiExecuteSlotted`, `PipelineBatchSlotted`) had no such intercept — including `PipelineBatchSlotted`, the one arm real client traffic (`handler_monoio`/`handler_sharded`'s pipelined remote-write path) actually uses. 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: the copy landed in the SAME db as the source (data corruption — confirmed pre-fix via manual probe, 4/20 remote COPYs landed in the wrong db). 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 extracted verbatim out of the `Execute` arm into a new shared helper, `src/shard/spsc_two_db::try_two_db_intercept`, and every arm now calls it before falling through to the generic write path: - `Some(response)` means MOVE, or COPY with a DB clause targeting a different db — the caller treats it as the command's final response (no generic dispatch, no COW intercept, no auto-index hooks — mirrors the Execute arm's pre-existing behavior). - `None` means anything else (not MOVE/COPY, or a same-db COPY with no DB clause) — falls through to the generic single-db path as before. Persistence (WAL/AOF) stays in each arm's own per-command block, gated on `matches!(response, Frame::Integer(1))` — stricter than the batch arms' generic `!Error` condition, since a same-db MOVE or a COPY of a missing key returns `Integer(0)` and must not persist. Borrow discipline: the batch arms (`MultiExecute`, `PipelineBatch`, and their `*Slotted` variants) previously took one `&mut Database` borrow for the whole loop. The intercept needs `&mut [Database]` (both src and dst) for just the MOVE/COPY command, so the per-db borrow is now taken per-iteration, inside the loop, after the intercept check — the one-time `refresh_now_from_cache` call and any other per-batch setup that doesn't need the borrow stays outside. Cross-db COPY duplicates the value (unlike net-zero MOVE), so the intercept still runs `spsc_eviction_gate` against the destination db before `copy_core`, using the same per-shard elastic budget as any other write. Two things found while wiring this that are explicitly OUT of scope, with CHANGELOG "found, not fixed" notes: 1. Cross-shard COPY's destination key is only reliably `GET`-able when it hash-tag-colocates with the source key. Moon shards purely by key hash with no db-awareness; COPY executes on whichever shard owns the SOURCE key (the routing key for the whole command) and writes the destination into that SAME shard's db space — it does not relocate the write to wherever the destination key's own hash would 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 (same key name, so `hash(key)` is unchanged). This is a pre-existing property of the per-key-hash sharding model applied to a renamed-destination COPY — the Execute arm this fix mirrors has the identical property, not a regression introduced here. Fixing it needs a real cross-shard two-phase protocol, well beyond "apply the intercept to more arms." 2. Moon's eviction gate is per-DATABASE, not per-shard-aggregate-across-dbs: `try_evict_if_needed_budget` measures the ONE db being written to, not the sum across all 16 logical dbs a shard holds. This is uniform across every SPSC write arm (see `spsc_handler.rs:619`'s pre-existing single-db gate call for a normal write) — not something this fix introduces — but it did surface a stale test premise (see below). Red/green TDD: - New suite `tests/spsc_two_db.rs` (3 cases, both runtimes): pipelined cross-shard `COPY ... DB n`, 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). RED confirmed via a git-stash of just the tracked source changes (spsc_two_db.rs is new/ untracked, so stashing spsc_handler.rs + mod.rs alone reverts to the pre-fix state): MOVE hit the loud wrong error, COPY silently same-db copied. GREEN after restoring the stash. The COPY case uses `{i}`-hash-tagged src/dst key pairs (see out-of-scope finding 1) — an earlier untagged draft of this test only passed for the ~1-in-4 keys where `hash(dst)` happened to collide with `hash(src)`. - Existing `tests/oom_bypass_closure.rs` case E (cross-db COPY under memory pressure) initially regressed to 0/300 OOM against this fix — traced to out-of-scope finding 2 above: case E's pre-fix pass relied on the same-db-copy bug landing writes in db 0 (already loaded from phase 1), tripping the per-db gate as a side effect of the routing bug, not because the destination db was under its own pressure. Fixed by adding a ballast phase that pre-loads db 1 with the same keys/sizes as db 0 before the OOM phase, so the destination-db gate has something real to trip on (76/300 OOM fully fixed). Verified this isolates the destination-db gate specifically: temporarily disabling only the `spsc_eviction_gate` call inside `try_two_db_intercept` (keeping the destination-db routing intact) reproduces 0/300 OOM — confirmed RED for the right reason, then restored. Verification: cargo fmt, cargo clippy -- -D warnings (default features, clean), CARGO_TARGET_DIR=target-tokio cargo clippy --no-default-features --features runtime-tokio,jemalloc -- -D warnings (clean), cargo test --lib (3732 passed, 1 ignored), CARGO_TARGET_DIR=target-tokio cargo test --no-default-features --features runtime-tokio,jemalloc --lib (3075 passed, 1 ignored), cargo test --test oom_bypass_closure (8/8) and cargo test --test spsc_two_db (3/3) on both a monoio release-fast binary and a tokio release-fast binary. author: Tin Dang --- CHANGELOG.md | 130 +++++---- src/shard/mod.rs | 4 + src/shard/spsc_handler.rs | 435 ++++++++++++++++++++++-------- src/shard/spsc_two_db.rs | 123 +++++++++ tests/oom_bypass_closure.rs | 81 ++++-- tests/spsc_two_db.rs | 519 ++++++++++++++++++++++++++++++++++++ 6 files changed, 1107 insertions(+), 185 deletions(-) create mode 100644 src/shard/spsc_two_db.rs create mode 100644 tests/spsc_two_db.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 5cf249b8..db7b97d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -76,63 +76,89 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 (`ConnectionContext::{shard_databases, runtime_config, shard_id, spill_sender, spill_file_id, disk_offload_dir}`); only test-only callers pass `LuaEvictionCtx::disabled()`. -- **Cross-db `COPY ... DB n` eviction confirmed/hardened.** `COPY` carries - the generic `W` write flag, so on the cross-shard dispatch arms - `handler_monoio` actually uses for remote writes (`ExecuteSlotted`, - `PipelineBatchSlotted`) it already hit the same generic `is_write` - eviction gate as any other write command — no COPY-specific bypass exists - there. The plain `Execute` arm's `MOVE`/`COPY DB` two-database intercept - (special-cased ahead of the generic path because both need two `&mut - Database` borrows at once — used by `handler_sharded`'s single-key remote - dispatch) *was* missed, since it returns before reaching the generic gate; - it now runs `spsc_eviction_gate` against the destination db before - `copy_core`. Same-shard `MOVE` is left ungated (net-zero — the key leaves - the source db as it lands in the destination, same rationale as `DEL`). - This Execute-arm dest-gate is inspection-verified as a verbatim mirror of - the tested generic gate; no wire-level case isolates the Execute arm - specifically (the arm is only reachable via `handler_sharded`'s - single-key remote dispatch, not the pipelined path any test in this suite - drives) — see the case E caveat below. -- **Found, NOT fixed (out of scope): cross-shard remote `COPY ... DB n` - silently ignores the DB clause.** The two-database intercept above exists - only in the plain `Execute` arm, not `ExecuteSlotted`/`PipelineBatchSlotted` - — so when the source key hashes to a different shard than the connection - (the common case under `handler_monoio`), `COPY ... DB n` falls through to - the generic single-db `key_extra::copy` and silently performs a same-db - copy instead of a cross-db one (confirmed via manual probe: 4/20 remote - `COPY`s landed in the wrong db). This is a data-correctness bug independent - of `--maxmemory` — memory growth from the (wrong) same-db copy is still - eviction-gated by the generic path above, so it is not a new OOM bypass, - but the DB clause itself does not do what it says. `MOVE` shares the same - structural setup (its two-database intercept also lives only in the plain - `Execute` arm) and likely has the same DB-target defect on the cross-shard - path — not independently verified here. Fixing it requires replicating the - two-database intercept across every `ShardMessage` arm (or restructuring it - as a shared helper reachable from all of them) — out of scope for this - fix, flagged as a follow-up covering both `COPY` and `MOVE`. +- **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 (E: 0/300 OOM, GREEN - after: 104/300). Case E pipelines `COPY`, which dispatches through the - same `ExecuteSlotted`/`PipelineBatchSlotted` generic-gate path as case B — - its RED/GREEN swing rides the *generic* gate, not the Execute-arm - dest-gate above; it confirms COPY drives that gate to OOM, not that the - Execute-arm path is independently exercised. 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. + 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) diff --git a/src/shard/mod.rs b/src/shard/mod.rs index b35f6d91..6904cda9 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; diff --git a/src/shard/spsc_handler.rs b/src/shard/spsc_handler.rs index cdc806f3..d96b7f1a 100644 --- a/src/shard/spsc_handler.rs +++ b/src/shard/spsc_handler.rs @@ -47,7 +47,7 @@ use super::shared_databases::ShardDatabases; /// 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. -fn spsc_eviction_gate( +pub(super) fn spsc_eviction_gate( db: &mut Database, db_idx: usize, shard_databases: &Arc, @@ -549,107 +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. - // M2 fix: 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, mirroring the - // is_write gate below. - 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); - if evict_active { - if let Err(oom) = 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, - ) - }, - ) - }) - } - }; + 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 = @@ -674,7 +595,11 @@ 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 @@ -843,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). @@ -860,6 +785,64 @@ 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 @@ -973,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). @@ -990,6 +976,65 @@ 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 @@ -1164,6 +1209,56 @@ 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. @@ -1284,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). @@ -1301,6 +1396,64 @@ 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 @@ -1415,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). @@ -1432,6 +1588,65 @@ 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 diff --git a/src/shard/spsc_two_db.rs b/src/shard/spsc_two_db.rs new file mode 100644 index 00000000..5e5ec374 --- /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/tests/oom_bypass_closure.rs b/tests/oom_bypass_closure.rs index b60f48a2..d88edaee 100644 --- a/tests/oom_bypass_closure.rs +++ b/tests/oom_bypass_closure.rs @@ -469,27 +469,33 @@ fn test_case_d_readonly_eval_not_blocked_at_oom() { } // --------------------------------------------------------------------------- -// Case E: cross-db `COPY ... DB n` under memory pressure. `COPY` carries the -// generic `W` (write) flag (src/command/metadata.rs), so on the cross-shard -// dispatch arms handler_monoio actually uses for remote writes -// (`ExecuteSlotted`/`PipelineBatchSlotted`) it hits the SAME generic -// `is_write` eviction gate cases B/C already exercise — there is no -// COPY-specific bypass on that path. This case is a targeted regression -// check that COPY specifically (not just SET) drives that gate to a real -// OOM under cross-shard load, and locks in the dest-db gate committed in -// spsc_handler.rs's plain `Execute` arm (used by the single-key remote -// dispatch path) as a defensive duplicate. +// 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. // -// NOTE — separate, pre-existing, out-of-scope bug found while writing this -// case: cross-shard remote dispatch of `COPY ... DB n` silently ignores the -// DB clause and performs a same-db copy instead (the MOVE/COPY two-database -// intercept in spsc_handler.rs only exists in the plain `Execute` arm, not -// `ExecuteSlotted`/`PipelineBatchSlotted`, so remote COPY falls through to -// the generic single-db `key_extra::copy`). This is a data-correctness bug, -// independent of eviction — confirmed via manual probe (4/20 remote COPYs -// landed in the wrong db). Distinct dst keys below deliberately still land -// in db 0 for the ~3/4 of keys dispatched remotely; this case asserts only -// on the OOM behavior, not on which db the copy actually landed in. +// 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] @@ -526,6 +532,34 @@ fn test_case_e_cross_db_copy_oom() { — 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. @@ -547,9 +581,10 @@ fn test_case_e_cross_db_copy_oom() { let other_count = N - oom_count - ok_count; // Threshold picked from observed data, same methodology as case B: - // pre-fix (git-stashed) run = 0/300 OOM; post-fix = 104/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. + // 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 \ diff --git a/tests/spsc_two_db.rs b/tests/spsc_two_db.rs new file mode 100644 index 00000000..dc922338 --- /dev/null +++ b/tests/spsc_two_db.rs @@ -0,0 +1,519 @@ +//! 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`); the other +//! arms are covered by direct unit tests of the shared +//! `try_two_db_intercept` helper in `src/shard/spsc_two_db.rs`, which every +//! arm (including `PipelineBatchSlotted`) calls verbatim — so the wire-level +//! coverage here transitively validates the same logic the dead/secondary +//! arms use. +//! +//! 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" + ); + } +} From 410bf2219c503096784f26a2d7b6a7d26c1a684e Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Mon, 6 Jul 2026 17:40:50 +0700 Subject: [PATCH 9/9] docs(shard): fix stale test-coverage claim in spsc_two_db.rs module doc MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Self-review after the Gap A commit (5278adb9) caught a false statement in `tests/spsc_two_db.rs`'s module doc: it claimed the dead/admin-only SPSC arms (`MultiExecute`, `ExecuteSlotted`, `MultiExecuteSlotted`) "are covered by direct unit tests of the shared `try_two_db_intercept` helper" — no such unit tests exist; `src/shard/spsc_two_db.rs` has no `#[cfg(test)]` module. The wire-level cases in this suite do exercise `PipelineBatchSlotted`, which calls `try_two_db_intercept` verbatim (the same helper every other arm calls), so that coverage transitively validates the shared logic — but it does not independently exercise the dead arms' own call sites, and no direct unit test does either. Rewrote the doc comment to state only what is actually true, and to flag the missing direct unit test as a documented gap rather than implying it's covered. No production code change; test-only doc correction. author: Tin Dang --- tests/spsc_two_db.rs | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/spsc_two_db.rs b/tests/spsc_two_db.rs index dc922338..bf50683b 100644 --- a/tests/spsc_two_db.rs +++ b/tests/spsc_two_db.rs @@ -20,12 +20,13 @@ //! 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`); the other -//! arms are covered by direct unit tests of the shared -//! `try_two_db_intercept` helper in `src/shard/spsc_two_db.rs`, which every -//! arm (including `PipelineBatchSlotted`) calls verbatim — so the wire-level -//! coverage here transitively validates the same logic the dead/secondary -//! arms use. +//! 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