Skip to content

Commit

Permalink
Merge branch 'main' into users/padgupta/add_flag_to_control_tlscheck
Browse files Browse the repository at this point in the history
  • Loading branch information
@badrishc
badrishc authored May 30, 2024
2 parents f956d60 + b34fc99 commit f9bf85f
Show file tree
Hide file tree
Showing 8 changed files with 128 additions and 49 deletions.
4 changes: 2 additions & 2 deletions libs/server/API/GarnetApiObjectCommands.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,8 +113,8 @@ public GarnetStatus SortedSetRemoveRange(byte[] key, ArgSlice input, out ObjectO
=> storageSession.SortedSetRemoveRange(key, input, out output, ref objectContext);

/// <inheritdoc />
public GarnetStatus SortedSetRank(byte[] key, ArgSlice input, out ObjectOutputHeader output)
=> storageSession.SortedSetRank(key, input, out output, ref objectContext);
public GarnetStatus SortedSetRank(byte[] key, ArgSlice input, ref GarnetObjectStoreOutput outputFooter)
=> storageSession.SortedSetRank(key, input, ref outputFooter, ref objectContext);

/// <inheritdoc />
public GarnetStatus SortedSetRandomMember(byte[] key, ArgSlice input, ref GarnetObjectStoreOutput outputFooter)
Expand Down
4 changes: 2 additions & 2 deletions libs/server/API/GarnetWatchApi.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -136,10 +136,10 @@ public GarnetStatus SortedSetScores(byte[] key, ArgSlice input, ref GarnetObject
}

/// <inheritdoc />
public GarnetStatus SortedSetRank(byte[] key, ArgSlice input, out ObjectOutputHeader output)
public GarnetStatus SortedSetRank(byte[] key, ArgSlice input, ref GarnetObjectStoreOutput outputFooter)
{
garnetApi.WATCH(key, StoreType.Object);
return garnetApi.SortedSetRank(key, input, out output);
return garnetApi.SortedSetRank(key, input, ref outputFooter);
}

/// <inheritdoc />
Expand Down
4 changes: 2 additions & 2 deletions libs/server/API/IGarnetApi.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1085,9 +1085,9 @@ public interface IGarnetReadApi
/// </summary>
/// <param name="key"></param>
/// <param name="input"></param>
/// <param name="output"></param>
/// <param name="outputFooter"></param>
/// <returns></returns>
GarnetStatus SortedSetRank(byte[] key, ArgSlice input, out ObjectOutputHeader output);
GarnetStatus SortedSetRank(byte[] key, ArgSlice input, ref GarnetObjectStoreOutput outputFooter);

/// <summary>
/// Returns a random element from the sorted set key.
Expand Down
4 changes: 2 additions & 2 deletions libs/server/Objects/SortedSet/SortedSetObject.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -212,7 +212,7 @@ public override unsafe bool Operate(ref SpanByte input, ref SpanByteAndMemory ou
SortedSetIncrement(_input, input.Length, ref output);
break;
case SortedSetOperation.ZRANK:
SortedSetRank(_input, input.Length, _output);
SortedSetRank(_input, input.Length, ref output);
break;
case SortedSetOperation.ZRANGE:
SortedSetRange(_input, input.Length, ref output);
Expand All @@ -239,7 +239,7 @@ public override unsafe bool Operate(ref SpanByte input, ref SpanByteAndMemory ou
SortedSetReverseRange(_input, input.Length, ref output);
break;
case SortedSetOperation.ZREVRANK:
SortedSetReverseRank(_input, input.Length, _output);
SortedSetReverseRank(_input, input.Length, ref output);
break;
case SortedSetOperation.ZREMRANGEBYLEX:
SortedSetRemoveRangeByLex(_input, input.Length, _output);
Expand Down
106 changes: 85 additions & 21 deletions libs/server/Objects/SortedSet/SortedSetObjectImpl.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -336,9 +336,9 @@ private void SortedSetIncrement(byte* input, int length, ref SpanByteAndMemory o
}
}

private void SortedSetRank(byte* input, int length, byte* output)
private void SortedSetRank(byte* input, int length, ref SpanByteAndMemory output)
{
GetRank(input, length, output);
GetRank(input, length, ref output);
}

private void SortedSetRange(byte* input, int length, ref SpanByteAndMemory output)
Expand Down Expand Up @@ -565,9 +565,9 @@ private void SortedSetReverseRange(byte* input, int length, ref SpanByteAndMemor
SortedSetRange(input, length, ref output);
}

private void SortedSetReverseRank(byte* input, int length, byte* output)
private void SortedSetReverseRank(byte* input, int length, ref SpanByteAndMemory output)
{
GetRank(input, length, output, ascending: false);
GetRank(input, length, ref output, ascending: false);
}

private void SortedSetRemoveRangeByLex(byte* input, int length, byte* output)
Expand Down Expand Up @@ -810,31 +810,95 @@ private void GetRangeOrCountByLex(byte* input, int length, byte* output, SortedS
/// <param name="length"></param>
/// <param name="output"></param>
/// <param name="ascending"></param>
private void GetRank(byte* input, int length, byte* output, bool ascending = true)
private void GetRank(byte* input, int length, ref SpanByteAndMemory output, bool ascending = true)
{
//ZRANK key member
var _input = (ObjectInputHeader*)input;
var _output = (ObjectOutputHeader*)output;

var member = new Span<byte>(input + sizeof(ObjectInputHeader), _input->count).ToArray();
var input_startptr = input + sizeof(ObjectInputHeader);
var input_currptr = input_startptr;
var withScore = false;

*_output = default;
var isMemory = false;
MemoryHandle ptrHandle = default;
byte* ptr = output.SpanByte.ToPointer();
var curr = ptr;
var end = curr + output.Length;
var error = false;

if (!sortedSetDict.TryGetValue(member, out var score))
{
_output->opsDone = -1;
return;
}
else
ObjectOutputHeader _output = default;
try
{
var rank = 0;
foreach (var item in sortedSet)
if (!RespReadUtils.ReadByteArrayWithLengthHeader(out var member, ref input_currptr, input + length))
return;

if (_input->count == 3) // ZRANK key member WITHSCORE
{
if (item.Item2.SequenceEqual(member))
break;
rank++;
if (!RespReadUtils.ReadByteArrayWithLengthHeader(out var token, ref input_currptr, input + length))
return;

if (Encoding.ASCII.GetString(token).ToUpperInvariant() == "WITHSCORE")
{
withScore = true;
}
else
{
while (!RespWriteUtils.WriteError(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR, ref curr, end))
ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
error = true;
}
}

if (!error)
{
if (!sortedSetDict.TryGetValue(member, out var score))
{
while (!RespWriteUtils.WriteNull(ref curr, end))
ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
}
else
{
var rank = 0;
foreach (var item in sortedSet)
{
if (item.Item2.SequenceEqual(member))
break;
rank++;
}

if (!ascending)
rank = sortedSet.Count - rank - 1;

if (withScore)
{
while (!RespWriteUtils.WriteArrayLength(2, ref curr, end)) // rank and score
ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);

while (!RespWriteUtils.WriteInteger(rank, ref curr, end))
ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);

while (!RespWriteUtils.WriteAsciiBulkString(score.ToString(CultureInfo.InvariantCulture), ref curr, end))
ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
}
else
{
while (!RespWriteUtils.WriteInteger(rank, ref curr, end))
ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
}
}
}
_output->opsDone = ascending ? rank : (sortedSet.Count - rank) - 1;

_output.bytesDone = (int)(input_currptr - input_startptr);
_output.countDone = _input->count;
_output.opsDone = 1;
}
finally
{
while (!RespWriteUtils.WriteDirect(ref _output, ref curr, end))
ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);

if (isMemory) ptrHandle.Dispose();

output.Length = (int)(curr - ptr);
}
}

