Add LSET command (#274)

* LSET draft. * Add case tests. Add doc. Fix TryBytesToInt. * Update libs/server/Objects/List/ListObjectImpl.cs Co-authored-by: Paulus Pärssinen <[email protected]> * Update libs/server/Objects/List/ListObjectImpl.cs Co-authored-by: Paulus Pärssinen <[email protected]> * Update libs/server/Objects/List/ListObjectImpl.cs Co-authored-by: Paulus Pärssinen <[email protected]> * Rename to _in_* , _out_ * modify CanDoLSETbasic test * Add case tests. Rename variables. * Add RESP_ERR_GENERIC_INDEX_OUT_RANGE. --------- Co-authored-by: Paulus Pärssinen <[email protected]> Co-authored-by: Tal Zaccai <[email protected]>
microsoft · Apr 23, 2024 · df94b6d · df94b6d
1 parent a4dc14f
commit df94b6d
Show file tree

Hide file tree

Showing 16 changed files with 280 additions and 6 deletions.
diff --git a/libs/common/NumUtils.cs b/libs/common/NumUtils.cs
@@ -542,8 +542,9 @@ public static unsafe bool TryBytesToInt(byte* source, int len, out int result)
         {
             bool fNeg = (*source == '-');
             var beg = fNeg ? source + 1 : source;
+            var start = fNeg ? 1 : 0;
             result = 0;
-            for (int i = 0; i < len; ++i)
+            for (int i = start; i < len; ++i)
             {
                 if (!(source[i] >= 48 && source[i] <= 57))
                 {

diff --git a/libs/server/API/GarnetApiObjectCommands.cs b/libs/server/API/GarnetApiObjectCommands.cs
@@ -236,6 +236,10 @@ public GarnetStatus ListIndex(byte[] key, ArgSlice input, ref GarnetObjectStoreO
         public GarnetStatus ListRemove(byte[] key, ArgSlice input, out ObjectOutputHeader output)
             => storageSession.ListRemove(key, input, out output, ref objectContext);
 
+        /// <inheritdoc />
+        public GarnetStatus ListSet(byte[] key, ArgSlice input, ref GarnetObjectStoreOutput outputFooter)
+            => storageSession.ListSet(key, input, ref outputFooter, ref objectContext);
+
         #endregion
 
         #region Set Methods

diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs
@@ -733,6 +733,15 @@ public interface IGarnetApi : IGarnetReadApi, IGarnetAdvancedApi
         /// <returns></returns>
         GarnetStatus ListRemove(byte[] key, ArgSlice input, out ObjectOutputHeader output);
 
+        /// <summary>
+        /// Sets the list element at index to element.
+        /// </summary>
+        /// <param name="key"></param>
+        /// <param name="input"></param>
+        /// <param name="output"></param>
+        /// <returns></returns>
+        GarnetStatus ListSet(byte[] key, ArgSlice input, ref GarnetObjectStoreOutput output);
+
         #endregion
 
         #region Hash Methods

diff --git a/libs/server/Objects/List/ListObject.cs b/libs/server/Objects/List/ListObject.cs
@@ -31,6 +31,7 @@ public enum ListOperation : byte
         LREM,
         RPOPLPUSH,
         LMOVE,
+        LSET,
     }
 
     /// <summary>
@@ -174,6 +175,9 @@ public override unsafe bool Operate(ref SpanByte input, ref SpanByteAndMemory ou
                     case ListOperation.LREM:
                         ListRemove(_input, input.Length, _output);
                         break;
+                    case ListOperation.LSET:
+                        ListSet(_input, input.Length, ref output);
+                        break;
 
                     default:
                         throw new GarnetException($"Unsupported operation {(ListOperation)_input[0]} in ListObject.Operate");

diff --git a/libs/server/Objects/List/ListObjectImpl.cs b/libs/server/Objects/List/ListObjectImpl.cs
@@ -359,7 +359,7 @@ private void ListPop(byte* input, ref SpanByteAndMemory output, bool fDelAtHead)
                         list.RemoveLast();
                     }
 
-                    this.UpdateSize(node.Value, false);
+                    UpdateSize(node.Value, false);
                     while (!RespWriteUtils.WriteBulkString(node.Value, ref curr, end))
                         ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
 
@@ -381,5 +381,81 @@ private void ListPop(byte* input, ref SpanByteAndMemory output, bool fDelAtHead)
                 output.Length = (int)(curr - ptr);
             }
         }
+
+        private void ListSet(byte* input, int length, ref SpanByteAndMemory output)
+        {
+            var isMemory = false;
+            MemoryHandle ptrHandle = default;
+            var output_startptr = output.SpanByte.ToPointer();
+            var output_currptr = output_startptr;
+            var output_end = output_currptr + output.Length;
+
+            ObjectOutputHeader _output = default;
+
+            try
+            {
+                if (list.Count == 0)
+                {
+                    while (!RespWriteUtils.WriteError(CmdStrings.RESP_ERR_GENERIC_NOSUCHKEY, ref output_currptr, output_end))
+                        ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref output_startptr, ref ptrHandle, ref output_currptr, ref output_end);
+                    return;
+                }
+
+                byte* input_startptr = input + sizeof(ObjectInputHeader);
+                byte* input_currptr = input_startptr;
+                byte* input_end = input + length;
+
+                byte* indexParam = default;
+                var indexParamSize = 0;
+
+                // index
+                if (!RespReadUtils.ReadPtrWithLengthHeader(ref indexParam, ref indexParamSize, ref input_currptr, input_end))
+                    return;
+
+                if (NumUtils.TryBytesToInt(indexParam, indexParamSize, out var index) == false)
+                {
+                    while (!RespWriteUtils.WriteError(CmdStrings.RESP_ERR_GENERIC_VALUE_IS_NOT_INTEGER, ref output_currptr, output_end))
+                        ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref output_startptr, ref ptrHandle, ref output_currptr, ref output_end);
+                    return;
+                }
+
+                index = index < 0 ? list.Count + index : index;
+
+                if (index > list.Count - 1 || index < 0)
+                {
+                    while (!RespWriteUtils.WriteError(CmdStrings.RESP_ERR_GENERIC_INDEX_OUT_RANGE, ref output_currptr, output_end))
+                        ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref output_startptr, ref ptrHandle, ref output_currptr, ref output_end);
+                    return;
+                }
+
+                // element
+                if (!RespReadUtils.ReadByteArrayWithLengthHeader(out var element, ref input_currptr, input_end))
+                    return;
+
+                var targetNode = index == 0 ? list.First
+                    : (index == list.Count - 1 ? list.Last
+                        : list.Nodes().ElementAtOrDefault(index));
+
+                UpdateSize(targetNode.Value, false);
+                targetNode.Value = element;
+                UpdateSize(targetNode.Value);
+
+                while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_OK, ref output_currptr, output_end))
+                    ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref output_startptr, ref ptrHandle, ref output_currptr, ref output_end);
+
+                // Write bytes parsed from input and count done, into output footer
+                _output.bytesDone = (int)(input_currptr - input_startptr);
+                _output.countDone = 1;
+                _output.opsDone = 1;
+            }
+            finally
+            {
+                while (!RespWriteUtils.WriteDirect(ref _output, ref output_currptr, output_end))
+                    ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref output_startptr, ref ptrHandle, ref output_currptr, ref output_end);
+
+                if (isMemory) ptrHandle.Dispose();
+                output.Length = (int)(output_currptr - output_startptr);
+            }
+        }
     }
 }
