Add RangeIndex cluster migration with sketch protection

Migration plumbing:
- Chunked serializer (pure state machine, no I/O) + async MigrationReader wrapper
- Chunked deserializer with FileStream I/O for receiving migration data
- TransmitRangeIndexAsync source-side driver with configurable chunk size
- RangeIndexMigrationReceiveState receiver with state machine
- Remove redundant content-length prefix in migration wire format
- SnapshotRangeIndexAndCreateReader factory on RangeIndexManager

Bug fixes:
- Remove RI+cluster startup guard (GarnetServer.cs)
- Fix round-trip migration: Publish deletes existing data file
- Fix Publish registration: accept InPlaceUpdated/CopyUpdated
- Fix serializer empty-buffer bug in FileData phase
- Add missing buffer-too-small guard for KeyHeader phase
- TransmitRangeIndexAsync catches all exceptions (never throws)

Sketch protection (SLOTS path):
- All RI keys added to sketch in one batch
- TRANSMITTING epoch barrier blocks writes during snapshot+transmit
- DELETING epoch barrier blocks reads+writes during delete
- try/finally ensures sketch cleanup on failure
- TODO: Vector Sets have the same unprotected pattern

Sketch protection (KEYS path):
- RI keys already in sketch from user enumeration
- Mark transmitted keys for DeleteKeysAsync()

Tests:
- 26 unit tests for serializer/deserializer
- 11 cluster migration integration tests (SingleBySlot, ByKeys,
  ManyBySlot, WhileModifying, MigrateBack, LargeTree, ChunkSize,
  StressAsync)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: Tiago Napoli

libs/client/ClientSession/GarnetClientSessionIncremental.cs

@@ -42,6 +42,12 @@ public enum MigrationRecordSpanType : byte
         /// Bespoke encoding for Vector Set indexes.
         /// </summary>
         VectorSetIndex = 3,
+
+        /// <summary>
+        /// Chunked serialization stream for a RangeIndex key during migration.
+        /// The receiver uses a state machine to track the in-progress stream.
+        /// </summary>
+        SerializedRangeIndexStream = 4,
     }
 
     public sealed unsafe partial class GarnetClientSession : IServerHook, IMessageConsumer

libs/cluster/Server/Migration/MigrateOperation.cs

@@ -22,15 +22,20 @@ internal sealed partial class MigrateOperation
             public StoreScan storeScan;
 
             private readonly ConcurrentDictionary<byte[], byte[]> vectorSetsIndexKeysToMigrate;
+            private readonly ConcurrentDictionary<byte[], byte[]> rangeIndexKeysToMigrate;
 
             readonly MigrateSession session;
             readonly GarnetClientSession gcs;
             readonly LocalServerSession localServerSession;
 
             public GarnetClientSession Client => gcs;
 
+            public LocalServerSession LocalSession => localServerSession;
+
             public IEnumerable<KeyValuePair<byte[], byte[]>> VectorSets => vectorSetsIndexKeysToMigrate;
 
+            public IEnumerable<KeyValuePair<byte[], byte[]>> RangeIndexes => rangeIndexKeysToMigrate;
+
             public void ThrowIfCancelled() => session._cts.Token.ThrowIfCancellationRequested();
 
             public bool Contains(int slot) => session._sslots.Contains(slot);
@@ -46,6 +51,8 @@ public bool ContainsNamespace(ReadOnlySpan<byte> namespaceBytes)
             public void EncounteredVectorSet(byte[] key, byte[] value)
             => vectorSetsIndexKeysToMigrate.TryAdd(key, value);
 
+            public void EncounteredRangeIndex(byte[] key, byte[] value) => rangeIndexKeysToMigrate.TryAdd(key, value);
+
             public MigrateOperation(MigrateSession session, Sketch sketch = null, int batchSize = 1 << 18)
             {
                 this.session = session;
@@ -55,6 +62,7 @@ public MigrateOperation(MigrateSession session, Sketch sketch = null, int batchS
                 storeScan = new StoreScan(this);
                 keysToDelete = [];
                 vectorSetsIndexKeysToMigrate = new(ByteArrayComparer.Instance);
+                rangeIndexKeysToMigrate = new(ByteArrayComparer.Instance);
             }
 
             public async ValueTask<bool> InitializeAsync()
@@ -127,7 +135,7 @@ public async Task<bool> TransmitSlotsAsync()
                 return true;
             }
 