Expand Down
30 changes: 13 additions & 17 deletions libs/server/Resp/Objects/SortedSetCommands.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -917,47 +917,43 @@ private unsafe bool SortedSetRank<TGarnetApi>(int count, byte* ptr, SortedSetOpe
return true;
}

byte* memberPtr = null;
int memberSize = 0;

// Read member parameter
if (!RespReadUtils.ReadPtrWithLengthHeader(ref memberPtr, ref memberSize, ref ptr, recvBufferPtr + bytesRead))
return false;

// Prepare input
var inputPtr = (ObjectInputHeader*)(memberPtr - sizeof(ObjectInputHeader));
var inputPtr = (ObjectInputHeader*)(ptr - sizeof(ObjectInputHeader));

// Save old values on buffer for possible revert
var save = *inputPtr;

// Prepare length of header in input buffer
var inputLength = memberSize + sizeof(ObjectInputHeader);
var inputLength = (int)(recvBufferPtr + bytesRead - (byte*)inputPtr);

// Prepare header in input buffer
inputPtr->header.type = GarnetObjectType.SortedSet;
inputPtr->header.flags = 0;
inputPtr->header.SortedSetOp = op;
inputPtr->count = memberSize;
inputPtr->count = count;
inputPtr->done = 0;

var status = storageApi.SortedSetRank(key, new ArgSlice((byte*)inputPtr, inputLength), out ObjectOutputHeader output);
// Prepare GarnetObjectStore output
var outputFooter = new GarnetObjectStoreOutput { spanByteAndMemory = new SpanByteAndMemory(dcurr, (int)(dend - dcurr)) };

