cleanup migrate keys

Author: vazois

libs/cluster/Server/Migration/MigrateSession.cs

@@ -78,8 +78,6 @@ internal sealed unsafe partial class MigrateSession : IDisposable
         /// </summary>
         public LimitedFixedBufferPool GetNetworkPool => clusterProvider.migrationManager.GetNetworkPool;
 
-        readonly GarnetClientSession _gcs;
-
         /// <summary>
         /// Check for overlapping slots between migrate sessions
         /// </summary>
@@ -94,15 +92,10 @@ public bool Overlap(MigrateSession session)
         readonly TransferOption transferOption;
 
         /// <summary>
-        /// MigrateSlotsScan for background slot migrate tasks
+        /// MigrateTask for background slot migrate tasks
         /// </summary>
         readonly MigrateTask[] migrateTasks;
 
-        /// <summary>
-        /// LocalServerSessions for background slot migrate tasks
-        /// </summary>
-        readonly LocalServerSession[] localServerSessions;
-
         /// <summary>
         /// MigrateSession Constructor
         /// </summary>
@@ -154,25 +147,16 @@ internal MigrateSession(
 
             Status = MigrateState.PENDING;
 
-            // Single key value size + few bytes for command header and arguments
-            _gcs = GetGarnetClient();
-
             if (transferOption == TransferOption.SLOTS)
             {
                 migrateTasks = new MigrateTask[clusterProvider.serverOptions.ParallelMigrateTasks];
-                localServerSessions = new LocalServerSession[clusterProvider.serverOptions.ParallelMigrateTasks];
                 for (var i = 0; i < migrateTasks.Length; i++)
-                {
                     migrateTasks[i] = new MigrateTask(this);
-                    localServerSessions[i] = new LocalServerSession(clusterProvider.storeWrapper);
-                }
             }
             else
             {
                 migrateTasks = new MigrateTask[1];
                 migrateTasks[0] = new MigrateTask(this, sketch: sketch);
-                localServerSessions = new LocalServerSession[1];
-                localServerSessions[0] = new LocalServerSession(clusterProvider.storeWrapper);
             }
         }
 
@@ -197,26 +181,22 @@ public void Dispose()
             if (!_disposed.TryWriteLock()) return;
             _cts?.Cancel();
             _cts?.Dispose();
-            _gcs.Dispose();
 
             for (var i = 0; i < migrateTasks.Length; i++)
-            {
                 migrateTasks[i].Dispose();
-                localServerSessions[i].Dispose();
-            }
         }
 
