LogRecord and SpanByte changes

benchmark/BDN.benchmark/Custom/CustomTxnSet.cs

@@ -22,15 +22,15 @@ sealed class CustomTxnSet : CustomTransactionProcedure
         /// </summary>
         public const string CommandName = "CTXNSET";
 
-        ArgSlice setA;
-        ArgSlice setB;
-        ArgSlice setC;
-        ArgSlice setD;
+        PinnedSpanByte setA;
+        PinnedSpanByte setB;
+        PinnedSpanByte setC;
+        PinnedSpanByte setD;
 
-        ArgSlice valueA;
-        ArgSlice valueB;
-        ArgSlice valueC;
-        ArgSlice valueD;
+        PinnedSpanByte valueA;
+        PinnedSpanByte valueB;
+        PinnedSpanByte valueC;
+        PinnedSpanByte valueD;
 
         /// <summary>
         ///  CTXNSET key1 key2 key3 key4 value1 value2 value3 value4

libs/client/ClientSession/GarnetClientSession.cs

@@ -530,12 +530,17 @@ private unsafe void Flush()
                     Dispose();
                     throw;
                 }
-                networkSender.GetResponseObject();
-                offset = networkSender.GetResponseObjectHead();
-                end = networkSender.GetResponseObjectTail();
+                ResetOffset();
             }
         }
 
+        private unsafe void ResetOffset()
+        {
+            networkSender.GetResponseObject();
+            offset = networkSender.GetResponseObjectHead();
+            end = networkSender.GetResponseObjectTail();
+        }
+
         /// <inheritdoc />
         public bool TryCreateMessageConsumer(Span<byte> bytesReceived, INetworkSender networkSender, out IMessageConsumer session)
             => throw new NotSupportedException();

libs/client/ClientSession/GarnetClientSessionIncremental.cs

@@ -23,7 +23,7 @@ public sealed unsafe partial class GarnetClientSession : IServerHook, IMessageCo
         IncrementalSendType ist;
         bool isMainStore;
         byte* curr, head;
-        int keyValuePairCount;
+        int recordCount;
         TaskCompletionSource<string> currTcsIterationTask = null;
 
         /// <summary>
@@ -42,15 +42,37 @@ public sealed unsafe partial class GarnetClientSession : IServerHook, IMessageCo
         public bool NeedsInitialization => curr == null;
 
         /// <summary>
-        /// Flush and initialize buffers/parameters used for migrate command
+        /// Return a <see cref="Span{_byte_}"/> of all remaining available space in the network buffer.
+        /// </summary>
+        public PinnedSpanByte GetAvailableNetworkBufferSpan() => PinnedSpanByte.FromPinnedPointer(curr, (int)(end - curr));
+
+        public void IncrementRecordDirect(int size)
+        {
+            ++recordCount;
+            curr += size;
+        }
+
+        private void EnsureTcsIsEnqueued()
+        {
+            // See comments in SetClusterMigrateHeader() as to why this is decoupled from the header initialization.
+            if (recordCount > 0 && currTcsIterationTask == null)
+            {
+                currTcsIterationTask = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
+                tcsQueue.Enqueue(currTcsIterationTask);
+            }
+        }
+
+        /// <summary>
+        /// Flush and initialize buffers/parameters used for Migrate and Replica commands
         /// </summary>
         /// <param name="iterationProgressFreq"></param>
         public void InitializeIterationBuffer(TimeSpan iterationProgressFreq = default)
         {
+            EnsureTcsIsEnqueued();
             Flush();
             currTcsIterationTask = null;
             curr = head = null;
-            keyValuePairCount = 0;
+            recordCount = 0;
             this.iterationProgressFreq = default ? TimeSpan.FromSeconds(5) : iterationProgressFreq;
         }
 
