diff --git a/libs/server/Objects/SortedSet/SortedSetObject.cs b/libs/server/Objects/SortedSet/SortedSetObject.cs index c2e746473f..aa0e520f46 100644 --- a/libs/server/Objects/SortedSet/SortedSetObject.cs +++ b/libs/server/Objects/SortedSet/SortedSetObject.cs @@ -27,18 +27,12 @@ public enum SortedSetOperation : byte ZINCRBY, ZRANK, ZRANGE, - ZRANGEBYLEX, - ZRANGEBYSCORE, - ZRANGESTORE, GEOADD, GEOHASH, GEODIST, GEOPOS, GEOSEARCH, GEOSEARCHSTORE, - ZREVRANGE, - ZREVRANGEBYLEX, - ZREVRANGEBYSCORE, ZREVRANK, ZREMRANGEBYLEX, ZREMRANGEBYRANK, @@ -51,6 +45,38 @@ public enum SortedSetOperation : byte ZMSCORE } + /// + /// Options for specifying the range in sorted set operations. + /// + [Flags] + public enum SortedSetRangeOpts : byte + { + /// + /// No options specified. + /// + None = 0, + /// + /// Range by score. + /// + ByScore = 1, + /// + /// Range by lexicographical order. + /// + ByLex = 1 << 1, + /// + /// Reverse the range order. + /// + Reverse = 1 << 2, + /// + /// Store the result. + /// + Store = 1 << 3, + /// + /// Include scores in the result. + /// + WithScores = 1 << 4 + } + [Flags] public enum SortedSetAddOption { @@ -259,14 +285,6 @@ public override unsafe bool Operate(ref ObjectInput input, ref GarnetObjectStore case SortedSetOperation.ZRANK: SortedSetRank(ref input, ref output.SpanByteAndMemory); break; - case SortedSetOperation.ZRANGE: - case SortedSetOperation.ZRANGESTORE: - SortedSetRange(ref input, ref output.SpanByteAndMemory); - break; - case SortedSetOperation.ZRANGEBYLEX: - case SortedSetOperation.ZRANGEBYSCORE: - SortedSetRange(ref input, ref output.SpanByteAndMemory); - break; case SortedSetOperation.GEOADD: GeoAdd(ref input, ref output.SpanByteAndMemory); break; @@ -283,11 +301,7 @@ public override unsafe bool Operate(ref ObjectInput input, ref GarnetObjectStore case SortedSetOperation.GEOSEARCHSTORE: GeoSearch(ref input, ref output.SpanByteAndMemory); break; - case SortedSetOperation.ZREVRANGE: - SortedSetRange(ref input, ref output.SpanByteAndMemory); - break; - case SortedSetOperation.ZREVRANGEBYLEX: - case SortedSetOperation.ZREVRANGEBYSCORE: + case SortedSetOperation.ZRANGE: SortedSetRange(ref input, ref output.SpanByteAndMemory); break; case SortedSetOperation.ZREVRANK: diff --git a/libs/server/Objects/SortedSet/SortedSetObjectImpl.cs b/libs/server/Objects/SortedSet/SortedSetObjectImpl.cs index 837a62d45b..638878f427 100644 --- a/libs/server/Objects/SortedSet/SortedSetObjectImpl.cs +++ b/libs/server/Objects/SortedSet/SortedSetObjectImpl.cs @@ -416,6 +416,7 @@ private void SortedSetRange(ref ObjectInput input, ref SpanByteAndMemory output) //ZRANGE key min max [BYSCORE|BYLEX] [REV] [LIMIT offset count] [WITHSCORES] //ZRANGEBYSCORE key min max [WITHSCORES] [LIMIT offset count] //ZREVRANGEBYSCORE key max min [WITHSCORES] [LIMIT offset count] + var rangeOpts = (SortedSetRangeOpts)input.arg2; var count = input.parseState.Count; var respProtocolVersion = input.arg1; @@ -436,30 +437,13 @@ private void SortedSetRange(ref ObjectInput input, ref SpanByteAndMemory output) var maxSpan = input.parseState.GetArgSliceByRef(currIdx++).ReadOnlySpan; // read the rest of the arguments - ZRangeOptions options = new(); - switch (input.header.SortedSetOp) + ZRangeOptions options = new() { - case SortedSetOperation.ZRANGEBYLEX: - options.ByLex = true; - break; - case SortedSetOperation.ZRANGESTORE: - options.WithScores = true; - break; - case SortedSetOperation.ZRANGEBYSCORE: - options.ByScore = true; - break; - case SortedSetOperation.ZREVRANGE: - options.Reverse = true; - break; - case SortedSetOperation.ZREVRANGEBYSCORE: - options.ByScore = true; - options.Reverse = true; - break; - case SortedSetOperation.ZREVRANGEBYLEX: - options.ByLex = true; - options.Reverse = true; - break; - } + ByScore = (rangeOpts & SortedSetRangeOpts.ByScore) != 0, + ByLex = (rangeOpts & SortedSetRangeOpts.ByLex) != 0, + Reverse = (rangeOpts & SortedSetRangeOpts.Reverse) != 0, + WithScores = (rangeOpts & SortedSetRangeOpts.WithScores) != 0 || (rangeOpts & SortedSetRangeOpts.Store) != 0 + }; if (count > 2) { diff --git a/libs/server/Resp/Objects/SortedSetCommands.cs b/libs/server/Resp/Objects/SortedSetCommands.cs index 44e3d9854d..d0c3626fc5 100644 --- a/libs/server/Resp/Objects/SortedSetCommands.cs +++ b/libs/server/Resp/Objects/SortedSetCommands.cs @@ -159,20 +159,34 @@ private unsafe bool SortedSetRange(RespCommand command, ref TGarnetA var sbKey = parseState.GetArgSliceByRef(0).SpanByte; var keyBytes = sbKey.ToByteArray(); - var op = - command switch - { - RespCommand.ZRANGE => SortedSetOperation.ZRANGE, - RespCommand.ZREVRANGE => SortedSetOperation.ZREVRANGE, - RespCommand.ZRANGEBYLEX => SortedSetOperation.ZRANGEBYLEX, - RespCommand.ZRANGEBYSCORE => SortedSetOperation.ZRANGEBYSCORE, - RespCommand.ZREVRANGEBYLEX => SortedSetOperation.ZREVRANGEBYLEX, - RespCommand.ZREVRANGEBYSCORE => SortedSetOperation.ZREVRANGEBYSCORE, - _ => throw new Exception($"Unexpected {nameof(SortedSetOperation)}: {command}") - }; + var rangeOpts = SortedSetRangeOpts.None; - var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = op }; - var input = new ObjectInput(header, ref parseState, startIdx: 1, arg1: respProtocolVersion); + switch (command) + { + case RespCommand.ZRANGE: + break; + case RespCommand.ZREVRANGE: + rangeOpts = SortedSetRangeOpts.Reverse; + break; + case RespCommand.ZRANGEBYLEX: + rangeOpts = SortedSetRangeOpts.ByLex; + break; + case RespCommand.ZRANGEBYSCORE: + rangeOpts = SortedSetRangeOpts.ByScore; + break; + case RespCommand.ZREVRANGEBYLEX: + rangeOpts = SortedSetRangeOpts.ByLex | SortedSetRangeOpts.Reverse; + break; + case RespCommand.ZREVRANGEBYSCORE: + rangeOpts = SortedSetRangeOpts.ByScore | SortedSetRangeOpts.Reverse; + break; + case RespCommand.ZRANGESTORE: + rangeOpts = SortedSetRangeOpts.Store; + break; + } + + var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = SortedSetOperation.ZRANGE }; + var input = new ObjectInput(header, ref parseState, startIdx: 1, arg1: respProtocolVersion, arg2: (int)rangeOpts); var outputFooter = new GarnetObjectStoreOutput { SpanByteAndMemory = new SpanByteAndMemory(dcurr, (int)(dend - dcurr)) }; @@ -208,8 +222,8 @@ private unsafe bool SortedSetRangeStore(ref TGarnetApi storageApi) var dstKey = parseState.GetArgSliceByRef(0); var srcKey = parseState.GetArgSliceByRef(1); - var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = SortedSetOperation.ZRANGESTORE }; - var input = new ObjectInput(header, ref parseState, startIdx: 2, arg1: respProtocolVersion); + var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = SortedSetOperation.ZRANGE }; + var input = new ObjectInput(header, ref parseState, startIdx: 2, arg1: respProtocolVersion, arg2: (int)SortedSetRangeOpts.Store); var status = storageApi.SortedSetRangeStore(dstKey, srcKey, ref input, out int result); diff --git a/libs/server/Storage/Session/ObjectStore/SortedSetOps.cs b/libs/server/Storage/Session/ObjectStore/SortedSetOps.cs index d3f90cb840..050178bc40 100644 --- a/libs/server/Storage/Session/ObjectStore/SortedSetOps.cs +++ b/libs/server/Storage/Session/ObjectStore/SortedSetOps.cs @@ -466,37 +466,18 @@ public unsafe GarnetStatus SortedSetRange(ArgSlice key, ArgSlice return GarnetStatus.NOTFOUND; } - ReadOnlySpan operation = default; - var sortedOperation = SortedSetOperation.ZRANGE; - switch (sortedSetOrderOperation) + var rangeOpts = sortedSetOrderOperation switch { - case SortedSetOrderOperation.ByScore: - sortedOperation = SortedSetOperation.ZRANGEBYSCORE; - operation = "BYSCORE"u8; - break; - case SortedSetOrderOperation.ByLex: - sortedOperation = SortedSetOperation.ZRANGE; - operation = "BYLEX"u8; - break; - case SortedSetOrderOperation.ByRank: - if (reverse) - sortedOperation = SortedSetOperation.ZREVRANGE; - operation = default; - break; - } + SortedSetOrderOperation.ByScore => SortedSetRangeOpts.ByScore, + SortedSetOrderOperation.ByLex => SortedSetRangeOpts.ByLex, + _ => SortedSetRangeOpts.None + }; var arguments = new List { min, max }; - // Operation order - if (!operation.IsEmpty) - { - arguments.Add(scratchBufferManager.CreateArgSlice(operation)); - } - - // Reverse - if (sortedOperation != SortedSetOperation.ZREVRANGE && reverse) + if (reverse) { - arguments.Add(scratchBufferManager.CreateArgSlice("REV"u8)); + rangeOpts |= SortedSetRangeOpts.Reverse; } // Limit parameter @@ -517,9 +498,9 @@ public unsafe GarnetStatus SortedSetRange(ArgSlice key, ArgSlice parseState.InitializeWithArguments([.. arguments]); // Prepare the input - var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = sortedOperation }; + var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = SortedSetOperation.ZRANGE }; var inputArg = 2; // Default RESP server protocol version - var input = new ObjectInput(header, ref parseState, arg1: inputArg); + var input = new ObjectInput(header, ref parseState, arg1: inputArg, arg2: (int)rangeOpts); var outputFooter = new GarnetObjectStoreOutput { SpanByteAndMemory = new SpanByteAndMemory(null) }; var status = ReadObjectStoreOperationWithOutput(key.ToArray(), ref input, ref objectContext, ref outputFooter); diff --git a/libs/server/Transaction/TxnKeyManager.cs b/libs/server/Transaction/TxnKeyManager.cs index 4df5679563..96d3ea68c5 100644 --- a/libs/server/Transaction/TxnKeyManager.cs +++ b/libs/server/Transaction/TxnKeyManager.cs @@ -90,8 +90,8 @@ internal int GetKeys(RespCommand command, int inputCount, out ReadOnlySpan RespCommand.ZINCRBY => SortedSetObjectKeys(SortedSetOperation.ZINCRBY, inputCount), RespCommand.ZRANK => SortedSetObjectKeys(SortedSetOperation.ZRANK, inputCount), RespCommand.ZRANGE => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount), - RespCommand.ZRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZRANGEBYLEX, inputCount), - RespCommand.ZRANGEBYSCORE => SortedSetObjectKeys(SortedSetOperation.ZRANGEBYSCORE, inputCount), + RespCommand.ZRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount), + RespCommand.ZRANGEBYSCORE => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount), RespCommand.ZREVRANK => SortedSetObjectKeys(SortedSetOperation.ZREVRANK, inputCount), RespCommand.ZREMRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZREMRANGEBYLEX, inputCount), RespCommand.ZREMRANGEBYRANK => SortedSetObjectKeys(SortedSetOperation.ZREMRANGEBYRANK, inputCount), @@ -105,9 +105,9 @@ internal int GetKeys(RespCommand command, int inputCount, out ReadOnlySpan RespCommand.GEODIST => SortedSetObjectKeys(SortedSetOperation.GEODIST, inputCount), RespCommand.GEOPOS => SortedSetObjectKeys(SortedSetOperation.GEOPOS, inputCount), RespCommand.GEOSEARCH => SortedSetObjectKeys(SortedSetOperation.GEOSEARCH, inputCount), - RespCommand.ZREVRANGE => SortedSetObjectKeys(SortedSetOperation.ZREVRANGE, inputCount), - RespCommand.ZREVRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZREVRANGEBYLEX, inputCount), - RespCommand.ZREVRANGEBYSCORE => SortedSetObjectKeys(SortedSetOperation.ZREVRANGEBYSCORE, inputCount), + RespCommand.ZREVRANGE => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount), + RespCommand.ZREVRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount), + RespCommand.ZREVRANGEBYSCORE => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount), RespCommand.LINDEX => ListObjectKeys((byte)ListOperation.LINDEX), RespCommand.LINSERT => ListObjectKeys((byte)ListOperation.LINSERT), RespCommand.LLEN => ListObjectKeys((byte)ListOperation.LLEN), @@ -210,7 +210,6 @@ private int SortedSetObjectKeys(SortedSetOperation command, int inputCount) SortedSetOperation.ZINCRBY => SingleKey(1, true, LockType.Exclusive), SortedSetOperation.ZRANK => SingleKey(1, true, LockType.Exclusive), SortedSetOperation.ZRANGE => SingleKey(1, true, LockType.Shared), - SortedSetOperation.ZRANGEBYSCORE => SingleKey(1, true, LockType.Shared), SortedSetOperation.ZREVRANK => SingleKey(1, true, LockType.Exclusive), SortedSetOperation.ZREMRANGEBYLEX => SingleKey(1, true, LockType.Exclusive), SortedSetOperation.ZREMRANGEBYRANK => SingleKey(1, true, LockType.Exclusive), @@ -224,9 +223,6 @@ private int SortedSetObjectKeys(SortedSetOperation command, int inputCount) SortedSetOperation.GEODIST => SingleKey(1, true, LockType.Shared), SortedSetOperation.GEOPOS => SingleKey(1, true, LockType.Shared), SortedSetOperation.GEOSEARCH => SingleKey(1, true, LockType.Shared), - SortedSetOperation.ZREVRANGE => SingleKey(1, true, LockType.Shared), - SortedSetOperation.ZREVRANGEBYLEX => SingleKey(1, true, LockType.Shared), - SortedSetOperation.ZREVRANGEBYSCORE => SingleKey(1, true, LockType.Shared), _ => -1 }; }