From 6e971b43d78a49251dcf1ad84afbbdb693986c72 Mon Sep 17 00:00:00 2001 From: Badrish Chandramouli Date: Wed, 29 Jan 2025 19:45:20 -0800 Subject: [PATCH] Add SortedSetCount to IGarnetApi Use it in test custom transaction (stored proc) --- libs/server/API/GarnetApiObjectCommands.cs | 4 ++ libs/server/API/GarnetWatchApi.cs | 7 ++++ libs/server/API/IGarnetApi.cs | 10 +++++ .../Session/ObjectStore/SortedSetOps.cs | 40 +++++++++++++++++++ .../Extensions/SortedSetCountTxn.cs | 31 ++++++++++++++ test/Garnet.test/RespCustomCommandTests.cs | 19 +++++++++ 6 files changed, 111 insertions(+) create mode 100644 test/Garnet.test/Extensions/SortedSetCountTxn.cs diff --git a/libs/server/API/GarnetApiObjectCommands.cs b/libs/server/API/GarnetApiObjectCommands.cs index c54c86ad7b..d07988a211 100644 --- a/libs/server/API/GarnetApiObjectCommands.cs +++ b/libs/server/API/GarnetApiObjectCommands.cs @@ -82,6 +82,10 @@ public GarnetStatus SortedSetMPop(ReadOnlySpan 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); + /// + public GarnetStatus SortedSetCount(ArgSlice key, ArgSlice minScore, ArgSlice maxScore, out int numElements) + => storageSession.SortedSetCount(key, minScore, maxScore, out numElements, ref objectContext); + /// public GarnetStatus SortedSetCount(byte[] key, ref ObjectInput input, ref GarnetObjectStoreOutput output) => storageSession.SortedSetCount(key, ref input, ref output, ref objectContext); diff --git a/libs/server/API/GarnetWatchApi.cs b/libs/server/API/GarnetWatchApi.cs index 762148bde1..e95c491e34 100644 --- a/libs/server/API/GarnetWatchApi.cs +++ b/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); } + /// + 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); + } + /// public GarnetStatus SortedSetCount(byte[] key, ref ObjectInput input, ref GarnetObjectStoreOutput output) { diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs index 0aa7ae3246..249f59d8df 100644 --- a/libs/server/API/IGarnetApi.cs +++ b/libs/server/API/IGarnetApi.cs @@ -1271,6 +1271,16 @@ public interface IGarnetReadApi /// GarnetStatus SortedSetScores(byte[] key, ref ObjectInput input, ref GarnetObjectStoreOutput outputFooter); + /// + /// Returns the number of elements in the sorted set at key with a score between min and max. + /// + /// Key + /// Min Score + /// Max score + /// Number of elements + /// + GarnetStatus SortedSetCount(ArgSlice key, ArgSlice minScore, ArgSlice maxScore, out int numElements); + /// /// Returns the number of elements in the sorted set at key with a score between min and max. /// diff --git a/libs/server/Storage/Session/ObjectStore/SortedSetOps.cs b/libs/server/Storage/Session/ObjectStore/SortedSetOps.cs index d20588effc..d3f90cb840 100644 --- a/libs/server/Storage/Session/ObjectStore/SortedSetOps.cs +++ b/libs/server/Storage/Session/ObjectStore/SortedSetOps.cs @@ -892,6 +892,46 @@ public GarnetStatus SortedSetPop(byte[] key, ref ObjectInput inp where TObjectContext : ITsavoriteContext => RMWObjectStoreOperationWithOutput(key, ref input, ref objectStoreContext, ref outputFooter); + /// + /// Returns the number of elements in the sorted set at key with a score between min and max. + /// + /// + /// + /// + /// + /// + /// + /// + public unsafe GarnetStatus SortedSetCount(ArgSlice key, ArgSlice minScore, ArgSlice maxScore, out int numElements, ref TObjectContext objectContext) + where TObjectContext : ITsavoriteContext + { + 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; + } + /// /// Returns the number of elements in the sorted set at key with a score between min and max. /// diff --git a/test/Garnet.test/Extensions/SortedSetCountTxn.cs b/test/Garnet.test/Extensions/SortedSetCountTxn.cs new file mode 100644 index 0000000000..bf6f8520d1 --- /dev/null +++ b/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 api, ref CustomProcedureInput input) + { + int offset = 0; + AddKey(GetNextArg(ref input, ref offset), LockType.Exclusive, true); + return true; + } + + public override unsafe void Main(TGarnetApi api, ref CustomProcedureInput input, ref MemoryResult 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()); + } + } +} \ No newline at end of file diff --git a/test/Garnet.test/RespCustomCommandTests.cs b/test/Garnet.test/RespCustomCommandTests.cs index 7e4752c3e0..41c69853c4 100644 --- a/test/Garnet.test/RespCustomCommandTests.cs +++ b/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() {