[Compatibility] Added SMISMEMBER command (#713)

* 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]>
microsoft · Oct 18, 2024 · a4b47b4 · a4b47b4
1 parent d69545c
commit a4b47b4
Show file tree

Hide file tree

Showing 21 changed files with 361 additions and 10 deletions.
diff --git a/libs/resources/RespCommandsDocs.json b/libs/resources/RespCommandsDocs.json
@@ -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",

diff --git a/libs/resources/RespCommandsInfo.json b/libs/resources/RespCommandsInfo.json
@@ -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",

diff --git a/libs/server/API/GarnetApiObjectCommands.cs b/libs/server/API/GarnetApiObjectCommands.cs
@@ -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);

diff --git a/libs/server/API/GarnetWatchApi.cs b/libs/server/API/GarnetWatchApi.cs
@@ -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)
         {

diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs
@@ -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.

diff --git a/libs/server/Objects/Set/SetObject.cs b/libs/server/Objects/Set/SetObject.cs
@@ -26,6 +26,7 @@ public enum SetOperation : byte
         SMOVE,
         SRANDMEMBER,
         SISMEMBER,
+        SMISMEMBER,
         SUNION,
         SUNIONSTORE,
         SDIFF,
@@ -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;

diff --git a/libs/server/Objects/Set/SetObjectImpl.cs b/libs/server/Objects/Set/SetObjectImpl.cs
@@ -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;

diff --git a/libs/server/Resp/Objects/SetCommands.cs b/libs/server/Resp/Objects/SetCommands.cs
@@ -2,6 +2,7 @@
 // Licensed under the MIT license.
 
 using System;
+using System.Diagnostics;
 using Garnet.common;
 using Tsavorite.core;
 
@@ -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
@@ -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))

diff --git a/libs/server/Resp/Parser/RespCommand.cs b/libs/server/Resp/Parser/RespCommand.cs
@@ -59,6 +59,7 @@ public enum RespCommand : byte
         SINTER,
         SISMEMBER,
         SMEMBERS,
+        SMISMEMBER,
         SRANDMEMBER,
         SSCAN,
         STRLEN,
@@ -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:
@@ -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;
                                 }

diff --git a/libs/server/Resp/RespServerSession.cs b/libs/server/Resp/RespServerSession.cs
@@ -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),

diff --git a/libs/server/Storage/Session/ObjectStore/Common.cs b/libs/server/Storage/Session/ObjectStore/Common.cs
@@ -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>

diff --git a/libs/server/Storage/Session/ObjectStore/SetOps.cs b/libs/server/Storage/Session/ObjectStore/SetOps.cs
@@ -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>