Skip to content

Commit

Permalink
[Compatibility] Added SMISMEMBER command (#713)
Browse files Browse the repository at this point in the history
* Added SMISMEMBER command

* Added doc for SMISMEMBER

* Code format fix

* Fixed test issue

* Fixed test issue finally

* Finial fix for the previous final fix for test case

* Added Garnet API

* Fixed code format

* Fixed build issue

---------

Co-authored-by: Yoganand Rajasekaran <[email protected]>
  • Loading branch information
Vijay-Nirmal and yrajas authored Oct 18, 2024
1 parent d69545c commit a4b47b4
Show file tree
Hide file tree
Showing 21 changed files with 361 additions and 10 deletions.
23 changes: 23 additions & 0 deletions libs/resources/RespCommandsDocs.json
Original file line number Diff line number Diff line change
Expand Up @@ -4747,6 +4747,29 @@
}
]
},
{
"Command": "SMISMEMBER",
"Name": "SMISMEMBER",
"Summary": "Determines whether multiple members belong to a set.",
"Group": "Set",
"Complexity": "O(N) where N is the number of elements being checked for membership",
"Arguments": [
{
"TypeDiscriminator": "RespCommandKeyArgument",
"Name": "KEY",
"DisplayText": "key",
"Type": "Key",
"KeySpecIndex": 0
},
{
"TypeDiscriminator": "RespCommandBasicArgument",
"Name": "MEMBER",
"DisplayText": "member",
"Type": "String",
"ArgumentFlags": "Multiple"
}
]
},
{
"Command": "SMOVE",
"Name": "SMOVE",
Expand Down
25 changes: 25 additions & 0 deletions libs/resources/RespCommandsInfo.json
Original file line number Diff line number Diff line change
Expand Up @@ -3479,6 +3479,31 @@
}
]
},
{
"Command": "SMISMEMBER",
"Name": "SMISMEMBER",
"Arity": -3,
"Flags": "Fast, ReadOnly",
"FirstKey": 1,
"LastKey": 1,
"Step": 1,
"AclCategories": "Fast, Read, Set",
"KeySpecifications": [
{
"BeginSearch": {
"TypeDiscriminator": "BeginSearchIndex",
"Index": 1
},
"FindKeys": {
"TypeDiscriminator": "FindKeysRange",
"LastKey": 0,
"KeyStep": 1,
"Limit": 0
},
"Flags": "RO, Access"
}
]
},
{
"Command": "SMOVE",
"Name": "SMOVE",
Expand Down
4 changes: 4 additions & 0 deletions libs/server/API/GarnetApiObjectCommands.cs
Original file line number Diff line number Diff line change
Expand Up @@ -301,6 +301,10 @@ public GarnetStatus SetMembers(byte[] key, ref ObjectInput input, ref GarnetObje
public GarnetStatus SetIsMember(byte[] key, ref ObjectInput input, ref GarnetObjectStoreOutput outputFooter)
=> storageSession.SetIsMember(key, ref input, ref outputFooter, ref objectContext);

/// <inheritdoc />
public GarnetStatus SetIsMember(ArgSlice key, ArgSlice[] members, out int[] result)
=> storageSession.SetIsMember(key, members, out result, ref objectContext);

/// <inheritdoc />
public GarnetStatus SetPop(ArgSlice key, out ArgSlice member)
=> storageSession.SetPop(key, out member, ref objectContext);
Expand Down
7 changes: 7 additions & 0 deletions libs/server/API/GarnetWatchApi.cs
Original file line number Diff line number Diff line change
Expand Up @@ -262,6 +262,13 @@ public GarnetStatus SetIsMember(byte[] key, ref ObjectInput input, ref GarnetObj
return garnetApi.SetIsMember(key, ref input, ref outputFooter);
}

/// <inheritdoc />
public GarnetStatus SetIsMember(ArgSlice key, ArgSlice[] members, out int[] result)
{
garnetApi.WATCH(key, StoreType.Object);
return garnetApi.SetIsMember(key, members, out result);
}

/// <inheritdoc />
public GarnetStatus SetMembers(byte[] key, ref ObjectInput input, ref GarnetObjectStoreOutput outputFooter)
{
Expand Down
8 changes: 8 additions & 0 deletions libs/server/API/IGarnetApi.cs
Original file line number Diff line number Diff line change
Expand Up @@ -1358,6 +1358,14 @@ public interface IGarnetReadApi
/// <returns></returns>
GarnetStatus SetIsMember(byte[] key, ref ObjectInput input, ref GarnetObjectStoreOutput outputFooter);

/// <summary>
/// Returns whether each member is a member of the set stored at key.
/// </summary>
/// <param name="key"></param>
/// <param name="members"></param>
/// <returns></returns>
GarnetStatus SetIsMember(ArgSlice key, ArgSlice[] members, out int[] result);

/// <summary>
/// Iterates over the members of the Set with the given key using a cursor,
/// a match pattern and count parameters.
Expand Down
4 changes: 4 additions & 0 deletions libs/server/Objects/Set/SetObject.cs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ public enum SetOperation : byte
SMOVE,
SRANDMEMBER,
SISMEMBER,
SMISMEMBER,
SUNION,
SUNIONSTORE,
SDIFF,
Expand Down Expand Up @@ -129,6 +130,9 @@ public override unsafe bool Operate(ref ObjectInput input, ref SpanByteAndMemory
case SetOperation.SISMEMBER:
SetIsMember(ref input, ref output);
break;
case SetOperation.SMISMEMBER:
SetMultiIsMember(ref input, ref output);
break;
case SetOperation.SREM:
SetRemove(ref input, _output);
break;
Expand Down
39 changes: 39 additions & 0 deletions libs/server/Objects/Set/SetObjectImpl.cs
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,45 @@ private void SetIsMember(ref ObjectInput input, ref SpanByteAndMemory output)
}
}

private void SetMultiIsMember(ref ObjectInput input, ref SpanByteAndMemory output)
{
var isMemory = false;
MemoryHandle ptrHandle = default;
var ptr = output.SpanByte.ToPointer();

var curr = ptr;
var end = curr + output.Length;

ObjectOutputHeader _output = default;
try
{
var totalCount = input.parseState.Count - input.parseStateFirstArgIdx;
while (!RespWriteUtils.WriteArrayLength(totalCount, ref curr, end))
ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);

var argCurr = input.parseStateFirstArgIdx;
while (argCurr < input.parseState.Count)
{
var member = input.parseState.GetArgSliceByRef(argCurr).SpanByte.ToByteArray();
var isMember = set.Contains(member);

while (!RespWriteUtils.WriteInteger(isMember ? 1 : 0, ref curr, end))
ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);

argCurr++;
}
_output.result1 = totalCount;
}
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);
}
}

