Skip to content

Commit 04a949b

Browse files
prvykTalZaccai
prvyk
and
TalZaccai
authored
Avoid unnecessary/buggy readHead access during command processing (#1336)
* Make ListReadKeysWithCount use idiomatic code instead of needless readHead and direct input parsing. * MakeUpperPtr should use length from ArgSlice instead of mistaken bytesRead calculation. * Remove GetOperationDirection entirely and use sessionparsestate extension. * Update ListCommands.cs Format fix * Address review comment --------- Co-authored-by: Tal Zaccai <[email protected]>
1 parent 4c5305d commit 04a949b

File tree

6 files changed

+42
-92
lines changed

6 files changed

+42
-92
lines changed

‎libs/server/Resp/Objects/ListCommands.cs‎

Lines changed: 16 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -226,10 +226,7 @@ private unsafe bool ListPopMultiple<TGarnetApi>(ref TGarnetApi storageApi)
226226
}
227227

228228
// Get the direction
229-
var dir = parseState.GetArgSliceByRef(currTokenId++);
230-
var popDirection = GetOperationDirection(dir);
231-
232-
if (popDirection == OperationDirection.Unknown)
229+
if (!parseState.TryGetOperationDirection(currTokenId++, out var popDirection))
233230
{
234231
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR);
235232
}
@@ -355,8 +352,12 @@ private unsafe bool ListBlockingMove()
355352

356353
var srcKey = parseState.GetArgSliceByRef(0);
357354
var dstKey = parseState.GetArgSliceByRef(1);
358-
var srcDir = parseState.GetArgSliceByRef(2);
359-
var dstDir = parseState.GetArgSliceByRef(3);
355+
356+
if (!parseState.TryGetOperationDirection(2, out var srcDir) ||
357+
!parseState.TryGetOperationDirection(3, out var dstDir))
358+
{
359+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR);
360+
}
360361

361362
if (!parseState.TryGetTimeout(4, out var timeout, out var error))
362363
{
@@ -379,27 +380,26 @@ private bool ListBlockingPopPush()
379380

380381
var srcKey = parseState.GetArgSliceByRef(0);
381382
var dstKey = parseState.GetArgSliceByRef(1);
382-
var rightOption = ArgSlice.FromPinnedSpan(CmdStrings.RIGHT);
383-
var leftOption = ArgSlice.FromPinnedSpan(CmdStrings.LEFT);
384383

385384
if (!parseState.TryGetTimeout(2, out var timeout, out var error))
386385
{
387386
return AbortWithErrorMessage(error);
388387
}
389388

390-
return ListBlockingMove(srcKey, dstKey, rightOption, leftOption, timeout);
389+
return ListBlockingMove(srcKey, dstKey, OperationDirection.Right,
390+
OperationDirection.Left, timeout);
391391
}
392392

393-
private bool ListBlockingMove(ArgSlice srcKey, ArgSlice dstKey, ArgSlice srcDir, ArgSlice dstDir, double timeout)
393+
private bool ListBlockingMove(ArgSlice srcKey, ArgSlice dstKey,
394+
OperationDirection sourceDirection,
395+
OperationDirection destinationDirection,
396+
double timeout)
394397
{
395398
var cmdArgs = new ArgSlice[] { default, default, default };
396399

397400
// Read destination key
398401
cmdArgs[0] = dstKey;
399402

400-
var sourceDirection = GetOperationDirection(srcDir);
401-
var destinationDirection = GetOperationDirection(dstDir);
402-
403403
if (sourceDirection == OperationDirection.Unknown || destinationDirection == OperationDirection.Unknown)
404404
{
405405
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR);
@@ -770,13 +770,8 @@ private bool ListMove<TGarnetApi>(ref TGarnetApi storageApi)
770770
var srcKey = parseState.GetArgSliceByRef(0);
771771
var dstKey = parseState.GetArgSliceByRef(1);
772772

773-
var srcDirSlice = parseState.GetArgSliceByRef(2);
774-
var dstDirSlice = parseState.GetArgSliceByRef(3);
775-
776-
var sourceDirection = GetOperationDirection(srcDirSlice);
777-
var destinationDirection = GetOperationDirection(dstDirSlice);
778-
779-
if (sourceDirection == OperationDirection.Unknown || destinationDirection == OperationDirection.Unknown)
773+
if (!parseState.TryGetOperationDirection(2, out var sourceDirection) ||
774+
!parseState.TryGetOperationDirection(3, out var destinationDirection))
780775
{
781776
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR);
782777
}
@@ -964,10 +959,7 @@ private unsafe bool ListBlockingPopMultiple()
964959
var cmdArgs = new ArgSlice[2];
965960