var status = storageApi.SortedSetRank(key, new ArgSlice((byte*)inputPtr, inputLength), ref outputFooter);

// Reset input buffer
*inputPtr = save;
switch (status)
{
case GarnetStatus.OK:
// Process output
if (output.opsDone != -1)
while (!RespWriteUtils.WriteInteger(output.opsDone, ref dcurr, dend))
SendAndReset();
else
while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_ERRNOTFOUND, ref dcurr, dend))
SendAndReset();
var objOutputHeader = ProcessOutputWithHeader(outputFooter.spanByteAndMemory);
ptr += objOutputHeader.bytesDone;
if (objOutputHeader.bytesDone == 0)
return false;
break;

case GarnetStatus.NOTFOUND:
while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_ERRNOTFOUND, ref dcurr, dend))
SendAndReset();
ReadLeftToken(count - 1, ref ptr);
break;
}
}
Expand Down
6 changes: 3 additions & 3 deletions libs/server/Storage/Session/ObjectStore/SortedSetOps.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -851,12 +851,12 @@ public GarnetStatus SortedSetRemoveRange<TObjectContext>(byte[] key, ArgSlice in
/// <typeparam name="TObjectContext"></typeparam>
/// <param name="key"></param>
/// <param name="input"></param>
/// <param name="output"></param>
/// <param name="outputFooter"></param>
/// <param name="objectContext"></param>
/// <returns></returns>
public GarnetStatus SortedSetRank<TObjectContext>(byte[] key, ArgSlice input, out ObjectOutputHeader output, ref TObjectContext objectContext)
public GarnetStatus SortedSetRank<TObjectContext>(byte[] key, ArgSlice input, ref GarnetObjectStoreOutput outputFooter, ref TObjectContext objectContext)
where TObjectContext : ITsavoriteContext<byte[], IGarnetObject, SpanByte, GarnetObjectStoreOutput, long>
=> ReadObjectStoreOperation(key, input, out output, ref objectContext);
=> ReadObjectStoreOperationWithOutput(key, input, ref objectContext, ref outputFooter);

/// <summary>
/// Returns a random member from the sorted set key.
Expand Down
19 changes: 19 additions & 0 deletions test/Garnet.test/RespSortedSetTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1181,10 +1181,20 @@ public void CanDoZRankLC(int bytesSent)
var actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
Assert.AreEqual(expectedResponse, actualValue);

response = lightClientRequest.SendCommandChunks("ZRANK board Tom INVALIDOPTION", bytesSent);
expectedResponse = "-ERR syntax error\r\n";
actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
Assert.AreEqual(expectedResponse, actualValue);

response = lightClientRequest.SendCommandChunks("ZRANK board Tom", bytesSent);
expectedResponse = ":2\r\n";
actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
Assert.AreEqual(expectedResponse, actualValue);

response = lightClientRequest.SendCommandChunks("ZRANK board Tom withscore", bytesSent);
expectedResponse = "*2\r\n:2\r\n$3\r\n560\r\n";
actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
Assert.AreEqual(expectedResponse, actualValue);
}

[Test]
Expand All @@ -1204,6 +1214,11 @@ public void CanDoZRevRankLC(int bytesSent)
var actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
Assert.AreEqual(expectedResponse, actualValue);

response = lightClientRequest.SendCommandChunks("ZREVRANK board Tom INVALIDOPTION", bytesSent);
expectedResponse = "-ERR syntax error\r\n";
actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
Assert.AreEqual(expectedResponse, actualValue);

response = lightClientRequest.SendCommandChunks("ZREVRANK board Tom", bytesSent);
expectedResponse = ":0\r\n";
actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
Expand All @@ -1219,6 +1234,10 @@ public void CanDoZRevRankLC(int bytesSent)
actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
Assert.AreEqual(expectedResponse, actualValue);

response = lightClientRequest.SendCommandChunks("ZREVRANK board Dave WITHSCORE", bytesSent);
expectedResponse = "*2\r\n:2\r\n$3\r\n340\r\n";
actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
Assert.AreEqual(expectedResponse, actualValue);
}

[Test]
Expand Down

0 comments on commit f9bf85f

Please sign in to comment.