Improvements in command handling #5

libs/cluster/Server/ClusterManagerSlotState.cs

@@ -12,7 +12,7 @@
 
 namespace Garnet.cluster
 {
-    using BasicGarnetApi = GarnetApi<BasicContext<SpanByte, SpanByte, SpanByte, SpanByteAndMemory, long, MainSessionFunctions,
+    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,

libs/cluster/Server/ClusterProvider.cs

@@ -15,7 +15,7 @@
 
 namespace Garnet.cluster
 {
-    using BasicGarnetApi = GarnetApi<BasicContext<SpanByte, SpanByte, SpanByte, SpanByteAndMemory, long, MainSessionFunctions,
+    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,

libs/cluster/Server/Migration/MigrateSessionKeys.cs

@@ -33,20 +33,11 @@ private bool MigrateKeysFromMainStore()
                 TryTransitionState(KeyMigrationStatus.MIGRATING);
                 WaitForConfigPropagation();
 
-                // 4 byte length of input
-                // 1 byte RespCommand
-                // 1 byte RespInputFlags
-                var inputSize = sizeof(int) + RespInputHeader.Size;
-                var pbCmdInput = stackalloc byte[inputSize];
-
                 ////////////////
                 // Build Input//
                 ////////////////
-                var pcurr = pbCmdInput;
-                *(int*)pcurr = inputSize - sizeof(int);
-                pcurr += sizeof(int);
-                // 1. Header
-                ((RespInputHeader*)pcurr)->SetHeader(RespCommandAccessor.MIGRATE, 0);
+                var input = new RawStringInput();
+                input.header.SetHeader(RespCommandAccessor.MIGRATE, 0);
 
                 foreach (var pair in _keys.GetKeys())
                 {
@@ -57,7 +48,7 @@ private bool MigrateKeysFromMainStore()
                     var key = pair.Key.SpanByte;
 
                     // Read value for key
-                    var status = localServerSession.BasicGarnetApi.Read_MainStore(ref key, ref Unsafe.AsRef<SpanByte>(pbCmdInput), ref o);
+                    var status = localServerSession.BasicGarnetApi.Read_MainStore(ref key, ref input, ref o);
 
                     // Check if found in main store
                     if (status == GarnetStatus.NOTFOUND)

libs/cluster/Session/ClusterSession.cs

@@ -12,7 +12,7 @@
 
 namespace Garnet.cluster
 {
-    using BasicGarnetApi = GarnetApi<BasicContext<SpanByte, SpanByte, SpanByte, SpanByteAndMemory, long, MainSessionFunctions,
+    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,

libs/common/RespReadUtils.cs

@@ -101,28 +101,71 @@ private static bool TryReadUlong(ref byte* ptr, byte* end, out ulong value, out
 
         /// <summary>
         /// Tries to read a signed 64-bit integer from a given ASCII-encoded input stream.
+        /// This method will throw if an overflow occurred.
         /// </summary>
         /// <param name="ptr">Pointer to the beginning of the ASCII encoded input string.</param>
         /// <param name="end">The end of the string to parse.</param>
         /// <param name="value">If parsing was successful, contains the parsed long value.</param>
         /// <param name="bytesRead">If parsing was successful, contains the number of bytes that were parsed.</param>
+        /// <param name="allowLeadingZeros">True if leading zeros allowed</param>
         /// <returns>
         /// True if a long was successfully parsed, false if the input string did not start with
         /// a valid integer or the end of the string was reached before finishing parsing.
         /// </returns>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public static bool TryReadLong(ref byte* ptr, byte* end, out long value, out ulong bytesRead)
+        public static bool TryReadLong(ref byte* ptr, byte* end, out long value, out ulong bytesRead, bool allowLeadingZeros = true)
+        {
+            var parseSuccessful = TryReadLongSafe(ref ptr, end, out value, out bytesRead, out var signRead,
+                out var overflow, allowLeadingZeros);
+
+            if (parseSuccessful) return true;
+
+            if (overflow)
+            {
+                var digitsRead = signRead ? bytesRead - 1 : bytesRead;
+                RespParsingException.ThrowIntegerOverflow(ptr - digitsRead, (int)digitsRead);
+                return false;
+            }
+
+            return false;
+        }
+
+        /// <summary>
+        /// Tries to read a signed 64-bit integer from a given ASCII-encoded input stream.
+        /// </summary>
+        /// <param name="ptr">Pointer to the beginning of the ASCII encoded input string.</param>
+        /// <param name="end">The end of the string to parse.</param>
+        /// <param name="value">If parsing was successful, contains the parsed long value.</param>
+        /// <param name="bytesRead">If parsing was successful, contains the number of bytes that were parsed.</param>
+        /// <param name="signRead">True if +/- sign was read during parsing</param>
+        /// <param name="overflow">True if overflow occured during parsing</param>
+        /// <param name="allowLeadingZeros">True if leading zeros allowed</param>
+        /// <returns>
+        /// True if a long was successfully parsed, false if the input string did not start with
+        /// a valid integer or the end of the string was reached before finishing parsing.
+        /// </returns>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static bool TryReadLongSafe(ref byte* ptr, byte* end, out long value, out ulong bytesRead, out bool signRead, out bool overflow, bool allowLeadingZeros = true)
         {
             bytesRead = 0;
             value = 0;
+            overflow = false;
 
             // Parse optional leading sign
-            if (TryReadSign(ptr, out var negative))
+            signRead = TryReadSign(ptr, out var negative);
+            if (signRead)
             {
                 ptr++;
                 bytesRead = 1;
             }
 
+            if (!allowLeadingZeros)
+            {
+                // Do not allow leading zeros
+                if (end - ptr > 1 && *ptr == '0')
+                    return false;
+            }
+
             // Parse digits as ulong
             if (!TryReadUlong(ref ptr, end, out var number, out var digitsRead))
             {
@@ -134,7 +177,8 @@ public static bool TryReadLong(ref byte* ptr, byte* end, out long value, out ulo
             {
                 if (number > ((ulong)long.MaxValue) + 1)
                 {
-                    RespParsingException.ThrowIntegerOverflow(ptr - digitsRead, (int)digitsRead);
+                    overflow = true;
+                    return false;
                 }
 
                 value = -1 - (long)(number - 1);
@@ -143,7 +187,8 @@ public static bool TryReadLong(ref byte* ptr, byte* end, out long value, out ulo
             {
                 if (number > long.MaxValue)
                 {
-                    RespParsingException.ThrowIntegerOverflow(ptr - digitsRead, (int)digitsRead);
+                    overflow = true;
+                    return false;
                 }
                 value = (long)number;
             }
@@ -153,25 +198,61 @@ public static bool TryReadLong(ref byte* ptr, byte* end, out long value, out ulo
             return true;
         }
 
+        /// <summary>
+        /// Tries to read a signed 32-bit integer from a given ASCII-encoded input stream.
+        /// This method will throw if an overflow occurred.
+        /// </summary>
+        /// <param name="ptr">Pointer to the beginning of the ASCII encoded input string.</param>
+        /// <param name="end">The end of the string to parse.</param>
+        /// <param name="value">If parsing was successful, contains the parsed int value.</param>
+        /// <param name="bytesRead">If parsing was successful, contains the number of bytes that were parsed.</param>
+        /// <param name="allowLeadingZeros">True if leading zeros allowed</param>
+        /// <returns>
+        /// True if an int was successfully parsed, false if the input string did not start with
+        /// a valid integer or the end of the string was reached before finishing parsing.
+        /// </returns>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static bool TryReadInt(ref byte* ptr, byte* end, out int value, out ulong bytesRead, bool allowLeadingZeros = true)
+        {
+            var parseSuccessful = TryReadIntSafe(ref ptr, end, out value, out bytesRead, out var signRead,
+                out var overflow, allowLeadingZeros);
+
+            if (parseSuccessful) return true;
+
+            if (overflow)
+            {
+                var digitsRead = signRead ? bytesRead - 1 : bytesRead;
+                RespParsingException.ThrowIntegerOverflow(ptr - digitsRead, (int)digitsRead);
+                return false;
+            }
+
+            return false;
+        }
+
         /// <summary>
         /// Tries to read a signed 32-bit integer from a given ASCII-encoded input stream.
         /// </summary>
         /// <param name="ptr">Pointer to the beginning of the ASCII encoded input string.</param>
         /// <param name="end">The end of the string to parse.</param>
         /// <param name="value">If parsing was successful, contains the parsed int value.</param>
         /// <param name="bytesRead">If parsing was successful, contains the number of bytes that were parsed.</param>
+        /// <param name="signRead">True if +/- sign was read during parsing</param>
+        /// <param name="overflow">True if overflow occured during parsing</param>
+        /// <param name="allowLeadingZeros">True if leading zeros allowed</param>
         /// <returns>
         /// True if an int was successfully parsed, false if the input string did not start with
         /// a valid integer or the end of the string was reached before finishing parsing.
         /// </returns>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public static bool TryReadInt(ref byte* ptr, byte* end, out int value, out ulong bytesRead)
+        public static bool TryReadIntSafe(ref byte* ptr, byte* end, out int value, out ulong bytesRead, out bool signRead, out bool overflow, bool allowLeadingZeros = true)
         {
             bytesRead = 0;
             value = 0;
+            overflow = false;
 
             // Parse optional leading sign
-            if (TryReadSign(ptr, out var negative))
+            signRead = TryReadSign(ptr, out var negative);
+            if (signRead)
             {
                 ptr++;
                 bytesRead = 1;
@@ -188,7 +269,8 @@ public static bool TryReadInt(ref byte* ptr, byte* end, out int value, out ulong
             {
                 if (number > ((ulong)int.MaxValue) + 1)
                 {
-                    RespParsingException.ThrowIntegerOverflow(ptr - digitsRead, (int)digitsRead);
+                    overflow = true;
+                    return false;
                 }
 
                 value = (int)(0 - (long)number);
@@ -197,7 +279,8 @@ public static bool TryReadInt(ref byte* ptr, byte* end, out int value, out ulong
             {
                 if (number > int.MaxValue)
                 {
-                    RespParsingException.ThrowIntegerOverflow(ptr - digitsRead, (int)digitsRead);
+                    overflow = true;
+                    return false;
                 }
                 value = (int)number;
             }

libs/server/AOF/AofProcessor.cs

@@ -37,7 +37,7 @@ public sealed unsafe partial class AofProcessor
         /// <summary>
         /// Session for main store
         /// </summary>
-        readonly BasicContext<SpanByte, SpanByte, SpanByte, SpanByteAndMemory, long, MainSessionFunctions, MainStoreFunctions, MainStoreAllocator> basicContext;
+        readonly BasicContext<SpanByte, SpanByte, RawStringInput, SpanByteAndMemory, long, MainSessionFunctions, MainStoreFunctions, MainStoreAllocator> basicContext;
 
         /// <summary>
         /// Session for object store
@@ -259,40 +259,113 @@ private unsafe bool ReplayOp(byte* entryPtr)
                     ObjectStoreDelete(objectStoreBasicContext, entryPtr);
                     break;
                 case AofEntryType.StoredProcedure:
-                    ref var input = ref Unsafe.AsRef<SpanByte>(entryPtr + sizeof(AofHeader));
-                    respServerSession.RunTransactionProc(header.type, new ArgSlice(ref input), ref output);
+                    RunStoredProc(header.type, entryPtr);
                     break;
                 default:
                     throw new GarnetException($"Unknown AOF header operation type {header.opType}");
             }
             return true;
         }
 
-        static unsafe void StoreUpsert(BasicContext<SpanByte, SpanByte, SpanByte, SpanByteAndMemory, long, MainSessionFunctions, MainStoreFunctions, MainStoreAllocator> basicContext, byte* ptr)
+        unsafe void RunStoredProc(byte id, byte* ptr)
         {
-            ref var key = ref Unsafe.AsRef<SpanByte>(ptr + sizeof(AofHeader));
-            ref var input = ref Unsafe.AsRef<SpanByte>(ptr + sizeof(AofHeader) + key.TotalSize);
-            ref var value = ref Unsafe.AsRef<SpanByte>(ptr + sizeof(AofHeader) + key.TotalSize + input.TotalSize);
+            var curr = ptr;
+            var parseStateCount = *(int*)curr;
+            curr += sizeof(int);
+
+            var parseState = new SessionParseState();
+            if (parseStateCount > 0)
+            {
+                ArgSlice[] parseStateBuffer = default;
+                parseState.Initialize(ref parseStateBuffer, parseStateCount);
+
+                for (var i = 0; i < parseStateCount; i++)
+                {
+                    ref var sbArgument = ref Unsafe.AsRef<SpanByte>(curr);
+                    parseStateBuffer[i] = new ArgSlice(ref sbArgument);
+                    curr += sbArgument.TotalSize;
+                }
+            }
+
+            respServerSession.RunTransactionProc(id, ref parseState, ref output);
+        }
+
+        static unsafe void StoreUpsert(BasicContext<SpanByte, SpanByte, RawStringInput, SpanByteAndMemory, long, MainSessionFunctions, MainStoreFunctions, MainStoreAllocator> basicContext, byte* ptr)
+        {
+            var curr = ptr + sizeof(AofHeader);
+            ref var key = ref Unsafe.AsRef<SpanByte>(curr);
+            curr += key.TotalSize;
+
+            // Reconstructing RawStringInput
+
+            // input
+            ref var sbInput = ref Unsafe.AsRef<SpanByte>(curr);
+            ref var input = ref Unsafe.AsRef<RawStringInput>(sbInput.ToPointer());
+            curr += sbInput.TotalSize;
+
+            // Reconstructing parse state
+            var parseStateCount = input.parseState.Count;
+
+            if (parseStateCount > 0)
+            {
+                ArgSlice[] parseStateBuffer = default;
+                input.parseState.Initialize(ref parseStateBuffer, parseStateCount);
+
+                for (var i = 0; i < parseStateCount; i++)
+                {
+                    ref var sbArgument = ref Unsafe.AsRef<SpanByte>(curr);
+                    parseStateBuffer[i] = new ArgSlice(ref sbArgument);
+                    curr += sbArgument.TotalSize;
+                }
+            }
+
+            ref var value = ref Unsafe.AsRef<SpanByte>(curr);
 
             SpanByteAndMemory output = default;
             basicContext.Upsert(ref key, ref input, ref value, ref output);
             if (!output.IsSpanByte)
                 output.Memory.Dispose();
         }
 
-        static unsafe void StoreRMW(BasicContext<SpanByte, SpanByte, SpanByte, SpanByteAndMemory, long, MainSessionFunctions, MainStoreFunctions, MainStoreAllocator> basicContext, byte* ptr)
+        static unsafe void StoreRMW(BasicContext<SpanByte, SpanByte, RawStringInput, SpanByteAndMemory, long, MainSessionFunctions, MainStoreFunctions, MainStoreAllocator> basicContext, byte* ptr)
         {
-            byte* pbOutput = stackalloc byte[32];
-            ref var key = ref Unsafe.AsRef<SpanByte>(ptr + sizeof(AofHeader));
-            ref var input = ref Unsafe.AsRef<SpanByte>(ptr + sizeof(AofHeader) + key.TotalSize);
+            var curr = ptr + sizeof(AofHeader);
+            ref var key = ref Unsafe.AsRef<SpanByte>(curr);
+            curr += key.TotalSize;
+
+            // Reconstructing RawStringInput
+
+            // input
+            ref var sbInput = ref Unsafe.AsRef<SpanByte>(curr);
+            ref var input = ref Unsafe.AsRef<RawStringInput>(sbInput.ToPointer());
+            curr += sbInput.TotalSize;
+
+            // Reconstructing parse state
+            var parseStateCount = input.parseState.Count;
+
+            if (parseStateCount > 0)
+            {
+                ArgSlice[] parseStateBuffer = default;
+                input.parseState.Initialize(ref parseStateBuffer, parseStateCount);
+
+                for (var i = 0; i < parseStateCount; i++)
+                {
+                    ref var sbArgument = ref Unsafe.AsRef<SpanByte>(curr);
+                    parseStateBuffer[i] = new ArgSlice(ref sbArgument);
+                    curr += sbArgument.TotalSize;
+                }
+            }
+
+            var pbOutput = stackalloc byte[32];
             var output = new SpanByteAndMemory(pbOutput, 32);
+
             if (basicContext.RMW(ref key, ref input, ref output).IsPending)
                 basicContext.CompletePending(true);
             if (!output.IsSpanByte)
                 output.Memory.Dispose();
         }
 
-        static unsafe void StoreDelete(BasicContext<SpanByte, SpanByte, SpanByte, SpanByteAndMemory, long, MainSessionFunctions, MainStoreFunctions, MainStoreAllocator> basicContext, byte* ptr)
+        static unsafe void StoreDelete(BasicContext<SpanByte, SpanByte, RawStringInput, SpanByteAndMemory, long, MainSessionFunctions, MainStoreFunctions, MainStoreAllocator> basicContext, byte* ptr)
         {
             ref var key = ref Unsafe.AsRef<SpanByte>(ptr + sizeof(AofHeader));
             basicContext.Delete(ref key);