diff --git a/libs/server/API/GarnetApi.cs b/libs/server/API/GarnetApi.cs index b036fdb3bd..5ce31b873c 100644 --- 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) /// - 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); + + /// + public GarnetStatus Increment(ArgSlice key, out long output, long incrementCount = 1) + => storageSession.Increment(key, out output, incrementCount, ref context); + + /// + 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 index 765dac6812..6ab965bd2d 100644 --- 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); + /// + public GarnetStatus SortedSetRank(ArgSlice key, ArgSlice member, bool reverse, out long? rank) + => storageSession.SortedSetRank(key, member, reverse, out rank, ref objectContext); + /// 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 index c5033fd28c..f9b1dff955 100644 --- 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); } + /// + 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); + } + /// 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 index e96df482e5..756715205b 100644 --- a/libs/server/API/IGarnetApi.cs +++ b/libs/server/API/IGarnetApi.cs @@ -185,6 +185,24 @@ public interface IGarnetApi : IGarnetReadApi, IGarnetAdvancedApi /// /// GarnetStatus Increment(ArgSlice key, ArgSlice input, ref ArgSlice output); + + /// + /// Increment (INCR, INCRBY) + /// + /// + /// + /// + /// + GarnetStatus Increment(ArgSlice key, out long output, long incrementCount = 1); + + /// + /// Decrement (DECR, DECRBY) + /// + /// + /// + /// + /// + GarnetStatus Decrement(ArgSlice key, out long output, long decrementCount = 1); #endregion #region DELETE @@ -1089,6 +1107,17 @@ public interface IGarnetReadApi /// GarnetStatus SortedSetRank(byte[] key, ArgSlice input, ref GarnetObjectStoreOutput outputFooter); + /// + /// 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 + /// + /// + /// + /// + /// + /// + GarnetStatus SortedSetRank(ArgSlice key, ArgSlice member, bool reverse, out long? rank); + /// /// Returns a random element from the sorted set key. /// diff --git a/libs/server/Storage/Session/MainStore/MainStoreOps.cs b/libs/server/Storage/Session/MainStore/MainStoreOps.cs index 19da8e8d86..3f6940bffe 100644 --- a/libs/server/Storage/Session/MainStore/MainStoreOps.cs +++ b/libs/server/Storage/Session/MainStore/MainStoreOps.cs @@ -834,6 +834,41 @@ public GarnetStatus Increment(ArgSlice key, ArgSlice input, ref ArgSli return GarnetStatus.OK; } + public unsafe GarnetStatus Increment(ArgSlice key, out long output, long increment, ref TContext context) + where TContext : ITsavoriteContext + { + 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 index 4a8f473d03..3ede0bba31 100644 --- a/libs/server/Storage/Session/ObjectStore/SortedSetOps.cs +++ b/libs/server/Storage/Session/ObjectStore/SortedSetOps.cs @@ -675,6 +675,50 @@ public unsafe GarnetStatus SortedSetScan(ArgSlice key, long curs } + /// + /// Returns the rank of member in the sorted set, the scores in the sorted set are ordered from high to low + /// The key of the sorted set + /// The member to get the rank + /// If true, the rank is calculated from low to high + /// The rank of the member (null if the member does not exist) + /// + /// + public unsafe GarnetStatus SortedSetRank(ArgSlice key, ArgSlice member, bool reverse, out long? rank, ref TObjectContext objectStoreContext) + where TObjectContext : ITsavoriteContext + { + 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; + } + /// /// Adds all the specified members with the specified scores to the sorted set stored at key. /// Current members get the score updated and reordered.