[Compatibility] Added PUBSUB CHANNELS, NUMPAT and NUMSUB commands

libs/resources/RespCommandsDocs.json

@@ -3858,6 +3858,54 @@
       }
     ]
   },
+  {
+    "Command": "PUBSUB",
+    "Name": "PUBSUB",
+    "Summary": "A container for Pub/Sub commands.",
+    "Group": "PubSub",
+    "Complexity": "Depends on subcommand.",
+    "SubCommands": [
+      {
+        "Command": "PUBSUB_NUMPAT",
+        "Name": "PUBSUB|NUMPAT",
+        "Summary": "Returns a count of unique pattern subscriptions.",
+        "Group": "PubSub",
+        "Complexity": "O(1)"
+      },
+      {
+        "Command": "PUBSUB_CHANNELS",
+        "Name": "PUBSUB|CHANNELS",
+        "Summary": "Returns the active channels.",
+        "Group": "PubSub",
+        "Complexity": "O(N) where N is the number of active channels, and assuming constant time pattern matching (relatively short channels and patterns)",
+        "Arguments": [
+          {
+            "TypeDiscriminator": "RespCommandBasicArgument",
+            "Name": "PATTERN",
+            "DisplayText": "pattern",
+            "Type": "Pattern",
+            "ArgumentFlags": "Optional"
+          }
+        ]
+      },
+      {
+        "Command": "PUBSUB_NUMSUB",
+        "Name": "PUBSUB|NUMSUB",
+        "Summary": "Returns a count of subscribers to channels.",
+        "Group": "PubSub",
+        "Complexity": "O(N) for the NUMSUB subcommand, where N is the number of requested channels",
+        "Arguments": [
+          {
+            "TypeDiscriminator": "RespCommandBasicArgument",
+            "Name": "CHANNEL",
+            "DisplayText": "channel",
+            "Type": "String",
+            "ArgumentFlags": "Optional, Multiple"
+          }
+        ]
+      }
+    ]
+  },
   {
     "Command": "PUNSUBSCRIBE",
     "Name": "PUNSUBSCRIBE",

libs/resources/RespCommandsInfo.json

@@ -2805,6 +2805,35 @@
     "Flags": "Fast, Loading, PubSub, Stale",
     "AclCategories": "Fast, PubSub"
   },