private void SetRemove(ref ObjectInput input, byte* output)
{
var _output = (ObjectOutputHeader*)output;
Expand Down
42 changes: 36 additions & 6 deletions libs/server/Resp/Objects/SetCommands.cs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
// Licensed under the MIT license.

using System;
using System.Diagnostics;
using Garnet.common;
using Tsavorite.core;

Expand Down Expand Up @@ -374,20 +375,33 @@ private unsafe bool SetMembers<TGarnetApi>(ref TGarnetApi storageApi)
return true;
}

private unsafe bool SetIsMember<TGarnetApi>(ref TGarnetApi storageApi)
private unsafe bool SetIsMember<TGarnetApi>(RespCommand cmd, ref TGarnetApi storageApi)
where TGarnetApi : IGarnetApi
{
if (parseState.Count != 2)
Debug.Assert(cmd == RespCommand.SISMEMBER || cmd == RespCommand.SMISMEMBER);

var isSingle = cmd == RespCommand.SISMEMBER;
if (isSingle)
{
return AbortWithWrongNumberOfArguments("SISMEMBER");
if (parseState.Count != 2)
{
return AbortWithWrongNumberOfArguments("SISMEMBER");
}
}
else
{
if (parseState.Count < 2)
{
return AbortWithWrongNumberOfArguments("SMISMEMBER");
}
}

// Get the key
var sbKey = parseState.GetArgSliceByRef(0).SpanByte;
var keyBytes = sbKey.ToByteArray();

// Prepare input
var header = new RespInputHeader(GarnetObjectType.Set) { SetOp = SetOperation.SISMEMBER };
var header = new RespInputHeader(GarnetObjectType.Set) { SetOp = isSingle ? SetOperation.SISMEMBER : SetOperation.SMISMEMBER };
var input = new ObjectInput(header, ref parseState, 1);

// Prepare GarnetObjectStore output
Expand All @@ -402,8 +416,24 @@ private unsafe bool SetIsMember<TGarnetApi>(ref TGarnetApi storageApi)
ProcessOutputWithHeader(outputFooter.spanByteAndMemory);
break;
case GarnetStatus.NOTFOUND:
while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_RETURN_VAL_0, ref dcurr, dend))
SendAndReset();
if (isSingle)
{
while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_RETURN_VAL_0, ref dcurr, dend))
SendAndReset();
}
else
{
var count = parseState.Count - 1; // Remove key
while (!RespWriteUtils.WriteArrayLength(count, ref dcurr, dend))
SendAndReset();

for (var i = 0; i < count; i++)
{
while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_RETURN_VAL_0, ref dcurr, dend))
SendAndReset();
}
}