966961
// Get the direction
967-
var dir = parseState.GetArgSliceByRef(currTokenId++);
968-
var popDirection = GetOperationDirection(dir);
969-
970-
if (popDirection == OperationDirection.Unknown)
962+
if (!parseState.TryGetOperationDirection(currTokenId++, out var popDirection))
971963
{
972964
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR);
973965
}

‎libs/server/Resp/Objects/ObjectStoreUtils.cs‎

Lines changed: 0 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -99,27 +99,5 @@ private bool AbortWithErrorMessage(string format, params object[] args)
9999
{
100100
return AbortWithErrorMessage(Encoding.ASCII.GetBytes(string.Format(format, args)));
101101
}
102-
103-
/// <summary>
104-
/// Tries to parse the input as "LEFT" or "RIGHT" and returns the corresponding OperationDirection.
105-
/// If parsing fails, returns OperationDirection.Unknown.
106-
/// </summary>
107-
/// <param name="input">The input to parse.</param>
108-
/// <returns>The parsed OperationDirection, or OperationDirection.Unknown if parsing fails.</returns>
109-
public OperationDirection GetOperationDirection(ArgSlice input)
110-
{
111-
// Optimize for the common case
112-
if (input.ReadOnlySpan.SequenceEqual("LEFT"u8))
113-
return OperationDirection.Left;
114-
if (input.ReadOnlySpan.SequenceEqual("RIGHT"u8))
115-
return OperationDirection.Right;
116-
// Rare case: try making upper case and retry
117-
MakeUpperCase(input.ptr);
118-
if (input.ReadOnlySpan.SequenceEqual("LEFT"u8))
119-
return OperationDirection.Left;
120-
if (input.ReadOnlySpan.SequenceEqual("RIGHT"u8))
121-
return OperationDirection.Right;
122-
return OperationDirection.Unknown;
123-
}
124102
}
125103
}