@@ -59,117 +81,64 @@ public void InitializeIterationBuffer(TimeSpan iterationProgressFreq = default)
         /// </summary>
         public Task<string> SendAndResetIterationBuffer()
         {
-            if (keyValuePairCount == 0) return null;
+            Task<string> task = null;
+            if (recordCount == 0)
+            {
+                // No records to Flush(), but we need to reset buffer offsets as we may have written a header due to the need to initialize the buffer
+                // before passing it to Tsavorite as the output SpanByteAndMemory.SpanByte for Read().
+                ResetOffset();
+                goto done;
+            }
 
             Debug.Assert(end - curr >= 2);
             *curr++ = (byte)'\r';
             *curr++ = (byte)'\n';
 
             // Payload format = [$length\r\n][number of keys (4 bytes)][raw key value pairs]\r\n
             var size = (int)(curr - 2 - head - (ExtraSpace - 4));
-            TrackIterationProgress(keyValuePairCount, size);
+            TrackIterationProgress(recordCount, size);
             var success = RespWriteUtils.TryWritePaddedBulkStringLength(size, ExtraSpace - 4, ref head, end);
             Debug.Assert(success);
 
             // Number of key value pairs in payload
-            *(int*)head = keyValuePairCount;
+            *(int*)head = recordCount;
 
             // Reset offset and flush buffer
             offset = curr;
+            EnsureTcsIsEnqueued();
             Flush();
             Interlocked.Increment(ref numCommands);
 
             // Return outstanding task and reset current tcs
-            var task = currTcsIterationTask.Task;
+            task = currTcsIterationTask.Task;
             currTcsIterationTask = null;
+            recordCount = 0;
+
+        done:
             curr = head = null;
-            keyValuePairCount = 0;
             return task;
         }
 
         /// <summary>
-        /// Try write key value pair for main store directly to the client buffer
+        /// Try to write the span for the entire record directly to the client buffer
         /// </summary>
-        /// <param name="key"></param>
-        /// <param name="value"></param>
-        /// <param name="task"></param>
-        /// <returns></returns>
-        public bool TryWriteKeyValueSpanByte(ref SpanByte key, ref SpanByte value, out Task<string> task)
+        public bool TryWriteRecordSpan(ReadOnlySpan<byte> recordSpan, out Task<string> task)
         {
-            task = null;
-            // Try write key value pair directly to client buffer
-            if (!WriteSerializedSpanByte(ref key, ref value))
+            // We include space for newline at the end, to be added before sending
+            var totalLen = recordSpan.TotalSize() + 2;
+            if (totalLen > (int)(end - curr))
             {
-                // If failed to write because no space left send outstanding data and retrieve task
-                // Caller is responsible for retrying
+                // If there is no space left, send outstanding data and return the send-completion task.
+                // Caller is responsible for waiting for task completion and retrying.
                 task = SendAndResetIterationBuffer();
                 return false;
             }
 
-            keyValuePairCount++;
-            return true;
-
-            bool WriteSerializedSpanByte(ref SpanByte key, ref SpanByte value)
-            {
-                var totalLen = key.TotalSize + value.TotalSize + 2 + 2;
-                if (totalLen > (int)(end - curr))
-                    return false;
-
-                key.CopyTo(curr);
-                curr += key.TotalSize;
-                value.CopyTo(curr);
-                curr += value.TotalSize;
-                return true;
-            }
-        }
-
-        /// <summary>
-        /// Try write key value pair for object store directly to the client buffer
-        /// </summary>
-        /// <param name="key"></param>
-        /// <param name="value"></param>
-        /// <param name="expiration"></param>
-        /// <param name="task"></param>
-        /// <returns></returns>
-        public bool TryWriteKeyValueByteArray(byte[] key, byte[] value, long expiration, out Task<string> task)
-        {
+            recordSpan.SerializeTo(curr);
+            curr += recordSpan.TotalSize();
+            ++recordCount;
             task = null;
-            // Try write key value pair directly to client buffer
-            if (!WriteSerializedKeyValueByteArray(key, value, expiration))
-            {
-                // If failed to write because no space left send outstanding data and retrieve task
-                // Caller is responsible for retrying
-                task = SendAndResetIterationBuffer();
-                return false;
-            }
-
-            keyValuePairCount++;
             return true;
-
-            bool WriteSerializedKeyValueByteArray(byte[] key, byte[] value, long expiration)
-            {
-                // We include space for newline at the end, to be added before sending
-                int totalLen = 4 + key.Length + 4 + value.Length + 8 + 2;
-                if (totalLen > (int)(end - curr))
-                    return false;
-
-                *(int*)curr = key.Length;
-                curr += 4;
-                fixed (byte* keyPtr = key)
-                    Buffer.MemoryCopy(keyPtr, curr, key.Length, key.Length);
-                curr += key.Length;
-
-                *(int*)curr = value.Length;
-                curr += 4;
-                fixed (byte* valPtr = value)
-                    Buffer.MemoryCopy(valPtr, curr, value.Length, value.Length);
-                curr += value.Length;
-
-                *(long*)curr = expiration;
-                curr += 8;
-
-                return true;
-            }
         }
 
         long lastLog = 0;

libs/client/ClientSession/GarnetClientSessionMigrationExtensions.cs