+  {
+    "Command": "PUBSUB",
+    "Name": "PUBSUB",
+    "Arity": -2,
+    "AclCategories": "Slow",
+    "SubCommands": [
+      {
+        "Command": "PUBSUB_CHANNELS",
+        "Name": "PUBSUB|CHANNELS",
+        "Arity": -2,
+        "Flags": "Loading, PubSub, Stale",
+        "AclCategories": "PubSub, Slow"
+      },
+      {
+        "Command": "PUBSUB_NUMPAT",
+        "Name": "PUBSUB|NUMPAT",
+        "Arity": 2,
+        "Flags": "Loading, PubSub, Stale",
+        "AclCategories": "PubSub, Slow"
+      },
+      {
+        "Command": "PUBSUB_NUMSUB",
+        "Name": "PUBSUB|NUMSUB",
+        "Arity": -2,
+        "Flags": "Loading, PubSub, Stale",
+        "AclCategories": "PubSub, Slow"
+      }
+    ]
+  },
   {
     "Command": "PUNSUBSCRIBE",
     "Name": "PUNSUBSCRIBE",

libs/server/PubSub/SubscribeBroker.cs

@@ -2,8 +2,10 @@
 // Licensed under the MIT license.
 
 using System;
+using System.Buffers;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 using Garnet.common;
@@ -408,6 +410,164 @@ public unsafe void Publish(byte* key, byte* value, int valueLength, bool ascii =
             log.Enqueue(logEntryBytes);
         }
 
+        /// <summary>
+        /// Retrieves the collection of channels that have active subscriptions.
+        /// </summary>
+        /// <returns>The collection of channels.</returns>
+        public unsafe void Channels(ref ObjectInput input, ref SpanByteAndMemory output)
+        {
+            var isMemory = false;
+            MemoryHandle ptrHandle = default;
+            var ptr = output.SpanByte.ToPointer();
+
+            var curr = ptr;
+            var end = curr + output.Length;
+
+            try
+            {
+                if (subscriptions is null || subscriptions.Count == 0)
+                {
+                    while (!RespWriteUtils.WriteEmptyArray(ref curr, end))
+                        ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
+                    return;
+                }
+
+                if (input.parseState.Count == 0)
+                {
+                    while (!RespWriteUtils.WriteArrayLength(subscriptions.Count, ref curr, end))
+                        ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
+
+                    foreach (var key in subscriptions.Keys)
+                    {
+                        while (!RespWriteUtils.WriteSimpleString(key.AsSpan().Slice(sizeof(int)), ref curr, end))
+                            ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
+                    }
+                    return;
+                }
+
+                var totalArrayHeaderLen = 0;
+                while (!RespWriteUtils.WriteArrayLength(subscriptions.Count, ref curr, end, out var _, out totalArrayHeaderLen))
+                    ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
+
+                var noOfFoundChannels = 0;
+                var pattern = input.parseState.GetArgSliceByRef(0).SpanByte;
+                var patternPtr = pattern.ToPointer() - sizeof(int);
+                *(int*)patternPtr = pattern.Length;
+
+                foreach (var key in subscriptions.Keys)
+                {
+                    fixed (byte* keyPtr = key)
+                    {
+                        var refKeyPtr = keyPtr;
+                        var _patternPtr = patternPtr;
+                        if (keySerializer.Match(ref keySerializer.ReadKeyByRef(ref refKeyPtr), true, ref keySerializer.ReadKeyByRef(ref _patternPtr), true))
+                        {
+                            while (!RespWriteUtils.WriteSimpleString(key.AsSpan().Slice(sizeof(int)), ref curr, end))
+                                ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
+                            noOfFoundChannels++;
+                        }
+                    }
+                }
+
+                if (noOfFoundChannels == 0)
+                {
+                    curr = ptr;
+                    while (!RespWriteUtils.WriteEmptyArray(ref curr, end))
+                        ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
+                    return;
+                }
+
+                var newTotalArrayHeaderLen = 0;
+                var _ptr = ptr;
+                RespWriteUtils.WriteArrayLength(noOfFoundChannels, ref _ptr, end, out var _, out newTotalArrayHeaderLen); // ReallocateOutput is not needed here as there should be always be available space in the output buffer as we have already written the max array length
+
+                Debug.Assert(totalArrayHeaderLen >= newTotalArrayHeaderLen, "newTotalArrayHeaderLen can't be bigger than totalArrayHeaderLen as we have already written max array lenght in the buffer");
+                if (totalArrayHeaderLen != newTotalArrayHeaderLen)
+                {
+                    var remainingLength = (curr - ptr) - totalArrayHeaderLen;
+                    Buffer.MemoryCopy(ptr + totalArrayHeaderLen, ptr + newTotalArrayHeaderLen, remainingLength, remainingLength);
+                    curr = curr - (totalArrayHeaderLen - newTotalArrayHeaderLen);
+                }
+            }
+            finally
+            {
+                if (isMemory)
+                    ptrHandle.Dispose();
+                output.Length = (int)(curr - ptr);
+            }
+        }
+
+        /// <summary>
+        /// Retrieves the number of pattern subscriptions.
+        /// </summary>
+        /// <returns>The number of pattern subscriptions.</returns>
+        public int NumPatternSubscriptions()
+        {
+            return prefixSubscriptions?.Count ?? 0;
+        }
+
+        /// <summary>
+        /// PUBSUB NUMSUB
+        /// </summary>
+        /// <param name="output"></param>
+        /// <returns></returns>
+        public unsafe void NumSubscriptions(ref ObjectInput input, ref SpanByteAndMemory output)
+        {
+            var isMemory = false;
+            MemoryHandle ptrHandle = default;
+            var ptr = output.SpanByte.ToPointer();
+
+            var curr = ptr;
+            var end = curr + output.Length;
+
+            try
+            {
+                var numOfChannels = input.parseState.Count;
+                if (subscriptions is null || numOfChannels == 0)
+                {
+                    while (!RespWriteUtils.WriteEmptyArray(ref curr, end))
+                        ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
+                    return;
+                }
+
+                while (!RespWriteUtils.WriteArrayLength(numOfChannels * 2, ref curr, end))
+                    ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
+
+                var currChannelIdx = input.parseStateFirstArgIdx;
+                while (currChannelIdx < numOfChannels)
+                {
+                    var channelArg = input.parseState.GetArgSliceByRef(currChannelIdx);
+                    var channelSpan = channelArg.SpanByte;
+                    var channelPtr = channelSpan.ToPointer() - sizeof(int);  // Memory would have been already pinned
+                    *(int*)channelPtr = channelSpan.Length;
+
+                    while (!RespWriteUtils.WriteSimpleString(channelArg.ReadOnlySpan, ref curr, end))
+                        ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
+
+                    var channel = new Span<byte>(channelPtr, channelSpan.Length + sizeof(int)).ToArray();
+
+                    if (subscriptions.TryGetValue(channel, out var subscriptionDict))
+                    {
+                        while (!RespWriteUtils.WriteInteger(subscriptionDict.Count, ref curr, end))
+                            ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
+                    }
+                    else
+                    {
+                        while (!RespWriteUtils.WriteInteger(0, ref curr, end))
+                            ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
+                    }
+
+                    currChannelIdx++;
+                }
+            }
+            finally
+            {
+                if (isMemory)
+                    ptrHandle.Dispose();
+                output.Length = (int)(curr - ptr);
+            }
+        }
+
         /// <inheritdoc />
         public void Dispose()
         {

libs/server/Resp/CmdStrings.cs

@@ -104,6 +104,10 @@ static partial class CmdStrings
         public static ReadOnlySpan<byte> rank => "rank"u8;
         public static ReadOnlySpan<byte> MAXLEN => "MAXLEN"u8;
         public static ReadOnlySpan<byte> maxlen => "maxlen"u8;
+        public static ReadOnlySpan<byte> PUBSUB => "PUBSUB"u8;
+        public static ReadOnlySpan<byte> CHANNELS => "CHANNELS"u8;
+        public static ReadOnlySpan<byte> NUMPAT => "NUMPAT"u8;
+        public static ReadOnlySpan<byte> NUMSUB => "NUMSUB"u8;
 
         /// <summary>
         /// Response strings

libs/server/Resp/Parser/RespCommand.cs

@@ -178,6 +178,11 @@ public enum RespCommand : ushort
 
         PING,
 
+        // Pub/Sub commands
+        PUBSUB,
+        PUBSUB_CHANNELS,
+        PUBSUB_NUMPAT,
+        PUBSUB_NUMSUB,
         PUBLISH,
         SUBSCRIBE,
         PSUBSCRIBE,
@@ -1962,6 +1967,35 @@ private RespCommand SlowParseCommand(ref int count, ref ReadOnlySpan<byte> speci
                     return RespCommand.MODULE_LOADCS;
                 }
             }
+            else if (command.SequenceEqual(CmdStrings.PUBSUB))
+            {
+                Span<byte> subCommand = GetCommand(out var gotSubCommand);
+                if (!gotSubCommand)
+                {
+                    success = false;
+                    return RespCommand.NONE;
+                }
+
+                count--;
+                AsciiUtils.ToUpperInPlace(subCommand);
+                if (subCommand.SequenceEqual(CmdStrings.CHANNELS))
+                {
+                    return RespCommand.PUBSUB_CHANNELS;
+                }
+                else if (subCommand.SequenceEqual(CmdStrings.NUMSUB))
+                {
+                    return RespCommand.PUBSUB_NUMSUB;
+                }
+                else if (subCommand.SequenceEqual(CmdStrings.NUMPAT))
+                {
+                    return RespCommand.PUBSUB_NUMPAT;
+                }
+                else
+                {
+                    success = false;
+                    return RespCommand.NONE;
+                }
+            }
             else
             {
                 // Custom commands should have never been set when we reach this point