-            public async Task<bool> TransmitKeysAsync(Dictionary<byte[], byte[]> vectorSetKeysToIgnore)
+            public async Task<bool> TransmitKeysAsync(Dictionary<byte[], byte[]> vectorSetKeysToIgnore, Dictionary<byte[], byte[]> rangeIndexKeysToIgnore)
             {
                 // Use this for both stores; main store will just use the SpanByteAndMemory directly. We want it to be outside iterations
                 // so we can reuse the SpanByteAndMemory.Memory across iterations.
@@ -136,7 +144,8 @@ public async Task<bool> TransmitKeysAsync(Dictionary<byte[], byte[]> vectorSetKe
                 var output = new UnifiedOutput();
 
 #if NET9_0_OR_GREATER
-                var ignoreLookup = vectorSetKeysToIgnore.GetAlternateLookup<ReadOnlySpan<byte>>();
+                var vectorSetIgnoreLookup = vectorSetKeysToIgnore.GetAlternateLookup<ReadOnlySpan<byte>>();
+                var rangeIndexIgnoreLookup = rangeIndexKeysToIgnore.GetAlternateLookup<ReadOnlySpan<byte>>();
 #endif
 
                 try
@@ -154,19 +163,23 @@ public async Task<bool> TransmitKeysAsync(Dictionary<byte[], byte[]> vectorSetKe
 
                         var spanByte = keys[i].Item1;
 
-                        // Don't transmit if a Vector Set
-                        var isVectorSet =
-                            vectorSetKeysToIgnore.Count > 0 &&
+                        // Don't transmit VectorSet or RangeIndex keys — these are handled out-of-band
+                        var shouldSkip =
+                            (vectorSetKeysToIgnore.Count > 0 &&
 #if NET9_0_OR_GREATER
-                            ignoreLookup.ContainsKey(spanByte.ReadOnlySpan);
+                            vectorSetIgnoreLookup.ContainsKey(spanByte.ReadOnlySpan)) ||
 #else
-                                vectorSetKeysToIgnore.ContainsKey(spanByte.ToArray());
+                            vectorSetKeysToIgnore.ContainsKey(spanByte.ToArray())) ||
+#endif
+                            (rangeIndexKeysToIgnore.Count > 0 &&
+#if NET9_0_OR_GREATER
+                            rangeIndexIgnoreLookup.ContainsKey(spanByte.ReadOnlySpan));
+#else
+                            rangeIndexKeysToIgnore.ContainsKey(spanByte.ToArray()));
 #endif
 
-                        if (isVectorSet)
-                        {
+                        if (shouldSkip)
                             continue;
-                        }
 
                         if (!await session.WriteOrSendRecordAsync(gcs, localServerSession, keys[i].Item1, ref input, ref output, out var status).ConfigureAwait(false))
                             return false;
@@ -285,6 +298,19 @@ public void DeleteVectorSet(PinnedSpanByte key)
 
                 session.logger?.LogDebug("Deleting Vector Set {key} after migration: {delRes}", System.Text.Encoding.UTF8.GetString(key), delRes);
             }
+
+            /// <summary>
+            /// Delete a RangeIndex after migration if _copyOption is not set.
+            /// </summary>
+            public void DeleteRangeIndex(PinnedSpanByte key)
+            {
+                if (session._copyOption)
+                    return;
+
+                var delRes = localServerSession.BasicGarnetApi.DELETE(key);
+
+                session.logger?.LogDebug("Deleting RangeIndex {key} after migration: {delRes}", System.Text.Encoding.UTF8.GetString(key), delRes);
+            }
         }
     }
 }

libs/cluster/Server/Migration/MigrateScanFunctions.cs

