fix: SPOP with count bug (#239)

* fix SPOP with count bug * fix SPOP with count bug --------- Co-authored-by: Tal Zaccai <talzacc@microsoft.com>
microsoft · Apr 10, 2024 · a149323 · a149323
1 parent b311568
commit a149323
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 0 deletions.
diff --git a/libs/server/Objects/Set/SetObjectImpl.cs b/libs/server/Objects/Set/SetObjectImpl.cs
@@ -221,6 +221,8 @@ private void SetPop(byte* input, int length, ref SpanByteAndMemory output)
                             ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
                         countDone++;
                     }
+
+                    countDone += count - countDone;
                 }
                 else if (count == int.MinValue) // no count parameter is present, we just pop and return a random item of the set
                 {

diff --git a/test/Garnet.test/RespSetTest.cs b/test/Garnet.test/RespSetTest.cs
@@ -465,6 +465,37 @@ public void CanDoSPOPWithCountCommandLC()
             Assert.AreEqual(2, arrLen);
         }
 
+        [Test]
+        public void CanDoSPOPWithMoreCountThanSetSizeCommandLC()
+        {
+            CreateLongSet();
+
+            var lightClientRequest = TestUtils.CreateRequest();
+
+            var response = lightClientRequest.SendCommand("SPOP myset 10", 5);
+
+            var strResponse = Encoding.ASCII.GetString(response);
+            Assert.AreEqual('*', strResponse[0]);
+
+            var arrLenEndIdx = strResponse.IndexOf("\r\n", StringComparison.InvariantCultureIgnoreCase);
+            Assert.IsTrue(arrLenEndIdx > 1);
+
+            var strArrLen = Encoding.ASCII.GetString(response).Substring(1, arrLenEndIdx - 1);
+            Assert.IsTrue(int.TryParse(strArrLen, out var arrLen));
+            Assert.IsTrue(arrLen == 5);
+
+            var lightClientRequest2 = TestUtils.CreateRequest();
+            var response2 = lightClientRequest2.SendCommand("SADD myset one");
+            var expectedResponse = ":1\r\n";
+            strResponse = Encoding.ASCII.GetString(response2).Substring(0, expectedResponse.Length);
+            Assert.AreEqual(expectedResponse, strResponse);
+
+            response2 = lightClientRequest2.SendCommand("SCARD myset");
+            expectedResponse = ":1\r\n";
+            strResponse = Encoding.ASCII.GetString(response2).Substring(0, expectedResponse.Length);
+            Assert.AreEqual(expectedResponse, strResponse);
+        }
+
         [Test]
         public void MultiWithNonExistingSet()
         {