diff --git a/libs/server/API/GarnetApiObjectCommands.cs b/libs/server/API/GarnetApiObjectCommands.cs index 481eeacfa6..cb73fe002a 100644 --- a/libs/server/API/GarnetApiObjectCommands.cs +++ b/libs/server/API/GarnetApiObjectCommands.cs @@ -299,6 +299,9 @@ public GarnetStatus SetScan(ArgSlice key, long cursor, string match, int count, /// public GarnetStatus SetDiff(ArgSlice[] keys, out HashSet members) => storageSession.SetDiff(keys, out members, ref objectContext); + + public GarnetStatus SetDiffStore(byte[] key, ArgSlice[] keys, out int count) + => storageSession.SetDiffStore(key, keys, out count, ref objectContext); #endregion #region Hash Methods diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs index fb63908b3b..a2d031f8bd 100644 --- a/libs/server/API/IGarnetApi.cs +++ b/libs/server/API/IGarnetApi.cs @@ -521,6 +521,15 @@ public interface IGarnetApi : IGarnetReadApi, IGarnetAdvancedApi /// GarnetStatus SetPop(byte[] key, ArgSlice input, ref GarnetObjectStoreOutput outputFooter); + /// + /// Diff result store. + /// Returns the number of result set. + /// + /// destination + /// + /// + /// + public GarnetStatus SetDiffStore(byte[] key, ArgSlice[] keys, out int count); #endregion #region List Methods diff --git a/libs/server/Objects/Set/SetObject.cs b/libs/server/Objects/Set/SetObject.cs index 0c6bba3ec2..9e4d71df76 100644 --- a/libs/server/Objects/Set/SetObject.cs +++ b/libs/server/Objects/Set/SetObject.cs @@ -24,6 +24,7 @@ public enum SetOperation : byte SCARD, SSCAN, SDIFF, + SDIFFSTORE, } diff --git a/libs/server/Resp/Objects/SetCommands.cs b/libs/server/Resp/Objects/SetCommands.cs index 45bad56dee..5608b44576 100644 --- a/libs/server/Resp/Objects/SetCommands.cs +++ b/libs/server/Resp/Objects/SetCommands.cs @@ -501,5 +501,56 @@ private unsafe bool SetDiff(int count, byte* ptr, ref TGarnetApi sto return true; } + + private unsafe bool SetDiffStore(int count, byte* ptr, ref TGarnetApi storageApi) + where TGarnetApi : IGarnetApi + { + ptr += 17; + + if (count < 3) + { + return AbortWithWrongNumberOfArguments("SDIFFSTORE", count); + } + + // Get the key + if (!RespReadUtils.ReadByteArrayWithLengthHeader(out var key, ref ptr, recvBufferPtr + bytesRead)) + return false; + + if (NetworkSingleKeySlotVerify(key, false)) + { + var bufSpan = new ReadOnlySpan(recvBufferPtr, bytesRead); + if (!DrainCommands(bufSpan, count)) + return false; + return true; + } + + var keys = new ArgSlice[count - 2]; + for (var i = 0; i < count - 2; i++) + { + keys[i] = default; + if (!RespReadUtils.ReadPtrWithLengthHeader(ref keys[i].ptr, ref keys[i].length, ref ptr, recvBufferPtr + bytesRead)) + return false; + } + + if (NetworkKeyArraySlotVerify(ref keys, true)) + { + var bufSpan = new ReadOnlySpan(recvBufferPtr, bytesRead); + if (!DrainCommands(bufSpan, count)) return false; + return true; + } + + var status = storageApi.SetDiffStore(key, keys, out var output); + + if (status == GarnetStatus.OK) + { + while (!RespWriteUtils.WriteInteger(output, ref dcurr, dend)) + SendAndReset(); + } + + // Move input head + readHead = (int)(ptr - recvBufferPtr); + + return true; + } } } \ No newline at end of file diff --git a/libs/server/Resp/RespCommand.cs b/libs/server/Resp/RespCommand.cs index 98f4cddfc8..db6cf428c2 100644 --- a/libs/server/Resp/RespCommand.cs +++ b/libs/server/Resp/RespCommand.cs @@ -665,6 +665,9 @@ private RespCommand FastParseCommand(byte* ptr) //[$5|SDIFF|] = 14 bytes = 8 (long) + 2 (ushort) if (*(long*)ptr == 5064654379396445476L && *(ushort*)(ptr + 8) == 3398 && *(ptr + 10) == 10) return (RespCommand.Set, (byte)SetOperation.SDIFF); + //[$10|SDIFFSTORE|] = + if (*(long*)ptr == 5279435965821104420L && *(ushort*)(ptr + 8) == 17990 && *(ptr + 16) == 10) + return (RespCommand.Set, (byte)SetOperation.SDIFFSTORE); #endregion } diff --git a/libs/server/Resp/RespCommandsInfo.cs b/libs/server/Resp/RespCommandsInfo.cs index da5ab60f66..d6ddaa3f58 100644 --- a/libs/server/Resp/RespCommandsInfo.cs +++ b/libs/server/Resp/RespCommandsInfo.cs @@ -207,6 +207,7 @@ public static RespCommandsInfo findCommand(RespCommand cmd, byte subCmd = 0) {(byte)SetOperation.SPOP, new RespCommandsInfo("SPOP", RespCommand.Set, -2, null, (byte)SetOperation.SPOP) }, {(byte)SetOperation.SSCAN, new RespCommandsInfo("SSCAN", RespCommand.Set, -3, null, (byte)SetOperation.SSCAN) }, {(byte)SetOperation.SDIFF, new RespCommandsInfo("SDIFF", RespCommand.Set, -2, null, (byte)SetOperation.SDIFF) }, + {(byte)SetOperation.SDIFFSTORE, new RespCommandsInfo("SDIFFSTORE", RespCommand.Set, -3, null, (byte)SetOperation.SDIFFSTORE) } }; private static readonly Dictionary customCommandsInfoMap = new Dictionary diff --git a/libs/server/Resp/RespInfo.cs b/libs/server/Resp/RespInfo.cs index 663ce502ed..3352aceae3 100644 --- a/libs/server/Resp/RespInfo.cs +++ b/libs/server/Resp/RespInfo.cs @@ -38,7 +38,7 @@ public static HashSet GetCommands() // Pub/sub "PUBLISH", "SUBSCRIBE", "PSUBSCRIBE", "UNSUBSCRIBE", "PUNSUBSCRIBE", // Set - "SADD", "SREM", "SPOP", "SMEMBERS", "SCARD", "SSCAN", "SDIFF", + "SADD", "SREM", "SPOP", "SMEMBERS", "SCARD", "SSCAN", "SDIFF", "SDIFFSTORE", //Scan ops "DBSIZE", "KEYS","SCAN", // Geospatial commands diff --git a/libs/server/Resp/RespServerSession.cs b/libs/server/Resp/RespServerSession.cs index c2d3263a38..93ef03082c 100644 --- a/libs/server/Resp/RespServerSession.cs +++ b/libs/server/Resp/RespServerSession.cs @@ -496,6 +496,7 @@ private bool ProcessArrayCommands(ref TGarnetApi storageApi) (RespCommand.Set, (byte)SetOperation.SPOP) => SetPop(count, ptr, ref storageApi), (RespCommand.Set, (byte)SetOperation.SSCAN) => ObjectScan(count, ptr, GarnetObjectType.Set, ref storageApi), (RespCommand.Set, (byte)SetOperation.SDIFF) => SetDiff(count, ptr, ref storageApi), + (RespCommand.Set, (byte)SetOperation.SDIFFSTORE) => SetDiffStore(count, ptr, ref storageApi), _ => ProcessOtherCommands(count, ref storageApi), }; return success; diff --git a/libs/server/Storage/Session/ObjectStore/SetOps.cs b/libs/server/Storage/Session/ObjectStore/SetOps.cs index 19e081d72b..167da0d802 100644 --- a/libs/server/Storage/Session/ObjectStore/SetOps.cs +++ b/libs/server/Storage/Session/ObjectStore/SetOps.cs @@ -440,7 +440,7 @@ public GarnetStatus SetPop(byte[] key, ArgSlice input, ref Garne /// /// /// - public GarnetStatus SetDiff(ArgSlice[] keys,out HashSet members, ref TObjectContext objectContext) + public GarnetStatus SetDiff(ArgSlice[] keys, out HashSet members, ref TObjectContext objectContext) where TObjectContext : ITsavoriteContext { members = default; @@ -468,5 +468,62 @@ public GarnetStatus SetDiff(ArgSlice[] keys,out HashSet return GarnetStatus.OK; } + + /// + /// Diff result store. + /// Returns the number of result set. + /// + /// + /// destination + /// + /// + /// + /// + public GarnetStatus SetDiffStore(byte[] key, ArgSlice[] keys,out int count, ref TObjectContext objectContext) + where TObjectContext : ITsavoriteContext + { + count = default; + + if (key.Length == 0 || keys.Length == 0) + return GarnetStatus.OK; + + var diffSet = _setDiff(keys, ref objectContext); + + var asKey = scratchBufferManager.CreateArgSlice(key); + var asMembers = new ArgSlice[diffSet.Count]; + for (var i = 0; i < diffSet.Count; i++) + { + asMembers[i] = scratchBufferManager.CreateArgSlice(diffSet.ElementAt(i)); + } + + var status = SetAdd(asKey, [.. asMembers], out var saddCount, ref objectContext); + count = saddCount; + return status; + } + + private HashSet _setDiff(ArgSlice[] keys, ref TObjectContext objectContext) + where TObjectContext : ITsavoriteContext + { + var result = new HashSet(); + + var status = GET(keys[0].Bytes, out var first, ref objectContext); + if (status == GarnetStatus.OK) + { + result.UnionWith(((SetObject)first.garnetObject).Set); + } + + for (var i = 1; i < keys.Length; i++) + { + status = GET(keys[i].Bytes, out var next, ref objectContext); + if (status == GarnetStatus.OK) + { + var nextSet = ((SetObject)next.garnetObject).Set; + var interItems = result.Intersect(nextSet, nextSet.Comparer); + result.ExceptWith(interItems); + } + } + + return result; + } } } \ No newline at end of file diff --git a/libs/server/Transaction/TxnKeyManager.cs b/libs/server/Transaction/TxnKeyManager.cs index 30ace244bf..fcf1c8b42e 100644 --- a/libs/server/Transaction/TxnKeyManager.cs +++ b/libs/server/Transaction/TxnKeyManager.cs @@ -48,7 +48,7 @@ internal int GetKeys(RespCommand command, int inputCount, out ReadOnlySpan RespCommand.SortedSet => SortedSetObjectKeys(subCommand, inputCount), RespCommand.List => ListObjectKeys(subCommand), RespCommand.Hash => HashObjectKeys(subCommand), - RespCommand.Set => SetObjectKeys(subCommand), + RespCommand.Set => SetObjectKeys(subCommand, inputCount), RespCommand.GET => SingleKey(1, false, LockType.Shared), RespCommand.SET => SingleKey(1, false, LockType.Exclusive), RespCommand.GETRANGE => SingleKey(1, false, LockType.Shared), @@ -172,7 +172,7 @@ private int HashObjectKeys(byte subCommand) }; } - private int SetObjectKeys(byte subCommand) + private int SetObjectKeys(byte subCommand, int inputCount) { return subCommand switch { @@ -181,6 +181,8 @@ private int SetObjectKeys(byte subCommand) (byte)SetOperation.SREM => SingleKey(1, true, LockType.Exclusive), (byte)SetOperation.SCARD => SingleKey(1, true, LockType.Exclusive), (byte)SetOperation.SPOP => SingleKey(1, true, LockType.Exclusive), + (byte)SetOperation.SDIFF => ListKeys(inputCount, true, LockType.Shared), + (byte)SetOperation.SDIFFSTORE => ListKeys(inputCount, true, LockType.Exclusive), _ => -1 }; } diff --git a/test/Garnet.test/RespSetTest.cs b/test/Garnet.test/RespSetTest.cs index dc7fcb81de..461654f837 100644 --- a/test/Garnet.test/RespSetTest.cs +++ b/test/Garnet.test/RespSetTest.cs @@ -4,6 +4,7 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Net; using System.Text; using NUnit.Framework; using StackExchange.Redis; @@ -260,8 +261,8 @@ public void CanDoSdiff() var key2 = "kye2"; var key2Value = new RedisValue[] { "c" }; - db.SetAdd(key1, key1Value); - db.SetAdd(key2, key2Value); + _ = db.SetAdd(key1, key1Value); + _ = db.SetAdd(key2, key2Value); var result = (RedisResult[])db.Execute("SDIFF", key1, key2); Assert.AreEqual(3, result.Length); @@ -275,7 +276,7 @@ public void CanDoSdiff() var key3 = "key3"; var key3Value = new RedisValue[] { "a", "c", "e" }; - db.SetAdd(key3, key3Value); + _ = db.SetAdd(key3, key3Value); result = (RedisResult[])db.Execute("SDIFF", key1, key2, key3); Assert.AreEqual(2, result.Length); @@ -286,7 +287,6 @@ public void CanDoSdiff() Assert.IsFalse(Array.Exists(result, t => t.ToString().Equals("c"))); Assert.IsFalse(Array.Exists(result, t => t.ToString().Equals("e"))); } - #endregion @@ -460,7 +460,23 @@ public void CanDoSdiffLC() lightClientRequest.SendCommand("SADD key3 a c e"); var response = lightClientRequest.SendCommand("SDIFF key1 key2 key3"); var expectedResponse = "*2\r\n$1\r\nb\r\n$1\r\nd\r\n"; - Assert.AreEqual(response.AsSpan().Slice(0,expectedResponse.Length).ToArray(), expectedResponse); + Assert.AreEqual(expectedResponse, response.AsSpan().Slice(0,expectedResponse.Length).ToArray()); + } + + [Test] + public void CanDoSdiffStoreLC() + { + var lightClientRequest = TestUtils.CreateRequest(); + lightClientRequest.SendCommand("SADD key1 a b c d"); + lightClientRequest.SendCommand("SADD key2 c"); + lightClientRequest.SendCommand("SADD key3 a c e"); + var response = lightClientRequest.SendCommand("SDIFFSTORE key key1 key2 key3"); + var expectedResponse = ":2\r\n"; + Assert.AreEqual(expectedResponse, response.AsSpan().Slice(0, expectedResponse.Length).ToArray()); + + var membersResponse = lightClientRequest.SendCommand("SMEMBERS key"); + expectedResponse = "*2\r\n$1\r\nb\r\n$1\r\nd\r\n"; + Assert.AreEqual(expectedResponse, membersResponse.AsSpan().Slice(0, expectedResponse.Length).ToArray()); } #endregion @@ -506,6 +522,21 @@ public void CanDoSdiffWhenKeyDoesNotExisting() var strResponse = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length); Assert.AreEqual(expectedResponse, strResponse); } + + [Test] + public void CanDoSdiffStoreWhenMemberKeysNotExisting() + { + using var lightClientRequest = TestUtils.CreateRequest(); + var response = lightClientRequest.SendCommand("SDIFFSTORE key key1 key2 key3"); + var expectedResponse = ":0\r\n"; + var strResponse = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length); + Assert.AreEqual(expectedResponse, strResponse); + + var membersResponse = lightClientRequest.SendCommand("SMEMBERS key"); + expectedResponse = "*0\r\n"; + strResponse = Encoding.ASCII.GetString(membersResponse).Substring(0, expectedResponse.Length); + Assert.AreEqual(expectedResponse, strResponse); + } #endregion