diff --git a/libs/server/Resp/CmdStrings.cs b/libs/server/Resp/CmdStrings.cs
@@ -123,7 +123,7 @@ public static ReadOnlySpan<byte> GetConfig(ReadOnlySpan<byte> key)
         public static ReadOnlySpan<byte> RESP_ERR_GENERIC_REGISTERCS_UNSUPPORTED_CLASS => "ERR unable to register one or more unsupported classes."u8;
         public static ReadOnlySpan<byte> RESP_ERR_GENERIC_VALUE_IS_NOT_INTEGER => "ERR value is not an integer or out of range."u8;
         public static ReadOnlySpan<byte> RESP_ERR_GENERIC_UKNOWN_SUBCOMMAND => "ERR Unknown subcommand. Try LATENCY HELP."u8;
-
+        public static ReadOnlySpan<byte> RESP_ERR_GENERIC_INDEX_OUT_RANGE => "ERR index out of range"u8;
         /// <summary>
         /// Response string templates
         /// </summary>

diff --git a/libs/server/Resp/Objects/ListCommands.cs b/libs/server/Resp/Objects/ListCommands.cs
@@ -778,5 +778,72 @@ private unsafe bool ListMove<TGarnetApi>(int count, ArgSlice sourceKey, ArgSlice
 
             return storageApi.ListMove(sourceKey, destinationKey, sourceDirection, destinationDirection, out node);
         }