@@ -172,8 +172,10 @@ public Task<string> SetSlotRange(Memory<byte> state, string nodeid, List<(int, i
         /// <param name="isMainStore"></param>
         public void SetClusterMigrateHeader(string sourceNodeId, bool replace, bool isMainStore)
         {
-            currTcsIterationTask = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
-            tcsQueue.Enqueue(currTcsIterationTask);
+            // For Migration we send the (curr - end) buffer as the SpanByteAndMemory.SpanByte output to Tsavorite. Thus we must
+            // initialize the header first, so we have curr properly positioned, but we cannot yet enqueue currTcsIterationTask.
+            // Therefore we defer this until the actual Flush(), when we know we have records to send. This is not a concern for
+            // Replication, because it uses an iterator and thus knows it has a record before it initializes the header.
             curr = offset;
             this.isMainStore = isMainStore;
             this.ist = IncrementalSendType.MIGRATE;
@@ -257,23 +259,24 @@ public Task<string> CompleteMigrate(string sourceNodeId, bool replace, bool isMa
 
             // Payload format = [$length\r\n][number of keys (4 bytes)][raw key value pairs]\r\n
             var size = (int)(curr - 2 - head - (ExtraSpace - 4));
-            TrackIterationProgress(keyValuePairCount, size, completed: true);
+            TrackIterationProgress(recordCount, size, completed: true);
             var success = RespWriteUtils.TryWritePaddedBulkStringLength(size, ExtraSpace - 4, ref head, end);
             Debug.Assert(success);
 
             // Number of key value pairs in payload
-            *(int*)head = keyValuePairCount;
+            *(int*)head = recordCount;
 
             // Reset offset and flush buffer
             offset = curr;
+            EnsureTcsIsEnqueued();
             Flush();
             Interlocked.Increment(ref numCommands);
 
             // Return outstanding task and reset current tcs
             var task = currTcsIterationTask.Task;
             currTcsIterationTask = null;
             curr = head = null;
-            keyValuePairCount = 0;
+            recordCount = 0;
             return task;
         }
     }

libs/client/ClientSession/GarnetClientSessionReplicationExtensions.cs

@@ -410,6 +410,8 @@ public Task<string> ExecuteAttachSync(byte[] syncMetadata)
         /// <param name="isMainStore"></param>
         public void SetClusterSyncHeader(string sourceNodeId, bool isMainStore)
         {
+            // Unlike Migration, where we don't know at the time of header initialization if we have a record or not, in Replication 
+            // we know we have a record at the time this is called, so we can initialize it directly.
             currTcsIterationTask = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
             tcsQueue.Enqueue(currTcsIterationTask);
             curr = offset;

libs/cluster/Server/ClusterManagerSlotState.cs

@@ -12,12 +12,12 @@
 
 namespace Garnet.cluster
 {
-    using BasicGarnetApi = GarnetApi<BasicContext<SpanByte, SpanByte, RawStringInput, SpanByteAndMemory, long, MainSessionFunctions,
-            /* MainStoreFunctions */ StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>,
-            SpanByteAllocator<StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>>>,
-        BasicContext<byte[], IGarnetObject, ObjectInput, GarnetObjectStoreOutput, long, ObjectSessionFunctions,
-            /* ObjectStoreFunctions */ StoreFunctions<byte[], IGarnetObject, ByteArrayKeyComparer, DefaultRecordDisposer<byte[], IGarnetObject>>,
-            GenericAllocator<byte[], IGarnetObject, StoreFunctions<byte[], IGarnetObject, ByteArrayKeyComparer, DefaultRecordDisposer<byte[], IGarnetObject>>>>>;
+    using BasicGarnetApi = GarnetApi<BasicContext<RawStringInput, SpanByteAndMemory, long, MainSessionFunctions,
+            /* MainStoreFunctions */ StoreFunctions<SpanByteComparer, SpanByteRecordDisposer>,
+            SpanByteAllocator<StoreFunctions<SpanByteComparer, SpanByteRecordDisposer>>>,
+        BasicContext<ObjectInput, GarnetObjectStoreOutput, long, ObjectSessionFunctions,
+            /* ObjectStoreFunctions */ StoreFunctions<SpanByteComparer, DefaultRecordDisposer>,
+            ObjectAllocator<StoreFunctions<SpanByteComparer, DefaultRecordDisposer>>>>;
 
     /// <summary>
     /// Cluster manager
@@ -477,12 +477,12 @@ public void TryResetSlotState(HashSet<int> slots)
         public static unsafe void DeleteKeysInSlotsFromMainStore(BasicGarnetApi BasicGarnetApi, HashSet<int> slots)
         {
             using var iter = BasicGarnetApi.IterateMainStore();
-            while (iter.GetNext(out _))
+            while (iter.GetNext())
             {
-                ref var key = ref iter.GetKey();
-                var s = HashSlotUtils.HashSlot(ref key);
+                var key = iter.Key;
+                var s = HashSlotUtils.HashSlot(key);
                 if (slots.Contains(s))
-                    _ = BasicGarnetApi.DELETE(ref key, StoreType.Main);
+                    _ = BasicGarnetApi.DELETE(PinnedSpanByte.FromPinnedSpan(key), StoreType.Main);
             }
         }
 
@@ -494,13 +494,12 @@ public static unsafe void DeleteKeysInSlotsFromMainStore(BasicGarnetApi BasicGar
         public static unsafe void DeleteKeysInSlotsFromObjectStore(BasicGarnetApi BasicGarnetApi, HashSet<int> slots)
         {
             using var iterObject = BasicGarnetApi.IterateObjectStore();
-            while (iterObject.GetNext(out _))
+            while (iterObject.GetNext())
             {
-                ref var key = ref iterObject.GetKey();
-                ref var value = ref iterObject.GetValue();
+                var key = iterObject.Key;
                 var s = HashSlotUtils.HashSlot(key);
                 if (slots.Contains(s))
-                    _ = BasicGarnetApi.DELETE(key, StoreType.Object);
+                    _ = BasicGarnetApi.DELETE(PinnedSpanByte.FromPinnedSpan(key), StoreType.Object);
             }
         }
     }

libs/cluster/Server/ClusterProvider.cs

@@ -15,12 +15,12 @@
 
 namespace Garnet.cluster
 {
-    using BasicGarnetApi = GarnetApi<BasicContext<SpanByte, SpanByte, RawStringInput, SpanByteAndMemory, long, MainSessionFunctions,
-            /* MainStoreFunctions */ StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>,
-            SpanByteAllocator<StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>>>,
-        BasicContext<byte[], IGarnetObject, ObjectInput, GarnetObjectStoreOutput, long, ObjectSessionFunctions,
-            /* ObjectStoreFunctions */ StoreFunctions<byte[], IGarnetObject, ByteArrayKeyComparer, DefaultRecordDisposer<byte[], IGarnetObject>>,
-            GenericAllocator<byte[], IGarnetObject, StoreFunctions<byte[], IGarnetObject, ByteArrayKeyComparer, DefaultRecordDisposer<byte[], IGarnetObject>>>>>;
+    using BasicGarnetApi = GarnetApi<BasicContext<RawStringInput, SpanByteAndMemory, long, MainSessionFunctions,
+            /* MainStoreFunctions */ StoreFunctions<SpanByteComparer, SpanByteRecordDisposer>,
+            SpanByteAllocator<StoreFunctions<SpanByteComparer, SpanByteRecordDisposer>>>,
+        BasicContext<ObjectInput, GarnetObjectStoreOutput, long, ObjectSessionFunctions,
+            /* ObjectStoreFunctions */ StoreFunctions<SpanByteComparer, DefaultRecordDisposer>,
+            ObjectAllocator<StoreFunctions<SpanByteComparer, DefaultRecordDisposer>>>>;
 
     /// <summary>
     /// Cluster provider

libs/cluster/Server/Migration/MigrateSessionKeyManager.cs

@@ -4,7 +4,7 @@
 using System;
 using System.Runtime.CompilerServices;
 using Garnet.common;
-using Garnet.server;
+using Tsavorite.core;
 
 namespace Garnet.cluster
 {
@@ -59,14 +59,14 @@ private void TryTransitionState(KeyMigrationStatus status)
         /// <param name="readOnly"></param>
         /// <returns></returns>
         /// <exception cref="GarnetException"></exception>
-        public bool CanAccessKey(ref ArgSlice key, int slot, bool readOnly)
+        public bool CanAccessKey(ref PinnedSpanByte key, int slot, bool readOnly)
         {
             // Skip operation check since this session is not responsible for migrating the associated slot
             if (!_sslots.Contains(slot))
                 return true;
 
             // If key is not queued for migration then
-            if (!_keys.TryGetValue(ref key, out var state))
+            if (!_keys.TryGetValue(key, out var state))
                 return true;
 
             // NOTE:
@@ -86,8 +86,8 @@ public bool CanAccessKey(ref ArgSlice key, int slot, bool readOnly)
         /// </summary>
         /// <param name="key"></param>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public bool AddKey(ref ArgSlice key)
-            => _keys.TryAdd(ref key, KeyMigrationStatus.QUEUED);
+        public bool AddKey(PinnedSpanByte key)
+            => _keys.TryAdd(key, KeyMigrationStatus.QUEUED);
 
         /// <summary>
         /// Clear keys from dictionary