@@ -50,7 +50,12 @@ public bool Reader<TSourceLogRecord>(in TSourceLogRecord srcLogRecord, RecordMet
                     // Check if key belongs to slot that is being migrated and if it can be added to our buffer
                     if (migrateOperation.Contains(slot))
                     {
-                        if (srcLogRecord.RecordType == VectorManager.RecordType)
+                        if (srcLogRecord.RecordType == RangeIndexManager.RangeIndexRecordType)
+                        {
+                            // RangeIndex keys need out-of-band migration (snapshot + chunks)
+                            migrateOperation.EncounteredRangeIndex(key.ToArray(), srcLogRecord.ValueSpan.ToArray());
+                        }
+                        else if (srcLogRecord.RecordType == VectorManager.RecordType)
                         {
                             // We can't delete the vector set _yet_ nor can we migrate it, 
                             // we just need to remember it to migrate once the associated namespaces are all moved over

libs/cluster/Server/Migration/MigrateSession.RangeIndex.cs

@@ -0,0 +1,134 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+using System;
+using System.Buffers;
+using System.Collections.Generic;
+using System.Threading.Tasks;
+using Garnet.client;
+using Garnet.server;
+using Microsoft.Extensions.Logging;
+using Tsavorite.core;
+
+namespace Garnet.cluster
+{
+    /// <summary>
+    /// RangeIndex migration support: source-side transmit driver.
+    /// </summary>
+    internal sealed partial class MigrateSession : IDisposable
+    {
+        /// <summary>
+        /// Transmit a single RangeIndex key to the destination node.
+        /// Uses <see cref="RangeIndexManager.SnapshotRangeIndexAndCreateReader"/> to obtain an async
+        /// migration reader that snapshots and streams the BfTree data.
+        /// Forces a flush and awaits ACK. Does not delete the source key — the caller
+        /// is responsible for deletion in the appropriate sketch status phase.
+        /// </summary>
+        private async Task<bool> TransmitRangeIndexAsync(MigrateOperation mo, byte[] keyBytes, byte[] stubBytes, int chunkSize = RangeIndexManager.DefaultMigrationChunkSize)
+        {
+            var rangeIndexManager = clusterProvider.storeWrapper.DefaultDatabase.RangeIndexManager;
+            var gcs = mo.Client;
+
+            RangeIndexMigrationReader reader;
+            try
+            {
+                reader = rangeIndexManager.SnapshotRangeIndexAndCreateReader(mo.LocalSession, keyBytes, stubBytes, chunkSize);
+            }
+            catch (Exception ex)
+            {
+                logger?.LogError(ex, "TransmitRangeIndex: failed to snapshot BfTree for key");
+                return false;
+            }
+
+            var buffer = ArrayPool<byte>.Shared.Rent(chunkSize);
+            try
+            {
+                using (reader)
+                {
+                    while (!reader.IsComplete)
+                    {
+                        var payloadLen = await reader.ReadNextChunkAsync(buffer).ConfigureAwait(false);
+                        if (payloadLen == 0)
+                        {
+                            logger?.LogError("TransmitRangeIndex: reader returned zero-length payload with a {Size}-byte buffer", chunkSize);
+                            return false;
+                        }
+
+                        if (!await WriteOrSendRecordSpanAsync(gcs, MigrationRecordSpanType.SerializedRangeIndexStream, buffer.AsSpan(0, payloadLen)).ConfigureAwait(false))
+                        {
+                            logger?.LogError("TransmitRangeIndex: failed to write chunk");
+                            return false;
+                        }
+                    }
+
+                    // Force flush and await ACK
+                    if (!await HandleMigrateTaskResponseAsync(gcs.SendAndResetIterationBuffer()).ConfigureAwait(false))
+                    {
+                        logger?.LogError("TransmitRangeIndex: flush failed");
+                        return false;
+                    }
+
+                    return true;
+                }
+            }
+            catch (Exception ex)
+            {
+                logger?.LogError(ex, "TransmitRangeIndex: unexpected error during transmission");
+                return false;
+            }
+            finally
+            {
+                ArrayPool<byte>.Shared.Return(buffer);
+            }
+        }
+
+        /// <summary>
+        /// Migrate a batch of RangeIndex keys with sketch protection.
+        /// Adds all keys to the sketch, transitions through TRANSMITTING → DELETING → MIGRATED
+        /// with epoch barriers, ensuring concurrent operations are properly gated.
+        /// </summary>
+        private async Task<bool> MigrateRangeIndexKeysAsync(MigrateOperation mo, Dictionary<byte[], byte[]> rangeIndexKeys)
+        {
+            logger?.LogWarning("Migrating {count} RangeIndex keys", rangeIndexKeys.Count);
+
+            // Add all RI keys to sketch during INITIALIZING (no gating yet)
+            mo.sketch.Clear();
+            mo.sketch.SetStatus(SketchStatus.INITIALIZING);
+            foreach (var (key, _) in rangeIndexKeys)
+                mo.sketch.TryHashAndStore(key);
+
+            // Block writes during snapshot + transmit
+            mo.sketch.SetStatus(SketchStatus.TRANSMITTING);
+            await WaitForConfigPropagationAsync().ConfigureAwait(false);
+
+            try
+            {
+                foreach (var (key, stubBytes) in rangeIndexKeys)
+                {
+                    if (!await TransmitRangeIndexAsync(mo, key, stubBytes).ConfigureAwait(false))
+                    {
+                        logger?.LogError("Failed to migrate RangeIndex key");
+                        return false;
+                    }
+                }
+
+                // Block reads + writes during delete
+                mo.sketch.SetStatus(SketchStatus.DELETING);
+                await WaitForConfigPropagationAsync().ConfigureAwait(false);
+
+                foreach (var (key, _) in rangeIndexKeys)
+                {
+                    var pinnedKey = PinnedSpanByte.FromPinnedSpan(key);
+                    mo.DeleteRangeIndex(pinnedKey);
+                }
+
+                return true;
+            }
+            finally
+            {
+                // Always clean up the sketch, even on failure, to unblock client operations
+                mo.sketch.Clear();
+            }
+        }
+    }
+}

libs/cluster/Server/Migration/MigrateSessionKeys.cs

@@ -35,8 +35,13 @@ private async Task<bool> MigrateKeysFromStoreAsync()
                 await WaitForConfigPropagationAsync().ConfigureAwait(false);
 
                 // Discover Vector Sets linked namespaces
+                var allKeys = migrateTask.sketch.Keys.Select(t => t.Item1.ToArray());
                 var indexesToMigrate = new Dictionary<byte[], byte[]>(ByteArrayComparer.Instance);
-                _namespaces = clusterProvider.storeWrapper.DefaultDatabase.VectorManager.GetNamespacesForKeys(clusterProvider.storeWrapper, migrateTask.sketch.Keys.Select(t => t.Item1.ToArray()), indexesToMigrate);
+                _namespaces = clusterProvider.storeWrapper.DefaultDatabase.VectorManager.GetNamespacesForKeys(clusterProvider.storeWrapper, allKeys, indexesToMigrate);
+
+                // Discover RangeIndex keys upfront
+                var rangeIndexKeysToMigrate = clusterProvider.storeWrapper.DefaultDatabase.RangeIndexManager?.GetRangeIndexKeysForMigration(clusterProvider.storeWrapper, allKeys)
+                    ?? new Dictionary<byte[], byte[]>(ByteArrayComparer.Instance);
 
                 // If we have any namespaces, that implies Vector Sets, and if we have any of THOSE
                 // we need to reserve destination sets on the other side
@@ -46,8 +51,8 @@ private async Task<bool> MigrateKeysFromStoreAsync()
                     return false;
                 }
 
-                // Transmit keys from store
-                if (!await migrateTask.TransmitKeysAsync(indexesToMigrate).ConfigureAwait(false))
+                // Transmit keys from store (skipping VectorSet and RangeIndex keys, which are handled out-of-band)
+                if (!await migrateTask.TransmitKeysAsync(indexesToMigrate, rangeIndexKeysToMigrate).ConfigureAwait(false))
                 {
                     logger?.LogError("Failed transmitting keys from store");
                     return false;
@@ -122,6 +127,33 @@ private async Task<bool> MigrateKeysFromStoreAsync()
                         return false;
                     }
                 }
+
+                // Migrate RangeIndex keys (snapshot + chunk stream).
+                // Keys are already in the sketch (added by caller during key enumeration),
+                // so they're protected by the TRANSMITTING status. Mark for deletion so
+                // DeleteKeysAsync() handles them in the DELETING sketch status sequence.
+                if (rangeIndexKeysToMigrate.Count > 0)
+                {
+                    logger?.LogWarning("Migrating {count} RangeIndex keys via KEYS path", rangeIndexKeysToMigrate.Count);
+
+                    foreach (var (key, stubBytes) in rangeIndexKeysToMigrate)
+                    {
+                        if (!await TransmitRangeIndexAsync(migrateTask, key, stubBytes).ConfigureAwait(false))
+                        {
+                            logger?.LogError("Failed to migrate RangeIndex key via KEYS path");
+                            return false;
+                        }
+                    }
+
+                    // Mark all transmitted RI keys in the sketch for deletion by DeleteKeysAsync()
+                    var keys = migrateTask.sketch.Keys;
+                    for (var i = 0; i < keys.Count; i++)
+                    {
+                        if (rangeIndexKeysToMigrate.ContainsKey(keys[i].Item1.ToArray()))
+                            keys[i] = (keys[i].Item1, true);
+                    }
+                }
+
                 // Final cleanup, which will also delete Vector Sets
                 await DeleteKeysAsync().ConfigureAwait(false);
             }

