From e307c3eb7e96a389a5a9640d1bf061ae74723d29 Mon Sep 17 00:00:00 2001 From: prvyk Date: Wed, 13 Aug 2025 13:39:44 +0300 Subject: [PATCH 1/5] Make ListReadKeysWithCount use idiomatic code instead of needless readHead and direct input parsing. --- libs/server/Resp/RespServerSession.cs | 37 ------------------------ libs/server/Transaction/TxnKeyManager.cs | 19 +++++++----- 2 files changed, 11 insertions(+), 45 deletions(-) diff --git a/libs/server/Resp/RespServerSession.cs b/libs/server/Resp/RespServerSession.cs index 6217bf9cbe0..9615b44a66c 100644 --- a/libs/server/Resp/RespServerSession.cs +++ b/libs/server/Resp/RespServerSession.cs @@ -1188,43 +1188,6 @@ ReadOnlySpan GetUpperCaseCommand(out bool success) return result; } - public ArgSlice GetCommandAsArgSlice(out bool success) - { - if (bytesRead - readHead < 6) - { - success = false; - return default; - } - - Debug.Assert(*(recvBufferPtr + readHead) == '$'); - int psize = *(recvBufferPtr + readHead + 1) - '0'; - readHead += 2; - while (*(recvBufferPtr + readHead) != '\r') - { - psize = psize * 10 + *(recvBufferPtr + readHead) - '0'; - if (bytesRead - readHead < 1) - { - success = false; - return default; - } - readHead++; - } - if (bytesRead - readHead < 2 + psize + 2) - { - success = false; - return default; - } - Debug.Assert(*(recvBufferPtr + readHead + 1) == '\n'); - - var result = new ArgSlice(recvBufferPtr + readHead + 2, psize); - Debug.Assert(*(recvBufferPtr + readHead + 2 + psize) == '\r'); - Debug.Assert(*(recvBufferPtr + readHead + 2 + psize + 1) == '\n'); - - readHead += 2 + psize + 2; - success = true; - return result; - } - /// /// Attempt to kill this session. /// diff --git a/libs/server/Transaction/TxnKeyManager.cs b/libs/server/Transaction/TxnKeyManager.cs index 07707909f0c..d43ce7f2b57 100644 --- a/libs/server/Transaction/TxnKeyManager.cs +++ b/libs/server/Transaction/TxnKeyManager.cs @@ -378,7 +378,7 @@ private int ListObjectKeys(RespCommand command) RespCommand.LINSERT => SingleKey(StoreType.Object, LockType.Exclusive), RespCommand.LLEN => SingleKey(StoreType.Object, LockType.Shared), RespCommand.LMOVE => ListKeys(2, StoreType.Object, LockType.Exclusive), - RespCommand.LMPOP => ListReadKeysWithCount(LockType.Exclusive), + RespCommand.LMPOP => ListReadKeysWithCount(LockType.Exclusive, mandatoryArgs: 1), RespCommand.LPOP => SingleKey(StoreType.Object, LockType.Exclusive), RespCommand.LPOS => SingleKey(StoreType.Object, LockType.Shared), RespCommand.LPUSH => SingleKey(StoreType.Object, LockType.Exclusive), @@ -487,17 +487,20 @@ private int ListKeys(int inputCount, StoreType storeType, LockType type) /// /// Returns a list of keys for LMPOP command /// - private int ListReadKeysWithCount(LockType type, bool isObject = true) + private int ListReadKeysWithCount(LockType type, bool isObject = true, int mandatoryArgs = 0) { - var numKeysArg = respSession.GetCommandAsArgSlice(out var success); - if (!success) return -2; + if (respSession.parseState.Count == 0) + return -2; - if (!NumUtils.TryParse(numKeysArg.ReadOnlySpan, out int numKeys)) return -2; + if (!respSession.parseState.TryGetInt(0, out var numKeys)) + return -2; - for (var i = 0; i < numKeys; i++) + if (numKeys + mandatoryArgs + 1 > respSession.parseState.Count) + return -2; + + for (var i = 1; i < numKeys + 1; i++) { - var key = respSession.GetCommandAsArgSlice(out success); - if (!success) return -2; + var key = respSession.parseState.GetArgSliceByRef(i); SaveKeyEntryToLock(key, isObject, type); SaveKeyArgSlice(key); } From 0387a36cb3b730cffb4ad50987f534ef676d18dd Mon Sep 17 00:00:00 2001 From: prvyk Date: Wed, 13 Aug 2025 13:41:25 +0300 Subject: [PATCH 2/5] MakeUpperPtr should use length from ArgSlice instead of mistaken bytesRead calculation. --- libs/server/Resp/Objects/ObjectStoreUtils.cs | 2 +- libs/server/Resp/Parser/RespCommand.cs | 2 +- libs/server/Resp/RespServerSession.cs | 3 +-- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/libs/server/Resp/Objects/ObjectStoreUtils.cs b/libs/server/Resp/Objects/ObjectStoreUtils.cs index bf265a7d0b3..ad70b6afa0d 100644 --- a/libs/server/Resp/Objects/ObjectStoreUtils.cs +++ b/libs/server/Resp/Objects/ObjectStoreUtils.cs @@ -114,7 +114,7 @@ public OperationDirection GetOperationDirection(ArgSlice input) if (input.ReadOnlySpan.SequenceEqual("RIGHT"u8)) return OperationDirection.Right; // Rare case: try making upper case and retry - MakeUpperCase(input.ptr); + MakeUpperCase(input.ptr, input.length); if (input.ReadOnlySpan.SequenceEqual("LEFT"u8)) return OperationDirection.Left; if (input.ReadOnlySpan.SequenceEqual("RIGHT"u8)) diff --git a/libs/server/Resp/Parser/RespCommand.cs b/libs/server/Resp/Parser/RespCommand.cs index 1d81898348d..9d4224d56c8 100644 --- a/libs/server/Resp/Parser/RespCommand.cs +++ b/libs/server/Resp/Parser/RespCommand.cs @@ -2768,7 +2768,7 @@ private RespCommand ArrayParseCommand(bool writeErrorOnFailure, ref int count, r var ptr = recvBufferPtr + readHead; // See if input command is all upper-case. If not, convert and try fast parse pass again. - if (MakeUpperCase(ptr)) + if (MakeUpperCase(ptr, bytesRead - readHead)) { cmd = FastParseCommand(out count); if (cmd != RespCommand.NONE) diff --git a/libs/server/Resp/RespServerSession.cs b/libs/server/Resp/RespServerSession.cs index 9615b44a66c..aeb3e966c36 100644 --- a/libs/server/Resp/RespServerSession.cs +++ b/libs/server/Resp/RespServerSession.cs @@ -657,7 +657,7 @@ private void ProcessMessages() } // Make first command in string as uppercase - private bool MakeUpperCase(byte* ptr) + private bool MakeUpperCase(byte* ptr, int len) { // Assume most commands are already upper case. // Assume most commands are 2-8 bytes long. @@ -672,7 +672,6 @@ private bool MakeUpperCase(byte* ptr) // Note that _all_ of these bytes are <= 95 in the common case // and there's no need to scan the whole string in those cases. - var len = bytesRead - readHead; if (len >= 12) { var cmdLen = (uint)(*(ptr + 5) - '2'); From b4d774676db5ce6c76db8082709fa7c074a43997 Mon Sep 17 00:00:00 2001 From: prvyk Date: Thu, 14 Aug 2025 11:28:39 +0300 Subject: [PATCH 3/5] Remove GetOperationDirection entirely and use sessionparsestate extension. --- libs/server/Resp/Objects/ListCommands.cs | 42 +++++++++----------- libs/server/Resp/Objects/ObjectStoreUtils.cs | 22 ---------- libs/server/SessionParseStateExtensions.cs | 15 +++++++ libs/server/Transaction/TxnKeyManager.cs | 1 + 4 files changed, 34 insertions(+), 46 deletions(-) diff --git a/libs/server/Resp/Objects/ListCommands.cs b/libs/server/Resp/Objects/ListCommands.cs index 69a6df8e0b9..e0f25d8d4e6 100644 --- a/libs/server/Resp/Objects/ListCommands.cs +++ b/libs/server/Resp/Objects/ListCommands.cs @@ -226,10 +226,7 @@ private unsafe bool ListPopMultiple(ref TGarnetApi storageApi) } // Get the direction - var dir = parseState.GetArgSliceByRef(currTokenId++); - var popDirection = GetOperationDirection(dir); - - if (popDirection == OperationDirection.Unknown) + if (!parseState.TryGetOperationDirection(currTokenId++, out var popDirection)) { return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR); } @@ -355,8 +352,13 @@ private unsafe bool ListBlockingMove() var srcKey = parseState.GetArgSliceByRef(0); var dstKey = parseState.GetArgSliceByRef(1); - var srcDir = parseState.GetArgSliceByRef(2); - var dstDir = parseState.GetArgSliceByRef(3); + + if (!parseState.TryGetOperationDirection(2, out var srcDir) || + !parseState.TryGetOperationDirection(3, out var dstDir) + ) + { + return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR); + } if (!parseState.TryGetTimeout(4, out var timeout, out var error)) { @@ -379,27 +381,26 @@ private bool ListBlockingPopPush() var srcKey = parseState.GetArgSliceByRef(0); var dstKey = parseState.GetArgSliceByRef(1); - var rightOption = ArgSlice.FromPinnedSpan(CmdStrings.RIGHT); - var leftOption = ArgSlice.FromPinnedSpan(CmdStrings.LEFT); if (!parseState.TryGetTimeout(2, out var timeout, out var error)) { return AbortWithErrorMessage(error); } - return ListBlockingMove(srcKey, dstKey, rightOption, leftOption, timeout); + return ListBlockingMove(srcKey, dstKey, OperationDirection.Right, + OperationDirection.Left, timeout); } - private bool ListBlockingMove(ArgSlice srcKey, ArgSlice dstKey, ArgSlice srcDir, ArgSlice dstDir, double timeout) + private bool ListBlockingMove(ArgSlice srcKey, ArgSlice dstKey, + OperationDirection sourceDirection, + OperationDirection destinationDirection, + double timeout) { var cmdArgs = new ArgSlice[] { default, default, default }; // Read destination key cmdArgs[0] = dstKey; - var sourceDirection = GetOperationDirection(srcDir); - var destinationDirection = GetOperationDirection(dstDir); - if (sourceDirection == OperationDirection.Unknown || destinationDirection == OperationDirection.Unknown) { return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR); @@ -770,13 +771,9 @@ private bool ListMove(ref TGarnetApi storageApi) var srcKey = parseState.GetArgSliceByRef(0); var dstKey = parseState.GetArgSliceByRef(1); - var srcDirSlice = parseState.GetArgSliceByRef(2); - var dstDirSlice = parseState.GetArgSliceByRef(3); - - var sourceDirection = GetOperationDirection(srcDirSlice); - var destinationDirection = GetOperationDirection(dstDirSlice); - - if (sourceDirection == OperationDirection.Unknown || destinationDirection == OperationDirection.Unknown) + if (!parseState.TryGetOperationDirection(2, out var sourceDirection) || + !parseState.TryGetOperationDirection(3, out var destinationDirection) + ) { return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR); } @@ -964,10 +961,7 @@ private unsafe bool ListBlockingPopMultiple() var cmdArgs = new ArgSlice[2]; // Get the direction - var dir = parseState.GetArgSliceByRef(currTokenId++); - var popDirection = GetOperationDirection(dir); - - if (popDirection == OperationDirection.Unknown) + if (!parseState.TryGetOperationDirection(currTokenId++, out var popDirection)) { return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR); } diff --git a/libs/server/Resp/Objects/ObjectStoreUtils.cs b/libs/server/Resp/Objects/ObjectStoreUtils.cs index ad70b6afa0d..3704a633c16 100644 --- a/libs/server/Resp/Objects/ObjectStoreUtils.cs +++ b/libs/server/Resp/Objects/ObjectStoreUtils.cs @@ -99,27 +99,5 @@ private bool AbortWithErrorMessage(string format, params object[] args) { return AbortWithErrorMessage(Encoding.ASCII.GetBytes(string.Format(format, args))); } - - /// - /// Tries to parse the input as "LEFT" or "RIGHT" and returns the corresponding OperationDirection. - /// If parsing fails, returns OperationDirection.Unknown. - /// - /// The input to parse. - /// The parsed OperationDirection, or OperationDirection.Unknown if parsing fails. - public OperationDirection GetOperationDirection(ArgSlice input) - { - // Optimize for the common case - if (input.ReadOnlySpan.SequenceEqual("LEFT"u8)) - return OperationDirection.Left; - if (input.ReadOnlySpan.SequenceEqual("RIGHT"u8)) - return OperationDirection.Right; - // Rare case: try making upper case and retry - MakeUpperCase(input.ptr, input.length); - if (input.ReadOnlySpan.SequenceEqual("LEFT"u8)) - return OperationDirection.Left; - if (input.ReadOnlySpan.SequenceEqual("RIGHT"u8)) - return OperationDirection.Right; - return OperationDirection.Unknown; - } } } \ No newline at end of file diff --git a/libs/server/SessionParseStateExtensions.cs b/libs/server/SessionParseStateExtensions.cs index 4ab87a3c92c..9a5f6ef8520 100644 --- a/libs/server/SessionParseStateExtensions.cs +++ b/libs/server/SessionParseStateExtensions.cs @@ -614,6 +614,21 @@ internal static bool TryGetManagerType(this SessionParseState parseState, int id return true; } + internal static bool TryGetOperationDirection(this SessionParseState parseState, int idx, out OperationDirection value) + { + value = OperationDirection.Unknown; + var sbArg = parseState.GetArgSliceByRef(idx).ReadOnlySpan; + + if (sbArg.EqualsUpperCaseSpanIgnoringCase("LEFT"u8)) + value = OperationDirection.Left; + else if (sbArg.EqualsUpperCaseSpanIgnoringCase("RIGHT"u8)) + value = OperationDirection.Right; + else + return false; + + return true; + } + /// /// Parse sorted set add option from parse state at specified index /// diff --git a/libs/server/Transaction/TxnKeyManager.cs b/libs/server/Transaction/TxnKeyManager.cs index d43ce7f2b57..24a5d98438b 100644 --- a/libs/server/Transaction/TxnKeyManager.cs +++ b/libs/server/Transaction/TxnKeyManager.cs @@ -378,6 +378,7 @@ private int ListObjectKeys(RespCommand command) RespCommand.LINSERT => SingleKey(StoreType.Object, LockType.Exclusive), RespCommand.LLEN => SingleKey(StoreType.Object, LockType.Shared), RespCommand.LMOVE => ListKeys(2, StoreType.Object, LockType.Exclusive), + // Account for mandatory LEFT/RIGHT argument RespCommand.LMPOP => ListReadKeysWithCount(LockType.Exclusive, mandatoryArgs: 1), RespCommand.LPOP => SingleKey(StoreType.Object, LockType.Exclusive), RespCommand.LPOS => SingleKey(StoreType.Object, LockType.Shared), From fee2a973ba4203f5e54026c98a4473b07764ed6e Mon Sep 17 00:00:00 2001 From: Tal Zaccai Date: Thu, 14 Aug 2025 17:15:03 -0700 Subject: [PATCH 4/5] Update ListCommands.cs Format fix --- libs/server/Resp/Objects/ListCommands.cs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/libs/server/Resp/Objects/ListCommands.cs b/libs/server/Resp/Objects/ListCommands.cs index e0f25d8d4e6..ef5a5365472 100644 --- a/libs/server/Resp/Objects/ListCommands.cs +++ b/libs/server/Resp/Objects/ListCommands.cs @@ -354,8 +354,7 @@ private unsafe bool ListBlockingMove() var dstKey = parseState.GetArgSliceByRef(1); if (!parseState.TryGetOperationDirection(2, out var srcDir) || - !parseState.TryGetOperationDirection(3, out var dstDir) - ) + !parseState.TryGetOperationDirection(3, out var dstDir)) { return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR); } @@ -772,8 +771,7 @@ private bool ListMove(ref TGarnetApi storageApi) var dstKey = parseState.GetArgSliceByRef(1); if (!parseState.TryGetOperationDirection(2, out var sourceDirection) || - !parseState.TryGetOperationDirection(3, out var destinationDirection) - ) + !parseState.TryGetOperationDirection(3, out var destinationDirection)) { return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR); } @@ -1032,4 +1030,4 @@ private unsafe bool ListBlockingPopMultiple() return true; } } -} \ No newline at end of file +} From ba649739435e12fa222bd0e1f769b0066da7c1ba Mon Sep 17 00:00:00 2001 From: prvyk Date: Fri, 15 Aug 2025 03:29:09 +0300 Subject: [PATCH 5/5] Address review comment --- libs/server/Resp/Objects/ListCommands.cs | 2 +- libs/server/Transaction/TxnKeyManager.cs | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/libs/server/Resp/Objects/ListCommands.cs b/libs/server/Resp/Objects/ListCommands.cs index ef5a5365472..58eb305e2bc 100644 --- a/libs/server/Resp/Objects/ListCommands.cs +++ b/libs/server/Resp/Objects/ListCommands.cs @@ -1030,4 +1030,4 @@ private unsafe bool ListBlockingPopMultiple() return true; } } -} +} \ No newline at end of file diff --git a/libs/server/Transaction/TxnKeyManager.cs b/libs/server/Transaction/TxnKeyManager.cs index 24a5d98438b..e309e8157e1 100644 --- a/libs/server/Transaction/TxnKeyManager.cs +++ b/libs/server/Transaction/TxnKeyManager.cs @@ -378,8 +378,7 @@ private int ListObjectKeys(RespCommand command) RespCommand.LINSERT => SingleKey(StoreType.Object, LockType.Exclusive), RespCommand.LLEN => SingleKey(StoreType.Object, LockType.Shared), RespCommand.LMOVE => ListKeys(2, StoreType.Object, LockType.Exclusive), - // Account for mandatory LEFT/RIGHT argument - RespCommand.LMPOP => ListReadKeysWithCount(LockType.Exclusive, mandatoryArgs: 1), + RespCommand.LMPOP => ListReadKeysWithCount(LockType.Exclusive), RespCommand.LPOP => SingleKey(StoreType.Object, LockType.Exclusive), RespCommand.LPOS => SingleKey(StoreType.Object, LockType.Shared), RespCommand.LPUSH => SingleKey(StoreType.Object, LockType.Exclusive), @@ -488,7 +487,7 @@ private int ListKeys(int inputCount, StoreType storeType, LockType type) /// /// Returns a list of keys for LMPOP command /// - private int ListReadKeysWithCount(LockType type, bool isObject = true, int mandatoryArgs = 0) + private int ListReadKeysWithCount(LockType type, bool isObject = true) { if (respSession.parseState.Count == 0) return -2; @@ -496,7 +495,7 @@ private int ListReadKeysWithCount(LockType type, bool isObject = true, int manda if (!respSession.parseState.TryGetInt(0, out var numKeys)) return -2; - if (numKeys + mandatoryArgs + 1 > respSession.parseState.Count) + if (numKeys + 1 > respSession.parseState.Count) return -2; for (var i = 1; i < numKeys + 1; i++)