diff --git a/.add/state.json b/.add/state.json index 88a7864b..ecb94d41 100644 --- a/.add/state.json +++ b/.add/state.json @@ -1,7 +1,7 @@ { "project": "moon", "stage": "production", - "active_task": "consistency-dispatch-gaps", + "active_task": "shardslice-migration", "active_milestone": "v1-shared-nothing", "tasks": { "hotpath-lock-quickwins": { @@ -33,6 +33,21 @@ "created": "2026-06-11T10:38:27+00:00", "updated": "2026-06-12T02:21:52+00:00", "flag_verified": true + }, + "shardslice-migration": { + "title": "Wire ShardSlice: thread-local shared-nothing storage becomes the live path", + "phase": "done", + "gate": "RISK-ACCEPTED", + "milestone": "v1-shared-nothing", + "depends_on": [], + "created": "2026-06-12T03:23:36+00:00", + "updated": "2026-06-12T17:04:31+00:00", + "flag_verified": true, + "waiver": { + "owner": "Tin Dang", + "ticket": "follow-up task: cross-shard-read-acceleration (observe)", + "expires": "2026-08-01" + } } }, "milestones": { @@ -46,7 +61,7 @@ } }, "created": "2026-06-11T03:18:21+00:00", - "updated": "2026-06-12T02:21:52+00:00", + "updated": "2026-06-12T17:04:31+00:00", "setup": { "locked": true, "locked_at": "2026-06-11T03:28:00+00:00", diff --git a/.add/tasks/shardslice-migration/TASK.md b/.add/tasks/shardslice-migration/TASK.md new file mode 100644 index 00000000..06ff04da --- /dev/null +++ b/.add/tasks/shardslice-migration/TASK.md @@ -0,0 +1,737 @@ +# TASK: Wire ShardSlice: thread-local shared-nothing storage becomes the live path + +slug: shardslice-migration · created: 2026-06-12 · stage: production · risk: high · autonomy: conservative +phase: done + + +> One file = one task. Fill sections top-to-bottom; the `add` skill drives each phase. +> When a phase is unclear, read its book chapter in `.add/docs/` (linked per section). +> The phase marker above is the single source of truth — keep it in sync via `add.py phase`. + +--- + +## 1 · SPECIFY — the rules ▸ docs/03-step-1-specify.md + +Feature: ShardSlice Phase 4 — `init_shard` wired at shard startup on BOTH runtimes, +thread-local shared-nothing storage becomes the live path, and the RwLock/Mutex +wrappers in ShardDatabases are deleted. This is the milestone's "without +per-command global locks" goal made real: today `slice::init_shard` has ZERO +production call sites, `is_initialized()` is always false, and every one of the +~130 dual-branch decision points takes the lock-based `else` branch. + +Ground truth (3 parallel scouts, 2026-06-12, file:line evidenced): +- BOTH runtimes already pin connection tasks to dedicated per-shard OS threads: + monoio = one runtime per `shard-N` thread (main.rs:1303–1347, + monoio_impl.rs:60–66); tokio = `new_current_thread` + `LocalSet` + + `spawn_local` (tokio_impl.rs:57–65, conn_accept.rs:263). NO work-stealing + anywhere — `!Send` ShardSlice fits both runtimes natively. +- Structural gap at startup: databases are moved into `ShardDatabases::new` + (main.rs:1226–1230, embedded.rs:278) BEFORE shard threads spawn; `Shard::run` + has nothing to hand `init_shard`. Boot ownership must be restructured so each + shard thread receives its `Box<[Database]>` + stores and calls `init_shard` + at the top of `run()`, before the first accept. +- Slice-vs-lock branch divergences (must be fixed, not collapsed blindly): + WS.DROP key cleanup (handler_sharded/mod.rs:116-area: slice deletes on conn's + shard, lock routes to {wsid} owner), ALL SIX MQ arms in both runtimes' + write.rs (slice conn-local, lock owner-routed — comments name this task), + TXN.COMMIT MQ materialize in both txn.rs (same). Everything else (~125 + branch points across 16 files) is semantically equivalent and collapses. +- WAL appends are channel SENDS (`MpscSender`, wal_append_txs) — Send by + design; cross-shard `wal_append(owner, …)` stays legal under slice with NO + message variant. The senders array stays in the residual shared struct. +- Cross-thread direct-LOCK consumers that block wrapper deletion (the real + Phase 4 work beyond wiring): (1) AOF rewrite fold — `aof-writer-N` threads + take a foreign shard's write locks (aof/rewrite.rs:497–502 carries an + explicit "Phase 4 MUST revisit" warning); (2) Prometheus memory publisher + (metrics_setup.rs:1396–1429, admin tokio thread locks read_db + vector + + graph across ALL shards every 15s); (3) MEMORY DOCTOR + (command/server_admin.rs:385–416, same pattern on demand). All other + background work is own-shard on the shard thread and already Phase-2d-gated. +- SPSC gaps needing new ShardMessage variants (after the WAL-sender and + global-WS simplifications): an MQ domain hop (GraphCommand precedent — one + variant serving CREATE/PUSH/POP/ACK/DLQLEN/TRIGGER on the owner, since + Execute cannot express durable-flag/DLQ logic), a TXN MQ-intent materialize + hop, and a WS.DROP prefix-cleanup hop. +- ~30 test/bench `ShardDatabases::new` construct sites break when its shape + changes; Phase-3 leftovers (aggregate_memory → read_memory_sum at + persistence_tick.rs:344,499; vector/graph memory atomics for metrics) become + prerequisites for deleting the wrappers. + +Framings weighed: Full Phase 4, both runtimes (CHOSEN — user decision +2026-06-12: wire unconditionally, delete wrappers, collapse all dual branches +in this task) · flag-gated cutover with wrapper deletion deferred (offered, +declined) · monoio-only with tokio keeping locks (offered, declined — moot +anyway: tokio is already thread-pinned). + +Must: + + - M1 (wiring): `init_shard(ShardSlice::new(…))` runs at the top of + `Shard::run()` on every shard thread, both runtimes, both entry points + (main.rs AND embedded.rs), before the first connection is accepted. + `is_initialized()` is true on every shard thread in production. Boot + ownership restructured so the slice receives the real `Box<[Database]>`, + vector/text/graph stores, intents, registries, and the + `memory_publisher(shard_id)` atomic — not copies. + - M2 (owner-routing the divergent branches): the five divergence groups + (WS.DROP cleanup; MQ CREATE/PUSH/POP/ACK/DLQLEN/TRIGGER ×2 runtimes; + TXN.COMMIT MQ materialize ×2) reach the OWNING shard under slice via new + SPSC domain hops (MqCommand-style message + MQ-intent materialize + WS + prefix-cleanup), mirroring the GraphCommand precedent. The consistency + suite must hold 197/197 at shards=1/4/12 with slice live — same bar as + the consistency-dispatch-gaps task. + - M3 (workspace registry): the slot-0 Mutex trick is replaced by ONE + process-global `Mutex` owned OUTSIDE the slices + (control-plane state, rare ops — sharding it buys nothing); the per-shard + `workspace_registry` field leaves ShardSlice/ShardSliceInit; WS WAL + records keep their shard-0 stream via the (still-shared) WAL senders. + - M4 (cross-thread consumers): AOF rewrite no longer takes foreign shard + locks — the fold is redesigned as a shard-thread-cooperative snapshot + handoff (message the shard, it produces the frozen view); the Prometheus + publisher and MEMORY DOCTOR read published per-shard atomics (extend the + memory_per_shard pattern to vector/text/graph memory) instead of locking + all shards; the two Phase-3 leftover aggregate_memory call sites switch + to read_memory_sum. + - M5 (wrapper deletion): the RwLock/Mutex wrappers around per-shard state + are DELETED from ShardDatabases; ~130 dual branches collapse to the slice + path; ShardDatabases shrinks to genuinely-shared state only (SPSC mesh + handles, WAL senders, published atomics, the global workspace registry, + cluster/pubsub handles). The ~30 test/bench construct sites are updated + to the new shape. uring_handler (MOON_URING=1 bridge, runs synchronously + ON the shard thread) switches from write_db locks to with_shard. + - M6 (verification bar): cargo test --lib + all integration suites green on + both runtimes (macOS + Linux VM, MOON_BIN pinned); consistency sweep exit + 0 at 1/4/12; loom models still pass; fuzz targets unaffected; fmt/clippy + clean on both CI feature sets. + - M7 (performance): no bench-swf cell regresses vs current main (same + REQS, idle VM, wakefloor-style comparison); the measured multi-shard + delta is recorded as the milestone's "measured scaling improvement" + evidence — improvement is recorded, not gated (user decision: that + no-regression + honest measurement is the bar). + +Reject: + + - A shard thread serving commands without init_shard having run -> fail + fast at STARTUP (panic/abort with shard id), never a per-command + `with_shard` panic at runtime ("uninitialized_shard"). + - Nested `with_shard`/`with_shard_db` (RefCell double-borrow) introduced by + the migration -> forbidden; closures stay flat, cross-shard work exits + the closure before hopping ("slice_reentrancy"). + - Any `unsafe`, `Send`/`Sync` impl, or pointer trick to share a ShardSlice + across threads -> forbidden; the `!Send` compile error IS the design + ("slice_send_violation"). + - Weakening any test, the consistency suite, or protocol semantics to make + the cutover pass -> forbidden ("test_weakening"). + - A background/admin path quietly re-acquiring a foreign-shard lock through + a retained wrapper -> forbidden once M5 lands; the wrappers must be GONE, + not bypassed ("wrapper_resurrection"). + - Holding a `with_shard` borrow across an `.await` -> forbidden (same rule + as locks across await; the borrow is synchronous by construction) + ("borrow_across_await"). + +After: + + - Production servers (monoio AND tokio, main AND embedded) run with + thread-local ShardSlices live on every shard thread; per-command access + to shard-local state takes zero lock operations. + - ShardDatabases contains no RwLock/Mutex wrappers; the + type still exists but only as the shared-infrastructure handle. + - WS/MQ/TXN-MQ behave identically from every connection at every shard + count (the cdg suite proves it), now via SPSC hops instead of cross-shard + locks. + - AOF rewrite, metrics, and MEMORY DOCTOR work without touching another + thread's state directly. + - bench-swf shows no regression; the multi-shard delta is recorded in §6 + as milestone evidence. + +Assumptions — lowest-confidence first: + + ⚠ The AOF-rewrite redesign is containable inside this task — lowest + confidence because it is a durability path crossing thread ownership + (aof/rewrite.rs:497's own comment defers it to "a future Phase 4"), and + a cooperative-snapshot protocol (freeze + handoff) touches WAL/AOF + correctness under concurrent writes; if wrong: data-loss-grade bug or + the task splits, leaving wrapper deletion (M5) blocked. + ⚠ The three scouts' inventory is COMPLETE — i.e. no un-inventoried + cross-thread Database/store access hides outside the ~130 gated sites + + 3 named consumers (candidates: blocking-command wakeups, Lua, response + slots, connection migration, replication) — lower confidence because + 213 call sites were migrated across Phases 2a–2f by different waves; + if wrong: runtime panic (thread-local miss) or silent wrong-shard data + under slice, surfacing only under load. + - [ ] Tokio's LocalSet pinning makes !Send slices safe under tokio with no + runtime changes (scout-verified: current_thread + spawn_local, + conn_accept.rs:263) — confirm with a tokio slice smoke test early. + - [ ] WAL cross-shard appends via MpscSender stay correct under slice + (senders are Send; the drain side is the owner's event loop). + - [ ] An MqCommand SPSC hop can carry all six MQ subcommands' semantics + (durable flag, DLQ routing, trigger debounce) — the GraphCommand + precedent suggests yes. + - [ ] No-regression (M7) is achievable: the slice path strictly removes + uncontended parking_lot acquisitions from the hot path. + + + + +--- + +## 2 · SCENARIOS — pass/fail cases ▸ docs/04-step-2-scenarios.md + + + +```gherkin +Scenario: ssm1 slice live on every shard thread at startup (M1) + Given a 4-shard server started fresh (each runtime; main and embedded entry) + When the server begins accepting connections + Then every shard thread logs ShardSlice initialization exactly once before + its first accept, and PING/SET/GET succeed from fresh connections + And a probe command confirms the slice path is serving (no lock-path + counters increment / the lock wrappers no longer exist to be counted) + +Scenario: ssm2 WS/MQ owner-routing survives the cutover (M2) + Given a 4-shard server with slice live + When the cdg6 cross-connection tests (a–g) and the full consistency suite run + Then cdg6 passes 7/7 and the suite reports 197/197 at shards=1, 4, and 12 + And MQ CREATE on one connection is PUSH/POP/ACK/DLQLEN-able from fresh + connections landing on every other shard (now via the MQ SPSC hop) + +Scenario: ssm3 workspace registry is process-global (M3) + Given a 4-shard server with slice live and a workspace created + When WS AUTH / WS LIST run from 12 fresh connections, then the server + restarts with appendonly yes + Then every connection sees the workspace before AND after restart + (WAL replay lands in the global registry) + And ShardSlice no longer carries a workspace_registry field + +Scenario: ssm4a AOF rewrite without foreign locks (M4) + Given a 4-shard server with slice live, appendonly yes, and 10k keys written + When BGREWRITEAOF runs while writes continue, then the server restarts + Then the rewrite completes, no aof-writer thread touches another thread's + Database directly (code-shape pin: the rewrite fold takes a + shard-produced frozen snapshot), and all keys survive the restart + +Scenario: ssm4b metrics and MEMORY DOCTOR are lock-free observers (M4) + Given a 4-shard server with slice live under active write traffic + When the Prometheus endpoint is scraped and MEMORY DOCTOR runs repeatedly + Then both return plausible non-zero memory figures sourced from published + per-shard atomics + And neither path contains a read_db/vector_store/graph_store lock acquisition + (code-shape pin) + +Scenario: ssm5 the wrappers are gone (M5) + Given the migrated codebase + When grepping ShardDatabases for RwLock / Mutex / + Mutex / RwLock / per-shard registry Mutexes + and for `slice::is_initialized()` + Then zero wrapper fields remain, zero dual-branch gates remain in + production code (the gate fn may survive only for tests), and all + ~30 test construct sites compile against the new shape + And uring_handler accesses shard state via with_shard, not write_db + +Scenario: ssm6 full matrix green (M6) + Given the migrated codebase + When the full verification matrix runs (lib tests both runtimes macOS+VM, + integration suites with MOON_BIN pinned, loom, fmt, clippy both CI + feature sets, consistency sweep) + Then everything passes with zero skipped/weakened tests + +Scenario: ssm7 perf bar (M7) + Given bench-swf on an idle Linux VM (load < 4), same REQS as the baseline + When slice-live HEAD is compared cell-by-cell against current main + Then no cell regresses beyond run-to-run noise, and the s1/s4 P1/P16 + deltas are recorded in §6 as the milestone scaling evidence + +Scenario: reject uninitialized_shard + Given a shard event loop entered without init_shard (constructible only in + a test harness after M1) + When the loop would begin serving + Then the process fails fast at startup with the shard id in the error + And no command is ever answered from an uninitialized thread + (no per-command with_shard panic reachable over the wire) + +Scenario: reject slice_reentrancy + Given the migrated handlers + When with_shard is entered while another with_shard borrow is live + (the slice.rs unit pin) + Then it panics with the named recursive-borrow message in tests + And no production call path constructs that nesting (review + the cross- + shard hops all exit the closure before sending) + +Scenario: reject slice_send_violation + Given ShardSlice's _not_send marker + When any code attempts to move/share a slice across threads + Then compilation fails (the marker stays; no Send/Sync impl, no unsafe + added anywhere in the migration) + +Scenario: reject test_weakening + Given the frozen suites (wire_reachability_red, cross_shard_consistency_red, + scripts/test-consistency.sh assertions) + When the migration lands + Then their assertions are byte-identical to pre-task state (diff review) + +Scenario: reject wrapper_resurrection + Given M5 is complete + When any later commit in this task adds a lock around per-shard state back + into ShardDatabases or a helper reintroducing cross-thread access + Then the code-shape pin (ssm5 grep test) fails CI + +Scenario: reject borrow_across_await + Given the migrated handlers + When reviewing every with_shard/with_shard_db call site + Then no closure spans an .await (structurally impossible — the closure is + sync — but verified: no async block captures the guard-like borrow) +``` + + + + + +--- + +## 3 · CONTRACT — freeze the shape ▸ docs/05-step-3-contract.md + +No wire-protocol change: every RESP command behaves byte-identically (the cdg +suite + scripts/test-consistency.sh assertions are the wire contract and stay +byte-identical — reject test_weakening). The frozen shape here is the +*internal* surface: boot handoff, new SPSC messages, the residual shared +struct, and the AOF cooperative-snapshot protocol. + +### C1 — Boot handoff (M1) + +```rust +// src/shard/shared_databases.rs — construction now RETURNS the per-shard +// slice packages instead of swallowing the per-shard state behind locks. +impl ShardDatabases { + /// Same argument list as today. WAL/AOF recovery replays into the + /// Databases BEFORE this call (single-threaded boot — no locks needed). + pub fn new(/* unchanged args */) -> (ShardDatabases, Vec); +} +``` + +- `main.rs` and `embedded.rs` destructure; each shard thread's spawn closure + receives its `ShardSliceInit` BY MOVE; the FIRST statement on the shard + thread (both runtimes, before runtime block-on / first accept) is + `crate::shard::slice::init_shard(init)`. +- The event loop asserts `slice::is_initialized()` once before entering its + accept/drain loop and aborts the process with the shard id otherwise + ("uninitialized_shard" → startup abort, never a per-command panic). +- The ~30 test/bench construct sites destructure the new tuple; test harnesses + that need lock-style access are migrated to drive a real shard thread or use + the slice directly on the test thread. + +### C2 — New ShardMessage variants (src/shard/dispatch.rs) + +```rust +/// MQ domain hop (GraphCommand precedent): conn handler computes +/// owner = key_to_shard(effective_queue_key); owner executes the FULL arm +/// (durable flag, DLQ stream creation, trigger debounce) against its slice. +MqCommand(Box), +pub struct MqCommandPayload { + pub db_index: usize, + /// Workspace prefix ("{wsid}:" or empty) — owner derives effective keys + /// exactly as the inline arm does today (incl. DLQ sibling keys). + pub key_prefix: bytes::Bytes, + /// Full original MQ.* frame (CREATE/PUSH/POP/ACK/DLQLEN/TRIGGER). + pub command: std::sync::Arc, + pub reply_tx: channel::OneshotSender, +} + +/// TXN.COMMIT MQ-intent materialize hop. Sender groups txn.mq_intents by +/// owner shard; owner applies get_stream_mut → durable check → add, exactly +/// the txn.rs:160–167 fold. MqIntent is the existing type +/// (src/transaction/mod.rs:84 — queue_key + fields, Clone). +MqTxnMaterialize { + db_index: usize, + intents: Vec, + reply_tx: channel::OneshotSender<()>, +}, + +/// WS.DROP best-effort key cleanup. prefix = "{}:"; +/// owner = key_to_shard(prefix); db pinned to 0 (current behavior, both +/// branches). Reply = deleted count (logged, not surfaced to the client). +WsDropCleanup { + prefix: bytes::Bytes, + reply_tx: channel::OneshotSender, +}, + +/// AOF rewrite cooperative snapshot (C4). Shard builds the expired-filtered +/// fold of ALL its dbs and replies; the writer thread never touches shard +/// state. +AofFold { + reply_tx: channel::OneshotSender, +}, +/// (entries, base_timestamp) per db index — same fold shape rewrite.rs +/// builds today (feeds rdb::save_snapshot_to_bytes unchanged). +pub struct AofFoldSnapshot { + pub dbs: Vec<(Vec<(CompactKey, Entry)>, u32)>, +} +``` + +Senders: conn handlers use the existing mesh producers (`ctx.dispatch_tx`); +the aof-writer threads get dedicated external mesh producers — the SAME +mechanism replication's master already uses (`shard_producers`, +replication/master.rs:123–129). Enum-size discipline: boxed where the inline +payload would breach the 256-byte cap (MqCommand); the others fit inline. + +### C3 — Global workspace registry (M3) + +```rust +// src/shard/shared_databases.rs — ONE process-global registry replaces the +// per-shard Box<[Mutex<…>]> array (slot-0 trick retired). +pub workspace_registry: parking_lot::Mutex>>, +/// Accessor keeps PR #173's no-param signature — call sites unchanged. +pub fn workspace_registry(&self) -> MutexGuard<'_, Option>>; +``` + +- `workspace_registry` field LEAVES `ShardSlice` + `ShardSliceInit`. +- WS WAL records keep the shard-0 stream via `wal_append(0, …)` (senders are + Send; unchanged). WAL replay populates the global registry at boot. + +### C4 — AOF cooperative-snapshot protocol (M4, replaces fold phases 2–5) + +1. Writer phase-1 drain → fsync → ack (unchanged). +2. Writer pushes `AofFold` into its shard's external mesh producer, then + BLOCKS on the oneshot reply (aof-writer-N is a plain OS thread). +3. Shard event loop processes `AofFold` atomically BETWEEN commands: builds + the snapshot, replies. The single-threaded loop replaces RwLock mutual + exclusion. +4. **Exactly-once invariant transfer** (the load-bearing line): every command + the shard processed before the fold enqueued its AOF append + (`try_send_append`) before the fold reply was sent — happens-before — so + the writer's post-reply mid-drain captures ALL pre-snapshot appends into + the OLD incr; post-snapshot appends land in the NEW incr. The in-guard + append invariant (rewrite.rs:480–502) becomes the in-event-loop-order + invariant; that doc comment is rewritten to state this. +5. Phases 6–8 (base write, manifest advance, outcome barrier, ShardDoneGuard) + unchanged. +- Failure: mesh push fails or the oneshot drops (shard shutdown) → + `AofError::RewriteFailed` → existing abort path keeps the old generation. +- Known limitation (channel saturation during phase 6) is pre-existing and + carries over unchanged — explicitly NOT widened by this redesign. + +### C5 — Published store-memory atomics (M4) + +```rust +// src/shard/shared_databases.rs +pub struct ShardStoreMemory { + pub vector: AtomicUsize, + pub text: AtomicUsize, + pub graph: AtomicUsize, +} +pub store_memory_per_shard: Box<[Arc]>, +``` + +- `ShardSliceInit`/`ShardSlice` gain `store_memory: Arc`; + the shard refreshes it on the existing 100ms persistence/elastic tick. +- Prometheus publisher (metrics_setup.rs:1396–1429), MEMORY DOCTOR + (server_admin.rs:385–416), and the two persistence_tick.rs leftovers + (344, 499 → read_memory_sum) read ONLY published atomics — zero lock + acquisitions (ssm4b code-shape pin). Figures may lag ≤1 tick (observability + paths; acceptable, documented at the readers). +- `memory_per_shard` (KV) and `elastic_budgets` keep their existing shapes. + +### C6 — Residual ShardDatabases (M5) + +KEPT: `wal_append_txs`, `num_shards`, `db_count`, `memory_per_shard`, +`elastic_budgets`, `workspace_registry` (C3, global), `store_memory_per_shard` +(C5). DELETED fields: `shards`, `vector_stores`, `text_stores`, +`graph_stores`, `temporal_registries`, `temporal_kv_indexes`, +`workspace_registries`, `durable_queue_registries`, `trigger_registries`, +`kv_write_intents`, `deferred_hnsw_inserts`. DELETED methods: `write_db`, +`read_db`, `all_shard_dbs`, per-shard store/registry accessors, +`aggregate_memory`. `SwapDb` handling moves onto the slice. ~130 +`is_initialized()` dual branches collapse to the slice arm; the gate fn may +survive only for tests (ssm5 grep pin). + +### Contracted responses for §1 rejects + +| Reject | Contracted response | +|---|---| +| uninitialized_shard | startup abort with shard id (C1 assert) before first accept; no wire-reachable with_shard panic | +| slice_reentrancy | RefCell double-borrow panic (existing slice.rs unit pin); hops exit the closure before sending — review + ssm-reject test | +| slice_send_violation | compile error — `PhantomData>` stays; zero new `unsafe`/`Send` impls (UNSAFE_POLICY) | +| test_weakening | §6 diff gate: frozen-suite assertions byte-identical pre/post | +| wrapper_resurrection | code-shape test (`tests/shardslice_shape.rs` grep pins from ssm5) fails on any reintroduced wrapper/gate | +| borrow_across_await | closures are sync by type; grep pin: no `with_shard` call inside an async block capturing the borrow | + +Names follow existing glossary: ShardMessage, ShardSliceInit, init_shard, +with_shard, key_to_shard. New names introduced here: MqCommandPayload, +MqTxnMaterialize, WsDropCleanup, AofFold, AofFoldSnapshot, ShardStoreMemory, +store_memory_per_shard. + +Least-sure flag surfaced at freeze: +⚠ [contract] C4's AOF cooperative-snapshot redesign is containable in this + task — because it moves a durability fold across thread ownership and the + exactly-once proof rests on a new ordering argument (append enqueued before + fold-reply ⇒ mid-drain captures it); if wrong: data-loss-grade bug or M4 + splits out, blocking M5 wrapper deletion. +⚠ [spec] The three scouts' cross-thread-consumer inventory is complete + (~130 gated sites + 3 named consumers, nothing hidden) — because 213 call + sites were migrated across six earlier waves by different hands; if wrong: + runtime thread-local panic or silent wrong-shard data surfacing only under + load. + +Status: FROZEN @ v1 — approved by Tin Dang 2026-06-12 (both ⚠ flags accepted: +AOF cooperative-snapshot containable in-task; scout inventory complete). + + +--- + +## 4 · TESTS — failing-first suite (red) ▸ docs/06-step-4-tests.md + +Coverage target: every scenario ssm1–ssm7 + 6 rejects has exactly one test; the +frozen suites (cdg, consistency sweep) are reused untouched as ssm2's oracle. +Plan (one test per scenario, asserting behavior not internals): + + - test_ssm1_slice_live_at_startup (`tests/shardslice_live.rs`, NEW): spawn a + real 4-shard server (MOON_BIN pinned, fresh --dir), assert one + "ShardSlice initialized" log line per shard BEFORE the port accepts, then + PING/SET/GET from fresh conns. Runs under both runtimes (feature matrix). + - test_ssm2_owner_routing: REUSED oracles — cdg6 a–g + (tests/cross_shard_consistency_red.rs, byte-identical) + consistency + sweep 197/197 at 1/4/12. Red today in the sense that they pass via the + LOCK path; post-cutover they must pass via the slice path — the + cutover-specific red signal is ssm5's shape test. + - test_ssm3_ws_global_registry (`tests/shardslice_live.rs`): WS CREATE → + AUTH/LIST from 12 fresh conns + bound/unbound key visibility (GREEN + oracle). Restart half split into its own RED test (below). Shape half + (no workspace_registry field in slice) lives in the shape test. + - test_ssm3_ws_registry_survives_restart (`tests/shardslice_live.rs`, RED): + WS CREATE with appendonly yes → restart same dir → WS LIST must still + show it. RED for a CONFIRMED pre-existing bug found while writing the + suite: replay_workspace_wal (shared_databases.rs:289) scans shard-{id}/ + but WAL v3 segments live in shard-{id}/wal-v3/ (recovery.rs:361) — zero + workspace records ever replay. Build fixes it under frozen C3 ("WAL + replay populates the global registry at boot"). + - test_ssm4a_aof_fold_cooperative (`tests/shardslice_live.rs`): appendonly + yes, 2000 keys, BGREWRITEAOF under a concurrent INCR writer, exactly-once + (ctr == recorded successful-INCR count), restart survival. Two variants + of one helper: shards=1 (default-gated path) and shards=4 + + --experimental-per-shard-rewrite (test_ssm4a_fold_4shard_experimental — + drives do_rewrite_per_shard, the EXACT fold C4 redesigns). Both GREEN + oracles that must stay green through the cutover. + - test_ssm4b_observer_lockfree (`tests/shardslice_shape.rs`): grep pins — + metrics_setup/server_admin/persistence_tick contain no + read_db/write_db/all_shard_dbs/store-lock calls; runtime probe: scrape + + MEMORY DOCTOR return non-zero figures under load. + - test_ssm5_wrappers_gone (`tests/shardslice_shape.rs`, NEW — the PRIMARY + RED TEST): asserts ShardDatabases source has zero + RwLock/Mutex/Mutex/RwLock/ + per-shard registry Mutex fields, zero `is_initialized()` gates in + src/server+src/command+src/shard production paths, and uring_handler + uses with_shard not write_db. RED today on every assertion. + - test_ssm6_full_matrix: not a test file — the §6 verify checklist (lib + tests both runtimes macOS+VM, integration MOON_BIN pinned, loom, fmt, + clippy both feature sets, sweep ×3). + - test_ssm7_perf_bar: not a test file — bench-swf cell-by-cell vs main on + idle VM, recorded in §6. + - test_reject_uninitialized_shard (`tests/shardslice_live.rs` + unit): + event-loop entry without init aborts at startup naming the shard id + (unit-level: the assert fn; the wire-level absence is covered by ssm1). + - test_reject_slice_reentrancy: existing slice.rs double-borrow unit pin + stays byte-identical + grep pin: no nested with_shard in src. + - test_reject_slice_send_violation: compile-fail doc-test on ShardSlice + (already in slice.rs — kept) + audit-unsafe.sh unchanged (zero new + unsafe). + - test_reject_test_weakening: §6 diff gate — git diff of frozen suites + empty. + - test_reject_wrapper_resurrection: same shape test as ssm5 (it IS the + permanent guard). + - test_reject_borrow_across_await (`tests/shardslice_shape.rs`): grep pin — + no `.await` inside any with_shard/with_shard_db closure body. + + +Tests live in: `tests/shardslice_shape.rs` · `tests/shardslice_live.rs` · +reused frozen oracles in `tests/cross_shard_consistency_red.rs` + +`scripts/test-consistency.sh` (byte-identical). MUST run red +(shape test fails on every wrapper-field assertion; live tests fail on the +missing init log) before Build. + +RED-SUITE RECORD (verified 2026-06-12, macOS, MOON_BIN=target/debug/moon): +- shape: 4 RED / 1 GREEN — test_ssm5_wrappers_gone + + test_reject_wrapper_resurrection (RwLock at + shared_databases.rs:24,747 etc.), test_ssm4b_observer_lockfree + (metrics_setup.rs:1411, server_admin.rs:398, persistence_tick.rs ×7), + test_ssm3_shape_no_ws_field_in_slice (slice.rs:87,128,154); + test_reject_borrow_across_await GREEN. +- live: 2 RED / 4 GREEN — test_ssm1_slice_live_at_startup (no "ShardSlice + initialized" marker), test_ssm3_ws_registry_survives_restart (wal-v3 + replay-dir bug); GREEN: ssm3 cross-conn oracle, ssm4a ×2 (1-shard + + 4-shard experimental), reject wire guard. 3.2s wall. +- BUILD CONSTRAINT pinned in the shape test: C1's startup assert must be + `slice::assert_initialized(shard_id)` (helper INSIDE slice.rs) — a raw + `is_initialized()` call at the event loop would trip the ssm5 pin, which + stays strict (zero call sites outside slice.rs). + + + + +--- + +## 5 · BUILD — AI writes code ▸ docs/07-step-5-build.md + +Safety rule (feature-specific): AOF fold exactly-once — every acked write +survives restart exactly once; no append may straddle the snapshot boundary +into the wrong incr generation. +Code lives in: `./src/` +Constraints: do NOT change any test or the contract; allow-list packages only; ask if unclear. + +### BUILD RECORD (2026-06-12, branch feat/shardslice-migration) + +Waves (parallel Sonnet subagents, orchestrator-reviewed; commits in order): +- f31f6cd bundle · 1e9dc8d A2 WS-WAL replay fix · 498fc19 A1 dispatch + variants · 7d2f271 B owner-routing · 4de32a8 E1 gate collapse · + ee7da49 fold bounded-drain · 6f32f90 E2 structural cutover (C1+C6) · + 8b2d87c C4 exact fold boundary + H1 fsync barrier. +- NOTE: seven foreign HYBRID FILTER commits (4e9c03c..e969cf0) from a + concurrent session are interleaved on this branch + their uncommitted + edits in tree; left untouched (not this task's work). + +Build-phase findings fixed beyond plan (each caught by frozen oracles or +orchestrator review, never by weakening a test): +1. Slice re-entrancy (frozen reject hit live): drain_spsc_shared was + wrapped in with_shard → BorrowMutError; arms now take flat borrows. +2. AOF fold deadlock: fold channels existed but were never wired in + main.rs; unwired pools now abort rewrites cleanly. +3. ±1 double-apply (ssm4a flake): cross-shard PipelineBatch[Slotted] + AOF appends were handler-deferred → escaped pending_aof_count → + replayed onto a base already containing them. Appends moved into the + SPSC arms (before response-slot fill). +4. H1 regression from (3): appendfsync=always lost its fsync ack for + cross-shard writes → new AofWriterPool::fsync_barrier (zero-length + AppendSync rendezvous, F2-bounded), one per target shard per batch. +5. P0 in (4) caught in orchestrator review: a len=0 framed record would + be read as corruption by replay_incr_framed and brick boot. Writers + + rewrite mid-drain now skip the disk write for barriers (fsync+ack + only); reader skips len=0 defensively. Regression tests added. +6. Pre-existing crash-matrix harness flake (unique_port()+1 SO_REUSEPORT + collision between the two parallel tests): +1 dropped. +7. Throttle a fix-iteration added to the frozen aof_fold_exactly_once + writer loop was REVERTED (test-weakening); tight loop passes. +8. E2's test migration missed 29 call sites in 28 tokio-gated integration + tests (never compiled under default monoio features → invisible to all + macOS runs). All migrated to the tuple API + slice_init run() arg; + compile-verified on both feature sets. +9. Slice re-entrancy #2, found by the first-ever post-cutover run of the + tokio integration suites: the cross-store TXN undo-capture re-entered + with_shard from inside the write closure (handler_sharded via + with_shard_db; handler_monoio nested inside its own with_shard) → + BorrowMutError panic on any TXN write. Fixed with disjoint NLL field + borrows (s.kv_write_intents directly); handler_sharded's do_write now + takes the whole ShardSlice. txn_kv_wiring 12/12 under tokio + (MOON_BIN pinned to the Linux binary per the Mach-O host-proxy trap). + + + +--- + +## 6 · VERIFY — evidence + non-functional review ▸ docs/08-step-6-verify.md + +- [x] all tests pass — EVIDENCE (2026-06-12, eb5d664): + - frozen oracles: shape 5/5 · live 6/6 serial (ssm1 marker+ids, ssm3 ×2, + ssm4a 1-shard + 4-shard experimental, wire guard) · ssm4a 10/10 serial + repeat at full INCR rate · cdg 7/7 + - lib: 3594 (monoio, macOS) · 2955+ (tokio; full VM integration sweep + all targets green, MOON_BIN=target-linux/debug/moon) + - crash_matrix_per_shard_bgrewriteaof 2/2 ×5 parallel + ×3 serial + - loom_response_slot 4/4 + - clippy -D warnings: macOS ×2 featuresets + VM tokio (incl. the + never-before-compiled cfg(linux+tokio) uring_handler) · fmt clean + - consistency sweep: 197/197 PASSED @ shards=1/4/12 (VM-local clone of + eb5d664, release build, 2026-06-12 23:49 — incl. TXN/GRAPH/FT/WS/MQ + cross-shard groups) +- [x] coverage did not decrease — red suite all flipped green; +4 new + regression tests (len-0 barrier reader/drain, H1-BARRIER ×3 in pool.rs) +- [x] no test or contract was altered during build — §3 untouched since + freeze; one frozen-test VIOLATION (agent throttle in aof_fold_exactly_once) + was caught and REVERTED, tight loop green; crash-matrix port fix is + harness isolation (strengthens), not weakening +- [x] concurrency / timing of the risky operation is safe — C4 exactly-once + verified 10/10 serial under saturating INCR load; slice re-entrancy guard + caught both nesting defects (event-loop drain, TXN intent capture), both + fixed with flat/disjoint borrows; no lock held across .await + (shape test test_reject_borrow_across_await green) +- [x] no exposed secrets, injection openings, or unexpected dependencies — + no new deps; no new unsafe (PhantomData> design holds); + parser paths untouched except replay_incr_framed len=0 skip (defensive, + fuzz targets unaffected — no new decode surface) +- [x] layering & dependencies follow CONVENTIONS.md — slice access only via + with_shard/with_shard_db; is_initialized confined to slice.rs (shape pin) +- [x] a person reviewed and approved the change — Tin Dang, 2026-06-13 + (gate: RISK-ACCEPTED + follow-up; PR shape: clean cherry-picked branch; + PR creation approved) + +### M7 bench-swf evidence (2026-06-13, VM idle load<1, REQS=200k, fresh +### server per rep ×3, eb5d664 vs main 3e376a1, both VM-local release) +- s1/P1: SET 305k vs 294k · GET 316k vs 306k — parity (slice ≥ main) +- s1/P16: SET 1.48M vs 1.61M (−5..8%, borderline vs noise) · GET 2.64M + vs 2.64M — parity +- s4 vs main ROUTED (--cross-shard-fast-path off; architecture-fair): + P1 SET 196k vs 187k · P1 GET 187k vs 186k · P16 SET 1.77M vs 1.81M · + P16 GET 1.99M vs 1.78M (+12%) · c1 SET 33.5k vs 33.7k — parity or + slice ahead in every routed cell +- s4 vs main DEFAULT (fast path ON — user-visible): read cells regress + by mechanism: c1 GET ~190k → ~26-34k (−85%; 4µs lock-read → ~47µs SPSC + round trip) · P1 GET ~224k → ~187k (−16%). The cross-thread RwLock + read_db fast path is definitionally incompatible with shared-nothing + (deleting those locks IS this task). This materializes the ⚠ flag + accepted at freeze; strictly exceeds the ssm7 "no cell beyond noise" + bar → escalated to the human gate rather than auto-passed. + Mitigations: {hash-tag} co-location (zero hops), --shards 1 guidance + (already CLAUDE.md), follow-up task candidate: slice-native cross-shard + read acceleration (read batching / snapshot reads) without locks. +- Methodology note: the first back-to-back run (both binaries, all cells + sequentially against long-lived servers) showed phantom −20..−40% on + s4 pipelined cells; fresh-server-per-rep ×3 showed parity. Recorded so + future benches use the per-rep protocol. + +### Deep checks — do not skim (fill the path that applies; the resolver judges which) +- [x] WIRING (code) — fsync_barrier: called from both handlers' cross-shard + batch paths (handler_sharded/mod.rs:1667, handler_monoio/mod.rs:1828), + policy-gated in pool.rs:313; fold channels: create_aof_fold_channels (mesh) → main.rs + wiring → pool.set_fold_channels → spsc AofFold arm → rewrite recv; + slice_init: main.rs/embedded.rs destructure → Shard::run → init_shard + (28 test harnesses migrated to same shape) +- [x] DEAD-CODE (code) — uring_handler shard_id underscored w/ comment; + E2 removed Cold variant + wrapper methods wholesale; clippy -D warnings + clean on all three featureset×OS combos confirms no orphans +- [x] SEMANTIC (prose / non-code) — fsync_barrier doc corrected during + review (claimed "reader skips len=0" before the writer-side skip existed; + doc now matches the implemented mechanism) + +### GATE RECORD +Outcome: RISK-ACCEPTED +Risk: cross-shard read fast-path cells regress vs main's default config +(single-client cross-shard GET 4µs → ~47µs / −85%; P1 GET −16%) — the +cross-thread RwLock read_db path is definitionally incompatible with +shared-nothing. All routed-vs-routed cells parity or better; s1 parity. +Non-security. Materializes the ⚠ flag accepted at the §3 freeze. +If RISK-ACCEPTED -> owner: Tin Dang · ticket: follow-up task +"cross-shard read acceleration (lock-free)" to be opened at observe · +expires: next perf milestone (M7 successor) +Reviewed by: Tin Dang (gate decision via session, "RISK-ACCEPTED + +follow-up") · date: 2026-06-13 + + + +--- + +## 7 · OBSERVE — feed the next loop ▸ docs/09-the-loop.md + +Watch (reuse scenarios as monitors): +Spec delta for the next loop: + +### Competency deltas +What did this loop teach the foundation? One line each, tagged by competency +(`DDD · SDD · UDD · TDD · ADD`), status `open`, with evidence. See the `add` skill's `deltas.md`. + diff --git a/CHANGELOG.md b/CHANGELOG.md index 09be9f92..af0a9684 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed — Shared-nothing thread-local shard storage (PR #175) + +Shard storage moved from `RwLock`/`Mutex`-wrapped shared databases to +thread-local **ShardSlices** owned exclusively by each shard's event +loop, on both runtimes (monoio and tokio). The lock wrappers are +deleted; all cross-shard access routes through the owner shard's SPSC +ring. Routed multi-shard workloads are parity or better (up to +12% on +pipelined GET at 4 shards); the cross-shard *read fast-path* (direct +RwLock read of another shard's data) is definitionally gone and those +cells regress — accepted with a follow-up task for lock-free cross-shard +read acceleration. BGREWRITEAOF on every layout now uses the C4 +cooperative-fold protocol: the AOF writer requests an atomic snapshot +from the owning shard and drains pre-snapshot appends with an exact +bound, so concurrent writes are neither lost nor double-applied. This +also fixes BGREWRITEAOF being silently inoperative on the top-level AOF +layout (`--shards 1`), where the rewrite previously errored internally +while the client was told it had started. + ### Added — HYBRID FT.SEARCH `FILTER` push-down on both branches (PR #174) `FT.SEARCH ... HYBRID` now accepts an optional `FILTER` clause that is diff --git a/src/admin/metrics_setup.rs b/src/admin/metrics_setup.rs index cde30fce..7525eea5 100644 --- a/src/admin/metrics_setup.rs +++ b/src/admin/metrics_setup.rs @@ -1391,41 +1391,31 @@ pub fn spawn_moon_memory_publisher() { /// Collect per-subsystem resident bytes and emit all 7 /// `moon_memory_bytes{kind=...}` series plus `moon_rss_bytes`. /// -/// Called every 15 s by `spawn_moon_memory_publisher`. Allocation-free in -/// the steady state — all label values are `&'static str`. +/// Called every 15 s by `spawn_moon_memory_publisher`. Lock-free: reads from +/// per-shard published atomics (C5 / M4). Figures lag at most one 100ms tick. fn update_moon_memory_bytes() { + use std::sync::atomic::Ordering; + let rss = get_rss_bytes() as usize; let mut dashtable: usize = 0; let mut hnsw: usize = 0; - let mut sealed: usize = 0; - #[cfg_attr(not(feature = "graph"), allow(unused_mut))] + let sealed: usize = 0; // combined into hnsw from vector atomic (C5) let mut csr: usize = 0; let wal: usize = 0; // WalWriterV3 is stack-owned; not reachable here let mut backlog: usize = 0; if let Some(shard_dbs) = get_global_shard_databases() { - let num_shards = shard_dbs.num_shards(); - for shard_id in 0..num_shards { - // Database + DashTable (DB 0 — the hot database) - let db_guard = shard_dbs.read_db(shard_id, 0); - dashtable += db_guard.resident_bytes(); - dashtable += db_guard.data().resident_bytes(); - drop(db_guard); - - // VectorStore: (mutable/hnsw, immutable/sealed) - let vs = shard_dbs.vector_store(shard_id); - let (m, i) = vs.resident_bytes(); - hnsw += m; - sealed += i; - drop(vs); - - // GraphStore (CSR) - #[cfg(feature = "graph")] - { - let gs = shard_dbs.graph_store_read(shard_id); - csr += gs.resident_bytes(); - } + // KV memory: sum of per-shard published atomics. Lock-free. + // C5 / M4: `read_memory_sum()` replaces per-shard `read_db(…)` locks. + dashtable = shard_dbs.read_memory_sum(); + + // Store memory: sum published per-shard vector/graph atomics. + // Values are refreshed by each shard's 100ms tick (publish_store_memory). + for mem in shard_dbs.store_memory_per_shard.iter() { + hnsw += mem.vector.load(Ordering::Relaxed); + // graph is cfg-gated at publish time; the atomic is always present. + csr += mem.graph.load(Ordering::Relaxed); } } diff --git a/src/command/persistence.rs b/src/command/persistence.rs index 6a80f049..dfb390ce 100644 --- a/src/command/persistence.rs +++ b/src/command/persistence.rs @@ -441,7 +441,7 @@ mod tests { // Wrap as a TopLevel pool to match the post-2e-β helper signature. let (tx, _rx) = crate::runtime::channel::mpsc_bounded::(1); let pool = AofWriterPool::top_level(tx); - let shard_dbs = crate::shard::shared_databases::ShardDatabases::new(vec![vec![ + let (shard_dbs, _inits) = crate::shard::shared_databases::ShardDatabases::new(vec![vec![ crate::storage::Database::new(), ]]); @@ -505,7 +505,7 @@ mod tests { let _guard = GATE_TEST_LOCK.lock(); let (tx, _rx) = crate::runtime::channel::mpsc_bounded::(1); let pool = AofWriterPool::top_level(tx); - let shard_dbs = crate::shard::shared_databases::ShardDatabases::new(vec![vec![ + let (shard_dbs, _inits) = crate::shard::shared_databases::ShardDatabases::new(vec![vec![ crate::storage::Database::new(), ]]); diff --git a/src/command/server_admin.rs b/src/command/server_admin.rs index 5f5e4c58..c6166847 100644 --- a/src/command/server_admin.rs +++ b/src/command/server_admin.rs @@ -383,37 +383,34 @@ fn memory_doctor() -> Frame { let rss = crate::admin::metrics_setup::get_rss_bytes() as usize; let vsz = get_vsz_bytes(); - // ── Gather per-subsystem resident bytes ────────────────────────────── - let mut dashtable_bytes: usize = 0; - let mut hnsw_bytes: usize = 0; - let mut sealed_bytes: usize = 0; - #[cfg_attr(not(feature = "graph"), allow(unused_mut))] - let mut csr_bytes: usize = 0; + // ── Gather per-subsystem resident bytes (C5 / M4 — lock-free atomics) ── + // KV and store memory are read from published per-shard atomics. Figures + // lag at most one 100ms tick — acceptable for an on-demand diagnostic. + use std::sync::atomic::Ordering; + let dashtable_bytes: usize; + let hnsw_bytes: usize; + let sealed_bytes: usize = 0; // combined into hnsw_bytes from vector atomic + #[cfg_attr(not(feature = "graph"), allow(unused_variables))] + let csr_bytes: usize; let wal_bytes: usize = 0; if let Some(shard_dbs) = crate::admin::metrics_setup::get_global_shard_databases() { - let num_shards = shard_dbs.num_shards(); - for shard_id in 0..num_shards { - // Database + DashTable (DB 0 only — the hot database) - let db_guard = shard_dbs.read_db(shard_id, 0); - dashtable_bytes += db_guard.resident_bytes(); - dashtable_bytes += db_guard.data().resident_bytes(); - drop(db_guard); - - // VectorStore: (mutable/hnsw, immutable/sealed) - let vs = shard_dbs.vector_store(shard_id); - let (mutable, immutable) = vs.resident_bytes(); - hnsw_bytes += mutable; - sealed_bytes += immutable; - drop(vs); - - // GraphStore (CSR) - #[cfg(feature = "graph")] - { - let gs = shard_dbs.graph_store_read(shard_id); - csr_bytes += gs.resident_bytes(); - } + // KV memory: sum of per-shard published atomics. Lock-free. + dashtable_bytes = shard_dbs.read_memory_sum(); + + // Store memory: sum published per-shard vector/graph atomics. + let mut vec_total = 0usize; + let mut csr_total = 0usize; + for mem in shard_dbs.store_memory_per_shard.iter() { + vec_total += mem.vector.load(Ordering::Relaxed); + csr_total += mem.graph.load(Ordering::Relaxed); } + hnsw_bytes = vec_total; + csr_bytes = csr_total; + } else { + dashtable_bytes = 0; + hnsw_bytes = 0; + csr_bytes = 0; } // Replication backlog via global state (same pattern as INFO replication). diff --git a/src/main.rs b/src/main.rs index 29a8c59e..9545e758 100644 --- a/src/main.rs +++ b/src/main.rs @@ -59,7 +59,7 @@ use moon::runtime::channel; use moon::runtime::{RuntimeFactoryImpl, traits::RuntimeFactory}; use moon::server; use moon::shard::Shard; -use moon::shard::mesh::{CHANNEL_BUFFER_SIZE, ChannelMesh}; +use moon::shard::mesh::{CHANNEL_BUFFER_SIZE, ChannelMesh, create_aof_fold_channels}; use moon::shard::shared_databases::ShardDatabases; use tracing::info; @@ -601,7 +601,19 @@ fn main() -> anyhow::Result<()> { std::process::exit(2); } - let aof_pool: Option> = if config.appendonly == "yes" { + // C4 TopLevel staging variables: populated in the `else` branch of the pool + // construction block when layout==TopLevel, consumed in the fold-channel + // wiring block and the shard-spawn loop below. + let mut toplevel_fold_consumer_staging: Option< + ringbuf::HeapCons, + > = None; + let mut toplevel_pool_fold_producer: Option< + std::sync::Arc>>, + > = None; + let mut toplevel_pool_fold_notifier: Option> = + None; + + let mut aof_pool: Option> = if config.appendonly == "yes" { let fsync = FsyncPolicy::from_str(&config.appendfsync); // PerShard writers required when num_shards >= 2 AND we'll have a // PerShard manifest at runtime. Two cases produce PerShard: @@ -657,6 +669,20 @@ fn main() -> anyhow::Result<()> { let (tx, rx) = channel::mpsc_bounded::(10_000); let aof_token = cancel_token.child_token(); let aof_file_path = PathBuf::from(&config.dir).join(&config.appendfilename); + + // C4 TopLevel: create fold channels for shard 0 here (before the + // writer spawn) so they can be moved into the closure. The consumer + // is stored in `toplevel_fold_cons_staging` and later merged into shard + // 0's consumers vec in the shard-spawn loop below. + // Notifier: use shard 0's mesh notifier (same handle the shard event + // loop is woken by) so the AofFold push wakes the right thread. + let (mut tl_fold_producers, mut tl_fold_consumers) = create_aof_fold_channels(1, 4); + let tl_fold_producer = tl_fold_producers.remove(0); + let tl_fold_consumer = tl_fold_consumers.remove(0); + let tl_fold_notifier = mesh.take_notify(0); + // fold channels for the writer task + let writer_fold_channels = Some((tl_fold_producer.clone(), tl_fold_notifier.clone())); + // Legacy single-writer thread. Each shard clones the outer // `aof_pool` Arc; sender lifetime is governed by the pool's Drop. std::thread::Builder::new() @@ -664,11 +690,23 @@ fn main() -> anyhow::Result<()> { .spawn(move || { RuntimeFactoryImpl::block_on_local( "aof-writer".to_string(), - aof::aof_writer_task(rx, aof_file_path, fsync, aof_token), + aof::aof_writer_task( + rx, + aof_file_path, + fsync, + aof_token, + writer_fold_channels, + ), ); }) .expect("failed to spawn AOF writer thread"); info!("AOF enabled (TopLevel, fsync: {:?})", fsync); + + // Stash consumer + pool-side channels for wiring in the blocks below. + toplevel_fold_consumer_staging = Some(tl_fold_consumer); + toplevel_pool_fold_producer = Some(tl_fold_producer); + toplevel_pool_fold_notifier = Some(tl_fold_notifier); + Some(AofWriterPool::top_level_with_policy( tx, fsync, @@ -787,6 +825,61 @@ fn main() -> anyhow::Result<()> { // Collect all notifiers before spawning shard threads let all_notifiers = mesh.all_notifiers(); + // C4: Wire AOF fold channels for the cooperative snapshot protocol. + // + // Two paths: + // + // A) PerShard layout (num_shards >= 2): fold channels are created here and + // wired into the pool via set_fold_channels. Arc::get_mut is safe because + // no shard threads have spawned yet (first Arc clone happens ~line 1290). + // fold_consumers[i] is merged into shard i's consumers vec before run(). + // + // B) TopLevel layout (--shards 1): fold channels were already created AND + // WIRED INTO THE WRITER TASK in the TopLevel writer-spawn branch above. + // The consumer side was stashed in `toplevel_fold_consumer_staging`. + // Here we only wire the pool's set_fold_channels and collect the consumer + // into fold_consumers so the shard-spawn loop can prepend it. + let mut fold_consumers: Option>> = + None; + if let Some(ref mut pool_arc) = aof_pool { + let layout = pool_arc.layout(); + if layout == moon::persistence::aof_manifest::AofLayout::PerShard { + // Path A: PerShard — create fold channels and wire now. + let (fold_producers, fold_cons) = create_aof_fold_channels(num_shards, 4); + // SAFETY: Arc::get_mut is valid here — no clones exist yet; the first + // shard_aof_pool clone happens inside the shard-spawn loop below. + if let Some(pool_mut) = std::sync::Arc::get_mut(pool_arc) { + pool_mut.set_fold_channels(fold_producers, all_notifiers.clone()); + fold_consumers = Some(fold_cons); + } else { + tracing::error!( + "C4 wiring (PerShard): Arc::get_mut failed — fold channels not wired. \ + BGREWRITEAOF will abort cleanly (old generation remains authoritative)." + ); + } + } else { + // Path B: TopLevel — channels were created + wired at writer-spawn time. + // Wire the producer/notifier into the pool and collect the consumer. + if let (Some(prod), Some(notif)) = ( + toplevel_pool_fold_producer.take(), + toplevel_pool_fold_notifier.take(), + ) { + // SAFETY: Arc::get_mut is valid here — no clones exist yet. + if let Some(pool_mut) = std::sync::Arc::get_mut(pool_arc) { + pool_mut.set_fold_channels(vec![prod], vec![notif]); + } else { + tracing::error!( + "C4 wiring (TopLevel): Arc::get_mut failed — pool fold channels \ + not wired. BGREWRITEAOF will abort cleanly." + ); + } + } + if let Some(cons) = toplevel_fold_consumer_staging.take() { + fold_consumers = Some(vec![cons]); + } + } + } + // Create admin SPSC channels for the console gateway (one per shard). #[cfg(feature = "console")] let mut admin_consumers = { @@ -1227,32 +1320,37 @@ fn main() -> anyhow::Result<()> { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = ShardDatabases::new(all_dbs); // Recover graph stores from persistence (CSR segments + metadata + WAL replay). + // These run on ShardSliceInit (mutate pre-shard state before handoff to threads). #[cfg(feature = "graph")] if let Some(ref dir) = persistence_dir { let dir_path = std::path::Path::new(dir); - shard_databases.recover_graph_stores(dir_path); - shard_databases.replay_graph_wal(dir_path); + moon::shard::shared_databases::recover_graph_stores(&mut slice_inits, dir_path); + moon::shard::shared_databases::replay_graph_wal( + &mut slice_inits, + dir_path, + config.databases, + ); } // Replay temporal WAL records (not gated on graph feature — temporal KV is core). if let Some(ref dir) = persistence_dir { let dir_path = std::path::Path::new(dir); - shard_databases.replay_temporal_wal(dir_path); + moon::shard::shared_databases::replay_temporal_wal(&mut slice_inits, dir_path); } - // Replay workspace WAL records (not gated on graph feature — workspaces are core). + // Replay workspace WAL records — uses shared registry, takes Arc. if let Some(ref dir) = persistence_dir { let dir_path = std::path::Path::new(dir); - shard_databases.replay_workspace_wal(dir_path); + moon::shard::shared_databases::replay_workspace_wal(&shard_databases, dir_path); } // Replay MQ WAL records (cursor-rollback for durable queues). if let Some(ref dir) = persistence_dir { let dir_path = std::path::Path::new(dir); - shard_databases.replay_mq_wal(dir_path); + moon::shard::shared_databases::replay_mq_wal(&mut slice_inits, dir_path); } // All shards recovered — mark server as ready for /readyz. @@ -1280,6 +1378,28 @@ fn main() -> anyhow::Result<()> { }); consumers.push(admin_cons); } + // C4: Prepend AOF fold consumer for this shard (aof-writer-N -> shard SPSC). + // The consumer receives ShardMessage::AofFold pushed by the per-shard AOF + // writer thread; the shard event loop processes it between commands and + // replies with a frozen snapshot. Only wired for PerShard pools. + // + // PRIORITY: the fold consumer is inserted at index 0 (before command + // consumers) so that AofFold is never starved by MAX_DRAIN_PER_CYCLE. + // Under sustained write load the command consumer can saturate the + // 256-message drain cap every cycle, which would prevent the fold + // consumer (appended last) from ever being reached — the AofFold reply + // is never sent, the per-shard writer blocks on recv_blocking() forever, + // and all other shards' writers wait on await_outcome() indefinitely + // (permanent deadlock until SIGKILL). Putting the fold consumer first + // costs ≤ 2 extra try_pop calls per cycle (ring capacity is 4) and + // eliminates the starvation window. + if let Some(ref mut fold_cons_vec) = fold_consumers { + let fold_cons = std::mem::replace(&mut fold_cons_vec[id], { + use ringbuf::traits::Split; + ringbuf::HeapRb::new(1).split().1 + }); + consumers.insert(0, fold_cons); + } let conn_rx = mesh.take_conn_rx(id); let shard_cancel = cancel_token.clone(); let shard_aof_pool = aof_pool.clone(); @@ -1299,6 +1419,8 @@ fn main() -> anyhow::Result<()> { let shard_pubsub_registries = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + // C1: hand the shard its ShardSlice initializer — consumed by init_shard inside run(). + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("shard-{}", id)) @@ -1340,6 +1462,7 @@ fn main() -> anyhow::Result<()> { shard_pubsub_registries, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ) .await; }); diff --git a/src/persistence/aof/mod.rs b/src/persistence/aof/mod.rs index 90d1cde0..b43b47f2 100644 --- a/src/persistence/aof/mod.rs +++ b/src/persistence/aof/mod.rs @@ -181,14 +181,25 @@ pub enum AofMessage { RewriteSharded(Arc), /// [F6] Trigger a per-shard AOF rewrite (compaction) in the PerShard /// layout. Sent to EVERY per-shard writer at once. Each writer folds its - /// own shard (drain → lock → snapshot → write new base+incr at the - /// coordinator's `new_seq` → reopen), then decrements the shared + /// own shard (drain → AofFold SPSC → snapshot → write new base+incr at + /// the coordinator's `new_seq` → reopen), then decrements the shared /// `PerShardRewriteCoord`; the last writer commits the manifest once /// (single seq flip) and prunes the old generation. The synchronized seq /// + single commit are what make multi-shard BGREWRITEAOF crash-safe. + /// + /// `fold_producer` / `fold_notifier` are per-shard SPSC handles used for + /// the C4 cooperative snapshot (ShardMessage::AofFold). Each writer + /// receives the producer/notifier for ITS own shard only. RewritePerShard { shard_dbs: Arc, coord: Arc, + /// SPSC producer into this shard's event-loop ring buffer. + /// Wrapped in `Mutex` so the AOF writer thread (not the shard thread) + /// can push the AofFold message safely. + fold_producer: + Arc>>, + /// Notifier that wakes the shard event loop after an SPSC push. + fold_notifier: Arc, }, /// Shut down the AOF writer task gracefully. Shutdown, diff --git a/src/persistence/aof/pool.rs b/src/persistence/aof/pool.rs index f97d34e5..f549d08b 100644 --- a/src/persistence/aof/pool.rs +++ b/src/persistence/aof/pool.rs @@ -26,6 +26,19 @@ pub struct AofWriterPool { /// the authoritative manifest fresh at rewrite time. `None` for TopLevel /// pools and test pools that never rewrite. base_dir: Option, + /// C4: per-shard SPSC fold producers for AofFold cooperative snapshot. + /// + /// `fold_producers[i]` is the SPSC producer into shard `i`'s event-loop + /// ring. Wrapped in `Arc` so the AOF writer thread (not the shard + /// thread) can push `ShardMessage::AofFold` safely. Set only for PerShard + /// pools that wire up fold channels; `None` for TopLevel pools and test + /// pools that use the legacy lock-based path or never rewrite. + fold_producers: Option< + Vec>>>, + >, + /// C4: per-shard Notify handles that wake the shard event loops after + /// an AofFold push. Parallel to `fold_producers`. + fold_notifiers: Option>>, } impl AofWriterPool { @@ -52,6 +65,8 @@ impl AofWriterPool { fsync_policy, fsync_timeout, base_dir: None, + fold_producers: None, + fold_notifiers: None, }) } @@ -81,6 +96,8 @@ impl AofWriterPool { fsync_policy, fsync_timeout, base_dir: None, + fold_producers: None, + fold_notifiers: None, }) } @@ -104,9 +121,105 @@ impl AofWriterPool { fsync_policy, fsync_timeout, base_dir: Some(base_dir), + fold_producers: None, + fold_notifiers: None, + }) + } + + /// C4: same as [`Self::per_shard_with_base_dir`] but also wires the per-shard + /// SPSC fold channels used by the cooperative AOF snapshot (ShardSlice path). + /// + /// `fold_producers[i]` MUST be a producer into shard `i`'s SPSC ring that has + /// a corresponding `HeapCons` already installed in the shard's consumer list — + /// created via [`crate::shard::mesh::create_aof_fold_channels`] and pushed into + /// each shard's `consumers` vec before `shard.run()` is called. + /// + /// The pool clones the `Arc`s; callers retain no live reference after passing. + pub fn per_shard_with_fold_channels( + senders: Vec>, + fsync_policy: FsyncPolicy, + fsync_timeout: Duration, + base_dir: PathBuf, + fold_producers: Vec< + Arc>>, + >, + fold_notifiers: Vec>, + ) -> Arc { + debug_assert!( + senders.len() >= 2, + "per_shard pool needs >=2 writers; use top_level for single-writer" + ); + debug_assert_eq!( + senders.len(), + fold_producers.len(), + "fold_producers must have one entry per shard" + ); + debug_assert_eq!( + senders.len(), + fold_notifiers.len(), + "fold_notifiers must have one entry per shard" + ); + Arc::new(Self { + senders, + layout: crate::persistence::aof_manifest::AofLayout::PerShard, + fsync_policy, + fsync_timeout, + base_dir: Some(base_dir), + fold_producers: Some(fold_producers), + fold_notifiers: Some(fold_notifiers), }) } + /// C4: Wire fold channels into a pool that was built with + /// [`Self::per_shard_with_base_dir`] (the production path in `main.rs`). + /// + /// `producers[i]` MUST be the SPSC producer into shard `i`'s ring and + /// `notifiers[i]` the matching Notify handle — both created by + /// [`crate::shard::mesh::create_aof_fold_channels`] and the corresponding + /// consumers already merged into shard `i`'s consumer vec before + /// `shard.run()` is called. + /// + /// **Idempotent-warn:** calling this a second time logs an error and + /// returns without overwriting the existing channels (double-set is a + /// startup-configuration bug, not a runtime error worth crashing over). + /// + /// **Panics (debug only):** if `producers.len() != self.senders.len()` or + /// `notifiers.len() != self.senders.len()` — the lengths must be identical + /// to the shard count established at construction time. + pub fn set_fold_channels( + &mut self, + producers: Vec< + Arc>>, + >, + notifiers: Vec>, + ) { + if self.fold_producers.is_some() || self.fold_notifiers.is_some() { + error!( + "set_fold_channels called twice on AofWriterPool (shard count {}). \ + The first set of channels is kept. This is a startup-configuration \ + bug — check the main.rs wiring.", + self.senders.len() + ); + return; + } + debug_assert_eq!( + producers.len(), + self.senders.len(), + "set_fold_channels: producers len {} != shard count {}", + producers.len(), + self.senders.len(), + ); + debug_assert_eq!( + notifiers.len(), + self.senders.len(), + "set_fold_channels: notifiers len {} != shard count {}", + notifiers.len(), + self.senders.len(), + ); + self.fold_producers = Some(producers); + self.fold_notifiers = Some(notifiers); + } + /// Returns the configured fsync policy. Hot-path callers read this to /// decide between the fast (`try_send_append`) and durable /// (`try_send_append_sync`) write paths. @@ -162,6 +275,59 @@ impl AofWriterPool { } } + /// Durability barrier for cross-shard pipelined writes under + /// `appendfsync=always` (H1 fix, C4-FOLD-FIX follow-up). + /// + /// **Why this exists:** the C4-FOLD-FIX moved AOF appends for cross-shard + /// writes into the SPSC arm (fire-and-forget `try_send_append`) so the + /// append is in the AOF channel *before* the response slot is filled — + /// required to keep AofFold's `pending_aof_count` accurate. The former + /// handler code that awaited `try_send_append_durable` was removed to + /// prevent double-apply on restart. However that removal also stripped the + /// H1 durability guarantee: under `appendfsync=always` the connection + /// handler was replying `+OK` without confirming fsync. + /// + /// **What this does:** after all cross-shard responses have been collected, + /// the handler calls this method ONCE per distinct target shard. Under + /// `Always` it enqueues a zero-length `AppendSync` into the target shard's + /// AOF writer channel and awaits the fsync ack with the same F2 + /// bounded-await as `try_send_append_durable`. Because the writer processes + /// messages in order, an acked barrier proves all prior `Append` messages + /// to that shard are also durable. Under `EverySec`/`No` it returns + /// `Ok(())` immediately (zero cost — one `match` on a field). + /// + /// **Zero-length AppendSync in the writer:** the per-shard writers (and + /// the rewrite mid-drain) recognize an empty payload as a barrier and + /// write NOTHING to disk — they fsync and ack only. This matters: a len=0 + /// framed header would otherwise be parsed by `replay_incr_framed` as a + /// corrupt entry (empty RESP payload → `Ok(None)` → RewriteFailed) and + /// brick the next boot. `replay_incr_framed` additionally skips len=0 + /// records as defense-in-depth. + /// + /// **Failure handling:** returns `Err(AofAck)` on the Always path if the + /// writer task has died, the channel is full, or the fsync timed out. + /// Callers MUST overwrite all cross-shard write responses with + /// `Frame::Error(AOF_FSYNC_ERR)` on `Err` — identical to the local-shard + /// durable path. + #[inline] + pub async fn fsync_barrier(&self, shard_id: usize) -> Result<(), AofAck> { + match self.fsync_policy { + FsyncPolicy::Always => { + // Enqueue a zero-length AppendSync. The writer will fsync all + // preceding Append messages (ordered channel) then ack Synced. + let rx = self.try_send_append_sync(shard_id, 0, Bytes::new()); + // F2 bounded await — same semantics as try_send_append_durable. + match Self::await_ack(rx, self.fsync_timeout).await { + AckOutcome::Ack(AofAck::Synced) => Ok(()), + AckOutcome::Ack(other) => Err(other), + AckOutcome::Disconnected => Err(AofAck::WriteFailed), + AckOutcome::TimedOut => Err(AofAck::FsyncFailed), + } + } + FsyncPolicy::EverySec | FsyncPolicy::No => Ok(()), + } + } + /// Await an `AppendSync` ack receiver under a bounded timeout (F2). /// /// `timeout == Duration::ZERO` preserves the legacy unbounded await @@ -236,9 +402,20 @@ impl AofWriterPool { /// replication-state handle pass 0. #[inline] pub fn try_send_append(&self, shard_id: usize, lsn: u64, bytes: Bytes) { - let _ = self + if let Err(e) = self .sender(shard_id) - .try_send(AofMessage::Append { lsn, bytes }); + .try_send(AofMessage::Append { lsn, bytes }) + { + warn!( + "AOF append dropped for shard {} (lsn {}): channel {}", + shard_id, + lsn, + match e { + flume::TrySendError::Full(_) => "full", + flume::TrySendError::Disconnected(_) => "disconnected", + } + ); + } } /// Synchronous (fsync-before-ack) append for `appendfsync=always` @@ -443,11 +620,51 @@ impl AofWriterPool { let n_shards = self.senders.len(); let shared_manifest = Arc::new(parking_lot::Mutex::new(manifest)); let coord = PerShardRewriteCoord::new(shared_manifest, current_seq, n_shards); + // C4 deadlock guard: fold channels MUST be wired before a per-shard + // rewrite is attempted. If they are absent the AofFold push would + // succeed into a 1-slot ring whose consumer was dropped at construction + // time, causing the writer to block forever on `recv_blocking` while + // the bounded AOF append channel (10_000) fills and silently drops + // every subsequent write (test_ssm4a_fold_4shard_experimental: 2016 of + // 272988 INCRs survived restart). Abort cleanly instead: log an error, + // mark the coord failed, and return — the old generation stays + // authoritative (same abort path as a disconnected writer channel above). + let (Some(fold_producers_vec), Some(fold_notifiers_vec)) = + (self.fold_producers.as_ref(), self.fold_notifiers.as_ref()) + else { + error!( + "F6 per-shard rewrite: fold channels not wired (pool built with \ + per_shard_with_base_dir — call set_fold_channels after mesh \ + creation). Rewrite aborted; old generation remains authoritative. \ + This is a startup-configuration bug." + ); + coord.mark_failed(); + for _ in 0..n_shards { + coord.shard_done(); + } + return Err(AofPoolSendError::SendFailed); + }; + for (idx, s) in self.senders.iter().enumerate() { + // Fail fast on a length mismatch rather than zip-truncating: a + // shard that never receives RewritePerShard never calls + // shard_done(), wedging the coordinator forever. + #[allow(clippy::unwrap_used)] // len == senders len, debug_asserted at set_fold_channels + let fold_producer = fold_producers_vec + .get(idx) + .cloned() + .expect("fold_producers len == senders len (debug_asserted at set_fold_channels)"); + #[allow(clippy::unwrap_used)] // len == senders len, debug_asserted at set_fold_channels + let fold_notifier = fold_notifiers_vec + .get(idx) + .cloned() + .expect("fold_notifiers len == senders len (debug_asserted at set_fold_channels)"); // Blocking send for guaranteed delivery — see the doc comment. if s.send(AofMessage::RewritePerShard { shard_dbs: shard_dbs.clone(), coord: coord.clone(), + fold_producer, + fold_notifier, }) .is_err() { @@ -535,7 +752,7 @@ mod pool_tests { tmp.path().to_path_buf(), ); - let shard_dbs = crate::shard::shared_databases::ShardDatabases::new(vec![ + let (shard_dbs, _inits) = crate::shard::shared_databases::ShardDatabases::new(vec![ vec![crate::storage::Database::new()], vec![crate::storage::Database::new()], ]); @@ -586,7 +803,7 @@ mod pool_tests { // Shard 0's append file starts on the OLD incr (where the live writer // appends before a rewrite). Empty databases keep the fold trivial; the // reopen behaviour is independent of key count. - let shard_dbs = crate::shard::shared_databases::ShardDatabases::new(vec![ + let (shard_dbs, _inits) = crate::shard::shared_databases::ShardDatabases::new(vec![ vec![crate::storage::Database::new()], vec![crate::storage::Database::new()], ]); @@ -597,10 +814,35 @@ mod pool_tests { .unwrap(); let (_tx, rx) = channel::mpsc_bounded::(4); + // Create a dummy fold channel whose ring is pre-filled so try_push(AofFold) + // fails immediately — do_rewrite_per_shard returns Err before reaching + // recv_blocking, which would hang forever with no shard thread consuming. + // Capacity 1, pre-filled with Shutdown so the single slot is taken. + use ringbuf::traits::{Producer as _, Split as _}; + let (mut fold_prod, fold_cons) = + ringbuf::HeapRb::::new(1).split(); + let _ = fold_prod.try_push(crate::shard::dispatch::ShardMessage::Shutdown); + let fold_producer = std::sync::Arc::new(parking_lot::Mutex::new(fold_prod)); + let fold_notifier = std::sync::Arc::new(crate::runtime::channel::Notify::new()); + // Keep fold_cons alive so the producer is not dropped before do_rewrite_per_shard runs. + let _fold_cons = fold_cons; + AOF_REWRITE_IN_PROGRESS.store(true, Ordering::SeqCst); // The fold itself succeeds; aborting is a coordinator decision, so this // returns Ok even though the rewrite is discarded. - do_rewrite_per_shard(0, &shard_dbs, &mut file, &rx, &coord).unwrap(); + // NOTE: With ShardSlice live, do_rewrite_per_shard sends AofFold and blocks + // on the reply. In this unit test there's no shard event loop consuming the + // SPSC ring, so the fold will error out. The test verifies the abort path, + // which is triggered by the fold guard's error handling. + let _ = do_rewrite_per_shard( + 0, + &shard_dbs, + &mut file, + &rx, + &coord, + &fold_producer, + &fold_notifier, + ); // Abort kept the old generation committed and pruned the new-gen incr. assert_eq!(manifest.lock().seq, old_seq, "abort must keep old_seq"); @@ -924,7 +1166,7 @@ mod pool_tests { .append(true) .open(&incr) .unwrap(); - let mut outcome = drain_pending_appends_framed(&rx0, &mut file).unwrap(); + let mut outcome = drain_pending_appends_framed(&rx0, &mut file, usize::MAX).unwrap(); // CONTRACT: drained + parked, NOT yet acked. assert_eq!(outcome.drained, 1, "the AppendSync was drained"); @@ -944,6 +1186,51 @@ mod pool_tests { ); } + /// H1-BARRIER: a zero-length AppendSync (fsync barrier) drained during a + /// rewrite must write NOTHING to the incr file — a len=0 framed header is + /// rejected by `replay_incr_framed` as corruption and bricks the next boot. + /// It still counts toward `drained` (the fold's `pending_aof_count` is a + /// channel-message count) and its ack still parks for the boundary fsync. + #[test] + fn drain_framed_barrier_writes_no_bytes_but_parks_ack() { + let tmp = tempfile::tempdir().unwrap(); + let incr = tmp.path().join("incr.aof"); + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard(vec![tx0, tx1]); + + // A real append followed by a barrier (empty payload). + pool.try_send_append(0, 5, Bytes::from_static(b"*1\r\n$4\r\nPING\r\n")); + let barrier_recv = pool.try_send_append_sync(0, 0, Bytes::new()); + + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&incr) + .unwrap(); + let mut outcome = drain_pending_appends_framed(&rx0, &mut file, usize::MAX).unwrap(); + + assert_eq!(outcome.drained, 2, "both messages count toward drained"); + assert_eq!(outcome.pending_acks.len(), 1, "barrier ack parked"); + + // On disk: ONLY the real append's framed record (12-byte header + 14 + // payload bytes). The barrier must leave no trace. + file.sync_data().unwrap(); + let on_disk = std::fs::read(&incr).unwrap(); + assert_eq!( + on_disk.len(), + 12 + 14, + "barrier must write no header/payload; got {} bytes", + on_disk.len() + ); + + outcome.fulfill_acks(true); + assert_eq!( + barrier_recv.recv_blocking().expect("ack resolves"), + AofAck::Synced + ); + } + /// Issue #140 failure path: if the rewrite-boundary fsync FAILS, a drained /// AppendSync must resolve `FsyncFailed`, never `Synced`. Exercises the /// non-framed `drain_pending_appends` — the DEFAULT `--shards 1` rewrite @@ -1379,4 +1666,105 @@ mod pool_tests { result ); } + + // ----------------------------------------------------------------------- + // H1-BARRIER-1: fsync_barrier under EverySec must be a zero-cost noop — + // no message enqueued, channel stays empty, returns Ok(()) immediately. + // ----------------------------------------------------------------------- + #[test] + fn fsync_barrier_everysec_is_noop() { + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard_with_policy( + vec![tx0, tx1], + FsyncPolicy::EverySec, + Duration::ZERO, + ); + + let result = futures::executor::block_on(pool.fsync_barrier(0)); + + assert!( + result.is_ok(), + "fsync_barrier under EverySec must return Ok immediately, got {:?}", + result + ); + // No AppendSync must have been enqueued — channel stays empty. + assert!( + rx0.try_recv().is_err(), + "fsync_barrier under EverySec must NOT enqueue any message" + ); + } + + // ----------------------------------------------------------------------- + // H1-BARRIER-2: fsync_barrier under Always with a writer that acks Synced + // must return Ok(()). + // ----------------------------------------------------------------------- + #[cfg(feature = "runtime-tokio")] + #[tokio::test] + async fn fsync_barrier_always_writer_acks_synced_returns_ok() { + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard_with_policy( + vec![tx0, tx1], + FsyncPolicy::Always, + Duration::from_millis(500), + ); + + // Spawn a mock writer that receives the AppendSync barrier and acks Synced. + tokio::spawn(async move { + if let Ok(AofMessage::AppendSync { bytes, ack, .. }) = rx0.recv_async().await { + // The barrier sends a zero-length payload — verify it. + assert!( + bytes.is_empty(), + "fsync_barrier must send zero-length AppendSync, got {} bytes", + bytes.len() + ); + let _ = ack.send(AofAck::Synced); + } + }); + + let result = pool.fsync_barrier(0).await; + assert_eq!( + result, + Ok(()), + "fsync_barrier with Synced ack must return Ok" + ); + } + + // ----------------------------------------------------------------------- + // H1-BARRIER-3: fsync_barrier under Always with dead writer (ack dropped) + // must return Err(WriteFailed). Mirrors + // try_send_append_durable_always_writer_dead_returns_write_failed. + // ----------------------------------------------------------------------- + #[test] + fn fsync_barrier_always_dead_writer_returns_write_failed() { + let (tx0, rx0) = channel::mpsc_bounded::(4); + let (tx1, _rx1) = channel::mpsc_bounded::(4); + let pool = AofWriterPool::per_shard_with_policy( + vec![tx0, tx1], + FsyncPolicy::Always, + Duration::ZERO, // legacy unbounded await — disconnect resolves it + ); + + // Spawn a thread that pulls the AppendSync but drops ack — simulating + // a crashed writer. + let handle = std::thread::spawn(move || match rx0.recv() { + Ok(AofMessage::AppendSync { ack, .. }) => drop(ack), + other => panic!("unexpected message: {:?}", other.is_ok()), + }); + + let result = futures::executor::block_on(pool.fsync_barrier(0)); + + handle.join().expect("ack dropper thread"); + + assert!( + result.is_err(), + "fsync_barrier with dead writer must return Err, got Ok" + ); + assert_eq!( + result.unwrap_err(), + AofAck::WriteFailed, + "dead writer must resolve to WriteFailed" + ); + } } diff --git a/src/persistence/aof/rewrite.rs b/src/persistence/aof/rewrite.rs index c052cbd8..2c3e3f5c 100644 --- a/src/persistence/aof/rewrite.rs +++ b/src/persistence/aof/rewrite.rs @@ -318,7 +318,7 @@ impl DrainOutcome { /// durability ordering (ack strictly AFTER the bytes are durable) for every /// `do_rewrite_*` drain site. #[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] -fn sync_and_fulfill_drain( +pub(crate) fn sync_and_fulfill_drain( outcome: &mut DrainOutcome, file: &mut std::fs::File, incr_path: PathBuf, @@ -342,70 +342,91 @@ fn sync_and_fulfill_drain( } } +/// Bounded variant of [`drain_pending_appends`]: drain at most `max_drain` +/// messages (RAW RESP, no framing). Pass [`usize::MAX`] for unbounded. +/// +/// The TopLevel fold's phase-3 mid-drain passes `pending_aof_count` as the +/// bound (C4-DRAIN-BOUND): draining more would consume post-snapshot appends +/// that belong in the NEW incr, and those bytes would then be lost when +/// `manifest.advance()` deletes the old incr at the end of the fold. #[cfg(feature = "runtime-monoio")] -pub(crate) fn drain_pending_appends( +pub(crate) fn drain_pending_appends_bounded( rx: &channel::MpscReceiver, file: &mut std::fs::File, + max_drain: usize, ) -> Result { use std::io::Write; let mut outcome = DrainOutcome::default(); - while let Ok(msg) = rx.try_recv() { - match msg { - // BGREWRITEAOF drain runs on the TopLevel writer (monoio) only; - // PerShard rewrite is RFC step 6. Legacy v1 disk format → ignore lsn. - AofMessage::Append { - bytes: data, - lsn: _, - } => { - file.write_all(&data).map_err(|e| AofError::Io { - path: PathBuf::from(""), - source: e, - })?; - outcome.drained += 1; - } - // AppendSync during a rewrite drain: bytes are written and counted, - // but the ack is PARKED until the caller's post-drain boundary fsync - // (issue #140) — acking `Synced` here would report durability before - // the bytes are fsynced. If the write itself fails the `?` propagates - // the error and the parked ack is dropped with the outcome — the - // caller observes `RecvError`, which it treats as failure. - AofMessage::AppendSync { - bytes: data, - lsn: _, - ack, - } => { - file.write_all(&data).map_err(|e| AofError::Io { - path: PathBuf::from(""), - source: e, - })?; - outcome.drained += 1; - outcome.pending_acks.push(ack); - } - AofMessage::Shutdown => { - outcome.shutdown_requested = true; - } - AofMessage::Rewrite(_) - | AofMessage::RewriteSharded(_) - | AofMessage::RewritePerShard { .. } => { - // Already rewriting — drop redundant request. - } + while outcome.drained < max_drain { + match rx.try_recv() { + Ok(msg) => match msg { + AofMessage::Append { + bytes: data, + lsn: _, + } => { + file.write_all(&data).map_err(|e| AofError::Io { + path: PathBuf::from(""), + source: e, + })?; + outcome.drained += 1; + } + AofMessage::AppendSync { + bytes: data, + lsn: _, + ack, + } => { + file.write_all(&data).map_err(|e| AofError::Io { + path: PathBuf::from(""), + source: e, + })?; + outcome.drained += 1; + outcome.pending_acks.push(ack); + } + AofMessage::Shutdown => { + outcome.shutdown_requested = true; + } + AofMessage::Rewrite(_) + | AofMessage::RewriteSharded(_) + | AofMessage::RewritePerShard { .. } => { + // Already rewriting — drop redundant request. + } + }, + Err(flume::TryRecvError::Empty) => break, + Err(flume::TryRecvError::Disconnected) => break, } } Ok(outcome) } -/// [F6] Drain a per-shard writer's queued appends into its OLD incr file using -/// the framed `[u64 lsn LE][u32 len LE][RESP bytes]` on-disk format that -/// per-shard recovery expects. +#[cfg(feature = "runtime-monoio")] +pub(crate) fn drain_pending_appends( + rx: &channel::MpscReceiver, + file: &mut std::fs::File, +) -> Result { + drain_pending_appends_bounded(rx, file, usize::MAX) +} + +/// [F6] Drain at most `max_drain` pending [`AofMessage::Append`] / +/// [`AofMessage::AppendSync`] messages from `rx` into `file` using the framed +/// `[u64 lsn LE][u32 len LE][RESP bytes]` on-disk format that per-shard +/// recovery expects. +/// +/// Pass [`usize::MAX`] for an unbounded drain (captures all currently-queued +/// messages). The fold phase-3 mid-drain passes +/// [`AofFoldSnapshot::pending_aof_count`] as the bound to avoid an infinite +/// loop under sustained high write load: all pre-snapshot appends are +/// guaranteed to be in the channel already (shard event loop is single- +/// threaded), so draining exactly that many captures every pre-snapshot append +/// without consuming post-snapshot ones. /// /// This is the per-shard twin of [`drain_pending_appends`] (which writes the /// legacy TopLevel raw-RESP format). Correctness depends on the framing -/// matching `replay_per_shard`'s reader — an unframed write here would make the -/// drained appends unparseable on restart. +/// matching `replay_per_shard`'s reader. #[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] pub(crate) fn drain_pending_appends_framed( rx: &channel::MpscReceiver, file: &mut std::fs::File, + max_drain: usize, ) -> Result { use std::io::Write; let mut outcome = DrainOutcome::default(); @@ -416,38 +437,49 @@ pub(crate) fn drain_pending_appends_framed( file.write_all(&header)?; file.write_all(data) }; - while let Ok(msg) = rx.try_recv() { - match msg { - AofMessage::Append { lsn, bytes: data } => { - write_framed(file, lsn, &data).map_err(|e| AofError::Io { - path: PathBuf::from(""), - source: e, - })?; - outcome.drained += 1; - } - AofMessage::AppendSync { - lsn, - bytes: data, - ack, - } => { - write_framed(file, lsn, &data).map_err(|e| AofError::Io { - path: PathBuf::from(""), - source: e, - })?; - outcome.drained += 1; - // Park the ack until the caller's post-drain boundary fsync - // (issue #140); resolved Synced/FsyncFailed by - // `sync_and_fulfill_drain`. Mirrors `drain_pending_appends`. - outcome.pending_acks.push(ack); - } - AofMessage::Shutdown => { - outcome.shutdown_requested = true; - } - AofMessage::Rewrite(_) - | AofMessage::RewriteSharded(_) - | AofMessage::RewritePerShard { .. } => { - // Already rewriting this shard — drop redundant request. - } + while outcome.drained < max_drain { + match rx.try_recv() { + Ok(msg) => match msg { + AofMessage::Append { lsn, bytes: data } => { + write_framed(file, lsn, &data).map_err(|e| AofError::Io { + path: PathBuf::from(""), + source: e, + })?; + outcome.drained += 1; + } + AofMessage::AppendSync { + lsn, + bytes: data, + ack, + } => { + // H1-BARRIER: a zero-length AppendSync is an fsync barrier + // (pool::fsync_barrier) — it must produce NO on-disk record + // (a len=0 framed header reads as corruption on replay) but + // still counts toward `drained` (sender.len() counted it) + // and its ack still parks for the boundary fsync. + if !data.is_empty() { + write_framed(file, lsn, &data).map_err(|e| AofError::Io { + path: PathBuf::from(""), + source: e, + })?; + } + outcome.drained += 1; + // Park the ack until the caller's post-drain boundary fsync + // (issue #140); resolved Synced/FsyncFailed by + // `sync_and_fulfill_drain`. Mirrors `drain_pending_appends`. + outcome.pending_acks.push(ack); + } + AofMessage::Shutdown => { + outcome.shutdown_requested = true; + } + AofMessage::Rewrite(_) + | AofMessage::RewriteSharded(_) + | AofMessage::RewritePerShard { .. } => { + // Already rewriting this shard — drop redundant request. + } + }, + Err(flume::TryRecvError::Empty) => break, + Err(flume::TryRecvError::Disconnected) => break, } } Ok(outcome) @@ -522,72 +554,163 @@ pub(crate) fn drain_pending_appends_framed( #[cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] pub(crate) fn do_rewrite_per_shard( shard_id: u16, - shard_dbs: &crate::shard::shared_databases::ShardDatabases, + _shard_dbs: &crate::shard::shared_databases::ShardDatabases, file: &mut std::fs::File, rx: &channel::MpscReceiver, coord: &PerShardRewriteCoord, + fold_producer: &parking_lot::Mutex>, + fold_notifier: &std::sync::Arc, ) -> Result<(), MoonError> { + use ringbuf::traits::Producer; // Panic/early-error safety: guarantees `shard_done` runs on EVERY exit // (success via `complete()`, `?`-error or panic-unwind via `Drop`). The // phase-8 `await_outcome` barrier makes that a liveness requirement, so // callers MUST NOT call `shard_done` after invoking this function — the // guard owns the single decrement for all exits. See `ShardDoneGuard`. let guard = ShardDoneGuard::new(coord); - - let sidx = shard_id as usize; - let all_shards = shard_dbs.all_shard_dbs(); - if sidx >= all_shards.len() { - return Err(AofError::RewriteFailed { - detail: format!( - "do_rewrite_per_shard: shard {} out of range ({} shards)", - sidx, - all_shards.len() - ), - } - .into()); - } + let _fold_t0 = std::time::Instant::now(); // Phase 1: drain pre-rewrite queued appends into old incr (framed), fsync, // then resolve their parked AppendSync acks (issue #140). - let mut pre_drain = drain_pending_appends_framed(rx, file)?; + // + // C4-DRAIN-BOUND (phase 1): snapshot `rx.len()` NOW — all messages currently + // in the channel are pre-fold appends. New appends that arrive AFTER this + // snapshot were enqueued concurrently with the fold and belong in the NEW incr; + // draining them here would pull post-fold appends into the OLD incr, causing + // them to be replayed on top of a base that already contains them (double-apply). + // Without this bound, under sustained high write load the drain loops forever + // because the INCR producer keeps the channel perpetually non-empty. + let pre_drain_bound = rx.len(); + let mut pre_drain = drain_pending_appends_framed(rx, file, pre_drain_bound)?; sync_and_fulfill_drain(&mut pre_drain, file, PathBuf::from(""))?; + info!( + "F6 shard {} phase1 done: drained {} appends ({:.1}ms)", + shard_id, + pre_drain.drained, + _fold_t0.elapsed().as_secs_f64() * 1000.0 + ); - // Phase 2: acquire write locks on this shard's db(s) (db_idx ascending). - let shard_locks = &all_shards[sidx]; - let guards: Vec<_> = shard_locks.iter().map(|lock| lock.write()).collect(); - - // Phase 3: drain appends that completed between phase 1 and phase 2, fsync, - // then resolve their parked AppendSync acks (issue #140). - let mut mid_drain = drain_pending_appends_framed(rx, file)?; - sync_and_fulfill_drain(&mut mid_drain, file, PathBuf::from(""))?; + // Phases 2-5 (C4 cooperative snapshot — ShardSlice is the live store): + // + // Instead of acquiring RwLock write guards on the shard's databases (which + // the AOF writer thread cannot do because ShardSlice is thread-local/!Send), + // we send an AofFold message to the shard's SPSC ring and block on the + // oneshot reply. The shard event loop processes AofFold atomically between + // commands (single-threaded cooperative runtime), providing the same mutual + // exclusion that RwLock write guards gave. See ShardMessage::AofFold and + // its handler in spsc_handler.rs. + // + // Ordering discipline (exactly-once guarantee preserved): + // 1. Phase 1 above drained pre-fold appends into OLD incr. + // 2. AofFold push notifies the shard; the shard event loop will drain its + // SPSC ring before processing AofFold, so any command that was + // in-flight when Phase 1 finished but not yet queued to the AOF channel + // will be processed (and its append enqueued) before the snapshot. + // 3. Phase 3 (mid-drain below) captures appends enqueued after Phase 1 + // but before the shard built the snapshot — same as the old mid-drain. + // 4. Snapshot is replied after all prior mutations are applied; no + // command can mutate the shard while it is building the snapshot. + let (reply_tx, reply_rx) = + crate::runtime::channel::oneshot::(); + { + let mut prod = fold_producer.lock(); + prod.try_push(crate::shard::dispatch::ShardMessage::AofFold { reply_tx }) + .map_err(|_| AofError::RewriteFailed { + detail: format!( + "do_rewrite_per_shard: shard {} AofFold SPSC ring full — fold aborted", + shard_id + ), + })?; + } + fold_notifier.notify_one(); - // Phase 4: snapshot this shard's databases under the locks. - let now_ms = current_time_ms(); - let mut snapshot: Vec<( - Vec<( - crate::storage::compact_key::CompactKey, - crate::storage::entry::Entry, - )>, - u32, - )> = Vec::with_capacity(guards.len()); - for guard in &guards { - let base_ts = guard.base_timestamp(); - let mut entries = Vec::new(); - for (key, entry) in guard.data().iter() { - if !entry.is_expired_at(base_ts, now_ms) { - entries.push((key.clone(), entry.clone())); + // Phase 4: collect the cooperative snapshot (blocks until the shard replies). + // Use recv_blocking since do_rewrite_per_shard runs on a dedicated std::thread + // (the per-shard AOF writer), not inside an async executor. + // + // C4 ordering invariant (exactly-once guarantee — frozen contract §3 C4): + // Every command the shard processed BEFORE building its snapshot had already + // enqueued its AOF append (via `try_send_append`) before the shard sent the + // reply — happens-before, by the single-threaded event-loop order. Therefore + // the mid-drain (phase 3) executed AFTER this recv_blocking captures ALL + // pre-snapshot appends into the OLD incr; post-snapshot appends land in the + // NEW incr. Running mid-drain BEFORE recv_blocking (the prior buggy order) + // misses appends that arrive while the shard is building its snapshot — + // they would land in the NEW incr and be replayed on top of a base that + // already contains them → double-apply on recovery (observed: 2016 of + // 272988 INCRs survived restart in test_ssm4a_fold_4shard_experimental). + // Wait for the shard event loop to build and reply with the snapshot. + // Poll with try_recv + sleep so we can log a warning if the reply is slow + // (helps diagnose starvation of the fold consumer under high write load). + let fold_snapshot = { + let wait_start = std::time::Instant::now(); + let mut warned = false; + loop { + match reply_rx.try_recv() { + Ok(snap) => break snap, + Err(flume::TryRecvError::Empty) => { + if !warned && wait_start.elapsed() >= std::time::Duration::from_millis(500) { + warned = true; + warn!( + "do_rewrite_per_shard: shard {} waiting for AofFold snapshot ({:.1}s elapsed) — \ + shard event loop may be stalled or fold consumer not draining", + shard_id, + wait_start.elapsed().as_secs_f64() + ); + } + std::thread::sleep(std::time::Duration::from_millis(1)); + } + Err(flume::TryRecvError::Disconnected) => { + return Err(AofError::RewriteFailed { + detail: format!( + "do_rewrite_per_shard: shard {} AofFold reply channel dropped (shard shut down?)", + shard_id + ), + }.into()); + } } } - snapshot.push((entries, base_ts)); - } + }; + let pending_aof_count = fold_snapshot.pending_aof_count; + let snapshot = fold_snapshot.dbs; + info!( + "F6 shard {} snapshot received: {} dbs, {} pre-snapshot pending ({:.1}ms total)", + shard_id, + snapshot.len(), + pending_aof_count, + _fold_t0.elapsed().as_secs_f64() * 1000.0 + ); - // Phase 5: release locks before the expensive disk write. - drop(guards); + // Phase 3: drain appends that arrived while the shard was building the snapshot, + // fsync, then resolve their parked AppendSync acks (issue #140). + // + // MUST run AFTER receiving the snapshot (see C4 ordering invariant above): the + // shard's event loop enqueues every pre-snapshot append before sending the reply, + // so a post-reply drain captures them all into the OLD incr. + // + // C4-DRAIN-BOUND: drain at most `pending_aof_count` messages — the exact number + // of pre-snapshot appends the shard reported before building its snapshot. This + // prevents an infinite drain loop under sustained high write load where new + // (post-snapshot) appends arrive faster than we can drain them. + let mut mid_drain = drain_pending_appends_framed(rx, file, pending_aof_count)?; + sync_and_fulfill_drain(&mut mid_drain, file, PathBuf::from(""))?; + info!( + "F6 shard {} phase3 done: drained {} mid-appends ({:.1}ms total)", + shard_id, + mid_drain.drained, + _fold_t0.elapsed().as_secs_f64() * 1000.0 + ); // Phase 6: write new base, advance THIS shard's manifest entry (no seq // commit), reopen to the new incr. The manifest lock is held only for the // brief, await-free advance_shard call. let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&snapshot)?; + info!( + "F6 shard {} rdb serialized: {} bytes ({:.1}ms total)", + shard_id, + rdb_bytes.len(), + _fold_t0.elapsed().as_secs_f64() * 1000.0 + ); let new_incr = { let mut m = coord.manifest.lock(); m.advance_shard(shard_id, coord.new_seq, &rdb_bytes)? @@ -739,75 +862,155 @@ pub(crate) fn do_rewrite_single( Ok(()) } -/// Multi-part rewrite: snapshot all shards → merged RDB base → advance manifest. +/// TopLevel cooperative fold: snapshot shard 0 → merged RDB base → advance manifest. /// -/// See [`do_rewrite_single`] for the ordering rationale. The multi-shard variant -/// holds write locks on every (shard, db) pair simultaneously for the duration -/// of the snapshot. This creates a brief global write pause, but it is the only -/// way to guarantee a torn-free snapshot without per-message sequence numbers. +/// Implements the C4 cooperative-snapshot protocol for the TopLevel (--shards 1) +/// AOF layout. Mirrors [`do_rewrite_per_shard`] but operates on the single +/// TopLevel manifest instead of per-shard manifest entries. +/// +/// Correctness ordering (exactly-once for non-idempotent commands like INCR): +/// +/// 1. Drain queued appends (RAW RESP, not framed) into the OLD incr and fsync. +/// 2. Send `ShardMessage::AofFold` to shard 0's SPSC ring and block on the +/// cooperative snapshot reply. The shard event loop processes AofFold +/// atomically between commands, providing the same mutual exclusion the old +/// RwLock write guards gave. +/// 3. Mid-drain: drain exactly `pending_aof_count` pre-snapshot appends into +/// OLD incr and fsync. +/// 4. Write new base RDB → advance manifest → reopen `file` to new incr. +/// +/// If `fold_channels` is `None` (main.rs failed to wire them at startup), +/// abort cleanly with an error so the old generation stays authoritative. +/// +/// **RAW RESP format**: TopLevel incr files use plain RESP bytes (no [lsn][len] +/// framing). [`drain_pending_appends`] writes the correct format; never use +/// [`drain_pending_appends_framed`] here. #[cfg(feature = "runtime-monoio")] pub(crate) fn do_rewrite_sharded( - shard_dbs: &crate::shard::shared_databases::ShardDatabases, + _shard_dbs: &crate::shard::shared_databases::ShardDatabases, manifest: &mut crate::persistence::aof_manifest::AofManifest, file: &mut std::fs::File, rx: &channel::MpscReceiver, + fold_channels: Option<&( + Arc>>, + Arc, + )>, ) -> Result<(), MoonError> { - // Phase 1: drain pre-rewrite queued appends into old incr, fsync, then - // resolve their parked AppendSync acks (issue #140). + use ringbuf::traits::Producer; + + let _fold_t0 = std::time::Instant::now(); + + // C4 deadlock guard: fold channels MUST be wired. If absent, abort + // cleanly — the old generation stays authoritative. + let (fold_producer, fold_notifier) = match fold_channels { + Some(pair) => pair, + None => { + return Err(AofError::RewriteFailed { + detail: "do_rewrite_sharded (TopLevel): fold channels not wired at startup; \ + rewrite aborted — old generation stays authoritative. \ + Check main.rs fold-channel wiring." + .to_string(), + } + .into()); + } + }; + + // Phase 1: drain pre-rewrite queued appends into old incr (RAW RESP), fsync. let mut pre_drain = drain_pending_appends(rx, file)?; sync_and_fulfill_drain(&mut pre_drain, file, manifest.incr_path())?; + info!( + "TopLevel fold phase1 done: drained {} appends ({:.1}ms)", + pre_drain.drained, + _fold_t0.elapsed().as_secs_f64() * 1000.0 + ); - // Phase 2: acquire write locks on ALL (shard, db) pairs simultaneously. - // Lock order is (shard_idx, db_idx) ascending — must match anywhere else - // that acquires multiple locks to prevent deadlock (currently no other - // call site does, but the ordering discipline is documented for future - // maintainers). - let all_shards = shard_dbs.all_shard_dbs(); - let mut guards: Vec> = Vec::with_capacity(all_shards.len()); - for shard_locks in all_shards { - let mut shard_guards = Vec::with_capacity(shard_locks.len()); - for lock in shard_locks { - shard_guards.push(lock.write()); - } - guards.push(shard_guards); + // Phases 2-4 (C4 cooperative snapshot): + // + // Send AofFold to shard 0 and block on the cooperative snapshot reply. + // The shard event loop processes AofFold atomically between commands — + // same mutual exclusion the old RwLock write guards provided. + let (reply_tx, reply_rx) = + crate::runtime::channel::oneshot::(); + { + let mut prod = fold_producer.lock(); + prod.try_push(crate::shard::dispatch::ShardMessage::AofFold { reply_tx }) + .map_err(|_| AofError::RewriteFailed { + detail: "do_rewrite_sharded (TopLevel): AofFold SPSC ring full — fold aborted" + .to_string(), + })?; } + fold_notifier.notify_one(); - // Phase 3: drain appends completed between phase 1 and phase 2, fsync, then - // resolve their parked AppendSync acks (issue #140). - let mut mid_drain = drain_pending_appends(rx, file)?; - sync_and_fulfill_drain(&mut mid_drain, file, manifest.incr_path())?; - - // Phase 4: snapshot under locks. - let db_count = shard_dbs.db_count(); - let mut merged: Vec<( - Vec<( - crate::storage::compact_key::CompactKey, - crate::storage::entry::Entry, - )>, - u32, - )> = (0..db_count).map(|_| (Vec::new(), 0u32)).collect(); - let now_ms = current_time_ms(); - for shard_guards in &guards { - for (db_idx, guard) in shard_guards.iter().enumerate() { - let base_ts = guard.base_timestamp(); - if merged[db_idx].0.is_empty() { - merged[db_idx].1 = base_ts; - } - for (key, entry) in guard.data().iter() { - if !entry.is_expired_at(base_ts, now_ms) { - merged[db_idx].0.push((key.clone(), entry.clone())); + // Poll for the cooperative snapshot (blocks until shard 0 replies). + // Runs on the dedicated AOF writer thread (not inside an async executor), + // so recv_blocking / sleep are safe. + let fold_snapshot = { + let wait_start = std::time::Instant::now(); + let mut warned = false; + loop { + match reply_rx.try_recv() { + Ok(snap) => break snap, + Err(flume::TryRecvError::Empty) => { + if !warned && wait_start.elapsed() >= std::time::Duration::from_millis(500) { + warned = true; + warn!( + "do_rewrite_sharded (TopLevel): waiting for AofFold snapshot \ + ({:.1}s elapsed) — shard 0 event loop may be stalled", + wait_start.elapsed().as_secs_f64() + ); + } + std::thread::sleep(std::time::Duration::from_millis(1)); + } + Err(flume::TryRecvError::Disconnected) => { + return Err(AofError::RewriteFailed { + detail: "do_rewrite_sharded (TopLevel): AofFold reply channel dropped \ + (shard 0 shut down?)" + .to_string(), + } + .into()); } } } - } + }; + let pending_aof_count = fold_snapshot.pending_aof_count; + let snapshot = fold_snapshot.dbs; + info!( + "TopLevel fold snapshot received: {} dbs, {} pre-snapshot pending ({:.1}ms)", + snapshot.len(), + pending_aof_count, + _fold_t0.elapsed().as_secs_f64() * 1000.0 + ); - // Phase 5: release locks before the expensive disk write. - drop(guards); + // Phase 3: mid-drain — drain exactly `pending_aof_count` pre-snapshot appends + // into OLD incr (RAW RESP format), then fsync + fulfill parked AppendSync acks. + // + // MUST run AFTER receiving the snapshot (C4 ordering invariant — identical + // to do_rewrite_per_shard's phase 3 rationale). + // + // C4-DRAIN-BOUND: drain at most `pending_aof_count` messages — the exact count + // the shard reported before building its snapshot. Draining beyond this count + // would consume post-snapshot appends that belong in the NEW incr; those bytes + // would be silently lost when `manifest.advance()` deletes the old incr file + // at the end of phase 4. This mirrors the identical bound used by + // `drain_pending_appends_framed` in `do_rewrite_per_shard`. + let mut mid_drain = drain_pending_appends_bounded(rx, file, pending_aof_count)?; + sync_and_fulfill_drain(&mut mid_drain, file, manifest.incr_path())?; + info!( + "TopLevel fold phase3 done: drained {} mid-appends ({:.1}ms)", + mid_drain.drained, + _fold_t0.elapsed().as_secs_f64() * 1000.0 + ); - // Phase 6: write new base, advance manifest, reopen. - let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&merged)?; + // Phase 4 (= Phase 6 in do_rewrite_per_shard numbering): + // Write new base RDB, advance the manifest (bumps seq, writes manifest, + // deletes old files), reopen `file` to the new incr. + let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&snapshot)?; + info!( + "TopLevel fold rdb serialized: {} bytes ({:.1}ms)", + rdb_bytes.len(), + _fold_t0.elapsed().as_secs_f64() * 1000.0 + ); let new_incr = manifest.advance(&rdb_bytes)?; - *file = std::fs::OpenOptions::new() .create(true) .append(true) @@ -818,11 +1021,11 @@ pub(crate) fn do_rewrite_sharded( })?; info!( - "AOF rewrite complete (sharded): drained {}+{} pre-snapshot appends, seq={}", + "TopLevel AOF rewrite complete: drained {}+{} appends, seq={}", pre_drain.drained, mid_drain.drained, manifest.seq ); if pre_drain.shutdown_requested || mid_drain.shutdown_requested { - warn!("AOF writer: shutdown requested during rewrite (will honor on next recv)"); + warn!("TopLevel AOF writer: shutdown requested during rewrite (honored on next recv)"); } Ok(()) } @@ -875,37 +1078,243 @@ fn rewrite_aof_sync(db: &SharedDatabases, aof_path: &Path) -> Result<(), MoonErr Ok(()) } -/// Rewrite the AOF in sharded mode with RDB preamble. +/// TopLevel cooperative fold for the tokio runtime. +/// +/// Implements the C4 cooperative-snapshot protocol for the TopLevel (--shards 1) +/// AOF layout under `runtime-tokio`. Called by the tokio `aof_writer_task` when +/// it receives `AofMessage::RewriteSharded`. /// -/// Merges all shards' databases into a single RDB snapshot, writes it as -/// the AOF base file. New incremental writes are appended as RESP after. +/// Unlike the monoio path (which has a manifest and a proper incr file), the +/// tokio TopLevel writer uses the legacy single-file `aof_path`. The rewrite: +/// +/// 1. Drains pre-snapshot appends from `rx` into `old_file` (RAW RESP). +/// 2. Sends `ShardMessage::AofFold` to shard 0 and blocks for the snapshot. +/// 3. Mid-drains into `old_file`. +/// 4. Writes the RDB snapshot to `aof_path.tmp` then renames atomically to +/// `aof_path` — the same post-rewrite file the caller will reopen for appending. +/// +/// Returns `Err` (abort, old file unchanged) if fold channels are absent or the +/// shard reply channel is dropped. The caller clears `AOF_REWRITE_IN_PROGRESS` +/// regardless. +/// +/// **Detection note:** the tokio path writes to `aof_path` (e.g. `appendonly.aof`), +/// NOT to a manifest-stamped `moon.aof..base.rdb`. Test helpers that poll +/// for completion should check for `appendonly.aof` containing the RDB preamble +/// (`MOON` magic at offset 0) rather than a manifest-based file. #[allow(dead_code)] pub(crate) fn rewrite_aof_sharded_sync( - shard_dbs: &crate::shard::shared_databases::ShardDatabases, + _shard_dbs: &crate::shard::shared_databases::ShardDatabases, aof_path: &Path, + rx: &channel::MpscReceiver, + old_file: &mut std::fs::File, + fold_channels: Option<&( + Arc>>, + Arc, + )>, ) -> Result<(), MoonError> { - let db_count = shard_dbs.db_count(); - let now_ms = current_time_ms(); - let mut merged_dbs: Vec = (0..db_count).map(|_| Database::new()).collect(); + use ringbuf::traits::Producer; + use std::io::Write as _; - for shard_locks in shard_dbs.all_shard_dbs() { - for (db_idx, lock) in shard_locks.iter().enumerate() { - let guard = lock.read(); - for (key, entry) in guard.data().iter() { - if !entry.is_expired_at(guard.base_timestamp(), now_ms) { - merged_dbs[db_idx].set(key.to_bytes(), entry.clone()); + let _fold_t0 = std::time::Instant::now(); + + // C4 deadlock guard: fold channels MUST be wired. + let (fold_producer, fold_notifier) = match fold_channels { + Some(pair) => pair, + None => { + return Err(AofError::RewriteFailed { + detail: "rewrite_aof_sharded_sync (tokio TopLevel): fold channels not wired; \ + rewrite aborted — old file unchanged." + .to_string(), + } + .into()); + } + }; + + // Phase 1: pre-drain (RAW RESP) into old file, then fsync + resolve parked + // AppendSync acks (D2 fix: park acks here; fulfill after boundary fsync so + // clients are never told Synced before the bytes are durable). + // drain_pending_appends is #[cfg(runtime-monoio)] only; under tokio we do + // the minimal equivalent inline (read until channel is empty). + let mut pre_shutdown_requested = false; + { + let mut pre_acks: Vec> = Vec::new(); + while let Ok(msg) = rx.try_recv() { + match msg { + AofMessage::Append { bytes: data, .. } => { + old_file.write_all(&data).map_err(|e| AofError::Io { + path: aof_path.to_path_buf(), + source: e, + })?; + } + AofMessage::AppendSync { + bytes: data, ack, .. + } => { + old_file.write_all(&data).map_err(|e| AofError::Io { + path: aof_path.to_path_buf(), + source: e, + })?; + pre_acks.push(ack); + } + AofMessage::Shutdown => { + pre_shutdown_requested = true; + } + AofMessage::Rewrite(_) + | AofMessage::RewriteSharded(_) + | AofMessage::RewritePerShard { .. } => {} + } + } + let _ = old_file.flush(); + // Fulfill pre-drain acks after the boundary fsync (issue #140). + match old_file.sync_data() { + Ok(()) => { + for ack in pre_acks { + let _ = ack.send(AofAck::Synced); + } + } + Err(_) => { + for ack in pre_acks { + let _ = ack.send(AofAck::FsyncFailed); } } } } + info!( + "rewrite_aof_sharded_sync (tokio) phase1 done ({:.1}ms)", + _fold_t0.elapsed().as_secs_f64() * 1000.0 + ); + + // Phases 2-3: AofFold cooperative snapshot. + let (reply_tx, reply_rx) = + crate::runtime::channel::oneshot::(); + { + let mut prod = fold_producer.lock(); + prod.try_push(crate::shard::dispatch::ShardMessage::AofFold { reply_tx }) + .map_err(|_| AofError::RewriteFailed { + detail: "rewrite_aof_sharded_sync (tokio): AofFold SPSC ring full".to_string(), + })?; + } + fold_notifier.notify_one(); - let rdb_bytes = crate::persistence::rdb::save_to_bytes(&merged_dbs)?; + let fold_snapshot = { + let wait_start = std::time::Instant::now(); + let mut warned = false; + loop { + match reply_rx.try_recv() { + Ok(snap) => break snap, + Err(flume::TryRecvError::Empty) => { + if !warned && wait_start.elapsed() >= std::time::Duration::from_millis(500) { + warned = true; + warn!( + "rewrite_aof_sharded_sync (tokio): waiting for AofFold snapshot \ + ({:.1}s) — shard 0 may be stalled", + wait_start.elapsed().as_secs_f64() + ); + } + std::thread::sleep(std::time::Duration::from_millis(1)); + } + Err(flume::TryRecvError::Disconnected) => { + return Err(AofError::RewriteFailed { + detail: "rewrite_aof_sharded_sync (tokio): AofFold reply channel dropped" + .to_string(), + } + .into()); + } + } + } + }; + // D1 fix: capture pending_aof_count to bound the mid-drain (C4-DRAIN-BOUND). + // The shard reports exactly how many AOF messages were enqueued before it + // built the snapshot. Draining more would consume post-snapshot appends that + // must land in the new file (aof_path) AFTER the rename in phase 4. + let pending_aof_count = fold_snapshot.pending_aof_count; + let snapshot = fold_snapshot.dbs; + info!( + "rewrite_aof_sharded_sync (tokio) snapshot: {} dbs, {} pre-snapshot pending ({:.1}ms)", + snapshot.len(), + pending_aof_count, + _fold_t0.elapsed().as_secs_f64() * 1000.0 + ); + // Phase 3 mid-drain: drain exactly `pending_aof_count` pre-snapshot appends + // into OLD file (RAW RESP), then fsync + resolve parked AppendSync acks. + // + // C4-DRAIN-BOUND: must NOT drain beyond pending_aof_count. Post-snapshot + // appends MUST remain in the channel so the writer loop appends them to the + // reopened new file after this function returns. Draining them here would + // cause them to be written to old_file (old inode), which is superseded by + // the rename in phase 4 → those bytes are lost on restart (D1 bug). + let mut mid_shutdown_requested = false; + { + let mut mid_acks: Vec> = Vec::new(); + let mut drained = 0usize; + while drained < pending_aof_count { + match rx.try_recv() { + Ok(msg) => match msg { + AofMessage::Append { bytes: data, .. } => { + old_file.write_all(&data).map_err(|e| AofError::Io { + path: aof_path.to_path_buf(), + source: e, + })?; + drained += 1; + } + AofMessage::AppendSync { + bytes: data, ack, .. + } => { + old_file.write_all(&data).map_err(|e| AofError::Io { + path: aof_path.to_path_buf(), + source: e, + })?; + drained += 1; + mid_acks.push(ack); + } + AofMessage::Shutdown => { + mid_shutdown_requested = true; + } + AofMessage::Rewrite(_) + | AofMessage::RewriteSharded(_) + | AofMessage::RewritePerShard { .. } => {} + }, + Err(_) => break, + } + } + let _ = old_file.flush(); + // Fulfill mid-drain acks after the boundary fsync (issue #140). + match old_file.sync_data() { + Ok(()) => { + for ack in mid_acks { + let _ = ack.send(AofAck::Synced); + } + } + Err(_) => { + for ack in mid_acks { + let _ = ack.send(AofAck::FsyncFailed); + } + } + } + info!( + "rewrite_aof_sharded_sync (tokio) phase3 mid-drain done: {} messages ({:.1}ms)", + drained, + _fold_t0.elapsed().as_secs_f64() * 1000.0 + ); + } + + // Phase 4: write RDB snapshot to aof_path (atomic tmp → rename). + let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&snapshot)?; let tmp_path = aof_path.with_extension("aof.tmp"); - std::fs::write(&tmp_path, &rdb_bytes).map_err(|e| AofError::Io { - path: tmp_path.clone(), - source: e, - })?; + { + let mut f = std::fs::File::create(&tmp_path).map_err(|e| AofError::Io { + path: tmp_path.clone(), + source: e, + })?; + f.write_all(&rdb_bytes).map_err(|e| AofError::Io { + path: tmp_path.clone(), + source: e, + })?; + f.sync_data().map_err(|e| AofError::Io { + path: tmp_path.clone(), + source: e, + })?; + } std::fs::rename(&tmp_path, aof_path).map_err(|e| AofError::RewriteFailed { detail: format!( "rename {} -> {}: {}", @@ -916,9 +1325,16 @@ pub(crate) fn rewrite_aof_sharded_sync( })?; info!( - "AOF rewrite (sharded, RDB preamble) complete: {} bytes", - rdb_bytes.len() + "rewrite_aof_sharded_sync (tokio) complete: {} bytes ({:.1}ms)", + rdb_bytes.len(), + _fold_t0.elapsed().as_secs_f64() * 1000.0 ); + if pre_shutdown_requested || mid_shutdown_requested { + warn!( + "rewrite_aof_sharded_sync (tokio): shutdown requested during rewrite \ + (will honor on next recv)" + ); + } Ok(()) } diff --git a/src/persistence/aof/writer_task.rs b/src/persistence/aof/writer_task.rs index 18b687ee..5eac860d 100644 --- a/src/persistence/aof/writer_task.rs +++ b/src/persistence/aof/writer_task.rs @@ -1,7 +1,10 @@ //! AOF writer tasks: single-file and per-shard background append loops. #![allow(unused_imports, unused_variables, unreachable_code, clippy::empty_loop)] -use super::rewrite::{do_rewrite_per_shard, rewrite_aof_sharded_sync}; +use super::rewrite::{ + do_rewrite_per_shard, drain_pending_appends_framed, rewrite_aof_sharded_sync, + sync_and_fulfill_drain, +}; use super::*; // `do_rewrite_single` / `do_rewrite_sharded` exist only under the monoio runtime. #[cfg(feature = "runtime-monoio")] @@ -9,11 +12,22 @@ use super::rewrite::{do_rewrite_sharded, do_rewrite_single}; /// Background AOF writer task. Receives commands via mpsc channel and appends them /// to the AOF file. Handles fsync according to the configured policy. +/// +/// `fold_channels` — for the TopLevel `--shards 1` path: an optional +/// `(fold_producer, fold_notifier)` pair for shard 0, wired by `main.rs` via +/// `set_fold_channels` + extracted here. When `Some`, `do_rewrite_sharded` +/// uses the C4 cooperative-snapshot fold; when `None` (legacy, no AOF fold +/// wired), `do_rewrite_sharded` falls back to an error-abort so the old +/// generation stays authoritative rather than deadlocking. pub async fn aof_writer_task( rx: channel::MpscReceiver, aof_path: PathBuf, fsync: FsyncPolicy, cancel: CancellationToken, + fold_channels: Option<( + Arc>>, + Arc, + )>, ) { #[cfg(feature = "runtime-tokio")] use tokio::io::AsyncWriteExt; @@ -37,10 +51,6 @@ pub async fn aof_writer_task( let mut writer = tokio::io::BufWriter::new(file); #[cfg(feature = "runtime-tokio")] let mut last_fsync = Instant::now(); - #[cfg(feature = "runtime-tokio")] - let mut interval = tokio::time::interval(std::time::Duration::from_secs(1)); - #[cfg(feature = "runtime-tokio")] - interval.tick().await; // consume first tick // Monoio path: multi-part AOF (base RDB + incremental RESP) with sync I/O. // @@ -250,7 +260,18 @@ pub async fn aof_writer_task( error!("AOF pre-rewrite sync failed (seq {}): {}", manifest.seq, e); } } - match do_rewrite_sharded(&shard_dbs, &mut manifest, &mut file, &rx) { + // C4 TopLevel cooperative fold: pass the wired fold channels + // (producer + notifier for shard 0) so do_rewrite_sharded can + // use the AofFold SPSC protocol instead of the deleted RwLock + // path. `fold_channels` is `None` only if main.rs failed to + // wire them at startup (Arc::get_mut race — logged at boot). + match do_rewrite_sharded( + &shard_dbs, + &mut manifest, + &mut file, + &rx, + fold_channels.as_ref(), + ) { Ok(()) => { write_error = false; } @@ -274,11 +295,33 @@ pub async fn aof_writer_task( loop { #[cfg(feature = "runtime-tokio")] - tokio::select! { - msg = rx.recv_async() => { + { + // Bounded recv (EverySec durability): wake at least every 200ms even + // when idle so the flush deadline check after this select! is honored + // within its 1s bound. A long-lived `interval.tick()` select arm is + // fairness-starvable under sustained writes and unreliable when idle + // (see the per-shard writer below, which hit exactly that) — the + // bounded recv cannot starve. flume's recv future is drop-safe on + // the Elapsed branch (no message consumed on timeout). + let recv_result = tokio::select! { + r = tokio::time::timeout( + std::time::Duration::from_millis(200), + rx.recv_async(), + ) => r, + _ = cancel.cancelled() => { + let _ = writer.flush().await; + let _ = writer.get_ref().sync_data().await; + info!("AOF writer cancelled"); + break; + } + }; + if let Ok(msg) = recv_result { match msg { // TopLevel writer (tokio): legacy v1 plain RESP, lsn ignored. - Ok(AofMessage::Append { bytes: data, lsn: _ }) => { + Ok(AofMessage::Append { + bytes: data, + lsn: _, + }) => { if let Err(e) = writer.write_all(&data).await { error!("AOF write error: {}", e); continue; @@ -289,12 +332,17 @@ pub async fn aof_writer_task( let _ = writer.get_ref().sync_data().await; } FsyncPolicy::EverySec | FsyncPolicy::No => { - // EverySec handled by interval tick below; No does nothing + // EverySec handled by the deadline check after + // the select!; No does nothing } } } // AppendSync: write + fsync + ack, regardless of policy. - Ok(AofMessage::AppendSync { bytes: data, lsn: _, ack }) => { + Ok(AofMessage::AppendSync { + bytes: data, + lsn: _, + ack, + }) => { if let Err(e) = writer.write_all(&data).await { error!("AOF AppendSync write error: {}", e); let _ = ack.send(AofAck::WriteFailed); @@ -324,11 +372,12 @@ pub async fn aof_writer_task( .store(false, std::sync::atomic::Ordering::SeqCst); // Reopen file after rewrite (it was replaced) - let reopen_result: Result = tokio::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&aof_path) - .await; + let reopen_result: Result = + tokio::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&aof_path) + .await; match reopen_result { Ok(f) => { writer = tokio::io::BufWriter::new(f); @@ -338,26 +387,58 @@ pub async fn aof_writer_task( return; } } + // Back-date so the backlog drained right after the + // rewrite reaches disk within ~100ms + wake floor, + // not a full second later (mirrors the per-shard + // writer's post-rewrite back-dating). + last_fsync = Instant::now() - std::time::Duration::from_millis(900); } Ok(AofMessage::RewriteSharded(shard_dbs)) => { + // C4 TopLevel cooperative fold (tokio path): + // flush + sync the BufWriter, convert to std::fs::File + // for the sync fold (same pattern as tokio per-shard), + // then reopen for appending. let _ = writer.flush().await; let _ = writer.get_ref().sync_data().await; - if let Err(e) = rewrite_aof_sharded_sync(&shard_dbs, &aof_path) { + let mut sf = writer.into_inner().into_std().await; + if let Err(e) = rewrite_aof_sharded_sync( + &shard_dbs, + &aof_path, + &rx, + &mut sf, + fold_channels.as_ref(), + ) { error!("AOF rewrite (sharded) failed: {}", e); } + // Drop sf — caller will reopen aof_path below. + drop(sf); crate::command::persistence::AOF_REWRITE_IN_PROGRESS .store(false, std::sync::atomic::Ordering::SeqCst); - let reopen_result: Result = tokio::fs::OpenOptions::new() - .create(true).append(true).open(&aof_path).await; + let reopen_result: Result = + tokio::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&aof_path) + .await; match reopen_result { Ok(f) => writer = tokio::io::BufWriter::new(f), - Err(e) => { error!("Failed to reopen AOF after rewrite: {}", e); return; } + Err(e) => { + error!("Failed to reopen AOF after rewrite: {}", e); + return; + } } + // Back-date so the channel backlog that accumulated + // during the blocking fold reaches disk within ~100ms + // + wake floor — a SIGKILL shortly after rewrite + // completion must not take the tail with it. + last_fsync = Instant::now() - std::time::Duration::from_millis(900); } // [F6] TopLevel writer never owns per-shard files — routing // bug. Self-abort so the countdown completes + flag clears. Ok(AofMessage::RewritePerShard { coord, .. }) => { - warn!("AOF TopLevel writer received RewritePerShard — routing bug; aborting"); + warn!( + "AOF TopLevel writer received RewritePerShard — routing bug; aborting" + ); coord.mark_failed(); coord.shard_done(); } @@ -369,18 +450,17 @@ pub async fn aof_writer_task( } } } - _ = interval.tick(), if fsync == FsyncPolicy::EverySec => { - if last_fsync.elapsed() >= std::time::Duration::from_secs(1) { - let _ = writer.flush().await; - let _ = writer.get_ref().sync_data().await; - last_fsync = Instant::now(); - } - } - _ = cancel.cancelled() => { + // EverySec deadline: the oldest unflushed byte reaches disk at + // most ~1.2s after it was written (1s deadline + 200ms wake + // floor). tokio's BufWriter holds up to 8KB in userspace — a + // SIGKILL takes that tail with it, so the bound must hold even + // when the recv arm is saturated with messages. + if fsync == FsyncPolicy::EverySec + && last_fsync.elapsed() >= std::time::Duration::from_secs(1) + { let _ = writer.flush().await; let _ = writer.get_ref().sync_data().await; - info!("AOF writer cancelled"); - break; + last_fsync = Instant::now(); } } } @@ -625,26 +705,32 @@ pub async fn per_shard_aof_writer_task( let _ = ack.send(AofAck::WriteFailed); continue; } - let mut header = [0u8; 12]; - header[..8].copy_from_slice(&lsn.to_le_bytes()); - header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); - if let Err(e) = writer.write_all(&header).await { - error!( - "AOF AppendSync header write error shard {}: {}", - shard_id, e - ); - write_error = true; - let _ = ack.send(AofAck::WriteFailed); - continue; - } - if let Err(e) = writer.write_all(&data).await { - error!( - "AOF AppendSync write error shard {}: {}", - shard_id, e - ); - write_error = true; - let _ = ack.send(AofAck::WriteFailed); - continue; + // H1-BARRIER: a zero-length AppendSync is an fsync + // barrier (pool::fsync_barrier) — fsync + ack only, + // NO on-disk record. A len=0 framed header would make + // replay_incr_framed reject the file as corrupt. + if !data.is_empty() { + let mut header = [0u8; 12]; + header[..8].copy_from_slice(&lsn.to_le_bytes()); + header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); + if let Err(e) = writer.write_all(&header).await { + error!( + "AOF AppendSync header write error shard {}: {}", + shard_id, e + ); + write_error = true; + let _ = ack.send(AofAck::WriteFailed); + continue; + } + if let Err(e) = writer.write_all(&data).await { + error!( + "AOF AppendSync write error shard {}: {}", + shard_id, e + ); + write_error = true; + let _ = ack.send(AofAck::WriteFailed); + continue; + } } // Test-only: skip real fsync and return FsyncFailed // immediately when the fault-injection env var is set. @@ -687,7 +773,7 @@ pub async fn per_shard_aof_writer_task( // buffered appends are durable in the OLD incr, convert // `tokio::fs::File` -> `std::fs::File` for the sync fold, // then wrap the (reopened) file back into the BufWriter. - Ok(AofMessage::RewritePerShard { shard_dbs, coord }) => { + Ok(AofMessage::RewritePerShard { shard_dbs, coord, fold_producer, fold_notifier }) => { if let Err(e) = writer.flush().await { error!( "F6 tokio per-shard rewrite: shard {} pre-fold flush \ @@ -702,6 +788,7 @@ pub async fn per_shard_aof_writer_task( let mut sf = writer.into_inner().into_std().await; let res = do_rewrite_per_shard( shard_id, &shard_dbs, &mut sf, &rx, &coord, + &fold_producer, &fold_notifier, ); // `sf` is left pointing at the committed generation // by the fold's internal barrier: NEW incr on @@ -848,6 +935,8 @@ pub async fn per_shard_aof_writer_task( let mut last_fsync = Instant::now(); let mut write_error = false; + let mut _dbg_processed: u64 = 0; + let _dbg_start = Instant::now(); // Test-only fault injection: if MOON_TEST_AOF_FSYNC_FAIL=1 is set in // the environment at writer task startup, every AppendSync ack resolves // as FsyncFailed instead of Synced. Read once before the loop so there @@ -855,7 +944,11 @@ pub async fn per_shard_aof_writer_task( let fail_fsync_for_test = std::env::var("MOON_TEST_AOF_FSYNC_FAIL").as_deref() == Ok("1"); loop { - match rx.recv() { + // Use recv_timeout so the EverySec fsync fires even when no new + // Appends arrive after a fold (or when the client stops writing). + // Without a timeout, the writer blocks forever in rx.recv() and + // the 1s fsync window never fires → data loss on kill. + match rx.recv_timeout(std::time::Duration::from_millis(50)) { // AppendSync (monoio + PerShard): framed write + fsync + ack. Ok(AofMessage::AppendSync { lsn, @@ -866,26 +959,32 @@ pub async fn per_shard_aof_writer_task( let _ = ack.send(AofAck::WriteFailed); continue; } - let mut header = [0u8; 12]; - header[..8].copy_from_slice(&lsn.to_le_bytes()); - header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); - if let Err(e) = file.write_all(&header) { - error!( - "AOF AppendSync header write failed shard {} (seq {}): {}", - shard_id, manifest.seq, e - ); - write_error = true; - let _ = ack.send(AofAck::WriteFailed); - continue; - } - if let Err(e) = file.write_all(&data) { - error!( - "AOF AppendSync write failed shard {} (seq {}): {}", - shard_id, manifest.seq, e - ); - write_error = true; - let _ = ack.send(AofAck::WriteFailed); - continue; + // H1-BARRIER: a zero-length AppendSync is an fsync barrier + // (pool::fsync_barrier) — fsync + ack only, NO on-disk + // record. A len=0 framed header would make + // replay_incr_framed reject the file as corrupt. + if !data.is_empty() { + let mut header = [0u8; 12]; + header[..8].copy_from_slice(&lsn.to_le_bytes()); + header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); + if let Err(e) = file.write_all(&header) { + error!( + "AOF AppendSync header write failed shard {} (seq {}): {}", + shard_id, manifest.seq, e + ); + write_error = true; + let _ = ack.send(AofAck::WriteFailed); + continue; + } + if let Err(e) = file.write_all(&data) { + error!( + "AOF AppendSync write failed shard {} (seq {}): {}", + shard_id, manifest.seq, e + ); + write_error = true; + let _ = ack.send(AofAck::WriteFailed); + continue; + } } // Test-only: skip real fsync and return FsyncFailed // immediately when the fault-injection env var is set. @@ -914,6 +1013,7 @@ pub async fn per_shard_aof_writer_task( if write_error { continue; } + _dbg_processed += 1; let mut header = [0u8; 12]; header[..8].copy_from_slice(&lsn.to_le_bytes()); header[8..].copy_from_slice(&(data.len() as u32).to_le_bytes()); @@ -933,38 +1033,25 @@ pub async fn per_shard_aof_writer_task( write_error = true; continue; } - match fsync { - FsyncPolicy::Always => { - let t = Instant::now(); - if let Err(e) = file.flush().and_then(|_| file.sync_data()) { - error!( - "AOF sync failed shard {} (seq {}, always): {}", - shard_id, manifest.seq, e - ); - write_error = true; - } else { - crate::admin::metrics_setup::record_aof_fsync( - t.elapsed().as_micros() as u64, - ); - } - } - FsyncPolicy::EverySec => { - if last_fsync.elapsed() >= std::time::Duration::from_secs(1) { - let t = Instant::now(); - if let Err(e) = file.flush().and_then(|_| file.sync_data()) { - error!( - "AOF sync failed shard {} (seq {}, everysec): {}", - shard_id, manifest.seq, e - ); - } else { - crate::admin::metrics_setup::record_aof_fsync( - t.elapsed().as_micros() as u64, - ); - last_fsync = Instant::now(); - } - } + // Always policy fsyncs inline. EverySec and No policies + // rely on the proactive fsync check AFTER the match block, + // which runs after every message OR timeout. Keeping EverySec + // out of the Append arm prevents it from advancing last_fsync + // after a fold completes, which would delay the next proactive + // fsync by a full second and open the EverySec durability window. + if fsync == FsyncPolicy::Always { + let t = Instant::now(); + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!( + "AOF sync failed shard {} (seq {}, always): {}", + shard_id, manifest.seq, e + ); + write_error = true; + } else { + crate::admin::metrics_setup::record_aof_fsync( + t.elapsed().as_micros() as u64 + ); } - FsyncPolicy::No => {} } } Ok(AofMessage::Rewrite(_)) | Ok(AofMessage::RewriteSharded(_)) => { @@ -979,10 +1066,21 @@ pub async fn per_shard_aof_writer_task( // then signal the coordinator; the last shard commits the // manifest. On error the old generation stays authoritative // (advance_shard did not commit the seq). - Ok(AofMessage::RewritePerShard { shard_dbs, coord }) => { - if let Err(e) = - do_rewrite_per_shard(shard_id, &shard_dbs, &mut file, &rx, &coord) - { + Ok(AofMessage::RewritePerShard { + shard_dbs, + coord, + fold_producer, + fold_notifier, + }) => { + if let Err(e) = do_rewrite_per_shard( + shard_id, + &shard_dbs, + &mut file, + &rx, + &coord, + &fold_producer, + &fold_notifier, + ) { // The fold's ShardDoneGuard already marked the rewrite // failed and decremented on this error exit (committing // new_seq with a shard missing its new base would break @@ -995,8 +1093,42 @@ pub async fn per_shard_aof_writer_task( shard_id, e ); } + // EverySec post-fold drain+fsync: the fold runs synchronously + // and does NOT update `last_fsync`. Appends that arrived during + // the fold queue in the bounded AOF channel; they land on the NEW + // incr but are NOT fsynced until the EverySec timer fires. + // Strategy: drain what's currently in the channel and fsync now + // (covers appends that landed before this drain), then set + // `last_fsync` 900ms in the past so the proactive check below + // fires within the NEXT 100ms window (≤150ms total, since the + // recv_timeout is 50ms). That second fsync covers any appends + // that arrived between the drain and that window close. + // Combined, the two fsyncs bound the post-fold EverySec window + // to ≤150ms — well within the test's 1500ms kill margin. + if !write_error { + if let Ok(mut post_drain) = + drain_pending_appends_framed(&rx, &mut file, usize::MAX) + { + if let Err(e) = sync_and_fulfill_drain( + &mut post_drain, + &mut file, + std::path::PathBuf::from(""), + ) { + error!( + "F6 per-shard rewrite: shard {} post-fold fsync \ + failed: {}. EverySec window open until next Append.", + shard_id, e + ); + } else { + // Back-date last_fsync by 900ms: the proactive check + // (threshold=1s) fires within the next 100ms, covering + // any appends that arrived after the drain above. + last_fsync = Instant::now() - std::time::Duration::from_millis(900); + } + } + } } - Ok(AofMessage::Shutdown) | Err(_) => { + Ok(AofMessage::Shutdown) | Err(flume::RecvTimeoutError::Disconnected) => { if !write_error { if let Err(e) = file.flush().and_then(|_| file.sync_data()) { error!( @@ -1006,11 +1138,43 @@ pub async fn per_shard_aof_writer_task( } } info!( - "AOF writer shard {} shutting down (monoio, seq {})", - shard_id, manifest.seq + "AOF writer shard {} shutting down (monoio, seq {}, processed {} appends in {:.3}s)", + shard_id, + manifest.seq, + _dbg_processed, + _dbg_start.elapsed().as_secs_f64() ); break; } + // Timeout: no message in the 50ms window. Fall through to + // the EverySec proactive fsync below so queued-but-unfsynced + // appends (e.g. after a fold with no new writes) are durable + // within the everysec contract even if the client goes quiet. + Err(flume::RecvTimeoutError::Timeout) => {} + } + // EverySec proactive fsync — runs after every loop iteration + // (message processed OR timeout). This is the only path that + // guarantees the 1s fsync bound when no new Appends arrive + // (e.g. after BGREWRITEAOF completes and the client stops writing). + if fsync == FsyncPolicy::EverySec + && !write_error + && last_fsync.elapsed() >= std::time::Duration::from_secs(1) + { + tracing::debug!( + "AOF EverySec proactive fsync firing shard {} (elapsed={:.3}s)", + shard_id, + last_fsync.elapsed().as_secs_f64() + ); + let t = Instant::now(); + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!( + "AOF EverySec proactive sync failed shard {} (seq {}): {}", + shard_id, manifest.seq, e + ); + } else { + crate::admin::metrics_setup::record_aof_fsync(t.elapsed().as_micros() as u64); + last_fsync = Instant::now(); + } } } } diff --git a/src/persistence/aof_manifest/shard_replay.rs b/src/persistence/aof_manifest/shard_replay.rs index d3b3d5b1..08bb5849 100644 --- a/src/persistence/aof_manifest/shard_replay.rs +++ b/src/persistence/aof_manifest/shard_replay.rs @@ -254,6 +254,14 @@ fn replay_incr_framed( break; } + // H1-BARRIER defense-in-depth: a zero-length record carries no RESP + // command (the writer skips barrier records entirely, but tolerate one + // if it ever lands on disk) — skip it rather than reject the file. + if len == 0 { + offset = payload_end; + continue; + } + // Strip the OrderedAcrossShards flag to recover the true LSN. let is_ordered = raw_lsn & crate::persistence::aof::ORDERED_LSN_FLAG != 0; let lsn = raw_lsn & !crate::persistence::aof::ORDERED_LSN_FLAG; @@ -776,6 +784,27 @@ mod tests { ); } + #[test] + fn replay_incr_framed_skips_zero_length_barrier_record() { + // H1-BARRIER defense-in-depth: a len=0 record (an fsync barrier that + // leaked to disk) carries no RESP command. Replay must SKIP it and + // continue — not reject the file as corrupt, which would brick boot. + let mut bytes = frame_entry(7, b"*1\r\n$4\r\nPING\r\n"); + bytes.extend_from_slice(&frame_entry(0, b"")); // barrier: lsn=0, len=0 + bytes.extend_from_slice(&frame_entry(21, b"*1\r\n$6\r\nDBSIZE\r\n")); + + let mut dbs: Vec = vec![crate::storage::Database::new()]; + let engine = RecordingEngine::new(); + let mut ordered: Vec = Vec::new(); + let (count, max_lsn) = replay_incr_framed(0, &mut dbs, &bytes, &engine, &mut ordered) + .expect("zero-length record must not be treated as corruption"); + + assert_eq!(count, 2, "both real commands replay; barrier is skipped"); + assert_eq!(max_lsn, 37, "next-free offset = 21 + 16"); + let calls = engine.calls.borrow(); + assert_eq!(&*calls, &["PING".to_string(), "DBSIZE".to_string()]); + } + #[test] fn replay_incr_framed_truncated_header_is_crash_eof() { // One valid entry, then a partial 5-byte header (crash mid-write). diff --git a/src/replication/master.rs b/src/replication/master.rs index 83cbb913..608741a5 100644 --- a/src/replication/master.rs +++ b/src/replication/master.rs @@ -573,7 +573,7 @@ pub async fn handle_psync_inline_single_shard( client_offset: i64, mut stream: monoio::net::TcpStream, repl_state: Arc>, - shard_databases: Arc, + _shard_databases: Arc, dispatch_tx: Rc>>>, replica_addr: std::net::SocketAddr, ) -> anyhow::Result<()> { @@ -624,13 +624,11 @@ pub async fn handle_psync_inline_single_shard( // Clone — its internal DashTable + FT/graph indices are large). let mut rdb_buf: Vec = Vec::new(); { - let db_count = shard_databases.db_count(); - let mut guards = Vec::with_capacity(db_count); - for db_idx in 0..db_count { - guards.push(shard_databases.read_db(0, db_idx)); - } - let refs: Vec<&crate::storage::Database> = guards.iter().map(|g| &**g).collect(); - crate::persistence::redis_rdb::write_rdb_refs(&refs, &mut rdb_buf); + // Shard 0 is this thread's shard — use the thread-local slice. + crate::shard::slice::with_shard(|s| { + let refs: Vec<&crate::storage::Database> = s.databases.iter().collect(); + crate::persistence::redis_rdb::write_rdb_refs(&refs, &mut rdb_buf); + }); } let header = format!("${}\r\n", rdb_buf.len()); let (wr, _) = stream.write_all(header.into_bytes()).await; diff --git a/src/server/conn/blocking.rs b/src/server/conn/blocking.rs index e6d33661..b4702210 100644 --- a/src/server/conn/blocking.rs +++ b/src/server/conn/blocking.rs @@ -97,7 +97,7 @@ pub(crate) async fn handle_blocking_command( cmd: &[u8], args: &[Frame], selected_db: usize, - shard_databases: &std::sync::Arc, + _shard_databases: &std::sync::Arc, blocking_registry: &Rc>, shard_id: usize, num_shards: usize, @@ -120,14 +120,20 @@ pub(crate) async fn handle_blocking_command( // --- Non-blocking fast path: try to get data immediately --- { - let mut guard = shard_databases.write_db(shard_id, selected_db); - for key in &keys { - let result = try_immediate_pop(cmd, &mut guard, key, args); - if let Some(frame) = result { - return frame; + // Check immediate availability via the thread-local slice. + let maybe_frame = crate::shard::slice::with_shard_db(selected_db, |db| { + for key in &keys { + let result = try_immediate_pop(cmd, db, key, args); + if result.is_some() { + return result; + } } + None + }); + if let Some(frame) = maybe_frame { + return frame; } - drop(guard); // CRITICAL before await + // Borrow released at with_shard_db boundary — safe before await. } let deadline = if timeout_secs > 0.0 { @@ -331,7 +337,7 @@ pub(crate) async fn handle_blocking_command_monoio( cmd: &[u8], args: &[Frame], selected_db: usize, - shard_databases: &std::sync::Arc, + _shard_databases: &std::sync::Arc, blocking_registry: &Rc>, shard_id: usize, num_shards: usize, @@ -354,15 +360,19 @@ pub(crate) async fn handle_blocking_command_monoio( }; // --- Non-blocking fast path: try to get data immediately --- + // Use with_shard_db (thread-local ShardSlice) — no RwLock guard needed. { - let mut guard = shard_databases.write_db(shard_id, selected_db); - for key in &keys { - let result = try_immediate_pop(cmd, &mut guard, key, args); - if let Some(frame) = result { - return frame; + let immediate_result = crate::shard::slice::with_shard_db(selected_db, |db| { + for key in &keys { + if let Some(frame) = try_immediate_pop(cmd, db, key, args) { + return Some(frame); + } } + None + }); + if let Some(frame) = immediate_result { + return frame; } - drop(guard); // CRITICAL before await } let deadline = if timeout_secs > 0.0 { @@ -1250,32 +1260,42 @@ pub(crate) fn try_inline_dispatch( // proven safe just above. let consumed = key_end_crlf; let key_bytes = &buf[key_start..key_end]; - let guard = shard_databases.read_db(shard_id, selected_db); - // Hot-key sampling: the inline path bypasses both dispatchers, so it - // must feed the sketch itself. tick() is one relaxed fetch_add. - if guard.hot_keys().tick() { - guard.hot_keys().observe(key_bytes); + // Hot-key sampling + lookup via thread-local slice — no lock needed. + enum GetResult { + Found(Vec), + WrongType, + Miss, } - match guard.get_if_alive(key_bytes, now_ms) { - Some(entry) => match entry.value.as_bytes() { - Some(val) => { - write_buf.extend_from_slice(b"$"); - let mut itoa_buf = itoa::Buffer::new(); - write_buf.extend_from_slice(itoa_buf.format(val.len()).as_bytes()); - write_buf.extend_from_slice(b"\r\n"); - write_buf.extend_from_slice(val); - write_buf.extend_from_slice(b"\r\n"); - } + let (inline_result, cold_loc) = crate::shard::slice::with_shard_db(selected_db, |db| { + if db.hot_keys().tick() { + db.hot_keys().observe(key_bytes); + } + match db.get_if_alive(key_bytes, now_ms) { + Some(entry) => match entry.value.as_bytes() { + Some(val) => (GetResult::Found(val.to_vec()), None), + None => (GetResult::WrongType, None), + }, None => { - write_buf.extend_from_slice( - b"-WRONGTYPE Operation against a key holding the wrong kind of value\r\n", - ); + let loc = db.cold_lookup_location(key_bytes); + (GetResult::Miss, loc) } - }, - None => { - // Cold storage fallback - let cold_loc = guard.cold_lookup_location(key_bytes); - drop(guard); + } + }); + match inline_result { + GetResult::Found(val) => { + write_buf.extend_from_slice(b"$"); + let mut itoa_buf = itoa::Buffer::new(); + write_buf.extend_from_slice(itoa_buf.format(val.len()).as_bytes()); + write_buf.extend_from_slice(b"\r\n"); + write_buf.extend_from_slice(&val); + write_buf.extend_from_slice(b"\r\n"); + } + GetResult::WrongType => { + write_buf.extend_from_slice( + b"-WRONGTYPE Operation against a key holding the wrong kind of value\r\n", + ); + } + GetResult::Miss => { let cold = cold_loc.and_then(|(loc, shard_dir)| { crate::storage::tiered::cold_read::read_cold_entry_at(&shard_dir, loc, now_ms) }); @@ -1299,7 +1319,6 @@ pub(crate) fn try_inline_dispatch( return 1; } } - drop(guard); let _ = read_buf.split_to(consumed); return 1; } @@ -1348,29 +1367,31 @@ pub(crate) fn try_inline_dispatch( // We must not index into `buf` after this point — use `frozen` instead. let frozen = read_buf.split_to(consumed).freeze(); - // Eviction check + write under exclusive lock + // Eviction check + write via the thread-local slice (no lock needed). { let rt = runtime_config.read(); - let mut guard = shard_databases.write_db(shard_id, selected_db); let budget = shard_databases.elastic_budget(shard_id); - if crate::storage::eviction::try_evict_if_needed_budget(&mut guard, &rt, budget).is_err() { + let oom = crate::shard::slice::with_shard_db(selected_db, |db| { + crate::storage::eviction::try_evict_if_needed_budget(db, &rt, budget).is_err() + }); + drop(rt); + if oom { write_buf .extend_from_slice(b"-OOM command not allowed when used memory > 'maxmemory'\r\n"); return 1; } - drop(rt); let key = frozen.slice(key_start..key_end); let value = frozen.slice(val_start..val_end); - // Hot-key sampling: the inline path bypasses both dispatchers, so it - // must feed the sketch itself. tick() is one relaxed fetch_add. - if guard.hot_keys().tick() { - guard.hot_keys().observe(&key); - } - let mut entry = crate::storage::entry::Entry::new_string(value); - entry.set_last_access(guard.now()); - entry.set_access_counter(5); - guard.set(key, entry); + crate::shard::slice::with_shard_db(selected_db, |db| { + if db.hot_keys().tick() { + db.hot_keys().observe(&key); + } + let mut entry = crate::storage::entry::Entry::new_string(value.clone()); + entry.set_last_access(db.now()); + entry.set_access_counter(5); + db.set(key.clone(), entry); + }); } // AOF: reuse the frozen RESP bytes directly (Arc clone, zero-copy). diff --git a/src/server/conn/handler_monoio/dispatch.rs b/src/server/conn/handler_monoio/dispatch.rs index 69db1a8f..44708e05 100644 --- a/src/server/conn/handler_monoio/dispatch.rs +++ b/src/server/conn/handler_monoio/dispatch.rs @@ -168,35 +168,19 @@ pub(super) fn try_handle_evalsha( if !cmd.eq_ignore_ascii_case(b"EVALSHA") { return false; } - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - let db_count = s.databases.len(); - crate::scripting::handle_evalsha( - &ctx.lua, - &ctx.script_cache, - cmd_args, - &mut s.databases[conn.selected_db], - ctx.shard_id, - ctx.num_shards, - conn.selected_db, - db_count, - ) - }) - } else { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let db_count = ctx.shard_databases.db_count(); - let db = &mut guard; + let response = crate::shard::slice::with_shard(|s| { + let db_count = s.databases.len(); crate::scripting::handle_evalsha( &ctx.lua, &ctx.script_cache, cmd_args, - db, + &mut s.databases[conn.selected_db], ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, ) - }; + }); responses.push(response); true } @@ -216,26 +200,9 @@ pub(super) fn try_handle_eval( if !cmd.eq_ignore_ascii_case(b"EVAL") { return false; } - // Phase 2a: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - let db_count = s.databases.len(); - let db = &mut s.databases[conn.selected_db]; - crate::scripting::handle_eval( - &ctx.lua, - &ctx.script_cache, - cmd_args, - db, - ctx.shard_id, - ctx.num_shards, - conn.selected_db, - db_count, - ) - }) - } else { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let db_count = ctx.shard_databases.db_count(); - let db = &mut guard; + let response = crate::shard::slice::with_shard(|s| { + let db_count = s.databases.len(); + let db = &mut s.databases[conn.selected_db]; crate::scripting::handle_eval( &ctx.lua, &ctx.script_cache, @@ -246,7 +213,7 @@ pub(super) fn try_handle_eval( conn.selected_db, db_count, ) - }; + }); responses.push(response); true } @@ -671,27 +638,13 @@ pub(super) fn try_handle_info( if !cmd.eq_ignore_ascii_case(b"INFO") { return false; } - // Phase 2a: gate on is_initialized(); new path uses ShardSlice directly. - let response_text = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| { - let resp_frame = conn_cmd::info_readonly(db, cmd_args); - match resp_frame { - Frame::BulkString(b) => String::from_utf8_lossy(&b).to_string(), - _ => String::new(), - } - }) - } else { - let guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); - let text = { - let resp_frame = conn_cmd::info_readonly(&guard, cmd_args); - match resp_frame { - Frame::BulkString(b) => String::from_utf8_lossy(&b).to_string(), - _ => String::new(), - } - }; - drop(guard); - text - }; + let response_text = crate::shard::slice::with_shard_db(conn.selected_db, |db| { + let resp_frame = conn_cmd::info_readonly(db, cmd_args); + match resp_frame { + Frame::BulkString(b) => String::from_utf8_lossy(&b).to_string(), + _ => String::new(), + } + }); let mut response_text = response_text; if let Some(ref rs) = ctx.repl_state { if let Ok(rs_guard) = rs.try_read() { @@ -1202,64 +1155,34 @@ pub(super) fn try_handle_functions( return true; } if cmd.eq_ignore_ascii_case(b"FCALL") { - // Phase 2a: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - let db_count = s.databases.len(); - crate::command::functions::handle_fcall( - &func_registry.borrow(), - cmd_args, - &mut s.databases[conn.selected_db], - ctx.shard_id, - ctx.num_shards, - conn.selected_db, - db_count, - ) - }) - } else { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let db_count = ctx.shard_databases.db_count(); + let response = crate::shard::slice::with_shard(|s| { + let db_count = s.databases.len(); crate::command::functions::handle_fcall( &func_registry.borrow(), cmd_args, - &mut guard, + &mut s.databases[conn.selected_db], ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, ) - }; + }); responses.push(response); return true; } if cmd.eq_ignore_ascii_case(b"FCALL_RO") { - // Phase 2a: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - let db_count = s.databases.len(); - crate::command::functions::handle_fcall_ro( - &func_registry.borrow(), - cmd_args, - &mut s.databases[conn.selected_db], - ctx.shard_id, - ctx.num_shards, - conn.selected_db, - db_count, - ) - }) - } else { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let db_count = ctx.shard_databases.db_count(); + let response = crate::shard::slice::with_shard(|s| { + let db_count = s.databases.len(); crate::command::functions::handle_fcall_ro( &func_registry.borrow(), cmd_args, - &mut guard, + &mut s.databases[conn.selected_db], ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, ) - }; + }); responses.push(response); return true; } diff --git a/src/server/conn/handler_monoio/ft.rs b/src/server/conn/handler_monoio/ft.rs index 744cded9..4dc7040b 100644 --- a/src/server/conn/handler_monoio/ft.rs +++ b/src/server/conn/handler_monoio/ft.rs @@ -109,6 +109,9 @@ pub(super) async fn try_handle_ft_command( top_k, offset: limit_offset, count: limit_count, + // CHANGE D: thread the parsed FILTER tree (if any) into + // HybridQuery so execute_hybrid_search_local / scatter + // can apply it pre-RRF on both branches (CHANGE B/F). filter: partial.filter, }; // Phase 171 HYB-02 / SCAT-02: resolve AS_OF / TXN LSN @@ -117,7 +120,6 @@ pub(super) async fn try_handle_ft_command( let as_of_lsn = match resolve_ft_search_as_of_lsn( cmd_args, Some(&ctx.shard_databases), - ctx.shard_id, conn.active_cross_txn.as_ref(), ) { Ok(lsn) => lsn, @@ -222,46 +224,42 @@ pub(super) async fn try_handle_ft_command( // We use the TextIndex's own field_analyzers (same pipeline used at index time). type ParseResult = Result<(Vec, Option), String>; - let parse_result: ParseResult = { - let ts = ctx.shard_databases.text_store(ctx.shard_id); - match ts.get_index(&index_name) { + // Parse query via thread-local slice (text_store is !Send, no lock needed). + let parse_result: ParseResult = crate::shard::slice::with_shard(|s| { + match s.text_store.get_index(&index_name) { None => Err("ERR no such index".to_owned()), - Some(text_index) => { - match text_index.field_analyzers.first() { - None => Err("ERR index has no TEXT fields".to_owned()), - Some(analyzer) => { - // analyzer borrows text_index which borrows ts — all in this block. - let parsed = crate::command::vector_search::parse_text_query( - &query_str, analyzer, - ); - match parsed { - Err(e) => Err(e.to_owned()), - Ok(clause) => { - let field_idx = match &clause.field_name { - None => Ok(None), - Some(field_name) => match text_index - .text_fields - .iter() - .position(|f| { - f.field_name.as_ref().eq_ignore_ascii_case( - field_name.as_ref(), - ) - }) { + Some(text_index) => match text_index.field_analyzers.first() { + None => Err("ERR index has no TEXT fields".to_owned()), + Some(analyzer) => { + let parsed = crate::command::vector_search::parse_text_query( + &query_str, analyzer, + ); + match parsed { + Err(e) => Err(e.to_owned()), + Ok(clause) => { + let field_idx = match &clause.field_name { + None => Ok(None), + Some(field_name) => { + match text_index.text_fields.iter().position(|f| { + f.field_name + .as_ref() + .eq_ignore_ascii_case(field_name.as_ref()) + }) { Some(idx) => Ok(Some(idx)), None => Err(format!( "ERR unknown field '{}'", String::from_utf8_lossy(field_name) )), - }, - }; - field_idx.map(|idx| (clause.terms, idx)) - } + } + } + }; + field_idx.map(|idx| (clause.terms, idx)) } } } - } + }, } - }; // MutexGuard dropped here + }); // with_shard borrow released here let (query_terms, field_idx) = match parse_result { Ok(t) => t, @@ -323,7 +321,6 @@ pub(super) async fn try_handle_ft_command( match resolve_ft_search_as_of_lsn( cmd_args, Some(&ctx.shard_databases), - ctx.shard_id, conn.active_cross_txn.as_ref(), ) { Err(err_frame) => err_frame, @@ -368,76 +365,48 @@ pub(super) async fn try_handle_ft_command( return true; } if cmd.eq_ignore_ascii_case(b"FT.INFO") { - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::vector_search::ft_info(&s.vector_store, &s.text_store, cmd_args) - }) - } else { - let vs = ctx.shard_databases.vector_store(ctx.shard_id); - let ts = ctx.shard_databases.text_store(ctx.shard_id); - crate::command::vector_search::ft_info(&vs, &ts, cmd_args) - }; + let response = crate::shard::slice::with_shard(|s| { + crate::command::vector_search::ft_info(&s.vector_store, &s.text_store, cmd_args) + }); responses.push(response); return true; } if cmd.eq_ignore_ascii_case(b"FT._LIST") { - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::vector_search::ft_list(&s.vector_store) - }) - } else { - let vs = ctx.shard_databases.vector_store(ctx.shard_id); - crate::command::vector_search::ft_list(&vs) - }; + let response = crate::shard::slice::with_shard(|s| { + crate::command::vector_search::ft_list(&s.vector_store) + }); responses.push(response); return true; } if cmd.eq_ignore_ascii_case(b"FT.COMPACT") { - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::vector_search::ft_compact( - &mut s.vector_store, - &mut s.text_store, - cmd_args, - ) - }) - } else { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - let mut ts = ctx.shard_databases.text_store(ctx.shard_id); - crate::command::vector_search::ft_compact(&mut vs, &mut ts, cmd_args) - }; + let response = crate::shard::slice::with_shard(|s| { + crate::command::vector_search::ft_compact( + &mut s.vector_store, + &mut s.text_store, + cmd_args, + ) + }); responses.push(response); return true; } if cmd.eq_ignore_ascii_case(b"FT.CACHESEARCH") { - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::vector_search::cache_search::ft_cachesearch( - &mut s.vector_store, - cmd_args, - ) - }) - } else { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - crate::command::vector_search::cache_search::ft_cachesearch(&mut vs, cmd_args) - }; + let response = crate::shard::slice::with_shard(|s| { + crate::command::vector_search::cache_search::ft_cachesearch( + &mut s.vector_store, + cmd_args, + ) + }); responses.push(response); return true; } if cmd.eq_ignore_ascii_case(b"FT.CONFIG") { - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::vector_search::ft_config( - &mut s.vector_store, - &mut s.text_store, - cmd_args, - ) - }) - } else { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - let mut ts = ctx.shard_databases.text_store(ctx.shard_id); - crate::command::vector_search::ft_config(&mut vs, &mut ts, cmd_args) - }; + let response = crate::shard::slice::with_shard(|s| { + crate::command::vector_search::ft_config( + &mut s.vector_store, + &mut s.text_store, + cmd_args, + ) + }); responses.push(response); return true; } @@ -446,61 +415,19 @@ pub(super) async fn try_handle_ft_command( || cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") || cmd.eq_ignore_ascii_case(b"FT.EXPAND") { - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") { - crate::command::vector_search::recommend::ft_recommend( - &mut s.vector_store, - cmd_args, - Some(&mut s.databases[0]), - ) - } else if cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") { - #[cfg(feature = "graph")] - { - crate::command::vector_search::navigate::ft_navigate( - &mut s.vector_store, - Some(&s.graph_store), - cmd_args, - None, - ) - } - #[cfg(not(feature = "graph"))] - { - Frame::Error(Bytes::from_static( - b"ERR FT.NAVIGATE requires graph feature", - )) - } - } else { - // FT.EXPAND - #[cfg(feature = "graph")] - { - crate::command::vector_search::ft_expand(&s.graph_store, cmd_args) - } - #[cfg(not(feature = "graph"))] - { - Frame::Error(Bytes::from_static( - b"ERR FT.EXPAND requires graph feature", - )) - } - } - }) - } else { - let sdb = &ctx.shard_databases; - let mut vs = sdb.vector_store(ctx.shard_id); + let response = crate::shard::slice::with_shard(|s| { if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") { - let mut db_guard = sdb.write_db(ctx.shard_id, 0); crate::command::vector_search::recommend::ft_recommend( - &mut vs, + &mut s.vector_store, cmd_args, - Some(&mut *db_guard), + Some(&mut s.databases[0]), ) } else if cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") { #[cfg(feature = "graph")] { - let graph_guard = sdb.graph_store_read(ctx.shard_id); crate::command::vector_search::navigate::ft_navigate( - &mut vs, - Some(&graph_guard), + &mut s.vector_store, + Some(&s.graph_store), cmd_args, None, ) @@ -515,15 +442,14 @@ pub(super) async fn try_handle_ft_command( // FT.EXPAND #[cfg(feature = "graph")] { - let graph_guard = sdb.graph_store_read(ctx.shard_id); - crate::command::vector_search::ft_expand(&graph_guard, cmd_args) + crate::command::vector_search::ft_expand(&s.graph_store, cmd_args) } #[cfg(not(feature = "graph"))] { Frame::Error(Bytes::from_static(b"ERR FT.EXPAND requires graph feature")) } } - }; + }); responses.push(response); return true; } @@ -636,40 +562,22 @@ pub(super) async fn try_handle_ft_command( offset.saturating_add(count) } .max(1); - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - match s.text_store.get_index(&index_name) { - None => Frame::Error(Bytes::from_static( - b"ERR no such index", - )), - Some(text_index) => { - let results = crate::command::vector_search::ft_text_search::execute_query_on_index( - text_index, &clause, None, None, top_k, - ); - crate::command::vector_search::ft_text_search::build_text_response( - &results, offset, count, - ) - } - } - }) - } else { - let ts_guard = - ctx.shard_databases.text_store(ctx.shard_id); - let r = match ts_guard.get_index(&index_name) { - None => Frame::Error(Bytes::from_static( - b"ERR no such index", - )), - Some(text_index) => { - let results = crate::command::vector_search::ft_text_search::execute_query_on_index( + let response = crate::shard::slice::with_shard(|s| match s + .text_store + .get_index(&index_name) + { + None => Frame::Error(Bytes::from_static( + b"ERR no such index", + )), + Some(text_index) => { + let results = crate::command::vector_search::ft_text_search::execute_query_on_index( text_index, &clause, None, None, top_k, ); - crate::command::vector_search::ft_text_search::build_text_response( + crate::command::vector_search::ft_text_search::build_text_response( &results, offset, count, ) - } - }; - r - }; + } + }); let mut response = response; if let Some(ws_id) = conn.workspace_id.as_ref() { strip_workspace_prefix_from_response( @@ -704,130 +612,31 @@ pub(super) async fn try_handle_ft_command( let summarize_opts = crate::command::vector_search::parse_summarize_clause(cmd_args); let needs_db = highlight_opts.is_some() || summarize_opts.is_some(); - // Steps 2-10: acquire stores and execute. Result - // where Ok = response to push, Err = error frame with early-return. - // Phase 2a: gate on is_initialized(); new path uses ShardSlice directly. + // Steps 2-10: acquire stores and execute via ShardSlice. let text_search_result: Result = - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - // Step 2-3: resolve index from text_store. - let text_index = match s.text_store.get_index(&index_name) { - Some(idx) => idx, - None => { - return Err(Frame::Error(Bytes::from_static( - b"ERR no such index", - ))); - } - }; - // Step 4: ensure TEXT fields. - if text_index.text_fields.is_empty() { - return Err(Frame::Error(Bytes::from_static( - b"ERR index has no TEXT fields", - ))); - } - // Step 5: parse query. - let analyzer = match text_index.field_analyzers.first() { - Some(a) => a, - None => { - return Err(Frame::Error(Bytes::from_static( - b"ERR index has no TEXT fields", - ))); - } - }; - let clause = - match crate::command::vector_search::parse_text_query( - query_bytes.as_ref(), - analyzer, - ) { - Ok(c) => c, - Err(msg) => { - return Err(Frame::Error( - Bytes::copy_from_slice(msg.as_bytes()), - )); - } - }; - // Step 5b: resolve field_idx. - let field_idx = match &clause.field_name { - None => None, - Some(field_name) => { - match text_index.text_fields.iter().position(|f| { - f.field_name - .as_ref() - .eq_ignore_ascii_case(field_name.as_ref()) - }) { - Some(idx) => Some(idx), - None => { - let bad = field_name.clone(); - return Err(Frame::Error(Bytes::from( - format!( - "ERR unknown field '{}'", - String::from_utf8_lossy(&bad) - ), - ))); - } - } - } - }; - let query_terms = clause.terms; - // Step 10: execute. - let mut response = - crate::command::vector_search::execute_text_search_local( - &s.text_store, - &index_name, - field_idx, - &query_terms, - top_k, - offset, - count, - ); - // Step 9+10b: optional post-processing with db access. - if needs_db { - let db = &s.databases[0]; - let term_strings: Vec = query_terms - .iter() - .map(|qt| qt.text.clone()) - .collect(); - crate::command::vector_search::apply_post_processing( - &mut response, - &term_strings, - text_index, - db, - highlight_opts.as_ref(), - summarize_opts.as_ref(), - ); - } - Ok(response) - }) - } else { - // Step 2: acquire ts guard (single-shard monoio). - let ts_guard = ctx.shard_databases.text_store(ctx.shard_id); - // Step 3: resolve index. - let text_index = match ts_guard.get_index(&index_name) { + crate::shard::slice::with_shard(|s| { + // Step 2-3: resolve index from text_store. + let text_index = match s.text_store.get_index(&index_name) { Some(idx) => idx, None => { - return { - responses.push(Frame::Error(Bytes::from_static( - b"ERR no such index", - ))); - true - }; + return Err(Frame::Error(Bytes::from_static( + b"ERR no such index", + ))); } }; - // Step 4: ensure at least one TEXT field. + // Step 4: ensure TEXT fields. if text_index.text_fields.is_empty() { - responses.push(Frame::Error(Bytes::from_static( + return Err(Frame::Error(Bytes::from_static( b"ERR index has no TEXT fields", ))); - return true; } - // Step 5: parse the query via the index's first analyzer. + // Step 5: parse query. let analyzer = match text_index.field_analyzers.first() { Some(a) => a, None => { - responses.push(Frame::Error(Bytes::from_static( + return Err(Frame::Error(Bytes::from_static( b"ERR index has no TEXT fields", ))); - return true; } }; let clause = @@ -837,10 +646,9 @@ pub(super) async fn try_handle_ft_command( ) { Ok(c) => c, Err(msg) => { - responses.push(Frame::Error( - Bytes::copy_from_slice(msg.as_bytes()), - )); - return true; + return Err(Frame::Error(Bytes::copy_from_slice( + msg.as_bytes(), + ))); } }; // Step 5b: resolve field_idx. @@ -854,29 +662,22 @@ pub(super) async fn try_handle_ft_command( }) { Some(idx) => Some(idx), None => { - let bad_name = field_name.clone(); - responses.push(Frame::Error(Bytes::from( + let bad = field_name.clone(); + return Err(Frame::Error(Bytes::from( format!( "ERR unknown field '{}'", - String::from_utf8_lossy(&bad_name) + String::from_utf8_lossy(&bad) ), ))); - return true; } } } }; let query_terms = clause.terms; - // Step 9: acquire DB read guard iff post-processing needed. - let db_guard_opt = if needs_db { - Some(ctx.shard_databases.read_db(ctx.shard_id, 0)) - } else { - None - }; - // Step 10: execute + optional post-processing. + // Step 10: execute. let mut response = crate::command::vector_search::execute_text_search_local( - &ts_guard, + &s.text_store, &index_name, field_idx, &query_terms, @@ -884,22 +685,22 @@ pub(super) async fn try_handle_ft_command( offset, count, ); - if let Some(ref db_guard) = db_guard_opt { + // Step 9+10b: optional post-processing with db access. + if needs_db { + let db = &s.databases[0]; let term_strings: Vec = query_terms.iter().map(|qt| qt.text.clone()).collect(); crate::command::vector_search::apply_post_processing( &mut response, &term_strings, text_index, - db_guard, + db, highlight_opts.as_ref(), summarize_opts.as_ref(), ); } - drop(db_guard_opt); - drop(ts_guard); Ok(response) - }; + }); let mut response = match text_search_result { Ok(r) => r, Err(e) => { @@ -917,163 +718,32 @@ pub(super) async fn try_handle_ft_command( } } } - // Phase 2a: Option A — gate on is_initialized(). New path uses ShardSlice - // directly (no RwLock); old path uses ctx.shard_databases guards. Dead code - // until Phase 4 wires init_shard() at shard startup. - // - // FT.SEARCH AS_OF needs shard_databases for vector_store LSN resolution; - // in the new path we resolve it inside the closure via s.vector_store. - let response = if crate::shard::slice::is_initialized() { - // Phase 2a: AS_OF temporal resolution still uses ctx.shard_databases - // (temporal_registry is migrated in Phase 2b, not here). This is - // intentional — both paths share the same AS_OF resolution logic. - let as_of_lsn_result = if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { - resolve_ft_search_as_of_lsn( - cmd_args, - Some(&ctx.shard_databases), - ctx.shard_id, - conn.active_cross_txn.as_ref(), - ) - } else { - Ok(0u64) - }; - let as_of_lsn = match as_of_lsn_result { - Ok(lsn) => lsn, - Err(err_frame) => { - responses.push(err_frame); - return true; - } - }; - crate::shard::slice::with_shard(|s| { - if cmd.eq_ignore_ascii_case(b"FT.CREATE") { - crate::command::vector_search::ft_create( - &mut s.vector_store, - &mut s.text_store, - cmd_args, - ) - } else if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { - let has_session = cmd_args.iter().any(|a| { - if let Frame::BulkString(b) = a { - b.eq_ignore_ascii_case(b"SESSION") - } else { - false - } - }); - if has_session { - crate::command::vector_search::ft_search( - &mut s.vector_store, - cmd_args, - Some(&mut s.databases[0]), - Some(&s.text_store), - as_of_lsn, - ) - } else { - crate::command::vector_search::ft_search( - &mut s.vector_store, - cmd_args, - None, - Some(&s.text_store), - as_of_lsn, - ) - } - } else if cmd.eq_ignore_ascii_case(b"FT.DROPINDEX") { - crate::command::vector_search::ft_dropindex( - &mut s.vector_store, - &mut s.text_store, - Some(&mut s.databases[0]), - cmd_args, - ) - } else if cmd.eq_ignore_ascii_case(b"FT.INFO") { - crate::command::vector_search::ft_info(&s.vector_store, &s.text_store, cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT._LIST") { - crate::command::vector_search::ft_list(&s.vector_store) - } else if cmd.eq_ignore_ascii_case(b"FT.COMPACT") { - crate::command::vector_search::ft_compact( - &mut s.vector_store, - &mut s.text_store, - cmd_args, - ) - } else if cmd.eq_ignore_ascii_case(b"FT.CACHESEARCH") { - crate::command::vector_search::cache_search::ft_cachesearch( - &mut s.vector_store, - cmd_args, - ) - } else if cmd.eq_ignore_ascii_case(b"FT.CONFIG") { - crate::command::vector_search::ft_config( - &mut s.vector_store, - &mut s.text_store, - cmd_args, - ) - } else if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") { - crate::command::vector_search::recommend::ft_recommend( - &mut s.vector_store, - cmd_args, - Some(&mut s.databases[0]), - ) - } else if cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") { - #[cfg(feature = "graph")] - { - crate::command::vector_search::navigate::ft_navigate( - &mut s.vector_store, - Some(&s.graph_store), - cmd_args, - None, - ) - } - #[cfg(not(feature = "graph"))] - { - Frame::Error(Bytes::from_static( - b"ERR FT.NAVIGATE requires graph feature", - )) - } - } else if cmd.eq_ignore_ascii_case(b"FT.EXPAND") { - #[cfg(feature = "graph")] - { - crate::command::vector_search::ft_expand(&s.graph_store, cmd_args) - } - #[cfg(not(feature = "graph"))] - { - Frame::Error(Bytes::from_static(b"ERR FT.EXPAND requires graph feature")) - } - } else if cmd.eq_ignore_ascii_case(b"FT.INVALIDATE_RANGE") { - #[cfg(feature = "text-index")] - { - crate::command::vector_search::ft_invalidate_range( - &mut s.text_store, - cmd_args, - ) - } - #[cfg(not(feature = "text-index"))] - { - Frame::Error(Bytes::from_static( - b"ERR FT.INVALIDATE_RANGE requires text-index feature", - )) - } - } else { - Frame::Error(Bytes::from_static(b"ERR unknown FT.* command")) - } - }) + // AS_OF temporal resolution uses ctx.shard_databases (temporal_registry + // migrated in Phase 2b). Both paths share the same AS_OF resolution logic. + let as_of_lsn_result = if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { + resolve_ft_search_as_of_lsn( + cmd_args, + Some(&ctx.shard_databases), + conn.active_cross_txn.as_ref(), + ) } else { - let shard_databases_ref = &ctx.shard_databases; - let mut vs = shard_databases_ref.vector_store(ctx.shard_id); - let mut ts = shard_databases_ref.text_store(ctx.shard_id); + Ok(0u64) + }; + let as_of_lsn = match as_of_lsn_result { + Ok(lsn) => lsn, + Err(err_frame) => { + responses.push(err_frame); + return true; + } + }; + let response = crate::shard::slice::with_shard(|s| { if cmd.eq_ignore_ascii_case(b"FT.CREATE") { - crate::command::vector_search::ft_create(&mut vs, &mut ts, cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { - // Resolve AS_OF temporal clause + TXN snapshot precedence (TEMP-04, ACID-09). - let as_of_lsn = match resolve_ft_search_as_of_lsn( + crate::command::vector_search::ft_create( + &mut s.vector_store, + &mut s.text_store, cmd_args, - Some(shard_databases_ref), - ctx.shard_id, - conn.active_cross_txn.as_ref(), - ) { - Ok(lsn) => lsn, - Err(err_frame) => { - responses.push(err_frame); - return true; - } - }; - // Detect SESSION keyword to provide database access for sorted set tracking + ) + } else if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { let has_session = cmd_args.iter().any(|a| { if let Frame::BulkString(b) = a { b.eq_ignore_ascii_case(b"SESSION") @@ -1082,55 +752,62 @@ pub(super) async fn try_handle_ft_command( } }); if has_session { - let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); crate::command::vector_search::ft_search( - &mut vs, + &mut s.vector_store, cmd_args, - Some(&mut *db_guard), - Some(&*ts), + Some(&mut s.databases[0]), + Some(&s.text_store), as_of_lsn, ) } else { crate::command::vector_search::ft_search( - &mut vs, + &mut s.vector_store, cmd_args, None, - Some(&*ts), + Some(&s.text_store), as_of_lsn, ) } } else if cmd.eq_ignore_ascii_case(b"FT.DROPINDEX") { - let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); crate::command::vector_search::ft_dropindex( - &mut vs, - &mut ts, - Some(&mut *db_guard), + &mut s.vector_store, + &mut s.text_store, + Some(&mut s.databases[0]), cmd_args, ) } else if cmd.eq_ignore_ascii_case(b"FT.INFO") { - crate::command::vector_search::ft_info(&vs, &ts, cmd_args) + crate::command::vector_search::ft_info(&s.vector_store, &s.text_store, cmd_args) } else if cmd.eq_ignore_ascii_case(b"FT._LIST") { - crate::command::vector_search::ft_list(&vs) + crate::command::vector_search::ft_list(&s.vector_store) } else if cmd.eq_ignore_ascii_case(b"FT.COMPACT") { - crate::command::vector_search::ft_compact(&mut vs, &mut ts, cmd_args) + crate::command::vector_search::ft_compact( + &mut s.vector_store, + &mut s.text_store, + cmd_args, + ) } else if cmd.eq_ignore_ascii_case(b"FT.CACHESEARCH") { - crate::command::vector_search::cache_search::ft_cachesearch(&mut vs, cmd_args) + crate::command::vector_search::cache_search::ft_cachesearch( + &mut s.vector_store, + cmd_args, + ) } else if cmd.eq_ignore_ascii_case(b"FT.CONFIG") { - crate::command::vector_search::ft_config(&mut vs, &mut ts, cmd_args) + crate::command::vector_search::ft_config( + &mut s.vector_store, + &mut s.text_store, + cmd_args, + ) } else if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") { - let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); crate::command::vector_search::recommend::ft_recommend( - &mut vs, + &mut s.vector_store, cmd_args, - Some(&mut *db_guard), + Some(&mut s.databases[0]), ) } else if cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") { #[cfg(feature = "graph")] { - let graph_guard = shard_databases_ref.graph_store_read(ctx.shard_id); crate::command::vector_search::navigate::ft_navigate( - &mut vs, - Some(&graph_guard), + &mut s.vector_store, + Some(&s.graph_store), cmd_args, None, ) @@ -1144,8 +821,7 @@ pub(super) async fn try_handle_ft_command( } else if cmd.eq_ignore_ascii_case(b"FT.EXPAND") { #[cfg(feature = "graph")] { - let graph_guard = shard_databases_ref.graph_store_read(ctx.shard_id); - crate::command::vector_search::ft_expand(&graph_guard, cmd_args) + crate::command::vector_search::ft_expand(&s.graph_store, cmd_args) } #[cfg(not(feature = "graph"))] { @@ -1154,7 +830,7 @@ pub(super) async fn try_handle_ft_command( } else if cmd.eq_ignore_ascii_case(b"FT.INVALIDATE_RANGE") { #[cfg(feature = "text-index")] { - crate::command::vector_search::ft_invalidate_range(&mut ts, cmd_args) + crate::command::vector_search::ft_invalidate_range(&mut s.text_store, cmd_args) } #[cfg(not(feature = "text-index"))] { @@ -1165,7 +841,7 @@ pub(super) async fn try_handle_ft_command( } else { Frame::Error(Bytes::from_static(b"ERR unknown FT.* command")) } - }; + }); let mut response = response; if let Some(ws_id) = conn.workspace_id.as_ref() { strip_workspace_prefix_from_response(ws_id, cmd, &mut response); diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index 3c4c2709..6fc8ea25 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -19,7 +19,7 @@ use std::collections::HashMap; use std::rc::Rc; use crate::command::metadata; -use crate::command::{DispatchResult, dispatch, dispatch_read, is_dispatch_read_supported}; +use crate::command::{DispatchResult, dispatch, dispatch_read}; use crate::persistence::aof; use crate::protocol::Frame; use crate::shard::dispatch::key_to_shard; @@ -597,15 +597,9 @@ pub(crate) async fn handle_connection_sharded_monoio< std::collections::HashMap::new(); // Refresh time once per batch — sub-millisecond accuracy not needed per-command. - // Phase 2a: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| { - db.refresh_now_from_cache(&ctx.cached_clock); - }); - } else { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - guard.refresh_now_from_cache(&ctx.cached_clock); - } + crate::shard::slice::with_shard_db(conn.selected_db, |db| { + db.refresh_now_from_cache(&ctx.cached_clock); + }); // Batch-level eviction gate: snapshot `maxmemory != 0` once per batch // (and cache the spill-sender presence check). When neither is set — @@ -879,7 +873,7 @@ pub(crate) async fn handle_connection_sharded_monoio< if txn::try_handle_txn_begin(cmd, cmd_args, &mut conn, ctx, &mut responses) { continue; } - if txn::try_handle_txn_commit(cmd, cmd_args, &mut conn, ctx, &mut responses) { + if txn::try_handle_txn_commit(cmd, cmd_args, &mut conn, ctx, &mut responses).await { continue; } if txn::try_handle_txn_abort(cmd, cmd_args, &mut conn, ctx, &mut responses).await { @@ -896,12 +890,14 @@ pub(crate) async fn handle_connection_sharded_monoio< } // --- WS.* --- - if write::try_handle_ws_command(cmd, cmd_args, &mut conn, ctx, &mut responses) { + if write::try_handle_ws_command(cmd, cmd_args, &mut conn, ctx, &mut responses).await { continue; } // --- MQ.* --- - if write::try_handle_mq_command(cmd, cmd_args, &mut conn, ctx, &mut responses) { + if write::try_handle_mq_command(cmd, cmd_args, &frame, &mut conn, ctx, &mut responses) + .await + { continue; } @@ -976,17 +972,9 @@ pub(crate) async fn handle_connection_sharded_monoio< // --- MA2: KILL SNAPSHOT --- if cmd.eq_ignore_ascii_case(b"KILL") { - // Phase 2a: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::server_admin::kill_snapshot(&mut s.vector_store, cmd_args) - }) - } else { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - let response = crate::command::server_admin::kill_snapshot(&mut vs, cmd_args); - drop(vs); - response - }; + let response = crate::shard::slice::with_shard(|s| { + crate::command::server_admin::kill_snapshot(&mut s.vector_store, cmd_args) + }); responses.push(response); continue; } @@ -1002,79 +990,42 @@ pub(crate) async fn handle_connection_sharded_monoio< if let Some(sub_frame) = cmd_args.first() { if let Some(sub) = crate::command::helpers::extract_bytes(sub_frame) { if sub.eq_ignore_ascii_case(b"VECTOR") { - // Phase 2a: gate on is_initialized() - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::server_admin::vacuum_vector( - &mut s.vector_store, - &cmd_args[1..], - ) - }) - } else { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - let r = crate::command::server_admin::vacuum_vector( - &mut vs, + let response = crate::shard::slice::with_shard(|s| { + crate::command::server_admin::vacuum_vector( + &mut s.vector_store, &cmd_args[1..], - ); - drop(vs); - r - }; + ) + }); responses.push(response); continue; } #[cfg(feature = "graph")] if sub.eq_ignore_ascii_case(b"GRAPH") { - // Phase 2a: gate on is_initialized() - let response = if crate::shard::slice::is_initialized() { - let graph_merge_max = ctx.config.graph_merge_max_segments; - let graph_dead = ctx.config.graph_dead_edge_trigger; - crate::shard::slice::with_shard(|s| { - crate::command::server_admin::vacuum_graph( - &mut s.graph_store, - &cmd_args[1..], - graph_merge_max, - graph_dead, - ) - }) - } else { - let mut gs = ctx.shard_databases.graph_store_write(ctx.shard_id); - let r = crate::command::server_admin::vacuum_graph( - &mut gs, + let graph_merge_max = ctx.config.graph_merge_max_segments; + let graph_dead = ctx.config.graph_dead_edge_trigger; + let response = crate::shard::slice::with_shard(|s| { + crate::command::server_admin::vacuum_graph( + &mut s.graph_store, &cmd_args[1..], - ctx.config.graph_merge_max_segments, - ctx.config.graph_dead_edge_trigger, - ); - drop(gs); - r - }; + graph_merge_max, + graph_dead, + ) + }); responses.push(response); continue; } } } - // Phase 2a: gate on is_initialized() for generic VACUUM - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::server_admin::vacuum( - &mut s.vector_store, - None, - None, - cmd_args, - crate::command::server_admin::DEFAULT_VACUUM_PRUNE_MARGIN, - ) - }) - } else { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - let r = crate::command::server_admin::vacuum( - &mut vs, + // Unconditional slice path for generic VACUUM. + let response = crate::shard::slice::with_shard(|s| { + crate::command::server_admin::vacuum( + &mut s.vector_store, None, // manifest — not available in connection handler None, // wal_v3 — not available in connection handler cmd_args, crate::command::server_admin::DEFAULT_VACUUM_PRUNE_MARGIN, - ); - drop(vs); - r - }; + ) + }); responses.push(response); continue; } @@ -1084,23 +1035,13 @@ pub(crate) async fn handle_connection_sharded_monoio< if let Some(sub) = cmd_args.first() { if let Some(s) = crate::command::helpers::extract_bytes(sub) { if s.eq_ignore_ascii_case(b"RECLAMATION") { - // Phase 2a: gate on is_initialized() - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::server_admin::debug_reclamation( - &s.vector_store, - None, - None, - ) - }) - } else { - let vs = ctx.shard_databases.vector_store(ctx.shard_id); - let r = crate::command::server_admin::debug_reclamation( - &vs, None, None, - ); - drop(vs); - r - }; + let response = crate::shard::slice::with_shard(|s| { + crate::command::server_admin::debug_reclamation( + &s.vector_store, + None, + None, + ) + }); responses.push(response); continue; } @@ -1171,27 +1112,14 @@ pub(crate) async fn handle_connection_sharded_monoio< let response = match ksmv::parse_move_args(cmd_args, db_count) { Err(e) => e, Ok((_key, dst_db)) if dst_db == src_db => Frame::Integer(0), - Ok((key, dst_db)) => { - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - ksmv::with_two_slice_dbs( - &mut s.databases, - src_db, - dst_db, - |src, dst| ksmv::move_core(src, dst, &key), - ) - }) - } else { - // Lock ordering (lower index first) prevents deadlock with - // concurrent reverse MOVE from another connection. - ksmv::with_two_dbs_locked( - &ctx.shard_databases.all_shard_dbs()[ctx.shard_id], - src_db, - dst_db, - |src, dst| ksmv::move_core(src, dst, &key), - ) - } - } + Ok((key, dst_db)) => crate::shard::slice::with_shard(|s| { + ksmv::with_two_slice_dbs( + &mut s.databases, + src_db, + dst_db, + |src, dst| ksmv::move_core(src, dst, &key), + ) + }), }; // AOF only on actual success (:1). Matches handler_single. // H1 fix: durable path under `appendfsync=always` @@ -1238,41 +1166,22 @@ pub(crate) async fn handle_connection_sharded_monoio< } let response = match copy_result { Err(e) => e, - Ok(ca) => { - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - ksmv::with_two_slice_dbs( - &mut s.databases, - src_db, - ca.dst_db, - |src, dst| { - ksmv::copy_core( - src, - dst, - &ca.src_key, - &ca.dst_key, - ca.replace, - ) - }, + Ok(ca) => crate::shard::slice::with_shard(|s| { + ksmv::with_two_slice_dbs( + &mut s.databases, + src_db, + ca.dst_db, + |src, dst| { + ksmv::copy_core( + src, + dst, + &ca.src_key, + &ca.dst_key, + ca.replace, ) - }) - } else { - ksmv::with_two_dbs_locked( - &ctx.shard_databases.all_shard_dbs()[ctx.shard_id], - src_db, - ca.dst_db, - |src, dst| { - ksmv::copy_core( - src, - dst, - &ca.src_key, - &ca.dst_key, - ca.replace, - ) - }, - ) - } - } + }, + ) + }), }; // AOF only on actual success (:1). Matches handler_single // — `:0` (key absent / dst exists w/o REPLACE) is a no-op. @@ -1304,17 +1213,11 @@ pub(crate) async fn handle_connection_sharded_monoio< } if metadata::is_write(cmd) { - // WRITE PATH: eviction + dispatch under write lock. - // - // Phase 2a: Option A gating — when is_initialized() the new path - // uses ShardSlice directly (dead code until Phase 4 wires init_shard). - // The old path (else branch) is the current live path. + // WRITE PATH: eviction + dispatch via ShardSlice. // // Fast path: when neither maxmemory nor disk-offload is - // configured (default deployment + default bench), skip - // the `runtime_config.read()` acquire and the eviction - // call entirely — both are no-ops. Saves one RwLock - // lock pair per pipelined write. + // configured (default deployment + default bench), the + // eviction call is a no-op. // // Returns Ok((result, new_selected_db, hset_inserts)) on success // or Err(oom_frame) when OOM eviction fails (caller pushes + continues). @@ -1325,235 +1228,109 @@ pub(crate) async fn handle_connection_sharded_monoio< smallvec::SmallVec<[(bytes::Bytes, u64); 4]>, ), Frame, - > = 'write_path: { - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - let sel_db = conn.selected_db; - let db = &mut s.databases[sel_db]; - - // Eviction under the new path - if batch_eviction_active { - run_write_eviction_gate(ctx, db, sel_db)?; - } - - // KV undo-log capture (MUST precede dispatch) - if let Some(ref mut txn) = conn.active_cross_txn { - if cmd.eq_ignore_ascii_case(b"DEL") - || cmd.eq_ignore_ascii_case(b"UNLINK") - { - for arg in cmd_args.iter() { - if let Frame::BulkString(key_bytes) = arg { - if let Some(old_entry) = - db.get(key_bytes.as_ref()).cloned() - { - txn.kv_undo.record_delete( - key_bytes.clone(), - old_entry, - ); - let lsn = txn.snapshot_lsn; - let tid = txn.txn_id; - ctx.shard_databases - .kv_intents(ctx.shard_id) - .record_write(key_bytes.clone(), lsn, tid); - } - } - } - } else if let Some(key) = - crate::server::conn::shared::extract_primary_key( - cmd, cmd_args, - ) - { - let old_entry = db.get(key.as_ref()).cloned(); - let lsn = txn.snapshot_lsn; - let tid = txn.txn_id; - match old_entry { - None => txn.kv_undo.record_insert(key.clone()), - Some(entry) => { - txn.kv_undo.record_update(key.clone(), entry) - } - } - ctx.shard_databases.kv_intents(ctx.shard_id).record_write( - key.clone(), - lsn, - tid, - ); - } - } + > = crate::shard::slice::with_shard(|s| { + let sel_db = conn.selected_db; + let db = &mut s.databases[sel_db]; - // Dispatch - let mut new_sel_db = sel_db; - let result = dispatch(db, cmd, cmd_args, &mut new_sel_db, db_count); - let response_frame = match &result { - DispatchResult::Response(f) | DispatchResult::Quit(f) => { - f.clone() - } - }; - let is_error = matches!(response_frame, Frame::Error(_)); - - // HSET auto-index: disjoint field borrows (NLL) - // &mut s.vector_store + &mut s.text_store are separate - // from s.databases[sel_db] already borrowed above as `db`. - // Re-borrow db after dispatch since `db` was moved into dispatch. - let hset_inserts = if !is_error && cmd.eq_ignore_ascii_case(b"HSET") - { - if let Some(key) = - cmd_args.first().and_then(|f| extract_bytes(f)) - { - crate::shard::spsc_handler::auto_index_hset_public( - &mut s.vector_store, - &mut s.text_store, - key.as_ref(), - cmd_args, - ) - } else { - smallvec::SmallVec::new() - } - } else { - smallvec::SmallVec::new() - }; + if batch_eviction_active { + run_write_eviction_gate(ctx, db, sel_db)?; + } - // Blocking wakeup: re-borrow db by index (NLL) - if !is_error { - let needs_wake = cmd.eq_ignore_ascii_case(b"LPUSH") - || cmd.eq_ignore_ascii_case(b"RPUSH") - || cmd.eq_ignore_ascii_case(b"LMOVE") - || cmd.eq_ignore_ascii_case(b"ZADD"); - if needs_wake { - if let Some(key) = - cmd_args.first().and_then(|f| extract_bytes(f)) + // KV undo-log capture (MUST precede dispatch) + if let Some(ref mut txn) = conn.active_cross_txn { + if cmd.eq_ignore_ascii_case(b"DEL") + || cmd.eq_ignore_ascii_case(b"UNLINK") + { + for arg in cmd_args.iter() { + if let Frame::BulkString(key_bytes) = arg { + if let Some(old_entry) = db.get(key_bytes.as_ref()).cloned() { - let mut reg = ctx.blocking_registry.borrow_mut(); - let wake_db = &mut s.databases[new_sel_db]; - 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, wake_db, new_sel_db, &key, - ); - } else { - crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, wake_db, new_sel_db, &key, - ); - } + txn.kv_undo.record_delete(key_bytes.clone(), old_entry); + let lsn = txn.snapshot_lsn; + let tid = txn.txn_id; + // Direct field access — the outer with_shard + // closure already owns `s`; re-entering + // with_shard here panics (slice re-entrancy + // guard). `db` borrows s.databases only, so + // s.kv_write_intents is a disjoint field (NLL). + s.kv_write_intents.record_write( + key_bytes.clone(), + lsn, + tid, + ); } } } - - Ok((result, new_sel_db, hset_inserts)) - }) - } else { - let mut guard = - ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - if batch_eviction_active { - let evict_result = - run_write_eviction_gate(ctx, &mut guard, conn.selected_db); - if let Err(oom_frame) = evict_result { - drop(guard); - break 'write_path Err(oom_frame); + } else if let Some(key) = + crate::server::conn::shared::extract_primary_key(cmd, cmd_args) + { + let old_entry = db.get(key.as_ref()).cloned(); + let lsn = txn.snapshot_lsn; + let tid = txn.txn_id; + match old_entry { + None => txn.kv_undo.record_insert(key.clone()), + Some(entry) => txn.kv_undo.record_update(key.clone(), entry), } + // Direct field access — see DEL/UNLINK arm above. + s.kv_write_intents.record_write(key.clone(), lsn, tid); } + } - if let Some(ref mut txn) = conn.active_cross_txn { - if cmd.eq_ignore_ascii_case(b"DEL") - || cmd.eq_ignore_ascii_case(b"UNLINK") - { - for arg in cmd_args.iter() { - if let Frame::BulkString(key_bytes) = arg { - if let Some(old_entry) = - guard.get(key_bytes.as_ref()).cloned() - { - txn.kv_undo - .record_delete(key_bytes.clone(), old_entry); - let lsn = txn.snapshot_lsn; - let tid = txn.txn_id; - ctx.shard_databases - .kv_intents(ctx.shard_id) - .record_write(key_bytes.clone(), lsn, tid); - } - } - } - } else if let Some(key) = - crate::server::conn::shared::extract_primary_key(cmd, cmd_args) - { - let old_entry = guard.get(key.as_ref()).cloned(); - let lsn = txn.snapshot_lsn; - let tid = txn.txn_id; - match old_entry { - None => txn.kv_undo.record_insert(key.clone()), - Some(entry) => { - txn.kv_undo.record_update(key.clone(), entry) - } - } - ctx.shard_databases.kv_intents(ctx.shard_id).record_write( - key.clone(), - lsn, - tid, - ); - } + // Dispatch + let mut new_sel_db = sel_db; + let result = dispatch(db, cmd, cmd_args, &mut new_sel_db, db_count); + let response_frame = match &result { + DispatchResult::Response(f) | DispatchResult::Quit(f) => f.clone(), + }; + let is_error = matches!(response_frame, Frame::Error(_)); + + // HSET auto-index: disjoint field borrows (NLL) + // &mut s.vector_store + &mut s.text_store are separate + // from s.databases[sel_db] already borrowed above as `db`. + // Re-borrow db after dispatch since `db` was moved into dispatch. + let hset_inserts = if !is_error && cmd.eq_ignore_ascii_case(b"HSET") { + if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { + crate::shard::spsc_handler::auto_index_hset_public( + &mut s.vector_store, + &mut s.text_store, + key.as_ref(), + cmd_args, + ) + } else { + smallvec::SmallVec::new() } + } else { + smallvec::SmallVec::new() + }; - let result = dispatch( - &mut guard, - cmd, - cmd_args, - &mut conn.selected_db, - db_count, - ); - let new_sel_db = conn.selected_db; - let is_error = matches!( - result, - DispatchResult::Response(Frame::Error(_)) - | DispatchResult::Quit(Frame::Error(_)) - ); - - let hset_inserts = if !is_error && cmd.eq_ignore_ascii_case(b"HSET") { + // Blocking wakeup: re-borrow db by index (NLL) + if !is_error { + let needs_wake = cmd.eq_ignore_ascii_case(b"LPUSH") + || cmd.eq_ignore_ascii_case(b"RPUSH") + || cmd.eq_ignore_ascii_case(b"LMOVE") + || cmd.eq_ignore_ascii_case(b"ZADD"); + if needs_wake { if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - let mut ts = ctx.shard_databases.text_store(ctx.shard_id); - crate::shard::spsc_handler::auto_index_hset_public( - &mut vs, - &mut *ts, - key.as_ref(), - cmd_args, - ) - } else { - smallvec::SmallVec::new() - } - } else { - smallvec::SmallVec::new() - }; - - if !is_error { - let needs_wake = cmd.eq_ignore_ascii_case(b"LPUSH") - || cmd.eq_ignore_ascii_case(b"RPUSH") - || cmd.eq_ignore_ascii_case(b"LMOVE") - || cmd.eq_ignore_ascii_case(b"ZADD"); - if needs_wake { - if let Some(key) = - cmd_args.first().and_then(|f| extract_bytes(f)) + let mut reg = ctx.blocking_registry.borrow_mut(); + let wake_db = &mut s.databases[new_sel_db]; + if cmd.eq_ignore_ascii_case(b"LPUSH") + || cmd.eq_ignore_ascii_case(b"RPUSH") + || cmd.eq_ignore_ascii_case(b"LMOVE") { - let mut reg = ctx.blocking_registry.borrow_mut(); - if cmd.eq_ignore_ascii_case(b"LPUSH") - || cmd.eq_ignore_ascii_case(b"RPUSH") - || cmd.eq_ignore_ascii_case(b"LMOVE") - { - crate::blocking::wakeup::try_wake_list_waiter( - &mut reg, &mut guard, new_sel_db, &key, - ); - } else { - crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, &mut guard, new_sel_db, &key, - ); - } + crate::blocking::wakeup::try_wake_list_waiter( + &mut reg, wake_db, new_sel_db, &key, + ); + } else { + crate::blocking::wakeup::try_wake_zset_waiter( + &mut reg, wake_db, new_sel_db, &key, + ); } } } + } - drop(guard); - Ok((result, new_sel_db, hset_inserts)) - } // end else branch - }; // end 'write_path labeled block + Ok((result, new_sel_db, hset_inserts)) + }); // Unpack write result — OOM causes immediate continue let (result, new_selected_db, hset_inserts) = match write_result { @@ -1679,23 +1456,17 @@ pub(crate) async fn handle_connection_sharded_monoio< let my_txn_id = txn.txn_id; // Clone committed treemap to release vector_store lock // before acquiring kv_intents lock (lock ordering). - let committed = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - s.vector_store.txn_manager().committed_treemap().clone() - }) - } else { - let vs = ctx.shard_databases.vector_store(ctx.shard_id); - vs.txn_manager().committed_treemap().clone() - }; - let visible = { - let intents = ctx.shard_databases.kv_intents(ctx.shard_id); - intents.is_key_visible( + let committed = crate::shard::slice::with_shard(|s| { + s.vector_store.txn_manager().committed_treemap().clone() + }); + let visible = crate::shard::slice::with_shard(|s| { + s.kv_write_intents.is_key_visible( key.as_ref(), snapshot_lsn, my_txn_id, &committed, ) - }; + }); if !visible { responses.push(Frame::Null); continue; @@ -1709,26 +1480,11 @@ pub(crate) async fn handle_connection_sharded_monoio< conn.cmd_counter = conn.cmd_counter.wrapping_add(1); let sample_latency = (conn.cmd_counter & 0xF) == 0; let dispatch_start = sample_latency.then(std::time::Instant::now); - let (result, new_read_selected_db) = if crate::shard::slice::is_initialized() { - let mut sel_db = conn.selected_db; - let result = crate::shard::slice::with_shard_db(conn.selected_db, |db| { - dispatch_read(db, cmd, cmd_args, now_ms, &mut sel_db, db_count) - }); - (result, sel_db) - } else { - let guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); - let result = dispatch_read( - &guard, - cmd, - cmd_args, - now_ms, - &mut conn.selected_db, - db_count, - ); - let sel = conn.selected_db; - drop(guard); - (result, sel) - }; + let mut sel_db = conn.selected_db; + let result = crate::shard::slice::with_shard_db(conn.selected_db, |db| { + dispatch_read(db, cmd, cmd_args, now_ms, &mut sel_db, db_count) + }); + let new_read_selected_db = sel_db; conn.selected_db = new_read_selected_db; if let Ok(cmd_str) = std::str::from_utf8(cmd) { if let Some(start) = dispatch_start { @@ -1796,60 +1552,12 @@ pub(crate) async fn handle_connection_sharded_monoio< // all foreign-shard reads through SPSC (eliminates RwLock contention at // the cost of one extra channel round-trip per read command). // See docs/production-guide.md §Cross-shard fast path. - if !metadata::is_write(cmd) - && !remote_groups.contains_key(&target) - && is_dispatch_read_supported(cmd) - && ctx.config.cross_shard_fast_path_enabled() - { - crate::admin::metrics_setup::record_dispatch_cross_read_fastpath(); - // Time the RwLock acquisition to surface contention under load. - // Overhead when metrics are disabled: one Relaxed atomic load (≈ 1 ns). - let t0 = if crate::admin::metrics_setup::is_metrics_enabled() { - Some(std::time::Instant::now()) - } else { - None - }; - let guard = ctx.shard_databases.read_db(target, conn.selected_db); - if let Some(t0) = t0 { - let ns = t0.elapsed().as_nanos() as u64; - crate::admin::metrics_setup::record_dispatch_cross_read_fastpath_timed( - target, ns, - ); - } - let now_ms = ctx.cached_clock.ms(); - let result = dispatch_read( - &guard, - cmd, - cmd_args, - now_ms, - &mut conn.selected_db, - db_count, - ); - drop(guard); - let response = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => { - should_quit = true; - f - } - }; - // Client tracking for cross-shard reads - if conn.tracking_state.enabled && !conn.tracking_state.bcast { - if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { - ctx.tracking_table.borrow_mut().track_key( - client_id, - &key, - conn.tracking_state.noloop, - ); - } - } - let mut response = apply_resp3_conversion(cmd, response, conn.protocol_version); - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } + // E2: Cross-shard fast path disabled — ShardSlice is thread-local; + // foreign-shard data can only be read via SPSC hop. All cross-shard + // reads now route through the SPSC channel below regardless of + // is_dispatch_read_supported. The fast path can be re-enabled in a + // later wave when the cross-shard snapshot protocol is implemented. + // See shardslice-migration TASK.md § C6. // Cross-shard write: deferred SPSC dispatch. // When workspace rewriting occurred, rebuild the frame with // prefixed args so the target shard stores the correct key. @@ -2089,45 +1797,41 @@ pub(crate) async fn handle_connection_sharded_monoio< continue; } }; + // H1-BARRIER: collect write resp_idxs before consuming meta + // so we can overwrite them if the fsync barrier fails. + let mut write_resp_idxs: Vec = Vec::new(); for ((resp_idx, aof_bytes, cmd_name), resp) in meta.into_iter().zip(shard_responses) { - // AOF logging for successful remote writes. - // Owner shard is `target` (NOT ctx.shard_id) — under PerShard - // layout the write must land in the target shard's AOF file - // since that shard owns the mutated data. Mirrors the - // load-bearing fix at handler_sharded/mod.rs:1651. - if let Some(bytes) = aof_bytes { - if !matches!(resp, Frame::Error(_)) { - if let Some(ref pool) = ctx.aof_pool { - // Cross-shard write: LSN must be sourced - // using `target`'s shard_id so the - // per-shard offset increment lands on the - // shard that owns the mutated data. - let lsn = aof::AofWriterPool::issue_append_lsn( - &ctx.repl_state, - target, - bytes.len(), - ); - // H1: durable path under appendfsync=always. - if pool - .try_send_append_durable(target, lsn, bytes) - .await - .is_err() - { - let err = Frame::Error(Bytes::from_static(aof::AOF_FSYNC_ERR)); - let err = apply_resp3_conversion( - &cmd_name, - err, - conn.protocol_version, - ); - responses[resp_idx] = err; - continue; - } + // C4-FOLD-FIX: AOF append for cross-shard writes is now done + // inside the SPSC arm (PipelineBatch), BEFORE the oneshot reply + // is sent. Appending here (after awaiting the oneshot response) + // defers the append until after drain_spsc_shared returns, which + // makes AofFold's pending_aof_count undercount it → escape to + // new incr → double-apply on restart. The SPSC arm now owns the + // AOF write; aof_bytes below is used only for the barrier check. + let resp = apply_resp3_conversion(&cmd_name, resp, conn.protocol_version); + if aof_bytes.is_some() && !matches!(resp, Frame::Error(_)) { + write_resp_idxs.push(resp_idx); + } + responses[resp_idx] = resp; + } + + // H1-BARRIER (C4-FOLD-FIX follow-up): under appendfsync=always, + // call fsync_barrier once per target shard AFTER responses are + // collected. The SPSC arm enqueued the Append fire-and-forget; + // the barrier enqueues a zero-length AppendSync into the SAME + // shard channel. Because the writer processes messages in order, + // an acked barrier proves all prior Appends to this shard are on + // durable storage. Under EverySec/No this is a zero-cost noop. + if !write_resp_idxs.is_empty() { + if let Some(ref pool) = ctx.aof_pool { + if pool.fsync_barrier(target).await.is_err() { + for idx in write_resp_idxs { + responses[idx] = + Frame::Error(Bytes::from_static(aof::AOF_FSYNC_ERR)); } } } - let resp = apply_resp3_conversion(&cmd_name, resp, conn.protocol_version); - responses[resp_idx] = resp; } } } diff --git a/src/server/conn/handler_monoio/txn.rs b/src/server/conn/handler_monoio/txn.rs index e91a6ef9..61777df1 100644 --- a/src/server/conn/handler_monoio/txn.rs +++ b/src/server/conn/handler_monoio/txn.rs @@ -22,7 +22,7 @@ pub(super) fn try_handle_txn_begin( cmd: &[u8], cmd_args: &[Frame], conn: &mut ConnectionState, - ctx: &ConnectionContext, + _ctx: &ConnectionContext, responses: &mut Vec, ) -> bool { if !is_txn_begin(cmd, cmd_args) { @@ -31,8 +31,8 @@ pub(super) fn try_handle_txn_begin( match txn_begin_validate(conn.in_multi, conn.in_cross_txn()) { Ok(()) => { // Get next txn_id and snapshot_lsn from vector store's transaction manager - let mut vector_store = ctx.shard_databases.vector_store(ctx.shard_id); - let active = vector_store.txn_manager_mut().begin(); + let active = + crate::shard::slice::with_shard(|s| s.vector_store.txn_manager_mut().begin()); conn.active_cross_txn = Some(CrossStoreTxn::new(active.txn_id, active.snapshot_lsn)); responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); } @@ -42,7 +42,7 @@ pub(super) fn try_handle_txn_begin( } /// Handle TXN.COMMIT — returns `true` if the command was consumed. -pub(super) fn try_handle_txn_commit( +pub(super) async fn try_handle_txn_commit( cmd: &[u8], cmd_args: &[Frame], conn: &mut ConnectionState, @@ -60,10 +60,18 @@ pub(super) fn try_handle_txn_commit( // have been excluded from oldest_snapshot, allowing prune_committed to // advance past its LSN. Committing with a stale read set is undefined // behaviour — force the client to restart the transaction. - let mut vector_store = ctx.shard_databases.vector_store(ctx.shard_id); - if vector_store.txn_manager().is_killed(txn.txn_id) { - vector_store.txn_manager_mut().abort_killed(txn.txn_id); - drop(vector_store); + // Scope the vector_store guard so it is DROPPED before any .await + // in the MQ-materialize hop below (lock-across-await rule). + let killed = crate::shard::slice::with_shard(|s| { + if s.vector_store.txn_manager().is_killed(txn.txn_id) { + s.vector_store.txn_manager_mut().abort_killed(txn.txn_id); + true + } else { + s.vector_store.txn_manager_mut().commit(txn.txn_id); + false + } + }); // with_shard borrow released here + if killed { tracing::warn!( txn_id = txn.txn_id, "TXN.COMMIT rejected: snapshot was killed (snapshot too old)" @@ -74,29 +82,18 @@ pub(super) fn try_handle_txn_commit( responses.push(Frame::Error(msg.freeze())); return true; } - vector_store.txn_manager_mut().commit(txn.txn_id); - drop(vector_store); // Write XactCommit WAL record with committed KV state let txn_id = txn.txn_id; if !txn.kv_undo.is_empty() { - let payload = if crate::shard::slice::is_initialized() { + let payload = crate::shard::slice::with_shard_db(conn.selected_db as usize, |db| { crate::persistence::wal_v3::record::encode_xact_commit_payload( txn_id, txn.kv_undo.records(), db, ) - }) - } else { - let db_guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); - let p = crate::persistence::wal_v3::record::encode_xact_commit_payload( - txn_id, - txn.kv_undo.records(), - &*db_guard, - ); - p - }; + }); let mut wal_buf = Vec::new(); crate::persistence::wal_v3::record::write_wal_v3_record( &mut wal_buf, @@ -108,54 +105,58 @@ pub(super) fn try_handle_txn_commit( .wal_append(ctx.shard_id, bytes::Bytes::from(wal_buf)); } - // Release KV write intents from shard side-table - ctx.shard_databases - .kv_intents(ctx.shard_id) - .release_txn(txn_id); - - // Drain deferred HNSW inserts (post-commit hook). - // The drain prevents phantom neighbors on abort. - // Actual HNSW graph insertion happens during compaction, - // not at commit time (point is already in mutable segment). - let drain_count = ctx - .shard_databases - .hnsw_queue(ctx.shard_id) - .drain_for_txn(txn_id) - .count(); + // Release KV write intents and drain HNSW queue via thread-local slice. + let drain_count = crate::shard::slice::with_shard(|s| { + s.kv_write_intents.release_txn(txn_id); + s.deferred_hnsw_inserts.drain_for_txn(txn_id).count() + }); if drain_count > 0 { tracing::debug!(txn_id, count = drain_count, "Drained deferred HNSW inserts"); } - // Materialize MQ intents: enqueue deferred MQ.PUBLISH messages + // Materialize MQ intents: enqueue deferred MQ.PUBLISH messages. + // Owner-route via MqTxnMaterialize: self-owned intents fold locally; + // foreign intents hop to the owning shard and await ack before reply. if !txn.mq_intents.is_empty() { - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db as usize, |db| { - for intent in &txn.mq_intents { - if let Ok(Some(stream)) = db.get_stream_mut(&intent.queue_key) { - if stream.durable { - let msg_id = stream.next_auto_id(); - stream.add(msg_id, intent.fields.clone()); + use std::collections::HashMap; + let mut by_shard: HashMap> = + HashMap::new(); + for intent in txn.mq_intents.iter().cloned() { + let owner = + crate::shard::dispatch::key_to_shard(&intent.queue_key, ctx.num_shards); + by_shard.entry(owner).or_default().push(intent); + } + for (owner, intents) in by_shard { + if owner == ctx.shard_id { + // Self: apply locally via slice. + crate::shard::slice::with_shard_db(conn.selected_db as usize, |db| { + for intent in &intents { + if let Ok(Some(stream)) = db.get_stream_mut(&intent.queue_key) { + if stream.durable { + let msg_id = stream.next_auto_id(); + stream.add(msg_id, intent.fields.clone()); + } } } - } - }); - } else { - // Each queue lives on the shard owning its key, which - // may differ from the connection's shard (and per - // intent) — acquire the owner's db per queue. - for intent in &txn.mq_intents { - let owner = crate::shard::dispatch::key_to_shard( - &intent.queue_key, - ctx.num_shards, - ); - let mut db_guard = - ctx.shard_databases.write_db(owner, conn.selected_db); - if let Ok(Some(stream)) = db_guard.get_stream_mut(&intent.queue_key) { - if stream.durable { - let msg_id = stream.next_auto_id(); - stream.add(msg_id, intent.fields.clone()); - } - } + }); + } else { + // Foreign: send MqTxnMaterialize hop and await ack. + let (reply_tx, reply_rx) = crate::runtime::channel::oneshot(); + let msg = crate::shard::dispatch::ShardMessage::MqTxnMaterialize { + db_index: conn.selected_db as usize, + intents, + reply_tx, + }; + crate::shard::coordinator::spsc_send( + &ctx.dispatch_tx, + ctx.shard_id, + owner, + msg, + &ctx.spsc_notifiers, + ) + .await; + // Await the ack before replying OK to the client. + let _ = reply_rx.recv().await; } } } @@ -214,7 +215,7 @@ pub(super) async fn try_handle_txn_abort( pub(super) fn try_handle_temporal_snapshot_at( cmd: &[u8], cmd_args: &[Frame], - ctx: &ConnectionContext, + _ctx: &ConnectionContext, responses: &mut Vec, ) -> bool { if !is_temporal_snapshot_at(cmd) { @@ -223,16 +224,13 @@ pub(super) fn try_handle_temporal_snapshot_at( match validate_snapshot_at(cmd_args) { Ok(()) => { let wall_ms = capture_wall_ms(); - let lsn = { - let vector_store = ctx.shard_databases.vector_store(ctx.shard_id); - vector_store.txn_manager().current_lsn() - }; - { - let mut guard = ctx.shard_databases.temporal_registry(ctx.shard_id); - let registry = - guard.get_or_insert_with(|| Box::new(crate::temporal::TemporalRegistry::new())); + crate::shard::slice::with_shard(|s| { + let lsn = s.vector_store.txn_manager().current_lsn(); + let registry = s + .temporal_registry + .get_or_insert_with(|| Box::new(crate::temporal::TemporalRegistry::new())); registry.record(wall_ms, lsn); - } + }); responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); } Err(e) => responses.push(e), @@ -286,30 +284,12 @@ pub(super) async fn try_handle_temporal_invalidate( } } let wall_ms = capture_wall_ms(); - // Phase 2a: gate on is_initialized(); both branches are - // semantically identical, the new path uses + // Unconditional slice path: apply temporal invalidation via // ShardSlice::graph_store directly (no lock). - let (result, wal_records) = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - let gs = &mut s.graph_store; - let r = crate::command::temporal::apply_invalidate( - gs, - entity_id, - is_node, - &graph_name, - wall_ms, - ); - let recs = if r.is_ok() { - gs.drain_wal() - } else { - Vec::new() - }; - (r, recs) - }) - } else { - let mut gs = ctx.shard_databases.graph_store_write(ctx.shard_id); + let (result, wal_records) = crate::shard::slice::with_shard(|s| { + let gs = &mut s.graph_store; let r = crate::command::temporal::apply_invalidate( - &mut gs, + gs, entity_id, is_node, &graph_name, @@ -321,7 +301,7 @@ pub(super) async fn try_handle_temporal_invalidate( Vec::new() }; (r, recs) - }; + }); match result { Ok(()) => { for record in wal_records { diff --git a/src/server/conn/handler_monoio/write.rs b/src/server/conn/handler_monoio/write.rs index 1e02462e..7260b01e 100644 --- a/src/server/conn/handler_monoio/write.rs +++ b/src/server/conn/handler_monoio/write.rs @@ -5,8 +5,8 @@ use bytes::Bytes; use crate::command::mq::{ - ERR_MQ_NOT_DURABLE, ERR_MQ_UNKNOWN_SUB, parse_mq_subcommand, validate_mq_ack, - validate_mq_create, validate_mq_dlqlen, validate_mq_pop, validate_mq_publish, validate_mq_push, + ERR_MQ_UNKNOWN_SUB, parse_mq_subcommand, validate_mq_ack, validate_mq_create, + validate_mq_dlqlen, validate_mq_pop, validate_mq_publish, validate_mq_push, validate_mq_trigger, }; use crate::command::transaction::ERR_MULTI_TXN_CONFLICT; @@ -28,7 +28,7 @@ use crate::workspace::{WorkspaceId, is_ws_command}; use super::execute_transaction_sharded; /// Handle WS.* workspace commands. Returns `true` if consumed. -pub(super) fn try_handle_ws_command( +pub(super) async fn try_handle_ws_command( cmd: &[u8], cmd_args: &[Frame], conn: &mut ConnectionState, @@ -112,14 +112,23 @@ pub(super) fn try_handle_ws_command( ); ctx.shard_databases.wal_append(0, Bytes::from(wal_buf)); // Best-effort cleanup: delete all KV keys with ws prefix (WS-03). - // Phase 2a: gate on is_initialized(); new path uses ShardSlice. + // Owner-route the cleanup via WsDropCleanup hop. The {wsid} hash + // tag co-locates every workspace key on ONE shard. { let prefix = format!("{{{}}}:", ws_id.as_hex()); - if crate::shard::slice::is_initialized() { + let prefix_bytes = Bytes::from(prefix.into_bytes()); + let owner = crate::shard::dispatch::key_to_shard( + &prefix_bytes, + ctx.num_shards, + ); + if owner == ctx.shard_id { + // Self: execute locally (we ARE the owner). crate::shard::slice::with_shard_db(0, |db| { let keys_to_delete: Vec> = db .keys() - .filter(|k| k.as_bytes().starts_with(prefix.as_bytes())) + .filter(|k| { + k.as_bytes().starts_with(prefix_bytes.as_ref()) + }) .map(|k| k.as_bytes().to_vec()) .collect(); for key in &keys_to_delete { @@ -127,22 +136,21 @@ pub(super) fn try_handle_ws_command( } }); } else { - // The {wsid} hash tag co-locates every - // workspace key on ONE shard — clean up - // there, not on the connection's shard. - let owner = crate::shard::dispatch::key_to_shard( - prefix.as_bytes(), - ctx.num_shards, - ); - let mut db_guard = ctx.shard_databases.write_db(owner, 0); - let keys_to_delete: Vec> = db_guard - .keys() - .filter(|k| k.as_bytes().starts_with(prefix.as_bytes())) - .map(|k| k.as_bytes().to_vec()) - .collect(); - for key in &keys_to_delete { - db_guard.remove(key); - } + // Foreign: send WsDropCleanup hop to owner. + let (reply_tx, reply_rx) = crate::runtime::channel::oneshot(); + let msg = crate::shard::dispatch::ShardMessage::WsDropCleanup { + prefix: prefix_bytes, + reply_tx, + }; + crate::shard::coordinator::spsc_send( + &ctx.dispatch_tx, + ctx.shard_id, + owner, + msg, + &ctx.spsc_notifiers, + ) + .await; + let _ = reply_rx.recv().await; } } responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); @@ -257,9 +265,10 @@ pub(super) fn try_handle_ws_command( } /// Handle MQ.* message queue commands. Returns `true` if consumed. -pub(super) fn try_handle_mq_command( +pub(super) async fn try_handle_mq_command( cmd: &[u8], cmd_args: &[Frame], + frame: &Frame, conn: &mut ConnectionState, ctx: &ConnectionContext, responses: &mut Vec, @@ -277,7 +286,7 @@ pub(super) fn try_handle_mq_command( if sub.eq_ignore_ascii_case(b"CREATE") { match validate_mq_create(cmd_args) { - Ok((queue_key, max_delivery_count, _debounce_ms)) => { + Ok((queue_key, _max_delivery_count, _debounce_ms)) => { let effective_key = crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); // A durable queue lives on the shard that owns its key — the @@ -289,60 +298,20 @@ pub(super) fn try_handle_mq_command( // slice mode is never initialized yet (owner-routing there is // the shardslice-migration task). let owner = crate::shard::dispatch::key_to_shard(&effective_key, ctx.num_shards); - // Phase 2a: gate on is_initialized(); new path uses ShardSlice directly. - let create_result: Result<(), Frame> = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| { - match db.get_or_create_stream(&effective_key) { - Ok(stream) => { - stream.durable = true; - stream.max_delivery_count = max_delivery_count; - let group_name = Bytes::from_static(b"__mq_consumers"); - let _ = stream.create_group(group_name, StreamId::ZERO); - Ok(()) - } - Err(e) => Err(e), - } - }) - } else { - let mut db_guard = ctx.shard_databases.write_db(owner, conn.selected_db); - match db_guard.get_or_create_stream(&effective_key) { - Ok(stream) => { - stream.durable = true; - stream.max_delivery_count = max_delivery_count; - let group_name = Bytes::from_static(b"__mq_consumers"); - let _ = stream.create_group(group_name, StreamId::ZERO); - Ok(()) - } - Err(e) => Err(e), - } - }; - if let Err(e) = create_result { - responses.push(e); - return true; - } - - // Store config in the owning shard's registry - let config = - crate::mq::DurableStreamConfig::new(effective_key.clone(), max_delivery_count); + // Owner-routes via MqCommand hop (execute_mq_on_owner + // handles stream create, registry insert, and WAL in one step). { - let mut guard = ctx.shard_databases.durable_queue_registry(owner); - let reg = guard - .get_or_insert_with(|| Box::new(crate::mq::DurableQueueRegistry::new())); - reg.insert(effective_key.clone(), config); + let key_prefix = mq_ws_prefix(conn.workspace_id.as_ref()); + let response = mq_hop_or_local( + owner, + conn.selected_db, + ctx, + key_prefix, + std::sync::Arc::new(frame.clone()), + ) + .await; + responses.push(response); } - - // WAL: MqCreate record — owner's WAL so replay restores the - // registry on the shard that also holds the stream. - let payload = crate::mq::wal::encode_mq_create(&effective_key, max_delivery_count); - let mut wal_buf = Vec::new(); - crate::persistence::wal_v3::record::write_wal_v3_record( - &mut wal_buf, - 0, - crate::persistence::wal_v3::record::WalRecordType::MqCreate, - &payload, - ); - ctx.shard_databases.wal_append(owner, Bytes::from(wal_buf)); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); } Err(e) => responses.push(e), } @@ -351,79 +320,24 @@ pub(super) fn try_handle_mq_command( if sub.eq_ignore_ascii_case(b"PUSH") { match validate_mq_push(cmd_args) { - Ok((queue_key, fields)) => { + Ok((queue_key, _fields)) => { let effective_key = crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); // Owner-shard targeting — see MQ CREATE above. let owner = crate::shard::dispatch::key_to_shard(&effective_key, ctx.num_shards); - // Phase 2a: gate on is_initialized(); new path uses ShardSlice. - // trigger_registry stays via old path (Phase 2b scope, not migrated here). - let push_result: Result = - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| { - match db.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - if !stream.durable { - Err(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))) - } else { - let id = stream.next_auto_id(); - Ok(stream.add(id, fields)) - } - } - Ok(None) => { - Err(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))) - } - Err(e) => Err(e), - } - }) - } else { - let mut db_guard = ctx.shard_databases.write_db(owner, conn.selected_db); - match db_guard.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - if !stream.durable { - Err(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))) - } else { - let id = stream.next_auto_id(); - Ok(stream.add(id, fields)) - } - } - Ok(None) => Err(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))), - Err(e) => Err(e), - } - }; - match push_result { - Ok(msg_id) => { - // trigger_registry: not in Phase 2a scope — old path only. - // Owner's registry: its event-loop tick fires triggers. - { - let mut trig_guard = ctx.shard_databases.trigger_registry(owner); - if let Some(reg) = trig_guard.as_mut() { - let trig_key = if let Some(ws_id) = conn.workspace_id.as_ref() { - let ws_hex = ws_id.as_hex(); - let mut k = - Vec::with_capacity(ws_hex.len() + 1 + queue_key.len()); - k.extend_from_slice(ws_hex.as_bytes()); - k.push(b':'); - k.extend_from_slice(&queue_key); - Bytes::from(k) - } else { - queue_key.clone() - }; - if let Some(trig_entry) = reg.get_mut(&trig_key) { - if trig_entry.pending_fire_ms == 0 { - let fire_at = - ctx.cached_clock.ms() + trig_entry.debounce_ms; - trig_entry.pending_fire_ms = fire_at; - } - } - } - } - responses.push(Frame::BulkString(Bytes::from(format!( - "{}-{}", - msg_id.ms, msg_id.seq - )))); - } - Err(e) => responses.push(e), + // Owner-routes via MqCommand hop (execute_mq_on_owner + // handles stream push, trigger debounce in one step). + { + let key_prefix = mq_ws_prefix(conn.workspace_id.as_ref()); + let response = mq_hop_or_local( + owner, + conn.selected_db, + ctx, + key_prefix, + std::sync::Arc::new(frame.clone()), + ) + .await; + responses.push(response); } } Err(e) => responses.push(e), @@ -433,208 +347,27 @@ pub(super) fn try_handle_mq_command( if sub.eq_ignore_ascii_case(b"POP") { match validate_mq_pop(cmd_args) { - Ok((queue_key, count)) => { + Ok((queue_key, _count)) => { let effective_key = crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); // Owner-shard targeting — see MQ CREATE above. let owner = crate::shard::dispatch::key_to_shard(&effective_key, ctx.num_shards); - let group_name = Bytes::from_static(b"__mq_consumers"); - let consumer_name = Bytes::from_static(b"__mq_default"); + let _group_name = Bytes::from_static(b"__mq_consumers"); + let _consumer_name = Bytes::from_static(b"__mq_default"); - // Phase 2a: gate on is_initialized(); new path uses ShardSlice. - // All DB work (claim, DLQ routing) happens inside the closure; - // only owned data (result_frames) escapes. - let pop_frame: Result = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| { - let mdc = match db.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - if !stream.durable { - return Err(Frame::Error(Bytes::from_static( - ERR_MQ_NOT_DURABLE, - ))); - } - stream.max_delivery_count - } - Ok(None) => { - return Err(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); - } - Err(e) => return Err(e), - }; - let request_count = count + (mdc as usize); - let claimed = match db.get_stream_mut(&effective_key) { - Ok(Some(s)) => match s.read_group_new( - &group_name, - &consumer_name, - Some(request_count), - false, - ) { - Ok(entries) => entries, - Err(_) => return Ok(Frame::Array(vec![].into())), - }, - _ => { - return Err(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); - } - }; - let mut results = Vec::new(); - let mut dlq_entries: Vec<(StreamId, Vec<(Bytes, Bytes)>)> = Vec::new(); - let mut dlq_ack_ids: Vec = Vec::new(); - // Need stream again for PEL lookup - if let Ok(Some(stream)) = db.get_stream_mut(&effective_key) { - for (id, fields) in &claimed { - let delivery_count = stream - .groups - .get(group_name.as_ref()) - .and_then(|g| g.pel.get(id)) - .map(|pe| pe.delivery_count) - .unwrap_or(1); - if mdc > 0 && delivery_count >= mdc as u64 { - dlq_entries.push((*id, fields.clone())); - dlq_ack_ids.push(*id); - } else if results.len() < count { - results.push((*id, fields.clone())); - } - } - if !dlq_ack_ids.is_empty() { - let _ = stream.xack(&group_name, &dlq_ack_ids); - } - } - if !dlq_entries.is_empty() { - let dlq_key = { - let mut buf = Vec::with_capacity(effective_key.len() + 8); - buf.extend_from_slice(&effective_key); - buf.extend_from_slice(b"::mq:dlq"); - Bytes::from(buf) - }; - if let Ok(dlq_stream) = db.get_or_create_stream(&dlq_key) { - for (_id, fields) in dlq_entries { - let dlq_id = dlq_stream.next_auto_id(); - dlq_stream.add(dlq_id, fields); - } - } - } - let result_frames: Vec = results - .iter() - .map(|(id, fields)| { - let mut entry_frames = Vec::with_capacity(2); - entry_frames.push(Frame::BulkString(Bytes::from(format!( - "{}-{}", - id.ms, id.seq - )))); - let field_frames: Vec = fields - .iter() - .flat_map(|(f, v)| { - vec![ - Frame::BulkString(f.clone()), - Frame::BulkString(v.clone()), - ] - }) - .collect(); - entry_frames.push(Frame::Array(field_frames.into())); - Frame::Array(entry_frames.into()) - }) - .collect(); - Ok(Frame::Array(result_frames.into())) - }) - } else { - let mut db_guard = ctx.shard_databases.write_db(owner, conn.selected_db); - let mdc = match db_guard.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - if !stream.durable { - responses - .push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); - return true; - } - stream.max_delivery_count - } - Ok(None) => { - responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); - return true; - } - Err(e) => { - responses.push(e); - return true; - } - }; - let request_count = count + (mdc as usize); - let stream = match db_guard.get_stream_mut(&effective_key) { - Ok(Some(s)) => s, - _ => { - responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); - return true; - } - }; - let claimed = match stream.read_group_new( - &group_name, - &consumer_name, - Some(request_count), - false, - ) { - Ok(entries) => entries, - Err(_) => { - responses.push(Frame::Array(vec![].into())); - return true; - } - }; - let mut results = Vec::new(); - let mut dlq_entries: Vec<(StreamId, Vec<(Bytes, Bytes)>)> = Vec::new(); - let mut dlq_ack_ids: Vec = Vec::new(); - for (id, fields) in &claimed { - let delivery_count = stream - .groups - .get(group_name.as_ref()) - .and_then(|g| g.pel.get(id)) - .map(|pe| pe.delivery_count) - .unwrap_or(1); - if mdc > 0 && delivery_count >= mdc as u64 { - dlq_entries.push((*id, fields.clone())); - dlq_ack_ids.push(*id); - } else if results.len() < count { - results.push((*id, fields.clone())); - } - } - if !dlq_ack_ids.is_empty() { - let _ = stream.xack(&group_name, &dlq_ack_ids); - } - if !dlq_entries.is_empty() { - let dlq_key = { - let mut buf = Vec::with_capacity(effective_key.len() + 8); - buf.extend_from_slice(&effective_key); - buf.extend_from_slice(b"::mq:dlq"); - Bytes::from(buf) - }; - if let Ok(dlq_stream) = db_guard.get_or_create_stream(&dlq_key) { - for (_id, fields) in dlq_entries { - let dlq_id = dlq_stream.next_auto_id(); - dlq_stream.add(dlq_id, fields); - } - } - } - let result_frames: Vec = results - .iter() - .map(|(id, fields)| { - let mut entry_frames = Vec::with_capacity(2); - entry_frames.push(Frame::BulkString(Bytes::from(format!( - "{}-{}", - id.ms, id.seq - )))); - let field_frames: Vec = fields - .iter() - .flat_map(|(f, v)| { - vec![Frame::BulkString(f.clone()), Frame::BulkString(v.clone())] - }) - .collect(); - entry_frames.push(Frame::Array(field_frames.into())); - Frame::Array(entry_frames.into()) - }) - .collect(); - Ok(Frame::Array(result_frames.into())) - }; - match pop_frame { - Ok(frame) => responses.push(frame), - Err(e) => { - responses.push(e); - return true; - } + // Owner-routes via MqCommand hop (execute_mq_on_owner + // handles claim, DLQ routing in one step). + { + let key_prefix = mq_ws_prefix(conn.workspace_id.as_ref()); + let response = mq_hop_or_local( + owner, + conn.selected_db, + ctx, + key_prefix, + std::sync::Arc::new(frame.clone()), + ) + .await; + responses.push(response); } } Err(e) => responses.push(e), @@ -647,56 +380,26 @@ pub(super) fn try_handle_mq_command( Ok((queue_key, msg_ids)) => { let effective_key = crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); - let ids: Vec = msg_ids + let _ids: Vec = msg_ids .iter() .map(|(ms, seq)| StreamId { ms: *ms, seq: *seq }) .collect(); // Owner-shard targeting — see MQ CREATE above. let owner = crate::shard::dispatch::key_to_shard(&effective_key, ctx.num_shards); - // Phase 2a: gate on is_initialized(); new path uses ShardSlice. - // WAL append stays outside closure (wal_append_tx not in Phase 2a scope). - let ack_result: Result = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| { - match db.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - let group_name = Bytes::from_static(b"__mq_consumers"); - Ok(stream - .xack(&group_name, &ids) - .map(|c| c as i64) - .unwrap_or(0)) - } - _ => Ok(0i64), - } - }) - } else { - let mut db_guard = ctx.shard_databases.write_db(owner, conn.selected_db); - match db_guard.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - let group_name = Bytes::from_static(b"__mq_consumers"); - Ok(stream - .xack(&group_name, &ids) - .map(|c| c as i64) - .unwrap_or(0)) - } - _ => Ok(0i64), - } - }; - let acked_count = ack_result.unwrap_or(0); - if acked_count > 0 { - // Emit MqAck WAL record for each acked ID (WAL stays outside closure) - for (ms, seq) in &msg_ids { - let payload = crate::mq::wal::encode_mq_ack(&effective_key, *ms, *seq); - let mut wal_buf = Vec::new(); - crate::persistence::wal_v3::record::write_wal_v3_record( - &mut wal_buf, - 0, - crate::persistence::wal_v3::record::WalRecordType::MqAck, - &payload, - ); - ctx.shard_databases.wal_append(owner, Bytes::from(wal_buf)); - } + // Owner-routes via MqCommand hop (execute_mq_on_owner + // handles xack and WAL records in one step). + { + let key_prefix = mq_ws_prefix(conn.workspace_id.as_ref()); + let response = mq_hop_or_local( + owner, + conn.selected_db, + ctx, + key_prefix, + std::sync::Arc::new(frame.clone()), + ) + .await; + responses.push(response); } - responses.push(Frame::Integer(acked_count)); } Err(e) => responses.push(e), } @@ -711,28 +414,25 @@ pub(super) fn try_handle_mq_command( // Owner of the QUEUE key, not the dlq_key: POP creates the // DLQ stream in the same db as the queue it drains. let owner = crate::shard::dispatch::key_to_shard(&effective_key, ctx.num_shards); - let dlq_key = { + let _dlq_key = { let mut buf = Vec::with_capacity(effective_key.len() + 8); buf.extend_from_slice(&effective_key); buf.extend_from_slice(b"::mq:dlq"); Bytes::from(buf) }; - // Phase 2a: gate on is_initialized(); new path uses ShardSlice. - let len: i64 = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| { - match db.get_stream_mut(&dlq_key) { - Ok(Some(stream)) => stream.length as i64, - _ => 0i64, - } - }) - } else { - let mut db_guard = ctx.shard_databases.write_db(owner, conn.selected_db); - match db_guard.get_stream_mut(&dlq_key) { - Ok(Some(stream)) => stream.length as i64, - _ => 0i64, - } - }; - responses.push(Frame::Integer(len)); + // Owner-routes via MqCommand hop. + { + let key_prefix = mq_ws_prefix(conn.workspace_id.as_ref()); + let response = mq_hop_or_local( + owner, + conn.selected_db, + ctx, + key_prefix, + std::sync::Arc::new(frame.clone()), + ) + .await; + responses.push(response); + } } Err(e) => responses.push(e), } @@ -741,36 +441,26 @@ pub(super) fn try_handle_mq_command( if sub.eq_ignore_ascii_case(b"TRIGGER") { match validate_mq_trigger(cmd_args) { - Ok((queue_key, callback_cmd, debounce_ms)) => { + Ok((queue_key, _callback_cmd, _debounce_ms)) => { let effective_key = crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); - let trig_key = if let Some(ws_id) = conn.workspace_id.as_ref() { - let ws_hex = ws_id.as_hex(); - let mut k = Vec::with_capacity(ws_hex.len() + 1 + queue_key.len()); - k.extend_from_slice(ws_hex.as_bytes()); - k.push(b':'); - k.extend_from_slice(&queue_key); - Bytes::from(k) - } else { - queue_key.clone() - }; // Owner's registry: its event-loop tick fires triggers // (timers.rs documents the home shard as authoritative). let owner = crate::shard::dispatch::key_to_shard(&effective_key, ctx.num_shards); - let entry = crate::mq::TriggerEntry { - queue_key: effective_key, - callback_cmd, - debounce_ms, - last_fire_ms: 0, - pending_fire_ms: 0, - }; + // Owner-routes via MqCommand hop (execute_mq_on_owner + // registers the trigger in the owner's slice registry). { - let mut guard = ctx.shard_databases.trigger_registry(owner); - let reg = - guard.get_or_insert_with(|| Box::new(crate::mq::TriggerRegistry::new())); - reg.register(trig_key, entry); + let key_prefix = mq_ws_prefix(conn.workspace_id.as_ref()); + let response = mq_hop_or_local( + owner, + conn.selected_db, + ctx, + key_prefix, + std::sync::Arc::new(frame.clone()), + ) + .await; + responses.push(response); } - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); } Err(e) => responses.push(e), } @@ -801,6 +491,76 @@ pub(super) fn try_handle_mq_command( true } +// ── shardslice-migration Wave B1 helpers ───────────────────────────────────── + +/// Build the workspace prefix `"{ws_hex}:"` used as `MqCommandPayload.key_prefix`. +/// +/// When the connection is not workspace-bound this returns `Bytes::new()`, +/// matching `MqCommandPayload`'s "empty = no prefix" contract. +#[inline] +fn mq_ws_prefix(workspace_id: Option<&crate::workspace::WorkspaceId>) -> Bytes { + match workspace_id { + None => Bytes::new(), + Some(ws_id) => { + let ws_hex = ws_id.as_hex(); + // "{" + 32 hex + "}" + ":" = 35 bytes + let mut buf = Vec::with_capacity(35); + buf.push(b'{'); + buf.extend_from_slice(ws_hex.as_bytes()); + buf.push(b'}'); + buf.push(b':'); + Bytes::from(buf) + } + } +} + +/// Send an MQ command to the owning shard (or run locally if self). +/// +/// When `owner == ctx.shard_id` the command runs synchronously via +/// `mq_exec::execute_mq_on_owner` (same thread, slice live — no hop overhead). +/// Otherwise a `ShardMessage::MqCommand` is pushed to the SPSC ring and the +/// caller awaits the oneshot reply. +/// +/// `db_index` = `conn.selected_db`; `owner` = `key_to_shard(effective_key)`. +/// +/// Returns the `Frame` response to push to the client. +async fn mq_hop_or_local( + owner: usize, + db_index: usize, + ctx: &ConnectionContext, + key_prefix: Bytes, + command: std::sync::Arc, +) -> crate::protocol::Frame { + if owner == ctx.shard_id { + // Self-short-circuit: run directly on this shard's slice. + crate::shard::mq_exec::execute_mq_on_owner(db_index, key_prefix, command) + } else { + // Cross-shard hop via MqCommand SPSC message (GraphCommand precedent). + let (reply_tx, reply_rx) = crate::runtime::channel::oneshot(); + let payload = crate::shard::dispatch::MqCommandPayload { + db_index, + key_prefix, + command, + reply_tx, + }; + let msg = crate::shard::dispatch::ShardMessage::MqCommand(Box::new(payload)); + crate::shard::coordinator::spsc_send( + &ctx.dispatch_tx, + ctx.shard_id, + owner, + msg, + &ctx.spsc_notifiers, + ) + .await; + match reply_rx.recv().await { + Ok(f) => f, + Err(_) => crate::protocol::Frame::Error(bytes::Bytes::from_static( + b"ERR cross-shard MQ reply channel closed", + )), + } + } +} + /// Handle MULTI/EXEC/DISCARD commands. Returns `true` if consumed. pub(super) fn try_handle_multi_exec( cmd: &[u8], @@ -935,57 +695,31 @@ pub(super) async fn try_handle_graph_command( } } } - // Phase 2a: gate on is_initialized(); new path uses ShardSlice::graph_store directly. + // Unconditional slice path: dispatch to ShardSlice::graph_store directly. let (response, wal_records, cypher_intents, cypher_undo_ops) = - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - if crate::command::graph::is_graph_write_cmd(cmd) - || (cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") - && crate::command::graph::is_cypher_write_query(cmd_args)) - { - let gs = &mut s.graph_store; - let (resp, cypher_intents, undo_ops) = - if cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") { - crate::command::graph::graph_query_or_write(gs, cmd_args) - } else { - ( - crate::command::graph::dispatch_graph_write(gs, cmd, cmd_args), - Vec::new(), - Vec::new(), - ) - }; - let records = gs.drain_wal(); - (resp, records, cypher_intents, undo_ops) + crate::shard::slice::with_shard(|s| { + if crate::command::graph::is_graph_write_cmd(cmd) + || (cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") + && crate::command::graph::is_cypher_write_query(cmd_args)) + { + let gs = &mut s.graph_store; + let (resp, cypher_intents, undo_ops) = if cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") { + crate::command::graph::graph_query_or_write(gs, cmd_args) } else { - let gs = &s.graph_store; - let resp = crate::command::graph::dispatch_graph_read(gs, cmd, cmd_args); - (resp, Vec::new(), Vec::new(), Vec::new()) - } - }) - } else if crate::command::graph::is_graph_write_cmd(cmd) - || (cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") - && crate::command::graph::is_cypher_write_query(cmd_args)) - { - let mut gs = ctx.shard_databases.graph_store_write(ctx.shard_id); - let (resp, cypher_intents, undo_ops) = if cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") { - // Phase 167 (CYP-01/02): capture Cypher-created - // nodes/edges so TXN.ABORT can roll them back via - // CrossStoreTxn::record_graph. - crate::command::graph::graph_query_or_write(&mut gs, cmd_args) + ( + crate::command::graph::dispatch_graph_write(gs, cmd, cmd_args), + Vec::new(), + Vec::new(), + ) + }; + let records = gs.drain_wal(); + (resp, records, cypher_intents, undo_ops) } else { - ( - crate::command::graph::dispatch_graph_write(&mut gs, cmd, cmd_args), - Vec::new(), - Vec::new(), - ) - }; - let records = gs.drain_wal(); - (resp, records, cypher_intents, undo_ops) - } else { - let gs = ctx.shard_databases.graph_store_read(ctx.shard_id); - let resp = crate::command::graph::dispatch_graph_read(&gs, cmd, cmd_args); - (resp, Vec::new(), Vec::new(), Vec::new()) - }; + let gs = &s.graph_store; + let resp = crate::command::graph::dispatch_graph_read(gs, cmd, cmd_args); + (resp, Vec::new(), Vec::new(), Vec::new()) + } + }); // Phase 166: record graph intent for TXN rollback. // Captures explicit ADDNODE/ADDEDGE by response id plus // Phase 167 Cypher CREATE/MERGE via intents returned from diff --git a/src/server/conn/handler_sharded/dispatch.rs b/src/server/conn/handler_sharded/dispatch.rs index 84eeab36..8fc46076 100644 --- a/src/server/conn/handler_sharded/dispatch.rs +++ b/src/server/conn/handler_sharded/dispatch.rs @@ -316,16 +316,14 @@ pub(super) fn try_handle_info( if !cmd.eq_ignore_ascii_case(b"INFO") { return false; } - let guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); - let response_text = { - let resp_frame = conn_cmd::info_readonly(&guard, cmd_args); + // ShardSlice path: access the local shard's database via thread-local. + let mut response_text = crate::shard::slice::with_shard_db(conn.selected_db, |db| { + let resp_frame = conn_cmd::info_readonly(db, cmd_args); match resp_frame { Frame::BulkString(b) => String::from_utf8_lossy(&b).to_string(), _ => String::new(), } - }; - drop(guard); - let mut response_text = response_text; + }); if let Some(ref rs) = ctx.repl_state { if let Ok(rs_guard) = rs.try_read() { response_text.push_str(&crate::replication::handshake::build_info_replication( diff --git a/src/server/conn/handler_sharded/ft.rs b/src/server/conn/handler_sharded/ft.rs index 3c7a9ec3..5382682d 100644 --- a/src/server/conn/handler_sharded/ft.rs +++ b/src/server/conn/handler_sharded/ft.rs @@ -113,6 +113,9 @@ pub(super) async fn try_handle_ft_command( top_k, offset: limit_offset, count: limit_count, + // CHANGE D: thread the parsed FILTER tree (if any) into + // HybridQuery so execute_hybrid_search_local / scatter + // can apply it pre-RRF on both branches (CHANGE B/F). filter: partial.filter, }; // Phase 171 HYB-02 / SCAT-02: resolve AS_OF / @@ -121,7 +124,6 @@ pub(super) async fn try_handle_ft_command( let as_of_lsn = match resolve_ft_search_as_of_lsn( cmd_args, Some(&ctx.shard_databases), - ctx.shard_id, conn.active_cross_txn.as_ref(), ) { Ok(lsn) => lsn, @@ -223,9 +225,7 @@ pub(super) async fn try_handle_ft_command( // We use the TextIndex's own field_analyzers (same pipeline used at index time). type ParseResult = Result<(Vec, Option), String>; - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - // The inner closure scans the TextIndex via text_store; same logic in - // both branches. + // The inner closure scans the TextIndex via text_store. let parse_body = |ts: &crate::text::store::TextStore| -> ParseResult { match ts.get_index(&index_name) { None => Err("ERR no such index".to_owned()), @@ -261,14 +261,9 @@ pub(super) async fn try_handle_ft_command( }, } }; - let parse_result: ParseResult = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| parse_body(&s.text_store)) - } else { - let ts = ctx.shard_databases.text_store(ctx.shard_id); - let r = parse_body(&ts); - drop(ts); - r - }; + // Unconditional slice path: ShardSlice is always initialized. + let parse_result: ParseResult = + crate::shard::slice::with_shard(|s| parse_body(&s.text_store)); let (query_terms, field_idx) = match parse_result { Ok(t) => t, @@ -329,7 +324,6 @@ pub(super) async fn try_handle_ft_command( match resolve_ft_search_as_of_lsn( cmd_args, Some(&ctx.shard_databases), - ctx.shard_id, conn.active_cross_txn.as_ref(), ) { Err(err_frame) => err_frame, @@ -468,7 +462,6 @@ pub(super) async fn try_handle_ft_command( offset.saturating_add(count) } .max(1); - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. let lookup_body = |ts: &crate::text::store::TextStore| -> Frame { match ts.get_index(&index_name) { @@ -485,16 +478,10 @@ pub(super) async fn try_handle_ft_command( } } }; - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - lookup_body(&s.text_store) - }) - } else { - let ts_guard = ctx.shard_databases.text_store(ctx.shard_id); - let r = lookup_body(&ts_guard); - drop(ts_guard); - r - }; + // Unconditional slice path: ShardSlice is always initialized. + let response = crate::shard::slice::with_shard(|s| { + lookup_body(&s.text_store) + }); let mut response = response; if let Some(ws_id) = conn.workspace_id.as_ref() { strip_workspace_prefix_from_response( @@ -513,8 +500,7 @@ pub(super) async fn try_handle_ft_command( return true; } } - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - // text_store + read_db(0) accessed in ONE with_shard closure + // text_store + databases[0] accessed in ONE with_shard closure // (multi-resource arm) — text_index borrows from text_store and is // also passed to apply_post_processing alongside &Database. let (offset, count) = @@ -608,24 +594,12 @@ pub(super) async fn try_handle_ft_command( response }; - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - let db_opt: Option<&crate::storage::db::Database> = - if need_db { Some(&s.databases[0]) } else { None }; - text_search_body(&s.text_store, db_opt) - }) - } else { - let ts_guard = ctx.shard_databases.text_store(ctx.shard_id); - let db_guard_opt = if need_db { - Some(ctx.shard_databases.read_db(ctx.shard_id, 0)) - } else { - None - }; - let response = text_search_body(&ts_guard, db_guard_opt.as_deref()); - drop(db_guard_opt); - drop(ts_guard); - response - }; + // Unconditional slice path: ShardSlice is always initialized. + let response = crate::shard::slice::with_shard(|s| { + let db_opt: Option<&crate::storage::db::Database> = + if need_db { Some(&s.databases[0]) } else { None }; + text_search_body(&s.text_store, db_opt) + }); let mut response = response; if let Some(ws_id) = conn.workspace_id.as_ref() { strip_workspace_prefix_from_response(ws_id, cmd, &mut response); @@ -637,135 +611,33 @@ pub(super) async fn try_handle_ft_command( } } } - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - // All five companion stores (vector_store, text_store, databases[0], - // graph_store) are accessed from ONE with_shard closure to avoid - // re-entrant RefCell borrows — this is the canonical multi-resource arm. + // All companion stores (vector_store, text_store, databases[0], graph_store) + // accessed from ONE with_shard closure to avoid re-entrant RefCell borrows. // // Resolve AS_OF outside the closure (acquires shard_databases for the // resolver helper, which itself reads vector_store; doing this inside - // would re-enter with_shard on the new path). + // would re-enter with_shard). let as_of_lsn_opt: Option> = if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { Some(resolve_ft_search_as_of_lsn( cmd_args, Some(&ctx.shard_databases), - ctx.shard_id, conn.active_cross_txn.as_ref(), )) } else { None }; - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - if cmd.eq_ignore_ascii_case(b"FT.CREATE") { - crate::command::vector_search::ft_create( - &mut s.vector_store, - &mut s.text_store, - cmd_args, - ) - } else if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { - #[allow(clippy::unwrap_used)] // as_of_lsn_opt is Some when cmd == FT.SEARCH - match as_of_lsn_opt.unwrap() { - Err(err_frame) => err_frame, - Ok(as_of_lsn) => { - let has_session = cmd_args.iter().any(|a| { - if let Frame::BulkString(b) = a { - b.eq_ignore_ascii_case(b"SESSION") - } else { - false - } - }); - if has_session { - // Borrow databases[0] disjointly from vector_store/text_store. - let (vs, ts, dbs) = - (&mut s.vector_store, &s.text_store, &mut s.databases); - crate::command::vector_search::ft_search( - vs, - cmd_args, - Some(&mut dbs[0]), - Some(ts), - as_of_lsn, - ) - } else { - crate::command::vector_search::ft_search( - &mut s.vector_store, - cmd_args, - None, - Some(&s.text_store), - as_of_lsn, - ) - } - } - } - } else if cmd.eq_ignore_ascii_case(b"FT.DROPINDEX") { - let (vs, ts, dbs) = (&mut s.vector_store, &mut s.text_store, &mut s.databases); - crate::command::vector_search::ft_dropindex(vs, ts, Some(&mut dbs[0]), cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT.INFO") { - crate::command::vector_search::ft_info(&s.vector_store, &s.text_store, cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT._LIST") { - crate::command::vector_search::ft_list(&s.vector_store) - } else if cmd.eq_ignore_ascii_case(b"FT.COMPACT") { - crate::command::vector_search::ft_compact( - &mut s.vector_store, - &mut s.text_store, - cmd_args, - ) - } else if cmd.eq_ignore_ascii_case(b"FT.CACHESEARCH") { - crate::command::vector_search::cache_search::ft_cachesearch( - &mut s.vector_store, - cmd_args, - ) - } else if cmd.eq_ignore_ascii_case(b"FT.CONFIG") { - crate::command::vector_search::ft_config( - &mut s.vector_store, - &mut s.text_store, - cmd_args, - ) - } else if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") { - let (vs, dbs) = (&mut s.vector_store, &mut s.databases); - crate::command::vector_search::recommend::ft_recommend( - vs, - cmd_args, - Some(&mut dbs[0]), - ) - } else if cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") { - #[cfg(feature = "graph")] - { - crate::command::vector_search::navigate::ft_navigate( - &mut s.vector_store, - Some(&s.graph_store), - cmd_args, - None, - ) - } - #[cfg(not(feature = "graph"))] - { - Frame::Error(Bytes::from_static( - b"ERR FT.NAVIGATE requires graph feature", - )) - } - } else if cmd.eq_ignore_ascii_case(b"FT.EXPAND") { - #[cfg(feature = "graph")] - { - crate::command::vector_search::ft_expand(&s.graph_store, cmd_args) - } - #[cfg(not(feature = "graph"))] - { - Frame::Error(Bytes::from_static(b"ERR FT.EXPAND requires graph feature")) - } - } else { - Frame::Error(Bytes::from_static(b"ERR unknown FT.* command")) - } - }) - } else { - let shard_databases_ref = &ctx.shard_databases; - let mut vs = shard_databases_ref.vector_store(ctx.shard_id); - let mut ts = shard_databases_ref.text_store(ctx.shard_id); + // Unconditional slice path: ShardSlice is always initialized. + // All companion stores (vector_store, text_store, databases[0], graph_store) + // accessed from ONE with_shard closure to avoid re-entrant RefCell borrows. + let response = crate::shard::slice::with_shard(|s| { if cmd.eq_ignore_ascii_case(b"FT.CREATE") { - crate::command::vector_search::ft_create(&mut vs, &mut ts, cmd_args) + crate::command::vector_search::ft_create( + &mut s.vector_store, + &mut s.text_store, + cmd_args, + ) } else if cmd.eq_ignore_ascii_case(b"FT.SEARCH") { - // Resolve AS_OF temporal clause + TXN snapshot precedence (TEMP-04, ACID-09). #[allow(clippy::unwrap_used)] // as_of_lsn_opt is Some when cmd == FT.SEARCH match as_of_lsn_opt.unwrap() { Err(err_frame) => err_frame, @@ -778,57 +650,59 @@ pub(super) async fn try_handle_ft_command( } }); if has_session { - let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); + // Borrow databases[0] disjointly from vector_store/text_store. + let (vs, ts, dbs) = (&mut s.vector_store, &s.text_store, &mut s.databases); crate::command::vector_search::ft_search( - &mut vs, + vs, cmd_args, - Some(&mut *db_guard), - Some(&*ts), + Some(&mut dbs[0]), + Some(ts), as_of_lsn, ) } else { crate::command::vector_search::ft_search( - &mut vs, + &mut s.vector_store, cmd_args, None, - Some(&*ts), + Some(&s.text_store), as_of_lsn, ) } } } } else if cmd.eq_ignore_ascii_case(b"FT.DROPINDEX") { - let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); - crate::command::vector_search::ft_dropindex( - &mut vs, - &mut ts, - Some(&mut *db_guard), - cmd_args, - ) + let (vs, ts, dbs) = (&mut s.vector_store, &mut s.text_store, &mut s.databases); + crate::command::vector_search::ft_dropindex(vs, ts, Some(&mut dbs[0]), cmd_args) } else if cmd.eq_ignore_ascii_case(b"FT.INFO") { - crate::command::vector_search::ft_info(&vs, &ts, cmd_args) + crate::command::vector_search::ft_info(&s.vector_store, &s.text_store, cmd_args) } else if cmd.eq_ignore_ascii_case(b"FT._LIST") { - crate::command::vector_search::ft_list(&vs) + crate::command::vector_search::ft_list(&s.vector_store) } else if cmd.eq_ignore_ascii_case(b"FT.COMPACT") { - crate::command::vector_search::ft_compact(&mut vs, &mut ts, cmd_args) + crate::command::vector_search::ft_compact( + &mut s.vector_store, + &mut s.text_store, + cmd_args, + ) } else if cmd.eq_ignore_ascii_case(b"FT.CACHESEARCH") { - crate::command::vector_search::cache_search::ft_cachesearch(&mut vs, cmd_args) + crate::command::vector_search::cache_search::ft_cachesearch( + &mut s.vector_store, + cmd_args, + ) } else if cmd.eq_ignore_ascii_case(b"FT.CONFIG") { - crate::command::vector_search::ft_config(&mut vs, &mut ts, cmd_args) - } else if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") { - let mut db_guard = shard_databases_ref.write_db(ctx.shard_id, 0); - crate::command::vector_search::recommend::ft_recommend( - &mut vs, + crate::command::vector_search::ft_config( + &mut s.vector_store, + &mut s.text_store, cmd_args, - Some(&mut *db_guard), ) + } else if cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") { + let (vs, dbs) = (&mut s.vector_store, &mut s.databases); + crate::command::vector_search::recommend::ft_recommend(vs, cmd_args, Some(&mut dbs[0])) } else if cmd.eq_ignore_ascii_case(b"FT.NAVIGATE") { #[cfg(feature = "graph")] { - let graph_guard = shard_databases_ref.graph_store_read(ctx.shard_id); crate::command::vector_search::navigate::ft_navigate( - &mut vs, - Some(&graph_guard), + &mut s.vector_store, + Some(&s.graph_store), cmd_args, None, ) @@ -842,28 +716,17 @@ pub(super) async fn try_handle_ft_command( } else if cmd.eq_ignore_ascii_case(b"FT.EXPAND") { #[cfg(feature = "graph")] { - let graph_guard = shard_databases_ref.graph_store_read(ctx.shard_id); - crate::command::vector_search::ft_expand(&graph_guard, cmd_args) + crate::command::vector_search::ft_expand(&s.graph_store, cmd_args) } #[cfg(not(feature = "graph"))] { Frame::Error(Bytes::from_static(b"ERR FT.EXPAND requires graph feature")) } - } else if cmd.eq_ignore_ascii_case(b"FT.INVALIDATE_RANGE") { - #[cfg(feature = "text-index")] - { - crate::command::vector_search::ft_invalidate_range(&mut ts, cmd_args) - } - #[cfg(not(feature = "text-index"))] - { - Frame::Error(Bytes::from_static( - b"ERR FT.INVALIDATE_RANGE requires text-index feature", - )) - } } else { Frame::Error(Bytes::from_static(b"ERR unknown FT.* command")) } - }; + }); + let mut response = response; if let Some(ws_id) = conn.workspace_id.as_ref() { strip_workspace_prefix_from_response(ws_id, cmd, &mut response); diff --git a/src/server/conn/handler_sharded/mod.rs b/src/server/conn/handler_sharded/mod.rs index 668092cc..53d6c90f 100644 --- a/src/server/conn/handler_sharded/mod.rs +++ b/src/server/conn/handler_sharded/mod.rs @@ -381,7 +381,7 @@ pub(crate) async fn handle_connection_sharded_inner< // Per-batch dispatch-path accumulators — flushed once at end of // batch so we pay one global atomic per path instead of N. let mut local_dispatches: u32 = 0; - let mut cross_read_fast_dispatches: u32 = 0; + let cross_read_fast_dispatches: u32 = 0; // fast-path removed in wave-e2; always 0 let mut cross_spsc_dispatches: u32 = 0; for frame in batch { @@ -507,35 +507,20 @@ pub(crate) async fn handle_connection_sharded_inner< // --- Lua scripting: EVAL / EVALSHA --- if cmd.eq_ignore_ascii_case(b"EVAL") || cmd.eq_ignore_ascii_case(b"EVALSHA") { let db_count = ctx.shard_databases.db_count(); - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| { - if cmd.eq_ignore_ascii_case(b"EVAL") { - crate::scripting::handle_eval( - &ctx.lua, &ctx.script_cache, cmd_args, db, - ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, - ) - } else { - crate::scripting::handle_evalsha( - &ctx.lua, &ctx.script_cache, cmd_args, db, - ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, - ) - } - }) - } else { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + // Unconditional slice path: ShardSlice is always initialized. + let response = crate::shard::slice::with_shard_db(conn.selected_db, |db| { if cmd.eq_ignore_ascii_case(b"EVAL") { crate::scripting::handle_eval( - &ctx.lua, &ctx.script_cache, cmd_args, &mut guard, + &ctx.lua, &ctx.script_cache, cmd_args, db, ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, ) } else { crate::scripting::handle_evalsha( - &ctx.lua, &ctx.script_cache, cmd_args, &mut guard, + &ctx.lua, &ctx.script_cache, cmd_args, db, ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, ) } - }; + }); responses.push(response); continue; } @@ -749,41 +734,25 @@ pub(crate) async fn handle_connection_sharded_inner< } if cmd.eq_ignore_ascii_case(b"FCALL") { let db_count = ctx.shard_databases.db_count(); - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| { - crate::command::functions::handle_fcall( - &func_registry.borrow(), cmd_args, db, - ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, - ) - }) - } else { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + // Unconditional slice path: ShardSlice is always initialized. + let response = crate::shard::slice::with_shard_db(conn.selected_db, |db| { crate::command::functions::handle_fcall( - &func_registry.borrow(), cmd_args, &mut guard, + &func_registry.borrow(), cmd_args, db, ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, ) - }; + }); responses.push(response); continue; } if cmd.eq_ignore_ascii_case(b"FCALL_RO") { let db_count = ctx.shard_databases.db_count(); - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| { - crate::command::functions::handle_fcall_ro( - &func_registry.borrow(), cmd_args, db, - ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, - ) - }) - } else { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); + // Unconditional slice path: ShardSlice is always initialized. + let response = crate::shard::slice::with_shard_db(conn.selected_db, |db| { crate::command::functions::handle_fcall_ro( - &func_registry.borrow(), cmd_args, &mut guard, + &func_registry.borrow(), cmd_args, db, ctx.shard_id, ctx.num_shards, conn.selected_db, db_count, ) - }; + }); responses.push(response); continue; } @@ -838,7 +807,7 @@ pub(crate) async fn handle_connection_sharded_inner< if txn::try_handle_txn_begin(cmd, cmd_args, &mut conn, ctx, &mut responses) { continue; } - if txn::try_handle_txn_commit(cmd, cmd_args, &mut conn, ctx, &mut responses) { + if txn::try_handle_txn_commit(cmd, cmd_args, &mut conn, ctx, &mut responses).await { continue; } if txn::try_handle_txn_abort(cmd, cmd_args, &mut conn, ctx, &mut responses) @@ -864,12 +833,12 @@ pub(crate) async fn handle_connection_sharded_inner< } // --- WS.* --- - if write::try_handle_ws_command(cmd, cmd_args, &mut conn, ctx, &mut responses) { + if write::try_handle_ws_command(cmd, cmd_args, &mut conn, ctx, &mut responses).await { continue; } // --- MQ.* --- - if write::try_handle_mq_command(cmd, cmd_args, &mut conn, ctx, &mut responses) { + if write::try_handle_mq_command(cmd, cmd_args, &frame, &mut conn, ctx, &mut responses).await { continue; } @@ -983,18 +952,10 @@ pub(crate) async fn handle_connection_sharded_inner< // Routes directly to VectorStore's TransactionManager on the // local shard. Bypasses write-stall guards (admin command). if cmd.eq_ignore_ascii_case(b"KILL") { - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::server_admin::kill_snapshot(&mut s.vector_store, cmd_args) - }) - } else { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - let response = - crate::command::server_admin::kill_snapshot(&mut vs, cmd_args); - drop(vs); - response - }; + // Unconditional slice path: ShardSlice is always initialized. + let response = crate::shard::slice::with_shard(|s| { + crate::command::server_admin::kill_snapshot(&mut s.vector_store, cmd_args) + }); responses.push(response); continue; } @@ -1010,84 +971,44 @@ pub(crate) async fn handle_connection_sharded_inner< crate::command::helpers::extract_bytes(sub_frame) { if sub.eq_ignore_ascii_case(b"VECTOR") { - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::server_admin::vacuum_vector( - &mut s.vector_store, - &cmd_args[1..], - ) - }) - } else { - let mut vs = - ctx.shard_databases.vector_store(ctx.shard_id); - let r = - crate::command::server_admin::vacuum_vector( - &mut vs, - &cmd_args[1..], - ); - drop(vs); - r - }; + // Unconditional slice path: ShardSlice is always initialized. + let response = crate::shard::slice::with_shard(|s| { + crate::command::server_admin::vacuum_vector( + &mut s.vector_store, + &cmd_args[1..], + ) + }); responses.push(response); continue; } #[cfg(feature = "graph")] if sub.eq_ignore_ascii_case(b"GRAPH") { - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { - let graph_merge_max = ctx.config.graph_merge_max_segments; - let graph_dead = ctx.config.graph_dead_edge_trigger; - crate::shard::slice::with_shard(|s| { - crate::command::server_admin::vacuum_graph( - &mut s.graph_store, - &cmd_args[1..], - graph_merge_max, - graph_dead, - ) - }) - } else { - let mut gs = ctx - .shard_databases - .graph_store_write(ctx.shard_id); - let r = - crate::command::server_admin::vacuum_graph( - &mut gs, - &cmd_args[1..], - ctx.config.graph_merge_max_segments, - ctx.config.graph_dead_edge_trigger, - ); - drop(gs); - r - }; + // Unconditional slice path: ShardSlice is always initialized. + let graph_merge_max = ctx.config.graph_merge_max_segments; + let graph_dead = ctx.config.graph_dead_edge_trigger; + let response = crate::shard::slice::with_shard(|s| { + crate::command::server_admin::vacuum_graph( + &mut s.graph_store, + &cmd_args[1..], + graph_merge_max, + graph_dead, + ) + }); responses.push(response); continue; } } } - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::server_admin::vacuum( - &mut s.vector_store, - None, - None, - cmd_args, - crate::command::server_admin::DEFAULT_VACUUM_PRUNE_MARGIN, - ) - }) - } else { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - let r = crate::command::server_admin::vacuum( - &mut vs, - None, // manifest — not available in connection handler - None, // wal_v3 — not available in connection handler + // Unconditional slice path: ShardSlice is always initialized. + let response = crate::shard::slice::with_shard(|s| { + crate::command::server_admin::vacuum( + &mut s.vector_store, + None, + None, cmd_args, - crate::command::server_admin::DEFAULT_VACUUM_PRUNE_MARGIN, // see server_admin.rs - ); - drop(vs); - r - }; + crate::command::server_admin::DEFAULT_VACUUM_PRUNE_MARGIN, + ) + }); responses.push(response); continue; } @@ -1097,23 +1018,14 @@ pub(crate) async fn handle_connection_sharded_inner< if let Some(sub) = cmd_args.first() { if let Some(s) = crate::command::helpers::extract_bytes(sub) { if s.eq_ignore_ascii_case(b"RECLAMATION") { - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::server_admin::debug_reclamation( - &s.vector_store, - None, - None, - ) - }) - } else { - let vs = ctx.shard_databases.vector_store(ctx.shard_id); - let r = crate::command::server_admin::debug_reclamation( - &vs, None, None, - ); - drop(vs); - r - }; + // Unconditional slice path: ShardSlice is always initialized. + let response = crate::shard::slice::with_shard(|s| { + crate::command::server_admin::debug_reclamation( + &s.vector_store, + None, + None, + ) + }); responses.push(response); continue; } @@ -1193,21 +1105,12 @@ pub(crate) async fn handle_connection_sharded_inner< Err(e) => e, Ok((_key, dst_db)) if dst_db == src_db => Frame::Integer(0), Ok((key, dst_db)) => { - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - ksmv::with_two_slice_dbs(&mut s.databases, src_db, dst_db, |src, dst| { - ksmv::move_core(src, dst, &key) - }) + // Unconditional slice path: ShardSlice is always initialized. + crate::shard::slice::with_shard(|s| { + ksmv::with_two_slice_dbs(&mut s.databases, src_db, dst_db, |src, dst| { + ksmv::move_core(src, dst, &key) }) - } else { - // Lock ordering (lower index first) prevents deadlock - // with concurrent reverse MOVE from another connection. - ksmv::with_two_dbs_locked( - &ctx.shard_databases.all_shard_dbs()[ctx.shard_id], - src_db, dst_db, - |src, dst| ksmv::move_core(src, dst, &key), - ) - } + }) } }; // AOF only on actual success (:1). Matches handler_single @@ -1252,19 +1155,12 @@ pub(crate) async fn handle_connection_sharded_inner< let response = match copy_result { Err(e) => e, Ok(ca) => { - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - ksmv::with_two_slice_dbs(&mut s.databases, src_db, ca.dst_db, |src, dst| { - ksmv::copy_core(src, dst, &ca.src_key, &ca.dst_key, ca.replace) - }) + // Unconditional slice path: ShardSlice is always initialized. + crate::shard::slice::with_shard(|s| { + ksmv::with_two_slice_dbs(&mut s.databases, src_db, ca.dst_db, |src, dst| { + ksmv::copy_core(src, dst, &ca.src_key, &ca.dst_key, ca.replace) }) - } else { - ksmv::with_two_dbs_locked( - &ctx.shard_databases.all_shard_dbs()[ctx.shard_id], - src_db, ca.dst_db, - |src, dst| ksmv::copy_core(src, dst, &ca.src_key, &ca.dst_key, ca.replace), - ) - } + }) } }; // AOF only on actual success (:1). Matches handler_single @@ -1296,12 +1192,12 @@ pub(crate) async fn handle_connection_sharded_inner< if metadata::is_write(cmd) { // WRITE PATH: single lock acquisition for eviction + dispatch. // - // Phase 2f: gate on is_initialized(); new path uses ShardSlice - // directly. Eviction → KV undo-log capture → dispatch → wakeup - // run inside the with_shard_db closure so the &mut Database - // borrow stays disjoint from the post-drop HSET auto-index - // and DEL/UNLINK auto-delete paths (which need vector/text - // stores, separate ShardSlice fields). + // Unconditional slice path: ShardSlice is always initialized. + // Eviction → KV undo-log capture → dispatch → wakeup run inside + // the with_shard_db closure so the &mut Database borrow stays + // disjoint from the post-drop HSET auto-index and DEL/UNLINK + // auto-delete paths (which need vector/text stores, separate + // ShardSlice fields). let db_count = ctx.shard_databases.db_count(); // Returns Ok((response, sample_latency, dispatch_start)) on // success or Err(oom_frame) when eviction fails (caller @@ -1310,9 +1206,24 @@ pub(crate) async fn handle_connection_sharded_inner< (Frame, bool, Option), Frame, >; - let mut do_write = |db: &mut crate::storage::db::Database, + // Takes the whole ShardSlice (not just the Database): + // the KV undo-log capture below records write intents on + // s.kv_write_intents, and re-entering with_shard from + // inside a with_shard_db closure panics (slice + // re-entrancy guard). `db` borrows s.databases only, so + // s.kv_write_intents stays accessible (disjoint NLL + // field borrows). + let mut do_write = |s: &mut crate::shard::slice::ShardSlice, conn: &mut super::core::ConnectionState| -> WriteOutcome { + let db_len = s.databases.len(); + #[allow(clippy::unwrap_used)] // mirror with_shard_db's panic semantics + let db = s.databases.get_mut(conn.selected_db).unwrap_or_else(|| { + panic!( + "db_index {} out of bounds ({db_len} databases on shard {})", + conn.selected_db, s.shard_id + ) + }); let rt = ctx.runtime_config.read(); // Disk-offload: when a per-shard spill sender is wired // (ConnCtx populated by spawn_tokio_connection), evicted @@ -1359,8 +1270,10 @@ pub(crate) async fn handle_connection_sharded_inner< txn.kv_undo.record_delete(key_bytes.clone(), old_entry); let lsn = txn.snapshot_lsn; let tid = txn.txn_id; - ctx.shard_databases.kv_intents(ctx.shard_id) - .record_write(key_bytes.clone(), lsn, tid); + // Direct field access — `s` is this + // closure's own param; re-entering + // with_shard here panics. + s.kv_write_intents.record_write(key_bytes.clone(), lsn, tid); } } } @@ -1372,8 +1285,8 @@ pub(crate) async fn handle_connection_sharded_inner< None => txn.kv_undo.record_insert(key.clone()), Some(entry) => txn.kv_undo.record_update(key.clone(), entry), } - ctx.shard_databases.kv_intents(ctx.shard_id) - .record_write(key.clone(), lsn, tid); + // Direct field access — see DEL/UNLINK arm above. + s.kv_write_intents.record_write(key.clone(), lsn, tid); } } @@ -1403,14 +1316,9 @@ pub(crate) async fn handle_connection_sharded_inner< Ok((response, sample_latency, dispatch_start)) }; - let write_outcome: WriteOutcome = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| do_write(db, &mut conn)) - } else { - let mut guard = ctx.shard_databases.write_db(ctx.shard_id, conn.selected_db); - let r = do_write(&mut *guard, &mut conn); - drop(guard); - r - }; + // Unconditional slice path: ShardSlice is always initialized. + let write_outcome: WriteOutcome = + crate::shard::slice::with_shard(|s| do_write(s, &mut conn)); let (mut response, _sample_latency, dispatch_start): (Frame, bool, Option) = match write_outcome { @@ -1468,26 +1376,15 @@ pub(crate) async fn handle_connection_sharded_inner< .as_ref() .map(|t| t.txn_id) .unwrap_or(0); - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - // vector_store + text_store accessed in ONE with_shard closure - // (multi-resource arm) to avoid re-entrant RefCell borrow. - let inserted = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - if active_txn_id != 0 { - crate::shard::spsc_handler::auto_index_hset_public_txn(&mut s.vector_store, &mut s.text_store, &key, cmd_args, active_txn_id) - } else { - crate::shard::spsc_handler::auto_index_hset_public(&mut s.vector_store, &mut s.text_store, &key, cmd_args) - } - }) - } else { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); - let mut ts = ctx.shard_databases.text_store(ctx.shard_id); + // Unconditional slice path: vector_store + text_store accessed in ONE + // with_shard closure (multi-resource arm) to avoid re-entrant RefCell borrow. + let inserted = crate::shard::slice::with_shard(|s| { if active_txn_id != 0 { - crate::shard::spsc_handler::auto_index_hset_public_txn(&mut vs, &mut *ts, &key, cmd_args, active_txn_id) + crate::shard::spsc_handler::auto_index_hset_public_txn(&mut s.vector_store, &mut s.text_store, &key, cmd_args, active_txn_id) } else { - crate::shard::spsc_handler::auto_index_hset_public(&mut vs, &mut *ts, &key, cmd_args) + crate::shard::spsc_handler::auto_index_hset_public(&mut s.vector_store, &mut s.text_store, &key, cmd_args) } - }; + }); // Push one VectorIntent per (index_name, key_hash) so // TXN.ABORT (Plan 166-03) can tombstone via // MutableSegment::mark_deleted_by_key_hash. @@ -1504,23 +1401,14 @@ pub(crate) async fn handle_connection_sharded_inner< if !matches!(response, Frame::Error(_)) && (cmd.eq_ignore_ascii_case(b"DEL") || cmd.eq_ignore_ascii_case(b"UNLINK")) { - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - for arg in cmd_args.iter() { - if let Some(key) = extract_bytes(arg) { - s.vector_store.mark_deleted_for_key(key.as_ref()); - } - } - }); - } else { - let mut vs = ctx.shard_databases.vector_store(ctx.shard_id); + // Unconditional slice path: ShardSlice is always initialized. + crate::shard::slice::with_shard(|s| { for arg in cmd_args.iter() { if let Some(key) = extract_bytes(arg) { - vs.mark_deleted_for_key(key.as_ref()); + s.vector_store.mark_deleted_for_key(key.as_ref()); } } - } + }); } // H1: durable path under appendfsync=always. let mut aof_failed = false; @@ -1569,19 +1457,14 @@ pub(crate) async fn handle_connection_sharded_inner< let my_txn_id = txn.txn_id; // Clone committed treemap to release vector_store lock // before acquiring kv_intents lock (lock ordering). - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let committed = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - s.vector_store.txn_manager().committed_treemap().clone() - }) - } else { - let vs = ctx.shard_databases.vector_store(ctx.shard_id); - vs.txn_manager().committed_treemap().clone() - }; - let visible = { - let intents = ctx.shard_databases.kv_intents(ctx.shard_id); - intents.is_key_visible(key.as_ref(), snapshot_lsn, my_txn_id, &committed) - }; + // Unconditional slice path: ShardSlice is always initialized. + let committed = crate::shard::slice::with_shard(|s| { + s.vector_store.txn_manager().committed_treemap().clone() + }); + // ShardSlice path: kv_write_intents lives on the thread-local slice. + let visible = crate::shard::slice::with_shard(|s| { + s.kv_write_intents.is_key_visible(key.as_ref(), snapshot_lsn, my_txn_id, &committed) + }); if !visible { responses.push(Frame::Null); continue; @@ -1590,23 +1473,15 @@ pub(crate) async fn handle_connection_sharded_inner< } } - // READ PATH: shared lock — no contention with other shards' reads. - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. + // READ PATH: unconditional slice path — ShardSlice is always initialized. let now_ms = ctx.cached_clock.ms(); let db_count = ctx.shard_databases.db_count(); conn.cmd_counter = conn.cmd_counter.wrapping_add(1); let sample_latency = (conn.cmd_counter & 0xF) == 0; let dispatch_start = sample_latency.then(std::time::Instant::now); - let result = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| { - dispatch_read(db, cmd, cmd_args, now_ms, &mut conn.selected_db, db_count) - }) - } else { - let guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); - let r = dispatch_read(&guard, cmd, cmd_args, now_ms, &mut conn.selected_db, db_count); - drop(guard); - r - }; + let result = crate::shard::slice::with_shard_db(conn.selected_db, |db| { + dispatch_read(db, cmd, cmd_args, now_ms, &mut conn.selected_db, db_count) + }); if let Ok(cmd_str) = std::str::from_utf8(cmd) { if let Some(start) = dispatch_start { let elapsed_us = start.elapsed().as_micros() as u64; @@ -1663,48 +1538,13 @@ pub(crate) async fn handle_connection_sharded_inner< ))); continue; } - // SHARED-READ FAST PATH: cross-shard reads bypass SPSC dispatch entirely. - // By this point conn.in_multi is false (MULTI queuing happens earlier with `continue`). - // Read commands execute directly on the target shard's database via RwLock read guard, - // avoiding ~88us of two async scheduling hops through the SPSC channel. - // - // Guard: if there are already pending writes for this target shard in the - // current pipeline batch, we must NOT take the fast path -- the read would - // execute before the deferred writes, violating command ordering. Fall through - // to SPSC dispatch to preserve pipeline semantics. - if !metadata::is_write(cmd) && !remote_groups.contains_key(&target) { - cross_read_fast_dispatches = cross_read_fast_dispatches.saturating_add(1); - let guard = ctx.shard_databases.read_db(target, conn.selected_db); - let now_ms = ctx.cached_clock.ms(); - let db_count = ctx.shard_databases.db_count(); - let result = dispatch_read(&guard, cmd, cmd_args, now_ms, &mut conn.selected_db, db_count); - drop(guard); - let response = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => { should_quit = true; f } - }; - if matches!(response, Frame::Error(_)) { - if let Ok(cmd_str) = std::str::from_utf8(cmd) { - crate::admin::metrics_setup::record_command_error_cached( - cmd_str, - &mut conn.cached_metrics, - ); - } - } - // Client tracking for cross-shard reads - if conn.tracking_state.enabled && !conn.tracking_state.bcast { - if let Some(key) = cmd_args.first().and_then(|f| extract_bytes(f)) { - ctx.tracking_table.borrow_mut().track_key(client_id, &key, conn.tracking_state.noloop); - } - } - let mut response = apply_resp3_conversion(cmd, response, conn.protocol_version); - if let Some(ws_id) = conn.workspace_id.as_ref() { - strip_workspace_prefix_from_response(ws_id, cmd, &mut response); - } - responses.push(response); - continue; - } - // Cross-shard write: deferred SPSC dispatch. + // Cross-shard dispatch via SPSC. + // NOTE(wave-e3): The shared-read fast-path that bypassed SPSC for cross-shard + // reads (using RwLock read guards) is removed in Wave E2. With ShardSlice as + // the live store, cross-thread DB access is not safe. All cross-shard commands + // (reads and writes) now go through SPSC dispatch. Restore the fast path in + // wave-e3 via AofFold-style snapshot reply if the latency regression matters. + // Cross-shard dispatch (reads and writes): deferred SPSC dispatch. // When workspace rewriting occurred, rebuild the frame with // prefixed args so the target shard stores the correct key. let dispatch_frame = if rewritten.is_some() { @@ -1794,34 +1634,44 @@ pub(crate) async fn handle_connection_sharded_inner< let proto_ver = conn.protocol_version; for (meta, target) in reply_futures { let shard_responses = response_pool.future_for(target).await; + + // H1-BARRIER: collect (resp_idx, had_aof_bytes) pairs so + // we can overwrite write responses on fsync failure below. + // aof_bytes is Some for write commands, None for reads. + let mut write_resp_idxs: Vec = Vec::new(); for ((resp_idx, aof_bytes, cmd_name), resp) in meta.into_iter().zip(shard_responses) { - // AOF logging for successful remote writes. - // Owner shard is `target` (NOT ctx.shard_id) — under PerShard - // layout the write must land in the target shard's AOF file - // since that shard owns the mutated data. This was the - // pre-existing routing bug that motivated the per-shard AOF - // RFC (Option B): under TopLevel a single writer absorbed - // every cross-shard append, masking the wrong-owner write. - let mut resp_final = resp; - if let Some(bytes) = aof_bytes { - if !matches!(resp_final, Frame::Error(_)) { - if let Some(ref pool) = ctx.aof_pool { - // Cross-shard: LSN sourced for `target`. - let lsn = aof::AofWriterPool::issue_append_lsn(&ctx.repl_state, target, bytes.len()); - // H1: durable path under appendfsync=always. - if pool - .try_send_append_durable(target, lsn, bytes) - .await - .is_err() - { - resp_final = Frame::Error(Bytes::from_static( - aof::AOF_FSYNC_ERR, - )); - } + // C4-FOLD-FIX: AOF append for cross-shard writes is now done + // inside the SPSC arm (PipelineBatchSlotted / PipelineBatch), + // BEFORE the response slot is filled. Appending here (after + // awaiting the response) defers the append until after + // drain_spsc_shared returns, which means AofFold's + // pending_aof_count undercount it → escape to new incr → + // double-apply on restart. The SPSC arm now owns the AOF + // write; aof_bytes below is used only for the barrier check. + let converted = apply_resp3_conversion(&cmd_name, resp, proto_ver); + if aof_bytes.is_some() && !matches!(converted, Frame::Error(_)) { + write_resp_idxs.push(resp_idx); + } + responses[resp_idx] = converted; + } + + // H1-BARRIER (C4-FOLD-FIX follow-up): under appendfsync=always, + // call fsync_barrier once per target shard AFTER responses are + // collected. The SPSC arm enqueued the Append fire-and-forget; + // the barrier enqueues a zero-length AppendSync into the SAME + // shard channel. Because the writer processes messages in order, + // an acked barrier proves all prior Appends to this shard are on + // durable storage. Under EverySec/No this is a zero-cost noop. + if !write_resp_idxs.is_empty() { + if let Some(ref pool) = ctx.aof_pool { + if pool.fsync_barrier(target).await.is_err() { + for idx in write_resp_idxs { + responses[idx] = Frame::Error( + Bytes::from_static(aof::AOF_FSYNC_ERR), + ); } } } - responses[resp_idx] = apply_resp3_conversion(&cmd_name, resp_final, proto_ver); } } } diff --git a/src/server/conn/handler_sharded/txn.rs b/src/server/conn/handler_sharded/txn.rs index afd1cf4d..33983db2 100644 --- a/src/server/conn/handler_sharded/txn.rs +++ b/src/server/conn/handler_sharded/txn.rs @@ -22,7 +22,7 @@ pub(super) fn try_handle_txn_begin( cmd: &[u8], cmd_args: &[Frame], conn: &mut ConnectionState, - ctx: &ConnectionContext, + _ctx: &ConnectionContext, responses: &mut Vec, ) -> bool { if !is_txn_begin(cmd, cmd_args) { @@ -31,13 +31,9 @@ pub(super) fn try_handle_txn_begin( match txn_begin_validate(conn.in_multi, conn.in_cross_txn()) { Ok(()) => { // Get next txn_id and snapshot_lsn from vector store's transaction manager. - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let active = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| s.vector_store.txn_manager_mut().begin()) - } else { - let mut vector_store = ctx.shard_databases.vector_store(ctx.shard_id); - vector_store.txn_manager_mut().begin() - }; + // Unconditional slice path: ShardSlice is always initialized. + let active = + crate::shard::slice::with_shard(|s| s.vector_store.txn_manager_mut().begin()); conn.active_cross_txn = Some(CrossStoreTxn::new(active.txn_id, active.snapshot_lsn)); responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); } @@ -47,7 +43,7 @@ pub(super) fn try_handle_txn_begin( } /// Handle TXN.COMMIT — returns `true` if the command was consumed. -pub(super) fn try_handle_txn_commit( +pub(super) async fn try_handle_txn_commit( cmd: &[u8], cmd_args: &[Frame], conn: &mut ConnectionState, @@ -65,31 +61,16 @@ pub(super) fn try_handle_txn_commit( // have been excluded from oldest_snapshot, allowing prune_committed to // advance past its LSN. Committing with a stale read set is undefined // behaviour — force the client to restart the transaction. - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - // Closure returns true iff the snapshot was killed (caller emits MOONERR - // and short-circuits the commit path). - let was_killed = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - if s.vector_store.txn_manager().is_killed(txn.txn_id) { - s.vector_store.txn_manager_mut().abort_killed(txn.txn_id); - true - } else { - s.vector_store.txn_manager_mut().commit(txn.txn_id); - false - } - }) - } else { - let mut vector_store = ctx.shard_databases.vector_store(ctx.shard_id); - let killed = if vector_store.txn_manager().is_killed(txn.txn_id) { - vector_store.txn_manager_mut().abort_killed(txn.txn_id); + // Unconditional slice path: returns true iff the snapshot was killed. + let was_killed = crate::shard::slice::with_shard(|s| { + if s.vector_store.txn_manager().is_killed(txn.txn_id) { + s.vector_store.txn_manager_mut().abort_killed(txn.txn_id); true } else { - vector_store.txn_manager_mut().commit(txn.txn_id); + s.vector_store.txn_manager_mut().commit(txn.txn_id); false - }; - drop(vector_store); - killed - }; + } + }); if was_killed { tracing::warn!( txn_id = txn.txn_id, @@ -105,26 +86,14 @@ pub(super) fn try_handle_txn_commit( // Write XactCommit WAL record with committed KV state let txn_id = txn.txn_id; if !txn.kv_undo.is_empty() { - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let payload = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| { - crate::persistence::wal_v3::record::encode_xact_commit_payload( - txn_id, - txn.kv_undo.records(), - db, - ) - }) - } else { - let db_guard = ctx.shard_databases.read_db(ctx.shard_id, conn.selected_db); - let payload = - crate::persistence::wal_v3::record::encode_xact_commit_payload( - txn_id, - txn.kv_undo.records(), - &*db_guard, - ); - drop(db_guard); - payload - }; + // Unconditional slice path: ShardSlice is always initialized. + let payload = crate::shard::slice::with_shard_db(conn.selected_db, |db| { + crate::persistence::wal_v3::record::encode_xact_commit_payload( + txn_id, + txn.kv_undo.records(), + db, + ) + }); let mut wal_buf = Vec::new(); crate::persistence::wal_v3::record::write_wal_v3_record( &mut wal_buf, @@ -136,56 +105,78 @@ pub(super) fn try_handle_txn_commit( .wal_append(ctx.shard_id, bytes::Bytes::from(wal_buf)); } - // Release KV write intents from shard side-table - ctx.shard_databases - .kv_intents(ctx.shard_id) - .release_txn(txn_id); + // Release KV write intents from shard side-table (ShardSlice path). + crate::shard::slice::with_shard(|s| { + s.kv_write_intents.release_txn(txn_id); + }); // Drain deferred HNSW inserts (post-commit hook). // The drain prevents phantom neighbors on abort. // Actual HNSW graph insertion happens during compaction, // not at commit time (point is already in mutable segment). - let drain_count = ctx - .shard_databases - .hnsw_queue(ctx.shard_id) - .drain_for_txn(txn_id) - .count(); + let drain_count = crate::shard::slice::with_shard(|s| { + s.deferred_hnsw_inserts.drain_for_txn(txn_id).count() + }); if drain_count > 0 { tracing::debug!(txn_id, count = drain_count, "Drained deferred HNSW inserts"); } // Materialize MQ intents: enqueue deferred MQ.PUBLISH messages. - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. + // Unconditional slice path (Wave B2): group intents by owning shard. + // Self-shard intents applied locally; foreign groups sent via MqTxnMaterialize. if !txn.mq_intents.is_empty() { - let materialize = |db: &mut crate::storage::db::Database| { - for intent in &txn.mq_intents { - if let Ok(Some(stream)) = db.get_stream_mut(&intent.queue_key) { - if stream.durable { - let msg_id = stream.next_auto_id(); - stream.add(msg_id, intent.fields.clone()); - } - } + let mut self_intents: Vec = Vec::new(); + let mut foreign: std::collections::HashMap< + usize, + Vec, + > = std::collections::HashMap::new(); + for intent in txn.mq_intents.iter().cloned() { + let owner = + crate::shard::dispatch::key_to_shard(&intent.queue_key, ctx.num_shards); + if owner == ctx.shard_id { + self_intents.push(intent); + } else { + foreign.entry(owner).or_default().push(intent); } - }; - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, materialize); - } else { - // Each queue lives on the shard owning its key, which - // may differ from the connection's shard (and per - // intent) — acquire the owner's db per queue. - for intent in &txn.mq_intents { - let owner = crate::shard::dispatch::key_to_shard( - &intent.queue_key, - ctx.num_shards, - ); - let mut db_guard = - ctx.shard_databases.write_db(owner, conn.selected_db); - if let Ok(Some(stream)) = db_guard.get_stream_mut(&intent.queue_key) { - if stream.durable { - let msg_id = stream.next_auto_id(); - stream.add(msg_id, intent.fields.clone()); + } + // Apply self-shard intents synchronously (no borrow across .await). + if !self_intents.is_empty() { + crate::shard::slice::with_shard_db(conn.selected_db, |db| { + for intent in &self_intents { + if let Ok(Some(stream)) = db.get_stream_mut(&intent.queue_key) { + if stream.durable { + let msg_id = stream.next_auto_id(); + stream.add(msg_id, intent.fields.clone()); + } } } + }); + } + // Send MqTxnMaterialize to each foreign shard and await all acks. + for (owner, intents) in foreign { + let (reply_tx, reply_rx) = crate::runtime::channel::oneshot(); + let msg = crate::shard::dispatch::ShardMessage::MqTxnMaterialize { + db_index: conn.selected_db, + intents, + reply_tx, + }; + crate::shard::coordinator::spsc_send( + &ctx.dispatch_tx, + ctx.shard_id, + owner, + msg, + &ctx.spsc_notifiers, + ) + .await; + match reply_rx.recv().await { + Ok(()) => {} + Err(_) => { + tracing::warn!( + "TXN.COMMIT MQ materialize: reply channel closed \ + for shard {}", + owner + ); + } } } } @@ -244,7 +235,7 @@ pub(super) async fn try_handle_txn_abort( pub(super) fn try_handle_temporal_snapshot_at( cmd: &[u8], cmd_args: &[Frame], - ctx: &ConnectionContext, + _ctx: &ConnectionContext, responses: &mut Vec, ) -> bool { if !is_temporal_snapshot_at(cmd) { @@ -253,19 +244,14 @@ pub(super) fn try_handle_temporal_snapshot_at( match validate_snapshot_at(cmd_args) { Ok(()) => { let wall_ms = capture_wall_ms(); - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let lsn = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| s.vector_store.txn_manager().current_lsn()) - } else { - let vector_store = ctx.shard_databases.vector_store(ctx.shard_id); - vector_store.txn_manager().current_lsn() - }; - { - let mut guard = ctx.shard_databases.temporal_registry(ctx.shard_id); - let registry = - guard.get_or_insert_with(|| Box::new(crate::temporal::TemporalRegistry::new())); + // Unconditional slice path: ShardSlice is always initialized. + crate::shard::slice::with_shard(|s| { + let lsn = s.vector_store.txn_manager().current_lsn(); + let registry = s + .temporal_registry + .get_or_insert_with(|| Box::new(crate::temporal::TemporalRegistry::new())); registry.record(wall_ms, lsn); - } + }); responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); } Err(e) => responses.push(e), @@ -319,30 +305,11 @@ pub(super) async fn try_handle_temporal_invalidate( } } let wall_ms = capture_wall_ms(); - // Phase 2f: gate on is_initialized(); both branches are - // semantically identical, the new path uses - // ShardSlice::graph_store directly (no lock). - let (result, wal_records) = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - let gs = &mut s.graph_store; - let r = crate::command::temporal::apply_invalidate( - gs, - entity_id, - is_node, - &graph_name, - wall_ms, - ); - let recs = if r.is_ok() { - gs.drain_wal() - } else { - Vec::new() - }; - (r, recs) - }) - } else { - let mut gs = ctx.shard_databases.graph_store_write(ctx.shard_id); + // Unconditional slice path: ShardSlice is always initialized. + let (result, wal_records) = crate::shard::slice::with_shard(|s| { + let gs = &mut s.graph_store; let r = crate::command::temporal::apply_invalidate( - &mut gs, + gs, entity_id, is_node, &graph_name, @@ -354,7 +321,7 @@ pub(super) async fn try_handle_temporal_invalidate( Vec::new() }; (r, recs) - }; + }); match result { Ok(()) => { for record in wal_records { diff --git a/src/server/conn/handler_sharded/write.rs b/src/server/conn/handler_sharded/write.rs index ff9d6fbc..bf0c06da 100644 --- a/src/server/conn/handler_sharded/write.rs +++ b/src/server/conn/handler_sharded/write.rs @@ -5,8 +5,8 @@ use bytes::Bytes; use crate::command::mq::{ - ERR_MQ_NOT_DURABLE, ERR_MQ_UNKNOWN_SUB, parse_mq_subcommand, validate_mq_ack, - validate_mq_create, validate_mq_dlqlen, validate_mq_pop, validate_mq_publish, validate_mq_push, + ERR_MQ_UNKNOWN_SUB, parse_mq_subcommand, validate_mq_ack, validate_mq_create, + validate_mq_dlqlen, validate_mq_pop, validate_mq_publish, validate_mq_push, validate_mq_trigger, }; use crate::command::transaction::ERR_MULTI_TXN_CONFLICT; @@ -20,7 +20,6 @@ use crate::protocol::Frame; use crate::server::conn::core::{ConnectionContext, ConnectionState}; #[cfg(feature = "graph")] use crate::server::conn::util::extract_bytes; -use crate::storage::stream::StreamId; #[cfg(feature = "graph")] use crate::workspace::strip_workspace_prefix_from_response; use crate::workspace::{WorkspaceId, is_ws_command}; @@ -28,7 +27,7 @@ use crate::workspace::{WorkspaceId, is_ws_command}; use super::execute_transaction_sharded; /// Handle WS.* workspace commands. Returns `true` if consumed. -pub(super) fn try_handle_ws_command( +pub(super) async fn try_handle_ws_command( cmd: &[u8], cmd_args: &[Frame], conn: &mut ConnectionState, @@ -111,10 +110,15 @@ pub(super) fn try_handle_ws_command( ); ctx.shard_databases.wal_append(0, Bytes::from(wal_buf)); // Best-effort cleanup: delete all KV keys with ws - // prefix (WS-03). - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - let prefix = format!("{{{}}}:", ws_id.as_hex()); + // prefix (WS-03). Unconditional slice path. + let prefix = format!("{{{}}}:", ws_id.as_hex()); + // The {wsid} hash tag co-locates every workspace key on ONE shard — + // derive the owner shard from the prefix. + let cleanup_owner = + crate::shard::dispatch::key_to_shard(prefix.as_bytes(), ctx.num_shards); + // Unconditional slice path: ShardSlice is always initialized. + if cleanup_owner == ctx.shard_id { + // Owner is this shard — operate directly on the slice. crate::shard::slice::with_shard_db(0, |db| { let keys_to_delete: Vec> = db .keys() @@ -126,22 +130,35 @@ pub(super) fn try_handle_ws_command( } }); } else { - let prefix = format!("{{{}}}:", ws_id.as_hex()); - // The {wsid} hash tag co-locates every workspace - // key on ONE shard — clean up there, not on the - // connection's shard. - let owner = crate::shard::dispatch::key_to_shard( - prefix.as_bytes(), - ctx.num_shards, - ); - let mut db_guard = ctx.shard_databases.write_db(owner, 0); - let keys_to_delete: Vec> = db_guard - .keys() - .filter(|k| k.as_bytes().starts_with(prefix.as_bytes())) - .map(|k| k.as_bytes().to_vec()) - .collect(); - for key in &keys_to_delete { - db_guard.remove(key); + // Foreign shard: hop via WsDropCleanup message. + let prefix_bytes = Bytes::from(prefix.into_bytes()); + let (reply_tx, reply_rx) = crate::runtime::channel::oneshot(); + let msg = crate::shard::dispatch::ShardMessage::WsDropCleanup { + prefix: prefix_bytes, + reply_tx, + }; + crate::shard::coordinator::spsc_send( + &ctx.dispatch_tx, + ctx.shard_id, + cleanup_owner, + msg, + &ctx.spsc_notifiers, + ) + .await; + match reply_rx.recv().await { + Ok(count) => { + tracing::debug!( + "WS.DROP cleanup: deleted {} keys on shard {}", + count, + cleanup_owner + ); + } + Err(_) => { + tracing::warn!( + "WS.DROP cleanup: reply channel closed for shard {}", + cleanup_owner + ); + } } } responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); @@ -254,10 +271,78 @@ pub(super) fn try_handle_ws_command( true } +/// Build the workspace key-prefix bytes for MQ dispatch payloads. +/// +/// Returns `"{ws_hex}:"` as `Bytes` when the connection is workspace-bound, +/// or `Bytes::new()` otherwise. Mirrors the prefix that `workspace_key()` +/// prepends to queue keys — passed to `MqCommandPayload.key_prefix` so the +/// owner shard can reconstruct effective keys without re-deriving them. +#[inline] +fn mq_key_prefix(workspace_id: Option<&crate::workspace::WorkspaceId>) -> bytes::Bytes { + match workspace_id { + None => bytes::Bytes::new(), + Some(ws_id) => { + let ws_hex = ws_id.as_hex(); + let mut buf = Vec::with_capacity(ws_hex.len() + 3); // '{' + hex + '}' + ':' + buf.push(b'{'); + buf.extend_from_slice(ws_hex.as_bytes()); + buf.push(b'}'); + buf.push(b':'); + bytes::Bytes::from(buf) + } + } +} + +/// Dispatch one MQ.* command to its owning shard via the SPSC hop. +/// +/// If `owner == ctx.shard_id` (this shard owns the queue), executes +/// `execute_mq_on_owner` directly on the current thread (no channel round- +/// trip). Otherwise sends `ShardMessage::MqCommand` and awaits the reply. +/// +/// Mirrors the GraphCommand precedent in `try_handle_graph_command`. +async fn mq_dispatch_to_owner( + frame: &Frame, + key_prefix: bytes::Bytes, + owner: usize, + db_index: usize, + ctx: &ConnectionContext, +) -> Frame { + let command_arc = std::sync::Arc::new(frame.clone()); + if owner == ctx.shard_id { + // Self-shard: execute directly — no channel allocation needed. + crate::shard::mq_exec::execute_mq_on_owner(db_index, key_prefix, command_arc) + } else { + // Foreign shard: send via SPSC and await the oneshot reply. + let (reply_tx, reply_rx) = crate::runtime::channel::oneshot(); + let payload = crate::shard::dispatch::MqCommandPayload { + db_index, + key_prefix, + command: command_arc, + reply_tx, + }; + let msg = crate::shard::dispatch::ShardMessage::MqCommand(Box::new(payload)); + crate::shard::coordinator::spsc_send( + &ctx.dispatch_tx, + ctx.shard_id, + owner, + msg, + &ctx.spsc_notifiers, + ) + .await; + match reply_rx.recv().await { + Ok(f) => f, + Err(_) => Frame::Error(bytes::Bytes::from_static( + b"ERR cross-shard MQ reply channel closed", + )), + } + } +} + /// Handle MQ.* message queue commands. Returns `true` if consumed. -pub(super) fn try_handle_mq_command( +pub(super) async fn try_handle_mq_command( cmd: &[u8], cmd_args: &[Frame], + frame: &Frame, conn: &mut ConnectionState, ctx: &ConnectionContext, responses: &mut Vec, @@ -275,73 +360,23 @@ pub(super) fn try_handle_mq_command( if sub.eq_ignore_ascii_case(b"CREATE") { match validate_mq_create(cmd_args) { - Ok((queue_key, max_delivery_count, _debounce_ms)) => { + Ok((queue_key, _max_delivery_count, _debounce_ms)) => { let effective_key = crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); // A durable queue lives on the shard that owns its key — the // stream, registry entry, trigger entry, and WAL records must // all target `owner`, not the connection's shard, or fresh // connections landing elsewhere (SO_REUSEPORT) can't see the - // queue. Lock-mode ShardDatabases permits direct cross-shard - // access; the ShardSlice branches stay conn-local because - // slice mode is never initialized yet (owner-routing there is - // the shardslice-migration task). + // queue. Owner-routing via MqCommand hop. let owner = crate::shard::dispatch::key_to_shard(&effective_key, ctx.num_shards); - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let create_result: Result<(), Frame> = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| { - match db.get_or_create_stream(&effective_key) { - Ok(stream) => { - stream.durable = true; - stream.max_delivery_count = max_delivery_count; - let group_name = Bytes::from_static(b"__mq_consumers"); - let _ = stream.create_group(group_name, StreamId::ZERO); - Ok(()) - } - Err(e) => Err(e), - } - }) - } else { - let mut db_guard = ctx.shard_databases.write_db(owner, conn.selected_db); - let r = match db_guard.get_or_create_stream(&effective_key) { - Ok(stream) => { - stream.durable = true; - stream.max_delivery_count = max_delivery_count; - let group_name = Bytes::from_static(b"__mq_consumers"); - let _ = stream.create_group(group_name, StreamId::ZERO); - Ok(()) - } - Err(e) => Err(e), - }; - drop(db_guard); - r - }; - if let Err(e) = create_result { - responses.push(e); - return true; - } - - let config = - crate::mq::DurableStreamConfig::new(effective_key.clone(), max_delivery_count); + // Unconditional slice path: owner-route via MqCommand hop (Wave B2). { - let mut guard = ctx.shard_databases.durable_queue_registry(owner); - let reg = guard - .get_or_insert_with(|| Box::new(crate::mq::DurableQueueRegistry::new())); - reg.insert(effective_key.clone(), config); + let key_prefix = mq_key_prefix(conn.workspace_id.as_ref()); + let response = + mq_dispatch_to_owner(frame, key_prefix, owner, conn.selected_db, ctx).await; + responses.push(response); + return true; } - - // Owner's WAL so replay restores the registry on the shard - // that also holds the stream. - let payload = crate::mq::wal::encode_mq_create(&effective_key, max_delivery_count); - let mut wal_buf = Vec::new(); - crate::persistence::wal_v3::record::write_wal_v3_record( - &mut wal_buf, - 0, - crate::persistence::wal_v3::record::WalRecordType::MqCreate, - &payload, - ); - ctx.shard_databases.wal_append(owner, Bytes::from(wal_buf)); - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); } Err(e) => responses.push(e), } @@ -350,83 +385,18 @@ pub(super) fn try_handle_mq_command( if sub.eq_ignore_ascii_case(b"PUSH") { match validate_mq_push(cmd_args) { - Ok((queue_key, fields)) => { + Ok((queue_key, _fields)) => { let effective_key = crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); // Owner-shard targeting — see MQ CREATE above. let owner = crate::shard::dispatch::key_to_shard(&effective_key, ctx.num_shards); - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - // Push outcome: Ok(Some(msg_id)) = pushed; Ok(None) = not durable; Err(e) = stream error. - type PushOutcome = Result, Frame>; - let push_outcome: PushOutcome = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, |db| { - match db.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - if !stream.durable { - Ok(None) - } else { - let msg_id = stream.next_auto_id(); - let msg_id = stream.add(msg_id, fields); - Ok(Some(msg_id)) - } - } - Ok(None) => Ok(None), - Err(e) => Err(e), - } - }) - } else { - let mut db_guard = ctx.shard_databases.write_db(owner, conn.selected_db); - let r = match db_guard.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - if !stream.durable { - Ok(None) - } else { - let msg_id = stream.next_auto_id(); - let msg_id = stream.add(msg_id, fields); - Ok(Some(msg_id)) - } - } - Ok(None) => Ok(None), - Err(e) => Err(e), - }; - drop(db_guard); - r - }; - match push_outcome { - Ok(Some(msg_id)) => { - // Owner's registry: its event-loop tick fires triggers. - { - let mut trig_guard = ctx.shard_databases.trigger_registry(owner); - if let Some(reg) = trig_guard.as_mut() { - let trig_key = if let Some(ws_id) = conn.workspace_id.as_ref() { - let ws_hex = ws_id.as_hex(); - let mut k = - Vec::with_capacity(ws_hex.len() + 1 + queue_key.len()); - k.extend_from_slice(ws_hex.as_bytes()); - k.push(b':'); - k.extend_from_slice(&queue_key); - Bytes::from(k) - } else { - queue_key.clone() - }; - if let Some(trig_entry) = reg.get_mut(&trig_key) { - if trig_entry.pending_fire_ms == 0 { - let fire_at = - ctx.cached_clock.ms() + trig_entry.debounce_ms; - trig_entry.pending_fire_ms = fire_at; - } - } - } - } - responses.push(Frame::BulkString(Bytes::from(format!( - "{}-{}", - msg_id.ms, msg_id.seq - )))); - } - Ok(None) => { - responses.push(Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE))); - } - Err(e) => responses.push(e), + // Unconditional slice path: owner-route via MqCommand hop (Wave B2). + { + let key_prefix = mq_key_prefix(conn.workspace_id.as_ref()); + let response = + mq_dispatch_to_owner(frame, key_prefix, owner, conn.selected_db, ctx).await; + responses.push(response); + return true; } } Err(e) => responses.push(e), @@ -436,117 +406,15 @@ pub(super) fn try_handle_mq_command( if sub.eq_ignore_ascii_case(b"POP") { match validate_mq_pop(cmd_args) { - Ok((queue_key, count)) => { + Ok((queue_key, _count)) => { let effective_key = crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); // Owner-shard targeting — see MQ CREATE above. let owner = crate::shard::dispatch::key_to_shard(&effective_key, ctx.num_shards); - let group_name = Bytes::from_static(b"__mq_consumers"); - let consumer_name = Bytes::from_static(b"__mq_default"); - - // The POP body needs the same Database guard across multiple steps - // (mdc lookup, read_group_new, DLQ stream creation). Refactor into a - // closure that takes &mut Database and produces a Frame to push. - // - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let pop_body = |db: &mut crate::storage::db::Database| -> Frame { - let mdc = match db.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - if !stream.durable { - return Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE)); - } - stream.max_delivery_count - } - Ok(None) => { - return Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE)); - } - Err(e) => return e, - }; - - let request_count = count + (mdc as usize); - let stream = match db.get_stream_mut(&effective_key) { - Ok(Some(s)) => s, - _ => { - return Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE)); - } - }; - let claimed = match stream.read_group_new( - &group_name, - &consumer_name, - Some(request_count), - false, - ) { - Ok(entries) => entries, - Err(_) => { - return Frame::Array(vec![].into()); - } - }; - - let mut results = Vec::new(); - let mut dlq_entries: Vec<(StreamId, Vec<(Bytes, Bytes)>)> = Vec::new(); - let mut dlq_ack_ids: Vec = Vec::new(); - for (id, fields) in &claimed { - let delivery_count = stream - .groups - .get(group_name.as_ref()) - .and_then(|g| g.pel.get(id)) - .map(|pe| pe.delivery_count) - .unwrap_or(1); - if mdc > 0 && delivery_count >= mdc as u64 { - dlq_entries.push((*id, fields.clone())); - dlq_ack_ids.push(*id); - } else if results.len() < count { - results.push((*id, fields.clone())); - } - } - - if !dlq_ack_ids.is_empty() { - let _ = stream.xack(&group_name, &dlq_ack_ids); - } - - if !dlq_entries.is_empty() { - let dlq_key = { - let mut buf = Vec::with_capacity(effective_key.len() + 8); - buf.extend_from_slice(&effective_key); - buf.extend_from_slice(b"::mq:dlq"); - Bytes::from(buf) - }; - if let Ok(dlq_stream) = db.get_or_create_stream(&dlq_key) { - for (_id, fields) in dlq_entries { - let dlq_id = dlq_stream.next_auto_id(); - dlq_stream.add(dlq_id, fields); - } - } - } - - let result_frames: Vec = results - .iter() - .map(|(id, fields)| { - let mut entry_frames = Vec::with_capacity(2); - entry_frames.push(Frame::BulkString(Bytes::from(format!( - "{}-{}", - id.ms, id.seq - )))); - let field_frames: Vec = fields - .iter() - .flat_map(|(f, v)| { - vec![Frame::BulkString(f.clone()), Frame::BulkString(v.clone())] - }) - .collect(); - entry_frames.push(Frame::Array(field_frames.into())); - Frame::Array(entry_frames.into()) - }) - .collect(); - Frame::Array(result_frames.into()) - }; - - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, pop_body) - } else { - let mut db_guard = ctx.shard_databases.write_db(owner, conn.selected_db); - let r = pop_body(&mut *db_guard); - drop(db_guard); - r + // Unconditional slice path: owner-route via MqCommand hop (Wave B2). + let response = { + let key_prefix = mq_key_prefix(conn.workspace_id.as_ref()); + mq_dispatch_to_owner(frame, key_prefix, owner, conn.selected_db, ctx).await }; responses.push(response); } @@ -557,54 +425,18 @@ pub(super) fn try_handle_mq_command( if sub.eq_ignore_ascii_case(b"ACK") { match validate_mq_ack(cmd_args) { - Ok((queue_key, msg_ids)) => { + Ok((queue_key, _msg_ids)) => { let effective_key = crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); - let ids: Vec = msg_ids - .iter() - .map(|(ms, seq)| StreamId { ms: *ms, seq: *seq }) - .collect(); // Owner-shard targeting — see MQ CREATE above. let owner = crate::shard::dispatch::key_to_shard(&effective_key, ctx.num_shards); - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - // Closure returns Some(acked_count) on success, None on any error/miss. - let ack_body = |db: &mut crate::storage::db::Database| -> Option { - match db.get_stream_mut(&effective_key) { - Ok(Some(stream)) => { - let group_name = Bytes::from_static(b"__mq_consumers"); - match stream.xack(&group_name, &ids) { - Ok(acked_count) => Some(acked_count), - Err(_) => None, - } - } - Ok(None) => None, - Err(_) => None, - } - }; - let acked = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, ack_body) - } else { - let mut db_guard = ctx.shard_databases.write_db(owner, conn.selected_db); - let r = ack_body(&mut *db_guard); - drop(db_guard); - r - }; - match acked { - Some(acked_count) => { - for (ms, seq) in &msg_ids { - let payload = crate::mq::wal::encode_mq_ack(&effective_key, *ms, *seq); - let mut wal_buf = Vec::new(); - crate::persistence::wal_v3::record::write_wal_v3_record( - &mut wal_buf, - 0, - crate::persistence::wal_v3::record::WalRecordType::MqAck, - &payload, - ); - ctx.shard_databases.wal_append(owner, Bytes::from(wal_buf)); - } - responses.push(Frame::Integer(acked_count as i64)); - } - None => responses.push(Frame::Integer(0)), + // Unconditional slice path: owner-route via MqCommand hop (Wave B2). + { + let key_prefix = mq_key_prefix(conn.workspace_id.as_ref()); + let response = + mq_dispatch_to_owner(frame, key_prefix, owner, conn.selected_db, ctx).await; + responses.push(response); + return true; } } Err(e) => responses.push(e), @@ -620,28 +452,14 @@ pub(super) fn try_handle_mq_command( // Owner of the QUEUE key, not the dlq_key: POP creates the // DLQ stream in the same db as the queue it drains. let owner = crate::shard::dispatch::key_to_shard(&effective_key, ctx.num_shards); - let dlq_key = { - let mut buf = Vec::with_capacity(effective_key.len() + 8); - buf.extend_from_slice(&effective_key); - buf.extend_from_slice(b"::mq:dlq"); - Bytes::from(buf) - }; - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. - let dlq_body = |db: &mut crate::storage::db::Database| -> i64 { - match db.get_stream_mut(&dlq_key) { - Ok(Some(stream)) => stream.length as i64, - _ => 0i64, - } - }; - let len = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(conn.selected_db, dlq_body) - } else { - let mut db_guard = ctx.shard_databases.write_db(owner, conn.selected_db); - let r = dlq_body(&mut *db_guard); - drop(db_guard); - r - }; - responses.push(Frame::Integer(len)); + // Unconditional slice path: owner-route via MqCommand hop (Wave B2). + { + let key_prefix = mq_key_prefix(conn.workspace_id.as_ref()); + let response = + mq_dispatch_to_owner(frame, key_prefix, owner, conn.selected_db, ctx).await; + responses.push(response); + return true; + } } Err(e) => responses.push(e), } @@ -650,36 +468,20 @@ pub(super) fn try_handle_mq_command( if sub.eq_ignore_ascii_case(b"TRIGGER") { match validate_mq_trigger(cmd_args) { - Ok((queue_key, callback_cmd, debounce_ms)) => { + Ok((queue_key, _callback_cmd, _debounce_ms)) => { let effective_key = crate::workspace::workspace_key(conn.workspace_id.as_ref(), &queue_key); - let trig_key = if let Some(ws_id) = conn.workspace_id.as_ref() { - let ws_hex = ws_id.as_hex(); - let mut k = Vec::with_capacity(ws_hex.len() + 1 + queue_key.len()); - k.extend_from_slice(ws_hex.as_bytes()); - k.push(b':'); - k.extend_from_slice(&queue_key); - Bytes::from(k) - } else { - queue_key.clone() - }; // Owner's registry: its event-loop tick fires triggers // (timers.rs documents the home shard as authoritative). let owner = crate::shard::dispatch::key_to_shard(&effective_key, ctx.num_shards); - let entry = crate::mq::TriggerEntry { - queue_key: effective_key, - callback_cmd, - debounce_ms, - last_fire_ms: 0, - pending_fire_ms: 0, - }; + // Unconditional slice path: owner-route via MqCommand hop (Wave B2). { - let mut guard = ctx.shard_databases.trigger_registry(owner); - let reg = - guard.get_or_insert_with(|| Box::new(crate::mq::TriggerRegistry::new())); - reg.register(trig_key, entry); + let key_prefix = mq_key_prefix(conn.workspace_id.as_ref()); + let response = + mq_dispatch_to_owner(frame, key_prefix, owner, conn.selected_db, ctx).await; + responses.push(response); + return true; } - responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); } Err(e) => responses.push(e), } @@ -847,52 +649,31 @@ pub(super) async fn try_handle_graph_command( let is_write = crate::command::graph::is_graph_write_cmd(cmd) || (cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") && crate::command::graph::is_cypher_write_query(cmd_args)); - // Phase 2f: gate on is_initialized(); new path uses ShardSlice directly. + // Unconditional slice path: ShardSlice is always initialized. let (response, wal_records, cypher_intents, cypher_undo_ops) = - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - if is_write { - let (resp, cypher_intents, undo_ops) = if cmd - .eq_ignore_ascii_case(b"GRAPH.QUERY") - { - crate::command::graph::graph_query_or_write(&mut s.graph_store, cmd_args) - } else { - ( - crate::command::graph::dispatch_graph_write( - &mut s.graph_store, - cmd, - cmd_args, - ), - Vec::new(), - Vec::new(), - ) - }; - let records = s.graph_store.drain_wal(); - (resp, records, cypher_intents, undo_ops) + crate::shard::slice::with_shard(|s| { + if is_write { + let (resp, cypher_intents, undo_ops) = if cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") { + crate::command::graph::graph_query_or_write(&mut s.graph_store, cmd_args) } else { - let resp = - crate::command::graph::dispatch_graph_read(&s.graph_store, cmd, cmd_args); - (resp, Vec::new(), Vec::new(), Vec::new()) - } - }) - } else if is_write { - let mut gs = ctx.shard_databases.graph_store_write(ctx.shard_id); - let (resp, cypher_intents, undo_ops) = if cmd.eq_ignore_ascii_case(b"GRAPH.QUERY") { - crate::command::graph::graph_query_or_write(&mut gs, cmd_args) + ( + crate::command::graph::dispatch_graph_write( + &mut s.graph_store, + cmd, + cmd_args, + ), + Vec::new(), + Vec::new(), + ) + }; + let records = s.graph_store.drain_wal(); + (resp, records, cypher_intents, undo_ops) } else { - ( - crate::command::graph::dispatch_graph_write(&mut gs, cmd, cmd_args), - Vec::new(), - Vec::new(), - ) - }; - let records = gs.drain_wal(); - (resp, records, cypher_intents, undo_ops) - } else { - let gs = ctx.shard_databases.graph_store_read(ctx.shard_id); - let resp = crate::command::graph::dispatch_graph_read(&gs, cmd, cmd_args); - (resp, Vec::new(), Vec::new(), Vec::new()) - }; + let resp = + crate::command::graph::dispatch_graph_read(&s.graph_store, cmd, cmd_args); + (resp, Vec::new(), Vec::new(), Vec::new()) + } + }); // Phase 166: record graph intent for TXN rollback. if let Some(txn) = conn.active_cross_txn.as_mut() { let is_node = cmd.eq_ignore_ascii_case(b"GRAPH.ADDNODE"); diff --git a/src/server/conn/handler_single.rs b/src/server/conn/handler_single.rs index d1375a49..bbd604b5 100644 --- a/src/server/conn/handler_single.rs +++ b/src/server/conn/handler_single.rs @@ -1511,7 +1511,7 @@ pub async fn handle_connection( // TEMP-04: single-shard handler has no TemporalRegistry and no cross-store TXN. // Unified helper with shard_databases=None returns ERR on AS_OF (correct per // Plan 165-01 contract); non-AS_OF continues to return latest (as_of_lsn=0). - match resolve_ft_search_as_of_lsn(cmd_args, None, 0, None) { + match resolve_ft_search_as_of_lsn(cmd_args, None, None) { Err(err_frame) => err_frame, Ok(as_of_lsn) => { let has_session = cmd_args.iter().any(|a| { @@ -1733,7 +1733,7 @@ pub async fn handle_connection( // TEMP-04: single-shard handler has no TemporalRegistry and no cross-store TXN. // Unified helper with shard_databases=None returns ERR on AS_OF (correct per // Plan 165-01 contract); non-AS_OF continues to return latest (as_of_lsn=0). - match resolve_ft_search_as_of_lsn(d_args, None, 0, None) { + match resolve_ft_search_as_of_lsn(d_args, None, None) { Err(err_frame) => err_frame, Ok(as_of_lsn) => { let has_session = d_args.iter().any(|a| { @@ -2017,7 +2017,7 @@ pub async fn handle_connection( // Write run: guard is already write-locked. // TEMP-04: single-shard handler has no registry and no TXN; helper returns // ERR on AS_OF and Ok(0) otherwise (Plan 165-01 contract). - match resolve_ft_search_as_of_lsn(d_args, None, 0, None) { + match resolve_ft_search_as_of_lsn(d_args, None, None) { Err(err_frame) => err_frame, Ok(as_of_lsn) => crate::command::vector_search::ft_search(&mut *store, d_args, Some(&mut *guard), Some(&*ts_m3), as_of_lsn), } diff --git a/src/server/conn/shared.rs b/src/server/conn/shared.rs index 94163c58..5b221580 100644 --- a/src/server/conn/shared.rs +++ b/src/server/conn/shared.rs @@ -50,21 +50,23 @@ pub(crate) type SharedDatabases = Arc>>; pub(crate) fn resolve_ft_search_as_of_lsn( cmd_args: &[Frame], shard_databases: Option<&ShardDatabases>, - shard_id: usize, active_cross_txn: Option<&CrossStoreTxn>, ) -> Result { use crate::command::vector_search::ft_search::parse::parse_as_of_clause; const ERR_MSG: &[u8] = b"ERR no temporal snapshot registered for the given AS_OF timestamp; call TEMPORAL.SNAPSHOT_AT first"; if let Some(wall_ms) = parse_as_of_clause(cmd_args) { - let Some(shard_dbs) = shard_databases else { + // `shard_databases` is `None` for the single-shard tokio handler, which + // has no ShardDatabases in scope. Presence/absence is the guard; the + // actual registry is accessed through the thread-local ShardSlice below. + if shard_databases.is_none() { return Err(Frame::Error(Bytes::from_static(ERR_MSG))); - }; - let guard = shard_dbs.temporal_registry(shard_id); - return guard - .as_ref() - .and_then(|r| r.lsn_at(wall_ms)) - .ok_or_else(|| Frame::Error(Bytes::from_static(ERR_MSG))); + } + // Read temporal registry via thread-local slice (no lock needed on shard path). + return crate::shard::slice::with_shard(|s| { + s.temporal_registry.as_ref().and_then(|r| r.lsn_at(wall_ms)) + }) + .ok_or_else(|| Frame::Error(Bytes::from_static(ERR_MSG))); } Ok(active_cross_txn.map(|t| t.snapshot_lsn).unwrap_or(0)) } @@ -189,14 +191,12 @@ pub(crate) fn execute_transaction( /// keys only. Cross-shard transactions require distributed coordination (future work). pub(crate) fn execute_transaction_sharded( shard_databases: &std::sync::Arc, - shard_id: usize, + _shard_id: usize, command_queue: &[Frame], selected_db: usize, cached_clock: &CachedClock, ) -> Frame { - let mut guard = shard_databases.write_db(shard_id, selected_db); let db_count = shard_databases.db_count(); - guard.refresh_now_from_cache(cached_clock); let mut results = Vec::with_capacity(command_queue.len()); let mut selected = selected_db; @@ -212,7 +212,10 @@ pub(crate) fn execute_transaction_sharded( } }; - let result = dispatch(&mut guard, cmd, cmd_args, &mut selected, db_count); + let result = crate::shard::slice::with_shard_db(selected, |db| { + db.refresh_now_from_cache(cached_clock); + dispatch(db, cmd, cmd_args, &mut selected, db_count) + }); let response = match result { DispatchResult::Response(f) => f, DispatchResult::Quit(f) => f, @@ -221,13 +224,16 @@ pub(crate) fn execute_transaction_sharded( // Auto-index: if HSET succeeded, check for vector index match if cmd.eq_ignore_ascii_case(b"HSET") && !matches!(response, Frame::Error(_)) { if let Some(Frame::BulkString(key_bytes)) = cmd_args.first() { - let mut vs = shard_databases.vector_store(shard_id); - let mut ts = shard_databases.text_store(shard_id); - // Plan 166-01 return value is not consumed here — this is a - // non-txn-aware batch write path. Plan 166-02 wires txn paths. - let _ = crate::shard::spsc_handler::auto_index_hset_public( - &mut vs, &mut *ts, key_bytes, cmd_args, - ); + crate::shard::slice::with_shard(|s| { + // Plan 166-01 return value is not consumed here — this is a + // non-txn-aware batch write path. Plan 166-02 wires txn paths. + let _ = crate::shard::spsc_handler::auto_index_hset_public( + &mut s.vector_store, + &mut s.text_store, + key_bytes, + cmd_args, + ); + }); } } @@ -416,13 +422,18 @@ mod as_of_tests { /// Build a 1-shard / 1-db `ShardDatabases` with a registered binding /// `wall_ms=1_000 -> lsn=42` so tests can exercise the registry path. + /// + /// Also initialises the thread-local `ShardSlice` (via `reset_test_shard`) + /// with the temporal registry pre-populated, so `with_shard` queries in + /// `resolve_ft_search_as_of_lsn` find the binding. fn build_fixture() -> Arc { - let dbs = ShardDatabases::new(vec![vec![Database::new()]]); - { - let mut guard = dbs.temporal_registry(0); - let reg = guard.get_or_insert_with(|| Box::new(TemporalRegistry::new())); - reg.record(1_000, 42); - } + let (dbs, mut inits) = ShardDatabases::new(vec![vec![Database::new()]]); + let mut init = inits.remove(0); + // Pre-populate temporal registry on the ShardSliceInit before wiring it. + let mut reg = Box::new(TemporalRegistry::new()); + reg.record(1_000, 42); + init.temporal_registry = Some(reg); + crate::shard::slice::reset_test_shard(crate::shard::slice::ShardSlice::new(init)); dbs } @@ -464,7 +475,7 @@ mod as_of_tests { fn resolve_ft_search_as_of_lsn_default_returns_zero() { let fixture = build_fixture(); let args = ft_search_args(None); - let got = resolve_ft_search_as_of_lsn(&args, Some(&fixture), 0, None); + let got = resolve_ft_search_as_of_lsn(&args, Some(&fixture), None); assert_eq!(got, Ok(0)); } @@ -473,7 +484,7 @@ mod as_of_tests { let fixture = build_fixture(); let args = ft_search_args(None); let txn = CrossStoreTxn::new(1, 99); - let got = resolve_ft_search_as_of_lsn(&args, Some(&fixture), 0, Some(&txn)); + let got = resolve_ft_search_as_of_lsn(&args, Some(&fixture), Some(&txn)); assert_eq!(got, Ok(99)); } @@ -482,7 +493,7 @@ mod as_of_tests { let fixture = build_fixture(); let args = ft_search_args(Some(1_000)); let txn = CrossStoreTxn::new(1, 99); - let got = resolve_ft_search_as_of_lsn(&args, Some(&fixture), 0, Some(&txn)); + let got = resolve_ft_search_as_of_lsn(&args, Some(&fixture), Some(&txn)); // Registry binding at wall_ms=1_000 is lsn=42, NOT txn.snapshot_lsn=99. assert_eq!(got, Ok(42)); } @@ -492,7 +503,7 @@ mod as_of_tests { let fixture = build_fixture(); // wall_ms=500 precedes the only registered binding (1_000 -> 42). let args = ft_search_args(Some(500)); - let got = resolve_ft_search_as_of_lsn(&args, Some(&fixture), 0, None); + let got = resolve_ft_search_as_of_lsn(&args, Some(&fixture), None); match got { Err(Frame::Error(msg)) => assert_eq!(msg.as_ref(), ERR_BYTES), other => panic!("expected Err(Frame::Error(ERR_BYTES)), got {other:?}"), @@ -504,7 +515,7 @@ mod as_of_tests { // handler_single.rs has no ShardDatabases in scope -> Option::None. // AS_OF cannot be resolved without a registry; surface the same ERR. let args = ft_search_args(Some(1_000)); - let got = resolve_ft_search_as_of_lsn(&args, None, 0, None); + let got = resolve_ft_search_as_of_lsn(&args, None, None); match got { Err(Frame::Error(msg)) => assert_eq!(msg.as_ref(), ERR_BYTES), other => panic!("expected Err(Frame::Error(ERR_BYTES)), got {other:?}"), diff --git a/src/server/conn/tests.rs b/src/server/conn/tests.rs index 9bcd6460..f4d7a163 100644 --- a/src/server/conn/tests.rs +++ b/src/server/conn/tests.rs @@ -9,9 +9,18 @@ use crate::storage::Database; use crate::storage::entry::Entry; use bytes::{Bytes, BytesMut}; -/// Helper: create a single-shard, single-database ShardDatabases for testing. +/// Helper: create a single-shard, single-database ShardDatabases for testing +/// and initialize the thread-local ShardSlice so dispatch paths work. +/// +/// Each test must call this before any dispatch or write_db calls. +/// Uses `reset_test_shard` (test-only) so multiple tests on the same OS thread +/// each get a fresh, empty shard state. fn make_dbs() -> std::sync::Arc { - crate::shard::shared_databases::ShardDatabases::new(vec![vec![Database::new()]]) + let (shared, mut inits) = + crate::shard::shared_databases::ShardDatabases::new(vec![vec![Database::new()]]); + let init = inits.remove(0); + crate::shard::slice::reset_test_shard(crate::shard::slice::ShardSlice::new(init)); + shared } /// Helper: default runtime config for inline dispatch tests. @@ -22,13 +31,12 @@ fn make_rt_config() -> parking_lot::RwLock { #[test] fn test_inline_get_hit() { let dbs = make_dbs(); - { - let mut guard = dbs.write_db(0, 0); - guard.set( + crate::shard::slice::with_shard_db(0, |db| { + db.set( Bytes::from_static(b"foo"), Entry::new_string(Bytes::from_static(b"bar")), ); - } + }); let mut read_buf = BytesMut::from(&b"*2\r\n$3\r\nGET\r\n$3\r\nfoo\r\n"[..]); let mut write_buf = BytesMut::new(); let aof_pool: Option> = None; @@ -135,15 +143,16 @@ fn test_inline_set_executes_when_writes_enabled() { assert_eq!(&write_buf[..], b"+OK\r\n"); // Verify the key was actually set - let guard = dbs.read_db(0, 0); - // Test assertion: SET was just issued with a live TTL, so the key must - // exist and hold a string-typed value; an absent or wrong-typed value - // would indicate the inline SET path regressed. - #[allow(clippy::expect_used, clippy::unwrap_used)] - let entry = guard.get_if_alive(b"foo", 0).expect("key should exist"); - #[allow(clippy::unwrap_used)] - let value_bytes = entry.value.as_bytes().unwrap(); - assert_eq!(value_bytes, b"bar"); + crate::shard::slice::with_shard_db(0, |db| { + // Test assertion: SET was just issued with a live TTL, so the key must + // exist and hold a string-typed value; an absent or wrong-typed value + // would indicate the inline SET path regressed. + #[allow(clippy::expect_used, clippy::unwrap_used)] + let entry = db.get_if_alive(b"foo", 0).expect("key should exist"); + #[allow(clippy::unwrap_used)] + let value_bytes = entry.value.as_bytes().unwrap(); + assert_eq!(value_bytes, b"bar"); + }); } #[test] @@ -205,13 +214,12 @@ fn test_inline_fallthrough() { #[test] fn test_inline_mixed_batch() { let dbs = make_dbs(); - { - let mut guard = dbs.write_db(0, 0); - guard.set( + crate::shard::slice::with_shard_db(0, |db| { + db.set( Bytes::from_static(b"foo"), Entry::new_string(Bytes::from_static(b"bar")), ); - } + }); // GET foo followed by PING let mut read_buf = BytesMut::new(); read_buf.extend_from_slice(b"*2\r\n$3\r\nGET\r\n$3\r\nfoo\r\n"); @@ -242,13 +250,12 @@ fn test_inline_mixed_batch() { #[test] fn test_inline_case_insensitive() { let dbs = make_dbs(); - { - let mut guard = dbs.write_db(0, 0); - guard.set( + crate::shard::slice::with_shard_db(0, |db| { + db.set( Bytes::from_static(b"foo"), Entry::new_string(Bytes::from_static(b"baz")), ); - } + }); let mut read_buf = BytesMut::from(&b"*2\r\n$3\r\nget\r\n$3\r\nfoo\r\n"[..]); let mut write_buf = BytesMut::new(); let aof_pool: Option> = None; @@ -339,17 +346,16 @@ fn test_inline_set_with_aof_falls_through_when_writes_disabled() { #[test] fn test_inline_multiple_gets() { let dbs = make_dbs(); - { - let mut guard = dbs.write_db(0, 0); - guard.set( + crate::shard::slice::with_shard_db(0, |db| { + db.set( Bytes::from_static(b"a"), Entry::new_string(Bytes::from_static(b"1")), ); - guard.set( + db.set( Bytes::from_static(b"b"), Entry::new_string(Bytes::from_static(b"2")), ); - } + }); let mut read_buf = BytesMut::new(); read_buf.extend_from_slice(b"*2\r\n$3\r\nGET\r\n$1\r\na\r\n"); read_buf.extend_from_slice(b"*2\r\n$3\r\nGET\r\n$1\r\nb\r\n"); diff --git a/src/server/embedded.rs b/src/server/embedded.rs index 2cb19329..e7bb07ee 100644 --- a/src/server/embedded.rs +++ b/src/server/embedded.rs @@ -150,7 +150,7 @@ pub async fn run_embedded( .spawn(move || { RuntimeFactoryImpl::block_on_local( "embedded-moon-aof".to_string(), - aof::aof_writer_task(rx, aof_file_path, fsync, aof_token), + aof::aof_writer_task(rx, aof_file_path, fsync, aof_token, None), ); }) .context("embedded moon: failed to spawn AOF writer thread")?; @@ -275,20 +275,26 @@ pub async fn run_embedded( .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = ShardDatabases::new(all_dbs); // Auxiliary WAL replay paths (no-op when files do not exist). + // Graph + temporal + MQ replays operate on ShardSliceInit (mutate pre-shard state). + // Workspace replay operates on Arc (shared WorkspaceRegistry). #[cfg(feature = "graph")] if let Some(ref dir) = persistence_dir { let dir_path = std::path::Path::new(dir); - shard_databases.recover_graph_stores(dir_path); - shard_databases.replay_graph_wal(dir_path); + crate::shard::shared_databases::recover_graph_stores(&mut slice_inits, dir_path); + crate::shard::shared_databases::replay_graph_wal( + &mut slice_inits, + dir_path, + config.databases, + ); } if let Some(ref dir) = persistence_dir { let dir_path = std::path::Path::new(dir); - shard_databases.replay_temporal_wal(dir_path); - shard_databases.replay_workspace_wal(dir_path); - shard_databases.replay_mq_wal(dir_path); + crate::shard::shared_databases::replay_temporal_wal(&mut slice_inits, dir_path); + crate::shard::shared_databases::replay_workspace_wal(&shard_databases, dir_path); + crate::shard::shared_databases::replay_mq_wal(&mut slice_inits, dir_path); } // Readiness flag — `/readyz` is gated on this; harmless without admin port. @@ -317,6 +323,8 @@ pub async fn run_embedded( let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + // C1: hand the shard its ShardSlice initializer — consumed by init_shard inside run(). + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("embedded-moon-shard-{}", id)) @@ -349,6 +357,7 @@ pub async fn run_embedded( shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ) .await; }, diff --git a/src/server/listener.rs b/src/server/listener.rs index fa325e53..c2a9ea33 100644 --- a/src/server/listener.rs +++ b/src/server/listener.rs @@ -128,7 +128,13 @@ pub async fn run_with_shutdown( let aof_token = token.child_token(); let fsync = FsyncPolicy::from_str(&config.appendfsync); let aof_file_path = PathBuf::from(&config.dir).join(&config.appendfilename); - tokio::spawn(aof::aof_writer_task(rx, aof_file_path, fsync, aof_token)); + tokio::spawn(aof::aof_writer_task( + rx, + aof_file_path, + fsync, + aof_token, + None, + )); info!("AOF enabled with fsync policy: {:?}", fsync); Some(AofWriterPool::top_level_with_policy( tx, diff --git a/src/shard/coordinator.rs b/src/shard/coordinator.rs index eef3e8fb..2c957719 100644 --- a/src/shard/coordinator.rs +++ b/src/shard/coordinator.rs @@ -120,7 +120,6 @@ pub async fn coordinate_multi_key( /// (identical semantics to a remote MultiExecute leg, minus the hop). fn run_local( shard_databases: &Arc, - my_shard: usize, db_index: usize, cached_clock: &CachedClock, cmd: &[u8], @@ -128,19 +127,13 @@ fn run_local( ) -> Frame { let db_count = shard_databases.db_count(); let mut selected = db_index; - let mut run = |db: &mut crate::storage::Database| { + let run = |db: &mut crate::storage::Database| { db.refresh_now_from_cache(cached_clock); match cmd_dispatch(db, cmd, args, &mut selected, db_count) { DispatchResult::Response(f) | DispatchResult::Quit(f) => f, } }; - // Phase 2c: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(db_index, run) - } else { - let mut guard = shard_databases.write_db(my_shard, db_index); - run(&mut guard) - } + crate::shard::slice::with_shard_db(db_index, run) } /// Send one full command to a REMOTE shard and await its reply. @@ -187,14 +180,7 @@ async fn run_on_owner( Some((Frame::BulkString(c), rest)) => (c.clone(), rest), _ => return Frame::Error(Bytes::from_static(b"ERR invalid command format")), }; - run_local( - shard_databases, - my_shard, - db_index, - cached_clock, - &cmd, - args, - ) + run_local(shard_databases, db_index, cached_clock, &cmd, args) } else { let command = Frame::Array(command_parts.to_vec().into()); run_remote( @@ -240,14 +226,7 @@ async fn coordinate_bitop( // Single-shard server: straight to local dispatch — zero coordinator // overhead (no key vec, no owner hashing) on the 1-shard hot path. if num_shards == 1 { - return run_local( - shard_databases, - my_shard, - db_index, - cached_clock, - b"BITOP", - args, - ); + return run_local(shard_databases, db_index, cached_clock, b"BITOP", args); } if args.len() < 3 { return Frame::Error(Bytes::from_static( @@ -317,7 +296,6 @@ async fn coordinate_bitop( for (idx, key) in indexed { let reply = run_local( shard_databases, - my_shard, db_index, cached_clock, b"GET", @@ -440,14 +418,7 @@ async fn coordinate_copy( // Single-shard server: straight to local dispatch — zero coordinator // overhead on the 1-shard hot path. if num_shards == 1 { - return run_local( - shard_databases, - my_shard, - db_index, - cached_clock, - b"COPY", - args, - ); + return run_local(shard_databases, db_index, cached_clock, b"COPY", args); } if args.len() < 2 { return Frame::Error(Bytes::from_static( @@ -640,7 +611,7 @@ async fn coordinate_mget( my_shard: usize, num_shards: usize, db_index: usize, - shard_databases: &Arc, + _shard_databases: &Arc, dispatch_tx: &Rc>>>, spsc_notifiers: &[Arc], cached_clock: &CachedClock, @@ -663,16 +634,10 @@ async fn coordinate_mget( // Fast path: all keys on local shard -- use mget directly if groups.len() == 1 && groups.contains_key(&my_shard) { - // Phase 2c: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - return crate::shard::slice::with_shard_db(db_index, |db| { - db.refresh_now_from_cache(cached_clock); - crate::command::string::mget(db, args) - }); - } - let mut guard = shard_databases.write_db(my_shard, db_index); - guard.refresh_now_from_cache(cached_clock); - return crate::command::string::mget(&mut guard, args); + return crate::shard::slice::with_shard_db(db_index, |db| { + db.refresh_now_from_cache(cached_clock); + crate::command::string::mget(db, args) + }); } let total = args.len(); @@ -685,27 +650,10 @@ async fn coordinate_mget( if *shard_id == my_shard { // Local execution: GET each key directly - // Phase 2c: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(db_index, |db| { - db.refresh_now_from_cache(cached_clock); - for (orig_idx, key) in indexed_keys { - let entry = db.get(key); - let frame = match entry { - Some(e) => match e.value.as_bytes() { - Some(v) => Frame::BulkString(Bytes::copy_from_slice(v)), - None => Frame::Null, - }, - None => Frame::Null, - }; - results[*orig_idx] = Some(frame); - } - }); - } else { - let mut guard = shard_databases.write_db(my_shard, db_index); - guard.refresh_now_from_cache(cached_clock); + crate::shard::slice::with_shard_db(db_index, |db| { + db.refresh_now_from_cache(cached_clock); for (orig_idx, key) in indexed_keys { - let entry = guard.get(key); + let entry = db.get(key); let frame = match entry { Some(e) => match e.value.as_bytes() { Some(v) => Frame::BulkString(Bytes::copy_from_slice(v)), @@ -715,7 +663,7 @@ async fn coordinate_mget( }; results[*orig_idx] = Some(frame); } - } + }); } else { // Remote dispatch: batch of GET commands via MultiExecuteSlotted let (reply_tx, reply_rx) = channel::oneshot(); @@ -776,7 +724,7 @@ async fn coordinate_mset( my_shard: usize, num_shards: usize, db_index: usize, - shard_databases: &Arc, + _shard_databases: &Arc, dispatch_tx: &Rc>>>, spsc_notifiers: &[Arc], cached_clock: &CachedClock, @@ -813,37 +761,22 @@ async fn coordinate_mset( // Fast path: all keys on local shard if groups.len() == 1 && groups.contains_key(&my_shard) { - // Phase 2c: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - return crate::shard::slice::with_shard_db(db_index, |db| { - db.refresh_now_from_cache(cached_clock); - crate::command::string::mset(db, args) - }); - } - let mut guard = shard_databases.write_db(my_shard, db_index); - guard.refresh_now_from_cache(cached_clock); - return crate::command::string::mset(&mut guard, args); + return crate::shard::slice::with_shard_db(db_index, |db| { + db.refresh_now_from_cache(cached_clock); + crate::command::string::mset(db, args) + }); } let mut pending_shards: Vec>> = Vec::new(); for (shard_id, kv_pairs) in &groups { if *shard_id == my_shard { - // Phase 2c: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(db_index, |db| { - db.refresh_now_from_cache(cached_clock); - for (key, value) in kv_pairs { - db.set_string(key.clone(), value.clone()); - } - }); - } else { - let mut guard = shard_databases.write_db(my_shard, db_index); - guard.refresh_now_from_cache(cached_clock); + crate::shard::slice::with_shard_db(db_index, |db| { + db.refresh_now_from_cache(cached_clock); for (key, value) in kv_pairs { - guard.set_string(key.clone(), value.clone()); + db.set_string(key.clone(), value.clone()); } - } + }); } else { let (reply_tx, reply_rx) = channel::oneshot(); let commands: Vec<(Bytes, Frame)> = kv_pairs @@ -907,17 +840,10 @@ async fn coordinate_multi_del_or_exists( // Fast path: all keys on local shard if groups.len() == 1 && groups.contains_key(&my_shard) { let mut selected = db_index; - // Phase 2c: gate on is_initialized(); new path uses ShardSlice directly. - let result = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(db_index, |db| { - db.refresh_now_from_cache(cached_clock); - cmd_dispatch(db, cmd, args, &mut selected, db_count) - }) - } else { - let mut guard = shard_databases.write_db(my_shard, db_index); - guard.refresh_now_from_cache(cached_clock); - cmd_dispatch(&mut guard, cmd, args, &mut selected, db_count) - }; + let result = crate::shard::slice::with_shard_db(db_index, |db| { + db.refresh_now_from_cache(cached_clock); + cmd_dispatch(db, cmd, args, &mut selected, db_count) + }); return match result { DispatchResult::Response(f) => f, DispatchResult::Quit(f) => f, @@ -930,17 +856,10 @@ async fn coordinate_multi_del_or_exists( for (shard_id, key_args) in &groups { if *shard_id == my_shard { let mut selected = db_index; - // Phase 2c: gate on is_initialized(); new path uses ShardSlice directly. - let result = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(db_index, |db| { - db.refresh_now_from_cache(cached_clock); - cmd_dispatch(db, cmd, key_args, &mut selected, db_count) - }) - } else { - let mut guard = shard_databases.write_db(my_shard, db_index); - guard.refresh_now_from_cache(cached_clock); - cmd_dispatch(&mut guard, cmd, key_args, &mut selected, db_count) - }; + let result = crate::shard::slice::with_shard_db(db_index, |db| { + db.refresh_now_from_cache(cached_clock); + cmd_dispatch(db, cmd, key_args, &mut selected, db_count) + }); if let DispatchResult::Response(Frame::Integer(n)) = result { total_count += n; } @@ -1016,17 +935,10 @@ pub async fn coordinate_keys( { let db_count = shard_databases.db_count(); let mut selected = db_index; - // Phase 2c: gate on is_initialized(); new path uses ShardSlice directly. - let result = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(db_index, |db| { - db.refresh_now_from_cache(cached_clock); - cmd_dispatch(db, b"KEYS", args, &mut selected, db_count) - }) - } else { - let mut guard = shard_databases.write_db(my_shard, db_index); - guard.refresh_now_from_cache(cached_clock); - cmd_dispatch(&mut guard, b"KEYS", args, &mut selected, db_count) - }; + let result = crate::shard::slice::with_shard_db(db_index, |db| { + db.refresh_now_from_cache(cached_clock); + cmd_dispatch(db, b"KEYS", args, &mut selected, db_count) + }); if let DispatchResult::Response(Frame::Array(keys)) = result { all_keys.extend(keys); } @@ -1120,17 +1032,10 @@ pub async fn coordinate_scan( let scan_result = if target_shard_id == my_shard { let db_count = shard_databases.db_count(); let mut selected = db_index; - // Phase 2c: gate on is_initialized(); new path uses ShardSlice directly. - let result = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(db_index, |db| { - db.refresh_now_from_cache(cached_clock); - cmd_dispatch(db, b"SCAN", &scan_args, &mut selected, db_count) - }) - } else { - let mut guard = shard_databases.write_db(my_shard, db_index); - guard.refresh_now_from_cache(cached_clock); - cmd_dispatch(&mut guard, b"SCAN", &scan_args, &mut selected, db_count) - }; + let result = crate::shard::slice::with_shard_db(db_index, |db| { + db.refresh_now_from_cache(cached_clock); + cmd_dispatch(db, b"SCAN", &scan_args, &mut selected, db_count) + }); match result { DispatchResult::Response(f) => f, DispatchResult::Quit(f) => f, @@ -1202,7 +1107,7 @@ pub async fn coordinate_dbsize( my_shard: usize, num_shards: usize, db_index: usize, - shard_databases: &Arc, + _shard_databases: &Arc, dispatch_tx: &Rc>>>, spsc_notifiers: &[Arc], _response_pool: &(), // placeholder — coordinator uses oneshot internally @@ -1212,13 +1117,7 @@ pub async fn coordinate_dbsize( // Local shard { - // Phase 2c: gate on is_initialized(); new path uses ShardSlice directly. - let local_len = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(db_index, |db| db.len()) as i64 - } else { - let guard = shard_databases.read_db(my_shard, db_index); - guard.len() as i64 - }; + let local_len = crate::shard::slice::with_shard_db(db_index, |db| db.len()) as i64; total += local_len; } @@ -1262,7 +1161,7 @@ pub async fn coordinate_hotkeys( my_shard: usize, num_shards: usize, db_index: usize, - shard_databases: &Arc, + _shard_databases: &Arc, dispatch_tx: &Rc>>>, spsc_notifiers: &[Arc], _response_pool: &(), // placeholder — coordinator uses oneshot internally @@ -1271,12 +1170,7 @@ pub async fn coordinate_hotkeys( // Local shard: read the sketch directly. { - let local = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(db_index, |db| db.hot_keys().top(count)) - } else { - let guard = shard_databases.read_db(my_shard, db_index); - guard.hot_keys().top(count) - }; + let local = crate::shard::slice::with_shard_db(db_index, |db| db.hot_keys().top(count)); merged.extend(local.into_iter().map(|(k, c)| (k, c as i64))); } @@ -1428,29 +1322,14 @@ pub async fn scatter_vector_search_remote( dispatch_tx: &Rc>>>, spsc_notifiers: &[Arc], ) -> Frame { + let _ = shard_databases; // E2 removes // LOCAL: direct vector store access (avoids SPSC self-send). // Phase 171 SCAT-01: honor AS_OF on the coordinator's own shard by // routing through `search_local_filtered` with the resolved LSN rather // than the AS_OF-unaware `search_local` helper. - // Phase 2c: gate on is_initialized(); new path uses ShardSlice directly. - let local_result = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::vector_search::search_local_filtered( - &mut s.vector_store, - &index_name, - &query_blob, - k, - None, - 0, - usize::MAX, - None, - as_of_lsn, - ) - }) - } else { - let mut vs = shard_databases.vector_store(my_shard); + let local_result = crate::shard::slice::with_shard(|s| { crate::command::vector_search::search_local_filtered( - &mut vs, + &mut s.vector_store, &index_name, &query_blob, k, @@ -1460,7 +1339,7 @@ pub async fn scatter_vector_search_remote( None, as_of_lsn, ) - }; + }); // REMOTE: SPSC to all other shards let mut receivers = Vec::with_capacity(num_shards.saturating_sub(1)); @@ -1513,6 +1392,7 @@ pub async fn broadcast_vector_command( dispatch_tx: &Rc>>>, spsc_notifiers: &[Arc], ) -> Frame { + let _ = shard_databases; // E2 removes // REMOTE FIRST: send to all other shards via SPSC before local mutation. // This ensures we detect remote failures before committing locally, // avoiding partial index metadata across the cluster. @@ -1552,62 +1432,36 @@ pub async fn broadcast_vector_command( _ => false, }; - // Phase 2c: gate on is_initialized(); new path uses ShardSlice with - // disjoint borrows of vector_store / text_store / graph_store / databases[0]. - let local_result = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - // Split borrows so rustc sees `&mut s.vector_store`, - // `&mut s.text_store`, `&s.graph_store`, and optional - // `&mut s.databases[0]` as four disjoint fields. - let (db_slice, vs, ts); + // Split borrows so rustc sees `&mut s.vector_store`, + // `&mut s.text_store`, `&s.graph_store`, and optional + // `&mut s.databases[0]` as four disjoint fields. + let local_result = crate::shard::slice::with_shard(|s| { + let (db_slice, vs, ts); + #[cfg(feature = "graph")] + let graph_ref; + { + vs = &mut s.vector_store; + ts = &mut s.text_store; #[cfg(feature = "graph")] - let graph_ref; - // Reborrow each field through `&mut *s` so the compiler - // tracks them independently. { - vs = &mut s.vector_store; - ts = &mut s.text_store; - #[cfg(feature = "graph")] - { - graph_ref = &s.graph_store; - } - db_slice = &mut s.databases; + graph_ref = &s.graph_store; } - let db_opt = if is_dropindex { - db_slice.get_mut(0) - } else { - None - }; - crate::shard::spsc_handler::dispatch_vector_command( - vs, - ts, - #[cfg(feature = "graph")] - Some(graph_ref), - &command, - db_opt, - ) - }) - } else { - let mut vs = shard_databases.vector_store(my_shard); - let mut ts = shard_databases.text_store(my_shard); - #[cfg(feature = "graph")] - let graph_guard = shard_databases.graph_store_read(my_shard); - let mut db_guard; + db_slice = &mut s.databases; + } let db_opt = if is_dropindex { - db_guard = shard_databases.write_db(my_shard, 0); - Some(&mut *db_guard) + db_slice.get_mut(0) } else { None }; crate::shard::spsc_handler::dispatch_vector_command( - &mut vs, - &mut *ts, + vs, + ts, #[cfg(feature = "graph")] - Some(&graph_guard), + Some(graph_ref), &command, db_opt, ) - }; + }); local_result } @@ -1630,6 +1484,7 @@ pub async fn scatter_invalidate_range( dispatch_tx: &Rc>>>, spsc_notifiers: &[Arc], ) -> Frame { + let _ = shard_databases; // E2 removes // PARTIAL-STATE: sends fire to all remotes in parallel, then we collect // sequentially. If any remote returns an error after others already // applied their deletes, this call returns Frame::Error but the @@ -1673,31 +1528,16 @@ pub async fn scatter_invalidate_range( } // Execute locally and add to total. - let local = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::shard::spsc_handler::dispatch_vector_command( - &mut s.vector_store, - &mut s.text_store, - #[cfg(feature = "graph")] - Some(&s.graph_store), - &command, - None, - ) - }) - } else { - let mut vs = shard_databases.vector_store(my_shard); - let mut ts = shard_databases.text_store(my_shard); - #[cfg(feature = "graph")] - let graph_guard = shard_databases.graph_store_read(my_shard); + let local = crate::shard::slice::with_shard(|s| { crate::shard::spsc_handler::dispatch_vector_command( - &mut vs, - &mut *ts, + &mut s.vector_store, + &mut s.text_store, #[cfg(feature = "graph")] - Some(&graph_guard), + Some(&s.graph_store), &command, None, ) - }; + }); match local { Frame::Integer(n) => Frame::Integer(total.saturating_add(n)), @@ -1741,6 +1581,7 @@ pub async fn scatter_text_search( highlight_opts: Option, summarize_opts: Option, ) -> Frame { + let _ = shard_databases; // E2 removes // Extract plain term strings for DocFreq phase (only needs term text, not modifiers). let term_strings: Vec = query_terms.iter().map(|qt| qt.text.clone()).collect(); @@ -1749,44 +1590,12 @@ pub async fn scatter_text_search( // Local IDF is globally accurate with one shard — skip DFS pre-pass. // Apply HIGHLIGHT/SUMMARIZE post-processing after local search. // - // Phase 2c: gate on is_initialized(); new path holds text_store + - // databases[0] simultaneously via a single `with_shard` call to avoid - // a reentrant `with_shard*` panic. - let result = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - let ts = &s.text_store; - let mut r = - crate::command::vector_search::ft_text_search::execute_text_search_local( - ts, - &index_name, - field_idx, - &query_terms, - top_k, - offset, - count, - ); - if highlight_opts.is_some() || summarize_opts.is_some() { - // Get the text_index reference, then borrow databases[0] - // disjointly. Both fields are tracked independently by rustc. - if let Some(text_index) = ts.get_index(&index_name) { - if let Some(db) = s.databases.get_mut(0) { - crate::command::vector_search::ft_text_search::apply_post_processing( - &mut r, - &term_strings, - text_index, - db, - highlight_opts.as_ref(), - summarize_opts.as_ref(), - ); - } - } - } - r - }) - } else { - let ts = shard_databases.text_store(my_shard); + // text_store + databases[0] accessed simultaneously via a single + // `with_shard` call to avoid a reentrant `with_shard*` panic. + let result = crate::shard::slice::with_shard(|s| { + let ts = &s.text_store; let mut r = crate::command::vector_search::ft_text_search::execute_text_search_local( - &ts, + ts, &index_name, field_idx, &query_terms, @@ -1794,24 +1603,24 @@ pub async fn scatter_text_search( offset, count, ); - // Apply post-processing if requested — guards held, no .await below. if highlight_opts.is_some() || summarize_opts.is_some() { + // Get the text_index reference, then borrow databases[0] + // disjointly. Both fields are tracked independently by rustc. if let Some(text_index) = ts.get_index(&index_name) { - let db_guard = shard_databases.read_db(my_shard, 0); - crate::command::vector_search::ft_text_search::apply_post_processing( - &mut r, - &term_strings, - text_index, - &*db_guard, - highlight_opts.as_ref(), - summarize_opts.as_ref(), - ); - // db_guard drops here. + if let Some(db) = s.databases.get_mut(0) { + crate::command::vector_search::ft_text_search::apply_post_processing( + &mut r, + &term_strings, + text_index, + db, + highlight_opts.as_ref(), + summarize_opts.as_ref(), + ); + } } } r - // MutexGuard dropped here — no .await held - }; + }); return result; } @@ -1826,9 +1635,8 @@ pub async fn scatter_text_search( for shard_id in 0..num_shards { if shard_id == my_shard { // Local: extract df/N directly — no SPSC overhead. - // CRITICAL: block scope drops MutexGuard before any .await (RESEARCH Pitfall 2). - // Phase 2c: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { + // Shard slice released before any .await. + let response = crate::shard::slice::with_shard(|s| match s.text_store.get_index(&index_name) { Some(text_index) => { let mut items: Vec = Vec::new(); @@ -1845,28 +1653,7 @@ pub async fn scatter_text_search( Frame::Array(items.into()) } None => Frame::Error(Bytes::from_static(b"ERR unknown index")), - }) - } else { - let ts = shard_databases.text_store(shard_id); - match ts.get_index(&index_name) { - Some(text_index) => { - let mut items: Vec = Vec::new(); - for (field_idx_opt, terms) in &field_queries { - let fidx = field_idx_opt.unwrap_or(0); - let (term_dfs, n) = text_index.doc_freq_for_terms(fidx, terms); - for (term, df) in term_dfs { - items.push(Frame::BulkString(Bytes::from(term))); - items.push(Frame::Integer(i64::from(df))); - } - items.push(Frame::BulkString(Bytes::from_static(b"N"))); - items.push(Frame::Integer(i64::from(n))); - } - Frame::Array(items.into()) - } - None => Frame::Error(Bytes::from_static(b"ERR unknown index")), - } - // MutexGuard dropped here, before .await below - }; + }); local_doc_freq = Some(response); } else { let (reply_tx, reply_rx) = channel::oneshot(); @@ -1905,45 +1692,11 @@ pub async fn scatter_text_search( for shard_id in 0..num_shards { if shard_id == my_shard { - // Local: execute with global IDF directly — block scope drops guard. - // CRITICAL: guard dropped before any .await (RESEARCH Pitfall 2). - // Phase 2c: gate on is_initialized(); fold text_store + databases[0] - // into a single `with_shard` to avoid reentrant `with_shard*` panic. - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - match s.text_store.get_index(&index_name) { - Some(text_index) => { - let mut r = - crate::command::vector_search::ft_text_search::execute_text_search_with_global_idf( - text_index, - field_idx, - &query_terms, - &global_df, - global_n, - top_k, - 0, // each shard returns top_k; coordinator applies final offset - top_k, - ); - if highlight_opts.is_some() || summarize_opts.is_some() { - if let Some(db) = s.databases.get_mut(0) { - crate::command::vector_search::ft_text_search::apply_post_processing( - &mut r, - &term_strings, - text_index, - db, - highlight_opts.as_ref(), - summarize_opts.as_ref(), - ); - } - } - r - } - None => Frame::Error(Bytes::from_static(b"ERR unknown index")), - } - }) - } else { - let ts = shard_databases.text_store(shard_id); - match ts.get_index(&index_name) { + // Local: execute with global IDF directly. + // text_store + databases[0] folded into a single `with_shard` to + // avoid reentrant `with_shard*` panic. Slice released before .await. + let response = crate::shard::slice::with_shard(|s| { + match s.text_store.get_index(&index_name) { Some(text_index) => { let mut r = crate::command::vector_search::ft_text_search::execute_text_search_with_global_idf( @@ -1956,25 +1709,23 @@ pub async fn scatter_text_search( 0, // each shard returns top_k; coordinator applies final offset top_k, ); - // Apply HIGHLIGHT/SUMMARIZE while guards are held — sync, no .await. if highlight_opts.is_some() || summarize_opts.is_some() { - let db_guard = shard_databases.read_db(shard_id, 0); - crate::command::vector_search::ft_text_search::apply_post_processing( - &mut r, - &term_strings, - text_index, - &*db_guard, - highlight_opts.as_ref(), - summarize_opts.as_ref(), - ); - // db_guard drops here. + if let Some(db) = s.databases.get_mut(0) { + crate::command::vector_search::ft_text_search::apply_post_processing( + &mut r, + &term_strings, + text_index, + db, + highlight_opts.as_ref(), + summarize_opts.as_ref(), + ); + } } r } None => Frame::Error(Bytes::from_static(b"ERR unknown index")), } - // MutexGuard dropped here, before .await below - }; + }); local_search = Some(response); } else { let (reply_tx, reply_rx) = channel::oneshot(); @@ -2048,10 +1799,10 @@ pub async fn scatter_text_search_filter( dispatch_tx: &Rc>>>, spsc_notifiers: &[Arc], ) -> Frame { + let _ = shard_databases; // E2 removes // ── Single-shard fast path ──────────────────────────────────────────────── if num_shards == 1 { - // Phase 2c: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { + let response = crate::shard::slice::with_shard(|s| match s.text_store.get_index(&index_name) { None => Frame::Error(Bytes::from_static(b"ERR no such index")), Some(text_index) => { @@ -2068,28 +1819,7 @@ pub async fn scatter_text_search_filter( &results, offset, count, ) } - }) - } else { - let ts = shard_databases.text_store(my_shard); - match ts.get_index(&index_name) { - None => Frame::Error(Bytes::from_static(b"ERR no such index")), - Some(text_index) => { - let clause = crate::command::vector_search::ft_text_search::TextQueryClause { - field_name: None, - terms: Vec::new(), - filter: Some(filter), - }; - let results = - crate::command::vector_search::ft_text_search::execute_query_on_index( - text_index, &clause, None, None, top_k, - ); - crate::command::vector_search::ft_text_search::build_text_response( - &results, offset, count, - ) - } - } - // MutexGuard dropped here - }; + }); return response; } @@ -2100,29 +1830,8 @@ pub async fn scatter_text_search_filter( for shard_id in 0..num_shards { if shard_id == my_shard { - // Phase 2c: gate on is_initialized(); new path uses ShardSlice directly. - let response = if crate::shard::slice::is_initialized() { + let response = crate::shard::slice::with_shard(|s| match s.text_store.get_index(&index_name) { - None => Frame::Error(Bytes::from_static(b"ERR no such index")), - Some(text_index) => { - let clause = - crate::command::vector_search::ft_text_search::TextQueryClause { - field_name: None, - terms: Vec::new(), - filter: Some(filter.clone()), - }; - let results = - crate::command::vector_search::ft_text_search::execute_query_on_index( - text_index, &clause, None, None, top_k, - ); - crate::command::vector_search::ft_text_search::build_text_response( - &results, 0, top_k, - ) - } - }) - } else { - let ts = shard_databases.text_store(my_shard); - match ts.get_index(&index_name) { None => Frame::Error(Bytes::from_static(b"ERR no such index")), Some(text_index) => { let clause = @@ -2141,9 +1850,7 @@ pub async fn scatter_text_search_filter( &results, 0, top_k, ) } - } - // MutexGuard dropped here, before the next `.await` - }; + }); local_response = Some(response); } else { let (reply_tx, reply_rx) = channel::oneshot(); @@ -2297,7 +2004,11 @@ pub async fn coordinate_swapdb( b"ERR SWAPDB aborted: WAL enqueue failed (persistence backpressure)", )); } - shard_databases.swap_dbs(my_shard, a, b); + crate::shard::slice::with_shard(|s| { + if a != b { + s.databases.swap(a, b); + } + }); } // Await all-remote-shard acks before returning +OK. @@ -2368,7 +2079,11 @@ mod tests { dbs[0].set_string(Bytes::from_static(b"a"), Bytes::from_static(b"1")); dbs[0].set_string(Bytes::from_static(b"b"), Bytes::from_static(b"2")); - let shard_databases = ShardDatabases::new(vec![dbs]); + let (shard_databases, mut inits) = ShardDatabases::new(vec![dbs]); + // coordinate_mget uses with_shard_db for local keys; ShardSlice must be initialized. + crate::shard::slice::reset_test_shard(crate::shard::slice::ShardSlice::new( + inits.remove(0), + )); let dispatch_tx: Rc>>> = Rc::new(RefCell::new(Vec::new())); @@ -2408,7 +2123,10 @@ mod tests { async fn test_coordinate_mset_all_local() { use crate::storage::Database; let dbs = vec![Database::new()]; - let shard_databases = ShardDatabases::new(vec![dbs]); + let (shard_databases, mut inits) = ShardDatabases::new(vec![dbs]); + crate::shard::slice::reset_test_shard(crate::shard::slice::ShardSlice::new( + inits.remove(0), + )); let dispatch_tx: Rc>>> = Rc::new(RefCell::new(Vec::new())); @@ -2436,11 +2154,11 @@ mod tests { .await; assert_eq!(result, Frame::SimpleString(Bytes::from_static(b"OK"))); - // Verify keys were set - let mut guard = shard_databases.write_db(0, 0); - guard.refresh_now_from_cache(&cached_clock); - let entry = guard.get(b"x"); - assert!(entry.is_some()); + // Verify keys were set via ShardSlice path. + crate::shard::slice::with_shard_db(0, |db| { + let entry = db.get(b"x"); + assert!(entry.is_some()); + }); } #[cfg(feature = "runtime-tokio")] @@ -2452,7 +2170,11 @@ mod tests { dbs[0].set_string(Bytes::from_static(b"b"), Bytes::from_static(b"2")); dbs[0].set_string(Bytes::from_static(b"c"), Bytes::from_static(b"3")); - let shard_databases = ShardDatabases::new(vec![dbs]); + let (shard_databases, mut inits) = ShardDatabases::new(vec![dbs]); + // coordinate_multi_del_or_exists uses with_shard_db for local keys; ShardSlice must be initialized. + crate::shard::slice::reset_test_shard(crate::shard::slice::ShardSlice::new( + inits.remove(0), + )); let dispatch_tx: Rc>>> = Rc::new(RefCell::new(Vec::new())); @@ -2571,7 +2293,13 @@ mod tests { use crate::storage::Database; let dbs = vec![Database::new()]; - let shard_databases = Arc::new(ShardDatabases::new(vec![dbs])); + let (shard_databases_inner, mut inits) = ShardDatabases::new(vec![dbs]); + // scatter_text_search calls execute_text_search_local which uses with_shard; + // ShardSlice must be initialized on this thread. + crate::shard::slice::reset_test_shard(crate::shard::slice::ShardSlice::new( + inits.remove(0), + )); + let shard_databases = Arc::new(shard_databases_inner); // Empty dispatch_tx — any SPSC send would panic (no channels configured). let dispatch_tx: Rc>>> = diff --git a/src/shard/dispatch.rs b/src/shard/dispatch.rs index dac17795..169f6ab1 100644 --- a/src/shard/dispatch.rs +++ b/src/shard/dispatch.rs @@ -236,9 +236,11 @@ pub struct FtHybridPayload { /// Non-zero values are forwarded to the dense KNN stream for MVCC filtering. /// BM25 stream remains AS_OF-unaware until text-index MVCC ships (v0.2). pub as_of_lsn: u64, - /// CHANGE F — per-shard filter forwarded to the raw-streams executor so the - /// allowlist is computed from the shard-local text index and applied to all - /// three streams BEFORE coordinator RRF fusion. `None` ⇒ unfiltered. + /// CHANGE F: optional pre-RRF filter tree forwarded from the coordinator's + /// `HybridQuery.filter`. The per-shard executor (`execute_hybrid_search_local_raw_streams`) + /// evaluates it against its OWN shard-local `text_index` (doc_ids are + /// shard-local — the allowlist CANNOT be computed at the coordinator). + /// `None` → unfiltered (backward compat). pub filter: Option, pub reply_tx: channel::OneshotSender, } @@ -592,10 +594,110 @@ pub enum ShardMessage { b: usize, reply_tx: channel::OneshotSender<()>, }, + /// MQ domain hop (C2, shardslice-migration Wave A1). + /// + /// The connection handler computes `owner = key_to_shard(effective_queue_key)` + /// and sends this message to the owning shard. The owner executes the full + /// MQ.* arm (durable flag, DLQ stream creation on max-delivery exhaustion, + /// trigger debounce, WAL records) against its ShardSlice — mirroring the + /// GraphCommand precedent. All six subcommands (CREATE/PUSH/POP/ACK/ + /// DLQLEN/TRIGGER) are dispatched via `mq_exec::execute_mq_on_owner`. + /// + /// Boxed because the inline payload (4 fields: usize + Bytes + Arc + + /// OneshotSender) totals ~48 B, which would push the enum past the 64-byte + /// (1 cache-line) cap when combined with the discriminant and alignment + /// padding. Boxing keeps this variant at 8 B (pointer only). + MqCommand(Box), + /// TXN.COMMIT MQ-intent materialize hop (C2, shardslice-migration Wave A1). + /// + /// The TXN.COMMIT path groups `txn.mq_intents` by owner shard and sends one + /// `MqTxnMaterialize` per owning shard. The owner applies the fold + /// `get_stream_mut → durable-check → next_auto_id → add` exactly as + /// txn.rs:160–167 does today, then replies `()` to unblock the coordinator. + /// + /// Inline: `usize(8) + Vec(24) + OneshotSender(8)` = 40 B — fits within + /// the 64-byte cap even after discriminant + alignment. + MqTxnMaterialize { + db_index: usize, + intents: Vec, + reply_tx: channel::OneshotSender<()>, + }, + /// WS.DROP best-effort key cleanup hop (C2, shardslice-migration Wave A1). + /// + /// `prefix` = `"{}:"`. The connection handler computes + /// `owner = key_to_shard(prefix)` and sends this to the owning shard + /// (the hash-tag in the prefix guarantees co-location). The owner scans + /// db 0, removes every key starting with `prefix`, and replies with the + /// deleted count (logged by the caller; not surfaced to the client). + /// + /// Inline: `Bytes(16) + OneshotSender(8)` = 24 B — comfortably within the cap. + WsDropCleanup { + prefix: Bytes, + reply_tx: channel::OneshotSender, + }, + /// AOF cooperative-snapshot hop (C2/C4, shardslice-migration Wave A1). + /// + /// The AOF rewrite writer pushes this message into the shard's external + /// mesh producer and blocks on the oneshot reply. The shard event loop + /// processes it atomically between commands: it builds an expired-filtered + /// snapshot of ALL its dbs (same fold as `do_rewrite_per_shard` phase 4) + /// and replies with the frozen view. The writer thread never touches the + /// shard's state directly. + /// + /// Inline: `OneshotSender(8)` — smallest possible inline payload. + AofFold { + reply_tx: channel::OneshotSender, + }, /// Graceful shutdown signal. Shutdown, } +/// Payload for [`ShardMessage::MqCommand`]. +/// +/// Boxed via `Box` in the enum variant to keep `ShardMessage` +/// within the 64-byte (1 cache-line) cap. +pub struct MqCommandPayload { + /// Target database index (from the connection's `selected_db`). + pub db_index: usize, + /// Workspace prefix (`"{wsid}:"` when the connection is WS-bound, empty + /// otherwise). The owner derives effective keys via + /// `workspace_key(key_prefix, raw_key)`, identical to the inline arm. + pub key_prefix: Bytes, + /// Full original MQ.* frame (CREATE/PUSH/POP/ACK/DLQLEN/TRIGGER). + /// Wrapped in `Arc` to avoid deep-cloning on dispatch. + pub command: std::sync::Arc, + /// Oneshot channel: owner sends the response `Frame` back to the caller. + pub reply_tx: channel::OneshotSender, +} + +/// Snapshot payload for [`ShardMessage::AofFold`]. +/// +/// Produced by the shard thread and consumed by the AOF rewrite writer thread. +/// Shape mirrors the per-shard snapshot that `do_rewrite_per_shard` phase 4 +/// builds today: `(entries, base_timestamp)` per db index. The writer feeds +/// this directly to `rdb::save_snapshot_to_bytes` unchanged. +pub struct AofFoldSnapshot { + /// One element per db: `(live_entries, base_timestamp)`. + /// Entries are pre-filtered — expired entries (per `is_expired_at`) are + /// excluded at snapshot time by the shard thread. + pub dbs: Vec<( + Vec<( + crate::storage::compact_key::CompactKey, + crate::storage::entry::Entry, + )>, + u32, + )>, + /// Number of messages in the AOF channel at the instant the shard read + /// this value — BEFORE building the snapshot and BEFORE sending this reply. + /// + /// Because the shard event loop is single-threaded, no new commands execute + /// between this read and the snapshot reply being sent. Therefore ALL of + /// these messages are pre-snapshot appends. The AOF writer uses this as the + /// bound for its phase-3 mid-drain, preventing an infinite drain loop under + /// sustained high write load (where the channel never empties on its own). + pub pending_aof_count: usize, +} + // ShardMessage is Send because all fields are Send. The raw pointer in // ResponseSlotPtr is the only non-auto-Send field, and it has its own // localized unsafe impl Send with documented safety invariants. diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index 5e9653d6..15664b76 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -73,9 +73,16 @@ impl super::Shard { all_pubsub_registries: Vec>>, all_remote_sub_maps: Vec>>, affinity_tracker: Arc>, + slice_init: crate::shard::slice::ShardSliceInit, ) { let _shard_id = self.id; + // C1: Initialize thread-local ShardSlice before any command handling. + // MUST be called before the accept/drain loop — assert_initialized panics + // on the first accept if this is skipped. + crate::shard::slice::init_shard(crate::shard::slice::ShardSlice::new(slice_init)); + crate::shard::slice::assert_initialized(self.id); + // Publish disk-offload status for INFO moonstore (set once per shard, idempotent). crate::vector::metrics::MOONSTORE_DISK_OFFLOAD_ENABLED.store( server_config.disk_offload_enabled(), @@ -804,22 +811,10 @@ impl super::Shard { if let Some(ref vdir) = vector_persist_dir { let _ = std::fs::create_dir_all(vdir); - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - s.vector_store.set_persist_dir(vdir.clone()); - s.text_store.set_persist_dir(vdir.clone()); - }); - } else { - let mut vs = shard_databases.vector_store(shard_id); - vs.set_persist_dir(vdir.clone()); - drop(vs); - - // Text indexes share the same shard directory (different filename). - let mut ts = shard_databases.text_store(shard_id); - ts.set_persist_dir(vdir.clone()); - drop(ts); - } + crate::shard::slice::with_shard(|s| { + s.vector_store.set_persist_dir(vdir.clone()); + s.text_store.set_persist_dir(vdir.clone()); + }); } // Try loading saved index metadata (with compaction weights) from the vector persist dir. @@ -841,32 +836,7 @@ impl super::Shard { }); if let Some(ref metas) = metas { - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - info!( - "Shard {}: restoring {} vector index(es) from sidecar", - shard_id, - metas.len() - ); - for (meta, weight) in metas { - let name = meta.name.clone(); - if let Err(e) = s.vector_store.create_index(meta.clone()) { - tracing::warn!( - "Shard {}: failed to restore index '{}': {}", - shard_id, - String::from_utf8_lossy(&name), - e - ); - } else if *weight != 1.0 { - if let Some(idx) = s.vector_store.get_index_mut(&name) { - idx.set_compaction_weight(*weight); - } - } - } - }); - } else { - let mut vs = shard_databases.vector_store(shard_id); + crate::shard::slice::with_shard(|s| { info!( "Shard {}: restoring {} vector index(es) from sidecar", shard_id, @@ -874,7 +844,7 @@ impl super::Shard { ); for (meta, weight) in metas { let name = meta.name.clone(); - if let Err(e) = vs.create_index(meta.clone()) { + if let Err(e) = s.vector_store.create_index(meta.clone()) { tracing::warn!( "Shard {}: failed to restore index '{}': {}", shard_id, @@ -882,48 +852,18 @@ impl super::Shard { e ); } else if *weight != 1.0 { - // W3-deep: restore persisted compaction_weight (skip default to avoid - // redundant write on every startup when weight=1.0). - if let Some(idx) = vs.get_index_mut(&name) { + if let Some(idx) = s.vector_store.get_index_mut(&name) { idx.set_compaction_weight(*weight); } } } - drop(vs); // release VectorStore lock before scanning databases - } + }); } // Restore text indexes from sidecar metadata. #[cfg(feature = "text-index")] if let Some(ref text_metas) = text_metas { - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - info!( - "Shard {}: restoring {} text index(es) from sidecar", - shard_id, - text_metas.len() - ); - for meta in text_metas { - let text_index = crate::text::store::TextIndex::new( - meta.name.clone(), - meta.key_prefixes.clone(), - meta.text_fields.clone(), - meta.bm25_config, - ); - if let Err(e) = s.text_store.create_index(meta.name.clone(), text_index) - { - tracing::warn!( - "Shard {}: failed to restore text index '{}': {}", - shard_id, - String::from_utf8_lossy(&meta.name), - e - ); - } - } - }); - } else { - let mut ts = shard_databases.text_store(shard_id); + crate::shard::slice::with_shard(|s| { info!( "Shard {}: restoring {} text index(es) from sidecar", shard_id, @@ -936,7 +876,7 @@ impl super::Shard { meta.text_fields.clone(), meta.bm25_config, ); - if let Err(e) = ts.create_index(meta.name.clone(), text_index) { + if let Err(e) = s.text_store.create_index(meta.name.clone(), text_index) { tracing::warn!( "Shard {}: failed to restore text index '{}': {}", shard_id, @@ -945,8 +885,7 @@ impl super::Shard { ); } } - drop(ts); - } + }); } // Auto-reindex existing HASH keys that match vector or text index prefixes. @@ -955,7 +894,6 @@ impl super::Shard { let db_count = shard_databases.db_count(); let mut reindexed = 0usize; for db_idx in 0..db_count { - // Phase 2d: gate read on is_initialized(); new path uses ShardSlice directly. let collect_matching = |db: &crate::storage::Database| -> Vec<(Vec, Vec)> { let mut matching: Vec<(Vec, Vec)> = Vec::new(); for (key, entry) in db.data().iter() { @@ -1009,39 +947,21 @@ impl super::Shard { } matching }; - let matching = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(db_idx, |db| collect_matching(db)) - } else { - let guard = shard_databases.read_db(shard_id, db_idx); - collect_matching(&guard) - }; + let matching = + { crate::shard::slice::with_shard_db(db_idx, |db| collect_matching(db)) }; if !matching.is_empty() { - // Phase 2d: gate write on is_initialized(). - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - for (key, args) in &matching { - let _ = crate::shard::spsc_handler::auto_index_hset_public( - &mut s.vector_store, - &mut s.text_store, - key, - args, - ); - reindexed += 1; - } - }); - } else { - let mut vs = shard_databases.vector_store(shard_id); - let mut ts = shard_databases.text_store(shard_id); + crate::shard::slice::with_shard(|s| { for (key, args) in &matching { - // Plan 166-01: return discarded — this is a - // reindex path, not a txn write path. let _ = crate::shard::spsc_handler::auto_index_hset_public( - &mut vs, &mut *ts, key, args, + &mut s.vector_store, + &mut s.text_store, + key, + args, ); reindexed += 1; } - } + }); } } if reindexed > 0 { @@ -1190,40 +1110,23 @@ impl super::Shard { _ = spsc_notify_local.notified() => { crate::admin::metrics_setup::bump_spsc_notify_wake(); let mut pending_snapshot = None; - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - let hit_cap = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - spsc_handler::drain_spsc_shared( - &shard_databases, &mut consumers, &mut *pubsub_arc.write(), - &blocking_rc, &mut pending_snapshot, &mut snapshot_state, - &mut wal_writer, &mut wal_v3_writer, &repl_backlog, &mut replica_txs, - &repl_offsets, shard_id, &script_cache_rc, &cached_clock, - &mut pending_migrations, &mut s.vector_store, - &mut pending_cdc_subscribes, - &mut shard_manifest, - server_config.mvcc_committed_prune_margin, - server_config.graph_merge_max_segments, - server_config.graph_dead_edge_trigger, - &mut autovacuum_daemon, - aof_pool.as_ref(), // FIX-W1-2 - ) - }) - } else { - spsc_handler::drain_spsc_shared( - &shard_databases, &mut consumers, &mut *pubsub_arc.write(), - &blocking_rc, &mut pending_snapshot, &mut snapshot_state, - &mut wal_writer, &mut wal_v3_writer, &repl_backlog, &mut replica_txs, - &repl_offsets, shard_id, &script_cache_rc, &cached_clock, - &mut pending_migrations, &mut *shard_databases.vector_store(shard_id), - &mut pending_cdc_subscribes, - &mut shard_manifest, - server_config.mvcc_committed_prune_margin, - server_config.graph_merge_max_segments, - server_config.graph_dead_edge_trigger, - &mut autovacuum_daemon, - aof_pool.as_ref(), // FIX-W1-2 - ) - }; + // No outer with_shard wrapper — each arm in drain_spsc_shared + // takes its own flat borrow, eliminating the re-entrancy BorrowMutError + // that occurred when arms called with_shard inside an enclosing borrow. + let hit_cap = spsc_handler::drain_spsc_shared( + &shard_databases, &mut consumers, &mut *pubsub_arc.write(), + &blocking_rc, &mut pending_snapshot, &mut snapshot_state, + &mut wal_writer, &mut wal_v3_writer, &repl_backlog, &mut replica_txs, + &repl_offsets, shard_id, &script_cache_rc, &cached_clock, + &mut pending_migrations, + &mut pending_cdc_subscribes, + &mut shard_manifest, + server_config.mvcc_committed_prune_margin, + server_config.graph_merge_max_segments, + server_config.graph_dead_edge_trigger, + &mut autovacuum_daemon, + aof_pool.as_ref(), // FIX-W1-2 + ); if hit_cap { // M3: capped drain may have left a tail — re-arm immediately // instead of stranding it until the next periodic tick. @@ -1304,40 +1207,21 @@ impl super::Shard { next_file_id = next_file_id.max(spill_file_id.get()); let mut pending_snapshot = None; - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - let hit_cap = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - spsc_handler::drain_spsc_shared( - &shard_databases, &mut consumers, &mut *pubsub_arc.write(), - &blocking_rc, &mut pending_snapshot, &mut snapshot_state, - &mut wal_writer, &mut wal_v3_writer, &repl_backlog, &mut replica_txs, - &repl_offsets, shard_id, &script_cache_rc, &cached_clock, - &mut pending_migrations, &mut s.vector_store, - &mut pending_cdc_subscribes, - &mut shard_manifest, - server_config.mvcc_committed_prune_margin, - server_config.graph_merge_max_segments, - server_config.graph_dead_edge_trigger, - &mut autovacuum_daemon, - aof_pool.as_ref(), // FIX-W1-2 - ) - }) - } else { - spsc_handler::drain_spsc_shared( - &shard_databases, &mut consumers, &mut *pubsub_arc.write(), - &blocking_rc, &mut pending_snapshot, &mut snapshot_state, - &mut wal_writer, &mut wal_v3_writer, &repl_backlog, &mut replica_txs, - &repl_offsets, shard_id, &script_cache_rc, &cached_clock, - &mut pending_migrations, &mut *shard_databases.vector_store(shard_id), - &mut pending_cdc_subscribes, - &mut shard_manifest, - server_config.mvcc_committed_prune_margin, - server_config.graph_merge_max_segments, - server_config.graph_dead_edge_trigger, - &mut autovacuum_daemon, - aof_pool.as_ref(), // FIX-W1-2 - ) - }; + // No outer with_shard — each arm takes its own flat borrow. + let hit_cap = spsc_handler::drain_spsc_shared( + &shard_databases, &mut consumers, &mut *pubsub_arc.write(), + &blocking_rc, &mut pending_snapshot, &mut snapshot_state, + &mut wal_writer, &mut wal_v3_writer, &repl_backlog, &mut replica_txs, + &repl_offsets, shard_id, &script_cache_rc, &cached_clock, + &mut pending_migrations, + &mut pending_cdc_subscribes, + &mut shard_manifest, + server_config.mvcc_committed_prune_margin, + server_config.graph_merge_max_segments, + server_config.graph_dead_edge_trigger, + &mut autovacuum_daemon, + aof_pool.as_ref(), // FIX-W1-2 + ); if hit_cap { // M3: capped drain may have left a tail — re-arm immediately // instead of stranding it until the next periodic tick. @@ -1514,28 +1398,16 @@ impl super::Shard { timers::sync_wal_v3(&mut wal_v3_writer); // P3+MA1+MA2: prune committed + sweep zombies + kill old snapshots // + update RECL_MVCC_* + segment-stall. - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - timers::run_mvcc_sweep( - &mut s.vector_store, - #[cfg(feature = "graph")] - &mut s.graph_store, - server_config.mvcc_committed_prune_margin, - server_config.max_unflushed_immutable_segments, - server_config.mvcc_old_snapshot_threshold_secs, - ); - }); - } else { + crate::shard::slice::with_shard(|s| { timers::run_mvcc_sweep( - &mut *shard_databases.vector_store(shard_id), + &mut s.vector_store, #[cfg(feature = "graph")] - &mut *shard_databases.graph_store_write(shard_id), + &mut s.graph_store, server_config.mvcc_committed_prune_margin, server_config.max_unflushed_immutable_segments, server_config.mvcc_old_snapshot_threshold_secs, ); - } + }); // P6: ceiling-trigger — runs at 1s cadence to avoid the // read_dir syscall overhead of wal.stats() on every 1ms tick. if let (Some(ckpt_mgr), Some(page_cache_inst), Some(wal_v3), Some(manifest), Some(ctrl), Some(ctrl_path)) = @@ -1563,22 +1435,9 @@ impl super::Shard { if let Some(ref mut manifest) = shard_manifest { let shard_dir = server_config.effective_disk_offload_dir() .join(format!("shard-{}", shard_id)); - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - persistence_tick::check_warm_transitions( - &s.vector_store, - &shard_dir, - manifest, - server_config.segment_warm_after, - &mut next_file_id, - shard_id, - &mut wal_v3_writer, - ); - }); - } else { + crate::shard::slice::with_shard(|s| { persistence_tick::check_warm_transitions( - &*shard_databases.vector_store(shard_id), + &s.vector_store, &shard_dir, manifest, server_config.segment_warm_after, @@ -1586,28 +1445,19 @@ impl super::Shard { shard_id, &mut wal_v3_writer, ); - } + }); } } // Budget enforcement runs on every warm-check tick regardless of // disk-offload state: warm segments can accumulate from in-memory // compaction even without disk-offload enabled. - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - persistence_tick::enforce_warm_mmap_budget( - &s.vector_store, - &mut warm_mmap_budget, - shard_id, - ); - }); - } else { + crate::shard::slice::with_shard(|s| { persistence_tick::enforce_warm_mmap_budget( - &*shard_databases.vector_store(shard_id), + &s.vector_store, &mut warm_mmap_budget, shard_id, ); - } + }); } // Cold tier transition check (60s, disk-offload only) _ = cold_check_interval.0.tick() => { @@ -1615,28 +1465,16 @@ impl super::Shard { if let Some(ref mut manifest) = shard_manifest { let shard_dir = server_config.effective_disk_offload_dir() .join(format!("shard-{}", shard_id)); - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - persistence_tick::check_cold_transitions( - &s.vector_store, - &shard_dir, - manifest, - server_config.segment_cold_after, - &mut next_file_id, - shard_id, - ); - }); - } else { + crate::shard::slice::with_shard(|s| { persistence_tick::check_cold_transitions( - &*shard_databases.vector_store(shard_id), + &s.vector_store, &shard_dir, manifest, server_config.segment_cold_after, &mut next_file_id, shard_id, ); - } + }); } } } @@ -1706,29 +1544,11 @@ impl super::Shard { } // P4: Autovacuum daemon tick (default 30s interval). _ = autovacuum_interval.0.tick() => { - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - autovacuum_daemon.run_tick( - &mut s.vector_store, - #[cfg(feature = "graph")] - &mut s.graph_store, - shard_manifest.as_mut(), - wal_v3_writer.as_mut(), - server_config.max_wal_size_bytes(), - server_config.manifest_tombstone_retain_epochs, - server_config.manifest_tombstone_retain_secs, - server_config.max_unflushed_immutable_segments as usize, - server_config.graph_merge_max_segments, - server_config.graph_dead_edge_trigger, - false, - ); - }); - } else { + crate::shard::slice::with_shard(|s| { autovacuum_daemon.run_tick( - &mut *shard_databases.vector_store(shard_id), + &mut s.vector_store, #[cfg(feature = "graph")] - &mut *shard_databases.graph_store_write(shard_id), + &mut s.graph_store, shard_manifest.as_mut(), wal_v3_writer.as_mut(), server_config.max_wal_size_bytes(), @@ -1739,7 +1559,7 @@ impl super::Shard { server_config.graph_dead_edge_trigger, false, ); - } + }); } _ = shutdown.cancelled() => { info!("Shard {} shutting down", self.id); @@ -1756,35 +1576,23 @@ impl super::Shard { persistence_tick::force_checkpoint(ckpt_mgr, page_cache_inst, wal_v3, manifest, ctrl, ctrl_path, shard_id, server_config.manifest_tombstone_retain_epochs, server_config.manifest_tombstone_retain_secs); } // Persist graph store to disk on shutdown. - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. #[cfg(feature = "graph")] if let Some(ref dir) = persistence_dir { - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - if s.graph_store.graph_count() > 0 { - if let Err(e) = crate::graph::recovery::save_graph_store( - &s.graph_store, - std::path::Path::new(dir), - shard_id, - ) { - tracing::warn!( - "Shard {shard_id}: failed to save graph store on shutdown: {e}" - ); - } else { - info!("Shard {shard_id}: graph store saved to {dir}"); - } - } - }); - } else { - let gs = shard_databases.graph_store_read(shard_id); - if gs.graph_count() > 0 { - if let Err(e) = crate::graph::recovery::save_graph_store(&gs, std::path::Path::new(dir), shard_id) { - tracing::warn!("Shard {shard_id}: failed to save graph store on shutdown: {e}"); + crate::shard::slice::with_shard(|s| { + if s.graph_store.graph_count() > 0 { + if let Err(e) = crate::graph::recovery::save_graph_store( + &s.graph_store, + std::path::Path::new(dir), + shard_id, + ) { + tracing::warn!( + "Shard {shard_id}: failed to save graph store on shutdown: {e}" + ); } else { info!("Shard {shard_id}: graph store saved to {dir}"); } } - } + }); } if let Some(ref mut wal) = wal_writer { let _ = wal.shutdown(); @@ -1960,62 +1768,31 @@ impl super::Shard { // --- Every-wake body (mirrors the tokio notify arm): drain SPSC, // handle drain outputs, sweep the pending_wakers relay --- let mut pending_snapshot = None; - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - let hit_cap = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - spsc_handler::drain_spsc_shared( - &shard_databases, - &mut consumers, - &mut *pubsub_arc.write(), - &blocking_rc, - &mut pending_snapshot, - &mut snapshot_state, - &mut wal_writer, - &mut wal_v3_writer, - &repl_backlog, - &mut replica_txs, - &repl_offsets, - shard_id, - &script_cache_rc, - &cached_clock, - &mut pending_migrations, - &mut s.vector_store, - &mut pending_cdc_subscribes, - &mut shard_manifest, - server_config.mvcc_committed_prune_margin, - server_config.graph_merge_max_segments, - server_config.graph_dead_edge_trigger, - &mut autovacuum_daemon, - aof_pool.as_ref(), // FIX-W1-2 - ) - }) - } else { - spsc_handler::drain_spsc_shared( - &shard_databases, - &mut consumers, - &mut *pubsub_arc.write(), - &blocking_rc, - &mut pending_snapshot, - &mut snapshot_state, - &mut wal_writer, - &mut wal_v3_writer, - &repl_backlog, - &mut replica_txs, - &repl_offsets, - shard_id, - &script_cache_rc, - &cached_clock, - &mut pending_migrations, - &mut *shard_databases.vector_store(shard_id), - &mut pending_cdc_subscribes, - &mut shard_manifest, - server_config.mvcc_committed_prune_margin, - server_config.graph_merge_max_segments, - server_config.graph_dead_edge_trigger, - &mut autovacuum_daemon, - aof_pool.as_ref(), // FIX-W1-2 - ) - }; + // No outer with_shard — each arm takes its own flat borrow. + let hit_cap = spsc_handler::drain_spsc_shared( + &shard_databases, + &mut consumers, + &mut *pubsub_arc.write(), + &blocking_rc, + &mut pending_snapshot, + &mut snapshot_state, + &mut wal_writer, + &mut wal_v3_writer, + &repl_backlog, + &mut replica_txs, + &repl_offsets, + shard_id, + &script_cache_rc, + &cached_clock, + &mut pending_migrations, + &mut pending_cdc_subscribes, + &mut shard_manifest, + server_config.mvcc_committed_prune_margin, + server_config.graph_merge_max_segments, + server_config.graph_dead_edge_trigger, + &mut autovacuum_daemon, + aof_pool.as_ref(), // FIX-W1-2 + ); if hit_cap { // M3: the drain stopped at its per-cycle cap (or a snapshot // barrier) — re-arm immediately so the tail drains on the next @@ -2256,28 +2033,16 @@ impl super::Shard { timers::sync_wal_v3(&mut wal_v3_writer); // P3+MA1+MA2: MVCC committed prune + zombie sweep + kill old snapshots // + RECL_* + segment-stall. - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - timers::run_mvcc_sweep( - &mut s.vector_store, - #[cfg(feature = "graph")] - &mut s.graph_store, - server_config.mvcc_committed_prune_margin, - server_config.max_unflushed_immutable_segments, - server_config.mvcc_old_snapshot_threshold_secs, - ); - }); - } else { + crate::shard::slice::with_shard(|s| { timers::run_mvcc_sweep( - &mut *shard_databases.vector_store(shard_id), + &mut s.vector_store, #[cfg(feature = "graph")] - &mut *shard_databases.graph_store_write(shard_id), + &mut s.graph_store, server_config.mvcc_committed_prune_margin, server_config.max_unflushed_immutable_segments, server_config.mvcc_old_snapshot_threshold_secs, ); - } + }); if let ( Some(ckpt_mgr), Some(page_cache_inst), @@ -2316,22 +2081,9 @@ impl super::Shard { let shard_dir = server_config .effective_disk_offload_dir() .join(format!("shard-{}", shard_id)); - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - persistence_tick::check_warm_transitions( - &s.vector_store, - &shard_dir, - manifest, - server_config.segment_warm_after, - &mut next_file_id, - shard_id, - &mut wal_v3_writer, - ); - }); - } else { + crate::shard::slice::with_shard(|s| { persistence_tick::check_warm_transitions( - &*shard_databases.vector_store(shard_id), + &s.vector_store, &shard_dir, manifest, server_config.segment_warm_after, @@ -2339,26 +2091,17 @@ impl super::Shard { shard_id, &mut wal_v3_writer, ); - } + }); } } // Budget enforcement: runs on every warm-check tick. - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - persistence_tick::enforce_warm_mmap_budget( - &s.vector_store, - &mut warm_mmap_budget, - shard_id, - ); - }); - } else { + crate::shard::slice::with_shard(|s| { persistence_tick::enforce_warm_mmap_budget( - &*shard_databases.vector_store(shard_id), + &s.vector_store, &mut warm_mmap_budget, shard_id, ); - } + }); } // Cold tier check: every cold_poll_secs * 1000 ticks if monoio_tick_counter % (cold_poll_secs as u64 * 1000) == 0 { @@ -2368,28 +2111,16 @@ impl super::Shard { let shard_dir = server_config .effective_disk_offload_dir() .join(format!("shard-{}", shard_id)); - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - persistence_tick::check_cold_transitions( - &s.vector_store, - &shard_dir, - manifest, - server_config.segment_cold_after, - &mut next_file_id, - shard_id, - ); - }); - } else { + crate::shard::slice::with_shard(|s| { persistence_tick::check_cold_transitions( - &*shard_databases.vector_store(shard_id), + &s.vector_store, &shard_dir, manifest, server_config.segment_cold_after, &mut next_file_id, shard_id, ); - } + }); } } } @@ -2401,29 +2132,11 @@ impl super::Shard { if monoio_tick_counter % (autovacuum_interval_secs * 1000) == 0 && monoio_tick_counter > 0 { - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - autovacuum_daemon.run_tick( - &mut s.vector_store, - #[cfg(feature = "graph")] - &mut s.graph_store, - shard_manifest.as_mut(), - wal_v3_writer.as_mut(), - server_config.max_wal_size_bytes(), - server_config.manifest_tombstone_retain_epochs, - server_config.manifest_tombstone_retain_secs, - server_config.max_unflushed_immutable_segments as usize, - server_config.graph_merge_max_segments, - server_config.graph_dead_edge_trigger, - false, - ); - }); - } else { + crate::shard::slice::with_shard(|s| { autovacuum_daemon.run_tick( - &mut *shard_databases.vector_store(shard_id), + &mut s.vector_store, #[cfg(feature = "graph")] - &mut *shard_databases.graph_store_write(shard_id), + &mut s.graph_store, shard_manifest.as_mut(), wal_v3_writer.as_mut(), server_config.max_wal_size_bytes(), @@ -2434,7 +2147,7 @@ impl super::Shard { server_config.graph_dead_edge_trigger, false, ); - } + }); } // Cold-tier orphan sweep (P9): every orphan_sweep_interval_secs * 1000 ticks. // Matches the tokio select! branch above. Disabled when interval is 0. diff --git a/src/shard/mesh.rs b/src/shard/mesh.rs index c08bd0da..dc57ead6 100644 --- a/src/shard/mesh.rs +++ b/src/shard/mesh.rs @@ -192,6 +192,38 @@ pub fn create_admin_channels( (producers, consumers) } +/// Create N SPSC channels for the AOF fold cooperative snapshot protocol (C4). +/// +/// Returns `(producers, consumers)` where: +/// - `producers[i]` is the SPSC producer into shard `i`'s event-loop ring, +/// wrapped in `Arc` so the AOF writer thread (not the shard thread) +/// can push `ShardMessage::AofFold` safely. +/// - `consumers[i]` is the matching consumer, to be merged into shard `i`'s +/// consumers vec before `shard.run()` is called (same merge pattern as admin +/// channels — see `create_admin_channels`). +/// +/// `capacity` of 4 is sufficient: only one `AofFold` can be in flight per +/// shard at a time (the writer blocks on the reply before issuing another), so +/// the ring never holds more than 1 entry in practice; 4 provides headroom for +/// any transient burst. +pub fn create_aof_fold_channels( + num_shards: usize, + capacity: usize, +) -> ( + Vec>>>, + Vec>, +) { + let mut producers = Vec::with_capacity(num_shards); + let mut consumers = Vec::with_capacity(num_shards); + for _ in 0..num_shards { + let rb = HeapRb::new(capacity); + let (prod, cons) = rb.split(); + producers.push(Arc::new(parking_lot::Mutex::new(prod))); + consumers.push(cons); + } + (producers, consumers) +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/shard/mod.rs b/src/shard/mod.rs index bfc14d17..07f4f6b7 100644 --- a/src/shard/mod.rs +++ b/src/shard/mod.rs @@ -9,6 +9,8 @@ pub mod event_loop; /// MA5: maintenance-window scheduler (cron-style budget multipliers). pub mod maintenance_schedule; pub mod mesh; +/// C2 (shardslice-migration Wave A1): owner-side MQ.* execution on the shard thread. +pub(crate) mod mq_exec; pub mod numa; pub mod persistence_tick; pub mod remote_subscriber_map; @@ -359,7 +361,7 @@ mod tests { #[test] fn test_pubsub_fanout_via_spsc() { let mut pubsub = PubSubRegistry::new(); - let shard_databases = ShardDatabases::new(vec![vec![Database::new()]]); + let (shard_databases, _inits) = ShardDatabases::new(vec![vec![Database::new()]]); let (tx, rx) = rt_channel::mpsc_bounded::(16); let sub = Subscriber::new(tx, 42); @@ -384,7 +386,6 @@ mod tests { let blocking = Rc::new(RefCell::new(BlockingRegistry::new(0))); let script_cache = Rc::new(RefCell::new(crate::scripting::ScriptCache::new())); let clock = CachedClock::new(); - let mut vs = crate::vector::store::VectorStore::new(); let backlog = std::sync::Arc::new(parking_lot::Mutex::new(None)); spsc_handler::drain_spsc_shared( &shard_databases, @@ -402,7 +403,6 @@ mod tests { &script_cache, &clock, &mut Vec::new(), - &mut vs, &mut Vec::new(), &mut None, // shard_manifest — None in tests (no persistence_dir) 1000, // mvcc_prune_margin default @@ -424,7 +424,7 @@ mod tests { #[test] fn test_drain_spsc_respects_limit() { let mut pubsub = PubSubRegistry::new(); - let shard_databases = ShardDatabases::new(vec![vec![Database::new()]]); + let (shard_databases, _inits) = ShardDatabases::new(vec![vec![Database::new()]]); let rb = HeapRb::new(512); let (mut prod, cons) = rb.split(); @@ -448,7 +448,6 @@ mod tests { let blocking = Rc::new(RefCell::new(BlockingRegistry::new(0))); let script_cache = Rc::new(RefCell::new(crate::scripting::ScriptCache::new())); let clock = CachedClock::new(); - let mut vs = crate::vector::store::VectorStore::new(); let backlog = std::sync::Arc::new(parking_lot::Mutex::new(None)); spsc_handler::drain_spsc_shared( &shard_databases, @@ -466,7 +465,6 @@ mod tests { &script_cache, &clock, &mut Vec::new(), - &mut vs, &mut Vec::new(), &mut None, // shard_manifest — None in tests (no persistence_dir) 1000, // mvcc_prune_margin default @@ -525,7 +523,7 @@ mod tests { let config = RuntimeConfig::default(); let shard = Shard::new(0, 1, 1, config); - let shard_databases = ShardDatabases::new(vec![shard.databases]); + let (shard_databases, _inits) = ShardDatabases::new(vec![shard.databases]); let mut parse_bufs = std::collections::HashMap::new(); parse_bufs.insert(42u32, bytes::BytesMut::from(&b"partial"[..])); let mut inflight_sends = std::collections::HashMap::new(); @@ -563,7 +561,7 @@ mod tests { let config = RuntimeConfig::default(); let shard = Shard::new(0, 1, 1, config); - let shard_databases = ShardDatabases::new(vec![shard.databases]); + let (shard_databases, _inits) = ShardDatabases::new(vec![shard.databases]); let mut parse_bufs = std::collections::HashMap::new(); let mut inflight_sends = std::collections::HashMap::new(); diff --git a/src/shard/mq_exec.rs b/src/shard/mq_exec.rs new file mode 100644 index 00000000..bae20342 --- /dev/null +++ b/src/shard/mq_exec.rs @@ -0,0 +1,672 @@ +//! Owner-side MQ.* command execution for the ShardSlice (shardslice-migration C2). +//! +//! `execute_mq_on_owner` runs **on the shard thread** that owns the queue key, +//! after a `ShardMessage::MqCommand` hop. It faithfully mirrors ALL SIX subcommand +//! arms from `handler_sharded/write.rs` and `handler_monoio` (which are byte- +//! identical mirrors of each other). +//! +//! # Key derivation +//! +//! `MqCommandPayload.key_prefix` carries the pre-computed workspace prefix +//! `"{}:"` (35 bytes) when the connection is WS-bound, or an empty +//! `Bytes` otherwise. The owner derives the effective key as: +//! +//! ```text +//! effective_key = concat(key_prefix, raw_queue_key) +//! ``` +//! +//! The `trigger_key` used by PUSH/TRIGGER also needs the prefix but WITHOUT the +//! hash-tag braces: `ws_hex:raw_queue_key`. This is extracted from `key_prefix` +//! by stripping the leading `{` and the trailing `}:` characters (see +//! `derive_trig_key`). +//! +//! # WAL appends +//! +//! WAL records are written via the slice's `wal_append_tx` channel sender +//! (fire-and-forget, same as the existing `ShardDatabases::wal_append` path). +//! Only MQ.CREATE and MQ.ACK emit WAL records — matching the lock-path arms. +//! +//! # Mirror divergences found +//! +//! NONE: the two runtime mirrors (`handler_sharded/write.rs` and +//! `handler_monoio/write.rs`) are byte-identical for MQ logic. The cached-clock +//! read in the PUSH trigger-debounce arm is resolved by passing `now_ms` +//! explicitly (the owner shard's thread-local clock, same as the cached clock). + +use bytes::Bytes; + +use crate::command::mq::{ + ERR_MQ_NOT_DURABLE, ERR_MQ_UNKNOWN_SUB, parse_mq_subcommand, validate_mq_ack, + validate_mq_create, validate_mq_dlqlen, validate_mq_pop, validate_mq_push, validate_mq_trigger, +}; +use crate::mq::is_mq_command; +use crate::protocol::Frame; +use crate::storage::entry::current_time_ms; +use crate::storage::stream::StreamId; + +/// Execute one MQ.* subcommand on the owning shard thread against its ShardSlice. +/// +/// Called from `spsc_handler::handle_shard_message_shared` when +/// `ShardMessage::MqCommand` is received. Takes the three data fields directly +/// (the caller destructures `MqCommandPayload` and sends the `reply_tx` +/// separately) so no dummy channel allocation is needed. +/// +/// # Contract +/// +/// - Runs synchronously on the shard's OS thread. +/// - `with_shard` / `with_shard_db` calls inside this function are NOT +/// re-entrant: each subcommand handler uses exactly one `with_shard*` call +/// per logical step, and no call is nested. +/// - Allocations: only at result-building boundaries (`Vec::with_capacity`). +/// No `format!`, `to_string`, or needless `clone` on hot fields. +/// - WAL appends use fire-and-forget `try_send` via `wal_append_on_slice`. +pub(crate) fn execute_mq_on_owner( + db_index: usize, + key_prefix: Bytes, + command: std::sync::Arc, +) -> Frame { + // full_args[0] = command name (b"MQ"), full_args[1..] = subcommand + params. + // Mirrors try_handle_mq_command which receives cmd and cmd_args separately: + // cmd = b"MQ", cmd_args = full_args[1..]. + let full_args = match &*command { + Frame::Array(a) if a.len() >= 2 => a.as_slice(), + _ => { + return Frame::Error(Bytes::from_static(b"ERR invalid MQ command format")); + } + }; + + // Verify the command name is MQ. + let cmd_name = match &full_args[0] { + Frame::BulkString(b) => b.as_ref(), + _ => return Frame::Error(Bytes::from_static(b"ERR invalid command format")), + }; + if !is_mq_command(cmd_name) { + return Frame::Error(Bytes::from_static(ERR_MQ_UNKNOWN_SUB)); + } + + // cmd_args mirrors what try_handle_mq_command receives: subcommand at [0]. + let cmd_args = &full_args[1..]; + let sub = match parse_mq_subcommand(cmd_args) { + Ok(s) => s, + Err(e) => return e, + }; + + if sub.eq_ignore_ascii_case(b"CREATE") { + return handle_create(cmd_args, &key_prefix, db_index); + } + if sub.eq_ignore_ascii_case(b"PUSH") { + return handle_push(cmd_args, &key_prefix, db_index); + } + if sub.eq_ignore_ascii_case(b"POP") { + return handle_pop(cmd_args, &key_prefix, db_index); + } + if sub.eq_ignore_ascii_case(b"ACK") { + return handle_ack(cmd_args, &key_prefix, db_index); + } + if sub.eq_ignore_ascii_case(b"DLQLEN") { + return handle_dlqlen(cmd_args, &key_prefix, db_index); + } + if sub.eq_ignore_ascii_case(b"TRIGGER") { + return handle_trigger(cmd_args, &key_prefix, db_index); + } + + Frame::Error(Bytes::from_static(ERR_MQ_UNKNOWN_SUB)) +} + +// ── Effective-key derivation ────────────────────────────────────────────────── + +/// Build the effective queue key: `concat(key_prefix, raw_queue_key)`. +/// +/// When `key_prefix` is empty (no workspace binding), returns a copy of +/// `raw_queue_key`. Otherwise prepends the workspace prefix (`{ws_hex}:`). +/// +/// Allocation: one `Vec` per call; acceptable at command granularity. +#[inline] +fn effective_key(key_prefix: &Bytes, raw_key: &Bytes) -> Bytes { + if key_prefix.is_empty() { + raw_key.clone() + } else { + let mut buf = Vec::with_capacity(key_prefix.len() + raw_key.len()); + buf.extend_from_slice(key_prefix); + buf.extend_from_slice(raw_key); + Bytes::from(buf) + } +} + +/// Build the trigger registry key from `key_prefix` and `raw_queue_key`. +/// +/// The trigger registry indexes triggers by `ws_hex:queue_key` (without the +/// hash-tag braces), while `key_prefix = "{ws_hex}:"`. Strip the `{` at [0] +/// and the `}` at [key_prefix.len()-2] (keeping the trailing `:`): +/// +/// ```text +/// key_prefix = "{" + ws_hex(32) + "}:" → 35 bytes +/// trig_prefix = ws_hex(32) + ":" → 33 bytes = key_prefix[1..33] + ":" +/// ``` +/// +/// When `key_prefix` is empty (no workspace) the trigger key equals `raw_queue_key`. +#[inline] +fn derive_trig_key(key_prefix: &Bytes, raw_queue_key: &Bytes) -> Bytes { + if key_prefix.is_empty() { + raw_queue_key.clone() + } else { + // key_prefix = "{" + ws_hex(32) + "}:" + // We want: ws_hex(32) + ":" + raw_queue_key + // i.e. key_prefix[1..key_prefix.len()-1] + raw_queue_key + // key_prefix.len()-1 strips the trailing ":" — wait, we keep ":" + // key_prefix[1..] = ws_hex + "}:" → we want ws_hex + ":" + // → key_prefix[1..key_prefix.len()-1] gives "ws_hex}" — no. + // Let's be explicit: + // key_prefix bytes: b'{', hex×32, b'}', b':' + // We want: hex×32, b':', raw_queue_key... + // That is: &key_prefix[1..key_prefix.len()-1] + raw_queue_key + // key_prefix[1..len-1] = hex(32) + "}" — still has "}" + // + // Correct: we want key_prefix[1..len-2] + ":" + raw_queue_key + // key_prefix[1..len-2] = hex(32) + // then append ":" (1 byte) and raw_queue_key + let prefix_len = key_prefix.len(); + if prefix_len < 4 { + // Malformed prefix; fall back to raw key + return raw_queue_key.clone(); + } + // ws_hex bytes = key_prefix[1..prefix_len-2] + let ws_hex = &key_prefix[1..prefix_len - 2]; + let mut buf = Vec::with_capacity(ws_hex.len() + 1 + raw_queue_key.len()); + buf.extend_from_slice(ws_hex); + buf.push(b':'); + buf.extend_from_slice(raw_queue_key); + Bytes::from(buf) + } +} + +// ── WAL append helper ───────────────────────────────────────────────────────── + +/// Append a WAL record byte sequence on the current shard thread via the +/// slice's `wal_append_tx` channel (fire-and-forget). +/// +/// Mirrors `ShardDatabases::wal_append` semantics — `try_send` failures are +/// ignored (the channel is bounded; under extreme backpressure the record is +/// dropped, same as the existing lock-path). +#[inline] +fn wal_append_on_slice(record_bytes: bytes::Bytes) { + crate::shard::slice::with_shard(|s| { + if let Some(ref tx) = s.wal_append_tx { + let _ = tx.try_send(record_bytes); + } + }); +} + +// ── Subcommand handlers ─────────────────────────────────────────────────────── + +/// MQ.CREATE — owner creates the durable stream + registry entry. +/// +/// Mirrors handler_sharded/write.rs MQ CREATE arm (lock-path `else` branch). +fn handle_create(args: &[Frame], key_prefix: &Bytes, db_index: usize) -> Frame { + let (raw_key, max_delivery_count, _debounce_ms) = match validate_mq_create(args) { + Ok(v) => v, + Err(e) => return e, + }; + let eff_key = effective_key(key_prefix, &raw_key); + + // Create / configure the durable stream. + let create_result: Result<(), Frame> = crate::shard::slice::with_shard_db(db_index, |db| { + match db.get_or_create_stream(&eff_key) { + Ok(stream) => { + stream.durable = true; + stream.max_delivery_count = max_delivery_count; + let group_name = Bytes::from_static(b"__mq_consumers"); + let _ = stream.create_group(group_name, StreamId::ZERO); + Ok(()) + } + Err(e) => Err(e), + } + }); + if let Err(e) = create_result { + return e; + } + + // Register in the durable-queue registry (lazy-init on first CREATE). + let config = crate::mq::DurableStreamConfig::new(eff_key.clone(), max_delivery_count); + crate::shard::slice::with_shard(|s| { + let reg = s + .durable_queue_registry + .get_or_insert_with(|| Box::new(crate::mq::DurableQueueRegistry::new())); + reg.insert(eff_key.clone(), config); + }); + + // WAL: MqCreate record on the owner shard. + { + let payload = crate::mq::wal::encode_mq_create(&eff_key, max_delivery_count); + let mut wal_buf = Vec::new(); + crate::persistence::wal_v3::record::write_wal_v3_record( + &mut wal_buf, + 0, + crate::persistence::wal_v3::record::WalRecordType::MqCreate, + &payload, + ); + wal_append_on_slice(Bytes::from(wal_buf)); + } + + Frame::SimpleString(Bytes::from_static(b"OK")) +} + +/// MQ.PUSH — owner enqueues a message into the durable stream. +/// +/// Mirrors handler_sharded/write.rs MQ PUSH arm (lock-path `else` branch). +/// Fires any registered trigger's debounce timer (sets `pending_fire_ms` if +/// not already pending). +fn handle_push(args: &[Frame], key_prefix: &Bytes, db_index: usize) -> Frame { + let (raw_key, fields) = match validate_mq_push(args) { + Ok(v) => v, + Err(e) => return e, + }; + let eff_key = effective_key(key_prefix, &raw_key); + let trig_key = derive_trig_key(key_prefix, &raw_key); + + // Push into the stream. + type PushResult = Result, Frame>; + let push_result: PushResult = + crate::shard::slice::with_shard_db(db_index, |db| match db.get_stream_mut(&eff_key) { + Ok(Some(stream)) => { + if !stream.durable { + Ok(None) + } else { + let msg_id = stream.next_auto_id(); + let msg_id = stream.add(msg_id, fields); + Ok(Some(msg_id)) + } + } + Ok(None) => Ok(None), + Err(e) => Err(e), + }); + + match push_result { + Ok(Some(msg_id)) => { + // Debounce trigger: set pending_fire_ms if not already armed. + let now_ms = current_time_ms(); + crate::shard::slice::with_shard(|s| { + if let Some(ref mut reg) = s.trigger_registry { + if let Some(trig_entry) = reg.get_mut(&trig_key) { + if trig_entry.pending_fire_ms == 0 { + trig_entry.pending_fire_ms = now_ms + trig_entry.debounce_ms; + } + } + } + }); + + let mut buf = itoa::Buffer::new(); + let ms_str = buf.format(msg_id.ms); + let mut buf2 = itoa::Buffer::new(); + let seq_str = buf2.format(msg_id.seq); + let mut id_bytes = Vec::with_capacity(ms_str.len() + 1 + seq_str.len()); + id_bytes.extend_from_slice(ms_str.as_bytes()); + id_bytes.push(b'-'); + id_bytes.extend_from_slice(seq_str.as_bytes()); + Frame::BulkString(Bytes::from(id_bytes)) + } + Ok(None) => Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE)), + Err(e) => e, + } +} + +/// MQ.POP — owner claims messages, routing max-delivery entries to the DLQ. +/// +/// Mirrors handler_sharded/write.rs MQ POP arm (lock-path `else` branch). +fn handle_pop(args: &[Frame], key_prefix: &Bytes, db_index: usize) -> Frame { + let (raw_key, count) = match validate_mq_pop(args) { + Ok(v) => v, + Err(e) => return e, + }; + let eff_key = effective_key(key_prefix, &raw_key); + let group_name = Bytes::from_static(b"__mq_consumers"); + let consumer_name = Bytes::from_static(b"__mq_default"); + + // All POP logic runs in a single with_shard_db closure to avoid re-entrancy. + crate::shard::slice::with_shard_db(db_index, |db| { + // Step 1: read max_delivery_count. + let mdc = match db.get_stream_mut(&eff_key) { + Ok(Some(stream)) => { + if !stream.durable { + return Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE)); + } + stream.max_delivery_count + } + Ok(None) => return Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE)), + Err(e) => return e, + }; + + let request_count = count + (mdc as usize); + + // Step 2: read_group_new to claim entries. + let stream = match db.get_stream_mut(&eff_key) { + Ok(Some(s)) => s, + _ => return Frame::Error(Bytes::from_static(ERR_MQ_NOT_DURABLE)), + }; + let claimed = + match stream.read_group_new(&group_name, &consumer_name, Some(request_count), false) { + Ok(entries) => entries, + Err(_) => return Frame::Array(vec![].into()), + }; + + // Step 3: partition claimed into good entries and DLQ entries. + let mut results: Vec<(StreamId, Vec<(Bytes, Bytes)>)> = + Vec::with_capacity(count.min(claimed.len())); + let mut dlq_entries: Vec<(StreamId, Vec<(Bytes, Bytes)>)> = Vec::new(); + let mut dlq_ack_ids: Vec = Vec::new(); + + for (id, fields) in &claimed { + let delivery_count = stream + .groups + .get(group_name.as_ref()) + .and_then(|g| g.pel.get(id)) + .map(|pe| pe.delivery_count) + .unwrap_or(1); + if mdc > 0 && delivery_count >= mdc as u64 { + dlq_entries.push((*id, fields.clone())); + dlq_ack_ids.push(*id); + } else if results.len() < count { + results.push((*id, fields.clone())); + } + } + + // Step 4: ACK DLQ entries from the main stream PEL. + if !dlq_ack_ids.is_empty() { + let _ = stream.xack(&group_name, &dlq_ack_ids); + } + + // Step 5: append DLQ entries to the sibling DLQ stream. + if !dlq_entries.is_empty() { + let dlq_key = { + let mut buf = Vec::with_capacity(eff_key.len() + 8); + buf.extend_from_slice(&eff_key); + buf.extend_from_slice(b"::mq:dlq"); + Bytes::from(buf) + }; + if let Ok(dlq_stream) = db.get_or_create_stream(&dlq_key) { + for (_id, fields) in dlq_entries { + let dlq_id = dlq_stream.next_auto_id(); + dlq_stream.add(dlq_id, fields); + } + } + } + + // Step 6: build response frames. + let result_frames: Vec = results + .iter() + .map(|(id, fields)| { + let mut entry_frames = Vec::with_capacity(2); + let mut ms_buf = itoa::Buffer::new(); + let mut seq_buf = itoa::Buffer::new(); + let ms_str = ms_buf.format(id.ms); + let seq_str = seq_buf.format(id.seq); + let mut id_bytes = Vec::with_capacity(ms_str.len() + 1 + seq_str.len()); + id_bytes.extend_from_slice(ms_str.as_bytes()); + id_bytes.push(b'-'); + id_bytes.extend_from_slice(seq_str.as_bytes()); + entry_frames.push(Frame::BulkString(Bytes::from(id_bytes))); + let field_frames: Vec = fields + .iter() + .flat_map(|(f, v)| [Frame::BulkString(f.clone()), Frame::BulkString(v.clone())]) + .collect(); + entry_frames.push(Frame::Array(field_frames.into())); + Frame::Array(entry_frames.into()) + }) + .collect(); + Frame::Array(result_frames.into()) + }) +} + +/// MQ.ACK — owner acknowledges one or more message IDs. +/// +/// Mirrors handler_sharded/write.rs MQ ACK arm. Emits a WAL record per acked +/// message, same as the lock-path. +fn handle_ack(args: &[Frame], key_prefix: &Bytes, db_index: usize) -> Frame { + let (raw_key, msg_ids) = match validate_mq_ack(args) { + Ok(v) => v, + Err(e) => return e, + }; + let eff_key = effective_key(key_prefix, &raw_key); + let ids: Vec = msg_ids + .iter() + .map(|(ms, seq)| StreamId { ms: *ms, seq: *seq }) + .collect(); + + let ack_result = + crate::shard::slice::with_shard_db(db_index, |db| match db.get_stream_mut(&eff_key) { + Ok(Some(stream)) => { + let group_name = Bytes::from_static(b"__mq_consumers"); + match stream.xack(&group_name, &ids) { + Ok(count) => Some(count), + Err(_) => None, + } + } + _ => None, + }); + + match ack_result { + Some(acked_count) => { + // Emit one WAL record per acked message id. + for (ms, seq) in &msg_ids { + let payload = crate::mq::wal::encode_mq_ack(&eff_key, *ms, *seq); + let mut wal_buf = Vec::new(); + crate::persistence::wal_v3::record::write_wal_v3_record( + &mut wal_buf, + 0, + crate::persistence::wal_v3::record::WalRecordType::MqAck, + &payload, + ); + wal_append_on_slice(Bytes::from(wal_buf)); + } + Frame::Integer(acked_count as i64) + } + None => Frame::Integer(0), + } +} + +/// MQ.DLQLEN — owner returns the depth of the dead-letter queue stream. +/// +/// Mirrors handler_sharded/write.rs MQ DLQLEN arm. +fn handle_dlqlen(args: &[Frame], key_prefix: &Bytes, db_index: usize) -> Frame { + let raw_key = match validate_mq_dlqlen(args) { + Ok(v) => v, + Err(e) => return e, + }; + let eff_key = effective_key(key_prefix, &raw_key); + let dlq_key = { + let mut buf = Vec::with_capacity(eff_key.len() + 8); + buf.extend_from_slice(&eff_key); + buf.extend_from_slice(b"::mq:dlq"); + Bytes::from(buf) + }; + + let len = + crate::shard::slice::with_shard_db(db_index, |db| match db.get_stream_mut(&dlq_key) { + Ok(Some(stream)) => stream.length as i64, + _ => 0i64, + }); + Frame::Integer(len) +} + +/// MQ.TRIGGER — owner registers a debounced trigger in the trigger registry. +/// +/// Mirrors handler_sharded/write.rs MQ TRIGGER arm. +fn handle_trigger(args: &[Frame], key_prefix: &Bytes, db_index: usize) -> Frame { + let (raw_key, callback_cmd, debounce_ms) = match validate_mq_trigger(args) { + Ok(v) => v, + Err(e) => return e, + }; + let eff_key = effective_key(key_prefix, &raw_key); + let trig_key = derive_trig_key(key_prefix, &raw_key); + + let entry = crate::mq::TriggerEntry { + queue_key: eff_key, + callback_cmd, + debounce_ms, + last_fire_ms: 0, + pending_fire_ms: 0, + }; + + crate::shard::slice::with_shard(|s| { + let reg = s + .trigger_registry + .get_or_insert_with(|| Box::new(crate::mq::TriggerRegistry::new())); + reg.register(trig_key, entry); + }); + + // `db_index` is unused for TRIGGER (registry only) but kept in the + // signature for API uniformity. Suppress the unused-variable warning. + let _ = db_index; + + Frame::SimpleString(Bytes::from_static(b"OK")) +} + +// ── Unit tests ──────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + + use super::*; + use crate::shard::shared_databases::ShardStoreMemory; + use crate::shard::slice::{ShardSlice, ShardSliceInit, init_shard}; + use crate::storage::Database; + use crate::text::store::TextStore; + use crate::transaction::{DeferredHnswInserts, KvWriteIntents}; + use crate::vector::store::VectorStore; + + fn make_test_slice(db_count: usize) -> ShardSlice { + let databases: Box<[Database]> = (0..db_count).map(|_| Database::new()).collect(); + ShardSlice::new(ShardSliceInit { + shard_id: 0, + databases, + vector_store: VectorStore::new(), + text_store: TextStore::new(), + #[cfg(feature = "graph")] + graph_store: crate::graph::store::GraphStore::new(), + kv_write_intents: KvWriteIntents::new(), + deferred_hnsw_inserts: DeferredHnswInserts::new(), + temporal_registry: None, + temporal_kv_index: None, + durable_queue_registry: None, + trigger_registry: None, + wal_append_tx: None, + estimated_memory: Arc::new(AtomicUsize::new(0)), + store_memory: Arc::new(ShardStoreMemory { + vector: AtomicUsize::new(0), + text: AtomicUsize::new(0), + graph: AtomicUsize::new(0), + }), + }) + } + + // ── effective_key derivation ────────────────────────────────────────────── + + #[test] + fn effective_key_without_prefix() { + let prefix = Bytes::new(); + let raw = Bytes::from_static(b"myqueue"); + let eff = effective_key(&prefix, &raw); + assert_eq!(eff.as_ref(), b"myqueue"); + } + + #[test] + fn effective_key_with_prefix() { + // key_prefix = "{" + 32 hex chars + "}:" + let ws_hex = "0102030405060708090a0b0c0d0e0f10"; + let prefix_str = format!("{{{ws_hex}}}:"); + let prefix = Bytes::from(prefix_str.clone()); + let raw = Bytes::from_static(b"tasks"); + let eff = effective_key(&prefix, &raw); + let expected = format!("{prefix_str}tasks"); + assert_eq!(eff.as_ref(), expected.as_bytes()); + } + + // ── derive_trig_key ─────────────────────────────────────────────────────── + + #[test] + fn trig_key_without_prefix() { + let prefix = Bytes::new(); + let raw = Bytes::from_static(b"alerts"); + let trig = derive_trig_key(&prefix, &raw); + assert_eq!(trig.as_ref(), b"alerts"); + } + + #[test] + fn trig_key_with_prefix() { + // prefix = "{0102...10}:" (35 bytes) + let ws_hex = "0102030405060708090a0b0c0d0e0f10"; + let prefix = Bytes::from(format!("{{{ws_hex}}}:")); + let raw = Bytes::from_static(b"alerts"); + let trig = derive_trig_key(&prefix, &raw); + // expected: ws_hex + ":" + "alerts" + let expected = format!("{ws_hex}:alerts"); + assert_eq!(trig.as_ref(), expected.as_bytes()); + } + + // ── DLQLEN on empty queue ───────────────────────────────────────────────── + + #[test] + fn dlqlen_empty_queue_returns_zero() { + // Use a fresh OS thread so init_shard doesn't conflict with the test thread. + let result = std::thread::spawn(|| { + init_shard(make_test_slice(1)); + + // Build a fake MQ DLQLEN command frame. + let cmd = Arc::new(Frame::Array( + vec![ + Frame::BulkString(Bytes::from_static(b"MQ")), + Frame::BulkString(Bytes::from_static(b"DLQLEN")), + Frame::BulkString(Bytes::from_static(b"nosuchqueue")), + ] + .into(), + )); + execute_mq_on_owner(0, Bytes::new(), cmd) + }) + .join() + .expect("test thread panicked"); + + assert_eq!(result, Frame::Integer(0)); + } + + // ── MQ.CREATE + DLQLEN round-trip ───────────────────────────────────────── + + #[test] + fn create_then_dlqlen_zero() { + let result = std::thread::spawn(|| { + init_shard(make_test_slice(1)); + + // CREATE + let create_cmd = Arc::new(Frame::Array( + vec![ + Frame::BulkString(Bytes::from_static(b"MQ")), + Frame::BulkString(Bytes::from_static(b"CREATE")), + Frame::BulkString(Bytes::from_static(b"testq")), + ] + .into(), + )); + let create_result = execute_mq_on_owner(0, Bytes::new(), create_cmd); + assert_eq!( + create_result, + Frame::SimpleString(Bytes::from_static(b"OK")), + "CREATE must return +OK" + ); + + // DLQLEN on newly created queue — DLQ stream doesn't exist yet. + let dlqlen_cmd = Arc::new(Frame::Array( + vec![ + Frame::BulkString(Bytes::from_static(b"MQ")), + Frame::BulkString(Bytes::from_static(b"DLQLEN")), + Frame::BulkString(Bytes::from_static(b"testq")), + ] + .into(), + )); + execute_mq_on_owner(0, Bytes::new(), dlqlen_cmd) + }) + .join() + .expect("test thread panicked"); + + assert_eq!(result, Frame::Integer(0)); + } +} diff --git a/src/shard/persistence_tick.rs b/src/shard/persistence_tick.rs index d9a898a0..a435c379 100644 --- a/src/shard/persistence_tick.rs +++ b/src/shard/persistence_tick.rs @@ -41,7 +41,15 @@ pub(crate) fn handle_pending_snapshot( } else { snap_dir.join(format!("shard-{}.rrdshard", shard_id)) }; - let (segment_counts, base_timestamps) = shard_databases.snapshot_metadata(shard_id); + let (segment_counts, base_timestamps) = crate::shard::slice::with_shard(|s| { + let mut seg_counts = Vec::with_capacity(s.databases.len()); + let mut base_ts = Vec::with_capacity(s.databases.len()); + for db in s.databases.iter() { + seg_counts.push(db.data().segment_count()); + base_ts.push(db.base_timestamp()); + } + (seg_counts, base_ts) + }); let db_count = shard_databases.db_count(); let mut state = SnapshotState::new_from_metadata( shard_id as u16, @@ -89,7 +97,15 @@ pub(crate) fn check_auto_save_trigger( } else { std::path::PathBuf::from(dir).join(format!("shard-{}.rrdshard", shard_id)) }; - let (segment_counts, base_timestamps) = shard_databases.snapshot_metadata(shard_id); + let (segment_counts, base_timestamps) = crate::shard::slice::with_shard(|s| { + let mut seg_counts = Vec::with_capacity(s.databases.len()); + let mut base_ts = Vec::with_capacity(s.databases.len()); + for db in s.databases.iter() { + seg_counts.push(db.data().segment_count()); + base_ts.push(db.base_timestamp()); + } + (seg_counts, base_ts) + }); let db_count = shard_databases.db_count(); let mut state = SnapshotState::new_from_metadata( shard_id as u16, @@ -116,17 +132,12 @@ pub(crate) fn advance_snapshot_segment( shard_databases: &Arc, shard_id: usize, ) -> bool { + let _ = shard_id; // E2 removes if let Some(snap) = snapshot_state { let current_db = snap.current_db_index(); let db_count = shard_databases.db_count(); if current_db < db_count { - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(current_db, |db| snap.advance_one_segment_db(db)) - } else { - let guard = shard_databases.read_db(shard_id, current_db); - snap.advance_one_segment_db(&guard) - } + crate::shard::slice::with_shard_db(current_db, |db| snap.advance_one_segment_db(db)) } else { // All databases serialized, return true to trigger finalization true @@ -341,12 +352,39 @@ pub(crate) fn run_eviction_tick( { let rt = runtime_config.read(); if rt.maxmemory > 0 { - let used = shard_databases.aggregate_memory(shard_id); + // C5 / Phase 3: compute per-shard KV memory via ShardSlice without + // lock acquisitions (avoids per-DB read locks). + let used = crate::shard::slice::with_shard(|s| { + s.databases + .iter() + .map(|db| db.estimated_memory()) + .sum::() + }); shard_databases.publish_memory(shard_id, used); shard_databases.recompute_elastic_budget(shard_id, &rt); } } + // C5 / M4: publish vector/text/graph store memory for lock-free observers. + // Uses the existing lock path (Wave E collapses to slice). Runs every tick + // so Prometheus and MEMORY DOCTOR never see stale zero values for long. + // C5 / M4: publish vector/text/graph store-memory atomics via thread-local slice. + crate::shard::slice::with_shard(|s| { + use std::sync::atomic::Ordering; + let (mutable, immutable) = s.vector_store.resident_bytes(); + s.store_memory + .vector + .store(mutable + immutable, Ordering::Relaxed); + s.store_memory.text.store(0, Ordering::Relaxed); // TextStore has no aggregate API yet + #[cfg(feature = "graph")] + { + let graph_bytes = s.graph_store.resident_bytes(); + s.store_memory.graph.store(graph_bytes, Ordering::Relaxed); + } + #[cfg(not(feature = "graph"))] + s.store_memory.graph.store(0, Ordering::Relaxed); + }); + if server_config.disk_offload_enabled() && should_run_pressure_cascade(runtime_config, server_config, shard_databases, shard_id) { @@ -386,7 +424,7 @@ pub(crate) fn drain_and_shutdown_spill( // flush that the drain above did not see; apply them so those cold keys // are not lost (file on disk but never recorded in the manifest). let leftover = st.shutdown(); - apply_completion_vec(leftover, shard_manifest, shard_databases, shard_id); + apply_completion_vec(leftover, shard_manifest); tracing::info!("Shard {}: spill background thread shut down", shard_id); } } @@ -403,8 +441,10 @@ pub(crate) fn apply_spill_completions( shard_databases: &std::sync::Arc, shard_id: usize, ) { + let _ = shard_databases; // E2 removes + let _ = shard_id; // E2 removes let completions = spill_thread.drain_completions(); - apply_completion_vec(completions, shard_manifest, shard_databases, shard_id); + apply_completion_vec(completions, shard_manifest); } /// Apply a batch of spill completions: ONE manifest `add_file`+commit per file, @@ -413,8 +453,6 @@ pub(crate) fn apply_spill_completions( fn apply_completion_vec( completions: Vec, shard_manifest: &mut Option, - shard_databases: &std::sync::Arc, - shard_id: usize, ) { if completions.is_empty() { return; @@ -451,19 +489,11 @@ fn apply_completion_vec( slot_idx: entry.slot_idx, }; - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(entry.db_index, |db| { - if let Some(ref mut ci) = db.cold_index { - ci.insert(entry.key.clone(), location); - } - }); - } else { - let mut guard = shard_databases.write_db(shard_id, entry.db_index); - if let Some(ref mut ci) = guard.cold_index { - ci.insert(entry.key, location); + crate::shard::slice::with_shard_db(entry.db_index, |db| { + if let Some(ref mut ci) = db.cold_index { + ci.insert(entry.key.clone(), location); } - } + }); } } } @@ -496,7 +526,9 @@ pub(crate) fn should_run_pressure_cascade( elastic => elastic.min(rt.maxmemory), }; let threshold = (budget as f64 * server_config.disk_offload_threshold) as usize; - let used = shard_databases.aggregate_memory(shard_id); + // C5 / Phase 3: read the already-published per-shard KV memory (written + // earlier this same tick by `run_eviction_tick`). Lock-free Relaxed load. + let used = shard_databases.published_shard_memory(shard_id); used > threshold } @@ -542,27 +574,15 @@ pub(crate) fn handle_memory_pressure( let shard_dir = server_config .effective_disk_offload_dir() .join(format!("shard-{}", shard_id)); - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - let count = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - s.vector_store.try_warm_transitions_all( - &shard_dir, - manifest, - aggressive_threshold, - next_file_id, - wal_v3, - ) - }) - } else { - let vs = shard_databases.vector_store(shard_id); - vs.try_warm_transitions_all( + let count = crate::shard::slice::with_shard(|s| { + s.vector_store.try_warm_transitions_all( &shard_dir, manifest, aggressive_threshold, next_file_id, wal_v3, ) - }; + }); if count > 0 { tracing::info!( "Shard {}: memory pressure step 2 -- force-demoted {} segment(s) HOT->WARM", @@ -585,8 +605,9 @@ pub(crate) fn handle_memory_pressure( { let rt = runtime_config.read(); if rt.maxmemory > 0 { - // Compute aggregate BEFORE acquiring write locks (same pattern as handler_sharded). - let total_mem = shard_databases.aggregate_memory(shard_id); + // C5 / Phase 3: read the already-published per-shard KV memory + // (written earlier this same tick). Lock-free Relaxed load. + let total_mem = shard_databases.published_shard_memory(shard_id); // GAP-1: hot shards evict against their elastic budget (idle // siblings' donated headroom), not the static maxmemory/N. let budget = match shard_databases.elastic_budget(shard_id) { @@ -599,29 +620,13 @@ pub(crate) fn handle_memory_pressure( .effective_disk_offload_dir() .join(format!("shard-{}", shard_id)); - // Phase 2d: gate on is_initialized(); new path uses ShardSlice directly. - let use_slice = crate::shard::slice::is_initialized(); if let Some(spill_t) = spill_thread { // Async spill path: background thread does pwrite let sender = spill_t.sender(); for i in 0..db_count { - if use_slice { - crate::shard::slice::with_shard_db(i, |db| { - let _ = crate::storage::eviction::try_evict_if_needed_async_spill_with_total_budget( - db, - &rt, - &sender, - &shard_dir, - next_file_id, - total_mem, - i, - budget, - ); - }); - } else { - let mut guard = shard_databases.write_db(shard_id, i); + crate::shard::slice::with_shard_db(i, |db| { let _ = crate::storage::eviction::try_evict_if_needed_async_spill_with_total_budget( - &mut guard, + db, &rt, &sender, &shard_dir, @@ -630,36 +635,14 @@ pub(crate) fn handle_memory_pressure( i, budget, ); - } + }); } // Drop sender clone immediately to avoid shutdown deadlock drop(sender); } else { // Sync spill fallback for i in 0..db_count { - if use_slice { - crate::shard::slice::with_shard_db(i, |db| { - if let Some(ref mut manifest) = *shard_manifest { - let mut ctx = crate::storage::eviction::SpillContext { - shard_dir: &shard_dir, - manifest, - next_file_id, - }; - let _ = crate::storage::eviction::try_evict_if_needed_with_spill_and_total_budget( - db, - &rt, - Some(&mut ctx), - total_mem, - budget, - ); - } else { - let _ = crate::storage::eviction::try_evict_if_needed_with_spill_and_total_budget( - db, &rt, None, total_mem, budget, - ); - } - }); - } else { - let mut guard = shard_databases.write_db(shard_id, i); + crate::shard::slice::with_shard_db(i, |db| { if let Some(ref mut manifest) = *shard_manifest { let mut ctx = crate::storage::eviction::SpillContext { shard_dir: &shard_dir, @@ -667,7 +650,7 @@ pub(crate) fn handle_memory_pressure( next_file_id, }; let _ = crate::storage::eviction::try_evict_if_needed_with_spill_and_total_budget( - &mut guard, + db, &rt, Some(&mut ctx), total_mem, @@ -675,10 +658,10 @@ pub(crate) fn handle_memory_pressure( ); } else { let _ = crate::storage::eviction::try_evict_if_needed_with_spill_and_total_budget( - &mut guard, &rt, None, total_mem, budget, + db, &rt, None, total_mem, budget, ); } - } + }); } } } @@ -1311,7 +1294,10 @@ mod tests { let tmp = tempfile::tempdir().unwrap(); let snap_dir = tmp.path().to_path_buf(); let dbs = vec![vec![Database::new()]]; - let shared = ShardDatabases::new(dbs); + let (shared, mut inits) = ShardDatabases::new(dbs); + crate::shard::slice::reset_test_shard(crate::shard::slice::ShardSlice::new( + inits.remove(0), + )); let (tx, _rx) = channel::oneshot::>(); let mut snapshot_state: Option = None; @@ -1338,7 +1324,10 @@ mod tests { fn test_handle_pending_snapshot_zero_lsn_is_unknown() { let tmp = tempfile::tempdir().unwrap(); let dbs = vec![vec![Database::new()]]; - let shared = ShardDatabases::new(dbs); + let (shared, mut inits) = ShardDatabases::new(dbs); + crate::shard::slice::reset_test_shard(crate::shard::slice::ShardSlice::new( + inits.remove(0), + )); let (tx, _rx) = channel::oneshot::>(); let mut snapshot_state: Option = None; @@ -1363,7 +1352,10 @@ mod tests { let tmp = tempfile::tempdir().unwrap(); let dir = tmp.path().to_string_lossy().to_string(); let dbs = vec![vec![Database::new()]]; - let shared = ShardDatabases::new(dbs); + let (shared, mut inits) = ShardDatabases::new(dbs); + crate::shard::slice::reset_test_shard(crate::shard::slice::ShardSlice::new( + inits.remove(0), + )); // Trigger goes from epoch 0 → 5; helper observes 5 > last(0) and // creates a snapshot state. diff --git a/src/shard/scatter_aggregate.rs b/src/shard/scatter_aggregate.rs index bbbf933e..933e647a 100644 --- a/src/shard/scatter_aggregate.rs +++ b/src/shard/scatter_aggregate.rs @@ -50,43 +50,26 @@ pub async fn scatter_text_aggregate( dispatch_tx: &Rc>>>, spsc_notifiers: &[Arc], ) -> Frame { + let _ = shard_databases; // E2 removes // ─── Single-shard fast path ─────────────────────────────────────────── if num_shards == 1 { - // Phase 2e: gate on is_initialized(); new path holds vector_store + - // text_store + databases[0] simultaneously via a single `with_shard` - // call to avoid a reentrant `with_shard*` panic (RefCell double-borrow). - let result = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - // Every shard slice is constructed with at least one database - // (db 0); see Shard::new in src/shard/mod.rs. - #[allow(clippy::expect_used)] - let db = s - .databases - .first() - .expect("shard slice must have at least db 0"); - crate::command::vector_search::ft_aggregate::execute_local_full( - &mut s.vector_store, - &s.text_store, - &index_name, - &query, - &pipeline, - db, - ) - }) - } else { - let mut vs = shard_databases.vector_store(my_shard); - let ts = shard_databases.text_store(my_shard); - let db_guard = shard_databases.read_db(my_shard, 0); + let result = crate::shard::slice::with_shard(|s| { + // Every shard slice is constructed with at least one database + // (db 0); see Shard::new in src/shard/mod.rs. + #[allow(clippy::expect_used)] + let db = s + .databases + .first() + .expect("shard slice must have at least db 0"); crate::command::vector_search::ft_aggregate::execute_local_full( - &mut vs, - &ts, + &mut s.vector_store, + &s.text_store, &index_name, &query, &pipeline, - &db_guard, + db, ) - // guards drop at end of this block, before any .await below - }; + }); return result; } @@ -101,37 +84,21 @@ pub async fn scatter_text_aggregate( // round trip (CONTEXT.md D-08 follows scatter_text_search // template). // - // Phase 2e: gate on is_initialized(); new path holds text_store - // + databases[0] together via a single `with_shard` closure to - // avoid a reentrant `with_shard*` panic. - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - // See note on previous expect — db 0 is constructor-guaranteed. - #[allow(clippy::expect_used)] - let db = s - .databases - .first() - .expect("shard slice must have at least db 0"); - crate::command::vector_search::ft_aggregate::execute_local_partial( - &s.text_store, - &index_name, - &query, - &pipeline, - db, - ) - }) - } else { - let ts = shard_databases.text_store(shard_id); - let db_guard = shard_databases.read_db(shard_id, 0); + let response = crate::shard::slice::with_shard(|s| { + // See note on previous expect — db 0 is constructor-guaranteed. + #[allow(clippy::expect_used)] + let db = s + .databases + .first() + .expect("shard slice must have at least db 0"); crate::command::vector_search::ft_aggregate::execute_local_partial( - &ts, + &s.text_store, &index_name, &query, &pipeline, - &db_guard, + db, ) - // guards drop before the subsequent .await - }; + }); local_partial = Some(response); } else { let (reply_tx, reply_rx) = channel::oneshot(); @@ -259,7 +226,7 @@ mod tests { // SPSC dispatch would panic on index access, proving the fast path // was taken. Empty TextStore yields "ERR unknown index". let dbs = vec![Database::new()]; - let shard_databases = ShardDatabases::new(vec![dbs]); + let (shard_databases, _inits) = ShardDatabases::new(vec![dbs]); let dispatch_tx: Rc>>> = Rc::new(RefCell::new(Vec::new())); diff --git a/src/shard/scatter_hybrid.rs b/src/shard/scatter_hybrid.rs index 1d85b517..f5f6af05 100644 --- a/src/shard/scatter_hybrid.rs +++ b/src/shard/scatter_hybrid.rs @@ -63,6 +63,7 @@ pub async fn scatter_hybrid_search( dispatch_tx: &Rc>>>, spsc_notifiers: &[Arc], ) -> Frame { + let _ = shard_databases; // E2 removes let top_k = query.top_k; let k_per_stream = query.effective_k_per_stream(); @@ -73,26 +74,14 @@ pub async fn scatter_hybrid_search( // Phase 171 HYB-02 / SCAT-02: thread coordinator-resolved AS_OF LSN // through the single-shard fast path so callers routed here still // honor temporal snapshots. - // Phase 2e: gate on is_initialized(); new path acquires vector_store - // and text_store via a single `with_shard` closure to avoid reentrant - // RefCell double-borrow panics. - return if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - crate::command::vector_search::hybrid::execute_hybrid_search_local( - &mut s.vector_store, - &s.text_store, - &query, - as_of_lsn, - ) - }) - } else { - let mut vs = shard_databases.vector_store(my_shard); - let ts = shard_databases.text_store(my_shard); + return crate::shard::slice::with_shard(|s| { crate::command::vector_search::hybrid::execute_hybrid_search_local( - &mut vs, &ts, &query, as_of_lsn, + &mut s.vector_store, + &s.text_store, + &query, + as_of_lsn, ) - // guards drop at end of this block - }; + }); } // ─── Parse text query into QueryTerms using the index's own analyzer ─── @@ -105,38 +94,10 @@ pub async fn scatter_hybrid_search( ), Frame, >; - // Phase 2e: gate on is_initialized(); inner closure parses against the - // owned text_store. Both branches return a fully-owned ParseOutcome so + // Both branches return a fully-owned ParseOutcome so // no borrows of the shard slice escape. - let parse_outcome: ParseOutcome = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - match s.text_store.get_index(query.index_name.as_ref()) { - None => Err(Frame::Error(Bytes::from_static(b"ERR unknown index"))), - Some(text_index) => match text_index.field_analyzers.first() { - None => Err(Frame::Error(Bytes::from_static( - b"ERR index has no TEXT fields", - ))), - Some(analyzer) => { - match crate::command::vector_search::parse_text_query(text_bytes, analyzer) - { - Ok(clause) => { - let strings: Vec = - clause.terms.iter().map(|qt| qt.text.clone()).collect(); - Ok((clause.terms, strings)) - } - Err(e) => { - let mut msg = b"ERR ".to_vec(); - msg.extend_from_slice(e.as_bytes()); - Err(Frame::Error(Bytes::from(msg))) - } - } - } - }, - } - }) - } else { - let ts = shard_databases.text_store(my_shard); - match ts.get_index(query.index_name.as_ref()) { + let parse_outcome: ParseOutcome = crate::shard::slice::with_shard(|s| { + match s.text_store.get_index(query.index_name.as_ref()) { None => Err(Frame::Error(Bytes::from_static(b"ERR unknown index"))), Some(text_index) => match text_index.field_analyzers.first() { None => Err(Frame::Error(Bytes::from_static( @@ -158,7 +119,7 @@ pub async fn scatter_hybrid_search( } }, } - }; // MutexGuard dropped here — no .await held + }); // shard slice released here — borrow released before first async op let (query_terms, term_strings) = match parse_outcome { Ok(v) => v, Err(f) => return f, @@ -174,33 +135,9 @@ pub async fn scatter_hybrid_search( if shard_id == my_shard { // Local DFS — direct read, no SPSC overhead. Drop guard before .await. // - // Phase 2e: gate on is_initialized(); both branches build the - // same owned Frame result. The new path borrows `text_store` - // exclusively from the slice — no other companion store is - // accessed in this branch, so a one-store `with_shard` suffices. - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - match s.text_store.get_index(query.index_name.as_ref()) { - Some(text_index) => { - let mut items: Vec = Vec::new(); - for (field_idx_opt, terms) in &field_queries { - let fidx = field_idx_opt.unwrap_or(0); - let (term_dfs, n) = text_index.doc_freq_for_terms(fidx, terms); - for (term, df) in term_dfs { - items.push(Frame::BulkString(Bytes::from(term))); - items.push(Frame::Integer(i64::from(df))); - } - items.push(Frame::BulkString(Bytes::from_static(b"N"))); - items.push(Frame::Integer(i64::from(n))); - } - Frame::Array(FrameVec::from(items)) - } - None => Frame::Error(Bytes::from_static(b"ERR unknown index")), - } - }) - } else { - let ts = shard_databases.text_store(shard_id); - match ts.get_index(query.index_name.as_ref()) { + // Local DFS via slice — text_store accessed exclusively in the closure. + let response = crate::shard::slice::with_shard(|s| { + match s.text_store.get_index(query.index_name.as_ref()) { Some(text_index) => { let mut items: Vec = Vec::new(); for (field_idx_opt, terms) in &field_queries { @@ -217,7 +154,7 @@ pub async fn scatter_hybrid_search( } None => Frame::Error(Bytes::from_static(b"ERR unknown index")), } - }; // text_guard drops here + }); // shard slice released here local_dfs = Some(response); } else { let (reply_tx, reply_rx) = channel::oneshot(); @@ -256,43 +193,20 @@ pub async fn scatter_hybrid_search( for shard_id in 0..num_shards { if shard_id == my_shard { - // Phase 2e: gate on is_initialized(); new path holds vector_store - // + text_store through ONE `with_shard` closure to avoid reentrant - // RefCell double-borrow panics. + // vector_store + text_store acquired in one with_shard closure to + // avoid reentrant RefCell double-borrow panics. let sparse_pair = match (sparse_field.as_ref(), sparse_blob.as_ref()) { (Some(f), Some(b)) => Some((f, b)), _ => None, }; - let response = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - // Phase 171 HYB-02 / SCAT-02: forward as_of_lsn through the - // local raw-streams executor (symmetric with remote payload). - crate::command::vector_search::hybrid_multi::execute_hybrid_search_local_raw_streams( - &mut s.vector_store, - &s.text_store, - &query.index_name, - &query_terms, - &query.dense_field, - &query.dense_blob, - sparse_pair, - query.weights, - k_per_stream, - top_k, - &global_df, - global_n, - as_of_lsn, - query.filter.as_ref(), - ) - }) - } else { - let mut vs = shard_databases.vector_store(shard_id); - let ts = shard_databases.text_store(shard_id); - // Phase 171 HYB-02 / SCAT-02: forward as_of_lsn to the local - // raw-streams executor so the coordinator's own shard honors - // AS_OF on the dense branch (symmetric with remote payload below). + // Phase 171 HYB-02 / SCAT-02: forward as_of_lsn through the + // local raw-streams executor (symmetric with remote payload). + // CHANGE F: clone filter for local path (shard-local allowlist eval). + let filter_clone = query.filter.clone(); + let response = crate::shard::slice::with_shard(|s| { crate::command::vector_search::hybrid_multi::execute_hybrid_search_local_raw_streams( - &mut vs, - &ts, + &mut s.vector_store, + &s.text_store, &query.index_name, &query_terms, &query.dense_field, @@ -304,10 +218,9 @@ pub async fn scatter_hybrid_search( &global_df, global_n, as_of_lsn, - query.filter.as_ref(), + filter_clone.as_ref(), ) - // guards drop here - }; + }); local_hyb = Some(response); } else { let (reply_tx, reply_rx) = channel::oneshot(); @@ -324,6 +237,7 @@ pub async fn scatter_hybrid_search( global_df: global_df.clone(), global_n, as_of_lsn, + // CHANGE F: forward filter to remote shard payload. filter: query.filter.clone(), reply_tx, }; diff --git a/src/shard/shared_databases.rs b/src/shard/shared_databases.rs index 9caf6327..b85d8655 100644 --- a/src/shard/shared_databases.rs +++ b/src/shard/shared_databases.rs @@ -1,150 +1,136 @@ -use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use parking_lot::{Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard}; +use parking_lot::{Mutex, MutexGuard}; -#[cfg(feature = "graph")] -use crate::graph::store::GraphStore; -use crate::mq::{DurableQueueRegistry, TriggerRegistry}; use crate::storage::Database; -use crate::temporal::{TemporalKvIndex, TemporalRegistry}; -use crate::text::store::TextStore; -use crate::transaction::{DeferredHnswInserts, KvWriteIntents}; -use crate::vector::store::VectorStore; use crate::workspace::wal::{decode_workspace_create, decode_workspace_drop}; use crate::workspace::{WorkspaceId, WorkspaceMetadata, WorkspaceRegistry}; -/// Thread-safe wrapper over per-shard databases. +/// Published per-shard store-memory counters (C5 / M4). +/// +/// Each shard's 100ms tick writes vector/text/graph resident bytes into these +/// atomics. Cross-thread observers — Prometheus publisher (metrics_setup.rs) +/// and MEMORY DOCTOR (server_admin.rs) — read from these atomics with zero +/// lock acquisitions. +/// +/// Figures lag at most one tick (≤100 ms); this is documented and acceptable +/// for observability paths. +pub struct ShardStoreMemory { + /// Combined resident bytes of all VectorStore segments (mutable + immutable). + pub vector: AtomicUsize, + /// Resident bytes of TextStore indexes. + pub text: AtomicUsize, + /// Resident bytes of GraphStore CSR segments. + pub graph: AtomicUsize, +} + +/// Shared infrastructure handle — the residual cross-shard state after M5. +/// +/// Contains ONLY genuinely-shared handles. Per-shard data (databases, stores, +/// registries) was moved into `ShardSliceInit` packages returned by `new()`. +/// No per-shard `RwLock`/`Mutex` wrappers live here. /// -/// Each shard owns `db_count` databases (SELECT 0-15). The outer Vec is indexed -/// by shard_id, the inner Vec by db_index. All access goes through `read_db()` -/// (shared) or `write_db()` (exclusive) to enable cross-shard direct reads. +/// # Fields kept (C6) +/// - `wal_append_txs`: WAL append senders — `Send` by design; cross-shard +/// `wal_append(owner, …)` stays legal under slice with no message variant. +/// - `num_shards` / `db_count`: immutable configuration. +/// - `memory_per_shard` / `elastic_budgets`: published-atomic patterns. +/// - `workspace_registry`: single process-global registry (C3 / M3). +/// - `store_memory_per_shard`: published store-memory atomics (C5 / M4). pub struct ShardDatabases { - shards: Vec>>, - /// Per-shard VectorStore for FT.* commands in single-shard mode. - vector_stores: Vec>, - /// Per-shard TextStore for full-text search indexes. - text_stores: Vec>, - /// Per-shard GraphStore for GRAPH.* commands. - #[cfg(feature = "graph")] - graph_stores: Vec>, /// Per-shard WAL append channel sender. Connection handlers send serialized /// write commands here; the event loop drains into WAL v2/v3 on the 1ms tick. /// OnceLock: set once at event-loop startup (before connections are - /// accepted), then every hot-path read is lock-free (QW2, 2026-06 review - /// finding 1.8 — was Mutex>, one lock acquire per write command). + /// accepted), then every hot-path read is lock-free. wal_append_txs: Vec>>, - /// Per-shard TemporalRegistry for wall-clock-to-LSN bindings. - /// Lazy-init: None until first TEMPORAL.SNAPSHOT_AT call. - temporal_registries: Vec>>>, - /// Per-shard TemporalKvIndex for versioned KV reads. - /// Lazy-init: None until first TemporalUpsert WAL write. - temporal_kv_indexes: Vec>>>, - /// Per-shard WorkspaceRegistry for workspace metadata. - /// Lazy-init: None until first WS.CREATE call on this shard. - workspace_registries: Vec>>>, - /// Per-shard DurableQueueRegistry for MQ.* commands. - /// Lazy-init: None until first MQ.CREATE call on this shard. - durable_queue_registries: Vec>>>, - /// Per-shard TriggerRegistry for MQ.TRIGGER debounced callbacks. - /// Lazy-init: None until first MQ.TRIGGER call on this shard. - trigger_registries: Vec>>>, - /// Per-shard KV write-intent side-table for transactional MVCC. - /// Shard-global: visible to all connections for cross-txn visibility checks. - kv_write_intents: Vec>, - /// Per-shard deferred HNSW insert queue for post-commit processing. - /// Shard-global: survives connection drops before commit. - deferred_hnsw_inserts: Vec>, + /// Process-global WorkspaceRegistry (C3 / M3). + /// + /// Workspaces are control-plane objects looked up by every connection + /// regardless of which shard accepted it — a single Mutex is not a + /// hot-path concern (WS commands are rare). The per-shard array + /// (`workspace_registries`) is retired; all paths use this one field. + /// WAL records keep the shard-0 stream via `wal_append(0, …)` (unchanged). + /// Caller lazy-inits via `get_or_insert_with(|| Box::new(WorkspaceRegistry::new()))`. + workspace_registry: Mutex>>, num_shards: usize, db_count: usize, /// Per-shard memory publishers for lock-free cross-shard reads. /// - /// Each shard's event loop will call `memory_publisher(shard_id)` once at - /// startup to clone an `Arc` into `ShardSlice.estimated_memory`. The shard - /// then atomically publishes its estimated memory after write operations - /// (Phase 2). Cross-shard readers — maxmemory eviction - /// (`persistence_tick.rs:436,511`) and metrics scrape - /// (`metrics_setup.rs:1238`) — call `read_memory_sum()` for a lock-free - /// sum (Phase 3). Coexists with the existing `aggregate_memory()` path - /// until Phase 3 switches the eviction tick. + /// Each shard's event loop holds a clone of its `Arc` in + /// `ShardSlice.estimated_memory`. Cross-shard readers — maxmemory eviction + /// and metrics scrape — call `read_memory_sum()` for a lock-free sum. memory_per_shard: Vec>, /// Per-shard elastic memory budgets (GAP-1 hot-shard pooling). - /// - /// Recomputed by each shard's 100ms eviction tick from the - /// `memory_per_shard` snapshot (`recompute_elastic_budget`). `0` means - /// "no elastic budget yet" — readers fall back to the static - /// `maxmemory / num_shards`. Write paths read with one Relaxed load - /// (`elastic_budget`). elastic_budgets: Vec>, + /// Per-shard published store-memory atomics (C5 / M4). + /// + /// The shard's 100ms tick refreshes these via `with_shard`. Prometheus + /// publisher and MEMORY DOCTOR read them with zero lock acquisitions. + pub store_memory_per_shard: Box<[Arc]>, } impl ShardDatabases { /// Create from pre-restored database vectors (one `Vec` per shard). - pub fn new(shard_databases: Vec>) -> Arc { + /// + /// Returns the shared infrastructure handle AND one `ShardSliceInit` + /// package per shard. Each package is handed to its shard thread's closure + /// by move; the shard calls `slice::init_shard(init)` at startup. + /// + /// WAL/AOF recovery MUST happen BEFORE this call (on the raw per-shard + /// state via `recover_*` free functions in this module) so the databases + /// passed here are already fully restored. + pub fn new( + shard_databases: Vec>, + ) -> (Arc, Vec) { let num_shards = shard_databases.len(); let db_count = shard_databases.first().map_or(0, |v| v.len()); - let shards = shard_databases - .into_iter() - .map(|dbs| dbs.into_iter().map(RwLock::new).collect()) - .collect(); - let vector_stores = (0..num_shards) - .map(|_| Mutex::new(VectorStore::new())) - .collect(); - let text_stores = (0..num_shards) - .map(|_| Mutex::new(TextStore::new())) - .collect(); - #[cfg(feature = "graph")] - let graph_stores = (0..num_shards) - .map(|_| RwLock::new(GraphStore::new())) - .collect(); + let wal_append_txs = (0..num_shards) .map(|_| std::sync::OnceLock::new()) .collect(); - let temporal_registries = (0..num_shards).map(|_| Mutex::new(None)).collect(); - let temporal_kv_indexes = (0..num_shards).map(|_| Mutex::new(None)).collect(); - let workspace_registries = (0..num_shards).map(|_| Mutex::new(None)).collect(); - let durable_queue_registries = (0..num_shards).map(|_| Mutex::new(None)).collect(); - let trigger_registries = (0..num_shards).map(|_| Mutex::new(None)).collect(); - let kv_write_intents = (0..num_shards) - .map(|_| Mutex::new(KvWriteIntents::new())) - .collect(); - let deferred_hnsw_inserts = (0..num_shards) - .map(|_| Mutex::new(DeferredHnswInserts::new())) - .collect(); - let memory_per_shard = (0..num_shards) + let workspace_registry = Mutex::new(None); + let memory_per_shard: Vec> = (0..num_shards) .map(|_| Arc::new(AtomicUsize::new(0))) .collect(); - let elastic_budgets = (0..num_shards) + let elastic_budgets: Vec> = (0..num_shards) .map(|_| Arc::new(AtomicUsize::new(0))) .collect(); - Arc::new(Self { - shards, - vector_stores, - text_stores, - #[cfg(feature = "graph")] - graph_stores, + let store_memory_per_shard: Box<[Arc]> = (0..num_shards) + .map(|_| { + Arc::new(ShardStoreMemory { + vector: AtomicUsize::new(0), + text: AtomicUsize::new(0), + graph: AtomicUsize::new(0), + }) + }) + .collect::>() + .into_boxed_slice(); + + let shared = Arc::new(Self { wal_append_txs, - temporal_registries, - temporal_kv_indexes, - workspace_registries, - durable_queue_registries, - trigger_registries, - kv_write_intents, - deferred_hnsw_inserts, + workspace_registry, num_shards, db_count, - memory_per_shard, - elastic_budgets, - }) + memory_per_shard: memory_per_shard.clone(), + elastic_budgets: elastic_budgets.clone(), + store_memory_per_shard: store_memory_per_shard.clone(), + }); + + // Build one ShardSliceInit per shard, consuming the databases. The + // construction lives in slice.rs with the type it builds (the only + // per-shard state this module touches is the moment of handoff). + let inits = crate::shard::slice::ShardSliceInit::build_all( + shard_databases, + &memory_per_shard, + &store_memory_per_shard, + ); + + (shared, inits) } /// Set the WAL append channel sender for a shard. /// - /// Called once during event loop startup. Uses interior mutability via - /// unsafe transmutation of the Arc — safe because this is called exactly - /// once per shard before any connections are accepted. - /// Set the WAL append channel sender for a shard. /// Called once during event loop startup before connections are accepted. pub fn set_wal_append_tx( &self, @@ -186,268 +172,292 @@ impl ShardDatabases { } } - /// Acquire exclusive access to a shard's VectorStore. + /// Acquire the process-global WorkspaceRegistry lock (C3 / M3). + /// + /// Workspaces are control-plane objects looked up by every connection + /// regardless of which shard accepted it, so all paths — handlers, + /// uring intercept, WAL replay — share one registry. WS commands are + /// rare; a single mutex is not a hot-path concern. + /// Caller lazy-inits via `get_or_insert_with(|| Box::new(WorkspaceRegistry::new()))`. #[inline] - pub fn vector_store(&self, shard_id: usize) -> MutexGuard<'_, VectorStore> { - self.vector_stores[shard_id].lock() + pub fn workspace_registry(&self) -> MutexGuard<'_, Option>> { + self.workspace_registry.lock() } - /// Acquire exclusive access to a shard's TextStore. + /// Total number of shards. #[inline] - pub fn text_store(&self, shard_id: usize) -> MutexGuard<'_, TextStore> { - self.text_stores[shard_id].lock() + pub fn num_shards(&self) -> usize { + self.num_shards } - /// Acquire shared read access to a shard's GraphStore. - #[cfg(feature = "graph")] + /// Number of databases per shard (typically 16). #[inline] - pub fn graph_store_read(&self, shard_id: usize) -> RwLockReadGuard<'_, GraphStore> { - self.graph_stores[shard_id].read() + pub fn db_count(&self) -> usize { + self.db_count } - /// Acquire exclusive write access to a shard's GraphStore. - #[cfg(feature = "graph")] + /// Get the per-shard memory publisher `Arc`. + /// + /// Called once per shard at startup when constructing the `ShardSlice`. + /// The returned `Arc` is cloned into `ShardSlice.estimated_memory`; the + /// master copy lives here in `memory_per_shard[shard_id]`. #[inline] - pub fn graph_store_write(&self, shard_id: usize) -> RwLockWriteGuard<'_, GraphStore> { - self.graph_stores[shard_id].write() + pub fn memory_publisher(&self, shard_id: usize) -> Arc { + self.memory_per_shard[shard_id].clone() } - /// Acquire the per-shard TemporalRegistry lock. Caller must lazy-init - /// via `get_or_insert_with(|| Box::new(TemporalRegistry::new()))`. + /// Publish a shard's aggregate memory usage (GAP-1 / Phase 2 wiring). + /// + /// Called from the shard's 100ms eviction tick. One Relaxed store. #[inline] - pub fn temporal_registry( - &self, - shard_id: usize, - ) -> MutexGuard<'_, Option>> { - self.temporal_registries[shard_id].lock() + pub fn publish_memory(&self, shard_id: usize, used: usize) { + self.memory_per_shard[shard_id].store(used, Ordering::Relaxed); } - /// Acquire the per-shard TemporalKvIndex lock. Caller must lazy-init - /// via `get_or_insert_with(|| Box::new(TemporalKvIndex::new()))`. - #[inline] - pub fn temporal_kv_index( + /// Recompute and publish this shard's elastic memory budget (GAP-1). + pub fn recompute_elastic_budget( &self, shard_id: usize, - ) -> MutexGuard<'_, Option>> { - self.temporal_kv_indexes[shard_id].lock() + config: &crate::config::RuntimeConfig, + ) -> usize { + let base = config.maxmemory_per_shard(); + if base == 0 || self.num_shards <= 1 { + self.elastic_budgets[shard_id].store(0, Ordering::Relaxed); + return 0; + } + let used: Vec = self + .memory_per_shard + .iter() + .map(|a| a.load(Ordering::Relaxed)) + .collect(); + let budget = crate::storage::eviction::compute_elastic_budget(shard_id, base, &used); + self.elastic_budgets[shard_id].store(budget, Ordering::Relaxed); + budget } - /// Acquire the GLOBAL WorkspaceRegistry lock (slot 0 of the per-shard - /// array). Workspaces are control-plane objects looked up by every - /// connection regardless of which shard accepted it, so all paths — - /// handlers, uring intercept, WAL replay — share one registry. WS - /// commands are rare; a single mutex is not a hot-path concern. - /// Caller lazy-inits via `get_or_insert_with(|| Box::new(WorkspaceRegistry::new()))`. + /// This shard's current elastic budget; `0` = none published yet. #[inline] - pub fn workspace_registry(&self) -> MutexGuard<'_, Option>> { - self.workspace_registries[0].lock() + pub fn elastic_budget(&self, shard_id: usize) -> usize { + self.elastic_budgets[shard_id].load(Ordering::Relaxed) } - /// Acquire the per-shard DurableQueueRegistry lock. - /// Caller lazy-inits via `get_or_insert_with(|| Box::new(DurableQueueRegistry::new()))`. + /// Sum all per-shard memory publishers with `Relaxed` loads. Lock-free. #[inline] - pub fn durable_queue_registry( - &self, - shard_id: usize, - ) -> MutexGuard<'_, Option>> { - self.durable_queue_registries[shard_id].lock() + pub fn read_memory_sum(&self) -> usize { + self.memory_per_shard + .iter() + .map(|a| a.load(Ordering::Relaxed)) + .sum() } - /// Acquire the per-shard TriggerRegistry lock. - /// Caller lazy-inits via `get_or_insert_with(|| Box::new(TriggerRegistry::new()))`. + /// Read the last published KV memory for a single shard with one Relaxed load. #[inline] - pub fn trigger_registry( - &self, - shard_id: usize, - ) -> MutexGuard<'_, Option>> { - self.trigger_registries[shard_id].lock() + pub fn published_shard_memory(&self, shard_id: usize) -> usize { + self.memory_per_shard[shard_id].load(Ordering::Relaxed) } - /// Acquire the per-shard KV write-intent side-table lock. - /// Used by handler write path (record_write) and read path (is_key_visible). + /// Return a clone of the `Arc` for `shard_id`. + /// + /// Called once at shard startup to hand the `Arc` into `ShardSliceInit`. #[inline] - pub fn kv_intents(&self, shard_id: usize) -> MutexGuard<'_, KvWriteIntents> { - self.kv_write_intents[shard_id].lock() + pub fn store_memory_publisher(&self, shard_id: usize) -> Arc { + self.store_memory_per_shard[shard_id].clone() } +} - /// Acquire the per-shard deferred HNSW insert queue lock. - /// Used by TXN.COMMIT (drain_for_txn) and TXN.ABORT (discard_for_txn). - #[inline] - pub fn hnsw_queue(&self, shard_id: usize) -> MutexGuard<'_, DeferredHnswInserts> { - self.deferred_hnsw_inserts[shard_id].lock() - } +// ── Boot-time recovery free functions ───────────────────────────────────────── +// +// These functions operate on `&mut [ShardSliceInit]` (pre-packaged per-shard +// state) rather than on `ShardDatabases`. They are called single-threaded +// during server startup AFTER `ShardDatabases::new` returns the init packages +// and BEFORE shard threads are spawned — no locks needed. - /// Replay WAL WorkspaceCreate and WorkspaceDrop records to restore workspace registry. - /// - /// Called during server startup after graph and temporal WAL replay. - /// Scans per-shard WAL directories for v3 segment files and processes - /// WorkspaceCreate/WorkspaceDrop records to restore workspace metadata. - pub fn replay_workspace_wal(&self, persistence_dir: &std::path::Path) { - use crate::persistence::wal_v3::record::{WalRecord, WalRecordType}; - - for shard_id in 0..self.num_shards { - let wal_dir = persistence_dir.join(format!("shard-{}", shard_id)); - if !wal_dir.exists() { - continue; - } +/// Replay workspace WAL records into the process-global `WorkspaceRegistry`. +/// +/// Called during server startup after graph and temporal WAL replay. +/// Scans `shard-{id}/wal-v3/` for v3 segment files and processes +/// WorkspaceCreate/WorkspaceDrop records. All records populate the single +/// global registry (C3/M3). +pub fn replay_workspace_wal(shared: &Arc, persistence_dir: &std::path::Path) { + use crate::persistence::wal_v3::record::{WalRecord, WalRecordType, read_wal_v3_record}; + + let num_shards = shared.num_shards; + let mut total_create = 0u64; + let mut total_drop = 0u64; + + for shard_id in 0..num_shards { + // WAL v3 segments live in shard-{id}/wal-v3/ — matching recovery.rs:361. + let wal_dir = persistence_dir + .join(format!("shard-{}", shard_id)) + .join("wal-v3"); + if !wal_dir.exists() { + continue; + } - let mut create_count = 0u64; - let mut drop_count = 0u64; - - let on_command = &mut |record: &WalRecord| { - match record.record_type { - WalRecordType::WorkspaceCreate => { - if let Some((ws_bytes, name)) = decode_workspace_create(&record.payload) { - let ws_id = WorkspaceId::from_bytes(ws_bytes); - let meta = WorkspaceMetadata { - id: ws_id, - name: bytes::Bytes::from(name), - created_at: 0, // WAL doesn't store created_at; use 0 as placeholder - }; - // Global registry (slot 0) — see workspace_registry(). - let mut guard = self.workspace_registries[0].lock(); - let reg = - guard.get_or_insert_with(|| Box::new(WorkspaceRegistry::new())); - reg.insert(ws_id, meta); - create_count += 1; - } + let mut create_count = 0u64; + let mut drop_count = 0u64; + + let mut handle_record = |record_type: WalRecordType, payload: &[u8]| match record_type { + WalRecordType::WorkspaceCreate => { + if let Some((ws_bytes, name)) = decode_workspace_create(payload) { + let ws_id = WorkspaceId::from_bytes(ws_bytes); + let meta = WorkspaceMetadata { + id: ws_id, + name: bytes::Bytes::from(name), + created_at: 0, + }; + let mut guard = shared.workspace_registry.lock(); + let reg = guard.get_or_insert_with(|| Box::new(WorkspaceRegistry::new())); + reg.insert(ws_id, meta); + create_count += 1; + } + } + WalRecordType::WorkspaceDrop => { + if let Some(ws_bytes) = decode_workspace_drop(payload) { + let ws_id = WorkspaceId::from_bytes(ws_bytes); + let mut guard = shared.workspace_registry.lock(); + if let Some(reg) = guard.as_mut() { + reg.remove(&ws_id); } - WalRecordType::WorkspaceDrop => { - if let Some(ws_bytes) = decode_workspace_drop(&record.payload) { - let ws_id = WorkspaceId::from_bytes(ws_bytes); - let mut guard = self.workspace_registries[0].lock(); - if let Some(reg) = guard.as_mut() { - reg.remove(&ws_id); - } - drop_count += 1; + drop_count += 1; + } + } + _ => {} + }; + + let on_command = &mut |record: &WalRecord| match record.record_type { + WalRecordType::WorkspaceCreate | WalRecordType::WorkspaceDrop => { + handle_record(record.record_type, &record.payload); + } + WalRecordType::Command => { + if let Some(inner) = read_wal_v3_record(&record.payload) { + match inner.record_type { + WalRecordType::WorkspaceCreate | WalRecordType::WorkspaceDrop => { + handle_record(inner.record_type, &inner.payload); } + _ => {} } - _ => {} // Skip non-workspace records - } - }; - let on_fpi = &mut |_: &WalRecord| {}; - - // Scan WAL files in the shard directory - if let Ok(entries) = std::fs::read_dir(&wal_dir) { - let mut wal_files: Vec<_> = entries - .filter_map(|e| e.ok()) - .filter(|e| e.file_name().to_str().is_some_and(|n| n.ends_with(".wal"))) - .map(|e| e.path()) - .collect(); - wal_files.sort(); - - for wal_file in &wal_files { - let _ = crate::persistence::wal_v3::replay::replay_wal_v3_file( - wal_file, 0, on_command, on_fpi, - ); } } - - if create_count > 0 || drop_count > 0 { - tracing::info!( - "Shard {}: replayed {} WorkspaceCreate + {} WorkspaceDrop WAL records", - shard_id, - create_count, - drop_count, + _ => {} + }; + let on_fpi = &mut |_: &WalRecord| {}; + + if let Ok(entries) = std::fs::read_dir(&wal_dir) { + let mut wal_files: Vec<_> = entries + .filter_map(|e| e.ok()) + .filter(|e| e.file_name().to_str().is_some_and(|n| n.ends_with(".wal"))) + .map(|e| e.path()) + .collect(); + wal_files.sort(); + + for wal_file in &wal_files { + let _ = crate::persistence::wal_v3::replay::replay_wal_v3_file( + wal_file, 0, on_command, on_fpi, ); } } + + total_create += create_count; + total_drop += drop_count; + + if create_count > 0 || drop_count > 0 { + tracing::info!( + "Shard {}: replayed {} WorkspaceCreate + {} WorkspaceDrop WAL records \ + into global registry", + shard_id, + create_count, + drop_count, + ); + } } - /// Replay MQ WAL records to restore DurableQueueRegistry and apply cursor-rollback. - /// - /// This is the P2 CRITICAL path: after restart, unacknowledged messages must be - /// redelivered. The cursor-rollback resets `ConsumerGroup.last_delivered_id` to - /// the last-acked position so that MQ.POP will re-deliver unacked messages. - /// - /// **Cursor-rollback algorithm:** - /// For each durable queue, after WAL scan: - /// - If PEL is non-empty: rollback target = min(PEL keys) - 1 - /// (the ID just before the first unacked message) - /// - If PEL is empty: no rollback needed (all messages were acked) - /// - /// This handles out-of-order acks correctly: if messages 1,2,3,4,5 were delivered - /// but only 1,3,5 were acked, the PEL contains {2,4}. Rollback target = 1 - /// (min(PEL) - 1), so messages 2,3,4,5 are all redelivered. - pub fn replay_mq_wal(&self, persistence_dir: &std::path::Path) { - use crate::persistence::wal_v3::record::{WalRecord, WalRecordType}; - use crate::storage::stream::StreamId; - - for shard_id in 0..self.num_shards { - let wal_dir = persistence_dir.join(format!("shard-{}", shard_id)); - if !wal_dir.exists() { - continue; - } + if total_create > 0 || total_drop > 0 { + tracing::info!( + "Workspace WAL replay complete: {} creates, {} drops across all shards", + total_create, + total_drop, + ); + } +} - // Phase 1: Scan WAL to collect MqCreate configs and MqAck positions - let mut durable_configs: HashMap, u32> = HashMap::new(); - let mut ack_count = 0u64; +/// Replay MQ WAL records to restore DurableQueueRegistry and apply cursor-rollback. +/// +/// Operates on `&mut [ShardSliceInit]` — called single-threaded at boot +/// before shard threads are spawned. No locks needed. +pub fn replay_mq_wal( + inits: &mut [crate::shard::slice::ShardSliceInit], + persistence_dir: &std::path::Path, +) { + use std::collections::HashMap; + + use crate::persistence::wal_v3::record::{WalRecord, WalRecordType}; + use crate::storage::stream::StreamId; + + for init in inits.iter_mut() { + let shard_id = init.shard_id; + let wal_dir = persistence_dir.join(format!("shard-{}", shard_id)); + if !wal_dir.exists() { + continue; + } - let on_command = &mut |record: &WalRecord| match record.record_type { - WalRecordType::MqCreate => { - if let Some((queue_key, max_delivery_count)) = - crate::mq::wal::decode_mq_create(&record.payload) - { - durable_configs.insert(queue_key, max_delivery_count); - } - } - WalRecordType::MqAck => { - if crate::mq::wal::decode_mq_ack(&record.payload).is_some() { - ack_count += 1; - } + let mut durable_configs: HashMap, u32> = HashMap::new(); + let mut ack_count = 0u64; + + let on_command = &mut |record: &WalRecord| match record.record_type { + WalRecordType::MqCreate => { + if let Some((queue_key, max_delivery_count)) = + crate::mq::wal::decode_mq_create(&record.payload) + { + durable_configs.insert(queue_key, max_delivery_count); } - _ => {} - }; - let on_fpi = &mut |_: &WalRecord| {}; - - // Scan WAL files in the shard directory - if let Ok(entries) = std::fs::read_dir(&wal_dir) { - let mut wal_files: Vec<_> = entries - .filter_map(|e| e.ok()) - .filter(|e| e.file_name().to_str().is_some_and(|n| n.ends_with(".wal"))) - .map(|e| e.path()) - .collect(); - wal_files.sort(); - - for wal_file in &wal_files { - let _ = crate::persistence::wal_v3::replay::replay_wal_v3_file( - wal_file, 0, on_command, on_fpi, - ); + } + WalRecordType::MqAck => { + if crate::mq::wal::decode_mq_ack(&record.payload).is_some() { + ack_count += 1; } } + _ => {} + }; + let on_fpi = &mut |_: &WalRecord| {}; + + if let Ok(entries) = std::fs::read_dir(&wal_dir) { + let mut wal_files: Vec<_> = entries + .filter_map(|e| e.ok()) + .filter(|e| e.file_name().to_str().is_some_and(|n| n.ends_with(".wal"))) + .map(|e| e.path()) + .collect(); + wal_files.sort(); + + for wal_file in &wal_files { + let _ = crate::persistence::wal_v3::replay::replay_wal_v3_file( + wal_file, 0, on_command, on_fpi, + ); + } + } - // Phase 2: Restore DurableQueueRegistry from MqCreate records - if !durable_configs.is_empty() { - let mut guard = self.durable_queue_registries[shard_id].lock(); - let reg = - guard.get_or_insert_with(|| Box::new(crate::mq::DurableQueueRegistry::new())); - for (queue_key_bytes, max_delivery_count) in &durable_configs { - let key = bytes::Bytes::copy_from_slice(queue_key_bytes); - let config = - crate::mq::DurableStreamConfig::new(key.clone(), *max_delivery_count); - reg.insert(key, config); - } + if !durable_configs.is_empty() { + let reg = init + .durable_queue_registry + .get_or_insert_with(|| Box::new(crate::mq::DurableQueueRegistry::new())); + for (queue_key_bytes, max_delivery_count) in &durable_configs { + let key = bytes::Bytes::copy_from_slice(queue_key_bytes); + let config = crate::mq::DurableStreamConfig::new(key.clone(), *max_delivery_count); + reg.insert(key, config); } + } - // Phase 3: Cursor-rollback for each durable queue. - // The Stream and its __mq_consumers consumer group were already restored - // from RESP Command WAL records (regular KV replay path). - // We just need to reset last_delivered_id based on PEL state. - for (queue_key_bytes, max_dc) in &durable_configs { - let key_bytes = bytes::Bytes::copy_from_slice(queue_key_bytes); - let mut db_guard = self.write_db(shard_id, 0); - - // Look up the Stream entry via get_stream_mut - if let Ok(Some(stream)) = db_guard.get_stream_mut(&key_bytes) { - // Mark as durable (may not have been set during Command replay) + // Cursor-rollback for each durable queue using db 0. + for (queue_key_bytes, max_dc) in &durable_configs { + let key_bytes = bytes::Bytes::copy_from_slice(queue_key_bytes); + // db 0 is the first database in the slice. + if let Some(db) = init.databases.get_mut(0) { + if let Ok(Some(stream)) = db.get_stream_mut(&key_bytes) { stream.durable = true; stream.max_delivery_count = *max_dc; - // Find the __mq_consumers group let group_name = bytes::Bytes::from_static(b"__mq_consumers"); if let Some(group) = stream.groups.get_mut(&group_name) { - // Cursor-rollback: if PEL is non-empty, set last_delivered_id - // to the ID just before the first unacked message if let Some((min_pel_id, _)) = group.pel.iter().next() { let rollback_target = if min_pel_id.seq > 0 { StreamId { @@ -455,7 +465,6 @@ impl ShardDatabases { seq: min_pel_id.seq - 1, } } else if min_pel_id.ms > 0 { - // Edge case: seq=0, roll back to previous ms with max seq StreamId { ms: min_pel_id.ms - 1, seq: u64::MAX, @@ -478,473 +487,234 @@ impl ShardDatabases { group.last_delivered_id = rollback_target; } - // If PEL is empty: all messages were acked, no rollback needed } } } + } - if !durable_configs.is_empty() { - tracing::info!( - "Shard {}: replayed {} MQ queue configs, {} ack records", - shard_id, - durable_configs.len(), - ack_count, - ); - } + if !durable_configs.is_empty() { + tracing::info!( + "Shard {}: replayed {} MQ queue configs, {} ack records", + shard_id, + durable_configs.len(), + ack_count, + ); } } +} - /// Recover graph stores from persistence for all shards. - /// - /// Called once after construction during server startup. Loads graph - /// metadata and CSR segments from the persistence directory. - #[cfg(feature = "graph")] - pub fn recover_graph_stores(&self, persistence_dir: &std::path::Path) { - for shard_id in 0..self.num_shards { - match crate::graph::recovery::recover_graph_store(persistence_dir, shard_id) { - Ok(Some(result)) => { - if result.store.graph_count() > 0 { - tracing::info!( - "Shard {}: recovered {} graph(s) ({} segments loaded, {} skipped)", - shard_id, - result.store.graph_count(), - result.segments_loaded, - result.segments_skipped, - ); - } - *self.graph_stores[shard_id].write() = result.store; - } - Ok(None) => { - // No graph metadata — clean start, nothing to recover. - } - Err(e) => { - tracing::error!("Shard {}: graph recovery failed: {}", shard_id, e); +/// Recover graph stores from persistence for all shards. +/// +/// Operates on `&mut [ShardSliceInit]` — called single-threaded at boot. +#[cfg(feature = "graph")] +pub fn recover_graph_stores( + inits: &mut [crate::shard::slice::ShardSliceInit], + persistence_dir: &std::path::Path, +) { + for init in inits.iter_mut() { + let shard_id = init.shard_id; + match crate::graph::recovery::recover_graph_store(persistence_dir, shard_id) { + Ok(Some(result)) => { + if result.store.graph_count() > 0 { + tracing::info!( + "Shard {}: recovered {} graph(s) ({} segments loaded, {} skipped)", + shard_id, + result.store.graph_count(), + result.segments_loaded, + result.segments_skipped, + ); } + init.graph_store = result.store; + } + Ok(None) => {} + Err(e) => { + tracing::error!("Shard {}: graph recovery failed: {}", shard_id, e); } } } +} - /// Replay graph WAL commands into graph stores for all shards. - /// - /// Called after `recover_graph_stores` during startup. Reads per-shard - /// WAL files, filters graph commands via `GraphReplayCollector`, and - /// applies them to each shard's `GraphStore`. - #[cfg(feature = "graph")] - pub fn replay_graph_wal(&self, persistence_dir: &std::path::Path) { - use crate::persistence::replay::DispatchReplayEngine; - use crate::persistence::wal; - - for shard_id in 0..self.num_shards { - let wal_file = wal::wal_path(persistence_dir, shard_id); - if !wal_file.exists() { - continue; - } - // Create a temporary engine to collect graph commands from the WAL. - let engine = DispatchReplayEngine::new(); - // We need dummy databases for the KV replay path (graph commands - // are intercepted before KV dispatch, so these are unused). - let mut dummy_dbs: Vec = - (0..self.db_count).map(|_| Database::new()).collect(); - match wal::replay_wal(&mut dummy_dbs, &wal_file, &engine) { - Ok(_) => { - let graph_count = engine.graph_command_count(); - if graph_count > 0 { - let mut gs = self.graph_stores[shard_id].write(); - let applied = engine.replay_graph_commands(&mut gs); - tracing::info!( - "Shard {}: replayed {} graph WAL commands ({} applied)", - shard_id, - graph_count, - applied, - ); - } - } - Err(e) => { - tracing::error!("Shard {}: graph WAL replay failed: {}", shard_id, e); +/// Replay graph WAL commands into graph stores for all shards. +#[cfg(feature = "graph")] +pub fn replay_graph_wal( + inits: &mut [crate::shard::slice::ShardSliceInit], + persistence_dir: &std::path::Path, + db_count: usize, +) { + use crate::persistence::replay::DispatchReplayEngine; + use crate::persistence::wal; + + for init in inits.iter_mut() { + let shard_id = init.shard_id; + let wal_file = wal::wal_path(persistence_dir, shard_id); + if !wal_file.exists() { + continue; + } + let engine = DispatchReplayEngine::new(); + let mut dummy_dbs: Vec = (0..db_count).map(|_| Database::new()).collect(); + match wal::replay_wal(&mut dummy_dbs, &wal_file, &engine) { + Ok(_) => { + let graph_count = engine.graph_command_count(); + if graph_count > 0 { + let applied = engine.replay_graph_commands(&mut init.graph_store); + tracing::info!( + "Shard {}: replayed {} graph WAL commands ({} applied)", + shard_id, + graph_count, + applied, + ); } } + Err(e) => { + tracing::error!("Shard {}: graph WAL replay failed: {}", shard_id, e); + } } } +} - /// Replay temporal WAL records into per-shard TemporalKvIndex and GraphStore. - /// - /// Called at startup after graph recovery. TemporalUpsert (KV) records are - /// always replayed regardless of the `graph` feature. GraphTemporal records - /// are only processed when the `graph` feature is enabled. - pub fn replay_temporal_wal(&self, persistence_dir: &std::path::Path) { - #[cfg(feature = "graph")] - use crate::persistence::wal_v3::record::decode_graph_temporal; - use crate::persistence::wal_v3::record::{ - WalRecord, WalRecordType, decode_temporal_upsert, - }; +/// Replay temporal WAL records into per-shard TemporalKvIndex and GraphStore. +pub fn replay_temporal_wal( + inits: &mut [crate::shard::slice::ShardSliceInit], + persistence_dir: &std::path::Path, +) { + #[cfg(feature = "graph")] + use crate::persistence::wal_v3::record::decode_graph_temporal; + use crate::persistence::wal_v3::record::{WalRecord, WalRecordType, decode_temporal_upsert}; + + for init in inits.iter_mut() { + let shard_id = init.shard_id; + let wal_dir = persistence_dir.join(format!("shard-{}", shard_id)); + if !wal_dir.exists() { + continue; + } - for shard_id in 0..self.num_shards { - let wal_dir = persistence_dir.join(format!("shard-{}", shard_id)); - if !wal_dir.exists() { - continue; + let mut temporal_upsert_count = 0usize; + #[cfg(feature = "graph")] + let mut graph_temporal_count = 0usize; + + let on_command = &mut |record: &WalRecord| match record.record_type { + WalRecordType::TemporalUpsert => { + if let Some((key, valid_from, _system_from, value)) = + decode_temporal_upsert(&record.payload) + { + let idx = init + .temporal_kv_index + .get_or_insert_with(|| Box::new(crate::temporal::TemporalKvIndex::new())); + idx.record( + bytes::Bytes::copy_from_slice(key), + valid_from, + bytes::Bytes::copy_from_slice(value), + ); + temporal_upsert_count += 1; + } } - - let mut temporal_upsert_count = 0usize; #[cfg(feature = "graph")] - let mut graph_temporal_count = 0usize; - - let on_command = &mut |record: &WalRecord| { - match record.record_type { - WalRecordType::TemporalUpsert => { - if let Some((key, valid_from, _system_from, value)) = - decode_temporal_upsert(&record.payload) - { - let mut guard = self.temporal_kv_indexes[shard_id].lock(); - let idx = guard.get_or_insert_with(|| { - Box::new(crate::temporal::TemporalKvIndex::new()) - }); - idx.record( - bytes::Bytes::copy_from_slice(key), - valid_from, - bytes::Bytes::copy_from_slice(value), - ); - temporal_upsert_count += 1; - } - } - #[cfg(feature = "graph")] - WalRecordType::GraphTemporal => { - if let Some((entity_id, is_node, valid_to, _system_from)) = - decode_graph_temporal(&record.payload) - { - let mut gs = self.graph_stores[shard_id].write(); - for named_graph in gs.iter_graphs_mut() { - let found = if is_node { - let nk: crate::graph::types::NodeKey = - slotmap::KeyData::from_ffi(entity_id).into(); - if let Some(node) = named_graph.write_buf.get_node_mut(nk) { - node.valid_to = valid_to; - true - } else { - false - } - } else { - let ek: crate::graph::types::EdgeKey = - slotmap::KeyData::from_ffi(entity_id).into(); - if let Some(edge) = named_graph.write_buf.get_edge_mut(ek) { - edge.valid_to = valid_to; - true - } else { - false - } - }; - if found { - graph_temporal_count += 1; - break; - } + WalRecordType::GraphTemporal => { + if let Some((entity_id, is_node, valid_to, _system_from)) = + decode_graph_temporal(&record.payload) + { + for named_graph in init.graph_store.iter_graphs_mut() { + let found = if is_node { + let nk: crate::graph::types::NodeKey = + slotmap::KeyData::from_ffi(entity_id).into(); + if let Some(node) = named_graph.write_buf.get_node_mut(nk) { + node.valid_to = valid_to; + true + } else { + false + } + } else { + let ek: crate::graph::types::EdgeKey = + slotmap::KeyData::from_ffi(entity_id).into(); + if let Some(edge) = named_graph.write_buf.get_edge_mut(ek) { + edge.valid_to = valid_to; + true + } else { + false } + }; + if found { + graph_temporal_count += 1; + break; } } - _ => {} // Other record types handled by their respective replay paths - } - }; - let on_fpi = &mut |_: &WalRecord| {}; - - // Scan WAL files in the shard directory - if let Ok(entries) = std::fs::read_dir(&wal_dir) { - let mut wal_files: Vec<_> = entries - .filter_map(|e| e.ok()) - .filter(|e| e.file_name().to_str().is_some_and(|n| n.ends_with(".wal"))) - .map(|e| e.path()) - .collect(); - wal_files.sort(); - - for wal_file in &wal_files { - let _ = crate::persistence::wal_v3::replay::replay_wal_v3_file( - wal_file, 0, on_command, on_fpi, - ); } } - - #[cfg(feature = "graph")] - if temporal_upsert_count > 0 || graph_temporal_count > 0 { - tracing::info!( - "Shard {}: replayed {} TemporalUpsert + {} GraphTemporal WAL records", - shard_id, - temporal_upsert_count, - graph_temporal_count, - ); - } - #[cfg(not(feature = "graph"))] - if temporal_upsert_count > 0 { - tracing::info!( - "Shard {}: replayed {} TemporalUpsert WAL records", - shard_id, - temporal_upsert_count, + _ => {} + }; + let on_fpi = &mut |_: &WalRecord| {}; + + if let Ok(entries) = std::fs::read_dir(&wal_dir) { + let mut wal_files: Vec<_> = entries + .filter_map(|e| e.ok()) + .filter(|e| e.file_name().to_str().is_some_and(|n| n.ends_with(".wal"))) + .map(|e| e.path()) + .collect(); + wal_files.sort(); + + for wal_file in &wal_files { + let _ = crate::persistence::wal_v3::replay::replay_wal_v3_file( + wal_file, 0, on_command, on_fpi, ); } } - } - - /// Acquire a shared read lock on a specific database. - #[inline] - pub fn read_db(&self, shard_id: usize, db_index: usize) -> RwLockReadGuard<'_, Database> { - debug_assert!( - shard_id < self.shards.len(), - "shard_id {shard_id} out of bounds ({})", - self.shards.len() - ); - debug_assert!( - db_index < self.shards[shard_id].len(), - "db_index {db_index} out of bounds ({})", - self.shards[shard_id].len() - ); - self.shards[shard_id][db_index].read() - } - - /// Acquire an exclusive write lock on a specific database. - /// - /// Fast path: bounded spin with `try_write()` for the common case where - /// cross-shard read locks are held briefly (~51ns for DashTable lookup). - /// Slow path: falls back to blocking `write()` (OS-level parking) to avoid - /// infinite busy-spin when readers hold locks longer (e.g., KEYS, SCAN). - #[inline] - pub fn write_db(&self, shard_id: usize, db_index: usize) -> RwLockWriteGuard<'_, Database> { - debug_assert!( - shard_id < self.shards.len(), - "shard_id {shard_id} out of bounds ({})", - self.shards.len() - ); - debug_assert!( - db_index < self.shards[shard_id].len(), - "db_index {db_index} out of bounds ({})", - self.shards[shard_id].len() - ); - // Spin briefly — resolves in 1-3 iterations for typical cross-shard reads. - for _ in 0..32 { - if let Some(guard) = self.shards[shard_id][db_index].try_write() { - return guard; - } - std::hint::spin_loop(); - } - // Slow path: park the thread via OS futex instead of busy-spinning. - self.shards[shard_id][db_index].write() - } - - /// Non-blocking write lock attempt for async callers that can yield. - #[inline] - pub fn try_write_db( - &self, - shard_id: usize, - db_index: usize, - ) -> Option> { - self.shards[shard_id][db_index].try_write() - } - - /// Total number of shards. - #[inline] - pub fn num_shards(&self) -> usize { - self.num_shards - } - - /// Return a reference to all databases across all shards (for AOF rewrite). - /// Callers iterate shards × dbs and acquire read locks individually. - #[inline] - pub fn all_shard_dbs(&self) -> &[Vec>] { - &self.shards - } - - /// Number of databases per shard (typically 16). - #[inline] - pub fn db_count(&self) -> usize { - self.db_count - } - - /// Aggregate estimated memory across all databases in a shard. - /// - /// Acquires read locks briefly on each DB. Used for maxmemory eviction - /// decisions (Redis maxmemory is a server-wide limit, not per-DB). - /// - /// Phase 3 will replace call sites with `read_memory_sum()` to eliminate - /// these lock acquisitions. Both methods coexist during the transition. - pub fn aggregate_memory(&self, shard_id: usize) -> usize { - let mut total = 0usize; - for db_idx in 0..self.db_count { - let guard = self.read_db(shard_id, db_idx); - total += guard.estimated_memory(); - } - total - } - /// Get the per-shard memory publisher `Arc`. - /// - /// Called once per shard at startup when constructing the `ShardSlice`. - /// The returned `Arc` is cloned into `ShardSlice.estimated_memory`; the - /// master copy lives here in `memory_per_shard[shard_id]`. - /// - /// Phase 2 write paths will store into this atomic after mutations. - /// Phase 3 eviction and metrics paths will read via `read_memory_sum()`. - #[inline] - pub fn memory_publisher(&self, shard_id: usize) -> Arc { - self.memory_per_shard[shard_id].clone() - } - - /// Publish a shard's aggregate memory usage (GAP-1 / Phase 2 wiring). - /// - /// Called from the shard's 100ms eviction tick with the freshly computed - /// `aggregate_memory(shard_id)` value. One Relaxed store — siblings read - /// the snapshot in `recompute_elastic_budget`. - #[inline] - pub fn publish_memory(&self, shard_id: usize, used: usize) { - self.memory_per_shard[shard_id].store(used, Ordering::Relaxed); - } - - /// Recompute and publish this shard's elastic memory budget (GAP-1). - /// - /// Reads every shard's published usage (N Relaxed loads, once per 100ms - /// tick) and derives this shard's budget via - /// [`crate::storage::eviction::compute_elastic_budget`]: under-budget - /// siblings donate their headroom, hot shards split it evenly, and the - /// snapshot aggregate never exceeds `maxmemory`. Returns the budget it - /// stored. Single-shard or unlimited deployments publish `0` - /// ("use static budget") so readers take the unchanged fallback path. - pub fn recompute_elastic_budget( - &self, - shard_id: usize, - config: &crate::config::RuntimeConfig, - ) -> usize { - let base = config.maxmemory_per_shard(); - if base == 0 || self.num_shards <= 1 { - self.elastic_budgets[shard_id].store(0, Ordering::Relaxed); - return 0; - } - // Snapshot of every shard's published usage. SmallVec would save the - // alloc, but this runs 10×/s on a background tick — a Vec is fine. - let used: Vec = self - .memory_per_shard - .iter() - .map(|a| a.load(Ordering::Relaxed)) - .collect(); - let budget = crate::storage::eviction::compute_elastic_budget(shard_id, base, &used); - self.elastic_budgets[shard_id].store(budget, Ordering::Relaxed); - budget - } - - /// This shard's current elastic budget; `0` = none published yet (use - /// the static `maxmemory / num_shards`). One Relaxed load — cheap enough - /// for the per-write eviction check. - #[inline] - pub fn elastic_budget(&self, shard_id: usize) -> usize { - self.elastic_budgets[shard_id].load(Ordering::Relaxed) - } - - /// Sum all per-shard memory publishers with `Relaxed` loads. Lock-free. - /// - /// Returns a best-effort estimate of total server memory across all shards. - /// `Relaxed` ordering is intentional: the eviction tick and metrics scrape - /// do not require cross-thread synchronization with write operations — they - /// only need a recent-enough value for threshold decisions. The error - /// introduced by stale atomics is far smaller than the headroom kept below - /// `maxmemory`. - /// - /// Phase 3 will switch `persistence_tick.rs:436,511` and - /// `metrics_setup.rs:1238` from `aggregate_memory()` to this method. - #[inline] - pub fn read_memory_sum(&self) -> usize { - self.memory_per_shard - .iter() - .map(|a| a.load(Ordering::Relaxed)) - .sum() - } - - /// Collect snapshot metadata (segment counts, base timestamps) for a shard. - /// - /// Acquires brief read locks on each database to gather metadata needed - /// by `SnapshotState::new_from_metadata`. Called once at snapshot epoch start. - pub fn snapshot_metadata(&self, shard_id: usize) -> (Vec, Vec) { - let mut segment_counts = Vec::with_capacity(self.db_count); - let mut base_timestamps = Vec::with_capacity(self.db_count); - for i in 0..self.db_count { - let guard = self.read_db(shard_id, i); - segment_counts.push(guard.data().segment_count()); - base_timestamps.push(guard.base_timestamp()); + #[cfg(feature = "graph")] + if temporal_upsert_count > 0 || graph_temporal_count > 0 { + tracing::info!( + "Shard {}: replayed {} TemporalUpsert + {} GraphTemporal WAL records", + shard_id, + temporal_upsert_count, + graph_temporal_count, + ); } - (segment_counts, base_timestamps) - } - - /// Atomically swap two databases within a single shard. - /// - /// Acquires write locks in ascending index order (lower index first) to - /// prevent deadlock if two concurrent SWAPDB calls swap the same pair from - /// opposite directions. `std::mem::swap` exchanges the `Database` values - /// in-place while both locks are held. - /// - /// # Panics - /// - /// Panics (debug_assert) if `shard_id`, `a`, or `b` are out of bounds. - /// Callers must validate indices before calling. - pub fn swap_dbs(&self, shard_id: usize, a: usize, b: usize) { - debug_assert!(shard_id < self.shards.len(), "shard_id out of bounds"); - debug_assert!(a < self.db_count, "db index a out of bounds"); - debug_assert!(b < self.db_count, "db index b out of bounds"); - // Same-index swap is a no-op; short-circuit in release builds to - // avoid self-deadlocking on the second write() acquire (parking_lot - // RwLock is not reentrant). Callers normally short-circuit earlier, - // but defending here is cheap and prevents a stall on the SPSC path. - if a == b { - return; + #[cfg(not(feature = "graph"))] + if temporal_upsert_count > 0 { + tracing::info!( + "Shard {}: replayed {} TemporalUpsert WAL records", + shard_id, + temporal_upsert_count, + ); } - - let (lo, hi) = if a < b { (a, b) } else { (b, a) }; - // Acquire in ascending index order to prevent deadlock. - let mut guard_lo = self.shards[shard_id][lo].write(); - let mut guard_hi = self.shards[shard_id][hi].write(); - std::mem::swap(&mut *guard_lo, &mut *guard_hi); } } #[cfg(test)] mod tests { use super::*; + use crate::storage::Database; + + fn new_shared(shard_count: usize, db_per_shard: usize) -> Arc { + let dbs: Vec> = (0..shard_count) + .map(|_| (0..db_per_shard).map(|_| Database::new()).collect()) + .collect(); + let (shared, _inits) = ShardDatabases::new(dbs); + shared + } #[test] fn test_new_creates_correct_dimensions() { - let dbs = vec![ - vec![Database::new(), Database::new()], - vec![Database::new(), Database::new()], - vec![Database::new(), Database::new()], - ]; - let shared = ShardDatabases::new(dbs); + let shared = new_shared(3, 2); assert_eq!(shared.num_shards(), 3); assert_eq!(shared.db_count(), 2); } #[test] - fn test_read_db_returns_shared_guard() { - let dbs = vec![vec![Database::new()]]; - let shared = ShardDatabases::new(dbs); - let _guard1 = shared.read_db(0, 0); - let _guard2 = shared.read_db(0, 0); // Multiple readers OK - } - - #[test] - fn test_write_db_returns_exclusive_guard() { - let dbs = vec![vec![Database::new()]]; - let shared = ShardDatabases::new(dbs); - let _guard = shared.write_db(0, 0); - // Exclusive access -- would deadlock if we tried another write_db here + fn test_empty_shard_databases() { + let (shared, inits) = ShardDatabases::new(vec![]); + assert_eq!(shared.num_shards(), 0); + assert_eq!(shared.db_count(), 0); + assert!(inits.is_empty()); } #[test] - fn test_cross_shard_access() { + fn test_new_returns_correct_init_count() { let dbs = vec![vec![Database::new()], vec![Database::new()]]; - let shared = ShardDatabases::new(dbs); - // Can read from different shards concurrently - let _guard0 = shared.read_db(0, 0); - let _guard1 = shared.read_db(1, 0); - } - - #[test] - fn test_snapshot_metadata() { - let dbs = vec![vec![Database::new(), Database::new()]]; - let shared = ShardDatabases::new(dbs); - let (seg_counts, base_ts) = shared.snapshot_metadata(0); - assert_eq!(seg_counts.len(), 2); - assert_eq!(base_ts.len(), 2); + let (shared, inits) = ShardDatabases::new(dbs); + assert_eq!(shared.num_shards(), 2); + assert_eq!(inits.len(), 2); + assert_eq!(inits[0].shard_id, 0); + assert_eq!(inits[1].shard_id, 1); } // ── GAP-1: elastic budget publish/recompute ────────────────────────── @@ -958,19 +728,14 @@ mod tests { #[test] fn elastic_budget_defaults_to_zero_until_recomputed() { - let shared = ShardDatabases::new(vec![vec![Database::new()], vec![Database::new()]]); + let shared = new_shared(2, 1); assert_eq!(shared.elastic_budget(0), 0); assert_eq!(shared.elastic_budget(1), 0); } #[test] fn recompute_elastic_budget_hot_shard_borrows_idle_headroom() { - let shared = ShardDatabases::new(vec![ - vec![Database::new()], - vec![Database::new()], - vec![Database::new()], - vec![Database::new()], - ]); + let shared = new_shared(4, 1); let rt = rt_config(400, 4); // base = 100 per shard shared.publish_memory(0, 120); // hot shared.publish_memory(1, 10); @@ -985,120 +750,13 @@ mod tests { #[test] fn recompute_elastic_budget_disabled_for_single_shard_or_unlimited() { - let shared = ShardDatabases::new(vec![vec![Database::new()]]); + let shared = new_shared(1, 1); let rt = rt_config(400, 1); assert_eq!(shared.recompute_elastic_budget(0, &rt), 0); - let shared2 = ShardDatabases::new(vec![vec![Database::new()], vec![Database::new()]]); + let shared2 = new_shared(2, 1); let unlimited = rt_config(0, 2); assert_eq!(shared2.recompute_elastic_budget(0, &unlimited), 0); assert_eq!(shared2.elastic_budget(0), 0); } - - #[test] - fn test_empty_shard_databases() { - let dbs: Vec> = vec![]; - let shared = ShardDatabases::new(dbs); - assert_eq!(shared.num_shards(), 0); - assert_eq!(shared.db_count(), 0); - } - - // ── T2.1 SWAPDB: swap_dbs unit tests ───────────────────────────────────── - - /// Insert a string key into a Database via cmd_dispatch so we don't need - /// to know internal storage layout. - fn db_set_key(db: &mut Database, key: &[u8], value: &[u8]) { - use crate::protocol::Frame; - let mut selected = 0usize; - let args = crate::framevec![ - Frame::BulkString(bytes::Bytes::copy_from_slice(key)), - Frame::BulkString(bytes::Bytes::copy_from_slice(value)), - ]; - let _ = crate::command::dispatch(db, b"SET", &args, &mut selected, 16); - } - - #[test] - fn test_swap_dbs_exchanges_contents() { - // Shard 0 has 2 databases. Put "key_a" in db-0 and "key_b" in db-1. - let mut db0 = Database::new(); - let db1 = Database::new(); - db_set_key(&mut db0, b"key_a", b"val_a"); - - // Put key_b in db-1 using a temporary binding. - let mut db1_tmp = db1; - db_set_key(&mut db1_tmp, b"key_b", b"val_b"); - - let shared = ShardDatabases::new(vec![vec![db0, db1_tmp]]); - assert_eq!( - shared.read_db(0, 0).len(), - 1, - "db-0 should have 1 key before swap" - ); - assert_eq!( - shared.read_db(0, 1).len(), - 1, - "db-1 should have 1 key before swap" - ); - - shared.swap_dbs(0, 0, 1); - - // After swap: db-0 should have key_b, db-1 should have key_a. - assert_eq!( - shared.read_db(0, 0).len(), - 1, - "db-0 should still have 1 key after swap" - ); - assert_eq!( - shared.read_db(0, 1).len(), - 1, - "db-1 should still have 1 key after swap" - ); - - // Verify key moved: db-0 now has the key from the former db-1. - let mut dummy_selected = 0usize; - let get_args = crate::framevec![crate::protocol::Frame::BulkString( - bytes::Bytes::from_static(b"key_b") - )]; - { - let mut guard = shared.write_db(0, 0); - let result = - crate::command::dispatch(&mut *guard, b"GET", &get_args, &mut dummy_selected, 16); - match result { - crate::command::DispatchResult::Response(crate::protocol::Frame::BulkString(v)) => { - assert_eq!(v.as_ref(), b"val_b", "db-0 should have val_b after swap"); - } - crate::command::DispatchResult::Response(_) => { - panic!("expected BulkString(val_b) for key_b in db-0 after swap"); - } - crate::command::DispatchResult::Quit(_) => { - panic!("unexpected Quit result"); - } - } - } - } - - #[test] - fn test_swap_dbs_reverse_order_same_result() { - // Swapping (1, 0) should produce the same result as (0, 1). - let mut db0 = Database::new(); - let db1 = Database::new(); - db_set_key(&mut db0, b"alpha", b"a"); - - let shared = ShardDatabases::new(vec![vec![db0, db1]]); - assert_eq!(shared.read_db(0, 0).len(), 1); - assert_eq!(shared.read_db(0, 1).len(), 0); - - shared.swap_dbs(0, 1, 0); // reversed argument order - - assert_eq!( - shared.read_db(0, 0).len(), - 0, - "db-0 should be empty after swap" - ); - assert_eq!( - shared.read_db(0, 1).len(), - 1, - "db-1 should have 1 key after swap" - ); - } } diff --git a/src/shard/slice.rs b/src/shard/slice.rs index ee9e3cb3..ae6961e0 100644 --- a/src/shard/slice.rs +++ b/src/shard/slice.rs @@ -36,12 +36,12 @@ use std::sync::atomic::AtomicUsize; use crate::graph::store::GraphStore; use crate::mq::{DurableQueueRegistry, TriggerRegistry}; use crate::runtime::channel::MpscSender; +use crate::shard::shared_databases::ShardStoreMemory; use crate::storage::Database; use crate::temporal::{TemporalKvIndex, TemporalRegistry}; use crate::text::store::TextStore; use crate::transaction::{DeferredHnswInserts, KvWriteIntents}; use crate::vector::store::VectorStore; -use crate::workspace::WorkspaceRegistry; /// Per-shard owned aggregate. `!Send` by construction. /// @@ -82,9 +82,6 @@ pub struct ShardSlice { /// Temporal KV index for versioned KV reads. Lazy-init: `None` until the /// first `TemporalUpsert` WAL write on this shard. pub temporal_kv_index: Option>, - /// Workspace registry for workspace metadata. Lazy-init: `None` until the - /// first `WS.CREATE` call on this shard. - pub workspace_registry: Option>, /// Durable queue registry for MQ.* commands. Lazy-init: `None` until the /// first `MQ.CREATE` call on this shard. pub durable_queue_registry: Option>, @@ -100,6 +97,12 @@ pub struct ShardSlice { /// Phase 3 will switch the eviction tick to read this atomic instead of /// calling `aggregate_memory()`, which acquires per-DB read locks. pub estimated_memory: Arc, + /// Per-shard published store-memory atomics (C5 / M4). + /// + /// Shared with `ShardDatabases::store_memory_per_shard[shard_id]`. The + /// shard refreshes vector/text/graph bytes on its 100ms tick; cross-thread + /// observers (metrics, MEMORY DOCTOR) read without locks. + pub store_memory: Arc, /// Makes `ShardSlice` unconditionally `!Send` and `!Sync`. /// /// `Rc<()>` is neither `Send` nor `Sync`, so any type containing @@ -125,13 +128,55 @@ pub struct ShardSliceInit { pub deferred_hnsw_inserts: DeferredHnswInserts, pub temporal_registry: Option>, pub temporal_kv_index: Option>, - pub workspace_registry: Option>, pub durable_queue_registry: Option>, pub trigger_registry: Option>, pub wal_append_tx: Option>, /// A clone of `ShardDatabases::memory_per_shard[shard_id]`. The master /// `Arc` lives in `ShardDatabases`; this is a second owner. pub estimated_memory: Arc, + /// A clone of `ShardDatabases::store_memory_per_shard[shard_id]`. + pub store_memory: Arc, +} + +impl ShardSliceInit { + /// Build one init package per shard from pre-restored databases plus + /// clones of the shared masters' atomics (shardslice-migration C1). + /// + /// Called by `ShardDatabases::new` at boot; each package is handed to its + /// shard thread's spawn closure BY MOVE and consumed by `init_shard`. + /// Stores and registries start empty/lazy — recovery has already replayed + /// into `shard_databases` before this point, and store recovery (vector/ + /// text/graph) happens on the shard thread after `init_shard`. + pub fn build_all( + shard_databases: Vec>, + memory_per_shard: &[Arc], + store_memory_per_shard: &[Arc], + ) -> Vec { + shard_databases + .into_iter() + .enumerate() + .map(|(shard_id, dbs)| { + let databases: Box<[Database]> = dbs.into_boxed_slice(); + ShardSliceInit { + shard_id, + databases, + vector_store: VectorStore::new(), + text_store: TextStore::new(), + #[cfg(feature = "graph")] + graph_store: GraphStore::new(), + kv_write_intents: KvWriteIntents::new(), + deferred_hnsw_inserts: DeferredHnswInserts::new(), + temporal_registry: None, + temporal_kv_index: None, + durable_queue_registry: None, + trigger_registry: None, + wal_append_tx: None, + estimated_memory: memory_per_shard[shard_id].clone(), + store_memory: store_memory_per_shard[shard_id].clone(), + } + }) + .collect() + } } impl ShardSlice { @@ -151,11 +196,11 @@ impl ShardSlice { deferred_hnsw_inserts: init.deferred_hnsw_inserts, temporal_registry: init.temporal_registry, temporal_kv_index: init.temporal_kv_index, - workspace_registry: init.workspace_registry, durable_queue_registry: init.durable_queue_registry, trigger_registry: init.trigger_registry, wal_append_tx: init.wal_append_tx, estimated_memory: init.estimated_memory, + store_memory: init.store_memory, _not_send: PhantomData, } } @@ -180,6 +225,7 @@ thread_local! { /// Panics if called a second time on the same thread — double-initialization /// indicates a programming error in shard startup. pub fn init_shard(slice: ShardSlice) { + let shard_id = slice.shard_id; SHARD.with(|cell| { let mut guard = cell.borrow_mut(); if guard.is_some() { @@ -191,6 +237,26 @@ pub fn init_shard(slice: ShardSlice) { } *guard = Some(slice); }); + // C1 (shardslice-migration): one line per shard, before its first accept. + // test_ssm1_slice_live_at_startup pins this exact marker + shard id. + // The shard id is embedded directly in the message (not as a structured + // tracing field) so that the test can match "shard=N" as a contiguous + // substring — ANSI colour codes inserted around structured-field `=` + // separators break plain string matching of `"shard=0"`. + tracing::info!("ShardSlice initialized shard={}", shard_id); +} + +/// Test-only: forcibly replace the thread-local ShardSlice with a fresh one. +/// +/// Production code MUST use `init_shard` (which panics on double-init). +/// Tests often run multiple test functions on the same OS thread; this helper +/// lets each test reset to a clean state without triggering the double-init +/// guard. Only compiled in `#[cfg(test)]` — never available in production. +#[cfg(test)] +pub fn reset_test_shard(slice: ShardSlice) { + SHARD.with(|cell| { + *cell.borrow_mut() = Some(slice); + }); } // ── Accessor API ────────────────────────────────────────────────────────────── @@ -281,6 +347,33 @@ pub fn is_initialized() -> bool { }) } +/// Assert that `init_shard` has been called on the current thread; abort the +/// process with the shard id if it has not. +/// +/// Call this ONCE at the entry of the shard accept/drain loop, after +/// `init_shard(init)` is called and before the first connection is accepted. +/// This is the C1 startup-abort guard: an uninitialized shard fails fast at +/// startup (not per-command), so no command is ever answered on an +/// uninitialized thread. +/// +/// Using this helper keeps `is_initialized()` calls out of production +/// dispatch paths (the shape test forbids `is_initialized()` outside +/// slice.rs). +#[inline] +pub fn assert_initialized(shard_id: usize) { + if !is_initialized() { + // Startup abort: the shard event loop entered without init_shard. + // Panic here (before the first accept) so the process terminates + // before any command can be answered on an uninitialized thread. + panic!( + "ShardSlice not initialized on shard thread {} — \ + init_shard must be called before the accept/drain loop. \ + This is a startup-configuration bug, not a runtime error.", + shard_id + ); + } +} + // ── Unit tests ──────────────────────────────────────────────────────────────── #[cfg(test)] @@ -307,11 +400,15 @@ mod tests { deferred_hnsw_inserts: DeferredHnswInserts::new(), temporal_registry: None, temporal_kv_index: None, - workspace_registry: None, durable_queue_registry: None, trigger_registry: None, wal_append_tx: None, estimated_memory: Arc::new(AtomicUsize::new(0)), + store_memory: Arc::new(crate::shard::shared_databases::ShardStoreMemory { + vector: AtomicUsize::new(0), + text: AtomicUsize::new(0), + graph: AtomicUsize::new(0), + }), } } diff --git a/src/shard/spsc_handler.rs b/src/shard/spsc_handler.rs index b4420b2f..52ffd4d8 100644 --- a/src/shard/spsc_handler.rs +++ b/src/shard/spsc_handler.rs @@ -64,7 +64,6 @@ pub(crate) fn drain_spsc_shared( crate::shard::dispatch::RawSocketFd, crate::server::conn::affinity::MigratedConnectionState, )>, - vector_store: &mut VectorStore, pending_cdc_subscribes: &mut Vec, // P8: optional manifest for VACUUM manifest/WAL passes; None when no persistence_dir. shard_manifest: &mut Option, @@ -192,7 +191,6 @@ pub(crate) fn drain_spsc_shared( shard_id, script_cache, cached_clock, - vector_store, shard_manifest, mvcc_prune_margin, graph_merge_max_segments, @@ -224,7 +222,6 @@ pub(crate) fn drain_spsc_shared( shard_id, script_cache, cached_clock, - vector_store, shard_manifest, mvcc_prune_margin, graph_merge_max_segments, @@ -270,7 +267,6 @@ pub(crate) fn handle_shard_message_shared( shard_id: usize, script_cache: &Rc>, cached_clock: &CachedClock, - vector_store: &mut VectorStore, // P8: optional manifest for VACUUM manifest/WAL passes; None when no persistence_dir. shard_manifest: &mut Option, // P8: MVCC committed-prune margin from server config (default 1000). @@ -307,10 +303,10 @@ pub(crate) fn handle_shard_message_shared( // Intercept before cmd_dispatch so the console gateway's // ShardMessage::Execute path reaches the vector handlers. if cmd.len() > 3 && cmd[..3].eq_ignore_ascii_case(b"FT.") { - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. - // graph_store, write_db(0), and text_store acquired in one closure - // to avoid re-entrant with_shard calls (multi-resource arm). - let frame = if crate::shard::slice::is_initialized() { + // All slice fields (vector_store, text_store, graph_store, databases) + // acquired in one flat with_shard closure — no outer borrow active, + // so this cannot re-enter with_shard. + let frame = { crate::shard::slice::with_shard(|s| { // SESSION clause needs Database access for sorted set storage. // Only acquire write lock when SESSION keyword is present. @@ -331,7 +327,7 @@ pub(crate) fn handle_shard_message_shared( None }; dispatch_vector_command( - vector_store, + &mut s.vector_store, &mut s.text_store, #[cfg(feature = "graph")] Some(&s.graph_store), @@ -339,33 +335,6 @@ pub(crate) fn handle_shard_message_shared( db_opt, ) }) - } else { - #[cfg(feature = "graph")] - let graph_guard = shard_databases.graph_store_read(shard_id); - - let needs_db = cmd.eq_ignore_ascii_case(b"FT.RECOMMEND") - || cmd.eq_ignore_ascii_case(b"FT.AGGREGATE") - || cmd.eq_ignore_ascii_case(b"FT.DROPINDEX") - || ((cmd.eq_ignore_ascii_case(b"FT.SEARCH") - || cmd.eq_ignore_ascii_case(b"FT.NAVIGATE")) - && has_session_keyword(&command)); - let mut db_guard; - let db_opt = if needs_db { - db_guard = shard_databases.write_db(shard_id, 0); - Some(&mut *db_guard) - } else { - None - }; - - let mut text_guard = shard_databases.text_store(shard_id); - dispatch_vector_command( - vector_store, - &mut *text_guard, - #[cfg(feature = "graph")] - Some(&graph_guard), - &command, - db_opt, - ) }; let _ = reply_tx.send(frame); return; @@ -378,16 +347,19 @@ pub(crate) fn handle_shard_message_shared( if let Some(crate::protocol::Frame::BulkString(sub)) = args.first() { if sub.eq_ignore_ascii_case(b"VECTOR") { let idx_args = &args[1..]; - let frame = - crate::command::server_admin::vacuum_vector(vector_store, idx_args); + let frame = crate::shard::slice::with_shard(|s| { + crate::command::server_admin::vacuum_vector( + &mut s.vector_store, + idx_args, + ) + }); let _ = reply_tx.send(frame); return; } #[cfg(feature = "graph")] if sub.eq_ignore_ascii_case(b"GRAPH") { let graph_args = &args[1..]; - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. - let frame = if crate::shard::slice::is_initialized() { + let frame = { crate::shard::slice::with_shard(|s| { crate::command::server_admin::vacuum_graph( &mut s.graph_store, @@ -396,14 +368,6 @@ pub(crate) fn handle_shard_message_shared( graph_dead_edge_trigger, ) }) - } else { - let mut gs = shard_databases.graph_store_write(shard_id); - crate::command::server_admin::vacuum_graph( - &mut gs, - graph_args, - graph_merge_max_segments, - graph_dead_edge_trigger, - ) }; let _ = reply_tx.send(frame); return; @@ -415,8 +379,7 @@ pub(crate) fn handle_shard_message_shared( // GRAPH.* commands route to GraphStore. #[cfg(feature = "graph")] if cmd.len() > 6 && cmd[..6].eq_ignore_ascii_case(b"GRAPH.") { - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. - let (frame, wal_records) = if crate::shard::slice::is_initialized() { + let (frame, wal_records) = { crate::shard::slice::with_shard(|s| { let resp = crate::command::graph::dispatch_graph_command( &mut s.graph_store, @@ -425,11 +388,6 @@ pub(crate) fn handle_shard_message_shared( let records = s.graph_store.drain_wal(); (resp, records) }) - } else { - let mut gs = shard_databases.graph_store_write(shard_id); - let resp = crate::command::graph::dispatch_graph_command(&mut gs, &command); - let records = gs.drain_wal(); - (resp, records) }; for record in wal_records { shard_databases.wal_append(shard_id, bytes::Bytes::from(record)); @@ -442,7 +400,9 @@ pub(crate) fn handle_shard_message_shared( // Routes directly to VectorStore's TransactionManager; bypasses // write-stall guards (it is an admin command, never a data write). if cmd.eq_ignore_ascii_case(b"KILL") { - let frame = crate::command::server_admin::kill_snapshot(vector_store, args); + let frame = crate::shard::slice::with_shard(|s| { + crate::command::server_admin::kill_snapshot(&mut s.vector_store, args) + }); let _ = reply_tx.send(frame); return; } @@ -450,13 +410,15 @@ pub(crate) fn handle_shard_message_shared( // P8: VACUUM — manual reclamation across manifest, MVCC, and WAL. // Bypasses write-stall guards (reclaims, does not write data). if cmd.eq_ignore_ascii_case(b"VACUUM") { - let frame = crate::command::server_admin::vacuum( - vector_store, - shard_manifest.as_mut(), - wal_v3_writer.as_mut(), - args, - mvcc_prune_margin, - ); + let frame = crate::shard::slice::with_shard(|s| { + crate::command::server_admin::vacuum( + &mut s.vector_store, + shard_manifest.as_mut(), + wal_v3_writer.as_mut(), + args, + mvcc_prune_margin, + ) + }); let _ = reply_tx.send(frame); return; } @@ -465,13 +427,15 @@ pub(crate) fn handle_shard_message_shared( // Intercept here so it has access to manifest and WAL (read-only). if cmd.eq_ignore_ascii_case(b"DEBUG") { if let Some(sub) = args.first() { - if let Some(s) = crate::command::helpers::extract_bytes(sub) { - if s.eq_ignore_ascii_case(b"RECLAMATION") { - let frame = crate::command::server_admin::debug_reclamation( - vector_store, - shard_manifest.as_ref(), - wal_v3_writer.as_ref(), - ); + if let Some(sub_bytes) = crate::command::helpers::extract_bytes(sub) { + if sub_bytes.eq_ignore_ascii_case(b"RECLAMATION") { + let frame = crate::shard::slice::with_shard(|s| { + crate::command::server_admin::debug_reclamation( + &s.vector_store, + shard_manifest.as_ref(), + wal_v3_writer.as_ref(), + ) + }); let _ = reply_tx.send(frame); return; } @@ -515,24 +479,9 @@ pub(crate) fn handle_shard_message_shared( // 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). - if crate::shard::slice::is_initialized() { - 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) - }, - ) - }) - } else { - // Lock ordering (lower index first) prevents deadlock with - // handler_monoio/sharded connections on the same shard. - ksmv::with_two_dbs_locked( - &shard_databases.all_shard_dbs()[shard_id], + crate::shard::slice::with_shard(|s| { + ksmv::with_two_slice_dbs( + &mut s.databases, db_idx, dst_db, |src, dst| { @@ -541,7 +490,7 @@ pub(crate) fn handle_shard_message_shared( ksmv::move_core(src, dst, &key) }, ) - } + }) } }; if matches!(response, crate::protocol::Frame::Integer(1)) { @@ -570,28 +519,9 @@ 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. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - ksmv::with_two_slice_dbs( - &mut s.databases, - db_idx, - ca.dst_db, - |src, dst| { - src.refresh_now_from_cache(cached_clock); - dst.refresh_now_from_cache(cached_clock); - ksmv::copy_core( - src, - dst, - &ca.src_key, - &ca.dst_key, - ca.replace, - ) - }, - ) - }) - } else { - ksmv::with_two_dbs_locked( - &shard_databases.all_shard_dbs()[shard_id], + crate::shard::slice::with_shard(|s| { + ksmv::with_two_slice_dbs( + &mut s.databases, db_idx, ca.dst_db, |src, dst| { @@ -606,7 +536,7 @@ pub(crate) fn handle_shard_message_shared( ) }, ) - } + }) } }; if matches!(response, crate::protocol::Frame::Integer(1)) { @@ -630,10 +560,9 @@ 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); - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. // 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 crate::shard::slice::is_initialized() { + let frame = { if is_write { crate::shard::slice::with_shard_db(db_idx, |db| { cow_intercept(snapshot_state, db, db_idx, &command); @@ -703,7 +632,7 @@ pub(crate) fn handle_shard_message_shared( // Auto-index: if HSET succeeded and key matches a vector index prefix, // extract the vector field and append to mutable segment. - // text_store accessed here (same with_shard closure) to avoid re-entrancy. + // 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(_)) { @@ -715,7 +644,7 @@ pub(crate) fn handle_shard_message_shared( // VectorIntents on the active CrossStoreTxn. Discarded // here because this path is not txn-aware yet. let _ = auto_index_hset( - vector_store, + &mut s.vector_store, &mut s.text_store, key_bytes, args, @@ -726,118 +655,20 @@ pub(crate) fn handle_shard_message_shared( frame }) - } else { - if is_write { - let db_guard = shard_databases.read_db(shard_id, db_idx); - cow_intercept(snapshot_state, &db_guard, db_idx, &command); - drop(db_guard); - } - - let mut guard = shard_databases.write_db(shard_id, db_idx); - guard.refresh_now_from_cache(cached_clock); - let mut selected = db_idx; - let result = cmd_dispatch(&mut guard, 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 - if is_write && !matches!(frame, crate::protocol::Frame::Error(_)) { - let serialized = aof::serialize_command(&command); - wal_append_and_fanout( - &serialized, - wal_writer, - wal_v3_writer, - repl_backlog, - replica_txs, - repl_state, - shard_id, - aof_pool, // FIX-W1-2 - ); - } - - // 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 { - // For LMOVE, wake waiters on the destination key (args[1]), - // not the source key (args[0]). - 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, &mut guard, db_idx, &key, - ); - } else if cmd.eq_ignore_ascii_case(b"ZADD") { - crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, &mut guard, db_idx, &key, - ); - } else { - crate::blocking::wakeup::try_wake_stream_waiter( - &mut reg, &mut guard, db_idx, &key, - ); - } - } - } - } - - // Auto-index: if HSET succeeded and key matches a vector index prefix, - // extract the vector field and append to mutable segment. - 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() { - let mut ts = shard_databases.text_store(shard_id); - // 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(vector_store, &mut *ts, key_bytes, args, 0); - } - } - - // Auto-delete: if DEL/UNLINK succeeded and key matches a vector - // index prefix, mark stale vectors as deleted in matching indexes. - // Note: HDEL removes fields, not keys — it should NOT trigger vector - // deletion unless the entire key is removed. - if (cmd.eq_ignore_ascii_case(b"DEL") || cmd.eq_ignore_ascii_case(b"UNLINK")) - && !matches!(frame, crate::protocol::Frame::Error(_)) - { - for arg in args { - if let crate::protocol::Frame::BulkString(key_bytes) = arg { - vector_store.mark_deleted_for_key(key_bytes); - } - } - } - - drop(guard); - frame }; // Auto-delete is a vector_store-only operation; runs outside the gate. + // Each arm uses its own flat with_shard borrow — no outer borrow is active. if (cmd.eq_ignore_ascii_case(b"DEL") || cmd.eq_ignore_ascii_case(b"UNLINK")) && !matches!(frame, crate::protocol::Frame::Error(_)) { - for arg in args { - if let crate::protocol::Frame::BulkString(key_bytes) = arg { - vector_store.mark_deleted_for_key(key_bytes); + crate::shard::slice::with_shard(|s| { + for arg in args.iter() { + if let crate::protocol::Frame::BulkString(key_bytes) = arg { + s.vector_store.mark_deleted_for_key(key_bytes.as_ref()); + } } - } + }); } frame @@ -852,87 +683,7 @@ 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)); - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(db_idx, |guard| { - guard.refresh_now_from_cache(cached_clock); - for (_key, cmd_frame) in &commands { - let (cmd, args) = match extract_command_static(cmd_frame) { - Some(pair) => pair, - None => { - results.push(crate::protocol::Frame::Error( - bytes::Bytes::from_static(b"ERR invalid command format"), - )); - continue; - } - }; - - let is_write = metadata::is_write(cmd); - if is_write { - cow_intercept(snapshot_state, guard, db_idx, cmd_frame); - } - - let mut selected = db_idx; - let result = cmd_dispatch(guard, cmd, args, &mut selected, db_count); - let frame = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => f, - }; - - if is_write && !matches!(frame, crate::protocol::Frame::Error(_)) { - let serialized = aof::serialize_command(cmd_frame); - wal_append_and_fanout( - &serialized, - wal_writer, - wal_v3_writer, - repl_backlog, - replica_txs, - repl_state, - shard_id, - aof_pool, // FIX-W1-2 - ); - - 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, guard, db_idx, &key, - ); - } else if cmd.eq_ignore_ascii_case(b"ZADD") { - crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, guard, db_idx, &key, - ); - } else { - crate::blocking::wakeup::try_wake_stream_waiter( - &mut reg, guard, db_idx, &key, - ); - } - } - } - } - - results.push(frame); - } - }); - let _ = reply_tx.send(results); - } else { - let mut guard = shard_databases.write_db(shard_id, db_idx); + crate::shard::slice::with_shard_db(db_idx, |guard| { guard.refresh_now_from_cache(cached_clock); for (_key, cmd_frame) in &commands { let (cmd, args) = match extract_command_static(cmd_frame) { @@ -945,20 +696,18 @@ pub(crate) fn handle_shard_message_shared( } }; - // COW intercept for each write command in the batch let is_write = metadata::is_write(cmd); if is_write { - cow_intercept(snapshot_state, &guard, db_idx, cmd_frame); + cow_intercept(snapshot_state, guard, db_idx, cmd_frame); } let mut selected = db_idx; - let result = cmd_dispatch(&mut guard, cmd, args, &mut selected, db_count); + let result = cmd_dispatch(guard, 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 if is_write && !matches!(frame, crate::protocol::Frame::Error(_)) { let serialized = aof::serialize_command(cmd_frame); wal_append_and_fanout( @@ -972,7 +721,6 @@ pub(crate) fn handle_shard_message_shared( aof_pool, // FIX-W1-2 ); - // Wake blocked waiters for producer commands (same as Execute path) let needs_wake = cmd.eq_ignore_ascii_case(b"LPUSH") || cmd.eq_ignore_ascii_case(b"RPUSH") || cmd.eq_ignore_ascii_case(b"LMOVE") @@ -993,15 +741,15 @@ pub(crate) fn handle_shard_message_shared( || cmd.eq_ignore_ascii_case(b"LMOVE") { crate::blocking::wakeup::try_wake_list_waiter( - &mut reg, &mut guard, db_idx, &key, + &mut reg, guard, db_idx, &key, ); } else if cmd.eq_ignore_ascii_case(b"ZADD") { crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, &mut guard, db_idx, &key, + &mut reg, guard, db_idx, &key, ); } else { crate::blocking::wakeup::try_wake_stream_waiter( - &mut reg, &mut guard, db_idx, &key, + &mut reg, guard, db_idx, &key, ); } } @@ -1010,9 +758,8 @@ pub(crate) fn handle_shard_message_shared( results.push(frame); } - drop(guard); - let _ = reply_tx.send(results); - } // else branch end (is_initialized gate) + }); + let _ = reply_tx.send(results); } ShardMessage::PipelineBatch { db_index, @@ -1022,119 +769,10 @@ 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)); - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. // write_db and text_store (HSET auto-index) accessed in one with_shard // closure to avoid re-entrant borrow (multi-resource arm). - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - let guard = &mut s.databases[db_idx]; - guard.refresh_now_from_cache(cached_clock); - for cmd_frame in &commands { - let (cmd, args) = match extract_command_static(cmd_frame) { - Some(pair) => pair, - None => { - results.push(crate::protocol::Frame::Error( - bytes::Bytes::from_static(b"ERR invalid command format"), - )); - continue; - } - }; - - let is_write = metadata::is_write(cmd); - if is_write { - cow_intercept(snapshot_state, guard, db_idx, cmd_frame); - } - - let mut selected = db_idx; - let result = cmd_dispatch(guard, cmd, args, &mut selected, db_count); - let frame = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => f, - }; - - if is_write && !matches!(frame, crate::protocol::Frame::Error(_)) { - let serialized = aof::serialize_command(cmd_frame); - wal_append_and_fanout( - &serialized, - wal_writer, - wal_v3_writer, - repl_backlog, - replica_txs, - repl_state, - shard_id, - // FIX-W1-2 r2: PipelineBatch AOF is written by the - // connection handler coordinator AFTER collecting the - // shard response (handler_monoio/mod.rs:2004, - // handler_sharded/mod.rs:1703). Passing aof_pool here - // would cause a second write to the same shard's AOF - // file, doubling every cross-shard pipeline entry. - None, - ); - } - - // Auto-index: if HSET succeeded, check for vector index match. - // text_store accessed here (same with_shard closure) to avoid re-entrancy. - 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() - { - // Plan 166-01: Vec<(idx, key_hash)> return discarded - // here; Plan 166-02 threads it into CrossStoreTxn. - let _ = auto_index_hset( - vector_store, - &mut s.text_store, - key_bytes, - args, - 0, - ); - } - } - - // 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, guard, db_idx, &key, - ); - } else if cmd.eq_ignore_ascii_case(b"ZADD") { - crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, guard, db_idx, &key, - ); - } else { - crate::blocking::wakeup::try_wake_stream_waiter( - &mut reg, guard, db_idx, &key, - ); - } - } - } - } - - results.push(frame); - } - }); - let _ = reply_tx.send(results); - } else { - let mut guard = shard_databases.write_db(shard_id, db_idx); + crate::shard::slice::with_shard(|s| { + let guard = &mut s.databases[db_idx]; guard.refresh_now_from_cache(cached_clock); for cmd_frame in &commands { let (cmd, args) = match extract_command_static(cmd_frame) { @@ -1147,20 +785,18 @@ pub(crate) fn handle_shard_message_shared( } }; - // COW intercept for each write command in the batch let is_write = metadata::is_write(cmd); if is_write { - cow_intercept(snapshot_state, &guard, db_idx, cmd_frame); + cow_intercept(snapshot_state, guard, db_idx, cmd_frame); } let mut selected = db_idx; - let result = cmd_dispatch(&mut guard, cmd, args, &mut selected, db_count); + let result = cmd_dispatch(guard, 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 if is_write && !matches!(frame, crate::protocol::Frame::Error(_)) { let serialized = aof::serialize_command(cmd_frame); wal_append_and_fanout( @@ -1171,27 +807,35 @@ pub(crate) fn handle_shard_message_shared( replica_txs, repl_state, shard_id, - // FIX-W1-2 r2: PipelineBatch AOF is handled by the - // connection-handler coordinator after collecting the - // shard response (handler_monoio/mod.rs:2004). Passing - // aof_pool here would produce a duplicate AOF entry for - // every cross-shard pipeline command. - None, + // C4-FOLD-FIX: AOF append MUST happen here (in the SPSC arm, + // before the response is sent) so the append is already in the + // AOF channel when AofFold reads sender.len(). Moving the append + // to the connection handler (after awaiting the response) defers + // it until AFTER drain_spsc_shared returns, so AofFold's + // pending_aof_count undercount by ≥1 and that append escapes + // into the NEW incr → double-apply on restart (+1 after + // restart observed in test_ssm4a_fold_4shard_experimental). + // The handler_monoio cross-shard AOF write is removed to avoid + // the double-write that was the original reason for None. + aof_pool, // FIX-C4-FOLD ); } - // Auto-index: if HSET succeeded, check for vector index match + // Auto-index: if HSET succeeded, check for vector index match. + // text_store and vector_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() { - // Use the `vector_store` parameter (already locked by caller), - // NOT shard_databases.vector_store() which would deadlock - // (parking_lot::Mutex is non-reentrant). - let mut ts = shard_databases.text_store(shard_id); // Plan 166-01: Vec<(idx, key_hash)> return discarded // here; Plan 166-02 threads it into CrossStoreTxn. - let _ = auto_index_hset(vector_store, &mut *ts, key_bytes, args, 0); + let _ = auto_index_hset( + &mut s.vector_store, + &mut s.text_store, + key_bytes, + args, + 0, + ); } } @@ -1203,8 +847,6 @@ pub(crate) fn handle_shard_message_shared( || cmd.eq_ignore_ascii_case(b"ZADD") || cmd.eq_ignore_ascii_case(b"XADD"); if needs_wake { - // For LMOVE, wake waiters on the destination key (args[1]), - // not the source key (args[0]). let wake_key = if cmd.eq_ignore_ascii_case(b"LMOVE") { args.get(1) .and_then(|f| crate::server::connection::extract_bytes(f)) @@ -1219,15 +861,15 @@ pub(crate) fn handle_shard_message_shared( || cmd.eq_ignore_ascii_case(b"LMOVE") { crate::blocking::wakeup::try_wake_list_waiter( - &mut reg, &mut guard, db_idx, &key, + &mut reg, guard, db_idx, &key, ); } else if cmd.eq_ignore_ascii_case(b"ZADD") { crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, &mut guard, db_idx, &key, + &mut reg, guard, db_idx, &key, ); } else { crate::blocking::wakeup::try_wake_stream_waiter( - &mut reg, &mut guard, db_idx, &key, + &mut reg, guard, db_idx, &key, ); } } @@ -1236,9 +878,8 @@ pub(crate) fn handle_shard_message_shared( results.push(frame); } - drop(guard); - let _ = reply_tx.send(results); - } // else branch end (is_initialized gate) + }); + let _ = reply_tx.send(results); } ShardMessage::ExecuteSlotted { db_index, @@ -1262,8 +903,7 @@ pub(crate) fn handle_shard_message_shared( { let is_write = metadata::is_write(cmd); - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. - let frame = if crate::shard::slice::is_initialized() { + let frame = { if is_write { crate::shard::slice::with_shard_db(db_idx, |db| { cow_intercept(snapshot_state, db, db_idx, &command); @@ -1331,74 +971,6 @@ pub(crate) fn handle_shard_message_shared( frame }) - } else { - if is_write { - let db_guard = shard_databases.read_db(shard_id, db_idx); - cow_intercept(snapshot_state, &db_guard, db_idx, &command); - drop(db_guard); - } - - let mut guard = shard_databases.write_db(shard_id, db_idx); - guard.refresh_now_from_cache(cached_clock); - let mut selected = db_idx; - let result = cmd_dispatch(&mut guard, cmd, args, &mut selected, db_count); - let frame = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => f, - }; - - if is_write && !matches!(frame, crate::protocol::Frame::Error(_)) { - let serialized = aof::serialize_command(&command); - wal_append_and_fanout( - &serialized, - wal_writer, - wal_v3_writer, - repl_backlog, - replica_txs, - repl_state, - shard_id, - aof_pool, // FIX-W1-2 - ); - } - - 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, &mut guard, db_idx, &key, - ); - } else if cmd.eq_ignore_ascii_case(b"ZADD") { - crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, &mut guard, db_idx, &key, - ); - } else { - crate::blocking::wakeup::try_wake_stream_waiter( - &mut reg, &mut guard, db_idx, &key, - ); - } - } - } - } - - drop(guard); - frame }; // SAFETY: response_slot points to a valid ResponseSlot (see above). let slot = unsafe { &*response_slot.0 }; @@ -1413,86 +985,7 @@ 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)); - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(db_idx, |guard| { - guard.refresh_now_from_cache(cached_clock); - for (_key, cmd_frame) in &commands { - let (cmd, args) = match extract_command_static(cmd_frame) { - Some(pair) => pair, - None => { - results.push(crate::protocol::Frame::Error( - bytes::Bytes::from_static(b"ERR invalid command format"), - )); - continue; - } - }; - - let is_write = metadata::is_write(cmd); - if is_write { - cow_intercept(snapshot_state, guard, db_idx, cmd_frame); - } - - let mut selected = db_idx; - let result = cmd_dispatch(guard, cmd, args, &mut selected, db_count); - let frame = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => f, - }; - - if is_write && !matches!(frame, crate::protocol::Frame::Error(_)) { - let serialized = aof::serialize_command(cmd_frame); - wal_append_and_fanout( - &serialized, - wal_writer, - wal_v3_writer, - repl_backlog, - replica_txs, - repl_state, - shard_id, - aof_pool, // FIX-W1-2 - ); - - 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, guard, db_idx, &key, - ); - } else if cmd.eq_ignore_ascii_case(b"ZADD") { - crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, guard, db_idx, &key, - ); - } else { - crate::blocking::wakeup::try_wake_stream_waiter( - &mut reg, guard, db_idx, &key, - ); - } - } - } - } - - results.push(frame); - } - }); - } else { - let mut guard = shard_databases.write_db(shard_id, db_idx); + crate::shard::slice::with_shard_db(db_idx, |guard| { guard.refresh_now_from_cache(cached_clock); for (_key, cmd_frame) in &commands { let (cmd, args) = match extract_command_static(cmd_frame) { @@ -1507,11 +1000,11 @@ pub(crate) fn handle_shard_message_shared( let is_write = metadata::is_write(cmd); if is_write { - cow_intercept(snapshot_state, &guard, db_idx, cmd_frame); + cow_intercept(snapshot_state, guard, db_idx, cmd_frame); } let mut selected = db_idx; - let result = cmd_dispatch(&mut guard, cmd, args, &mut selected, db_count); + let result = cmd_dispatch(guard, cmd, args, &mut selected, db_count); let frame = match result { DispatchResult::Response(f) => f, DispatchResult::Quit(f) => f, @@ -1550,15 +1043,15 @@ pub(crate) fn handle_shard_message_shared( || cmd.eq_ignore_ascii_case(b"LMOVE") { crate::blocking::wakeup::try_wake_list_waiter( - &mut reg, &mut guard, db_idx, &key, + &mut reg, guard, db_idx, &key, ); } else if cmd.eq_ignore_ascii_case(b"ZADD") { crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, &mut guard, db_idx, &key, + &mut reg, guard, db_idx, &key, ); } else { crate::blocking::wakeup::try_wake_stream_waiter( - &mut reg, &mut guard, db_idx, &key, + &mut reg, guard, db_idx, &key, ); } } @@ -1567,8 +1060,7 @@ pub(crate) fn handle_shard_message_shared( results.push(frame); } - drop(guard); - } // else branch end (is_initialized gate) + }); // SAFETY: response_slot points to a valid ResponseSlot (see ExecuteSlotted). let slot = unsafe { &*response_slot.0 }; slot.fill(results); @@ -1581,115 +1073,9 @@ 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)); - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. // write_db and text_store (HSET auto-index) in one with_shard closure. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - let guard = &mut s.databases[db_idx]; - guard.refresh_now_from_cache(cached_clock); - for cmd_frame in &commands { - let (cmd, args) = match extract_command_static(cmd_frame) { - Some(pair) => pair, - None => { - results.push(crate::protocol::Frame::Error( - bytes::Bytes::from_static(b"ERR invalid command format"), - )); - continue; - } - }; - - let is_write = metadata::is_write(cmd); - if is_write { - cow_intercept(snapshot_state, guard, db_idx, cmd_frame); - } - - let mut selected = db_idx; - let result = cmd_dispatch(guard, cmd, args, &mut selected, db_count); - let frame = match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => f, - }; - - if is_write && !matches!(frame, crate::protocol::Frame::Error(_)) { - let serialized = aof::serialize_command(cmd_frame); - wal_append_and_fanout( - &serialized, - wal_writer, - wal_v3_writer, - repl_backlog, - replica_txs, - repl_state, - shard_id, - // FIX-W1-2 r2: PipelineBatchSlotted AOF is written by the - // connection-handler coordinator after collecting the shard - // response (handler_sharded/mod.rs:1703). Passing aof_pool - // here produces a duplicate AOF entry for every cross-shard - // pipeline command (double-write P0 bug). - None, - ); - } - - // Auto-index: if HSET succeeded, check for vector index match. - // text_store in same with_shard closure to avoid re-entrancy. - 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() - { - // Plan 166-01: Vec<(idx, key_hash)> return discarded - // here; Plan 166-02 threads it into CrossStoreTxn. - let _ = auto_index_hset( - vector_store, - &mut s.text_store, - key_bytes, - args, - 0, - ); - } - } - - 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, guard, db_idx, &key, - ); - } else if cmd.eq_ignore_ascii_case(b"ZADD") { - crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, guard, db_idx, &key, - ); - } else { - crate::blocking::wakeup::try_wake_stream_waiter( - &mut reg, guard, db_idx, &key, - ); - } - } - } - } - - results.push(frame); - } - }); - } else { - let mut guard = shard_databases.write_db(shard_id, db_idx); + crate::shard::slice::with_shard(|s| { + let guard = &mut s.databases[db_idx]; guard.refresh_now_from_cache(cached_clock); for cmd_frame in &commands { let (cmd, args) = match extract_command_static(cmd_frame) { @@ -1704,11 +1090,11 @@ pub(crate) fn handle_shard_message_shared( let is_write = metadata::is_write(cmd); if is_write { - cow_intercept(snapshot_state, &guard, db_idx, cmd_frame); + cow_intercept(snapshot_state, guard, db_idx, cmd_frame); } let mut selected = db_idx; - let result = cmd_dispatch(&mut guard, cmd, args, &mut selected, db_count); + let result = cmd_dispatch(guard, cmd, args, &mut selected, db_count); let frame = match result { DispatchResult::Response(f) => f, DispatchResult::Quit(f) => f, @@ -1724,25 +1110,36 @@ pub(crate) fn handle_shard_message_shared( replica_txs, repl_state, shard_id, - // FIX-W1-2 r2: PipelineBatchSlotted AOF (else branch — pre- - // ShardSlice path) is handled by handler_sharded/mod.rs:1703. - // Passing aof_pool here duplicates the AOF entry. - None, + // C4-FOLD-FIX: AOF append MUST happen here, before the + // response_slot is filled, so the append is already in the + // AOF channel when AofFold reads sender.len(). Deferring to + // the connection handler (after slot.fill wakes the handler + // task) means the append arrives AFTER drain_spsc_shared + // returns and AFTER AofFold's sender.len() snapshot, so + // pending_aof_count undercounts by ≥1 → that append escapes + // into the NEW incr → double-apply on restart (+1 observed + // in test_ssm4a_fold_4shard_experimental). The handler's + // cross-shard AOF write (handler_sharded/mod.rs) is removed + // to avoid the double-write this None guard was preventing. + aof_pool, // FIX-C4-FOLD ); } - // Auto-index: if HSET succeeded, check for vector index match + // Auto-index: if HSET succeeded, check for vector index match. + // vector_store and text_store in 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() { - // Use the `vector_store` parameter (already locked by caller), - // NOT shard_databases.vector_store() which would deadlock - // (parking_lot::Mutex is non-reentrant). - let mut ts = shard_databases.text_store(shard_id); // Plan 166-01: Vec<(idx, key_hash)> return discarded // here; Plan 166-02 threads it into CrossStoreTxn. - let _ = auto_index_hset(vector_store, &mut *ts, key_bytes, args, 0); + let _ = auto_index_hset( + &mut s.vector_store, + &mut s.text_store, + key_bytes, + args, + 0, + ); } } @@ -1767,15 +1164,15 @@ pub(crate) fn handle_shard_message_shared( || cmd.eq_ignore_ascii_case(b"LMOVE") { crate::blocking::wakeup::try_wake_list_waiter( - &mut reg, &mut guard, db_idx, &key, + &mut reg, guard, db_idx, &key, ); } else if cmd.eq_ignore_ascii_case(b"ZADD") { crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, &mut guard, db_idx, &key, + &mut reg, guard, db_idx, &key, ); } else { crate::blocking::wakeup::try_wake_stream_waiter( - &mut reg, &mut guard, db_idx, &key, + &mut reg, guard, db_idx, &key, ); } } @@ -1784,8 +1181,7 @@ pub(crate) fn handle_shard_message_shared( results.push(frame); } - drop(guard); - } // else branch end (is_initialized gate) + }); // SAFETY: response_slot points to a valid ResponseSlot (see ExecuteSlotted). let slot = unsafe { &*response_slot.0 }; slot.fill(results); @@ -1837,35 +1233,15 @@ pub(crate) fn handle_shard_message_shared( let mut reg = blocking_registry.borrow_mut(); reg.register(db_index, key.clone(), entry); // Check if data is already available (race: data arrived before registration). - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. - if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard_db(db_index, |guard| { - if guard.exists(&key) { - crate::blocking::wakeup::try_wake_list_waiter( - &mut reg, guard, db_index, &key, - ); - crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, guard, db_index, &key, - ); - crate::blocking::wakeup::try_wake_stream_waiter( - &mut reg, guard, db_index, &key, - ); - } - }); - } else { - let mut guard = shard_databases.write_db(shard_id, db_index); + crate::shard::slice::with_shard_db(db_index, |guard| { if guard.exists(&key) { - crate::blocking::wakeup::try_wake_list_waiter( - &mut reg, &mut guard, db_index, &key, - ); - crate::blocking::wakeup::try_wake_zset_waiter( - &mut reg, &mut guard, db_index, &key, - ); + crate::blocking::wakeup::try_wake_list_waiter(&mut reg, guard, db_index, &key); + crate::blocking::wakeup::try_wake_zset_waiter(&mut reg, guard, db_index, &key); crate::blocking::wakeup::try_wake_stream_waiter( - &mut reg, &mut guard, db_index, &key, + &mut reg, guard, db_index, &key, ); } - } + }); } ShardMessage::BlockCancel { wait_id } => { blocking_registry.borrow_mut().remove_wait(wait_id); @@ -1876,8 +1252,7 @@ pub(crate) fn handle_shard_message_shared( count, reply_tx, } => { - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. - let keys = if crate::shard::slice::is_initialized() { + let keys = { crate::shard::slice::with_shard_db(db_index, |db| { crate::cluster::migration::handle_get_keys_in_slot( std::slice::from_ref(db), @@ -1886,16 +1261,6 @@ pub(crate) fn handle_shard_message_shared( count, ) }) - } else { - let db_guard = shard_databases.read_db(shard_id, db_index); - let keys = crate::cluster::migration::handle_get_keys_in_slot( - std::slice::from_ref(&*db_guard), - 0, - slot, - count, - ); - drop(db_guard); - keys }; let _ = reply_tx.send(keys); } @@ -1918,17 +1283,20 @@ pub(crate) fn handle_shard_message_shared( // no-op and behavior matches the pre-171 path. Route through // `search_local_filtered` with AS_OF threaded in to apply MVCC // filtering against the committed treemap inside `search_local_raw`. - let response = vector_search::search_local_filtered( - vector_store, - &index_name, - &query_blob, - k, - None, - 0, - usize::MAX, - None, - as_of_lsn, - ); + // Flat with_shard borrow — no outer borrow active, no re-entrancy. + let response = crate::shard::slice::with_shard(|s| { + vector_search::search_local_filtered( + &mut s.vector_store, + &index_name, + &query_blob, + k, + None, + 0, + usize::MAX, + None, + as_of_lsn, + ) + }); let _ = reply_tx.send(response); } ShardMessage::DocFreq { @@ -1938,8 +1306,7 @@ pub(crate) fn handle_shard_message_shared( } => { // DFS Phase 1: collect per-term df + total N from this shard's TextIndex. // Returns crate::protocol::Frame::Array with interleaved [term, df, ..., "N", n] per field_query. - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. - let response = if crate::shard::slice::is_initialized() { + let response = { crate::shard::slice::with_shard(|s| match s.text_store.get_index(&index_name) { Some(text_index) => { let mut items: Vec = Vec::new(); @@ -1963,34 +1330,7 @@ pub(crate) fn handle_shard_message_shared( b"ERR unknown index", )), }) - } else { - let text_guard = shard_databases.text_store(shard_id); - let response = match text_guard.get_index(&index_name) { - Some(text_index) => { - let mut items: Vec = Vec::new(); - for (field_idx_opt, terms) in &field_queries { - let fidx = field_idx_opt.unwrap_or(0); - let (term_dfs, n) = text_index.doc_freq_for_terms(fidx, terms); - for (term, df) in term_dfs { - items.push(crate::protocol::Frame::BulkString(bytes::Bytes::from( - term, - ))); - items.push(crate::protocol::Frame::Integer(i64::from(df))); - } - items.push(crate::protocol::Frame::BulkString( - bytes::Bytes::from_static(b"N"), - )); - items.push(crate::protocol::Frame::Integer(i64::from(n))); - } - crate::protocol::Frame::Array(items.into()) - } - None => crate::protocol::Frame::Error(bytes::Bytes::from_static( - b"ERR unknown index", - )), - }; - drop(text_guard); - response - }; // else branch end (is_initialized gate) + }; let _ = reply_tx.send(response); } ShardMessage::TextSearch(payload) => { @@ -2015,9 +1355,8 @@ pub(crate) fn handle_shard_message_shared( // query_terms is Vec — fuzzy/prefix terms use the OR-union expansion path. // Extract plain strings for HIGHLIGHT/SUMMARIZE (needs analyzed term text only). let term_strings: Vec = query_terms.iter().map(|qt| qt.text.clone()).collect(); - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. // text_store and read_db(0) accessed in one with_shard closure (multi-resource). - let response = if crate::shard::slice::is_initialized() { + let response = { crate::shard::slice::with_shard(|s| match s.text_store.get_index(&index_name) { Some(text_index) => { let mut result = @@ -2048,52 +1387,13 @@ pub(crate) fn handle_shard_message_shared( b"ERR unknown index", )), }) - } else { - let text_guard = shard_databases.text_store(shard_id); - let r = match text_guard.get_index(&index_name) { - Some(text_index) => { - let mut result = - crate::command::vector_search::ft_text_search::execute_text_search_with_global_idf( - text_index, - field_idx, - &query_terms, - &global_df, - global_n, - top_k, - offset, - count, - ); - // Apply HIGHLIGHT/SUMMARIZE post-processing in-place. - // Safe to hold both text_guard and db_guard here — synchronous context, - // no .await points. Guards drop at end of this block. - if highlight_opts.is_some() || summarize_opts.is_some() { - let db_guard = shard_databases.read_db(shard_id, 0); - crate::command::vector_search::ft_text_search::apply_post_processing( - &mut result, - &term_strings, - text_index, - &*db_guard, - highlight_opts.as_ref(), - summarize_opts.as_ref(), - ); - // db_guard drops here. - } - result - } - None => crate::protocol::Frame::Error(bytes::Bytes::from_static( - b"ERR unknown index", - )), - }; - // text_guard drops here (end of block). - r }; let _ = reply_tx.send(response); } ShardMessage::VectorCommand { command, reply_tx } => { - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. - // graph_store, write_db(0), and text_store acquired in one closure - // to avoid re-entrant with_shard calls (multi-resource arm). - let response = if crate::shard::slice::is_initialized() { + // All slice fields (vector_store, text_store, graph_store, databases) + // acquired in one flat with_shard closure — no outer borrow active. + let response = { crate::shard::slice::with_shard(|s| { let cmd_bytes = extract_command_static(&command).map(|(c, _)| c); let is_dropindex = cmd_bytes @@ -2107,7 +1407,7 @@ pub(crate) fn handle_shard_message_shared( None }; dispatch_vector_command( - vector_store, + &mut s.vector_store, &mut s.text_store, #[cfg(feature = "graph")] Some(&s.graph_store), @@ -2115,34 +1415,6 @@ pub(crate) fn handle_shard_message_shared( db_opt, ) }) - } else { - #[cfg(feature = "graph")] - let graph_guard = shard_databases.graph_store_read(shard_id); - - // SESSION clause needs Database access for sorted set storage. - // FT.DROPINDEX with DD flag needs Database to delete indexed docs. - let cmd_bytes = extract_command_static(&command).map(|(c, _)| c); - let is_dropindex = cmd_bytes - .map(|c| c.eq_ignore_ascii_case(b"FT.DROPINDEX")) - .unwrap_or(false); - let has_session = has_session_keyword(&command); - let mut db_guard; - let db_opt = if has_session || is_dropindex { - db_guard = shard_databases.write_db(shard_id, 0); - Some(&mut *db_guard) - } else { - None - }; - - let mut text_guard = shard_databases.text_store(shard_id); - dispatch_vector_command( - vector_store, - &mut *text_guard, - #[cfg(feature = "graph")] - Some(&graph_guard), - &command, - db_opt, - ) }; let _ = reply_tx.send(response); } @@ -2150,19 +1422,13 @@ pub(crate) fn handle_shard_message_shared( ShardMessage::GraphCommand { command, reply_tx } => { // GraphCommand is dispatched via connection handlers using ShardDatabases, // not through SPSC. If we receive one here, dispatch it locally. - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. - let (response, wal_records) = if crate::shard::slice::is_initialized() { + let (response, wal_records) = { crate::shard::slice::with_shard(|s| { let resp = crate::command::graph::dispatch_graph_command(&mut s.graph_store, &command); let records = s.graph_store.drain_wal(); (resp, records) }) - } else { - let mut gs = shard_databases.graph_store_write(shard_id); - let resp = crate::command::graph::dispatch_graph_command(&mut gs, &command); - let records = gs.drain_wal(); - (resp, records) }; for record in wal_records { shard_databases.wal_append(shard_id, bytes::Bytes::from(record)); @@ -2180,7 +1446,7 @@ pub(crate) fn handle_shard_message_shared( // Multi-shard TXN.ABORT leg: this shard owns the graphs named in // these ops. Apply the rollback on the local store and append the // drained WAL records here so replay sees them on the owning shard. - let wal_records = if crate::shard::slice::is_initialized() { + let wal_records = { crate::shard::slice::with_shard(|s| { crate::transaction::abort::apply_graph_rollback( &mut s.graph_store, @@ -2189,14 +1455,6 @@ pub(crate) fn handle_shard_message_shared( &graph_intents, ) }) - } else { - let mut gs = shard_databases.graph_store_write(shard_id); - crate::transaction::abort::apply_graph_rollback( - &mut gs, - txn_id, - &graph_undo, - &graph_intents, - ) }; for record in wal_records { shard_databases.wal_append(shard_id, bytes::Bytes::from(record)); @@ -2215,8 +1473,7 @@ pub(crate) fn handle_shard_message_shared( snapshot_lsn, reply_tx, } = *payload; - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. - let response = if crate::shard::slice::is_initialized() { + let response = { crate::shard::slice::with_shard(|s| { crate::graph::cross_shard::handle_graph_traverse( &s.graph_store, @@ -2226,14 +1483,6 @@ pub(crate) fn handle_shard_message_shared( snapshot_lsn, ) }) - } else { - crate::graph::cross_shard::handle_graph_traverse( - &shard_databases.graph_store_read(shard_id), - &graph_name, - &node_ids, - edge_type_filter, - snapshot_lsn, - ) }; let _ = reply_tx.send(response); } @@ -2251,8 +1500,7 @@ pub(crate) fn handle_shard_message_shared( count, reply_tx, } = *payload; - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. - let response = if crate::shard::slice::is_initialized() { + let response = { crate::shard::slice::with_shard(|s| match s.text_store.get_index(&index_name) { Some(text_index) => { let clause = @@ -2273,32 +1521,6 @@ pub(crate) fn handle_shard_message_shared( b"ERR no such index", )), }) - } else { - let text_guard = shard_databases.text_store(shard_id); - let r = match text_guard.get_index(&index_name) { - Some(text_index) => { - let clause = - crate::command::vector_search::ft_text_search::TextQueryClause { - field_name: None, - terms: Vec::new(), - filter: Some(filter), - }; - let results = - crate::command::vector_search::ft_text_search::execute_query_on_index( - text_index, &clause, None, None, top_k, - ); - // Same response shape TextSearch emits — matches what - // the coordinator's merge_text_results consumes. - crate::command::vector_search::ft_text_search::build_text_response( - &results, offset, count, - ) - } - None => crate::protocol::Frame::Error(bytes::Bytes::from_static( - b"ERR no such index", - )), - }; - // text_guard dropped here - r }; let _ = reply_tx.send(response); } @@ -2315,9 +1537,8 @@ pub(crate) fn handle_shard_message_shared( pipeline, reply_tx, } = *payload; - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. // text_store and read_db(0) accessed in one with_shard closure (multi-resource). - let response = if crate::shard::slice::is_initialized() { + let response = { crate::shard::slice::with_shard(|s| { crate::command::vector_search::ft_aggregate::execute_local_partial( &s.text_store, @@ -2327,18 +1548,6 @@ pub(crate) fn handle_shard_message_shared( &s.databases[0], ) }) - } else { - let text_guard = shard_databases.text_store(shard_id); - let db_guard = shard_databases.read_db(shard_id, 0); - let r = crate::command::vector_search::ft_aggregate::execute_local_partial( - &text_guard, - &index_name, - &query, - &pipeline, - &db_guard, - ); - // guards dropped here - r }; let _ = reply_tx.send(response); } @@ -2365,8 +1574,7 @@ pub(crate) fn handle_shard_message_shared( filter, reply_tx, } = *payload; - // Phase 2b: gate on is_initialized(); new path uses ShardSlice. - let response = if crate::shard::slice::is_initialized() { + let response = { crate::shard::slice::with_shard(|s| { let sparse_pair = match (sparse_field.as_ref(), sparse_blob.as_ref()) { (Some(f), Some(b)) => Some((f, b)), @@ -2377,7 +1585,7 @@ pub(crate) fn handle_shard_message_shared( // applies MVCC filtering consistently across shards. // CHANGE F: forward the filter for per-shard pre-fusion filtering. crate::command::vector_search::hybrid_multi::execute_hybrid_search_local_raw_streams( - vector_store, + &mut s.vector_store, &s.text_store, &index_name, &query_terms, @@ -2393,35 +1601,6 @@ pub(crate) fn handle_shard_message_shared( filter.as_ref(), ) }) - } else { - let text_guard = shard_databases.text_store(shard_id); - let sparse_pair = match (sparse_field.as_ref(), sparse_blob.as_ref()) { - (Some(f), Some(b)) => Some((f, b)), - _ => None, - }; - // Phase 171 HYB-02 / SCAT-02: forward the coordinator-resolved - // AS_OF LSN into the raw-streams executor so the dense branch - // applies MVCC filtering consistently across shards. BM25 - // AS_OF coherent post-G-1 (v0.1.10); text-index MVCC upsert-chain - // pending Phase 178 (MVCC-01). - let r = crate::command::vector_search::hybrid_multi::execute_hybrid_search_local_raw_streams( - vector_store, - &text_guard, - &index_name, - &query_terms, - &dense_field, - &dense_blob, - sparse_pair, - weights, - k_per_stream, - top_k, - &global_df, - global_n, - as_of_lsn, - filter.as_ref(), - ); - // text_guard drops here - r }; let _ = reply_tx.send(response); } @@ -2464,6 +1643,113 @@ pub(crate) fn handle_shard_message_shared( // Notify the coordinator that this shard completed its swap. let _ = reply_tx.send(()); } + // ── C2: shardslice-migration Wave A1 ───────────────────────────────── + // + // These four arms are the owner-side execution legs for the new + // ShardMessage variants defined in §C2 of the frozen contract. + // They run AFTER init_shard is wired (Wave B) — until then the + // variants are never sent, so these arms are dead code at runtime + // today. These use ShardSlice directly: + // the owning shard's slice is the authoritative state once slice mode + // is live. + ShardMessage::MqCommand(payload) => { + // MQ domain hop: execute the full MQ.* subcommand arm against the + // owner's ShardSlice. All six subcommands (CREATE/PUSH/POP/ACK/ + // DLQLEN/TRIGGER) are dispatched through `mq_exec::execute_mq_on_owner`, + // which takes the three data fields directly and returns a Frame. + // Destructure first so reply_tx stays here for the send. + let crate::shard::dispatch::MqCommandPayload { + db_index, + key_prefix, + command, + reply_tx, + } = *payload; + let response = + crate::shard::mq_exec::execute_mq_on_owner(db_index, key_prefix, command); + // Ignore send failure: receiver dropped means the client disconnected. + let _ = reply_tx.send(response); + } + ShardMessage::MqTxnMaterialize { + db_index, + intents, + reply_tx, + } => { + // TXN.COMMIT MQ-intent materialize: fold deferred MQ.PUBLISH messages + // onto the owner shard. Mirrors txn.rs:160–167 exactly: + // for intent in intents: get_stream_mut → durable-check → add. + crate::shard::slice::with_shard_db(db_index, |db| { + for intent in &intents { + if let Ok(Some(stream)) = db.get_stream_mut(&intent.queue_key) { + if stream.durable { + let msg_id = stream.next_auto_id(); + stream.add(msg_id, intent.fields.clone()); + } + } + } + }); + // Ignore send failure: receiver dropped means the TXN coordinator + // has already given up (e.g. client disconnect mid-commit). + let _ = reply_tx.send(()); + } + ShardMessage::WsDropCleanup { prefix, reply_tx } => { + // WS.DROP best-effort key cleanup. The connection handler routes + // this to the shard that owns `prefix` (hash-tag co-location). + // db pinned to 0 — matches the lock-path behaviour in both runtimes. + let deleted_count = crate::shard::slice::with_shard_db(0, |db| { + let keys_to_delete: Vec> = db + .keys() + .filter(|k| k.as_bytes().starts_with(prefix.as_ref())) + .map(|k| k.as_bytes().to_vec()) + .collect(); + let count = keys_to_delete.len() as u64; + for key in &keys_to_delete { + db.remove(key); + } + count + }); + // Ignore send failure: caller logs the count but the drop already + // completed; losing the ack is harmless. + let _ = reply_tx.send(deleted_count); + } + ShardMessage::AofFold { reply_tx } => { + // AOF cooperative-snapshot (C4): build an expired-filtered snapshot + // of ALL databases on this shard and reply. The AOF rewrite writer + // blocks on the oneshot; the shard processes this between commands, + // providing the equivalent mutual-exclusion that the old RwLock write + // guards gave (single-threaded event loop = no concurrent mutations). + // + // C4-DRAIN-BOUND: Before building the snapshot, record how many + // messages are currently pending in this shard's AOF channel. Since + // the shard event loop is single-threaded, no new commands execute + // between this read and the reply being sent. Therefore ALL of these + // pending messages are pre-snapshot appends. The AOF writer uses this + // count as the phase-3 mid-drain bound, preventing an infinite drain + // loop under sustained high write load where the channel never empties. + let pending_aof_count = aof_pool.map(|p| p.sender(shard_id).len()).unwrap_or(0); + let now_ms = crate::storage::entry::current_time_ms(); + let snapshot = crate::shard::slice::with_shard(|s| { + let mut dbs = Vec::with_capacity(s.databases.len()); + for db in s.databases.iter() { + let base_ts = db.base_timestamp(); + let mut entries = Vec::new(); + for (key, entry) in db.data().iter() { + if !entry.is_expired_at(base_ts, now_ms) { + entries.push((key.clone(), entry.clone())); + } + } + dbs.push((entries, base_ts)); + } + crate::shard::dispatch::AofFoldSnapshot { + dbs, + pending_aof_count, + } + }); + // Ignore send failure: the AOF writer dropped its receiver + // (e.g. rewrite aborted) — the snapshot is simply discarded. + let _ = reply_tx.send(snapshot); + } + + // ───────────────────────────────────────────────────────────────────── ShardMessage::Shutdown => { info!("Received shutdown via SPSC"); } @@ -3389,7 +2675,8 @@ mod drain_cap_tests { /// which keeps every other dependency inert (no WAL, no snapshot). #[test] fn drain_cap_reports_possible_tail() { - let shard_databases = Arc::new(ShardDatabases::new(vec![vec![Database::new()]])); + let (shard_databases_inner, _inits) = ShardDatabases::new(vec![vec![Database::new()]]); + let shard_databases = Arc::new(shard_databases_inner); let rb = HeapRb::::new(512); let (mut prod, cons) = rb.split(); for i in 0..300u64 { @@ -3414,11 +2701,11 @@ mod drain_cap_tests { let script_cache = Rc::new(RefCell::new(crate::scripting::ScriptCache::new())); let clock = CachedClock::new(); let mut migrations = Vec::new(); - let mut vector_store = VectorStore::new(); let mut cdc = Vec::new(); let mut manifest = None; let mut autovacuum = crate::shard::autovacuum::AutovacuumDaemon::new(Default::default()); + // BlockCancel messages don't touch ShardSlice, so no init_shard needed. // First cycle: 300 queued > 256 cap -> drains exactly 256, reports tail. let hit_cap = drain_spsc_shared( &shard_databases, @@ -3436,7 +2723,6 @@ mod drain_cap_tests { &script_cache, &clock, &mut migrations, - &mut vector_store, &mut cdc, &mut manifest, 1000, @@ -3467,7 +2753,6 @@ mod drain_cap_tests { &script_cache, &clock, &mut migrations, - &mut vector_store, &mut cdc, &mut manifest, 1000, diff --git a/src/shard/timers.rs b/src/shard/timers.rs index 0fcef70a..50498bf4 100644 --- a/src/shard/timers.rs +++ b/src/shard/timers.rs @@ -18,8 +18,9 @@ use super::shared_databases::ShardDatabases; pub(crate) fn run_active_expiry(shard_databases: &Arc, shard_id: usize) { let db_count = shard_databases.db_count(); for i in 0..db_count { - let mut guard = shard_databases.write_db(shard_id, i); - crate::server::expiration::expire_cycle_direct(&mut guard); + crate::shard::slice::with_shard_db(i, |db| { + crate::server::expiration::expire_cycle_direct(db); + }); } // Update RSS gauge on shard 0 only, once per second (not every 100ms tick). // Gated by a simple counter to reduce /proc/self/statm open/read/close churn. @@ -48,8 +49,9 @@ pub(crate) fn run_eviction( let budget = shard_databases.elastic_budget(shard_id); let db_count = shard_databases.db_count(); for i in 0..db_count { - let mut guard = shard_databases.write_db(shard_id, i); - let _ = crate::storage::eviction::try_evict_if_needed_budget(&mut guard, &rt, budget); + crate::shard::slice::with_shard_db(i, |db| { + let _ = crate::storage::eviction::try_evict_if_needed_budget(db, &rt, budget); + }); } } } @@ -94,38 +96,40 @@ pub(crate) fn fire_pending_mq_triggers( now_ms: u64, pubsub_registry: &Arc>, ) { - let mut guard = shard_databases.trigger_registry(shard_id); - let Some(reg) = guard.as_mut() else { return }; - - // Collect keys of triggers ready to fire - let ready_keys = reg.fire_ready(now_ms); + // Collect keys of triggers ready to fire via with_shard (no lock needed on slice path). + let ready_keys: Vec = crate::shard::slice::with_shard(|s| { + let Some(ref mut reg) = s.trigger_registry else { + return Vec::new(); + }; + reg.fire_ready(now_ms) + }); if ready_keys.is_empty() { return; } - // Collect (channel, message) pairs while holding trigger lock, - // then publish after releasing it to avoid holding two locks. - let mut notifications: Vec<(bytes::Bytes, bytes::Bytes)> = Vec::with_capacity(ready_keys.len()); - - for key in &ready_keys { - if let Some(entry) = reg.get(key) { - // Build pub/sub channel name: "mq:trigger:{queue_key}" - let channel = { - let mut ch = Vec::with_capacity(11 + entry.queue_key.len()); - ch.extend_from_slice(b"mq:trigger:"); - ch.extend_from_slice(&entry.queue_key); - bytes::Bytes::from(ch) - }; - notifications.push((channel, entry.callback_cmd.clone())); + // Collect (channel, message) pairs via with_shard while not crossing await. + let notifications: Vec<(bytes::Bytes, bytes::Bytes)> = crate::shard::slice::with_shard(|s| { + let Some(ref mut reg) = s.trigger_registry else { + return Vec::new(); + }; + let mut out = Vec::with_capacity(ready_keys.len()); + for key in &ready_keys { + if let Some(entry) = reg.get(key) { + let channel = { + let mut ch = Vec::with_capacity(11 + entry.queue_key.len()); + ch.extend_from_slice(b"mq:trigger:"); + ch.extend_from_slice(&entry.queue_key); + bytes::Bytes::from(ch) + }; + out.push((channel, entry.callback_cmd.clone())); + } + reg.mark_fired(key, now_ms); } - // Mark as fired (updates last_fire_ms, clears pending_fire_ms) - reg.mark_fired(key, now_ms); - } + out + }); + let _ = shard_databases; // shared handle no longer needed on this path - // Release trigger registry lock before acquiring pubsub lock - drop(guard); - - // Publish each trigger notification via pub/sub + // Publish each trigger notification via pub/sub (outside with_shard — no re-entry). if !notifications.is_empty() { let mut pubsub = pubsub_registry.write(); for (channel, message) in ¬ifications { @@ -169,55 +173,48 @@ pub(crate) fn run_cold_orphan_sweep( let mut total = SweepStats::default(); for db_idx in 0..db_count { - let mut guard = shard_databases.write_db(shard_id, db_idx); - - // Skip databases without a cold index (disk-offload disabled or no spills). - if guard.cold_index.is_none() { - continue; - } + // Phase 1: collect orphan keys and pending-unlink flag inside with_shard_db. + let (orphan_keys, has_pending_unlink) = crate::shard::slice::with_shard_db(db_idx, |db| { + if db.cold_index.is_none() { + return (Vec::new(), false); + } + let orphan_keys: Vec = db + .cold_index + .as_ref() + .map(|ci| { + ci.iter() + .filter(|(key, _loc)| db.is_hot(key)) + .map(|(k, _)| k.clone()) + .collect() + }) + .unwrap_or_default(); + let has_pending_unlink = db + .cold_index + .as_ref() + .map(|ci| ci.has_pending_unlink()) + .unwrap_or(false); + (orphan_keys, has_pending_unlink) + }); - // Two-phase sweep to sidestep the &mut cold_index / &db borrow conflict: - // - // Phase 1: collect orphan keys (immutable borrow of cold_index + db). - // A cold entry is an orphan when the key is also present in the hot - // DashTable — meaning a hot write shadowed the spilled copy. - let orphan_keys: Vec = guard - .cold_index - .as_ref() - .map(|ci| { - ci.iter() - .filter(|(key, _loc)| guard.is_hot(key)) - .map(|(k, _)| k.clone()) - .collect() - }) - .unwrap_or_default(); - - // Skip only when there is nothing to do: no hot-shadowed orphan keys AND - // no zero-ref files queued for unlink. Files orphaned by re-eviction - // (insert overwrite) or promotion (remove) carry no hot∩cold key, so the - // drain must still run for them even when `orphan_keys` is empty. - let has_pending_unlink = guard - .cold_index - .as_ref() - .map(|ci| ci.has_pending_unlink()) - .unwrap_or(false); if orphan_keys.is_empty() && !has_pending_unlink { continue; } - // Phase 2: delete files + update cold_index (mutable borrow; db not borrowed). - let stats = guard.cold_index.as_mut().and_then(|ci| { - ci.sweep_known_orphans(orphan_keys, shard_dir, manifest.as_deref_mut()) - .map_err(|e| { - tracing::error!( - shard = shard_id, - db = db_idx, - err = %e, - "cold_orphan_sweep: manifest commit error", - ); - e - }) - .ok() + // Phase 2: delete files + update cold_index. + let stats = crate::shard::slice::with_shard_db(db_idx, |db| { + db.cold_index.as_mut().and_then(|ci| { + ci.sweep_known_orphans(orphan_keys, shard_dir, manifest.as_deref_mut()) + .map_err(|e| { + tracing::error!( + shard = shard_id, + db = db_idx, + err = %e, + "cold_orphan_sweep: manifest commit error", + ); + e + }) + .ok() + }) }); if let Some(s) = stats { diff --git a/src/shard/uring_handler.rs b/src/shard/uring_handler.rs index c50226b7..de00adb8 100644 --- a/src/shard/uring_handler.rs +++ b/src/shard/uring_handler.rs @@ -81,7 +81,9 @@ pub(crate) fn handle_uring_event( event: IoEvent, driver: &mut UringDriver, shard_databases: &Arc, - shard_id: usize, + // Unused since the ShardSlice cutover: this handler runs on the shard's + // own thread, so the thread-local slice (with_shard_db) IS this shard. + _shard_id: usize, parse_bufs: &mut std::collections::HashMap, inflight_sends: &mut std::collections::HashMap>, uring_listener_fd: Option, @@ -133,15 +135,21 @@ pub(crate) fn handle_uring_event( return; } - // Phase B: Dispatch all commands under a single write lock. - let responses: Vec = { - let mut guard = shard_databases.write_db(shard_id, 0); - let db_count = shard_databases.db_count(); - guard.refresh_now_from_cache(cached_clock); - let mut selected = 0usize; - let result: Vec<_> = batch - .iter() - .map(|frame| { + // Phase B: Dispatch all commands inside a single ShardSlice borrow. + // This handler runs synchronously ON the shard thread, so the + // thread-local slice borrow replaces the old per-batch write lock + // (shardslice-migration C6). The workspace-registry Mutex and + // wal_append channel sends inside the closure are independent of + // the slice borrow — no re-entrancy, no await. + let db_count = shard_databases.db_count(); + let responses: Vec = crate::shard::slice::with_shard_db( + 0, + |db| { + db.refresh_now_from_cache(cached_clock); + let mut selected = 0usize; + batch + .iter() + .map(|frame| { let (cmd, args) = match extract_command_static(frame) { Some(pair) => pair, None => { @@ -276,16 +284,15 @@ pub(crate) fn handle_uring_event( )); } - let result = cmd_dispatch(&mut guard, cmd, args, &mut selected, db_count); - match result { - DispatchResult::Response(f) => f, - DispatchResult::Quit(f) => f, - } - }) - .collect(); - drop(guard); - result - }; + let result = cmd_dispatch(db, cmd, args, &mut selected, db_count); + match result { + DispatchResult::Response(f) => f, + DispatchResult::Quit(f) => f, + } + }) + .collect() + }, + ); // Phase C: Serialize and send all responses (outside borrow). for response in responses { diff --git a/src/transaction/abort.rs b/src/transaction/abort.rs index e78cb0a5..e127be7a 100644 --- a/src/transaction/abort.rs +++ b/src/transaction/abort.rs @@ -90,7 +90,11 @@ use bytes::Bytes; /// `parking_lot` guards; no `.await` points are crossed while any guard is /// held. Follows Phase 161 lock ordering: each layer's guard is dropped /// before the next layer is accessed. +// `shard_databases` and `shard_id` are consumed only under `#[cfg(feature = "graph")]` +// (WAL record drain). Without the graph feature they are structurally unused; +// suppress the lint rather than removing a semantically load-bearing parameter. #[allow(clippy::needless_pass_by_value)] +#[cfg_attr(not(feature = "graph"), allow(unused_variables))] pub fn abort_cross_store_txn( shard_databases: &ShardDatabases, shard_id: usize, @@ -105,21 +109,22 @@ pub fn abort_cross_store_txn( // out of handler_sharded/handler_monoio TXN.ABORT. // ------------------------------------------------------------------ { - let mut db = shard_databases.write_db(shard_id, selected_db); - for record in txn.kv_undo.into_rollback_order() { - match record { - UndoRecord::Insert { key } => { - db.remove(&key); - } - UndoRecord::Update { key, old_entry } => { - db.set(key, old_entry); - } - UndoRecord::Delete { key, old_entry } => { - db.set(key, old_entry); + crate::shard::slice::with_shard_db(selected_db, |db| { + for record in txn.kv_undo.into_rollback_order() { + match record { + UndoRecord::Insert { key } => { + db.remove(&key); + } + UndoRecord::Update { key, old_entry } => { + db.set(key, old_entry); + } + UndoRecord::Delete { key, old_entry } => { + db.set(key, old_entry); + } } } - } - // db guard drops at scope end — releases write_db before next layer. + }); + // with_shard_db releases the borrow at the closure boundary — before next layer. } // ------------------------------------------------------------------ @@ -133,21 +138,15 @@ pub fn abort_cross_store_txn( #[cfg(feature = "graph")] { if !txn.graph_undo.is_empty() || !txn.graph_intents.is_empty() { - let wal_records = if crate::shard::slice::is_initialized() { - crate::shard::slice::with_shard(|s| { - apply_graph_rollback( - &mut s.graph_store, - txn_id, - &txn.graph_undo, - &txn.graph_intents, - ) - }) - } else { - let mut gs = shard_databases.graph_store_write(shard_id); - apply_graph_rollback(&mut gs, txn_id, &txn.graph_undo, &txn.graph_intents) - // gs guard drops at scope end — releases graph_store_write - // before the next layer (vector_store). - }; + // Unconditional slice path: ShardSlice is always initialized. + let wal_records = crate::shard::slice::with_shard(|s| { + apply_graph_rollback( + &mut s.graph_store, + txn_id, + &txn.graph_undo, + &txn.graph_intents, + ) + }); // Drain any WAL records produced by remove_node/remove_edge so // replay observes the rollback in the same ordering the live // write path would have produced. @@ -159,53 +158,39 @@ pub fn abort_cross_store_txn( // ------------------------------------------------------------------ // 3. Vector rollback — tombstone every mutable-HNSW entry appended - // during the transaction. We use the txn-scoped variant - // `mark_deleted_by_key_hash_after_lsn(key_hash, txn.snapshot_lsn)` - // so that earlier-committed rows that happen to share the same - // Redis key (same xxh64 `key_hash`) are NOT rolled back. The - // method sets each matching entry's `delete_lsn = insert_lsn`, - // which makes the entry invisible at every snapshot >= its own - // insert LSN (per MVCC visibility) — the row "never existed" from - // any reader's perspective, which is what TXN.ABORT requires. - // This is the core of ACID-08 (HNSW rollback). + // during the transaction. Uses with_shard so the thread-local + // VectorStore is accessed without a lock. // ------------------------------------------------------------------ { let txn_snapshot_lsn = txn.snapshot_lsn; - let mut vector_store = shard_databases.vector_store(shard_id); - - for intent in &txn.vector_intents { - let Some(idx) = vector_store.get_index_mut(&intent.index_name) else { - tracing::warn!( - txn_id, - index_name = ?intent.index_name, - point_id = intent.point_id, - "txn abort: vector index missing at rollback time, skipping intent", - ); - continue; - }; - let snap = idx.segments.load(); - // The threshold is the transaction's snapshot LSN — only rows - // appended inside the txn (insert_lsn > snapshot_lsn) get - // tombstoned. Preserves pre-txn rows sharing the same key_hash. - let count = snap - .mutable - .mark_deleted_by_key_hash_after_lsn(intent.point_id, txn_snapshot_lsn); - if count == 0 { - tracing::warn!( - txn_id, - index_name = ?intent.index_name, - point_id = intent.point_id, - "txn abort: mark_deleted_by_key_hash_after_lsn matched zero entries (rollback may leak)", - ); + crate::shard::slice::with_shard(|s| { + for intent in &txn.vector_intents { + let Some(idx) = s.vector_store.get_index_mut(&intent.index_name) else { + tracing::warn!( + txn_id, + index_name = ?intent.index_name, + point_id = intent.point_id, + "txn abort: vector index missing at rollback time, skipping intent", + ); + continue; + }; + let snap = idx.segments.load(); + let count = snap + .mutable + .mark_deleted_by_key_hash_after_lsn(intent.point_id, txn_snapshot_lsn); + if count == 0 { + tracing::warn!( + txn_id, + index_name = ?intent.index_name, + point_id = intent.point_id, + "txn abort: mark_deleted_by_key_hash_after_lsn matched zero entries (rollback may leak)", + ); + } } - } - - // Transition the TransactionManager into abort state — emits - // whatever XactAbort-equivalent bookkeeping the manager owns. - vector_store.txn_manager_mut().abort(txn_id); - - // LOCK-ORDER: drop vector_store before kv_intents - drop(vector_store); + // Transition the TransactionManager into abort state. + s.vector_store.txn_manager_mut().abort(txn_id); + // LOCK-ORDER: with_shard releases before kv_intents step. + }); } // ------------------------------------------------------------------ @@ -214,8 +199,10 @@ pub fn abort_cross_store_txn( // HNSW insertions queued for this txn (prevents phantom neighbors // from showing up post-compaction on a txn that never committed). // ------------------------------------------------------------------ - shard_databases.kv_intents(shard_id).release_txn(txn_id); - shard_databases.hnsw_queue(shard_id).discard_for_txn(txn_id); + crate::shard::slice::with_shard(|s| { + s.kv_write_intents.release_txn(txn_id); + s.deferred_hnsw_inserts.discard_for_txn(txn_id); + }); } /// Apply the graph half of a TXN.ABORT to one `GraphStore`. diff --git a/tests/adversarial_v0110_fix01_set_delete_rollback.rs b/tests/adversarial_v0110_fix01_set_delete_rollback.rs index 222a9e60..087099eb 100644 --- a/tests/adversarial_v0110_fix01_set_delete_rollback.rs +++ b/tests/adversarial_v0110_fix01_set_delete_rollback.rs @@ -161,7 +161,8 @@ async fn start_txn_server(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -176,6 +177,7 @@ async fn start_txn_server(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("fix01-shard-{}", id)) @@ -218,6 +220,7 @@ async fn start_txn_server(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/adversarial_v0110_fix02_err_path_intent.rs b/tests/adversarial_v0110_fix02_err_path_intent.rs index cee4d8b3..0f40efdf 100644 --- a/tests/adversarial_v0110_fix02_err_path_intent.rs +++ b/tests/adversarial_v0110_fix02_err_path_intent.rs @@ -164,7 +164,8 @@ async fn start_txn_server(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -179,6 +180,7 @@ async fn start_txn_server(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("fix02-shard-{}", id)) @@ -221,6 +223,7 @@ async fn start_txn_server(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/adversarial_v0110_fix03_simplestring_graph.rs b/tests/adversarial_v0110_fix03_simplestring_graph.rs index 6a79ce29..3e48b94e 100644 --- a/tests/adversarial_v0110_fix03_simplestring_graph.rs +++ b/tests/adversarial_v0110_fix03_simplestring_graph.rs @@ -166,7 +166,8 @@ async fn start_txn_server(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -181,6 +182,7 @@ async fn start_txn_server(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("fix03-shard-{}", id)) @@ -223,6 +225,7 @@ async fn start_txn_server(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/adversarial_v0110_fix04_shortest_path_call_parity.rs b/tests/adversarial_v0110_fix04_shortest_path_call_parity.rs index b37cba0d..dacbb4cc 100644 --- a/tests/adversarial_v0110_fix04_shortest_path_call_parity.rs +++ b/tests/adversarial_v0110_fix04_shortest_path_call_parity.rs @@ -150,7 +150,8 @@ async fn start_server(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -165,6 +166,7 @@ async fn start_server(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("fix04-shard-{}", id)) @@ -207,6 +209,7 @@ async fn start_server(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/adversarial_v0110_fix06_shortest_path_min_hops.rs b/tests/adversarial_v0110_fix06_shortest_path_min_hops.rs index 57fdc796..34f3e51f 100644 --- a/tests/adversarial_v0110_fix06_shortest_path_min_hops.rs +++ b/tests/adversarial_v0110_fix06_shortest_path_min_hops.rs @@ -155,7 +155,8 @@ async fn start_server(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -170,6 +171,7 @@ async fn start_server(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("fix06-shard-{}", id)) @@ -212,6 +214,7 @@ async fn start_server(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/adversarial_v0110_fix07_multihop_edge_var_reject.rs b/tests/adversarial_v0110_fix07_multihop_edge_var_reject.rs index 01f90dba..e85afe7d 100644 --- a/tests/adversarial_v0110_fix07_multihop_edge_var_reject.rs +++ b/tests/adversarial_v0110_fix07_multihop_edge_var_reject.rs @@ -157,7 +157,8 @@ async fn start_server(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -172,6 +173,7 @@ async fn start_server(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("fix07-shard-{}", id)) @@ -214,6 +216,7 @@ async fn start_server(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/crash_matrix_per_shard_bgrewriteaof.rs b/tests/crash_matrix_per_shard_bgrewriteaof.rs index 41e9de2e..37614b4e 100644 --- a/tests/crash_matrix_per_shard_bgrewriteaof.rs +++ b/tests/crash_matrix_per_shard_bgrewriteaof.rs @@ -278,7 +278,15 @@ fn bgrewriteaof_base_plus_incr_recovers_exact() { const PRE: i64 = 300; const POST: i64 = 200; - let port = unique_port().saturating_add(1); + // Plain unique_port(), NOT unique_port()+1: macOS allocates ephemeral + // ports sequentially, so when this test runs in parallel with the + // straddle test, `+1` lands exactly on the port the straddle test's own + // unique_port() call returns next. Both moons then bind the SAME port via + // SO_REUSEPORT (no bind error) and the kernel splits each test's + // redis-cli connections between the two servers — INCRs land on the + // wrong server and the rewrite asserts fire ("rewrite did not compact"). + // Two independent kernel allocations never return the same port. + let port = unique_port(); let dir = unique_dir("compose"); std::fs::create_dir_all(&dir).expect("create test dir"); diff --git a/tests/ft_search_as_of_boundary.rs b/tests/ft_search_as_of_boundary.rs index 7bc29d10..64dbfa29 100644 --- a/tests/ft_search_as_of_boundary.rs +++ b/tests/ft_search_as_of_boundary.rs @@ -155,7 +155,8 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -170,6 +171,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("boundary-shard-{}", id)) @@ -210,6 +212,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/ft_search_as_of_filter.rs b/tests/ft_search_as_of_filter.rs index c68e1e26..73016b4e 100644 --- a/tests/ft_search_as_of_filter.rs +++ b/tests/ft_search_as_of_filter.rs @@ -163,7 +163,8 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -178,6 +179,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("filter-shard-{}", id)) @@ -218,6 +220,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/ft_search_concurrent_readers.rs b/tests/ft_search_concurrent_readers.rs index f0ce851d..5fdbb2b9 100644 --- a/tests/ft_search_concurrent_readers.rs +++ b/tests/ft_search_concurrent_readers.rs @@ -152,7 +152,8 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -167,6 +168,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("concurrent-shard-{}", id)) @@ -207,6 +209,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/ft_search_multi_shard_as_of.rs b/tests/ft_search_multi_shard_as_of.rs index 776b9e92..48cef4b2 100644 --- a/tests/ft_search_multi_shard_as_of.rs +++ b/tests/ft_search_multi_shard_as_of.rs @@ -170,7 +170,8 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -185,6 +186,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("ft-mshard-asof-shard-{}", id)) @@ -227,6 +229,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/ft_search_temporal_parity.rs b/tests/ft_search_temporal_parity.rs index 67679a3f..69ff892a 100644 --- a/tests/ft_search_temporal_parity.rs +++ b/tests/ft_search_temporal_parity.rs @@ -183,7 +183,8 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -198,6 +199,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("ft-parity-shard-{}", id)) @@ -240,6 +242,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/hybrid_filter_backward_compat.rs b/tests/hybrid_filter_backward_compat.rs index 5a9b5cbe..35aaf8f8 100644 --- a/tests/hybrid_filter_backward_compat.rs +++ b/tests/hybrid_filter_backward_compat.rs @@ -153,7 +153,8 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -168,6 +169,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("filter-bc-shard-{}", id)) @@ -210,6 +212,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/hybrid_filter_multishard.rs b/tests/hybrid_filter_multishard.rs index 443d0c8f..e3f2248d 100644 --- a/tests/hybrid_filter_multishard.rs +++ b/tests/hybrid_filter_multishard.rs @@ -151,7 +151,8 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -166,6 +167,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("filter-ms-shard-{}", id)) @@ -208,6 +210,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/hybrid_filter_tag.rs b/tests/hybrid_filter_tag.rs index 3ad76fc9..9cf5802d 100644 --- a/tests/hybrid_filter_tag.rs +++ b/tests/hybrid_filter_tag.rs @@ -148,7 +148,8 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -163,6 +164,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("filter-tag-shard-{}", id)) @@ -205,6 +207,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/integration.rs b/tests/integration.rs index 10f986d0..5bcb0867 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -2749,7 +2749,8 @@ async fn start_sharded_server(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); // Spawn shard threads let mut shard_handles = Vec::with_capacity(num_shards); @@ -2765,6 +2766,7 @@ async fn start_sharded_server(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("test-shard-{}", id)) @@ -2806,6 +2808,7 @@ async fn start_sharded_server(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); @@ -3968,7 +3971,8 @@ async fn start_cluster_server() -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); // Spawn shard threads let mut shard_handles = Vec::with_capacity(num_shards); @@ -3985,6 +3989,7 @@ async fn start_cluster_server() -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("test-cluster-shard-{}", id)) @@ -4026,6 +4031,7 @@ async fn start_cluster_server() -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/kill_snapshot.rs b/tests/kill_snapshot.rs index 4fc946eb..db169cf2 100644 --- a/tests/kill_snapshot.rs +++ b/tests/kill_snapshot.rs @@ -172,7 +172,8 @@ async fn start_sharded_server(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -187,6 +188,7 @@ async fn start_sharded_server(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("kill-snap-shard-{}", id)) @@ -227,6 +229,7 @@ async fn start_sharded_server(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/lunaris_cypher_shortest_path.rs b/tests/lunaris_cypher_shortest_path.rs index 707d4baa..954e29e1 100644 --- a/tests/lunaris_cypher_shortest_path.rs +++ b/tests/lunaris_cypher_shortest_path.rs @@ -206,7 +206,8 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -221,6 +222,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("lunaris-v3-shard-{}", id)) @@ -263,6 +265,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/lunaris_cypher_temporal.rs b/tests/lunaris_cypher_temporal.rs index 1de442a5..3e17ff85 100644 --- a/tests/lunaris_cypher_temporal.rs +++ b/tests/lunaris_cypher_temporal.rs @@ -215,7 +215,8 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -230,6 +231,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("lunaris-v1-shard-{}", id)) @@ -272,6 +274,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/lunaris_hybrid_ft_search.rs b/tests/lunaris_hybrid_ft_search.rs index d812ed56..141a5fba 100644 --- a/tests/lunaris_hybrid_ft_search.rs +++ b/tests/lunaris_hybrid_ft_search.rs @@ -193,7 +193,8 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -208,6 +209,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("lunaris-v2-shard-{}", id)) @@ -250,6 +252,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/mq_integration.rs b/tests/mq_integration.rs index 47ff3261..4b39a5a4 100644 --- a/tests/mq_integration.rs +++ b/tests/mq_integration.rs @@ -159,7 +159,8 @@ async fn start_mq_server(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -174,6 +175,7 @@ async fn start_mq_server(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("mq-test-shard-{}", id)) @@ -216,6 +218,7 @@ async fn start_mq_server(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/pipeline_auto_index.rs b/tests/pipeline_auto_index.rs index 7ae7506e..5df42b8e 100644 --- a/tests/pipeline_auto_index.rs +++ b/tests/pipeline_auto_index.rs @@ -163,7 +163,8 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -178,6 +179,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("pipe-auto-shard-{}", id)) @@ -218,6 +220,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/quickwins_red.rs b/tests/quickwins_red.rs index 4caea526..1c8e5526 100644 --- a/tests/quickwins_red.rs +++ b/tests/quickwins_red.rs @@ -211,7 +211,7 @@ fn qw2_try_wal_append_required_false_on_full() { use moon::shard::shared_databases::ShardDatabases; use moon::storage::db::Database; - let dbs = ShardDatabases::new(vec![vec![Database::new()]]); + let (dbs, _inits) = ShardDatabases::new(vec![vec![Database::new()]]); // Persistence disabled -> no durability requirement -> true. assert!( diff --git a/tests/shardslice_live.rs b/tests/shardslice_live.rs new file mode 100644 index 00000000..a94639ed --- /dev/null +++ b/tests/shardslice_live.rs @@ -0,0 +1,934 @@ +//! ADD task `shardslice-migration` — live-server integration tests (phase: tests). +//! +//! Tests in this file cover scenarios ssm1, ssm3, ssm4a, and the +//! `reject_uninitialized_shard` wire guard. +//! +//! RED today: +//! - test_ssm1_slice_live_at_startup — FAILS because `init_shard` has zero +//! production call sites; no "ShardSlice initialized" line is emitted in +//! shard startup log output (frozen contract C1, M1). +//! +//! GREEN today (behavioral oracles, must keep passing through the cutover): +//! - test_ssm3_ws_global_registry — WS CREATE/AUTH/LIST survive restart; +//! works today via the per-shard lock path; must keep passing after cutover +//! via the process-global registry (contract C3). +//! - test_ssm4a_aof_fold_cooperative — exactly-once INCR across a +//! BGREWRITEAOF boundary; currently correct via RwLock fold; must keep +//! passing after the cooperative-snapshot redesign (contract C4). +//! - test_reject_uninitialized_shard_no_wire_panic — wire-level guard: no +//! junk command + normal SET/GET causes a with_shard panic reachable from +//! the network; the server stays alive. +//! +//! Run alone with: +//! MOON_BIN=$PWD/target/release/moon \ +//! cargo test --test shardslice_live 2>&1 | tail -40 + +#![allow(clippy::unwrap_used)] + +use std::io::{Read, Write}; +use std::net::{TcpStream, ToSocketAddrs}; +use std::process::{Child, Command}; +use std::time::{Duration, Instant}; + +// --------------------------------------------------------------------------- +// Binary resolution — honours MOON_BIN, falls back to release then debug. +// Pattern: tests/txn_kv_wiring.rs::find_moon_binary +// --------------------------------------------------------------------------- + +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 with `cargo build [--release]` or set \ + MOON_BIN=/path/to/moon." + ); +} + +// --------------------------------------------------------------------------- +// Port allocation — bind-to-0 trick, freed before server binds. +// --------------------------------------------------------------------------- + +fn free_port() -> u16 { + 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); + p +} + +// --------------------------------------------------------------------------- +// Server spawn helpers +// --------------------------------------------------------------------------- + +/// Spawn a moon server. +/// +/// stdout + stderr are written to log files under `dir` so callers can +/// inspect the captured output for log-line assertions. +/// +/// RUST_LOG=moon=info is set so startup info lines (including any future +/// "ShardSlice initialized" markers) are emitted. +/// +/// ALWAYS wrap the returned `Child` in a `ServerGuard` immediately so that +/// test failures do not leak the process. +fn spawn_moon(port: u16, dir: &std::path::Path, shards: u32, extra: &[&str]) -> Child { + let mut args: Vec = vec![ + "--port".into(), + port.to_string(), + "--dir".into(), + dir.to_string_lossy().into_owned(), + "--shards".into(), + shards.to_string(), + ]; + for &e in extra { + args.push(e.into()); + } + Command::new(find_moon_binary()) + .args(&args) + .stdout(std::fs::File::create(dir.join("moon.stdout.log")).expect("create stdout log")) + .stderr(std::fs::File::create(dir.join("moon.stderr.log")).expect("create stderr log")) + .env("RUST_LOG", "moon=info") + .spawn() + .unwrap_or_else(|e| { + panic!( + "Failed to spawn moon binary at '{}': {e}. Build with \ + `cargo build [--release]` or set MOON_BIN.", + find_moon_binary().display() + ) + }) +} + +/// RAII guard: kills the server process when dropped. +/// +/// Wrap the `Child` in this immediately after `spawn_moon` so that test +/// panics and early returns cannot leak the process. +struct ServerGuard(Child); + +impl Drop for ServerGuard { + fn drop(&mut self) { + let _ = self.0.kill(); + let _ = self.0.wait(); + } +} + +// --------------------------------------------------------------------------- +// Connection helpers (deadline-based, no fixed sleep) +// --------------------------------------------------------------------------- + +fn connect(port: u16, deadline: Duration) -> TcpStream { + let addr = format!("127.0.0.1:{port}") + .to_socket_addrs() + .expect("parse addr") + .next() + .expect("one addr"); + let start = Instant::now(); + loop { + match TcpStream::connect_timeout(&addr, Duration::from_millis(200)) { + Ok(s) => { + s.set_read_timeout(Some(Duration::from_secs(10))).ok(); + s.set_write_timeout(Some(Duration::from_secs(10))).ok(); + return s; + } + Err(_) if start.elapsed() < deadline => { + std::thread::sleep(Duration::from_millis(50)); + } + Err(e) => panic!("server never accepted on port {port}: {e}"), + } + } +} + +/// Poll until the server answers PING→PONG. Returns the live connection. +fn wait_ready(port: u16) -> TcpStream { + let mut s = connect(port, Duration::from_secs(30)); + let start = Instant::now(); + loop { + s.write_all(b"PING\r\n").expect("write PING"); + let mut buf = [0u8; 64]; + if let Ok(n) = s.read(&mut buf) { + if n > 0 && buf[..n].windows(4).any(|w| w == b"PONG") { + return s; + } + } + assert!( + start.elapsed() < Duration::from_secs(15), + "server accepted TCP but never answered PING on port {port}" + ); + std::thread::sleep(Duration::from_millis(100)); + s = connect(port, Duration::from_secs(5)); + } +} + +// --------------------------------------------------------------------------- +// Minimal RESP2 client (self-contained; copied helpers from +// cross_shard_consistency_red.rs; not imported across test files) +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone, PartialEq)] +enum Resp { + Simple(String), + Error(String), + Int(i64), + Bulk(Option>), + Array(Option>), +} + +impl Resp { + fn flat(&self) -> String { + match self { + Resp::Simple(s) | Resp::Error(s) => s.clone(), + Resp::Int(i) => i.to_string(), + Resp::Bulk(Some(b)) => String::from_utf8_lossy(b).into_owned(), + Resp::Bulk(None) => "".into(), + Resp::Array(Some(items)) => items.iter().map(Resp::flat).collect::>().join(" "), + Resp::Array(None) => "".into(), + } + } +} + +struct Conn { + s: TcpStream, + buf: Vec, + pos: usize, +} + +impl Conn { + fn new(s: TcpStream) -> Self { + Conn { + s, + buf: Vec::with_capacity(16 * 1024), + pos: 0, + } + } + + fn open(port: u16) -> Self { + Conn::new(connect(port, Duration::from_secs(10))) + } + + fn cmd(&mut self, parts: &[&[u8]]) -> Resp { + let mut req = Vec::with_capacity(128); + req.extend_from_slice(format!("*{}\r\n", parts.len()).as_bytes()); + for p in parts { + req.extend_from_slice(format!("${}\r\n", p.len()).as_bytes()); + req.extend_from_slice(p); + req.extend_from_slice(b"\r\n"); + } + self.s.write_all(&req).expect("write cmd"); + self.frame() + } + + fn cmd_s(&mut self, parts: &[&str]) -> Resp { + let v: Vec<&[u8]> = parts.iter().map(|p| p.as_bytes()).collect(); + self.cmd(&v) + } + + fn fill(&mut self) { + let mut chunk = [0u8; 16 * 1024]; + let n = self.s.read(&mut chunk).expect("read from server"); + assert!(n > 0, "connection closed mid-frame"); + self.buf.extend_from_slice(&chunk[..n]); + } + + fn line(&mut self) -> String { + loop { + if let Some(rel) = self.buf[self.pos..].windows(2).position(|w| w == b"\r\n") { + let line = + String::from_utf8_lossy(&self.buf[self.pos..self.pos + rel]).into_owned(); + self.pos += rel + 2; + return line; + } + self.fill(); + } + } + + fn exact(&mut self, n: usize) -> Vec { + while self.buf.len() - self.pos < n + 2 { + self.fill(); + } + let out = self.buf[self.pos..self.pos + n].to_vec(); + self.pos += n + 2; + out + } + + fn frame(&mut self) -> Resp { + if self.pos > 0 && self.pos == self.buf.len() { + self.buf.clear(); + self.pos = 0; + } + let line = self.line(); + let (tag, rest) = line.split_at(1); + match tag { + "+" => Resp::Simple(rest.to_string()), + "-" => Resp::Error(rest.to_string()), + ":" => Resp::Int(rest.parse().unwrap_or(0)), + "$" => { + let n: i64 = rest.parse().unwrap_or(-1); + if n < 0 { + Resp::Bulk(None) + } else { + Resp::Bulk(Some(self.exact(n as usize))) + } + } + "*" => { + let n: i64 = rest.parse().unwrap_or(-1); + if n < 0 { + Resp::Array(None) + } else { + let mut items = Vec::with_capacity(n as usize); + for _ in 0..n { + items.push(self.frame()); + } + Resp::Array(Some(items)) + } + } + other => panic!("unexpected RESP tag {other:?} in line {line:?}"), + } + } + + /// Send a pipelined batch of raw RESP bytes and read back `count` frames. + fn pipeline(&mut self, raw: &[u8], count: usize) -> Vec { + self.s.write_all(raw).expect("pipeline write"); + (0..count).map(|_| self.frame()).collect() + } +} + +// --------------------------------------------------------------------------- +// Log helpers: read the captured log files for a running server. +// --------------------------------------------------------------------------- + +fn read_log(dir: &std::path::Path) -> String { + let mut out = String::new(); + for name in &["moon.stdout.log", "moon.stderr.log"] { + if let Ok(content) = std::fs::read_to_string(dir.join(name)) { + out.push_str(&content); + } + } + out +} + +// --------------------------------------------------------------------------- +// Helper: detect a seq>1 base RDB file (proves BGREWRITEAOF physically ran). +// Mirrors crash_matrix_per_shard_bgrewriteaof.rs::compacted_base_exists. +// +// Recognises three layouts: +// PerShard: appendonlydir/shard-*/moon.aof..base.rdb (seq > 1) +// TopLevel: appendonlydir/moon.aof..base.rdb (seq > 1) +// (also searched directly under `dir` as a fallback) +// tokio-TopLevel (in-place rewrite): /appendonly.aof starts with +// the 4-byte RDB magic b"MOON" — proves a full rewrite +// completed; the pre-rewrite file is raw RESP (starts with '*' +// or '+'), so the magic check is sufficient. +// --------------------------------------------------------------------------- + +fn compacted_base_exists(dir: &std::path::Path) -> bool { + let is_seq_gt1 = |name: &str| -> bool { + if let Some(rest) = name.strip_prefix("moon.aof.") { + if let Some(seq_str) = rest.strip_suffix(".base.rdb") { + return seq_str.parse::().map(|s| s > 1).unwrap_or(false); + } + } + false + }; + + let aof_dir = dir.join("appendonlydir"); + + // TopLevel layout: base files live directly in appendonlydir/ (or `dir`). + for search_dir in &[&aof_dir, dir] { + if let Ok(files) = std::fs::read_dir(search_dir) { + for f in files.flatten() { + let name = f.file_name().to_string_lossy().to_string(); + if is_seq_gt1(&name) { + return true; + } + } + } + } + + // PerShard layout: appendonlydir/shard-*/moon.aof..base.rdb + if let Ok(shards) = std::fs::read_dir(&aof_dir) { + for entry in shards.flatten() { + let p = entry.path(); + if !p.is_dir() { + continue; + } + if let Ok(files) = std::fs::read_dir(&p) { + for f in files.flatten() { + let name = f.file_name().to_string_lossy().to_string(); + if is_seq_gt1(&name) { + return true; + } + } + } + } + } + + // tokio TopLevel (--shards 1) rewrites aof_path in-place: the legacy + // single AOF file gets atomically replaced with an RDB-preamble snapshot. + // No manifest-stamped base.rdb is produced. Detect by checking whether + // appendonly.aof (the default appendfilename) starts with b"MOON". + // Also check directly under appendonlydir/ as a belt-and-suspenders search. + for candidate in &[dir.join("appendonly.aof"), aof_dir.join("appendonly.aof")] { + if let Ok(mut f) = std::fs::File::open(candidate) { + use std::io::Read as _; + let mut magic = [0u8; 4]; + if f.read_exact(&mut magic).is_ok() && &magic == b"MOON" { + return true; + } + } + } + + false +} + +/// Return the `seq` field from the PerShard AOF manifest, or 0 if absent/unparseable. +fn aof_manifest_seq(dir: &std::path::Path) -> u64 { + let p = dir.join("appendonlydir").join("moon.aof.manifest"); + let text = match std::fs::read_to_string(&p) { + Ok(t) => t, + Err(_) => return 0, + }; + for line in text.lines() { + if let Some(rest) = line.strip_prefix("seq ") { + if let Ok(n) = rest.trim().parse::() { + return n; + } + } + } + 0 +} + +// =========================================================================== +// TESTS +// =========================================================================== + +// --------------------------------------------------------------------------- +// test_ssm1_slice_live_at_startup +// +// RED today: `init_shard` has ZERO production call sites. No +// "ShardSlice initialized" log line is emitted for any shard. This test +// FAILS on the log-marker assertion → the correct RED signal. +// +// Contract C1 (frozen): init_shard logs "ShardSlice initialized" (with the +// shard id) exactly once per shard thread, before the first connection is +// accepted. After the Build phase lands this will turn green. +// --------------------------------------------------------------------------- + +#[test] +fn test_ssm1_slice_live_at_startup() { + const SHARDS: u32 = 4; + let port = free_port(); + let dir = tempfile::tempdir().expect("tempdir"); + + // Use --appendonly no so the test is not slowed by WAL replay. + let _guard = ServerGuard(spawn_moon( + port, + dir.path(), + SHARDS, + &["--appendonly", "no"], + )); + + // Wait for the server to accept connections (all 4 shards have started + // and the first accept loop iteration has run). + drop(wait_ready(port)); + + // Give the shards a moment to flush their startup log lines to disk. + std::thread::sleep(Duration::from_millis(400)); + let log = read_log(dir.path()); + + // --- Contract C1 assertion (RED today) --- + // + // The frozen contract mandates ONE "ShardSlice initialized" log line per + // shard before its first accept. Today `init_shard` is never called in + // production so this marker is absent → the assertion fails → correct + // RED signal. + // + // The shard id must also appear on or near the same line. We check both + // substrings anywhere in the combined log; ordering is not asserted + // (they are on the same tracing event by the contract). + let has_marker = log.contains("ShardSlice initialized"); + assert!( + has_marker, + "expected log to contain \"ShardSlice initialized\" for at least one shard, \ + but the marker is absent.\n\ + RED: init_shard is not wired at shard startup (frozen contract C1, M1).\n\ + Log excerpt (last 3000 chars):\n{}", + &log[log.len().saturating_sub(3000)..] + ); + + // If the marker is present, verify each shard id 0..4 is covered. + // This part will also fail today (marker is absent above), but makes the + // per-shard coverage check explicit for when the Build phase lands. + for shard_id in 0..SHARDS { + // The log line shape required by C1: + // INFO moon::shard::event_loop: ShardSlice initialized shard=N + // We search for the marker AND the shard id independently; a + // properly-formatted tracing field "shard=N" or "shard N" both match. + let id_present = log.contains(&format!("shard={shard_id}")) + || log.contains(&format!("shard {shard_id}")); + assert!( + id_present, + "expected log to contain shard id {shard_id} alongside \ + \"ShardSlice initialized\". Log excerpt:\n{}", + &log[log.len().saturating_sub(3000)..] + ); + } + + // --- Basic wire sanity: server handles commands after startup --- + let mut c1 = Conn::open(port); + assert_eq!( + c1.cmd_s(&["PING"]), + Resp::Simple("PONG".into()), + "PING must return PONG" + ); + + let mut c2 = Conn::open(port); + assert_eq!( + c2.cmd_s(&["SET", "ssm1:key", "value"]), + Resp::Simple("OK".into()), + "SET must return OK" + ); + + let mut c3 = Conn::open(port); + assert_eq!( + c3.cmd_s(&["GET", "ssm1:key"]), + Resp::Bulk(Some(b"value".to_vec())), + "GET must return the value written by c2" + ); +} + +// --------------------------------------------------------------------------- +// test_ssm3_ws_global_registry +// +// GREEN today (behavioral oracle — must KEEP passing through the cutover). +// +// Workspace CREATE/AUTH/LIST are connection-independent. Today this works via +// the per-shard lock path (slot-0 Mutex trick). After the Build phase (M3) it +// works via the process-global WorkspaceRegistry. The assertions are identical +// — the cutover must not regress this. +// +// The restart half of ssm3 lives in test_ssm3_ws_registry_survives_restart +// below — RED today because of a pre-existing bug in `replay_workspace_wal` +// (shared_databases.rs:289): it scans `shard-{id}/` directly while WAL v3 +// segments live in `shard-{id}/wal-v3/`, so zero workspace records replay. +// Contract C3 freezes "WAL replay populates the global registry at boot", so +// the Build phase fixes that bug as part of the registry migration. +// +// The compile-shape half of ssm3 (ShardSliceInit loses its workspace_registry +// field) lives in tests/shardslice_shape.rs, not here. +// --------------------------------------------------------------------------- + +#[test] +fn test_ssm3_ws_global_registry() { + const SHARDS: u32 = 4; + let port = free_port(); + let dir = tempfile::tempdir().expect("tempdir"); + + // --- Create workspace, verify from 12 connections --- + let _guard = ServerGuard(spawn_moon( + port, + dir.path(), + SHARDS, + &["--appendonly", "yes"], + )); + drop(wait_ready(port)); + + // Create workspace from connection A. + let mut ca = Conn::open(port); + let ws_id = match ca.cmd_s(&["WS", "CREATE", "ssm3ws"]) { + Resp::Bulk(Some(id)) => String::from_utf8(id).expect("WS CREATE id must be utf8"), + other => panic!("WS CREATE reply {other:?}"), + }; + + // WS AUTH must succeed from 12 fresh connections (kernel spreads them + // across shard listeners on Linux; macOS may pin — both must pass). + // + // Today: works via the per-shard lock path. + // Post-cutover (M3): must work via the global registry. + // THIS IS THE PRIMARY BEHAVIORAL ORACLE — never weaken this assertion. + for i in 0..12 { + let mut c = Conn::open(port); + let r = c.cmd_s(&["WS", "AUTH", &ws_id]); + assert_eq!( + r, + Resp::Simple("OK".into()), + "conn {i}: WS AUTH must succeed from ANY shard connection (got {r:?})" + ); + let list_flat = c.cmd_s(&["WS", "LIST"]).flat(); + assert!( + list_flat.contains("ssm3ws"), + "conn {i}: WS LIST must include 'ssm3ws' (got: {list_flat})" + ); + } + + // Workspace-scoped key: bound write visible to bound GET, invisible to + // an unbound connection — same oracle as cdg6f. + let mut bound = Conn::open(port); + assert_eq!( + bound.cmd_s(&["WS", "AUTH", &ws_id]), + Resp::Simple("OK".into()), + "bound conn: WS AUTH must succeed" + ); + assert_eq!( + bound.cmd_s(&["SET", "ssm3:wkey", "wsval"]), + Resp::Simple("OK".into()), + "bound conn: SET inside workspace must succeed" + ); + assert_eq!( + bound.cmd_s(&["GET", "ssm3:wkey"]), + Resp::Bulk(Some(b"wsval".to_vec())), + "bound conn: GET must return workspace-scoped value" + ); + let mut unbound = Conn::open(port); + assert_eq!( + unbound.cmd_s(&["GET", "ssm3:wkey"]), + Resp::Bulk(None), + "unbound conn: workspace key must be invisible" + ); + // _guard dropped here → server killed +} + +// --------------------------------------------------------------------------- +// test_ssm3_ws_registry_survives_restart +// +// RED today: `replay_workspace_wal` (shared_databases.rs:289) joins +// `shard-{id}` directly, but WAL v3 segments are written under +// `shard-{id}/wal-v3/` (compare recovery.rs:361) — so workspace +// Create/Drop records are never replayed and the registry comes up empty +// after a restart. Frozen contract C3 mandates "WAL replay populates the +// global registry at boot"; the Build phase fixes the directory path as part +// of the global-registry migration, turning this green. +// --------------------------------------------------------------------------- + +#[test] +fn test_ssm3_ws_registry_survives_restart() { + const SHARDS: u32 = 4; + let dir = tempfile::tempdir().expect("tempdir"); + + // --- Round 1: create the workspace with WAL enabled --- + { + let port = free_port(); + let _guard = ServerGuard(spawn_moon( + port, + dir.path(), + SHARDS, + &["--appendonly", "yes"], + )); + drop(wait_ready(port)); + + let mut ca = Conn::open(port); + let ws_id = match ca.cmd_s(&["WS", "CREATE", "ssm3rs"]) { + Resp::Bulk(Some(id)) => String::from_utf8(id).expect("WS CREATE id must be utf8"), + other => panic!("WS CREATE reply {other:?}"), + }; + assert_eq!( + ca.cmd_s(&["WS", "AUTH", &ws_id]), + Resp::Simple("OK".into()), + "WS AUTH must succeed pre-restart" + ); + // Let the 1ms WAL flush tick + everysec window drain before the kill. + std::thread::sleep(Duration::from_millis(1500)); + // _guard dropped → server killed + } + + // --- Round 2: restart on the SAME dir; the workspace must be replayed --- + { + let port2 = free_port(); + let _guard2 = ServerGuard(spawn_moon( + port2, + dir.path(), + SHARDS, + &["--appendonly", "yes"], + )); + drop(wait_ready(port2)); + + let mut c = Conn::open(port2); + let list_flat = c.cmd_s(&["WS", "LIST"]).flat(); + assert!( + list_flat.contains("ssm3rs"), + "after restart, WS LIST must include 'ssm3rs' (got: {list_flat}).\n\ + RED: replay_workspace_wal scans shard-{{id}}/ but WAL v3 segments \ + live in shard-{{id}}/wal-v3/ — zero workspace records replayed \ + (pre-existing bug; fixed by Build under frozen contract C3)." + ); + } +} + +// --------------------------------------------------------------------------- +// test_ssm4a_aof_fold_cooperative +// +// GREEN today (behavioral oracle — must KEEP passing through the cutover). +// +// Exactly-once INCR across a BGREWRITEAOF boundary: +// 1. Write 2000 keys with pipelined SETs. +// 2. Start a writer thread doing INCR on "ssm4a:ctr". +// 3. Issue BGREWRITEAOF. +// 4. Wait ≤30s for the rewrite to settle (poll for seq>1 base file). +// 5. Stop the writer; record exact successful-INCR count. +// 6. Assert GET ssm4a:ctr == recorded count (exactly-once). +// 7. Restart on the same dir; assert counter and key sample survive. +// +// Two variants share one helper: +// - shards=1 (no flag): BGREWRITEAOF on shards>=2 + appendonly yes is gated +// (MULTI_SHARD_AOF_REWRITE_UNSAFE) without the experimental flag. +// - shards=4 + --experimental-per-shard-rewrite: drives +// do_rewrite_per_shard — the EXACT foreign-lock fold C4 redesigns into +// the cooperative snapshot. Post-cutover both use the new protocol; the +// assertions are unchanged. +// +// Total runtime target: < 60s per variant. +// --------------------------------------------------------------------------- + +fn aof_fold_exactly_once(shards: u32, extra: &[&str]) { + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + + let mut flags: Vec<&str> = vec!["--appendonly", "yes", "--appendfsync", "everysec"]; + flags.extend_from_slice(extra); + + let port = free_port(); + let dir = tempfile::tempdir().expect("tempdir"); + let recorded_count: i64; + + // --- Round 1: write data, trigger rewrite under concurrent writes --- + { + let _guard = ServerGuard(spawn_moon(port, dir.path(), shards, &flags)); + drop(wait_ready(port)); + + // Pipeline 2000 SET operations. + { + let mut bulk = Conn::open(port); + let mut req = Vec::with_capacity(80 * 2000); + for i in 0..2000_u32 { + let key = format!("ssm4a:k{i}"); + let val = format!("v{i}"); + req.extend_from_slice( + format!( + "*3\r\n$3\r\nSET\r\n${}\r\n{}\r\n${}\r\n{}\r\n", + key.len(), + key, + val.len(), + val + ) + .as_bytes(), + ); + } + let responses = bulk.pipeline(&req, 2000); + for (i, r) in responses.iter().enumerate() { + assert_eq!( + *r, + Resp::Simple("OK".into()), + "pipelined SET {i} failed: {r:?}" + ); + } + } + + // Spin up a writer thread doing INCR on the counter key. + let stop = Arc::new(AtomicBool::new(false)); + let success_count = Arc::new(AtomicU64::new(0)); + let writer_port = port; + let stop_clone = stop.clone(); + let count_clone = success_count.clone(); + let writer = std::thread::spawn(move || { + let mut c = Conn::open(writer_port); + while !stop_clone.load(Ordering::Relaxed) { + if matches!(c.cmd_s(&["INCR", "ssm4a:ctr"]), Resp::Int(_)) { + count_clone.fetch_add(1, Ordering::Relaxed); + } + } + }); + + // Let the writer accumulate some increments before the rewrite. + std::thread::sleep(Duration::from_millis(200)); + + // Issue BGREWRITEAOF. + let mut admin = Conn::open(port); + let bgrw = admin.cmd_s(&["BGREWRITEAOF"]); + // Accept "Background append only file rewriting started" (Simple) or + // "ERR Background AOF rewrite already in progress" (idempotent). + assert!( + matches!(&bgrw, Resp::Simple(_) | Resp::Error(_)), + "BGREWRITEAOF unexpected reply: {bgrw:?}" + ); + + // Poll for rewrite completion up to 30s. + // + // For the per-shard experimental path we wait for the manifest seq to + // advance: the coordinator atomically bumps seq only after ALL shards + // complete their folds. Waiting for any single shard's base file + // (compacted_base_exists) is too early — other shards may still be + // folding, so EverySec fsyncs on the new incr have not yet run, and + // killing the server in that window loses data. + // + // For the legacy single-writer path compacted_base_exists remains + // correct (there is no coordinator; one seq flip = one base file). + let uses_per_shard_rewrite = flags.contains(&"--experimental-per-shard-rewrite"); + let rewrite_deadline = Instant::now() + Duration::from_secs(30); + loop { + std::thread::sleep(Duration::from_millis(500)); + let done = if uses_per_shard_rewrite { + aof_manifest_seq(dir.path()) > 1 + } else { + compacted_base_exists(dir.path()) + }; + if done { + break; + } + if Instant::now() > rewrite_deadline { + panic!( + "BGREWRITEAOF did not complete within 30s — \ + no layout-appropriate compaction artifact appeared in {:?}.", + dir.path() + ); + } + } + + // Stop the writer and record the exact successful-INCR count. + stop.store(true, Ordering::Relaxed); + writer.join().expect("writer thread join"); + let n = success_count.load(Ordering::Relaxed) as i64; + assert!(n > 0, "writer must have succeeded at least once"); + + // Let the everysec flush window close so every acked INCR is durable. + std::thread::sleep(Duration::from_millis(1500)); + + // Read the counter — MUST equal the exact successful-INCR count. + // This is the exactly-once assertion: neither dropped nor double-applied. + let mut check = Conn::open(port); + let got = match check.cmd_s(&["GET", "ssm4a:ctr"]) { + Resp::Bulk(Some(b)) => String::from_utf8(b) + .expect("utf8") + .parse::() + .expect("parse i64"), + other => panic!("GET ssm4a:ctr reply {other:?}"), + }; + assert_eq!( + got, + n, + "exactly-once invariant: ctr ({got}) must equal recorded INCR count ({n}). \ + diff={}: positive=double-apply, negative=lost writes.", + got - n + ); + + recorded_count = n; + + // _guard dropped here → server killed (SIGKILL equivalent) + } + + // --- Round 2: restart on the SAME dir, assert durability --- + { + let port2 = free_port(); + let _guard2 = ServerGuard(spawn_moon(port2, dir.path(), shards, &flags)); + drop(wait_ready(port2)); + + let mut c = Conn::open(port2); + + // Counter must survive the restart with the same value. + let ctr_after = match c.cmd_s(&["GET", "ssm4a:ctr"]) { + Resp::Bulk(Some(b)) => String::from_utf8(b) + .expect("utf8") + .parse::() + .unwrap_or(-1), + other => panic!("after restart GET ssm4a:ctr reply {other:?}"), + }; + assert_eq!( + ctr_after, recorded_count, + "after restart, ssm4a:ctr must equal recorded count ({recorded_count}); \ + got {ctr_after}" + ); + + // A sample of the 2000 keys must exist. + for i in [0_u32, 499, 999, 1499, 1999] { + let key = format!("ssm4a:k{i}"); + let expected = format!("v{i}"); + let r = c.cmd_s(&["GET", &key]); + assert_eq!( + r, + Resp::Bulk(Some(expected.into_bytes())), + "after restart, key {key} must exist" + ); + } + } +} + +#[test] +fn test_ssm4a_aof_fold_cooperative() { + // shards=1: BGREWRITEAOF runs without the experimental flag. + aof_fold_exactly_once(1, &[]); +} + +#[test] +fn test_ssm4a_fold_4shard_experimental() { + // do_rewrite_per_shard (F6) — the exact foreign-lock fold contract C4 + // redesigns into the cooperative snapshot. GREEN today; must stay green + // through the cutover. + aof_fold_exactly_once(4, &["--experimental-per-shard-rewrite"]); +} + +// --------------------------------------------------------------------------- +// test_reject_uninitialized_shard_no_wire_panic +// +// GREEN today and must stay green forever (permanent wire guard). +// +// Sending junk-but-parseable commands + normal SET/GET must not cause a +// with_shard panic that is reachable over the network. The server must stay +// alive and answer PINGs on subsequent connections. +// +// The real uninitialized-abort assertion (C1: startup abort naming the +// shard id, before first accept) is unit-level in slice.rs's tests and in +// shardslice_shape.rs. This wire guard pins "no with_shard panic reachable +// over the wire" and is intended to be a permanent regression guard. +// --------------------------------------------------------------------------- + +#[test] +fn test_reject_uninitialized_shard_no_wire_panic() { + let port = free_port(); + let dir = tempfile::tempdir().expect("tempdir"); + let _guard = ServerGuard(spawn_moon(port, dir.path(), 4, &["--appendonly", "no"])); + drop(wait_ready(port)); + + // --- Send a burst of junk-but-parseable commands --- + let mut c = Conn::open(port); + + // *0\r\n is a valid RESP2 frame (zero-element array) that the server + // parses and rejects as "empty command" — exercises parse→dispatch without + // routing to any specific shard key. + c.s.write_all(b"*0\r\n").expect("write *0"); + let _ = c.frame(); // consume the error reply + + // A negative-arity array (*-1) is also valid RESP2 (null array). The + // server parses it and replies "ERR invalid command format" — consume + // that reply to keep the connection in sync. + c.s.write_all(b"*-1\r\n").expect("write *-1 (null array)"); + let _ = c.frame(); // consume the error reply + + // Normal SET/GET from the SAME connection must still work after the + // junk burst. + assert_eq!( + c.cmd_s(&["SET", "ssm_guard:k", "alive"]), + Resp::Simple("OK".into()), + "SET after junk burst must succeed — server must not have crashed" + ); + assert_eq!( + c.cmd_s(&["GET", "ssm_guard:k"]), + Resp::Bulk(Some(b"alive".to_vec())), + "GET after junk burst must return the set value" + ); + + // Open a SECOND fresh connection and confirm the server is still alive. + let mut c2 = Conn::open(port); + assert_eq!( + c2.cmd_s(&["PING"]), + Resp::Simple("PONG".into()), + "server must still answer PING from a fresh connection after the junk burst" + ); +} diff --git a/tests/shardslice_shape.rs b/tests/shardslice_shape.rs new file mode 100644 index 00000000..1c133216 --- /dev/null +++ b/tests/shardslice_shape.rs @@ -0,0 +1,550 @@ +//! ADD task `shardslice-migration` — code-shape pin tests (§4 TESTS). +//! +//! All tests in this file are **source-grep pins**: they read `.rs` files at +//! runtime and assert internal code shape. No network connections, no server +//! binary, no feature flags needed. +//! +//! Red/green split (as of task phase `tests`, pre-build): +//! RED — test_ssm5_wrappers_gone, test_ssm4b_observer_lockfree, +//! test_ssm3_shape_no_ws_field_in_slice, test_reject_wrapper_resurrection +//! GREEN — test_reject_borrow_across_await +//! +//! Running: +//! cargo test --test shardslice_shape + +use std::path::{Path, PathBuf}; + +// ── helpers ─────────────────────────────────────────────────────────────────── + +/// Read a source file relative to the crate root (CARGO_MANIFEST_DIR). +/// +/// Panics with the path on failure. +fn read_src(rel: &str) -> String { + let base = Path::new(env!("CARGO_MANIFEST_DIR")); + let full = base.join(rel); + std::fs::read_to_string(&full).unwrap_or_else(|e| panic!("cannot read {}: {e}", full.display())) +} + +/// Recursively collect all `.rs` files under `dir` (relative to manifest). +/// +/// Returns absolute paths. +fn rs_files_under(rel_dir: &str) -> Vec { + let base = Path::new(env!("CARGO_MANIFEST_DIR")).join(rel_dir); + let mut out = Vec::new(); + collect_rs(&base, &mut out); + out +} + +fn collect_rs(dir: &Path, out: &mut Vec) { + let rd = match std::fs::read_dir(dir) { + Ok(r) => r, + Err(_) => return, + }; + for entry in rd.flatten() { + let path = entry.path(); + if path.is_dir() { + collect_rs(&path, out); + } else if path.extension().is_some_and(|e| e == "rs") { + out.push(path); + } + } +} + +/// Return lines of `text` that match a simple substring, annotated with 1-based +/// line numbers. Skips any line that comes after a `#[cfg(test)]` block marker +/// (heuristic: once we see a line containing exactly `#[cfg(test)]` followed by +/// a `mod tests` the rest of the file is test code). This is intentionally +/// conservative: we only skip lines that appear AFTER a `mod tests {` line +/// (i.e., inside the test module). The heuristic is documented here because it +/// is approximate — it stops at the first `mod tests {` it finds. +fn grep_production_lines<'a>(text: &'a str, pattern: &str) -> Vec<(usize, &'a str)> { + // Simple heuristic: split at the first occurrence of a `#[cfg(test)]` + // attribute followed (within a few lines) by `mod tests`. We split the + // file at that point and only search the production half. + let production_text = split_off_test_module(text); + production_text + .lines() + .enumerate() + .filter(|(_, line)| line.contains(pattern)) + .map(|(i, line)| (i + 1, line)) + .collect() +} + +/// Heuristic: return the slice of `text` that precedes the `#[cfg(test)]` +/// module marker. If no such marker is found, returns all of `text`. +/// +/// Strategy: scan forward; when we see a line containing `#[cfg(test)]` and +/// the NEXT non-blank line contains `mod tests`, truncate there. +fn split_off_test_module(text: &str) -> &str { + let mut cfg_test_pos: Option = None; + let mut byte_offset = 0usize; + + for line in text.lines() { + let trimmed = line.trim(); + if trimmed == "#[cfg(test)]" { + cfg_test_pos = Some(byte_offset); + } + byte_offset += line.len() + 1; // +1 for the '\n' + } + + match cfg_test_pos { + Some(pos) => &text[..pos], + None => text, + } +} + +// ── Shared wrapper-shape checker (used by both ssm5 and reject_wrapper_resurrection) ─ + +fn assert_wrappers_gone() { + // ── 1. ShardDatabases must have no per-shard lock wrappers ─────────────── + + let shared_db = read_src("src/shard/shared_databases.rs"); + let production = split_off_test_module(&shared_db); + + // Field: shards: Vec>> (or any variant) + let shards_locks: Vec<_> = grep_production_lines(production, "RwLock") + .into_iter() + .filter(|(_, l)| { + // We want field declarations, not method return types. Accept lines + // that look like struct field declarations (contain `shards` or the + // literal type, inside the struct body — any non-comment occurrence + // in production is the failing signal). + !l.trim_start().starts_with("//") + }) + .collect(); + assert!( + shards_locks.is_empty(), + "src/shard/shared_databases.rs: found RwLock in production code \ + (wrapper field not yet deleted) — offending lines:\n{:?}", + shards_locks + ); + + // Field: vector_stores: Vec> + let vec_stores: Vec<_> = grep_production_lines(production, "Mutex") + .into_iter() + .filter(|(_, l)| !l.trim_start().starts_with("//")) + .collect(); + assert!( + vec_stores.is_empty(), + "src/shard/shared_databases.rs: found Mutex in production code \ + (wrapper field not yet deleted) — offending lines:\n{:?}", + vec_stores + ); + + // Field: text_stores: Vec> + let text_stores: Vec<_> = grep_production_lines(production, "Mutex") + .into_iter() + .filter(|(_, l)| !l.trim_start().starts_with("//")) + .collect(); + assert!( + text_stores.is_empty(), + "src/shard/shared_databases.rs: found Mutex in production code \ + (wrapper field not yet deleted) — offending lines:\n{:?}", + text_stores + ); + + // Field: graph_stores: Vec> or Vec> + let graph_stores: Vec<_> = production + .lines() + .enumerate() + .filter(|(_, l)| { + !l.trim_start().starts_with("//") + && (l.contains("RwLock") || l.contains("Mutex")) + }) + .map(|(i, l)| (i + 1, l)) + .collect(); + assert!( + graph_stores.is_empty(), + "src/shard/shared_databases.rs: found RwLock/Mutex in \ + production code (wrapper field not yet deleted) — offending lines:\n{:?}", + graph_stores + ); + + // Per-shard registry wrapper arrays (plural names indicate per-shard arrays) + for pattern in &[ + "temporal_registries", + "temporal_kv_indexes", + "workspace_registries", + "durable_queue_registries", + "trigger_registries", + "kv_write_intents", + "deferred_hnsw_inserts", + ] { + let hits: Vec<_> = grep_production_lines(production, pattern) + .into_iter() + .filter(|(_, l)| !l.trim_start().starts_with("//")) + .collect(); + assert!( + hits.is_empty(), + "src/shard/shared_databases.rs: found per-shard registry array '{}' in \ + production code (must be deleted in M5) — offending lines:\n{:?}", + pattern, + hits + ); + } + + // Deleted methods: write_db, read_db, all_shard_dbs, aggregate_memory + for method in &[ + "fn write_db", + "fn read_db", + "fn all_shard_dbs", + "fn aggregate_memory", + ] { + let hits: Vec<_> = grep_production_lines(production, method) + .into_iter() + .filter(|(_, l)| !l.trim_start().starts_with("//")) + .collect(); + assert!( + hits.is_empty(), + "src/shard/shared_databases.rs: method '{}' still present \ + (must be deleted in M5) — offending lines:\n{:?}", + method, + hits + ); + } + + // ── 2. No is_initialized() dual-branch gates in production dirs ────────── + // Scan src/server/, src/command/, src/shard/ (excluding slice.rs itself). + // Heuristic: skip the test-module tail of each file. + // + // NOTE for the Build phase: contract C1's once-per-shard startup assert + // must be implemented as `slice::assert_initialized(shard_id)` — a helper + // INSIDE slice.rs — not as a raw `slice::is_initialized()` call at the + // event loop, or this pin would flag it. The pin stays strict on purpose: + // zero is_initialized() call sites outside slice.rs. + + for dir in &["src/server", "src/command", "src/shard"] { + for path in rs_files_under(dir) { + // Exclude slice.rs itself — the gate fn is allowed to SURVIVE there. + if path.file_name().is_some_and(|n| n == "slice.rs") { + continue; + } + let text = std::fs::read_to_string(&path) + .unwrap_or_else(|e| panic!("cannot read {}: {e}", path.display())); + let prod = split_off_test_module(&text); + + // Look for is_initialized() in production code. + // Acceptable survivors: lines that are pure comments. + let hits: Vec<_> = prod + .lines() + .enumerate() + .filter(|(_, l)| { + let t = l.trim_start(); + !t.starts_with("//") && l.contains("is_initialized()") + }) + .map(|(i, l)| (i + 1, l)) + .collect(); + + assert!( + hits.is_empty(), + "{}:{}: is_initialized() dual-branch gate found in production code \ + (must be deleted in M5, C6) — offending lines:\n{:?}", + path.display(), + "production", + hits + ); + } + } + + // ── 3. uring_handler: no write_db(, must have with_shard ───────────────── + let uring = read_src("src/shard/uring_handler.rs"); + let uring_prod = split_off_test_module(&uring); + + let write_db_hits: Vec<_> = grep_production_lines(uring_prod, "write_db(") + .into_iter() + .filter(|(_, l)| !l.trim_start().starts_with("//")) + .collect(); + assert!( + write_db_hits.is_empty(), + "src/shard/uring_handler.rs: write_db( call still present \ + (must switch to with_shard per M5/C6) — offending lines:\n{:?}", + write_db_hits + ); + + let with_shard_hits: Vec<_> = grep_production_lines(uring_prod, "with_shard") + .into_iter() + .filter(|(_, l)| !l.trim_start().starts_with("//")) + .collect(); + assert!( + !with_shard_hits.is_empty(), + "src/shard/uring_handler.rs: no with_shard call found — \ + uring_handler must use with_shard (not write_db) after M5" + ); +} + +// ── Test 1: PRIMARY RED TEST — wrappers gone (ssm5) ───────────────────────── + +/// ssm5: ShardDatabases source has zero RwLock/Mutex/ +/// Mutex/RwLock/per-shard registry Mutex fields, zero +/// `is_initialized()` gates in production paths, and uring_handler uses +/// `with_shard` not `write_db`. +/// +/// RED today: every wrapper field and gate still exists. +#[test] +fn test_ssm5_wrappers_gone() { + assert_wrappers_gone(); +} + +// ── Test 2: ssm4b — metrics/admin paths are lock-free observers ────────────── + +/// ssm4b: metrics_setup, server_admin, and persistence_tick must not call +/// read_db/write_db/all_shard_dbs/aggregate_memory or access .vector_stores / +/// .text_stores / .graph_stores. After M4/C5 they read published atomics only. +/// +/// RED today: persistence_tick.rs calls aggregate_memory + write_db + read_db; +/// metrics_setup.rs calls read_db; server_admin.rs calls read_db. +#[test] +fn test_ssm4b_observer_lockfree() { + let forbidden = [ + "read_db(", + "write_db(", + "all_shard_dbs(", + "aggregate_memory(", + ".vector_stores", + ".text_stores", + ".graph_stores", + ]; + + // Paths to check: (relative path, human label) + // metrics_setup is under src/admin/ + let files = [ + ("src/admin/metrics_setup.rs", "src/admin/metrics_setup.rs"), + ("src/command/server_admin.rs", "src/command/server_admin.rs"), + ( + "src/shard/persistence_tick.rs", + "src/shard/persistence_tick.rs", + ), + ]; + + let mut all_violations: Vec = Vec::new(); + + for (rel, label) in &files { + let text = read_src(rel); + let prod = split_off_test_module(&text); + + for pat in &forbidden { + let hits: Vec<_> = prod + .lines() + .enumerate() + .filter(|(_, l)| { + let t = l.trim_start(); + !t.starts_with("//") && l.contains(pat) + }) + .map(|(i, l)| format!(" {}:{}: {}", label, i + 1, l.trim())) + .collect(); + all_violations.extend(hits); + } + } + + assert!( + all_violations.is_empty(), + "ssm4b: lock-acquiring / store-accessing calls found in observer paths \ + (must be replaced with published atomics per M4/C5):\n{}", + all_violations.join("\n") + ); +} + +// ── Test 3: ssm3 — workspace_registry shape ────────────────────────────────── + +/// ssm3 shape pin: +/// - slice.rs must NOT have `workspace_registry` as a field in ShardSlice +/// or ShardSliceInit (it leaves in M3). +/// - shared_databases.rs MUST have a single process-global field +/// `workspace_registry:` typed `Mutex>>`. +/// - shared_databases.rs must NOT have `workspace_registries:` (the per-shard +/// array — deleted in M3). +/// +/// RED today: slice.rs HAS `workspace_registry` in ShardSlice/ShardSliceInit; +/// shared_databases.rs has `workspace_registries` (array) and lacks the +/// global single-field form. +#[test] +fn test_ssm3_shape_no_ws_field_in_slice() { + let slice_text = read_src("src/shard/slice.rs"); + let slice_prod = split_off_test_module(&slice_text); + + // slice.rs must have NO `workspace_registry` struct field. + // We look for lines that declare this as a field (contain `workspace_registry:`) + // outside of a comment. + let slice_field_hits: Vec<_> = slice_prod + .lines() + .enumerate() + .filter(|(_, l)| { + let t = l.trim_start(); + !t.starts_with("//") && l.contains("workspace_registry:") + }) + .map(|(i, l)| (i + 1, l.trim())) + .collect(); + + assert!( + slice_field_hits.is_empty(), + "src/shard/slice.rs: workspace_registry field still present in ShardSlice \ + or ShardSliceInit (must be removed in M3) — offending lines:\n{:?}", + slice_field_hits + ); + + // shared_databases.rs must NOT have the plural array `workspace_registries:` + let shared_db_text = read_src("src/shard/shared_databases.rs"); + let shared_prod = split_off_test_module(&shared_db_text); + + let array_hits: Vec<_> = shared_prod + .lines() + .enumerate() + .filter(|(_, l)| { + let t = l.trim_start(); + !t.starts_with("//") && l.contains("workspace_registries") + }) + .map(|(i, l)| (i + 1, l.trim())) + .collect(); + + assert!( + array_hits.is_empty(), + "src/shard/shared_databases.rs: workspace_registries (per-shard array) still present \ + (must be replaced by single global `workspace_registry` field in M3) — \ + offending lines:\n{:?}", + array_hits + ); + + // shared_databases.rs MUST have a single global field `workspace_registry:` + // typed with `Mutex>>`. + // We check for the field name AND the type on the same or adjacent line. + // Simple heuristic: look for a line that matches the field declaration. + let global_field_present = shared_prod.lines().any(|l| { + let t = l.trim_start(); + !t.starts_with("//") + && l.contains("workspace_registry:") + && l.contains("Mutex>>") + }); + + assert!( + global_field_present, + "src/shard/shared_databases.rs: process-global `workspace_registry: \ + Mutex>>` field not found (required by M3/C3)" + ); +} + +// ── Test 4: reject borrow_across_await (GREEN today) ───────────────────────── + +/// reject borrow_across_await: no `.await` appears inside a `with_shard(` or +/// `with_shard_db(` closure body in any production source file. +/// +/// Implementation strategy: for each `.rs` file under `src/`, scan line by +/// line. When we encounter a line containing `with_shard(` or `with_shard_db(`, +/// we walk forward collecting lines until the call's paren depth returns to +/// zero. We assert that no `.await` appears in that span. +/// +/// Paren-depth counting is simple ASCII character scanning (`(` increments, +/// `)` decrements). We start counting from the `(` that opens the `with_shard` +/// call (the depth starts at 1 on that line). If the file ends before depth +/// reaches 0, we conservatively include all remaining lines in the span. +/// +/// Edge cases and limitations (documented per spec): +/// - Parentheses inside string literals are not handled; this is acceptable +/// because `with_shard` call sites do not contain string literals with +/// unbalanced parens in practice. +/// - Nested `with_shard` calls would be flagged by the reentrancy check at +/// runtime; the borrow-across-await check here is an additional structural pin. +/// +/// GREEN today: all with_shard closures are synchronous. +#[test] +fn test_reject_borrow_across_await() { + let mut violations: Vec = Vec::new(); + + for path in rs_files_under("src") { + let text = match std::fs::read_to_string(&path) { + Ok(t) => t, + Err(_) => continue, + }; + let prod = split_off_test_module(&text); + let lines: Vec<&str> = prod.lines().collect(); + + let mut i = 0; + while i < lines.len() { + let line = lines[i]; + // Only care about lines that open a with_shard call (not comments). + let trimmed = line.trim_start(); + if !trimmed.starts_with("//") + && (line.contains("with_shard(") || line.contains("with_shard_db(")) + { + // Find the call span by tracking paren depth. + // We count ALL parens from the first `with_shard(` on this line. + let call_start_col = line + .find("with_shard(") + .or_else(|| line.find("with_shard_db(")) + .unwrap_or(0); + // Count depth starting from this line at the opening paren. + let mut depth: i32 = 0; + let span_start = i; + let mut span_end = i; + let mut found_close = false; + + 'outer: for j in i..lines.len() { + let scan_line = if j == i { + &lines[j][call_start_col..] + } else { + lines[j] + }; + for ch in scan_line.chars() { + match ch { + '(' => depth += 1, + ')' => { + depth -= 1; + if depth == 0 { + span_end = j; + found_close = true; + break 'outer; + } + } + _ => {} + } + } + } + + if !found_close { + // Span extends to end of production section. + span_end = lines.len().saturating_sub(1); + } + + // Check the span for .await + for j in span_start..=span_end { + let span_line = lines[j]; + let span_t = span_line.trim_start(); + if !span_t.starts_with("//") && span_line.contains(".await") { + violations.push(format!( + "{}:{}: .await inside with_shard closure — line: {}", + path.display(), + j + 1, + span_line.trim() + )); + } + } + + // Advance past this call span to avoid double-counting. + i = span_end + 1; + continue; + } + i += 1; + } + } + + assert!( + violations.is_empty(), + "reject borrow_across_await: found .await inside with_shard/with_shard_db \ + closure — this is forbidden (borrow is sync by construction; an async \ + block capturing the guard is the target violation):\n{}", + violations.join("\n") + ); +} + +// ── Test 5: reject wrapper_resurrection — alias of ssm5 ────────────────────── + +/// Permanent guard: same shape test as ssm5. This is the CI gate that blocks +/// any future commit reintroducing a lock wrapper into ShardDatabases. +/// +/// Intentionally a separate test name so it shows up as a distinct failure in +/// CI with its own reject-scenario label. +/// +/// RED today (same reason as test_ssm5_wrappers_gone). +#[test] +fn test_reject_wrapper_resurrection() { + assert_wrappers_gone(); +} diff --git a/tests/txn_completeness_edge_cases.rs b/tests/txn_completeness_edge_cases.rs index 49a28390..eb9a0d21 100644 --- a/tests/txn_completeness_edge_cases.rs +++ b/tests/txn_completeness_edge_cases.rs @@ -164,7 +164,8 @@ async fn start_txn_server(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -179,6 +180,7 @@ async fn start_txn_server(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("edge-shard-{}", id)) @@ -219,6 +221,7 @@ async fn start_txn_server(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/txn_cypher_write_rollback.rs b/tests/txn_cypher_write_rollback.rs index e17f3a93..a364caf1 100644 --- a/tests/txn_cypher_write_rollback.rs +++ b/tests/txn_cypher_write_rollback.rs @@ -170,7 +170,8 @@ async fn start_txn_server(num_shards: usize, persistence_dir: &str) -> (u16, Can .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -185,6 +186,7 @@ async fn start_txn_server(num_shards: usize, persistence_dir: &str) -> (u16, Can let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("txn-cypher-shard-{}", id)) @@ -227,6 +229,7 @@ async fn start_txn_server(num_shards: usize, persistence_dir: &str) -> (u16, Can shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/txn_ft_search_snapshot.rs b/tests/txn_ft_search_snapshot.rs index 4c371fd8..0f592c99 100644 --- a/tests/txn_ft_search_snapshot.rs +++ b/tests/txn_ft_search_snapshot.rs @@ -168,7 +168,8 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -183,6 +184,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("txn-ft-shard-{}", id)) @@ -225,6 +227,7 @@ async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/txn_graph_wiring.rs b/tests/txn_graph_wiring.rs index a13b00cd..d365d9d6 100644 --- a/tests/txn_graph_wiring.rs +++ b/tests/txn_graph_wiring.rs @@ -180,7 +180,8 @@ async fn start_txn_server(num_shards: usize, persistence_dir: &str) -> (u16, Can .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -195,6 +196,7 @@ async fn start_txn_server(num_shards: usize, persistence_dir: &str) -> (u16, Can let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("txn-graph-shard-{}", id)) @@ -237,6 +239,7 @@ async fn start_txn_server(num_shards: usize, persistence_dir: &str) -> (u16, Can shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/txn_kv_wiring.rs b/tests/txn_kv_wiring.rs index 9816908a..b0654824 100644 --- a/tests/txn_kv_wiring.rs +++ b/tests/txn_kv_wiring.rs @@ -164,7 +164,8 @@ async fn start_txn_server(num_shards: usize, persistence_dir: &str) -> (u16, Can .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -179,6 +180,7 @@ async fn start_txn_server(num_shards: usize, persistence_dir: &str) -> (u16, Can let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("txn-test-shard-{}", id)) @@ -221,6 +223,7 @@ async fn start_txn_server(num_shards: usize, persistence_dir: &str) -> (u16, Can shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); diff --git a/tests/workspace_integration.rs b/tests/workspace_integration.rs index 2543ce79..4dd5aad9 100644 --- a/tests/workspace_integration.rs +++ b/tests/workspace_integration.rs @@ -152,7 +152,8 @@ async fn start_workspace_server(num_shards: usize) -> (u16, CancellationToken) { .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -167,6 +168,7 @@ async fn start_workspace_server(num_shards: usize) -> (u16, CancellationToken) { let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("ws-test-shard-{}", id)) @@ -209,6 +211,7 @@ async fn start_workspace_server(num_shards: usize) -> (u16, CancellationToken) { shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread"); @@ -375,7 +378,8 @@ async fn start_workspace_server_with_auth( .iter_mut() .map(|s| std::mem::take(&mut s.databases)) .collect(); - let shard_databases = moon::shard::shared_databases::ShardDatabases::new(all_dbs); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); let mut shard_handles = Vec::with_capacity(num_shards); for (id, mut shard) in shards.into_iter().enumerate() { @@ -390,6 +394,7 @@ async fn start_workspace_server_with_auth( let shard_pubsub_regs = all_pubsub_registries.clone(); let shard_remote_sub_maps = all_remote_sub_maps.clone(); let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); let handle = std::thread::Builder::new() .name(format!("ws-auth-shard-{}", id)) @@ -432,6 +437,7 @@ async fn start_workspace_server_with_auth( shard_pubsub_regs, shard_remote_sub_maps, shard_affinity, + shard_slice_init, ))); }) .expect("failed to spawn shard thread");