‎libs/server/Resp/Parser/RespCommand.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2768,7 +2768,7 @@ private RespCommand ArrayParseCommand(bool writeErrorOnFailure, ref int count, r
27682768
var ptr = recvBufferPtr + readHead;
27692769

27702770
// See if input command is all upper-case. If not, convert and try fast parse pass again.
2771-
if (MakeUpperCase(ptr))
2771+
if (MakeUpperCase(ptr, bytesRead - readHead))
27722772
{
27732773
cmd = FastParseCommand(out count);
27742774
if (cmd != RespCommand.NONE)

‎libs/server/Resp/RespServerSession.cs‎

Lines changed: 1 addition & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -657,7 +657,7 @@ private void ProcessMessages()
657657
}
658658

659659
// Make first command in string as uppercase
660-
private bool MakeUpperCase(byte* ptr)
660+
private bool MakeUpperCase(byte* ptr, int len)
661661
{
662662
// Assume most commands are already upper case.
663663
// Assume most commands are 2-8 bytes long.
@@ -672,7 +672,6 @@ private bool MakeUpperCase(byte* ptr)
672672
// Note that _all_ of these bytes are <= 95 in the common case
673673
// and there's no need to scan the whole string in those cases.
674674

675-
var len = bytesRead - readHead;
676675
if (len >= 12)
677676
{
678677
var cmdLen = (uint)(*(ptr + 5) - '2');
@@ -1188,43 +1187,6 @@ ReadOnlySpan<byte> GetUpperCaseCommand(out bool success)
11881187
return result;
11891188
}
11901189

1191-
public ArgSlice GetCommandAsArgSlice(out bool success)
1192-
{
1193-
if (bytesRead - readHead < 6)
1194-
{
1195-
success = false;
1196-
return default;
1197-
}
1198-
1199-
Debug.Assert(*(recvBufferPtr + readHead) == '$');
1200-
int psize = *(recvBufferPtr + readHead + 1) - '0';
1201-
readHead += 2;
1202-
while (*(recvBufferPtr + readHead) != '\r')
1203-
{
1204-
psize = psize * 10 + *(recvBufferPtr + readHead) - '0';
1205-
if (bytesRead - readHead < 1)
1206-
{
1207-
success = false;
1208-
return default;
1209-
}
1210-
readHead++;
1211-
}
1212-
if (bytesRead - readHead < 2 + psize + 2)
1213-
{
1214-
success = false;
1215-
return default;
1216-
}
1217-
Debug.Assert(*(recvBufferPtr + readHead + 1) == '\n');
1218-
1219-
var result = new ArgSlice(recvBufferPtr + readHead + 2, psize);
1220-
Debug.Assert(*(recvBufferPtr + readHead + 2 + psize) == '\r');
1221-
Debug.Assert(*(recvBufferPtr + readHead + 2 + psize + 1) == '\n');
1222-
1223-
readHead += 2 + psize + 2;
1224-
success = true;
1225-
return result;
1226-
}
1227-
12281190
/// <summary>
12291191
/// Attempt to kill this session.
12301192
///

‎libs/server/SessionParseStateExtensions.cs‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -614,6 +614,21 @@ internal static bool TryGetManagerType(this SessionParseState parseState, int id
614614
return true;
615615
}
616616

617+
internal static bool TryGetOperationDirection(this SessionParseState parseState, int idx, out OperationDirection value)
618+
{
619+
value = OperationDirection.Unknown;
620+
var sbArg = parseState.GetArgSliceByRef(idx).ReadOnlySpan;
621+
622+
if (sbArg.EqualsUpperCaseSpanIgnoringCase("LEFT"u8))
623+
value = OperationDirection.Left;
624+
else if (sbArg.EqualsUpperCaseSpanIgnoringCase("RIGHT"u8))
625+
value = OperationDirection.Right;
626+
else
627+
return false;
628+
629+
return true;
630+
}
631+
617632
/// <summary>
618633
/// Parse sorted set add option from parse state at specified index
619634
/// </summary>

‎libs/server/Transaction/TxnKeyManager.cs‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -489,15 +489,18 @@ private int ListKeys(int inputCount, StoreType storeType, LockType type)
489489
/// </summary>
490490
private int ListReadKeysWithCount(LockType type, bool isObject = true)
491491
{
492-
var numKeysArg = respSession.GetCommandAsArgSlice(out var success);
493-
if (!success) return -2;
492+
if (respSession.parseState.Count == 0)
493+
return -2;
494494

495-
if (!NumUtils.TryParse(numKeysArg.ReadOnlySpan, out int numKeys)) return -2;
495+
if (!respSession.parseState.TryGetInt(0, out var numKeys))
496+
return -2;
496497

497-
for (var i = 0; i < numKeys; i++)
498+
if (numKeys + 1 > respSession.parseState.Count)
499+
return -2;
500+
501+
for (var i = 1; i < numKeys + 1; i++)
498502
{
499-
var key = respSession.GetCommandAsArgSlice(out success);
500-
if (!success) return -2;
503+
var key = respSession.parseState.GetArgSliceByRef(i);
501504
SaveKeyEntryToLock(key, isObject, type);
502505
SaveKeyArgSlice(key);
503506
}

0 commit comments

Comments
 (0)