Skip to content

Commit

Permalink
Custom Command Registration Updates (#410)
Browse files Browse the repository at this point in the history 
* wip

* wip

* wip

* dotnet format

* bugfix

* Small fix to GarnetServer main

---------

Co-authored-by: Badrish Chandramouli <[email protected]>
  • Loading branch information
@TalZaccai @badrishc
TalZaccai and badrishc authored May 30, 2024
1 parent b34fc99 commit 4991f92
Show file tree
Hide file tree
Showing 12 changed files with 332 additions and 210 deletions.
3 changes: 2 additions & 1 deletion libs/server/Custom/CustomCommandManager.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,13 +48,14 @@ internal int Register(string name, int numParams, CommandType type, CustomRawStr
return id;
}

internal int Register(string name, int numParams, Func<CustomTransactionProcedure> proc)
internal int Register(string name, int numParams, Func<CustomTransactionProcedure> proc, RespCommandsInfo commandInfo = null)
{
int id = Interlocked.Increment(ref TransactionProcId) - 1;
if (id >= MaxRegistrations)
throw new Exception("Out of registration space");

transactionProcMap[id] = new CustomTransaction(name, (byte)id, numParams, proc);
if (commandInfo != null) customCommandsInfo.Add(name, commandInfo);
return id;
}

Expand Down
2 changes: 1 addition & 1 deletion libs/server/Custom/CustomCommandRegistration.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -212,7 +212,7 @@ public RegisterCustomTransactionProcedureProvider(CustomTransactionProcedure ins

public override void Register(CustomCommandManager customCommandManager)
{
customCommandManager.Register(this.RegisterArgs.Name, this.RegisterArgs.NumParams, () => this.Instance);
customCommandManager.Register(this.RegisterArgs.Name, this.RegisterArgs.NumParams, () => this.Instance, this.RegisterArgs.CommandInfo);
}
}
}
83 changes: 77 additions & 6 deletions libs/server/Resp/ArrayCommands.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -220,18 +220,51 @@ private bool NetworkMGET_SG<TGarnetApi>(int count, byte* ptr, ref TGarnetApi sto
/// Register all custom commands / transactions
/// </summary>
/// <param name="binaryPaths">Binary paths from which to load assemblies</param>
/// <param name="cmdInfoPath">Path of JSON file containing RespCommandsInfo for custom commands</param>
/// <param name="classNameToRegisterArgs">Mapping between class names to register and arguments required for registration</param>
/// <param name="customCommandManager">CustomCommandManager instance used to register commands</param>
/// <param name="errorMessage">If method returned false, contains ASCII encoded generic error string; otherwise <c>default</c></param>
/// <returns>A boolean value indicating whether registration of the custom commands was successful.</returns>
private bool TryRegisterCustomCommands(
IEnumerable<string> binaryPaths,
string cmdInfoPath,
Dictionary<string, List<RegisterArgsBase>> classNameToRegisterArgs,
CustomCommandManager customCommandManager,
out ReadOnlySpan<byte> errorMessage)
{
errorMessage = default;
var classInstances = new Dictionary<string, object>();
IReadOnlyDictionary<string, RespCommandsInfo> cmdNameToInfo = new Dictionary<string, RespCommandsInfo>();

if (cmdInfoPath != null)
{
// Check command info path, if specified
if (!File.Exists(cmdInfoPath))
{
errorMessage = CmdStrings.RESP_ERR_GENERIC_GETTING_CMD_INFO_FILE;
return false;
}

// Check command info path is in allowed paths
if (storeWrapper.serverOptions.ExtensionBinPaths.All(p => !FileUtils.IsFileInDirectory(cmdInfoPath, p)))
{
errorMessage = CmdStrings.RESP_ERR_GENERIC_CMD_INFO_FILE_NOT_IN_ALLOWED_PATHS;
return false;
}

var streamProvider = StreamProviderFactory.GetStreamProvider(FileLocationType.Local);
var commandsInfoProvider = RespCommandsInfoProviderFactory.GetRespCommandsInfoProvider();

var importSucceeded = commandsInfoProvider.TryImportRespCommandsInfo(cmdInfoPath,
streamProvider, out cmdNameToInfo, logger);

if (!importSucceeded)
{
errorMessage = CmdStrings.RESP_ERR_GENERIC_MALFORMED_COMMAND_INFO_JSON;
return false;
}
}

// Get all binary file paths from inputs binary paths
if (!FileUtils.TryGetFiles(binaryPaths, out var files, out _, [".dll", ".exe"],
SearchOption.AllDirectories))
Expand Down Expand Up @@ -317,6 +350,12 @@ private bool TryRegisterCustomCommands(
{
foreach (var args in classNameToArgs.Value)
{
// Add command info to register arguments, if exists
if (cmdNameToInfo.ContainsKey(args.Name))
{
args.CommandInfo = cmdNameToInfo[args.Name];
}

var registerApi =
RegisterCustomCommandProviderFactory.GetRegisterCustomCommandProvider(classInstances[classNameToArgs.Key], args);

Expand Down Expand Up @@ -349,25 +388,30 @@ private bool NetworkREGISTERCS(int count, byte* ptr, CustomCommandManager custom
{
var leftTokens = count;
var readPathsOnly = false;
var optionalParamsRead = 0;

var binaryPaths = new HashSet<string>();
string cmdInfoPath = default;

// Custom class name to arguments read from each sub-command
var classNameToRegisterArgs = new Dictionary<string, List<RegisterArgsBase>>();

ReadOnlySpan<byte> errorMsg = null;

if (leftTokens == 0)
if (leftTokens < 6)
errorMsg = CmdStrings.RESP_ERR_GENERIC_MALFORMED_REGISTERCS_COMMAND;

// Parse the REGISTERCS command - list of registration sub-commands followed by a list of paths to binary files / folders
// Syntax - REGISTERCS cmdType name numParams className [expTicks] [cmdType name numParams className [expTicks] ...] SRC path [path ...]
// Parse the REGISTERCS command - list of registration sub-commands
// followed by an optional path to JSON file containing an array of RespCommandsInfo objects,
// followed by a list of paths to binary files / folders
// Syntax - REGISTERCS cmdType name numParams className [expTicks] [cmdType name numParams className [expTicks] ...]
// [INFO path] SRC path [path ...]
RegisterArgsBase args = null;

while (leftTokens > 0)
{
byte* firstTokenPtr = null;
int firstTokenSize = 0;
var firstTokenSize = 0;

// Read first token of current sub-command or path
if (!RespReadUtils.ReadPtrWithLengthHeader(ref firstTokenPtr, ref firstTokenSize, ref ptr, recvBufferPtr + bytesRead))
Expand Down Expand Up @@ -395,6 +439,22 @@ private bool NetworkREGISTERCS(int count, byte* ptr, CustomCommandManager custom
{
args = new RegisterTxnArgs();
}
else if (tokenSpan.SequenceEqual(CmdStrings.INFO) ||
tokenSpan.SequenceEqual(CmdStrings.info))
{
// If first token is not a cmdType and no other sub-command is previously defined, command is malformed
if (classNameToRegisterArgs.Count == 0 || leftTokens == 0)
{
errorMsg = CmdStrings.RESP_ERR_GENERIC_MALFORMED_REGISTERCS_COMMAND;
break;
}

if (!RespReadUtils.ReadStringWithLengthHeader(out cmdInfoPath, ref ptr, recvBufferPtr + bytesRead))
return false;

leftTokens--;
continue;
}
else if (readPathsOnly || (tokenSpan.SequenceEqual(CmdStrings.SRC) ||
tokenSpan.SequenceEqual(CmdStrings.src)))
{
Expand All @@ -419,10 +479,11 @@ private bool NetworkREGISTERCS(int count, byte* ptr, CustomCommandManager custom
else
{
// Check optional parameters for previous sub-command
if (args is RegisterCmdArgs cmdArgs)
if (optionalParamsRead == 0 && args is RegisterCmdArgs cmdArgs)
{
var expTicks = NumUtils.BytesToLong(tokenSpan);
cmdArgs.ExpirationTicks = expTicks;
optionalParamsRead++;
continue;
}

Expand All @@ -431,6 +492,16 @@ private bool NetworkREGISTERCS(int count, byte* ptr, CustomCommandManager custom
break;
}

optionalParamsRead = 0;

// At this point we expect at least 6 remaining tokens -
// 3 more tokens for command definition + 2 for source definition
if (leftTokens < 5)
{
errorMsg = CmdStrings.RESP_ERR_GENERIC_MALFORMED_REGISTERCS_COMMAND;
break;
}

// Start reading the sub-command arguments
// Read custom command name
if (!RespReadUtils.ReadStringWithLengthHeader(out var name, ref ptr, recvBufferPtr + bytesRead))
Expand Down Expand Up @@ -469,7 +540,7 @@ private bool NetworkREGISTERCS(int count, byte* ptr, CustomCommandManager custom

// If no error is found, continue to try register custom commands in the server
if (errorMsg == null &&
TryRegisterCustomCommands(binaryPaths, classNameToRegisterArgs, customCommandManager, out errorMsg))
TryRegisterCustomCommands(binaryPaths, cmdInfoPath, classNameToRegisterArgs, customCommandManager, out errorMsg))
{
while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_OK, ref dcurr, dend))
SendAndReset();
Expand Down
3 changes: 3 additions & 0 deletions libs/server/Resp/CmdStrings.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,8 +121,11 @@ static partial class CmdStrings
public static ReadOnlySpan<byte> RESP_ERR_GENERIC_OFFSETOUTOFRANGE => "ERR offset is out of range"u8;
public static ReadOnlySpan<byte> RESP_ERR_GENERIC_CURSORVALUE => "ERR cursor value should be equal or greater than 0."u8;
public static ReadOnlySpan<byte> RESP_ERR_GENERIC_MALFORMED_REGISTERCS_COMMAND => "ERR malformed REGISTERCS command."u8;
public static ReadOnlySpan<byte> RESP_ERR_GENERIC_MALFORMED_COMMAND_INFO_JSON => "ERR malformed command info JSON."u8;
public static ReadOnlySpan<byte> RESP_ERR_GENERIC_GETTING_BINARY_FILES => "ERR unable to access one or more binary files."u8;
public static ReadOnlySpan<byte> RESP_ERR_GENERIC_GETTING_CMD_INFO_FILE => "ERR unable to access command info file."u8;
public static ReadOnlySpan<byte> RESP_ERR_GENERIC_BINARY_FILES_NOT_IN_ALLOWED_PATHS => "ERR one or more binary file are not contained in allowed paths."u8;
public static ReadOnlySpan<byte> RESP_ERR_GENERIC_CMD_INFO_FILE_NOT_IN_ALLOWED_PATHS => "ERR command info file is not contained in allowed paths."u8;
public static ReadOnlySpan<byte> RESP_ERR_GENERIC_LOADING_ASSEMBLIES => "ERR unable to load one or more assemblies."u8;
public static ReadOnlySpan<byte> RESP_ERR_GENERIC_ASSEMBLY_NOT_SIGNED => "ERR one or more assemblies loaded is not digitally signed."u8;
public static ReadOnlySpan<byte> RESP_ERR_GENERIC_INSTANTIATING_CLASS => "ERR unable to instantiate one or more classes from given assemblies."u8;
Expand Down
10 changes: 6 additions & 4 deletions libs/server/Servers/RegisterApi.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,18 +44,20 @@ public int NewCommand(string name, int numParams, CommandType type, CustomRawStr
/// <param name="name">Name of command</param>
/// <param name="numParams">Number of parameters</param>
/// <param name="proc">Custom stored procedure</param>
/// <param name="commandInfo">RESP command info</param>
/// <returns>ID of the registered command</returns>
public int NewTransactionProc(string name, int numParams, Func<CustomTransactionProcedure> proc)
=> provider.StoreWrapper.customCommandManager.Register(name, numParams, proc);
public int NewTransactionProc(string name, int numParams, Func<CustomTransactionProcedure> proc, RespCommandsInfo commandInfo = null)
=> provider.StoreWrapper.customCommandManager.Register(name, numParams, proc, commandInfo);

/// <summary>
/// Register transaction procedure with Garnet, with a variable number of parameters
/// </summary>
/// <param name="name">Name of command</param>
/// <param name="proc">Custom stored procedure</param>
/// <param name="commandInfo">RESP command info</param>
/// <returns>ID of the registered command</returns>
public int NewTransactionProc(string name, Func<CustomTransactionProcedure> proc)
=> provider.StoreWrapper.customCommandManager.Register(name, int.MaxValue, proc);
public int NewTransactionProc(string name, Func<CustomTransactionProcedure> proc, RespCommandsInfo commandInfo = null)
=> provider.StoreWrapper.customCommandManager.Register(name, int.MaxValue, proc, commandInfo);

/// <summary>
/// Register object type with server
Expand Down
45 changes: 26 additions & 19 deletions main/GarnetServer/Program.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,7 @@
// Licensed under the MIT license.

using System;
using System.Collections.Generic;
using System.Threading;
using Garnet.common;
using Garnet.server;

namespace Garnet
Expand All @@ -21,11 +19,7 @@ static void Main(string[] args)
using var server = new GarnetServer(args);

// Optional: register custom extensions
if (!TryRegisterExtensions(server))
{
Console.WriteLine("Unable to register server extensions.");
return;
}
RegisterExtensions(server);

// Start the server
server.Start();
Expand All @@ -43,10 +37,21 @@ static void Main(string[] args)
/// commands such as db.Execute in StackExchange.Redis. Example:
/// db.Execute("SETIFPM", key, value, prefix);
/// </summary>
static bool TryRegisterExtensions(GarnetServer server)
static void RegisterExtensions(GarnetServer server)
{
// Register custom command on raw strings (SETIFPM = "set if prefix match")
server.Register.NewCommand("SETIFPM", 2, CommandType.ReadModifyWrite, new SetIfPMCustomCommand());
// Add RESP command info to registration for command to appear when client runs COMMAND / COMMAND INFO
var setIfPmCmdInfo = new RespCommandsInfo
{
Name = "SETIFPM",
Arity = 4,
FirstKey = 1,
LastKey = 1,
Step = 1,
Flags = RespCommandFlags.DenyOom | RespCommandFlags.Write,
AclCategories = RespAclCategories.String | RespAclCategories.Write,
};
server.Register.NewCommand("SETIFPM", 2, CommandType.ReadModifyWrite, new SetIfPMCustomCommand(), setIfPmCmdInfo);

// Register custom command on raw strings (SETWPIFPGT = "set with prefix, if prefix greater than")
server.Register.NewCommand("SETWPIFPGT", 2, CommandType.ReadModifyWrite, new SetWPIFPGTCustomCommand());
Expand All @@ -60,7 +65,18 @@ static bool TryRegisterExtensions(GarnetServer server)
server.Register.NewCommand("MYDICTGET", 1, CommandType.Read, factory);

// Register stored procedure to run a transactional command
server.Register.NewTransactionProc("READWRITETX", 3, () => new ReadWriteTxn());
// Add RESP command info to registration for command to appear when client runs COMMAND / COMMAND INFO
var readWriteTxCmdInfo = new RespCommandsInfo
{
Name = "READWRITETX",
Arity = 4,
FirstKey = 1,
LastKey = 3,
Step = 1,
Flags = RespCommandFlags.DenyOom | RespCommandFlags.Write,
AclCategories = RespAclCategories.Write,
};
server.Register.NewTransactionProc("READWRITETX", 3, () => new ReadWriteTxn(), readWriteTxCmdInfo);

// Register stored procedure to run a transactional command
server.Register.NewTransactionProc("MSETPX", () => new MSetPxTxn());
Expand All @@ -74,15 +90,6 @@ static bool TryRegisterExtensions(GarnetServer server)
// Register sample transactional procedures
server.Register.NewTransactionProc("SAMPLEUPDATETX", 8, () => new SampleUpdateTxn());
server.Register.NewTransactionProc("SAMPLEDELETETX", 5, () => new SampleDeleteTxn());

return true;
}

private static bool TryGetRespCommandsInfo(string path, out IReadOnlyDictionary<string, RespCommandsInfo> commandsInfo)
{
var streamProvider = StreamProviderFactory.GetStreamProvider(FileLocationType.Local);
var commandsInfoProvider = RespCommandsInfoProviderFactory.GetRespCommandsInfoProvider();
return commandsInfoProvider.TryImportRespCommandsInfo(path, streamProvider, out commandsInfo);
}
}
}
1 change: 1 addition & 0 deletions test/Garnet.test.cluster/Garnet.test.cluster.csproj
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@

<ItemGroup>
<PackageReference Include="CommandLineParser" />
<PackageReference Include="Microsoft.CodeAnalysis" />
<PackageReference Include="Microsoft.Extensions.Configuration.Json" />
<PackageReference Include="Microsoft.Extensions.Configuration.Binder" />
<PackageReference Include="Microsoft.NET.Test.Sdk" />
Expand Down
18 changes: 9 additions & 9 deletions test/Garnet.test/CustomRespCommandsInfo.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,13 @@
{
"Command": "NONE",
"ArrayCommand": null,
"Name": "DELIFM",
"Arity": 3,
"Flags": "Write",
"Name": "MGETIFPM",
"Arity": -3,
"Flags": "ReadOnly",
"FirstKey": 1,
"LastKey": 1,
"LastKey": -1,
"Step": 1,
"AclCategories": "KeySpace, String, Write",
"AclCategories": "Read",
"Tips": null,
"KeySpecifications": null,
"SubCommands": null
Expand Down Expand Up @@ -44,21 +44,21 @@
{
"Command": "NONE",
"ArrayCommand": null,
"Name": "SETIFPM",
"Name": "READWRITETX",
"Arity": 4,
"Flags": "DenyOom, Write",
"FirstKey": 1,
"LastKey": 1,
"LastKey": 3,
"Step": 1,
"AclCategories": "String, Write",
"AclCategories": "Write",
"Tips": null,
"KeySpecifications": null,
"SubCommands": null
},
{
"Command": "NONE",
"ArrayCommand": null,
"Name": "SETWPIFPGT",
"Name": "SETIFPM",
"Arity": 4,
"Flags": "DenyOom, Write",
"FirstKey": 1,
Expand Down
Loading

0 comments on commit 4991f92

Please sign in to comment.