Align Migration.cs with current main APIs and update migration docs

Tiago Napoli · Copilot · Tiago Napoli · commit 9e457010494b · 2026-05-25T17:30:25.000-07:00
- Update RangeIndexManager.Migration.cs for 7 API changes from main:
  dataDir-&gt;riLogRoot, liveIndexes key nint-&gt;Guid, SnapshotToFile-&gt;CPR
  pattern, evicted tree paths, RecoverFromCprSnapshot, LogDataPathFor,
  stub.ResetFlags()
- Document TRYAGAIN behavior for RI commands during migration instead
  of spin-wait, with rationale (pipeline stall, client timeout, and
  connection reset risks)
- Add Redis protocol context and SE.Redis client handling details

Co-authored-by: Copilot &lt;223556219+Copilot@users.noreply.github.com&gt;
diff --git a/libs/server/Resp/RangeIndex/RangeIndexManager.Migration.cs b/libs/server/Resp/RangeIndex/RangeIndexManager.Migration.cs
@@ -6,6 +6,7 @@
 using System.IO;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
+using System.Threading;
 using Garnet.common;
 using Garnet.server.BfTreeInterop;
 using Microsoft.Extensions.Logging;
@@ -82,11 +83,11 @@ internal RangeIndexMigrationReader SnapshotRangeIndexAndCreateReader(LocalServer
 
         /// <summary>
         /// Derive a temporary file path for an in-progress inbound migration.
-        /// Format: {dataDir}/rangeindex/.migration-tmp/{guid}.bftree
+        /// Format: {riLogRoot}/.migration-tmp/{guid}.bftree
         /// </summary>
         internal string DeriveTempMigrationPath()
         {
-            var tmpDir = Path.Combine(dataDir ?? string.Empty, "rangeindex", ".migration-tmp");
+            var tmpDir = Path.Combine(riLogRoot ?? string.Empty, ".migration-tmp");
             Directory.CreateDirectory(tmpDir);
             return Path.Combine(tmpDir, $"{Guid.NewGuid():N}.bftree");
         }
@@ -99,7 +100,7 @@ internal unsafe bool PublishMigratedIndex(ReadOnlySpan<byte> keyBytes, ReadOnlyS
         {
             try
             {
-                var workingPath = DeriveWorkingPath(keyBytes);
+                var workingPath = LogDataPathFor(keyBytes);
                 Directory.CreateDirectory(Path.GetDirectoryName(workingPath)!);
 
                 // If a data file already exists (e.g., from a previous migration of the same key
@@ -111,21 +112,18 @@ internal unsafe bool PublishMigratedIndex(ReadOnlySpan<byte> keyBytes, ReadOnlyS
 
                 ref readonly var srcStub = ref ReadIndex(stubBytes);
 
-                var bfTree = BfTreeService.RecoverFromSnapshot(
+                var scratchPath = LogScratchPathFor(keyBytes);
+                var bfTree = BfTreeService.RecoverFromCprSnapshot(
                     workingPath,
-                    (StorageBackendType)srcStub.StorageBackend,
-                    srcStub.CacheSize,
-                    srcStub.MinRecordSize,
-                    srcStub.MaxRecordSize,
-                    srcStub.MaxKeyLen,
-                    srcStub.LeafPageSize);
+                    scratchPath,
+                    (StorageBackendType)srcStub.StorageBackend);
 
                 Span<byte> newStubBytes = stackalloc byte[IndexSizeBytes];
                 stubBytes.CopyTo(newStubBytes);
                 ref var newStub = ref Unsafe.As<byte, RangeIndexStub>(
                     ref MemoryMarshal.GetReference(newStubBytes));
                 newStub.TreeHandle = bfTree.NativePtr;
-                newStub.Flags = 0;
+                newStub.ResetFlags();
                 newStub.SerializationPhase = 0;
 
                 var parseState = new SessionParseState();
@@ -164,10 +162,10 @@ internal unsafe bool PublishMigratedIndex(ReadOnlySpan<byte> keyBytes, ReadOnlyS
         /// <summary>
         /// Source side: snapshot a BfTree for migration under an exclusive lock.
         /// Acquires the exclusive lock, re-reads the stub from the store to get a fresh
-        /// <c>TreeHandle</c>. If the tree is live, snapshots it to a temporary migration
-        /// file. If evicted, copies the existing flush/checkpoint snapshot file to a
-        /// temporary migration file. The temp file is safe from concurrent overwrites
-        /// by <see cref="SnapshotTreeForFlush"/> or the native BfTree.
+        /// <c>TreeHandle</c>. If the tree is live, takes a CPR snapshot and copies
+        /// the scratch file to a temporary migration file (following the same pattern as
+        /// <see cref="SnapshotForFlushViaCpr"/>). If evicted, copies the working
+        /// <c>data.bftree</c> file (same source as checkpoint pending entries).
         /// </summary>
         internal bool SnapshotForMigration(StorageSession session, ReadOnlySpan<byte> keyBytes, out string path, out long totalBytes)
         {
@@ -176,6 +174,7 @@ internal bool SnapshotForMigration(StorageSession session, ReadOnlySpan<byte> ke
 
             Span<byte> stubSpan = stackalloc byte[IndexSizeBytes];
             var keyHash = session.stringBasicContext.GetKeyHash((FixedSpanByteKey)PinnedSpanByte.FromPinnedSpan(keyBytes));
+            var hashPrefix = HashKeyToPrefix(keyBytes);
             var migrationPath = DeriveTempMigrationPath();
 
             rangeIndexLocks.AcquireExclusiveLock(keyHash, out var lockToken);
@@ -201,26 +200,33 @@ internal bool SnapshotForMigration(StorageSession session, ReadOnlySpan<byte> ke
                     return false;
                 }
 
-                if (stub.TreeHandle != nint.Zero && liveIndexes.TryGetValue(stub.TreeHandle, out var treeEntry))
+                var keyId = KeyId(keyBytes);
+                if (stub.TreeHandle != nint.Zero && liveIndexes.TryGetValue(keyId, out var treeEntry) && treeEntry.Tree != null)
                 {
-                    // Tree is live — snapshot to a temp migration file (safe from concurrent OnFlush overwrites)
-                    treeEntry.Tree.SnapshotToFile(migrationPath);
+                    // Tree is live — CPR snapshot + copy from scratch path (same pattern as SnapshotForFlushViaCpr)
+                    while (!treeEntry.TryClaimSnapshot()) Thread.Yield();
+                    try
+                    {
+                        BfTreeService.CprSnapshotByPtr(treeEntry.Tree.NativePtr);
+                        File.Copy(LogScratchPath(hashPrefix), migrationPath, overwrite: true);
+                    }
+                    finally
+                    {
+                        treeEntry.ReleaseSnapshot();
+                    }
                 }
                 else
                 {
-                    // Tree was evicted — copy the existing flush/checkpoint file to a temp migration file.
-                    // Tsavorite guarantees OnFlush runs before OnEvict, so a snapshot file must exist.
-                    var existingPath = stub.IsRecovered && recoveredCheckpointToken != Guid.Empty
-                        ? DeriveCheckpointPath(keyBytes, recoveredCheckpointToken)
-                        : DeriveFlushPath(keyBytes);
+                    // Tree was evicted — copy the working data.bftree file (same source as checkpoint pending entries)
+                    var dataPath = LogDataPath(hashPrefix);
 
-                    if (!File.Exists(existingPath))
+                    if (!File.Exists(dataPath))
                     {
-                        logger?.LogWarning("SnapshotForMigration: expected snapshot file not found: {Path}", existingPath);
+                        logger?.LogWarning("SnapshotForMigration: data.bftree not found: {Path}", dataPath);
                         return false;
                     }
 
-                    File.Copy(existingPath, migrationPath, overwrite: true);
+                    File.Copy(dataPath, migrationPath, overwrite: true);
                 }
             }
             finally
diff --git a/website/docs/dev/range-index-resp-api.md b/website/docs/dev/range-index-resp-api.md
@@ -2673,7 +2673,7 @@ and restores from the snapshot files at the expected paths.
 
 **Problem.** During slot migration, a stub-only transfer is insufficient: the 51-byte stub
 points to a process-local native `BfTree` whose on-disk data file lives outside Tsavorite
-(`{dataDir}/rangeindex/{key_hash}/data.bftree`). The target node has no access to that
+(`{riLogRoot}/{key_hash}/data.bftree`). The target node has no access to that
 file. We must ship the entire tree-data file alongside the stub, and the target must
 rebuild the native `BfTree` from it before publishing a usable stub into its main store.
 
@@ -2695,7 +2695,7 @@ Source side (async):
   TransmitRangeIndexAsync loop:
     reader.ReadNextChunkAsync(buffer)
       ├→ fileStream.ReadAsync() — async file read
-      └→ serializer.MoveNext(dest, fileData, out consumed) — sync framing
+      └→ serializer.SupplyFileData(buffer) + serializer.MoveNext(dest) — sync framing
     WriteOrSendRecordSpanAsync() — send over network
 
 Destination side (sync):
@@ -2757,21 +2757,24 @@ The stream format across one or more chunks:
 3. **Snapshot details.** `SnapshotForMigration` (in `RangeIndexManager.Migration.cs`):
    - Acquires the per-key exclusive lock.
    - Re-reads the stub from the main store (`RIGET`) under the lock to get a fresh `TreeHandle`.
-   - If the tree is live (`TreeHandle ≠ 0`, in `liveIndexes`): snapshots to
-     `{dataDir}/rangeindex/.migration-tmp/{guid}.bftree`.
-   - If the tree was evicted: copies `flush.bftree` or `snapshot.{token}.bftree` to the
-     same temp directory.
+   - If the tree is live (`TreeHandle ≠ 0`, in `liveIndexes`): takes a CPR snapshot
+     (`CprSnapshotByPtr`) and copies the scratch file to
+     `{riLogRoot}/.migration-tmp/{guid}.bftree`, using `TryClaimSnapshot`/`ReleaseSnapshot`
+     to serialize against concurrent flush/checkpoint snapshots (same pattern as
+     `SnapshotForFlushViaCpr`).
+   - If the tree was evicted: copies the working `data.bftree` file (same source as
+     checkpoint pending entries) to the same temp directory.
    - Releases the exclusive lock.
 
    The temp file has a unique GUID name, protecting it from concurrent `OnFlush` overwrites,
    checkpoint operations, and other migrations.
 
 4. **Serializer architecture.** `RangeIndexChunkedSerializer` is a pure state machine with
-   phases: `KeyHeader → KeyData → FileHeader → FileData → Trailer → Done`. It takes file
-   data as input via `MoveNext(dest, fileData, out consumed)` rather than reading files
-   itself. Properties `NeedsFileData` and `FileDataRemaining` tell the caller when to supply
-   file bytes. The `RangeIndexMigrationReader` wrapper handles async file reads and loops
-   to handle phase transitions within a single call.
+   phases: `KeyHeader → KeyData → FileHeader → FileData → Trailer → Done`. File data is
+   supplied via `SupplyFileData(Memory<byte>)` rather than passed directly to `MoveNext`.
+   The `NeedsFileData` property tells the caller when to supply more file bytes. The
+   `RangeIndexMigrationReader` wrapper handles async file reads and loops to handle phase
+   transitions within a single call.
 
 #### Source side — KEYS path
 
@@ -2797,21 +2800,79 @@ receiving a stream, only `SerializedRangeIndexStream` records are accepted.
 `RangeIndexChunkedDeserializer` is a state machine that:
 1. Parses the key length header and accumulates key bytes.
 2. Parses the file size header.
-3. Writes file bytes to a temp file (`{dataDir}/rangeindex/.migration-tmp/{guid}.bftree`),
+3. Writes file bytes to a temp file (`{riLogRoot}/.migration-tmp/{guid}.bftree`),
    updating an xxHash64 incrementally.
 4. Parses the trailer: validates checksum, extracts stub.
 
 On completion, `Publish()`:
-1. Moves the temp file to the working path (`{dataDir}/rangeindex/{key_hash}/data.bftree`).
+1. Moves the temp file to the working path (`{riLogRoot}/{key_hash}/data.bftree`).
    If a data file already exists (e.g., from a previous migration of the same key), it is
    deleted first — this enables round-trip migration (P0→P1→P0).
-2. Calls `BfTreeService.RecoverFromSnapshot()` to load the native tree.
-3. Rewrites the stub: `TreeHandle = bfTree.NativePtr`, clears `Flags` and `SerializationPhase`.
+2. Calls `BfTreeService.RecoverFromCprSnapshot()` to load the native tree.
+3. Rewrites the stub: `TreeHandle = bfTree.NativePtr`, calls `ResetFlags()` and clears `SerializationPhase`.
 4. Publishes via `RICREATE` RMW into the main store.
 5. Registers the tree via `RegisterIndex()`.
 
 On disposal (error or abort), the temp file is deleted.
 
+#### Write protection during migration
+
+RI commands (`RI.SET`, `RI.CREATE`, `RI.DEL`, `RI.GET`, `RI.SCAN`, etc.) are classified as
+data commands and have `KeySpecifications` in the command metadata JSON. This means the
+generic slot verification layer (`CanServeSlot` → `NetworkMultiKeySlotVerify`) automatically
+extracts the key, hashes it to a slot, and checks migration state — no special code is
+needed inside the RI command handlers.
+
+When a slot is in `MIGRATING` state, `CanOperateOnKey` calls
+`migrationManager.CanAccessKey(key, slot, readOnly)`, which probes the sketch. Behavior
+depends on the sketch status:
+
+| Sketch Status | Read commands (RI.GET, RI.SCAN…) | Write commands (RI.SET, RI.CREATE…) |
+|---|---|---|
+| `INITIALIZING` | ✅ Allowed | ✅ Allowed |
+| `TRANSMITTING` | ✅ Allowed | ❌ **Returns `-TRYAGAIN`** |
+| `DELETING` | ❌ Returns `-TRYAGAIN` | ❌ Returns `-TRYAGAIN` |
+| `MIGRATED` | ✅ Allowed | ✅ Allowed |
+
+During the `TRANSMITTING` phase — while the CPR snapshot is being taken and streamed —
+write commands return a `-TRYAGAIN` error immediately, signaling the client to retry later.
+Read commands remain unblocked during transmission. During the `DELETING` phase, both reads
+and writes return `-TRYAGAIN`.
+
+**Why TRYAGAIN instead of spin-wait.** The default Garnet behavior for regular keys during
+migration is to spin-wait (`Thread.Yield()` in a loop). This is acceptable for regular keys
+because individual key transmission is fast. However, RangeIndex migration can take
+significantly longer — the BfTree must be snapshotted via CPR and streamed as chunked data.
+A long spin-wait has cascading effects on multiplexed clients like StackExchange.Redis:
+
+- SE.Redis uses a single TCP connection with pipelined commands in FIFO order.
+- If the server spin-waits on command A, **all subsequent pipelined commands on that
+  connection also stall** — even for unrelated keys.
+- SE.Redis has a per-connection `SyncTimeout` (default 5s) / `AsyncTimeout`. Commands
+  exceeding the timeout are faulted with `TimeoutException` on the client side, but the
+  server-side command **continues executing** (fire-and-forget after timeout).
+- If nothing is read from the pipe for 4× the timeout (~20s), SE.Redis detects a "dead
+  socket" and **kills the connection entirely**, triggering a reconnect.
+
+Returning `-TRYAGAIN` avoids these problems: the response is immediate, the pipeline stays
+unblocked, and the client can implement retry logic with backoff.
+
+> **TRYAGAIN in the Redis protocol.** `-TRYAGAIN` is an official Redis cluster error
+> (see `cluster.c:1466`). Redis uses it for multi-key commands during migration when some
+> keys exist locally and some have already migrated — the command cannot be served atomically.
+> Our use extends this semantics: the key exists, but it cannot be served right now because
+> it is being serialized for migration. The error message and retry semantics are the same.
+>
+> **Client handling.** SE.Redis does not auto-retry on TRYAGAIN — it surfaces as a
+> `RedisServerException("TRYAGAIN ...")` that the application must catch and retry manually.
+> This is consistent with Redis behavior for TRYAGAIN (no client auto-retries it).
+>
+> **Comparison with Redis.** Redis is single-threaded, so migration and client commands are
+> naturally serialized — no spin-wait or TRYAGAIN is needed for single-key operations.
+> When a slot is `MIGRATING`, Redis serves the request if the key exists locally and returns
+> `-ASK` redirect if it doesn't. Redis only returns `-TRYAGAIN` for multi-key commands where
+> some keys have already migrated.
+
 #### Failure modes
 
 - **Transmit failure.** `TransmitRangeIndexAsync` catches all exceptions and returns `false`.
