Refactor TransactionManager key tracking to use SessionParseState and ScratchBufferAllocator

libs/cluster/Session/SlotVerification/ClusterSlotVerify.cs

@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-using System;
 using System.Runtime.CompilerServices;
 using System.Threading;
 using Garnet.common;
@@ -118,39 +117,17 @@ bool CanOperateOnKey(ref PinnedSpanByte key, int slot, bool readOnly)
             }
         }
 
-        ClusterSlotVerificationResult MultiKeySlotVerify(ClusterConfig config, ref Span<PinnedSpanByte> keys, bool readOnly, byte sessionAsking, int count)
-        {
-            var _end = count < 0 ? keys.Length : count;
-            var slot = HashSlotUtils.HashSlot(keys[0]);
-            var verifyResult = SingleKeySlotVerify(ref config, ref keys[0], readOnly, sessionAsking, slot);
-
-            for (var i = 1; i < _end; i++)
-            {
-                var _slot = HashSlotUtils.HashSlot(keys[i]);
-                var _verifyResult = SingleKeySlotVerify(ref config, ref keys[i], readOnly, sessionAsking, _slot);
-
-                // Check if slot changes between keys
-                if (_slot != slot)
-                    return new(SlotVerifiedState.CROSSSLOT, slot);
-
-                // Check if state of key changes
-                if (_verifyResult.state != verifyResult.state)
-                    return new(SlotVerifiedState.TRYAGAIN, slot);
-            }
-
-            return verifyResult;
-        }
-
-        ClusterSlotVerificationResult MultiKeySlotVerify(ClusterConfig config, ref SessionParseState parseState, ref ClusterSlotVerificationInput csvi)
+        ClusterSlotVerificationResult MultiKeySlotVerify(ClusterConfig config, ref SessionParseState parseState, ref ClusterSlotVerificationInput csvi, bool isTxn)
         {
             // Find the first valid key and initialize slot/result
             var specIndex = 0;
-            (int firstIdx, int lastIdx, int step) searchArgs = default;
-            while (specIndex < csvi.keySpecs.Length &&
+            // If slot verification is called from transaction manager, parse state contains consecutive keys so we can skip key search
+            (int firstIdx, int lastIdx, int step) searchArgs = isTxn ? (0, parseState.Count - 1, 1) : default;
+            while (specIndex < csvi.keySpecs?.Length &&
                    !parseState.TryGetKeySearchArgsFromSimpleKeySpec(csvi.keySpecs[specIndex], csvi.isSubCommand, out searchArgs))
                 specIndex++;
 
-            if (specIndex == csvi.keySpecs.Length)
+            if (specIndex == csvi.keySpecs?.Length && !isTxn)
                 return default;
 
             ref var firstKey = ref parseState.GetArgSliceByRef(searchArgs.firstIdx);
@@ -164,7 +141,7 @@ ClusterSlotVerificationResult MultiKeySlotVerify(ClusterConfig config, ref Sessi
                 return verifyResult;
 
             // Verify keys from remaining specs
-            for (specIndex++; specIndex < csvi.keySpecs.Length; specIndex++)
+            for (specIndex++; specIndex < csvi.keySpecs?.Length; specIndex++)
             {
                 if (!parseState.TryGetKeySearchArgsFromSimpleKeySpec(csvi.keySpecs[specIndex], csvi.isSubCommand, out searchArgs))
                     continue;

libs/cluster/Session/SlotVerification/RespClusterSlotVerify.cs

@@ -7,7 +7,6 @@
 using Garnet.common;
 using Garnet.server;
 using Microsoft.Extensions.Logging;
-using Tsavorite.core;
 
 namespace Garnet.cluster
 {
@@ -87,32 +86,6 @@ private void WriteClusterSlotVerificationMessage(ClusterConfig config, ClusterSl
                 SendAndReset(ref dcurr, ref dend);
         }
 
-        /// <summary>
-        /// Check if read/write is permitted on an array of keys and generate appropriate resp response.
-        /// </summary>
-        /// <param name="keys"></param>
-        /// <param name="readOnly"></param>
-        /// <param name="sessionAsking"></param>
-        /// <param name="dcurr"></param>
-        /// <param name="dend"></param>
-        /// <param name="count"></param>
-        /// <returns></returns>
-        public bool NetworkKeyArraySlotVerify(Span<PinnedSpanByte> keys, bool readOnly, byte sessionAsking, ref byte* dcurr, ref byte* dend, int count = -1)
-        {
-            // If cluster is not enabled or a transaction is running skip slot check
-            if (!clusterProvider.serverOptions.EnableCluster || txnManager.state == TxnState.Running)
-                return false;
-
-            var config = clusterProvider.clusterManager.CurrentConfig;
-            var vres = MultiKeySlotVerify(config, ref keys, readOnly, sessionAsking, count);
-
-            if (vres.state == SlotVerifiedState.OK)
-                return false;
-            else
-                WriteClusterSlotVerificationMessage(config, vres, ref dcurr, ref dend);
-            return true;
-        }
-
         /// <summary>
         /// Verify multi-key slot ownership
         /// </summary>
@@ -121,13 +94,13 @@ public bool NetworkKeyArraySlotVerify(Span<PinnedSpanByte> keys, bool readOnly,
         /// <param name="dcurr"></param>
         /// <param name="dend"></param>
         /// <returns></returns>
-        public unsafe bool NetworkMultiKeySlotVerify(ref SessionParseState parseState, ref ClusterSlotVerificationInput csvi, ref byte* dcurr, ref byte* dend)
+        public unsafe bool NetworkMultiKeySlotVerify(ref SessionParseState parseState, ref ClusterSlotVerificationInput csvi, ref byte* dcurr, ref byte* dend, bool isTxn = false)
         {
             // If cluster is not enabled or a transaction is running skip slot check
             if (!clusterProvider.serverOptions.EnableCluster || txnManager.state == TxnState.Running) return false;
 
             var config = clusterProvider.clusterManager.CurrentConfig;
-            var vres = MultiKeySlotVerify(config, ref parseState, ref csvi);
+            var vres = MultiKeySlotVerify(config, ref parseState, ref csvi, isTxn);
 
             if (vres.state == SlotVerifiedState.OK)
                 return false;
@@ -143,14 +116,15 @@ public unsafe bool NetworkMultiKeySlotVerify(ref SessionParseState parseState, r
         /// <param name="csvi"></param>
         /// <param name="dcurr"></param>
         /// <param name="dend"></param>
+        /// <param name="isTxn"></param>
         /// <returns></returns>
-        public unsafe bool NetworkMultiKeySlotVerifyNoResponse(ref SessionParseState parseState, ref ClusterSlotVerificationInput csvi, ref byte* dcurr, ref byte* dend)
+        public unsafe bool NetworkMultiKeySlotVerifyNoResponse(ref SessionParseState parseState, ref ClusterSlotVerificationInput csvi, ref byte* dcurr, ref byte* dend, bool isTxn = false)
         {
             // If cluster is not enabled or a transaction is running skip slot check
             if (!clusterProvider.serverOptions.EnableCluster || txnManager.state == TxnState.Running) return false;
 
             var config = clusterProvider.clusterManager.CurrentConfig;
-            var vres = MultiKeySlotVerify(config, ref parseState, ref csvi);
+            var vres = MultiKeySlotVerify(config, ref parseState, ref csvi, isTxn);
 
             return vres.state != SlotVerifiedState.OK;
         }

libs/server/Cluster/IClusterSession.cs

@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-using System;
 using Garnet.common;
 using Garnet.server.ACL;
 using Tsavorite.core;
@@ -86,20 +85,16 @@ public interface IClusterSession
         /// <param name="output"></param>
         public void WriteCachedSlotVerificationMessage(ref MemoryResult<byte> output);
 
-        /// <summary>
-        /// Key array slot verify (write result to network)
-        /// </summary>
-        unsafe bool NetworkKeyArraySlotVerify(Span<PinnedSpanByte> keys, bool readOnly, byte SessionAsking, ref byte* dcurr, ref byte* dend, int count = -1);
-
         /// <summary>
         /// Array slot verify (write result to network)
         /// </summary>
         /// <param name="parseState"></param>
         /// <param name="csvi"></param>
         /// <param name="dcurr"></param>
         /// <param name="dend"></param>
+        /// <param name="isTxn"></param>
         /// <returns></returns>
-        unsafe bool NetworkMultiKeySlotVerify(ref SessionParseState parseState, ref ClusterSlotVerificationInput csvi, ref byte* dcurr, ref byte* dend);
+        unsafe bool NetworkMultiKeySlotVerify(ref SessionParseState parseState, ref ClusterSlotVerificationInput csvi, ref byte* dcurr, ref byte* dend, bool isTxn = false);
 
         /// <summary>
         /// Array slot verify with no response
@@ -108,8 +103,9 @@ public interface IClusterSession
         /// <param name="csvi"></param>
         /// <param name="dcurr"></param>
         /// <param name="dend"></param>
+        /// <param name="isTxn"></param>
         /// <returns></returns>
-        unsafe bool NetworkMultiKeySlotVerifyNoResponse(ref SessionParseState parseState, ref ClusterSlotVerificationInput csvi, ref byte* dcurr, ref byte* dend);
+        unsafe bool NetworkMultiKeySlotVerifyNoResponse(ref SessionParseState parseState, ref ClusterSlotVerificationInput csvi, ref byte* dcurr, ref byte* dend, bool isTxn = false);
 
         /// <summary>
         /// Sets the <see cref="UserHandle"/> currently authenticated in this session (used for permission checks)

libs/server/Resp/Parser/SessionParseState.cs

@@ -26,6 +26,16 @@ public unsafe struct SessionParseState
         /// </summary>
         public int Count;
 
+        /// <summary>
+        /// Get the allocated capacity of the argument buffer
+        /// </summary>
+        public int Capacity { get; }
+
+        /// <summary>
+        /// Get a Span of the parsed parameters in the form an PinnedSpanByte
+        /// </summary>
+        public ReadOnlySpan<PinnedSpanByte> Parameters => new(bufferPtr, Count);
+
         /// <summary>
         /// Pointer to the slice of <see cref="rootBuffer"/> (which is always pinned) that is accessible within the range of this instance's arguments.
         /// </summary>
@@ -41,17 +51,13 @@ public unsafe struct SessionParseState
         /// </summary>
         PinnedSpanByte[] rootBuffer;
 
-        /// <summary>
-        /// Get a Span of the parsed parameters in the form an PinnedSpanByte
-        /// </summary>
-        public ReadOnlySpan<PinnedSpanByte> Parameters => new(bufferPtr, Count);
-
         private SessionParseState(ref PinnedSpanByte[] rootBuffer, int rootCount, ref PinnedSpanByte* bufferPtr, int count) : this()
         {
             this.rootBuffer = rootBuffer;
             this.rootCount = rootCount;
             this.bufferPtr = bufferPtr;
             this.Count = count;
+            this.Capacity = rootBuffer.Length;
         }
 
         /// <summary>

libs/server/Resp/RespServerSessionSlotVerify.cs

@@ -1,9 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-using System;
 using System.Diagnostics;
-using Tsavorite.core;
 
 namespace Garnet.server
 {
@@ -12,16 +10,6 @@ namespace Garnet.server
     /// </summary>
     internal sealed unsafe partial class RespServerSession : ServerSessionBase
     {
-        /// <summary>
-        /// This method is used to verify slot ownership for provided array of key argslices.
-        /// </summary>
-        /// <param name="keys">Array of key ArgSlice</param>
-        /// <param name="readOnly">Whether caller is going to perform a readonly or read/write operation</param>
-        /// <param name="count">Key count if different than keys array length</param>
-        /// <returns>True when ownership is verified, false otherwise</returns>
-        bool NetworkKeyArraySlotVerify(Span<PinnedSpanByte> keys, bool readOnly, int count = -1)
-            => clusterSession != null && clusterSession.NetworkKeyArraySlotVerify(keys, readOnly, SessionAsking, ref dcurr, ref dend, count);
-
         /// <summary>
         /// Validate if this command can be served based on the current slot assignment
         /// </summary>

libs/server/Transaction/TransactionManager.cs

@@ -68,6 +68,8 @@ public sealed unsafe partial class TransactionManager
         private readonly RespServerSession respSession;
         readonly FunctionsState functionsState;
         internal readonly ScratchBufferAllocator scratchBufferAllocator;
+        internal readonly ScratchBufferAllocator txnScratchBufferAllocator;
+        internal SessionParseState txnKeysParseState;
         private readonly TsavoriteLog appendOnlyFile;
         internal readonly WatchedKeysContainer watchContainer;
         private readonly StateMachineDriver stateMachineDriver;
@@ -135,7 +137,8 @@ internal TransactionManager(
 
             this.respSession = respSession;
 
-            watchContainer = new WatchedKeysContainer(initialSliceBufferSize, functionsState.watchVersionMap);
+            txnScratchBufferAllocator = new ScratchBufferAllocator();
+            watchContainer = new WatchedKeysContainer(initialSliceBufferSize, functionsState.watchVersionMap, txnScratchBufferAllocator);
             keyEntries = new TxnKeyEntries(initialSliceBufferSize, unifiedTransactionalContext);
             this.scratchBufferAllocator = scratchBufferAllocator;
 
@@ -149,7 +152,10 @@ internal TransactionManager(
 
             this.clusterEnabled = clusterEnabled;
             if (clusterEnabled)
-                keys = new PinnedSpanByte[initialKeyBufferSize];
+            {
+                txnKeysParseState.Initialize(initialKeyBufferSize);
+                txnKeysParseState.Count = 0;
+            }
 
             Reset(false);
         }
@@ -184,9 +190,13 @@ internal void Reset(bool isRunning)
             functionsState.StoredProcMode = false;
             this.PerformWrites = false;
 
-            // Reset cluster variables used for slot verification
-            this.saveKeyRecvBufferPtr = null;
-            this.keyCount = 0;
+            // Reset cluster key parse state
+            if (clusterEnabled)
+            {
+                txnKeysParseState.Count = 0;
+                saveKeyRecvBufferPtr = null;
+                txnScratchBufferAllocator.Reset();
+            }
         }
 
         internal bool RunTransactionProc(byte id, ref CustomProcedureInput procInput, CustomTransactionProcedure proc, ref MemoryResult<byte> output, bool isReplaying = false)
@@ -257,7 +267,6 @@ internal bool RunTransactionProc(byte id, ref CustomProcedureInput procInput, Cu
                 scratchBufferAllocator.Reset();
             }
 
-
             return true;
         }
 
@@ -328,13 +337,22 @@ internal void AddTransactionStoreType(StoreType storeType)
 
         internal string GetLockset() => keyEntries.GetLockset();
 
-        internal void GetKeysForValidation(byte* recvBufferPtr, out PinnedSpanByte[] keys, out int keyCount, out bool readOnly)
+        internal void GetSlotVerificationInput(byte* recvBufferPtr, byte sessionAsking, out ClusterSlotVerificationInput clusterSlotVerificationInput)
         {
-            UpdateRecvBufferPtr(recvBufferPtr);
+            // Copy keys if buffer changed since last queued command
+            if (recvBufferPtr != saveKeyRecvBufferPtr)
+            {
+                CopyExistingKeysToScratchBuffer();
+                saveKeyRecvBufferPtr = recvBufferPtr;
+            }
+
             watchContainer.SaveKeysToKeyList(this);
-            keys = this.keys;
-            keyCount = this.keyCount;
-            readOnly = keyEntries.IsReadOnly;
+            clusterSlotVerificationInput = new ClusterSlotVerificationInput
+            {
+                readOnly = keyEntries.IsReadOnly,
+                sessionAsking = sessionAsking,
+                // We don't specify key specs here as slot verification will know to iterate over all keys in this context
+            };
         }
 
         void BeginTransaction()

libs/server/Transaction/TxnClusterSlotCheck.cs

@@ -1,51 +1,51 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-using System;
+using System.Diagnostics;
 using Tsavorite.core;
 
 namespace Garnet.server
 {
     sealed unsafe partial class TransactionManager
     {
-        // Keys involved in the current transaction
-        PinnedSpanByte[] keys;
-        int keyCount;
-
-        internal byte* saveKeyRecvBufferPtr;
         readonly bool clusterEnabled;
+        internal byte* saveKeyRecvBufferPtr;
 
         /// <summary>
         /// Keep track of actual key accessed by command
         /// </summary>
-        /// <param name="argSlice"></param>
-        public void SaveKeyArgSlice(PinnedSpanByte argSlice)
+        /// <param name="keySlice"></param>
+        public void SaveKeyArgSlice(PinnedSpanByte keySlice)
         {
             // Execute method only if clusterEnabled
             if (!clusterEnabled) return;
-            // Grow the buffer if needed
-            if (keyCount >= keys.Length)
+
+            var count = txnKeysParseState.Count;
+
+            // Double the parse state buffer capacity if needed, and copy existing parameters to the extended buffer
+            if (count >= txnKeysParseState.Capacity)
             {
-                var oldKeys = keys;
-                keys = new PinnedSpanByte[keys.Length * 2];
-                Array.Copy(oldKeys, keys, oldKeys.Length);
+                var oldParams = txnKeysParseState.Parameters;
+                txnKeysParseState.Initialize(count * 2);
+                txnKeysParseState.SetArguments(0, oldParams);
             }
 
-            keys[keyCount++] = argSlice;
+            txnKeysParseState.Count = count + 1;
+            txnKeysParseState.SetArgument(count, keySlice);
         }
 
         /// <summary>
-        /// Update argslice ptr if input buffer has been resized
+        /// Copy all existing keys into <see cref="txnScratchBufferAllocator"/> so they are independent of the old receive buffer.
+        /// Called when the receive buffer has been reallocated since keys were last stored.
         /// </summary>
-        /// <param name="recvBufferPtr"></param>
-        public unsafe void UpdateRecvBufferPtr(byte* recvBufferPtr)
+        public void CopyExistingKeysToScratchBuffer()
         {
-            // Execute method only if clusterEnabled
-            if (!clusterEnabled) return;
-            if (recvBufferPtr != saveKeyRecvBufferPtr)
+            Debug.Assert(clusterEnabled);
+
+            for (var i = 0; i < txnKeysParseState.Count; i++)
             {
-                for (int i = 0; i < keyCount; i++)
-                    keys[i].ptr = recvBufferPtr + (keys[i].ptr - saveKeyRecvBufferPtr);
+                ref var key = ref txnKeysParseState.GetArgSliceByRef(i);
+                key = txnScratchBufferAllocator.CreateArgSlice(key.ReadOnlySpan);
             }
         }
     }