Skip to content

Commit

Permalink
[Compatibility] Added PUBSUB CHANNELS, NUMPAT and NUMSUB commands (#719)
Browse files Browse the repository at this point in the history 
* Added PUBSUB CHANNELS, NUMPAT and NUMSUB commands

* Code style fix

* Review comment fix

* Fixed build issue

* Review command fix

* Review command fix

---------

Co-authored-by: Tal Zaccai <[email protected]>
  • Loading branch information
@Vijay-Nirmal @TalZaccai
Vijay-Nirmal and TalZaccai authored Oct 28, 2024
1 parent 4e3b462 commit 6d489c6
Show file tree
Hide file tree
Showing 12 changed files with 580 additions and 24 deletions.
48 changes: 48 additions & 0 deletions libs/resources/RespCommandsDocs.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
29 changes: 29 additions & 0 deletions libs/resources/RespCommandsInfo.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
158 changes: 158 additions & 0 deletions libs/server/PubSub/SubscribeBroker.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -408,6 +410,162 @@ 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.WriteBulkString(key.AsSpan().Slice(sizeof(int)), ref curr, end))
ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
}
return;
}

// Below WriteArrayLength is primarily to move the start of the buffer to the max length that is required to write the array length. The actual length is written in the below line.
// This is done to avoid multiple two passes over the subscriptions or new array allocation if we use single pass over the subscriptions
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 endKeyPtr = keyPtr;
var _patternPtr = patternPtr;
if (keySerializer.Match(ref keySerializer.ReadKeyByRef(ref endKeyPtr), 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;
}

// Below code is to write the actual array length in the buffer
// And move the array elements to the start of the new array length if new array length is less than the max array length that we orginally write in the above line
var newTotalArrayHeaderLen = 0;
var _ptr = ptr;
// 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
_ = RespWriteUtils.WriteArrayLength(noOfFoundChannels, ref _ptr, end, out var _, out newTotalArrayHeaderLen);

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 = 0;
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.WriteBulkString(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();

subscriptions.TryGetValue(channel, out var subscriptionDict);
while (!RespWriteUtils.WriteInteger(subscriptionDict is null ? 0 : subscriptionDict.Count, 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()
{
Expand Down
5 changes: 5 additions & 0 deletions libs/server/Resp/CmdStrings.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -207,6 +211,7 @@ static partial class CmdStrings
public const string GenericParamShouldBeGreaterThanZero = "ERR {0} should be greater than 0";
public const string GenericUnknownClientType = "ERR Unknown client type '{0}'";
public const string GenericErrDuplicateFilter = "ERR Filter '{0}' defined multiple times";
public const string GenericPubSubCommandDisabled = "ERR {0} is disabled, enable it with --pubsub option.";

/// <summary>
/// Response errors while scripting
Expand Down
34 changes: 34 additions & 0 deletions libs/server/Resp/Parser/RespCommand.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,11 @@ public enum RespCommand : ushort

PING,

// Pub/Sub commands
PUBSUB,
PUBSUB_CHANNELS,
PUBSUB_NUMPAT,
PUBSUB_NUMSUB,
PUBLISH,
SUBSCRIBE,
PSUBSCRIBE,
Expand Down Expand Up @@ -1967,6 +1972,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
Expand Down
Loading

0 comments on commit 6d489c6

Please sign in to comment.