+
+        /// <summary>
+        /// Sets the list element at index to element
+        /// LSET key index element
+        /// </summary>
+        /// <typeparam name="TGarnetApi"></typeparam>
+        /// <param name="count"></param>
+        /// <param name="ptr"></param>
+        /// <param name="storageApi"></param>
+        /// <returns></returns>
+        public unsafe bool ListSet<TGarnetApi>(int count, byte* ptr, ref TGarnetApi storageApi)
+            where TGarnetApi : IGarnetApi
+        {
+            if (count != 3)
+            {
+                return AbortWithWrongNumberOfArguments("LSET", count);
+            }
+            else
+            {
+                // Get the key for List
+                if (!RespReadUtils.ReadByteArrayWithLengthHeader(out var key, ref ptr, recvBufferPtr + bytesRead))
+                    return false;
+
+                if (NetworkSingleKeySlotVerify(key, true))
+                {
+                    var bufSpan = new ReadOnlySpan<byte>(recvBufferPtr, bytesRead);
+                    if (!DrainCommands(bufSpan, count))
+                        return false;
+                    return true;
+                }
+
+                // Prepare input
+                var inputPtr = (ObjectInputHeader*)(ptr - sizeof(ObjectInputHeader));
+
+                // Save old values for possible revert
+                var save = *inputPtr;
+
+                var inputLength = (int)(recvBufferPtr + bytesRead - (byte*)inputPtr);
+
+                // Prepare header in input buffer
+                inputPtr->header.type = GarnetObjectType.List;
+                inputPtr->header.ListOp = ListOperation.LSET;
+                inputPtr->count = 0;
+                inputPtr->done = 0;
+
+                // Prepare GarnetObjectStore output
+                var outputFooter = new GarnetObjectStoreOutput { spanByteAndMemory = new SpanByteAndMemory(dcurr, (int)(dend - dcurr)) };
+
+                var statusOp = storageApi.ListSet(key, new ArgSlice((byte*)inputPtr, inputLength), ref outputFooter);
+
+                //restore input
+                *inputPtr = save;
+
+                switch (statusOp)
+                {
+                    case GarnetStatus.OK:
+                        //process output
+                        var objOutputHeader = ProcessOutputWithHeader(outputFooter.spanByteAndMemory);
+                        ptr += objOutputHeader.bytesDone;
+                        break;
+                }
+            }
+
+            // Move input head, write result to output
+            readHead = (int)(ptr - recvBufferPtr);
+            return true;
+        }
     }
 }
diff --git a/libs/server/Resp/RespCommand.cs b/libs/server/Resp/RespCommand.cs
@@ -391,6 +391,10 @@ static RespCommand MatchedNone(RespServerSession session, int oldReadHead)
                                         {
                                             return (RespCommand.List, (byte)ListOperation.LREM);
                                         }
+                                        else if (*(ulong*)(ptr + 2) == MemoryMarshal.Read<ulong>("\r\nLSET\r\n"u8))
+                                        {
+                                            return (RespCommand.List, (byte)ListOperation.LSET);
+                                        }
                                         break;
 
                                     case 'M':

diff --git a/libs/server/Resp/RespCommandsInfo.cs b/libs/server/Resp/RespCommandsInfo.cs
@@ -188,6 +188,7 @@ public static RespCommandsInfo findCommand(RespCommand cmd, byte subCmd = 0)
             {(byte)ListOperation.LINDEX,    new RespCommandsInfo("LINDEX",  RespCommand.List,    2, null, (byte)ListOperation.LINDEX)},
             {(byte)ListOperation.LINSERT,   new RespCommandsInfo("LINSERT", RespCommand.List,    4, null, (byte)ListOperation.LINSERT)},
             {(byte)ListOperation.LREM,      new RespCommandsInfo("LREM",    RespCommand.List,    3, null, (byte)ListOperation.LREM) },
+            {(byte)ListOperation.LSET,      new RespCommandsInfo("LSET",    RespCommand.List,    3, null, (byte)ListOperation.LSET) },
         };
 
         private static readonly Dictionary<byte, RespCommandsInfo> hashCommandsInfoMap = new Dictionary<byte, RespCommandsInfo>

