diff --git a/libs/server/Resp/Objects/ListCommands.cs b/libs/server/Resp/Objects/ListCommands.cs index 69a6df8e0b9..58eb305e2bc 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,12 @@ 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 +380,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 +770,8 @@ 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 +959,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 bf265a7d0b3..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); - 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/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 6217bf9cbe0..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'); @@ -1188,43 +1187,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/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 07707909f0c..e309e8157e1 100644 --- a/libs/server/Transaction/TxnKeyManager.cs +++ b/libs/server/Transaction/TxnKeyManager.cs @@ -489,15 +489,18 @@ private int ListKeys(int inputCount, StoreType storeType, LockType type) /// private int ListReadKeysWithCount(LockType type, bool isObject = true) { - 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 + 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); }