-        private bool CheckConnection()
+        private bool CheckConnection(GarnetClientSession client)
         {
-            bool status = true;
-            if (!_gcs.IsConnected)
+            var status = true;
+            if (!client.IsConnected)
             {
-                _gcs.Reconnect((int)_timeout.TotalMilliseconds);
+                client.Reconnect((int)_timeout.TotalMilliseconds);
                 if (_passwd != null)
                 {
                     try
                     {
-                        status = _gcs.Authenticate(_username, _passwd).ContinueWith(resp =>
+                        status = client.Authenticate(_username, _passwd).ContinueWith(resp =>
                         {
                             // Check if authenticate succeeded
                             if (!resp.Result.Equals("OK", StringComparison.Ordinal))
@@ -274,9 +254,10 @@ private bool CheckConnection()
         public bool TrySetSlotRanges(string nodeid, MigrateState state)
         {
             var status = false;
+            var client = migrateTasks[0].Client;
             try
             {
-                if (!CheckConnection())
+                if (!CheckConnection(client))
                     return false;
                 var stateBytes = state switch
                 {
@@ -286,7 +267,7 @@ public bool TrySetSlotRanges(string nodeid, MigrateState state)
                     _ => throw new Exception("Invalid SETSLOT Operation"),
                 };
 
-                status = _gcs.SetSlotRange(stateBytes, nodeid, _slotRanges).ContinueWith(resp =>
+                status = client.SetSlotRange(stateBytes, nodeid, _slotRanges).ContinueWith(resp =>
                 {
                     // Check if setslotsrange executed correctly
                     if (!resp.Result.Equals("OK", StringComparison.Ordinal))

libs/cluster/Server/Migration/MigrateSessionCommonUtils.cs

@@ -12,19 +12,10 @@ namespace Garnet.cluster
 {
     internal sealed unsafe partial class MigrateSession : IDisposable
     {
-        /// <summary>
-        /// Write to network buffer or send if full the previous payload and then write to network before the associated kv-pair.
-        /// </summary>
-        /// <param name="gcs"></param>
-        /// <param name="localServerSession"></param>
-        /// <param name="key"></param>
-        /// <param name="input"></param>
-        /// <param name="o"></param>
-        /// <returns></returns>
-        private bool WriteOrSendMainStoreKeyValuePair(GarnetClientSession gcs, LocalServerSession localServerSession, ref SpanByte key, ref RawStringInput input, ref SpanByteAndMemory o)
+        private bool WriteOrSendMainStoreKeyValuePair(GarnetClientSession gcs, LocalServerSession localServerSession, ref SpanByte key, ref RawStringInput input, ref SpanByteAndMemory o, out GarnetStatus status)
         {
             // Read value for key
-            var status = localServerSession.BasicGarnetApi.Read_MainStore(ref key, ref input, ref o);
+            status = localServerSession.BasicGarnetApi.Read_MainStore(ref key, ref input, ref o);
 
             // Skip if key NOTFOUND
             if (status == GarnetStatus.NOTFOUND)
@@ -64,13 +55,13 @@ bool WriteOrSendMainStoreKeyValuePair(GarnetClientSession gcs, ref SpanByte key,
             }
         }
 
-        private bool WriteOrSendObjectStoreKeyValuePair(GarnetClientSession gcs, LocalServerSession localServerSession, ref ArgSlice key, ref SpanByteAndMemory o)
+        private bool WriteOrSendObjectStoreKeyValuePair(GarnetClientSession gcs, LocalServerSession localServerSession, ref ArgSlice key, out GarnetStatus status)
         {
             var keyByteArray = key.ToArray();
 
             ObjectInput input = default;
             GarnetObjectStoreOutput value = default;
-            var status = localServerSessions[0].BasicGarnetApi.Read_ObjectStore(ref keyByteArray, ref input, ref value);
+            status = localServerSession.BasicGarnetApi.Read_ObjectStore(ref keyByteArray, ref input, ref value);
 
             // Skip if key NOTFOUND
             if (status == GarnetStatus.NOTFOUND)
@@ -103,54 +94,6 @@ bool WriteOrSendObjectStoreKeyValuePair(GarnetClientSession gcs, byte[] key, byt
             }
         }
 
-        /// <summary>
-        /// Write main store key-value pair directly to client buffer or flush buffer to make space and try again writing.
-        /// </summary>
-        /// <param name="key"></param>
-        /// <param name="value"></param>
-        /// <returns>True on success, else false</returns>
-        private bool WriteOrSendMainStoreKeyValuePair(ref SpanByte key, ref SpanByte value)
-        {
-            // Check if we need to initialize cluster migrate command arguments
-            if (_gcs.NeedsInitialization)
-                _gcs.SetClusterMigrateHeader(_sourceNodeId, _replaceOption, isMainStore: true);
-
-            // Try write serialized key value to client buffer
-            while (!_gcs.TryWriteKeyValueSpanByte(ref key, ref value, out var task))
-            {
-                // Flush key value pairs in the buffer
-                if (!HandleMigrateTaskResponse(task))
-                    return false;
-
-                // re-initialize cluster migrate command parameters
-                _gcs.SetClusterMigrateHeader(_sourceNodeId, _replaceOption, isMainStore: true);
-            }
-            return true;
-        }
-
-        /// <summary>
-        /// Write object store key-value pair directly to client buffer or flush buffer to make space and try again writing.
-        /// </summary>
-        /// <param name="key"></param>
-        /// <param name="value"></param>
-        /// <param name="expiration"></param>
-        /// <returns></returns>
-        private bool WriteOrSendObjectStoreKeyValuePair(byte[] key, byte[] value, long expiration)
-        {
-            // Check if we need to initialize cluster migrate command arguments
-            if (_gcs.NeedsInitialization)
-                _gcs.SetClusterMigrateHeader(_sourceNodeId, _replaceOption, isMainStore: false);
-
-            while (!_gcs.TryWriteKeyValueByteArray(key, value, expiration, out var task))
-            {
-                // Flush key value pairs in the buffer
-                if (!HandleMigrateTaskResponse(task))
-                    return false;
-                _gcs.SetClusterMigrateHeader(_sourceNodeId, _replaceOption, isMainStore: false);
-            }
-            return true;
-        }
-
         /// <summary>
         /// Handle response from migrate data task
         /// </summary>

libs/cluster/Server/Migration/MigrateSessionKeyAccess.cs

@@ -38,9 +38,9 @@ public bool CanAccessKey(ref ArgSlice key, int slot, bool readOnly)
                 return true;
 
             var state = SketchStatus.INITIALIZING;
-            foreach (var migrateScan in migrateTasks)
+            foreach (var migrateTask in migrateTasks)
             {
-                if (migrateScan.sketch.Probe(key.SpanByte, out state))
+                if (migrateTask.sketch.Probe(key.SpanByte, out state))
                     goto found;
             }
 

libs/cluster/Server/Migration/MigrateSessionKeys.cs

@@ -25,52 +25,20 @@ private bool MigrateKeysFromMainStore()
             var bufPtr = buffer.GetValidPointer();
             var bufPtrEnd = bufPtr + bufferSize;
             var o = new SpanByteAndMemory(bufPtr, (int)(bufPtrEnd - bufPtr));
+            var migrateTask = migrateTasks[0];
 
             try
             {
                 // Transition keys to MIGRATING status
-                migrateTasks[0].sketch.SetStatus(SketchStatus.TRANSMITTING);
+                migrateTask.sketch.SetStatus(SketchStatus.TRANSMITTING);
                 WaitForConfigPropagation();
 
-                ////////////////
-                // Build Input//
-                ////////////////
-                var input = new RawStringInput(RespCommandAccessor.MIGRATE);
-                var keys = migrateTasks[0].sketch.Keys;
-                for (var i = 0; i < migrateTasks[0].sketch.Keys.Count; i++)
+                // Transmit keys from main store
+                if (!migrateTask.TransmitKeys(StoreType.Main))
                 {
-                    if (keys[i].Item2) continue;
-                    var spanByte = keys[i].Item1.SpanByte;
-
-                    // Read value for key
-                    var status = localServerSessions[0].BasicGarnetApi.Read_MainStore(ref spanByte, ref input, ref o);
-
-                    // Skip if key NOTFOUND
-                    if (status == GarnetStatus.NOTFOUND)
-                        continue;
-
-                    // Get SpanByte from stack if any
-                    ref var value = ref o.SpanByte;
-                    if (!o.IsSpanByte)
-                    {
-                        // Reinterpret heap memory to SpanByte
-                        value = ref SpanByte.ReinterpretWithoutLength(o.Memory.Memory.Span);
-                    }
-
-                    // Write key to network buffer if it has not expired
-                    if (!ClusterSession.Expired(ref value) && !WriteOrSendMainStoreKeyValuePair(ref spanByte, ref value))
-                        return false;
-
-                    // Reset SpanByte for next read if any but don't dispose heap buffer as we might re-use it
-                    o.SpanByte = new SpanByte((int)(bufPtrEnd - bufPtr), (IntPtr)bufPtr);
-
-                    // Mark for deletion
-                    keys[i] = (keys[i].Item1, true);
-                }
-
-                // Flush data in client buffer
-                if (!HandleMigrateTaskResponse(_gcs.SendAndResetIterationBuffer()))
+                    logger?.LogError("Failed transmitting keys from main store");
                     return false;
+                }
 
                 DeleteKeys();
             }
@@ -93,40 +61,18 @@ private bool MigrateKeysFromMainStore()
         /// <returns>True on success, false otherwise</returns>
         private bool MigrateKeysFromObjectStore()
         {
-            // NOTE: Any keys not found in main store are automatically set to QUEUED before this method is called
-            // Transition all QUEUED to MIGRATING state
-            migrateTasks[0].sketch.SetStatus(SketchStatus.TRANSMITTING);
+            var migrateTask = migrateTasks[0];
+            // NOTE: Any keys not found in main store are automatically set to INITIALIZING before this method is called
+            // Transition all INITIALIZING to TRANSMITTING state
+            migrateTask.sketch.SetStatus(SketchStatus.TRANSMITTING);
             WaitForConfigPropagation();
 
-            var keys = migrateTasks[0].sketch.Keys;
-            for (var i = 0; i < migrateTasks[0].sketch.Keys.Count; i++)
+            // Transmit keys from object store
+            if (!migrateTask.TransmitKeys(StoreType.Object))
             {
-                if (keys[i].Item2) continue;
-                var keyByteArray = keys[i].Item1.ToArray();
-
-                ObjectInput input = default;
-                GarnetObjectStoreOutput value = default;
-                var status = localServerSessions[0].BasicGarnetApi.Read_ObjectStore(ref keyByteArray, ref input, ref value);
-
-                // Skip if key NOTFOUND
-                if (status == GarnetStatus.NOTFOUND)
-                    continue;
-
-                if (!ClusterSession.Expired(ref value.GarnetObject))
-                {
-                    var objectData = GarnetObjectSerializer.Serialize(value.GarnetObject);
-
-                    if (!WriteOrSendObjectStoreKeyValuePair(keyByteArray, objectData, value.GarnetObject.Expiration))
-                        return false;
-                }
-
-                // Mark for deletion
-                keys[i] = (keys[i].Item1, true);
-            }
-
-            // Flush data in client buffer
-            if (!HandleMigrateTaskResponse(_gcs.SendAndResetIterationBuffer()))
+                logger?.LogError("Failed transmitting keys from object store");
                 return false;
+            }
 
             // Delete keys if COPY option is false or transition KEYS from MIGRATING to MIGRATED status
             DeleteKeys();
@@ -138,51 +84,39 @@ private bool MigrateKeysFromObjectStore()
         /// </summary>
         private void DeleteKeys()
         {
-            if (_copyOption)
-                goto migrated;
-
-            // Transition to deleting to block read requests
-            migrateTasks[0].sketch.SetStatus(SketchStatus.DELETING);
+            var migrateTask = migrateTasks[0];
+            // Transition to deleting to block read requests                
+            migrateTask.sketch.SetStatus(SketchStatus.DELETING);
             WaitForConfigPropagation();
 
-            foreach (var pair in migrateTasks[0].sketch.Keys)
-            {
-                if (!pair.Item2) continue;
-                var spanByte = pair.Item1.SpanByte;
-                _ = localServerSessions[0].BasicGarnetApi.DELETE(ref spanByte);
-            }
+            // Delete keys
+            migrateTask.DeleteKeys();
 
-        migrated:
             // Transition to MIGRATED to release waiting operations
-            migrateTasks[0].sketch.SetStatus(SketchStatus.MIGRATED);
+            migrateTask.sketch.SetStatus(SketchStatus.MIGRATED);
             WaitForConfigPropagation();
         }
 
         /// <summary>
         /// Method used to migrate keys from main and object stores.
         /// This method is used to process the MIGRATE KEYS transfer option.
         /// </summary>
-        /// <param name="storeType"></param>
         /// <returns></returns>
-        public bool MigrateKeys(StoreType storeType = StoreType.All)
+        public bool MigrateKeys()
         {
             try
             {
-                if (!CheckConnection())
+                var migrateTask = migrateTasks[0];
+                if (!migrateTask.Initialize())
                     return false;
 
-                if (storeType is StoreType.All or StoreType.Main)
-                {
-                    // Migrate main store keys
-                    _gcs.InitializeIterationBuffer(clusterProvider.storeWrapper.loggingFrequency);
-                    if (!MigrateKeysFromMainStore())
-                        return false;
-                }
+                // Migrate main store keys
+                if (!MigrateKeysFromMainStore())
+                    return false;
 
                 // Migrate object store keys
-                if (!clusterProvider.serverOptions.DisableObjects && storeType is StoreType.All or StoreType.Object)
+                if (!clusterProvider.serverOptions.DisableObjects)
                 {
-                    _gcs.InitializeIterationBuffer(clusterProvider.storeWrapper.loggingFrequency);
                     if (!MigrateKeysFromObjectStore())
                         return false;
                 }