diff --git a/libs/server/Resp/RespInfo.cs b/libs/server/Resp/RespInfo.cs
@@ -28,7 +28,7 @@ public static HashSet<string> GetCommands()
                 // Sorted Set
                 "ZADD", "ZCARD", "ZPOPMAX", "ZSCORE", "ZREM", "ZCOUNT", "ZINCRBY", "ZRANK", "ZRANGE", "ZRANGEBYSCORE", "ZREVRANGE", "ZREVRANK", "ZREMRANGEBYLEX", "ZREMRANGEBYRANK", "ZREMRANGEBYSCORE", "ZLEXCOUNT", "ZPOPMIN", "ZRANDMEMBER", "ZDIFF", "ZSCAN", "ZMSCORE",
                 // List
-                "LPOP", "LPUSH", "RPOP", "RPUSH", "LLEN", "LTRIM", "LRANGE", "LINDEX", "LINSERT", "LREM", "RPOPLPUSH", "LMOVE", "LPUSHX", "RPUSHX",
+                "LPOP", "LPUSH", "RPOP", "RPUSH", "LLEN", "LTRIM", "LRANGE", "LINDEX", "LINSERT", "LREM", "RPOPLPUSH", "LMOVE", "LPUSHX", "RPUSHX", "LSET",
                 // Hash
                 "HSET", "HGET", "HMGET", "HMSET", "HDEL", "HLEN", "HEXISTS", "HGETALL", "HKEYS", "HVALS", "HINCRBY", "HINCRBYFLOAT", "HSETNX", "HRANDFIELD", "HSCAN", "HSTRLEN",
                 // Hyperloglog

diff --git a/libs/server/Resp/RespServerSession.cs b/libs/server/Resp/RespServerSession.cs
@@ -482,6 +482,7 @@ private bool ProcessArrayCommands<TGarnetApi>(RespCommand cmd, byte subcmd, int
                 (RespCommand.List, (byte)ListOperation.LREM) => ListRemove(count, ptr, ref storageApi),
                 (RespCommand.List, (byte)ListOperation.RPOPLPUSH) => ListRightPopLeftPush(count, ptr, ref storageApi),
                 (RespCommand.List, (byte)ListOperation.LMOVE) => ListMove(count, ptr, ref storageApi),
+                (RespCommand.List, (byte)ListOperation.LSET) => ListSet(count, ptr, ref storageApi),
                 // Hash Commands
                 (RespCommand.Hash, (byte)HashOperation.HSET) => HashSet(count, ptr, HashOperation.HSET, ref storageApi),
                 (RespCommand.Hash, (byte)HashOperation.HMSET) => HashSet(count, ptr, HashOperation.HMSET, ref storageApi),

diff --git a/libs/server/Storage/Session/ObjectStore/ListOps.cs b/libs/server/Storage/Session/ObjectStore/ListOps.cs
@@ -413,5 +413,17 @@ public unsafe GarnetStatus ListLength<TObjectContext>(byte[] key, ArgSlice input
              where TObjectContext : ITsavoriteContext<byte[], IGarnetObject, SpanByte, GarnetObjectStoreOutput, long>
              => ReadObjectStoreOperation(key, input, out output, ref objectStoreContext);
 
+        /// <summary>
+        /// Sets the list element at index to element.
+        /// </summary>
+        /// <typeparam name="TObjectContext"></typeparam>
+        /// <param name="key"></param>
+        /// <param name="input"></param>
+        /// <param name="outputFooter"></param>
+        /// <param name="objectStoreContext"></param>
+        /// <returns></returns>
+        public unsafe GarnetStatus ListSet<TObjectContext>(byte[] key, ArgSlice input, ref GarnetObjectStoreOutput outputFooter, ref TObjectContext objectStoreContext)
+            where TObjectContext : ITsavoriteContext<byte[], IGarnetObject, SpanByte, GarnetObjectStoreOutput, long>
+            => RMWObjectStoreOperationWithOutput(key, input, ref objectStoreContext, ref outputFooter);
     }
 }
diff --git a/libs/server/Transaction/TxnKeyManager.cs b/libs/server/Transaction/TxnKeyManager.cs
@@ -145,6 +145,7 @@ private int ListObjectKeys(byte subCommand)
                 (byte)ListOperation.LINDEX => SingleKey(1, true, LockType.Shared),
                 (byte)ListOperation.LINSERT => SingleKey(1, true, LockType.Exclusive),
                 (byte)ListOperation.LREM => SingleKey(1, true, LockType.Exclusive),
