Add SortedSetCount to IGarnetApi (#978)

* Add SortedSetCount to IGarnetApi

Use it in test custom transaction (stored proc)

* fix lock type

libs/server/API/GarnetApiObjectCommands.cs

@@ -82,6 +82,10 @@ public GarnetStatus SortedSetMPop(ReadOnlySpan<ArgSlice> keys, int count, bool l
         public GarnetStatus SortedSetPop(ArgSlice key, out (ArgSlice member, ArgSlice score)[] pairs, int count = 1, bool lowScoresFirst = true)
             => storageSession.SortedSetPop(key, count, lowScoresFirst, out pairs, ref objectContext);
 
+        /// <inheritdoc />
+        public GarnetStatus SortedSetCount(ArgSlice key, ArgSlice minScore, ArgSlice maxScore, out int numElements)
+            => storageSession.SortedSetCount(key, minScore, maxScore, out numElements, ref objectContext);
+
         /// <inheritdoc />
         public GarnetStatus SortedSetCount(byte[] key, ref ObjectInput input, ref GarnetObjectStoreOutput output)
             => storageSession.SortedSetCount(key, ref input, ref output, ref objectContext);

libs/server/API/GarnetWatchApi.cs

@@ -119,6 +119,13 @@ public GarnetStatus SortedSetLength(byte[] key, ref ObjectInput input, out Objec
             return garnetApi.SortedSetLength(key, ref input, out output);
         }
 
+        /// <inheritdoc />
+        public GarnetStatus SortedSetCount(ArgSlice key, ArgSlice minScore, ArgSlice maxScore, out int numElements)
+        {
+            garnetApi.WATCH(key, StoreType.Object);
+            return garnetApi.SortedSetCount(key, minScore, maxScore, out numElements);
+        }
+
         /// <inheritdoc />
         public GarnetStatus SortedSetCount(byte[] key, ref ObjectInput input, ref GarnetObjectStoreOutput output)
         {

libs/server/API/IGarnetApi.cs

@@ -1271,6 +1271,16 @@ public interface IGarnetReadApi
         /// <returns></returns>
         GarnetStatus SortedSetScores(byte[] key, ref ObjectInput input, ref GarnetObjectStoreOutput outputFooter);
 
+        /// <summary>
+        /// Returns the number of elements in the sorted set at key with a score between min and max.
+        /// </summary>
+        /// <param name="key">Key</param>
+        /// <param name="minScore">Min Score</param>
+        /// <param name="maxScore">Max score</param>
+        /// <param name="numElements">Number of elements</param>
+        /// <returns></returns>
+        GarnetStatus SortedSetCount(ArgSlice key, ArgSlice minScore, ArgSlice maxScore, out int numElements);
+
         /// <summary>
         /// Returns the number of elements in the sorted set at key with a score between min and max.
         /// </summary>

libs/server/Storage/Session/ObjectStore/SortedSetOps.cs

@@ -892,6 +892,46 @@ public GarnetStatus SortedSetPop<TObjectContext>(byte[] key, ref ObjectInput inp
              where TObjectContext : ITsavoriteContext<byte[], IGarnetObject, ObjectInput, GarnetObjectStoreOutput, long, ObjectSessionFunctions, ObjectStoreFunctions, ObjectStoreAllocator>
             => RMWObjectStoreOperationWithOutput(key, ref input, ref objectStoreContext, ref outputFooter);
 
+        /// <summary>
+        /// Returns the number of elements in the sorted set at key with a score between min and max.
+        /// </summary>
+        /// <typeparam name="TObjectContext"></typeparam>
+        /// <param name="key"></param>
+        /// <param name="minScore"></param>
+        /// <param name="maxScore"></param>
+        /// <param name="numElements"></param>
+        /// <param name="objectContext"></param>
+        /// <returns></returns>
+        public unsafe GarnetStatus SortedSetCount<TObjectContext>(ArgSlice key, ArgSlice minScore, ArgSlice maxScore, out int numElements, ref TObjectContext objectContext)
+             where TObjectContext : ITsavoriteContext<byte[], IGarnetObject, ObjectInput, GarnetObjectStoreOutput, long, ObjectSessionFunctions, ObjectStoreFunctions, ObjectStoreAllocator>
+        {
+            numElements = 0;
+            if (key.Length == 0)
+                return GarnetStatus.OK;
+
+            // Prepare the parse state
+            parseState.InitializeWithArguments(minScore, maxScore);
+
+            // Prepare the input
+            var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = SortedSetOperation.ZCOUNT };
+            var input = new ObjectInput(header, ref parseState);
+
+            const int outputContainerSize = 32; // 3 for HEADER + CRLF + 20 for ascii long
+            var outputContainer = stackalloc byte[outputContainerSize];
+            var outputFooter = new GarnetObjectStoreOutput { SpanByteAndMemory = new SpanByteAndMemory(outputContainer, outputContainerSize) };
+
+            var status = ReadObjectStoreOperationWithOutput(key.ToArray(), ref input, ref objectContext, ref outputFooter);
+
+            if (status == GarnetStatus.OK)
+            {
+                Debug.Assert(*outputContainer == (byte)':');
+                var read = TryProcessRespSimple64IntOutput(outputFooter, out var value);
+                numElements = read ? (int)value : default;
+                Debug.Assert(read);
+            }
+            return status;
+        }
+
         /// <summary>
         /// Returns the number of elements in the sorted set at key with a score between min and max.
         /// </summary>

test/Garnet.test/Extensions/SortedSetCountTxn.cs

@@ -0,0 +1,31 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+using Garnet.common;
+using Garnet.server;
+using Tsavorite.core;
+
+namespace Garnet
+{
+    sealed class SortedSetCountTxn : CustomTransactionProcedure
+    {
+        public override bool Prepare<TGarnetReadApi>(TGarnetReadApi api, ref CustomProcedureInput input)
+        {
+            int offset = 0;
+            AddKey(GetNextArg(ref input, ref offset), LockType.Shared, true);
+            return true;
+        }
+
+        public override unsafe void Main<TGarnetApi>(TGarnetApi api, ref CustomProcedureInput input, ref MemoryResult<byte> output)
+        {
+            int offset = 0;
+            var key = GetNextArg(ref input, ref offset);
+            var minScore = GetNextArg(ref input, ref offset);
+            var maxScore = GetNextArg(ref input, ref offset);
+
+            var status = api.SortedSetCount(key, minScore, maxScore, out int numElements);
+
+            WriteSimpleString(ref output, numElements.ToString());
+        }
+    }
+}

test/Garnet.test/RespCustomCommandTests.cs

@@ -1140,6 +1140,25 @@ public void RateLimiterTest()
             ClassicAssert.AreEqual("ALLOWED", result.ToString());
         }
 
+        [Test]
+        public void SortedSetCountTxn()
+        {
+            server.Register.NewTransactionProc("SORTEDSETCOUNT", () => new SortedSetCountTxn(), new RespCommandsInfo { Arity = 4 });
+
+            using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig());
+            var db = redis.GetDatabase(0);
+
+            // Add entries to sorted set
+            for (var i = 1; i < 10; i++)
+            {
+                db.SortedSetAdd("key1", $"field{i}", i);
+            }
+
+            // Run transaction to get count of elements in given range of sorted set
+            var result = db.Execute("SORTEDSETCOUNT", "key1", 5, 10);
+            ClassicAssert.AreEqual(5, (int)result);
+        }
+
         [Test]
         public void CustomProcInvokingCustomCmdTest()
         {