libs/cluster/Server/Migration/MigrateSessionSlots.cs

@@ -114,6 +114,20 @@ async Task<bool> CreateAndRunMigrateTasksAsync(long beginAddress, long tailAddre
                         return false;
                     }
 
+                    // Handle migration of discovered RangeIndex keys with sketch protection.
+                    // All RI keys are added to the sketch in one batch, then a single epoch
+                    // barrier per phase gates concurrent operations.
+                    var rangeIndexKeys = migrateOperation.SelectMany(static mo => mo.RangeIndexes).GroupBy(static g => g.Key, ByteArrayComparer.Instance).ToDictionary(static g => g.Key, g => g.First().Value, ByteArrayComparer.Instance);
+
+                    if (rangeIndexKeys.Count > 0)
+                    {
+                        if (!await MigrateRangeIndexKeysAsync(migrateOperation[0], rangeIndexKeys).ConfigureAwait(false))
+                            return false;
+                    }
+
+                    // TODO: Vector Set index keys are excluded from the sketch during scan (MigrateScanFunctions.cs),
+                    // so concurrent writes are not blocked during transmit and reads are not blocked during delete.
+                    // This should be fixed with the same sketch-protected batch pattern used for RangeIndex keys above.
                     // Handle migration of discovered Vector Set keys now that they're namespaces have been moved
                     var vectorSets = migrateOperation.SelectMany(static mo => mo.VectorSets).GroupBy(static g => g.Key, ByteArrayComparer.Instance).ToDictionary(static g => g.Key, g => g.First().Value, ByteArrayComparer.Instance);
 

