feat(replication): R2 multi-shard master PSYNC — merged RDB + per-record SELECT framing

CHANGELOG.md

@@ -6,6 +6,47 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added — R2: multi-shard master PSYNC (task #20, RFC 1B)
+
+A master running `--shards N` now serves full replication to a single-shard
+replica — previously PSYNC was rejected with `-ERR PSYNC across multiple
+shards is not yet supported`.
+
+- **Per-shard atomic snapshot legs.** A new `ShardMessage::PrepareReplicaSync`
+  fans out to every shard (own shard via the self queue, the rest over the
+  SPSC mesh). Each shard serializes its keyspace slice to an RDB *body*,
+  captures its replication offset, and registers the replica's live channel
+  in ONE synchronous stretch on its own thread — per shard, nothing can land
+  between "inside the snapshot" and "streamed live", so there is no backlog
+  catch-up leg and non-idempotent commands (INCR) can never double-apply.
+- **One merged Redis-format RDB.** The PSYNC task stitches the per-shard
+  bodies into a single valid RDB (`redis_rdb::write_rdb_merged`: header +
+  bodies + EOF/CRC64) and answers `+FULLRESYNC <replid> <Σ shard offsets>`
+  with one `$<len>` bulk — the replica's existing R0 loader needs no changes.
+  Index definitions ride once; graph content is sharded, so the snapshot
+  carries one `moon-graph-store` aux entry per shard and the replica imports
+  all of them (`install_graph_store_many`, `read_moon_aux_all`).
+- **Per-record SELECT framing on the merged wire.** N shard threads feed one
+  replica socket, so a shared "current db" context cannot exist: on
+  multi-shard masters every db-scoped record is fused with its own
+  `SELECT <db>` prefix (single channel send / backlog append pair / offset
+  advance) — no cross-shard interleave can split a SELECT from the write it
+  frames. Gated on the replica-attach hint; single-shard masters keep the
+  cheaper emit-on-change tracking.
+- **Partial resync degrades to full.** A replica's single scalar offset
+  cannot be mapped back onto N per-shard backlogs, so a multi-shard master
+  answers every PSYNC (any replid/offset) with `+FULLRESYNC`.
+- Overflow-kick (task #35), `REPLCONF ACK`, and `WAIT` all carry over: the
+  summed snapshot offset keeps `total_offset - base == bytes on wire`, so
+  WAIT/ACK math stays exact on multi-shard masters.
+- New e2e suite `tests/replication_multishard.rs`: 2/4/8-shard full resync +
+  live-stream convergence with INCR exactness, interleaved multi-db writers
+  with db-leak asserts, per-shard graph snapshot import, and the
+  partial→full degradation handshake.
+- Known limitation (unchanged from R1): master-side PSYNC requires
+  `runtime-monoio` (the default); the tokio build has no master-side PSYNC
+  intercept. Multi-shard *replicas* remain unsupported (`--shards 1`).
+
 ### Fixed — consistency/durability defects caught by the R1 gates (task #35)
 
 Three pre-existing data-integrity bugs surfaced by the new load/kill-9 gates

src/persistence/redis_rdb.rs

@@ -462,7 +462,17 @@ pub fn write_rdb_refs_with_moon_aux(
     for (key, value) in moon_aux {
         write_aux(buf, key, value);
     }
+    write_rdb_body_refs(databases, buf);
+    write_rdb_footer(buf);
+}
 
+/// Body-only writer: per-DB SELECTDB/RESIZEDB sections and entries — NO
+/// header, aux, EOF, or CRC. Building block for [`write_rdb_merged`]: a
+/// multi-shard master has each shard serialize its own keyspace slice with
+/// this, then stitches the bodies into ONE valid RDB. Repeated SELECTDB
+/// opcodes for the same db index (one per shard) are valid RDB — loaders
+/// treat SELECTDB as "switch current db" and accumulate entries.
+pub fn write_rdb_body_refs(databases: &[&Database], buf: &mut Vec<u8>) {
     let now_ms = current_time_ms();
 
     for (db_idx, db) in databases.iter().enumerate() {
@@ -486,7 +496,7 @@ pub fn write_rdb_refs_with_moon_aux(
         buf.push(RDB_OPCODE_SELECTDB);
         write_length(buf, db_idx as u64);
 
-        // RESIZEDB
+        // RESIZEDB (per-shard slice size — a presize hint only, safe to repeat)
         buf.push(RDB_OPCODE_RESIZEDB);
         write_length(buf, live.len() as u64);
         write_length(buf, expires_count as u64);
@@ -495,7 +505,21 @@ pub fn write_rdb_refs_with_moon_aux(
             write_rdb_entry(buf, key.as_bytes(), entry, base_ts);
         }
     }
+}
 
+/// Assemble ONE valid Redis-format RDB from pre-serialized per-shard bodies
+/// (each produced by [`write_rdb_body_refs`]): header + moon aux fields +
+/// concatenated bodies + EOF/CRC footer. The multi-shard PSYNC full resync
+/// (R2, task #20) sends this as a single `$<len>` bulk so a single-shard
+/// replica loads it with the unchanged R0 path.
+pub fn write_rdb_merged(moon_aux: &[(&[u8], &[u8])], bodies: &[Vec<u8>], buf: &mut Vec<u8>) {
+    write_rdb_header(buf);
+    for (key, value) in moon_aux {
+        write_aux(buf, key, value);
+    }
+    for body in bodies {
+        buf.extend_from_slice(body);
+    }
     write_rdb_footer(buf);
 }
 
@@ -551,6 +575,39 @@ pub fn read_moon_aux(data: &[u8], key: &[u8]) -> Option<Vec<u8>> {
     }
 }
 
+/// Like [`read_moon_aux`] but collects EVERY occurrence of `key` in the
+/// header aux block, in write order. A merged multi-shard snapshot
+/// ([`write_rdb_merged`]) carries one `moon-graph-store` aux entry PER shard —
+/// graph content is sharded, so the replica must import all of them, not just
+/// the first. Returns an empty Vec for foreign/truncated buffers or when the
+/// key never appears.
+pub fn read_moon_aux_all(data: &[u8], key: &[u8]) -> Vec<Vec<u8>> {
+    let mut out = Vec::new();
+    if data.len() < 9 || &data[..5] != REDIS_RDB_MAGIC || &data[5..9] != REDIS_RDB_VERSION {
+        return out;
+    }
+    let mut cursor = Cursor::new(data);
+    cursor.set_position(9); // magic + version
+    loop {
+        let mut opcode = [0u8; 1];
+        if cursor.read_exact(&mut opcode).is_err() {
+            return out;
+        }
+        if opcode[0] != RDB_OPCODE_AUX {
+            return out;
+        }
+        let Ok(k) = read_redis_string(&mut cursor) else {
+            return out;
+        };
+        let Ok(v) = read_redis_string(&mut cursor) else {
+            return out;
+        };
+        if k == key {
+            out.push(v);
+        }
+    }
+}
+
 /// Load an RDB file in Redis format into the provided databases.
 ///
 /// Verifies magic bytes, version, and CRC64 checksum.
@@ -767,6 +824,74 @@ mod tests {
         assert_eq!(loaded, 0);
     }
 
+    /// R2 (task #20): a merged multi-shard snapshot — per-shard bodies with
+    /// REPEATED SELECTDB sections for the same db — must load as one keyspace,
+    /// and repeated graph aux entries must all be readable in write order.
+    #[test]
+    fn merged_multishard_rdb_round_trip() {
+        // "Shard 0": keys in db0 and db2. "Shard 1": different keys, same dbs.
+        let mk = |pairs: &[(usize, &str, &str)]| {
+            let mut dbs = vec![Database::new(), Database::new(), Database::new()];
+            for (db_idx, k, v) in pairs {
+                dbs[*db_idx].set(
+                    Bytes::copy_from_slice(k.as_bytes()),
+                    Entry::new_string(Bytes::copy_from_slice(v.as_bytes())),
+                );
+            }
+            dbs
+        };
+        let shard0 = mk(&[(0, "a", "1"), (2, "c", "3")]);
+        let shard1 = mk(&[(0, "b", "2"), (2, "d", "4")]);
+
+        let mut body0 = Vec::new();
+        write_rdb_body_refs(&shard0.iter().collect::<Vec<_>>(), &mut body0);
+        let mut body1 = Vec::new();
+        write_rdb_body_refs(&shard1.iter().collect::<Vec<_>>(), &mut body1);
+
+        let mut merged = Vec::new();
+        write_rdb_merged(
+            &[
+                (MOON_AUX_VECTOR_DEFS, b"vec-defs"),
+                (MOON_AUX_GRAPH_STORE, b"graph-shard-0"),
+                (MOON_AUX_GRAPH_STORE, b"graph-shard-1"),
+            ],
+            &[body0, body1],
+            &mut merged,
+        );
+
+        // Single-occurrence reader still finds the first entry of each key.
+        assert_eq!(
+            read_moon_aux(&merged, MOON_AUX_VECTOR_DEFS).as_deref(),
+            Some(&b"vec-defs"[..])
+        );
+        // All-occurrences reader returns every shard's graph blob in order.
+        assert_eq!(
+            read_moon_aux_all(&merged, MOON_AUX_GRAPH_STORE),
+            vec![b"graph-shard-0".to_vec(), b"graph-shard-1".to_vec()]
+        );
+        assert!(read_moon_aux_all(&merged, b"moon-unknown").is_empty());
+
+        // The merged blob is ONE valid RDB (magic/version/CRC) whose repeated
+        // SELECTDB sections accumulate into a single keyspace.
+        let mut loaded = vec![Database::new(), Database::new(), Database::new()];
+        let n = load_rdb(&mut loaded, &merged).expect("merged RDB must load");
+        assert_eq!(n, 4);
+        let get = |db: &Database, k: &str| {
+            db.data()
+                .iter()
+                .find(|(key, _)| key.as_bytes() == k.as_bytes())
+                .map(|(_, e)| match e.as_redis_value() {
+                    RedisValueRef::String(s) => s.to_vec(),
+                    _ => panic!("expected string"),
+                })
+        };
+        assert_eq!(get(&loaded[0], "a").as_deref(), Some(&b"1"[..]));
+        assert_eq!(get(&loaded[0], "b").as_deref(), Some(&b"2"[..]));
+        assert_eq!(get(&loaded[2], "c").as_deref(), Some(&b"3"[..]));
+        assert_eq!(get(&loaded[2], "d").as_deref(), Some(&b"4"[..]));
+        assert!(get(&loaded[1], "a").is_none());
+    }
+
     #[test]
     fn read_moon_aux_rejects_foreign_and_truncated_buffers() {
         assert_eq!(read_moon_aux(b"", MOON_AUX_VECTOR_DEFS), None);

src/replication/apply.rs

@@ -460,8 +460,11 @@ pub(crate) fn load_snapshot(rdb: &[u8]) -> anyhow::Result<usize> {
     // header) carry the FT index DEFINITIONS; standard RDB loaders skip them.
     let vec_defs = redis_rdb::read_moon_aux(rdb, redis_rdb::MOON_AUX_VECTOR_DEFS);
     let text_defs = redis_rdb::read_moon_aux(rdb, redis_rdb::MOON_AUX_TEXT_DEFS);
+    // R2 (task #20): a multi-shard master's merged snapshot carries one
+    // graph-store aux entry PER shard (graph content is sharded) — collect
+    // them all; a single-shard snapshot yields exactly one.
     #[cfg(feature = "graph")]
-    let graph_blob = redis_rdb::read_moon_aux(rdb, redis_rdb::MOON_AUX_GRAPH_STORE);
+    let graph_blobs = redis_rdb::read_moon_aux_all(rdb, redis_rdb::MOON_AUX_GRAPH_STORE);
     match crate::shard::slice::try_with_shard(|s| {
         for db in s.databases.iter_mut() {
             db.clear();
@@ -472,13 +475,19 @@ pub(crate) fn load_snapshot(rdb: &[u8]) -> anyhow::Result<usize> {
         // (authoritative replace — an EMPTY blob drops replica-local graphs;
         // an ABSENT aux means a pre-graph-sync master, warn-and-keep).
         #[cfg(feature = "graph")]
-        match graph_blob.as_deref() {
-            Some(blob) => {
-                match crate::replication::graph_sync::install_graph_store(&mut s.graph_store, blob)
-                {
+        match &graph_blobs[..] {
+            blobs if !blobs.is_empty() => {
+                match crate::replication::graph_sync::install_graph_store_many(
+                    &mut s.graph_store,
+                    blobs,
+                ) {
                     Some(n) => {
                         if n > 0 {
-                            tracing::info!("replica snapshot: installed {} graph(s)", n);
+                            tracing::info!(
+                                "replica snapshot: installed {} graph(s) from {} shard blob(s)",
+                                n,
+                                blobs.len()
+                            );
                         }
                     }
                     None => {
@@ -488,7 +497,7 @@ pub(crate) fn load_snapshot(rdb: &[u8]) -> anyhow::Result<usize> {
                     }
                 }
             }
-            None => {
+            _ => {
                 if s.graph_store.graph_count() > 0 {
                     tracing::warn!(
                         "replica snapshot carried no graph-store aux but {} local graph(s) \

src/replication/graph_sync.rs

@@ -183,15 +183,39 @@ pub fn export_graph_store(store: &mut GraphStore) -> Vec<u8> {
 /// (store is left in whatever partial state was reached — the caller aborts
 /// the sync and the replica retries with a fresh full resync).
 pub fn install_graph_store(store: &mut GraphStore, blob: &[u8]) -> Option<usize> {
-    let mut cur = Cursor { data: blob, pos: 0 };
-    if cur.u8()? != FORMAT_VERSION {
-        return None;
+    drop_all_local_graphs(store);
+    install_graphs_additive(store, blob)
+}
+
+/// R2 (task #20): install a MULTI-SHARD snapshot's graph blobs — one per
+/// master shard (`read_moon_aux_all` order). Authoritative replace happens
+/// ONCE, then every blob installs additively. Graph names are disjoint across
+/// blobs (each graph lives on exactly one master shard); a duplicate name is
+/// malformed input and fails the install (`None`).
+pub fn install_graph_store_many(store: &mut GraphStore, blobs: &[Vec<u8>]) -> Option<usize> {
+    drop_all_local_graphs(store);
+    let mut total = 0usize;
+    for blob in blobs {
+        total += install_graphs_additive(store, blob)?;
     }
-    // Authoritative replace: drop everything local first.
+    Some(total)
+}
+
+/// Authoritative replace leg shared by both install entry points.
+fn drop_all_local_graphs(store: &mut GraphStore) {
     let local: Vec<Bytes> = store.list_graphs().into_iter().cloned().collect();
     for name in local {
         let _ = store.drop_graph(&name);
     }
+}
+
+/// Decode one export blob and create its graphs on top of whatever the store
+/// already holds. Callers handle the drop-local leg.
+fn install_graphs_additive(store: &mut GraphStore, blob: &[u8]) -> Option<usize> {
+    let mut cur = Cursor { data: blob, pos: 0 };
+    if cur.u8()? != FORMAT_VERSION {
+        return None;
+    }
 
     let graph_count = cur.u32()? as usize;
     for _ in 0..graph_count {