API: Add user-friendly methods for INCR/DECR and ZRANK/ZREVRANK (#434)

* API: Add user-friendly methods for INCR/DECR and ZRANK/ZREVRANK operations * Remove test transaction * Fix formatting * Fix build error * Apply requested changes to the new api methods --------- Co-authored-by: Badrish Chandramouli <[email protected]> Co-authored-by: Tal Zaccai <[email protected]>
microsoft · Jun 12, 2024 · 66ebaa7 · 66ebaa7
1 parent 2647813
commit 66ebaa7
Show file tree

Hide file tree

Showing 6 changed files with 128 additions and 1 deletion.
diff --git a/libs/server/API/GarnetApi.cs b/libs/server/API/GarnetApi.cs
@@ -179,8 +179,16 @@ public unsafe GarnetStatus PERSIST(ArgSlice key, StoreType storeType = StoreType
 
         #region Increment (INCR, INCRBY, DECR, DECRBY)
         /// <inheritdoc />
-        public unsafe GarnetStatus Increment(ArgSlice key, ArgSlice input, ref ArgSlice output)
+        public GarnetStatus Increment(ArgSlice key, ArgSlice input, ref ArgSlice output)
             => storageSession.Increment(key, input, ref output, ref context);
+
+        /// <inheritdoc />
+        public GarnetStatus Increment(ArgSlice key, out long output, long incrementCount = 1)
+            => storageSession.Increment(key, out output, incrementCount, ref context);
+
+        /// <inheritdoc />
+        public GarnetStatus Decrement(ArgSlice key, out long output, long decrementCount = 1)
+            => Increment(key, out output, -decrementCount);
         #endregion
 
         #region DELETE

diff --git a/libs/server/API/GarnetApiObjectCommands.cs b/libs/server/API/GarnetApiObjectCommands.cs
@@ -116,6 +116,10 @@ public GarnetStatus SortedSetRemoveRange(byte[] key, ArgSlice input, out ObjectO
         public GarnetStatus SortedSetRank(byte[] key, ArgSlice input, ref GarnetObjectStoreOutput outputFooter)
             => storageSession.SortedSetRank(key, input, ref outputFooter, ref objectContext);
 
+        /// <inheritdoc />
+        public GarnetStatus SortedSetRank(ArgSlice key, ArgSlice member, bool reverse, out long? rank)
+            => storageSession.SortedSetRank(key, member, reverse, out rank, ref objectContext);
+
         /// <inheritdoc />
         public GarnetStatus SortedSetRandomMember(byte[] key, ArgSlice input, ref GarnetObjectStoreOutput outputFooter)
             => storageSession.SortedSetRandomMember(key, input, ref outputFooter, ref objectContext);

diff --git a/libs/server/API/GarnetWatchApi.cs b/libs/server/API/GarnetWatchApi.cs
@@ -142,6 +142,13 @@ public GarnetStatus SortedSetRank(byte[] key, ArgSlice input, ref GarnetObjectSt
             return garnetApi.SortedSetRank(key, input, ref outputFooter);
         }
 
+        /// <inheritdoc />
+        public GarnetStatus SortedSetRank(ArgSlice key, ArgSlice member, bool reverse, out long? rank)
+        {
+            garnetApi.WATCH(key, StoreType.Object);
+            return garnetApi.SortedSetRank(key, member, reverse, out rank);
+        }
+
         /// <inheritdoc />
         public GarnetStatus SortedSetRange(ArgSlice key, ArgSlice min, ArgSlice max, SortedSetOrderOperation sortedSetOrderOperation, out ArgSlice[] elements, out string error, bool withScores = false, bool reverse = false, (string, int) limit = default)
         {

diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs
@@ -185,6 +185,24 @@ public interface IGarnetApi : IGarnetReadApi, IGarnetAdvancedApi
         /// <param name="output"></param>
         /// <returns></returns>
         GarnetStatus Increment(ArgSlice key, ArgSlice input, ref ArgSlice output);
+
+        /// <summary>
+        /// Increment (INCR, INCRBY)
+        /// </summary>
+        /// <param name="key"></param>
+        /// <param name="output"></param>
+        /// <param name="incrementCount"></param>
+        /// <returns></returns>
+        GarnetStatus Increment(ArgSlice key, out long output, long incrementCount = 1);
+
+        /// <summary>
+        /// Decrement (DECR, DECRBY)
+        /// </summary>
+        /// <param name="key"></param>
+        /// <param name="output"></param>
+        /// <param name="decrementCount"></param>
+        /// <returns></returns>
+        GarnetStatus Decrement(ArgSlice key, out long output, long decrementCount = 1);
         #endregion
 
         #region DELETE
@@ -1089,6 +1107,17 @@ public interface IGarnetReadApi
         /// <returns></returns>
         GarnetStatus SortedSetRank(byte[] key, ArgSlice input, ref GarnetObjectStoreOutput outputFooter);
 
+        /// <summary>
+        /// ZRANK: Returns the rank of member in the sorted set, the scores in the sorted set are ordered from low to high
+        /// ZREVRANK: Returns the rank of member in the sorted set, with the scores ordered from high to low
+        /// </summary>
+        /// <param name="key"></param>
+        /// <param name="member"></param>
+        /// <param name="reverse"></param>
+        /// <param name="rank"></param>
+        /// <returns></returns>
+        GarnetStatus SortedSetRank(ArgSlice key, ArgSlice member, bool reverse, out long? rank);
+
         /// <summary>
         /// Returns a random element from the sorted set key.
         /// </summary>

diff --git a/libs/server/Storage/Session/MainStore/MainStoreOps.cs b/libs/server/Storage/Session/MainStore/MainStoreOps.cs
@@ -834,6 +834,41 @@ public GarnetStatus Increment<TContext>(ArgSlice key, ArgSlice input, ref ArgSli
             return GarnetStatus.OK;
         }
 
+        public unsafe GarnetStatus Increment<TContext>(ArgSlice key, out long output, long increment, ref TContext context)
+            where TContext : ITsavoriteContext<SpanByte, SpanByte, SpanByte, SpanByteAndMemory, long, MainStoreFunctions>
+        {
+            var cmd = RespCommand.INCRBY;
+            if (increment < 0)
+            {
+                cmd = RespCommand.DECRBY;
+                increment = -increment;
+            }
+
+            var incrementNumDigits = NumUtils.NumDigitsInLong(increment);
+            var inputByteSize = RespInputHeader.Size + incrementNumDigits;
+            var input = stackalloc byte[inputByteSize];
+            ((RespInputHeader*)input)->cmd = cmd;
+            ((RespInputHeader*)input)->flags = 0;
+            var longOutput = input + RespInputHeader.Size;
+            NumUtils.LongToBytes(increment, incrementNumDigits, ref longOutput);
+
+            const int outputBufferLength = NumUtils.MaximumFormatInt64Length + 1;
+            byte* outputBuffer = stackalloc byte[outputBufferLength];
+
+            var _key = key.SpanByte;
+            var _input = SpanByte.FromPinnedPointer(input, inputByteSize);
+            var _output = new SpanByteAndMemory(outputBuffer, outputBufferLength);
+
+            var status = context.RMW(ref _key, ref _input, ref _output);
+            if (status.IsPending)
+                CompletePendingForSession(ref status, ref _output, ref context);
+            Debug.Assert(_output.IsSpanByte);
+            Debug.Assert(_output.Length == outputBufferLength);
+
+            output = NumUtils.BytesToLong(_output.Length, outputBuffer);
+            return GarnetStatus.OK;
+        }
+
         public void WATCH(ArgSlice key, StoreType type)
         {
             txnManager.Watch(key, type);

diff --git a/libs/server/Storage/Session/ObjectStore/SortedSetOps.cs b/libs/server/Storage/Session/ObjectStore/SortedSetOps.cs
@@ -675,6 +675,50 @@ public unsafe GarnetStatus SortedSetScan<TObjectContext>(ArgSlice key, long curs
 
         }
 
+        /// <summary>
+        /// Returns the rank of member in the sorted set, the scores in the sorted set are ordered from high to low
+        /// <param name="key">The key of the sorted set</param>
+        /// <param name="member">The member to get the rank</param>
+        /// <param name="reverse">If true, the rank is calculated from low to high</param>
+        /// <param name="rank">The rank of the member (null if the member does not exist)</param>
+        /// <param name="objectStoreContext"></param>
+        /// </summary>
+        public unsafe GarnetStatus SortedSetRank<TObjectContext>(ArgSlice key, ArgSlice member, bool reverse, out long? rank, ref TObjectContext objectStoreContext)
+            where TObjectContext : ITsavoriteContext<byte[], IGarnetObject, SpanByte, GarnetObjectStoreOutput, long, ObjectStoreFunctions>
+        {
+            rank = null;
+            if (key.Length == 0)
+                return GarnetStatus.OK;
+
+            var inputSlice = scratchBufferManager.FormatScratchAsResp(ObjectInputHeader.Size, member);
+            var rawInput = (ObjectInputHeader*)inputSlice.ptr;
+            rawInput->header.type = GarnetObjectType.SortedSet;
+            rawInput->header.flags = 0;
+            rawInput->header.SortedSetOp = reverse ? SortedSetOperation.ZREVRANK : SortedSetOperation.ZRANK;
+            rawInput->count = 1;
+            rawInput->done = 0;
+
+            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(), inputSlice, ref objectStoreContext, ref outputFooter);
+
+            if (status == GarnetStatus.OK)
+            {
+                Debug.Assert(*outputContainer == (byte)'$' || *outputContainer == (byte)':');
+                if (*outputContainer == (byte)':')
+                {
+                    // member exists -> read the rank
+                    bool read = RespReadUtils.Read64Int(out var rankValue, ref outputContainer, &outputContainer[outputContainerSize]);
+                    Debug.Assert(read);
+                    rank = rankValue;
+                }
+            }
+
+            return status;
+        }
+
         /// <summary>
         /// Adds all the specified members with the specified scores to the sorted set stored at key.
         /// Current members get the score updated and reordered.