libs/cluster/Session/ClusterSession.cs

@@ -52,6 +52,7 @@ internal sealed partial class ClusterSession : IClusterSession
 
         private StringBasicContext stringBasicContext;
         private VectorBasicContext vectorBasicContext;
+        private readonly RangeIndexMigrationReceiveState rangeIndexMigrationState;
 
         public ClusterSession(ClusterProvider clusterProvider, TransactionManager txnManager, IGarnetAuthenticator authenticator, UserHandle userHandle, GarnetSessionMetrics sessionMetrics, BasicGarnetApi basicGarnetApi, StringBasicContext stringBasicContext, VectorBasicContext vectorBasicContext, INetworkSender networkSender, ILogger logger = null)
         {
@@ -65,6 +66,8 @@ public ClusterSession(ClusterProvider clusterProvider, TransactionManager txnMan
             this.vectorBasicContext = vectorBasicContext;
             this.networkSender = networkSender;
             this.logger = logger;
+            var riManager = clusterProvider.storeWrapper.DefaultDatabase.RangeIndexManager;
+            rangeIndexMigrationState = riManager != null ? new RangeIndexMigrationReceiveState(riManager) : null;
         }
 
         public unsafe void ProcessClusterCommands(RespCommand command, VectorManager vectorManager, ref SessionParseState parseState, ref byte* dcurr, ref byte* dend)
@@ -197,6 +200,8 @@ public void UnsafeSetConfig(string replicaOf = null)
 
         public void Dispose()
         {
+            rangeIndexMigrationState?.Dispose();
+
             // Call dispose on ref of this session if this session is a replication task
             if (IsReplicating)
                 replicaReplayDriverStore?.Dispose();