Add a ScanCursor wrapper for store iteration (#669)

* Add a ScanCursor wrapper for IterateMainStore() and IterateObjectStore() to avoid tempKV
Remove a couple unused Iterate functions
Fix reversed Assert.AreEqual params in a test

* formatting fix

* Convert another ScanFunctions implementation to class

---------

Co-authored-by: Badrish Chandramouli <[email protected]>
Co-authored-by: Vasileios Zois <[email protected]>

Author: 3 people

libs/cluster/Server/Migration/MigrationKeyIterationFunctions.cs

@@ -14,7 +14,7 @@ internal sealed unsafe partial class MigrateSession : IDisposable
     {
         internal sealed class MigrationKeyIterationFunctions
         {
-            internal unsafe struct MainStoreGetKeysInSlots : IScanIteratorFunctions<SpanByte, SpanByte>
+            internal sealed unsafe class MainStoreGetKeysInSlots : IScanIteratorFunctions<SpanByte, SpanByte>
             {
                 MigrationScanIterator iterator;
 
@@ -90,7 +90,7 @@ public void OnStop(bool completed, long numberOfRecords) { }
                 public void OnException(Exception exception, long numberOfRecords) { }
             }
 
-            internal struct MigrationScanIterator
+            internal sealed class MigrationScanIterator
             {
                 readonly MigrateSession session;
                 readonly HashSet<int> slots;

libs/cluster/Session/ClusterKeyIterationFunctions.cs

@@ -13,8 +13,9 @@ internal sealed unsafe partial class ClusterSession : IClusterSession
     {
         internal static class ClusterKeyIterationFunctions
         {
-            internal struct MainStoreCountKeys : IScanIteratorFunctions<SpanByte, SpanByte>
+            internal sealed class MainStoreCountKeys : IScanIteratorFunctions<SpanByte, SpanByte>
             {
+                // This must be a class as it is passed through pending IO operations
                 internal int keyCount;
                 readonly int slot;
 
@@ -34,8 +35,9 @@ public void OnStop(bool completed, long numberOfRecords) { }
                 public void OnException(Exception exception, long numberOfRecords) { }
             }
 
-            internal struct ObjectStoreCountKeys : IScanIteratorFunctions<byte[], IGarnetObject>
+            internal sealed class ObjectStoreCountKeys : IScanIteratorFunctions<byte[], IGarnetObject>
             {
+                // This must be a class as it is passed through pending IO operations
                 internal int keyCount;
                 readonly int slot;
 

libs/server/Storage/Session/Common/ArrayKeyIterationFunctions.cs

@@ -111,15 +111,15 @@ internal unsafe bool DbScan(ArgSlice patternB, bool allKeys, long cursor, out lo
         }
 
         /// <summary>
-        /// Iterate the contents of the main store
+        /// Iterate the contents of the main store (push-based)
         /// </summary>
         /// <typeparam name="TScanFunctions"></typeparam>
         /// <param name="scanFunctions"></param>
         /// <param name="untilAddress"></param>
         /// <returns></returns>
         internal bool IterateMainStore<TScanFunctions>(ref TScanFunctions scanFunctions, long untilAddress = -1)
             where TScanFunctions : IScanIteratorFunctions<SpanByte, SpanByte>
-            => basicContext.Session.Iterate(ref scanFunctions, untilAddress);
+            => basicContext.Session.IterateLookup(ref scanFunctions, untilAddress);
 
         /// <summary>
         /// Iterate the contents of the main store (pull based)
@@ -128,15 +128,15 @@ internal ITsavoriteScanIterator<SpanByte, SpanByte> IterateMainStore()
             => basicContext.Session.Iterate();
 
         /// <summary>
-        /// Iterate the contents of the object store
+        /// Iterate the contents of the object store (push-based)
         /// </summary>
         /// <typeparam name="TScanFunctions"></typeparam>
         /// <param name="scanFunctions"></param>
         /// <param name="untilAddress"></param>
         /// <returns></returns>
         internal bool IterateObjectStore<TScanFunctions>(ref TScanFunctions scanFunctions, long untilAddress = -1)
             where TScanFunctions : IScanIteratorFunctions<byte[], IGarnetObject>
-            => objectStoreBasicContext.Session.Iterate(ref scanFunctions, untilAddress);
+            => objectStoreBasicContext.Session.IterateLookup(ref scanFunctions, untilAddress);
 
         /// <summary>
         /// Iterate the contents of the main store (pull based)

libs/storage/Tsavorite/cs/src/core/Allocator/AllocatorScan.cs

@@ -8,22 +8,6 @@
 
 namespace Tsavorite.core
 {
-    internal sealed class ScanCursorState<TKey, TValue>
-    {
-        internal IScanIteratorFunctions<TKey, TValue> functions;
-        internal long acceptedCount;    // Number of records pushed to and accepted by the caller
-        internal bool endBatch;         // End the batch (but return a valid cursor for the next batch, as of "count" records had been returned)
-        internal bool stop;             // Stop the operation (as if all records in the db had been returned)
-
-        internal void Initialize(IScanIteratorFunctions<TKey, TValue> scanIteratorFunctions)
-        {
-            functions = scanIteratorFunctions;
-            acceptedCount = 0;
-            endBatch = false;
-            stop = false;
-        }
-    }
-
     public abstract partial class AllocatorBase<TKey, TValue, TStoreFunctions, TAllocator> : IDisposable
         where TStoreFunctions : IStoreFunctions<TKey, TValue>
         where TAllocator : IAllocator<TKey, TValue, TStoreFunctions>
@@ -205,7 +189,8 @@ internal unsafe bool GetFromDiskAndPushToReader<TScanFunctions>(ref TKey key, re
         internal abstract bool ScanCursor<TScanFunctions>(TsavoriteKV<TKey, TValue, TStoreFunctions, TAllocator> store, ScanCursorState<TKey, TValue> scanCursorState, ref long cursor, long count, TScanFunctions scanFunctions, long endAddress, bool validateCursor)
             where TScanFunctions : IScanIteratorFunctions<TKey, TValue>;
 
-        private protected bool ScanLookup<TInput, TOutput, TScanFunctions, TScanIterator>(TsavoriteKV<TKey, TValue, TStoreFunctions, TAllocator> store, ScanCursorState<TKey, TValue> scanCursorState, ref long cursor, long count, TScanFunctions scanFunctions, TScanIterator iter, bool validateCursor)
+        private protected bool ScanLookup<TInput, TOutput, TScanFunctions, TScanIterator>(TsavoriteKV<TKey, TValue, TStoreFunctions, TAllocator> store,
+                ScanCursorState<TKey, TValue> scanCursorState, ref long cursor, long count, TScanFunctions scanFunctions, TScanIterator iter, bool validateCursor)
             where TScanFunctions : IScanIteratorFunctions<TKey, TValue>
             where TScanIterator : ITsavoriteScanIterator<TKey, TValue>, IPushScanIterator<TKey>
         {

libs/storage/Tsavorite/cs/src/core/Allocator/ScanCursorState.cs

@@ -0,0 +1,21 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+namespace Tsavorite.core
+{
+    internal sealed class ScanCursorState<TKey, TValue>
+    {
+        internal IScanIteratorFunctions<TKey, TValue> functions;
+        internal long acceptedCount;    // Number of records pushed to and accepted by the caller
+        internal bool endBatch;         // End the batch (but return a valid cursor for the next batch, as of "count" records had been returned)
+        internal bool stop;             // Stop the operation (as if all records in the db had been returned)
+
+        internal void Initialize(IScanIteratorFunctions<TKey, TValue> scanIteratorFunctions)
+        {
+            functions = scanIteratorFunctions;
+            acceptedCount = 0;
+            endBatch = false;
+            stop = false;
+        }
+    }
+}

libs/storage/Tsavorite/cs/src/core/ClientSession/ClientSession.cs

@@ -467,7 +467,7 @@ public ITsavoriteScanIterator<TKey, TValue> Iterate(long untilAddress = -1)
             => store.Iterate<TInput, TOutput, TContext, TFunctions>(functions, untilAddress);
 
         /// <summary>
-        /// Push iteration of all (distinct) live key-values stored in Tsavorite
+        /// Push iteration of all (distinct) live key-values stored in Tsavorite, using a temporary TsavoriteKV to ensure uniqueness
         /// </summary>
         /// <param name="scanFunctions">Functions receiving pushed records</param>
         /// <param name="untilAddress">Report records until this address (tail by default)</param>
@@ -476,6 +476,21 @@ public bool Iterate<TScanFunctions>(ref TScanFunctions scanFunctions, long until
             where TScanFunctions : IScanIteratorFunctions<TKey, TValue>
             => store.Iterate<TInput, TOutput, TContext, TFunctions, TScanFunctions>(functions, ref scanFunctions, untilAddress);
 
+        /// <summary>
+        /// Push iteration of all (distinct) live key-values stored in Tsavorite, using a lookup strategy to ensure uniqueness
+        /// </summary>
+        /// <param name="scanFunctions">Functions receiving pushed records</param>
+        /// <param name="untilAddress">Report records until this address (tail by default)</param>
+        /// <returns>True if Iteration completed; false if Iteration ended early due to one of the TScanIterator reader functions returning false</returns>
+        public bool IterateLookup<TScanFunctions>(ref TScanFunctions scanFunctions, long untilAddress = -1)
+            where TScanFunctions : IScanIteratorFunctions<TKey, TValue>
+        {
+            if (untilAddress == -1)
+                untilAddress = store.Log.TailAddress;
+            var cursor = 0L;
+            return ScanCursor(ref cursor, count: long.MaxValue, scanFunctions, endAddress: untilAddress);
+        }
+
         /// <summary>
         /// Push-scan the log from <paramref name="cursor"/> (which should be a valid address) and push up to <paramref name="count"/> records
         /// to the caller via <paramref name="scanFunctions"/> for each Key that is not found at a higher address.

libs/storage/Tsavorite/cs/src/core/Index/Tsavorite/Implementation/FindRecord.cs

@@ -49,6 +49,7 @@ internal bool TryFindRecordInMainLog(ref TKey key, ref OperationStackContext<TKe
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        // Return true if the record is found in the log, else false and an indication of whether we need to do IO to continue the search
         internal bool TryFindRecordInMainLogForConditionalOperation<TInput, TOutput, TContext, TSessionFunctionsWrapper>(TSessionFunctionsWrapper sessionFunctions,
                 ref TKey key, ref OperationStackContext<TKey, TValue, TStoreFunctions, TAllocator> stackCtx, long minAddress, out OperationStatus internalStatus, out bool needIO)
             where TSessionFunctionsWrapper : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>

libs/storage/Tsavorite/cs/src/core/Index/Tsavorite/TsavoriteIterator.cs

@@ -51,26 +51,6 @@ public bool Iterate<TInput, TOutput, TContext, TFunctions, TScanFunctions>(TFunc
             scanFunctions.OnStop(!stop, numRecords);
             return !stop;
         }
-
-        /// <summary>
-        /// Iterator for all (distinct) live key-values stored in Tsavorite
-        /// </summary>
-        /// <param name="untilAddress">Report records until this address (tail by default)</param>
-        /// <returns>Tsavorite iterator</returns>
-        [Obsolete("Invoke Iterate() on a client session (ClientSession), or use store.Iterate overload with Functions provided as parameter")]
-        public ITsavoriteScanIterator<TKey, TValue> Iterate(long untilAddress = -1)
-            => throw new TsavoriteException("Invoke Iterate() on a client session (ClientSession), or use store.Iterate overload with Functions provided as parameter");
-
-        /// <summary>
-        /// Iterator for all (distinct) live key-values stored in Tsavorite
-        /// </summary>
-        /// <param name="compactionFunctions">User provided compaction functions (see <see cref="ICompactionFunctions{Key, Value}"/>).</param>
-        /// <param name="untilAddress">Report records until this address (tail by default)</param>
-        /// <returns>Tsavorite iterator</returns>
-        [Obsolete("Invoke Iterate() on a client session (ClientSession), or use store.Iterate overload with Functions provided as parameter")]
-        public ITsavoriteScanIterator<TKey, TValue> Iterate<CompactionFunctions>(CompactionFunctions compactionFunctions, long untilAddress = -1)
-            where CompactionFunctions : ICompactionFunctions<TKey, TValue>
-            => throw new TsavoriteException("Invoke Iterate() on a client session (ClientSession), or use store.Iterate overload with Functions provided as parameter");
     }
 
     internal sealed class TsavoriteKVIterator<TKey, TValue, TInput, TOutput, TContext, TFunctions, TStoreFunctions, TAllocator> : ITsavoriteScanIterator<TKey, TValue>

test/Garnet.test.cluster/ClusterMigrateTests.cs

@@ -266,12 +266,12 @@ public void ClusterSimpleSlotInfo()
 
             var sourceIndex = context.clusterTestUtils.GetSourceNodeIndexFromSlot((ushort)slot, context.logger);
             var expectedKeyCount = context.clusterTestUtils.CountKeysInSlot(slot, context.logger);
-            ClassicAssert.AreEqual(expectedKeyCount, keyCount);
+            ClassicAssert.AreEqual(keyCount, expectedKeyCount);
             _ = context.clusterTestUtils.CountKeysInSlot(-1, context.logger);
             _ = context.clusterTestUtils.CountKeysInSlot(ushort.MaxValue, context.logger);
 
             var result = context.clusterTestUtils.GetKeysInSlot(sourceIndex, slot, expectedKeyCount, context.logger);
-            ClassicAssert.AreEqual(result.Count, keyCount);
+            ClassicAssert.AreEqual(keyCount, result.Count);
             _ = context.clusterTestUtils.GetKeysInSlot(-1, expectedKeyCount);
             _ = context.clusterTestUtils.GetKeysInSlot(ushort.MaxValue, expectedKeyCount);