add COMMAND COUNT

microsoft · Mar 29, 2024 · 8b4556b · 8b4556b
1 parent 03c50ca
commit 8b4556b
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 4 deletions.
diff --git a/libs/server/Resp/AdminCommands.cs b/libs/server/Resp/AdminCommands.cs
@@ -276,10 +276,20 @@ private bool ProcessAdminCommands<TGarnetApi>(ReadOnlySpan<byte> command, ReadOn
             {
                 if (count != 1)
                 {
-                    if (!DrainCommands(bufSpan, count - 1))
-                        return false;
-                    errorFlag = true;
-                    errorCmd = "command";
+                    var param = GetCommand(bufSpan, out var success);
+                    if (!success) return false;
+                    if (Encoding.ASCII.GetString(param) == "COUNT")
+                    {
+                        while (!RespWriteUtils.WriteDirect(Encoding.ASCII.GetBytes($":{RespInfo.GetCommandsCount()}\r\n"), ref dcurr, dend))
+                            SendAndReset();
+                    }
+                    else
+                    {
+                        if (!DrainCommands(bufSpan, count - 1))
+                            return false;
+                        errorFlag = true;
+                        errorCmd = "command";
+                    }
                 }
                 else
                 {

diff --git a/libs/server/Resp/RespInfo.cs b/libs/server/Resp/RespInfo.cs
@@ -10,6 +10,12 @@ namespace Garnet.server
     /// </summary>
     public static class RespInfo
     {
+        private static readonly ushort _commandsCount;
+        static RespInfo()
+        {
+            _commandsCount = (ushort)GetCommands().Count;
+        }
+
         /// <summary>
         /// Get set of RESP commands supported by Garnet server
         /// </summary>
@@ -47,5 +53,14 @@ public static HashSet<string> GetCommands()
                 "WATCH", "UNWATCH", "MULTI", "EXEC", "DISCARD",
             };
         }
+
+        /// <summary>
+        /// Get RESP commands count supported by Garnet server
+        /// </summary>
+        /// <returns></returns>
+        public static ushort GetCommandsCount()
+        {
+            return _commandsCount;
+        }
     }
 }
diff --git a/test/Garnet.test/RespAdminCommandsTests.cs b/test/Garnet.test/RespAdminCommandsTests.cs
@@ -6,6 +6,7 @@
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using Garnet.server;
 using NUnit.Framework;
 using StackExchange.Redis;
 
@@ -130,6 +131,15 @@ public void TimeWithReturnErrorTest()
             Assert.AreEqual(expectedResponse, actualValue);
         }
 
+        [Test]
+        public void CommandCountCommandTest()
+        {
+            using var lightClientRequest = TestUtils.CreateRequest();
+            var expectedResponse = $":{RespInfo.GetCommandsCount()}\r\n";
+            var response = lightClientRequest.SendCommand("COMMAND COUNT");
+            var actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
+        }
+
         #endregion
 
         #region SeClientTests
@@ -373,6 +383,16 @@ public async Task SeFlushDBTest([Values] bool async, [Values] bool unsafetruncat
             Assert.IsTrue(_value.IsNull);
         }
 
+        [Test]
+        public async Task SeCommandCountTest()
+        {
+            using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig(allowAdmin: true));
+            var db = redis.GetDatabase(0);
+            var expectedResponse = RespInfo.GetCommandsCount() + "";
+            var actualValue = (await db.ExecuteAsync("COMMAND", "COUNT")).ToString();
+
+            Assert.AreEqual(expectedResponse, actualValue);
+        }
         #endregion
     }
 }