break;
case GarnetStatus.WRONGTYPE:
while (!RespWriteUtils.WriteError(CmdStrings.RESP_ERR_WRONG_TYPE, ref dcurr, dend))
Expand Down
9 changes: 7 additions & 2 deletions libs/server/Resp/Parser/RespCommand.cs
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ public enum RespCommand : byte
SINTER,
SISMEMBER,
SMEMBERS,
SMISMEMBER,
SRANDMEMBER,
SSCAN,
STRLEN,
Expand Down Expand Up @@ -1331,10 +1332,14 @@ private RespCommand FastParseArrayCommand(ref int count, ref ReadOnlySpan<byte>
{
return RespCommand.SDIFFSTORE;
}
else if (*(ulong*)(ptr + 1) == MemoryMarshal.Read<ulong>("10\r\nEXPI"u8) && *(uint*)(ptr + 9) == MemoryMarshal.Read<uint>("RETIME\r\n"u8))
else if (*(ulong*)(ptr + 1) == MemoryMarshal.Read<ulong>("10\r\nEXPI"u8) && *(ulong*)(ptr + 9) == MemoryMarshal.Read<ulong>("RETIME\r\n"u8))
{
return RespCommand.EXPIRETIME;
}
else if (*(ulong*)(ptr + 1) == MemoryMarshal.Read<ulong>("10\r\nSMIS"u8) && *(ulong*)(ptr + 9) == MemoryMarshal.Read<ulong>("MEMBER\r\n"u8))
{
return RespCommand.SMISMEMBER;
}
break;

case 11:
Expand Down Expand Up @@ -1362,7 +1367,7 @@ private RespCommand FastParseArrayCommand(ref int count, ref ReadOnlySpan<byte>
{
return RespCommand.SINTERSTORE;
}
else if (*(ulong*)(ptr + 2) == MemoryMarshal.Read<ulong>("1\r\nPEXPI"u8) && *(uint*)(ptr + 10) == MemoryMarshal.Read<uint>("RETIME\r\n"u8))
else if (*(ulong*)(ptr + 2) == MemoryMarshal.Read<ulong>("1\r\nPEXPI"u8) && *(ulong*)(ptr + 10) == MemoryMarshal.Read<ulong>("RETIME\r\n"u8))
{
return RespCommand.PEXPIRETIME;
}
Expand Down
3 changes: 2 additions & 1 deletion libs/server/Resp/RespServerSession.cs
Original file line number Diff line number Diff line change
Expand Up @@ -663,7 +663,8 @@ private bool ProcessArrayCommands<TGarnetApi>(RespCommand cmd, ref TGarnetApi st
// Set Commands
RespCommand.SADD => SetAdd(ref storageApi),
RespCommand.SMEMBERS => SetMembers(ref storageApi),
RespCommand.SISMEMBER => SetIsMember(ref storageApi),
RespCommand.SISMEMBER => SetIsMember(cmd, ref storageApi),
RespCommand.SMISMEMBER => SetIsMember(cmd, ref storageApi),
RespCommand.SREM => SetRemove(ref storageApi),
RespCommand.SCARD => SetLength(ref storageApi),
RespCommand.SPOP => SetPop(ref storageApi),
Expand Down
56 changes: 56 additions & 0 deletions libs/server/Storage/Session/ObjectStore/Common.cs
Original file line number Diff line number Diff line change
Expand Up @@ -296,6 +296,62 @@ unsafe ArgSlice[] ProcessRespArrayOutput(GarnetObjectStoreOutput outputFooter, o
return elements;
}

/// <summary>
/// Converts an array of elements in RESP format to ArgSlice[] type
/// </summary>
/// <param name="outputFooter">The RESP format output object</param>
/// <param name="error">A description of the error, if there is any</param>
/// <returns></returns>
unsafe int[] ProcessRespIntegerArrayOutput(GarnetObjectStoreOutput outputFooter, out string error)
{
int[] elements = default;
error = default;

// For reading the elements in the outputFooter
byte* element = null;

var outputSpan = outputFooter.spanByteAndMemory.IsSpanByte ?
outputFooter.spanByteAndMemory.SpanByte.AsReadOnlySpan() : outputFooter.spanByteAndMemory.AsMemoryReadOnlySpan();

try
{
fixed (byte* outputPtr = outputSpan)
{
var refPtr = outputPtr;

if (*refPtr == '-')
{
if (!RespReadUtils.ReadErrorAsString(out error, ref refPtr, outputPtr + outputSpan.Length))
return default;
}
else if (*refPtr == '*')
{
// Get the number of elements
if (!RespReadUtils.ReadUnsignedArrayLength(out var arraySize, ref refPtr, outputPtr + outputSpan.Length))
return default;

// Create the argslice[]
elements = new int[arraySize];
for (int i = 0; i < elements.Length; i++)
{
element = null;
if (RespReadUtils.TryReadInt(ref refPtr, outputPtr + outputSpan.Length, out var number, out var _))
{
elements[i] = number;
}
}
}
}
}
finally
{
if (!outputFooter.spanByteAndMemory.IsSpanByte)
outputFooter.spanByteAndMemory.Memory.Dispose();
}

return elements;
}

/// <summary>
/// Processes RESP output as pairs of score and member.
/// </summary>
Expand Down
35 changes: 35 additions & 0 deletions libs/server/Storage/Session/ObjectStore/SetOps.cs
Original file line number Diff line number Diff line change
Expand Up @@ -711,6 +711,41 @@ public GarnetStatus SetIsMember<TObjectContext>(byte[] key, ref ObjectInput inpu
where TObjectContext : ITsavoriteContext<byte[], IGarnetObject, ObjectInput, GarnetObjectStoreOutput, long, ObjectSessionFunctions, ObjectStoreFunctions, ObjectStoreAllocator>
=> ReadObjectStoreOperationWithOutput(key, ref input, ref objectContext, ref outputFooter);

/// <summary>
/// Returns whether each member is a member of the set stored at key.
/// </summary>
/// <typeparam name="TObjectContext"></typeparam>
/// <param name="key"></param>
/// <param name="members"></param>
/// <param name="objectContext"></param>
/// <returns></returns>
public unsafe GarnetStatus SetIsMember<TObjectContext>(ArgSlice key, ArgSlice[] members, out int[] result, ref TObjectContext objectContext)
where TObjectContext : ITsavoriteContext<byte[], IGarnetObject, ObjectInput, GarnetObjectStoreOutput, long, ObjectSessionFunctions, ObjectStoreFunctions, ObjectStoreAllocator>
{
result = default;

if (key.Length == 0)
return GarnetStatus.OK;

var parseState = new SessionParseState();
parseState.InitializeWithArguments(members);

// Prepare the input
var input = new ObjectInput(new RespInputHeader
{
type = GarnetObjectType.Set,
SetOp = SetOperation.SMISMEMBER,
}, ref parseState, 0);

var outputFooter = new GarnetObjectStoreOutput { spanByteAndMemory = new SpanByteAndMemory(null) };
var status = ReadObjectStoreOperationWithOutput(key.ToArray(), ref input, ref objectContext, ref outputFooter);

if (status == GarnetStatus.OK)
result = ProcessRespIntegerArrayOutput(outputFooter, out _);

return status;
}

/// <summary>
/// Removes and returns one or more random members from the set at key.
/// </summary>
Expand Down
Loading

0 comments on commit a4b47b4

Please sign in to comment.