+                (byte)ListOperation.LSET => SingleKey(1, true, LockType.Exclusive),
                 _ => -1
             };
         }

diff --git a/test/Garnet.test/RespListTests.cs b/test/Garnet.test/RespListTests.cs
@@ -593,6 +593,27 @@ public void CanDoLRANGEcorrect()
             Assert.IsTrue(result[1].ToString().Equals("g"));
         }
 
+        [Test]
+        public void CanDoLSETbasic()
+        {
+            using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig());
+            var db = redis.GetDatabase(0);
+
+            var key = "mylist";
+            var values = new RedisValue[] { "one", "two", "three" };
+            var pushResult = db.ListRightPush(key, values);
+            Assert.AreEqual(3, pushResult);
+
+            db.ListSetByIndex(key, 0, "four");
+            db.ListSetByIndex(key, -2, "five");
+
+            var result = db.ListRange(key, 0, -1);
+            var strResult = result.Select(r => r.ToString()).ToArray();
+            Assert.AreEqual(3, result.Length);
+            var expected = new[] { "four", "five", "three" };
+            Assert.IsTrue(expected.SequenceEqual(strResult));
+        }
+
         #region GarnetClientTests
 
         [Test]
@@ -748,7 +769,6 @@ public async Task CanUseLMoveWithCancellationTokenGC()
 
             tokenSource.Dispose();
         }
-
         #endregion
 
         #region LightClientTests
@@ -869,6 +889,66 @@ public void CanSendErrorInWrongTypeLC()
             Assert.AreEqual(expectedResponse, actualValue);
         }
 
+        [Test]
+        public void CanDoLSETbasicLC()
+        {
+            using var lightClientRequest = TestUtils.CreateRequest();
+            _ = lightClientRequest.SendCommand("RPUSH mylist one two three");
+            var response = lightClientRequest.SendCommand("LSET mylist 0 four");
+            var expectedResponse = "+OK\r\n";
+            var actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
+            Assert.AreEqual(expectedResponse, actualValue);
+        }
+
+        [Test]
+        public void CanReturnErrorLSETWhenNosuchkey()
+        {
+            using var lightClientRequest = TestUtils.CreateRequest();
+            var response = lightClientRequest.SendCommand("LSET mylist 0 four");
+            var expectedResponse = "-ERR no such key\r\n";
+            var actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
+            Assert.AreEqual(expectedResponse, actualValue);
+        }
+
+        [Test]
+        public void CanReturnErrorLSETWhenIndexNotInteger()
+        {
+            using var lightClientRequest = TestUtils.CreateRequest();
+            _ = lightClientRequest.SendCommand("RPUSH mylist one two three");
+            var response = lightClientRequest.SendCommand("LSET mylist a four");
+            var expectedResponse = "-ERR value is not an integer or out of range.\r\n";
+            var actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
+            Assert.AreEqual(expectedResponse, actualValue);
+        }
+
+        [Test]
+        public void CanReturnErrorLSETWhenIndexOutRange()
+        {
+            using var lightClientRequest = TestUtils.CreateRequest();
+            _ = lightClientRequest.SendCommand("RPUSH mylist one two three");
+            var response = lightClientRequest.SendCommand("LSET mylist 10 four");
+            // 
+            var expectedResponse = "-ERR index out of range";
+            var actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
+            Assert.AreEqual(expectedResponse, actualValue);
+
+            response = lightClientRequest.SendCommand("LSET mylist -100 four");
+            expectedResponse = "-ERR index out of range";
+            actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
+            Assert.AreEqual(expectedResponse, actualValue);
+        }
+
+        [Test]
+        public void CanReturnErrorLSETWhenArgumentsWrong()
+        {
+            using var lightClientRequest = TestUtils.CreateRequest();
+            _ = lightClientRequest.SendCommand("RPUSH mylist one two three");
+            var response = lightClientRequest.SendCommand("LSET mylist a");
+            var expectedResponse = "-ERR wrong number of arguments for 'LSET'";
+            var actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
+            Assert.AreEqual(expectedResponse, actualValue);
+        }
+
         #endregion
 
         [Test]