From 8b4556b7310162e20a016d3d5bdbe0afe63d8200 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=B0=A2=E4=BD=B3=E9=BE=99?= <tisilent@163.com>
Date: Fri, 29 Mar 2024 13:22:23 +0800
Subject: [PATCH] add COMMAND COUNT

---
 libs/server/Resp/AdminCommands.cs          | 18 ++++++++++++++----
 libs/server/Resp/RespInfo.cs               | 15 +++++++++++++++
 test/Garnet.test/RespAdminCommandsTests.cs | 20 ++++++++++++++++++++
 3 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/libs/server/Resp/AdminCommands.cs b/libs/server/Resp/AdminCommands.cs
index 8dd697d19e..353e127584 100644
--- 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
index ba35b53687..f2e94cec13 100644
--- 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;
+        }
     }
 }
\ No newline at end of file
diff --git a/test/Garnet.test/RespAdminCommandsTests.cs b/test/Garnet.test/RespAdminCommandsTests.cs
index 93725a97c1..972fe97b1a 100644
--- 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
     }
 }
\ No newline at end of file