From eabadb5f85155277fe492ce00cfe719323aba78e Mon Sep 17 00:00:00 2001 From: Tal Zaccai Date: Tue, 22 Oct 2024 19:28:00 -0700 Subject: [PATCH 01/12] Expanding RespCommand enum to ushort (#736) * wip * fix * Fixed comments * Another fix * More fixes * format * Update libs/server/Resp/Parser/RespCommand.cs Co-authored-by: Meir Blachman --------- Co-authored-by: Meir Blachman --- libs/server/ACL/CommandPermissionSet.cs | 7 ++--- libs/server/API/GarnetApi.cs | 4 +-- libs/server/API/GarnetWatchApi.cs | 2 +- libs/server/API/IGarnetApi.cs | 4 +-- libs/server/Custom/CustomCommandManager.cs | 18 +++++++----- libs/server/Custom/CustomObjectCommand.cs | 2 +- libs/server/Custom/CustomRawStringCommand.cs | 4 +-- libs/server/Custom/CustomRespCommands.cs | 4 +-- libs/server/InputHeader.cs | 8 ++--- .../Objects/Types/GarnetObjectSerializer.cs | 4 +-- libs/server/Objects/Types/GarnetObjectType.cs | 10 +++++++ libs/server/Resp/Bitmap/BitmapCommands.cs | 20 ++++++------- .../Resp/Bitmap/BitmapManagerBitfield.cs | 10 +++---- libs/server/Resp/Parser/RespCommand.cs | 29 ++++++++++++++----- libs/server/Resp/RespCommandAccessor.cs | 2 +- libs/server/Resp/RespServerSession.cs | 2 +- .../Functions/MainStore/PrivateMethods.cs | 6 ++-- .../Storage/Functions/MainStore/RMWMethods.cs | 16 +++++----- .../Functions/MainStore/ReadMethods.cs | 8 ++--- .../Functions/MainStore/VarLenInputMethods.cs | 8 ++--- .../Functions/ObjectStore/PrivateMethods.cs | 2 +- .../Functions/ObjectStore/RMWMethods.cs | 10 +++---- .../Functions/ObjectStore/ReadMethods.cs | 2 +- .../Storage/Session/MainStore/BitmapOps.cs | 14 ++++----- test/Garnet.test/RespCommandTests.cs | 2 +- test/Garnet.test/TestProcedureBitmap.cs | 4 +-- 26 files changed, 114 insertions(+), 88 deletions(-) diff --git a/libs/server/ACL/CommandPermissionSet.cs b/libs/server/ACL/CommandPermissionSet.cs index 175a1ba5d9c..fc9f0cb6947 100644 --- a/libs/server/ACL/CommandPermissionSet.cs +++ b/libs/server/ACL/CommandPermissionSet.cs @@ -167,10 +167,9 @@ public bool IsEquivalentTo(CommandPermissionSet other) /// private static ushort GetCommandListLength() { - int commandCount = (int)Enum.GetValues().Where(static cmd => cmd != RespCommand.NONE && cmd != RespCommand.INVALID).Max(); - - int neededBits = commandCount; - int neededULongs = neededBits / 64; + // # of bits needed to represent all valid commands + var neededBits = (ushort)RespCommandExtensions.LastValidCommand + 1; + var neededULongs = neededBits / 64; if ((neededBits % 64) != 0) { diff --git a/libs/server/API/GarnetApi.cs b/libs/server/API/GarnetApi.cs index 4906d6d5ed0..34c91d00a5f 100644 --- a/libs/server/API/GarnetApi.cs +++ b/libs/server/API/GarnetApi.cs @@ -333,11 +333,11 @@ public GarnetStatus StringBitPosition(ref SpanByte key, ref RawStringInput input => storageSession.StringBitPosition(ref key, ref input, ref output, ref context); /// - public GarnetStatus StringBitField(ref SpanByte key, ref RawStringInput input, byte secondaryCommand, ref SpanByteAndMemory output) + public GarnetStatus StringBitField(ref SpanByte key, ref RawStringInput input, RespCommand secondaryCommand, ref SpanByteAndMemory output) => storageSession.StringBitField(ref key, ref input, secondaryCommand, ref output, ref context); /// - public GarnetStatus StringBitFieldReadOnly(ref SpanByte key, ref RawStringInput input, byte secondaryCommand, ref SpanByteAndMemory output) + public GarnetStatus StringBitFieldReadOnly(ref SpanByte key, ref RawStringInput input, RespCommand secondaryCommand, ref SpanByteAndMemory output) => storageSession.StringBitFieldReadOnly(ref key, ref input, secondaryCommand, ref output, ref context); /// diff --git a/libs/server/API/GarnetWatchApi.cs b/libs/server/API/GarnetWatchApi.cs index ee0896ac565..ef6d82266f3 100644 --- a/libs/server/API/GarnetWatchApi.cs +++ b/libs/server/API/GarnetWatchApi.cs @@ -473,7 +473,7 @@ public GarnetStatus StringBitPosition(ref SpanByte key, ref RawStringInput input } /// - public GarnetStatus StringBitFieldReadOnly(ref SpanByte key, ref RawStringInput input, byte secondaryCommand, ref SpanByteAndMemory output) + public GarnetStatus StringBitFieldReadOnly(ref SpanByte key, ref RawStringInput input, RespCommand secondaryCommand, ref SpanByteAndMemory output) { garnetApi.WATCH(new ArgSlice(ref key), StoreType.Main); return garnetApi.StringBitFieldReadOnly(ref key, ref input, secondaryCommand, ref output); diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs index a10d0add73b..e271407c63f 100644 --- a/libs/server/API/IGarnetApi.cs +++ b/libs/server/API/IGarnetApi.cs @@ -987,7 +987,7 @@ public interface IGarnetApi : IGarnetReadApi, IGarnetAdvancedApi /// /// /// - GarnetStatus StringBitField(ref SpanByte key, ref RawStringInput input, byte secondaryCommand, ref SpanByteAndMemory output); + GarnetStatus StringBitField(ref SpanByte key, ref RawStringInput input, RespCommand secondaryCommand, ref SpanByteAndMemory output); /// /// Performs arbitrary bitfield integer operations on strings. @@ -1622,7 +1622,7 @@ public interface IGarnetReadApi /// /// /// - GarnetStatus StringBitFieldReadOnly(ref SpanByte key, ref RawStringInput input, byte secondaryCommand, ref SpanByteAndMemory output); + GarnetStatus StringBitFieldReadOnly(ref SpanByte key, ref RawStringInput input, RespCommand secondaryCommand, ref SpanByteAndMemory output); #endregion diff --git a/libs/server/Custom/CustomCommandManager.cs b/libs/server/Custom/CustomCommandManager.cs index 07fa54224c1..409b5ba3444 100644 --- a/libs/server/Custom/CustomCommandManager.cs +++ b/libs/server/Custom/CustomCommandManager.cs @@ -12,8 +12,10 @@ namespace Garnet.server /// public class CustomCommandManager { - internal const byte StartOffset = 200; - internal const int MaxRegistrations = byte.MaxValue - StartOffset; + internal static readonly ushort StartOffset = (ushort)(RespCommandExtensions.LastValidCommand + 1); + internal static readonly int MaxRegistrations = ushort.MaxValue - StartOffset; + internal static readonly byte TypeIdStartOffset = (byte)(GarnetObjectTypeExtensions.LastObjectType + 1); + internal static readonly int MaxTypeRegistrations = (byte)(GarnetObjectTypeExtensions.FirstSpecialObjectType) - TypeIdStartOffset; internal readonly CustomRawStringCommand[] rawStringCommandMap; internal readonly CustomObjectCommandWrapper[] objectCommandMap; @@ -34,7 +36,7 @@ public class CustomCommandManager public CustomCommandManager() { rawStringCommandMap = new CustomRawStringCommand[MaxRegistrations]; - objectCommandMap = new CustomObjectCommandWrapper[MaxRegistrations]; + objectCommandMap = new CustomObjectCommandWrapper[MaxTypeRegistrations]; transactionProcMap = new CustomTransaction[MaxRegistrations]; // can increase up to byte.MaxValue customProcedureMap = new CustomProcedureWrapper[MaxRegistrations]; } @@ -45,7 +47,7 @@ internal int Register(string name, CommandType type, CustomRawStringFunctions cu if (id >= MaxRegistrations) throw new Exception("Out of registration space"); - rawStringCommandMap[id] = new CustomRawStringCommand(name, (byte)id, type, customFunctions, expirationTicks); + rawStringCommandMap[id] = new CustomRawStringCommand(name, (ushort)id, type, customFunctions, expirationTicks); if (commandInfo != null) CustomCommandsInfo.Add(name, commandInfo); if (commandDocs != null) CustomCommandsDocs.Add(name, commandDocs); return id; @@ -73,7 +75,7 @@ internal int RegisterType(CustomObjectFactory factory) do { type = Interlocked.Increment(ref ObjectTypeId) - 1; - if (type >= MaxRegistrations) + if (type >= MaxTypeRegistrations) throw new Exception("Out of registration space"); } while (objectCommandMap[type] != null); @@ -84,7 +86,7 @@ internal int RegisterType(CustomObjectFactory factory) internal void RegisterType(int objectTypeId, CustomObjectFactory factory) { - if (objectTypeId >= MaxRegistrations) + if (objectTypeId >= MaxTypeRegistrations) throw new Exception("Type is outside registration space"); if (ObjectTypeId <= objectTypeId) ObjectTypeId = objectTypeId + 1; @@ -106,7 +108,7 @@ internal void RegisterType(int objectTypeId, CustomObjectFactory factory) if (objectTypeId == -1) { objectTypeId = Interlocked.Increment(ref ObjectTypeId) - 1; - if (objectTypeId >= MaxRegistrations) + if (objectTypeId >= MaxTypeRegistrations) throw new Exception("Out of registration space"); objectCommandMap[objectTypeId] = new CustomObjectCommandWrapper((byte)objectTypeId, factory); } @@ -135,7 +137,7 @@ internal void RegisterType(int objectTypeId, CustomObjectFactory factory) if (objectTypeId == -1) { objectTypeId = Interlocked.Increment(ref ObjectTypeId) - 1; - if (objectTypeId >= MaxRegistrations) + if (objectTypeId >= MaxTypeRegistrations) throw new Exception("Out of registration space"); objectCommandMap[objectTypeId] = new CustomObjectCommandWrapper((byte)objectTypeId, factory); } diff --git a/libs/server/Custom/CustomObjectCommand.cs b/libs/server/Custom/CustomObjectCommand.cs index edf7610021f..fac189e48fc 100644 --- a/libs/server/Custom/CustomObjectCommand.cs +++ b/libs/server/Custom/CustomObjectCommand.cs @@ -24,6 +24,6 @@ internal CustomObjectCommand(string name, byte id, byte subid, CommandType type, this.functions = functions; } - internal RespCommand GetRespCommand() => (RespCommand)(id + CustomCommandManager.StartOffset); + internal GarnetObjectType GetObjectType() => (GarnetObjectType)(id + CustomCommandManager.TypeIdStartOffset); } } \ No newline at end of file diff --git a/libs/server/Custom/CustomRawStringCommand.cs b/libs/server/Custom/CustomRawStringCommand.cs index e4f7351e757..4cc4b5c4257 100644 --- a/libs/server/Custom/CustomRawStringCommand.cs +++ b/libs/server/Custom/CustomRawStringCommand.cs @@ -7,12 +7,12 @@ class CustomRawStringCommand { public readonly string NameStr; public readonly byte[] name; - public readonly byte id; + public readonly ushort id; public readonly CommandType type; public readonly CustomRawStringFunctions functions; public long expirationTicks; - internal CustomRawStringCommand(string name, byte id, CommandType type, CustomRawStringFunctions functions, long expirationTicks) + internal CustomRawStringCommand(string name, ushort id, CommandType type, CustomRawStringFunctions functions, long expirationTicks) { NameStr = name.ToUpperInvariant(); this.name = System.Text.Encoding.ASCII.GetBytes(NameStr); diff --git a/libs/server/Custom/CustomRespCommands.cs b/libs/server/Custom/CustomRespCommands.cs index 1426c6e1a28..d625f667338 100644 --- a/libs/server/Custom/CustomRespCommands.cs +++ b/libs/server/Custom/CustomRespCommands.cs @@ -132,14 +132,14 @@ private bool TryCustomRawStringCommand(RespCommand cmd, long expirat /// /// Custom object command /// - private bool TryCustomObjectCommand(RespCommand cmd, byte subid, CommandType type, ref TGarnetApi storageApi) + private bool TryCustomObjectCommand(GarnetObjectType objType, byte subid, CommandType type, ref TGarnetApi storageApi) where TGarnetApi : IGarnetAdvancedApi { var keyBytes = parseState.GetArgSliceByRef(0).SpanByte.ToByteArray(); // Prepare input - var header = new RespInputHeader(cmd) { SubId = subid }; + var header = new RespInputHeader(objType) { SubId = subid }; var input = new ObjectInput(header, ref parseState, 1); var output = new GarnetObjectStoreOutput { spanByteAndMemory = new SpanByteAndMemory(null) }; diff --git a/libs/server/InputHeader.cs b/libs/server/InputHeader.cs index 3779666538f..6625b37a267 100644 --- a/libs/server/InputHeader.cs +++ b/libs/server/InputHeader.cs @@ -38,7 +38,7 @@ public struct RespInputHeader /// /// Size of header /// - public const int Size = 2; + public const int Size = 3; internal const byte FlagMask = (byte)RespInputFlags.SetGet - 1; [FieldOffset(0)] @@ -47,7 +47,7 @@ public struct RespInputHeader [FieldOffset(0)] internal GarnetObjectType type; - [FieldOffset(1)] + [FieldOffset(2)] internal RespInputFlags flags; /// @@ -77,7 +77,7 @@ public RespInputHeader(GarnetObjectType type, RespInputFlags flags = 0) /// /// Command /// Flags - public void SetHeader(byte cmd, byte flags) + public void SetHeader(ushort cmd, byte flags) { this.cmd = (RespCommand)cmd; this.flags = (RespInputFlags)flags; @@ -347,7 +347,7 @@ public RawStringInput(RespCommand cmd, RespInputFlags flags = 0, long arg1 = 0) /// Command /// Flags /// General-purpose argument - public RawStringInput(byte cmd, byte flags = 0, long arg1 = 0) : + public RawStringInput(ushort cmd, byte flags = 0, long arg1 = 0) : this((RespCommand)cmd, (RespInputFlags)flags, arg1) { diff --git a/libs/server/Objects/Types/GarnetObjectSerializer.cs b/libs/server/Objects/Types/GarnetObjectSerializer.cs index 3c72473652b..25633712692 100644 --- a/libs/server/Objects/Types/GarnetObjectSerializer.cs +++ b/libs/server/Objects/Types/GarnetObjectSerializer.cs @@ -58,8 +58,8 @@ private IGarnetObject DeserializeInternal(BinaryReader binaryReader) private IGarnetObject CustomDeserialize(byte type, BinaryReader binaryReader) { - if (type < CustomCommandManager.StartOffset) return null; - return customCommands[type - CustomCommandManager.StartOffset].factory.Deserialize(type, binaryReader); + if (type < CustomCommandManager.TypeIdStartOffset) return null; + return customCommands[type - CustomCommandManager.TypeIdStartOffset].factory.Deserialize(type, binaryReader); } /// diff --git a/libs/server/Objects/Types/GarnetObjectType.cs b/libs/server/Objects/Types/GarnetObjectType.cs index 61d6122df5a..69ad2e793b6 100644 --- a/libs/server/Objects/Types/GarnetObjectType.cs +++ b/libs/server/Objects/Types/GarnetObjectType.cs @@ -29,6 +29,10 @@ public enum GarnetObjectType : byte /// Set, + // Any new object type inserted here should update GarnetObjectTypeExtensions.LastObjectType + + // Any new special type inserted here should update GarnetObjectTypeExtensions.FirstSpecialObjectType + /// /// Special type indicating EXPIRETIME command /// @@ -68,6 +72,12 @@ public enum GarnetObjectType : byte /// Indicating a Custom Object command /// All = 0xfb + } + + public static class GarnetObjectTypeExtensions + { + internal const GarnetObjectType LastObjectType = GarnetObjectType.Set; + internal const GarnetObjectType FirstSpecialObjectType = GarnetObjectType.ExpireTime; } } \ No newline at end of file diff --git a/libs/server/Resp/Bitmap/BitmapCommands.cs b/libs/server/Resp/Bitmap/BitmapCommands.cs index bfa7a6e0812..5f536cf802b 100644 --- a/libs/server/Resp/Bitmap/BitmapCommands.cs +++ b/libs/server/Resp/Bitmap/BitmapCommands.cs @@ -65,43 +65,43 @@ public struct BitFieldCmdArgs /// BITFIELD command /// [FieldOffset(0)] - public byte secondaryOpCode; + public RespCommand secondaryCommand; /// /// encoding info /// - [FieldOffset(1)] + [FieldOffset(sizeof(RespCommand))] public byte typeInfo; /// /// offset /// - [FieldOffset(2)] + [FieldOffset(sizeof(RespCommand) + sizeof(byte))] public long offset; /// /// value /// - [FieldOffset(10)] + [FieldOffset(sizeof(RespCommand) + sizeof(byte) + sizeof(long))] public long value; /// /// BitFieldOverflow enum /// - [FieldOffset(18)] + [FieldOffset(sizeof(RespCommand) + sizeof(byte) + (2 * sizeof(long)))] public byte overflowType; /// /// add a command to execute in bitfield /// - /// + /// /// /// /// /// - public BitFieldCmdArgs(byte secondaryOpCode, byte typeInfo, long offset, long value, byte overflowType) + public BitFieldCmdArgs(RespCommand secondaryCommand, byte typeInfo, long offset, long value, byte overflowType) { - this.secondaryOpCode = secondaryOpCode; + this.secondaryCommand = secondaryCommand; this.typeInfo = typeInfo; this.offset = offset; this.value = value; @@ -488,7 +488,7 @@ private bool StringBitField(ref TGarnetApi storageApi, bool readOnly for (var i = 0; i < secondaryCommandArgs.Count; i++) { - var opCode = (byte)secondaryCommandArgs[i].Item1; + var opCode = secondaryCommandArgs[i].Item1; var opArgs = secondaryCommandArgs[i].Item2; parseState.Initialize(opArgs.Length + (isOverflowTypeSet ? 1 : 0)); @@ -508,7 +508,7 @@ private bool StringBitField(ref TGarnetApi storageApi, bool readOnly var status = storageApi.StringBitField(ref sbKey, ref input, opCode, ref output); - if (status == GarnetStatus.NOTFOUND && opCode == (byte)RespCommand.GET) + if (status == GarnetStatus.NOTFOUND && opCode == RespCommand.GET) { while (!RespWriteUtils.WriteArrayItem(0, ref dcurr, dend)) SendAndReset(); diff --git a/libs/server/Resp/Bitmap/BitmapManagerBitfield.cs b/libs/server/Resp/Bitmap/BitmapManagerBitfield.cs index f33d1493cf6..5c60528dcf7 100644 --- a/libs/server/Resp/Bitmap/BitmapManagerBitfield.cs +++ b/libs/server/Resp/Bitmap/BitmapManagerBitfield.cs @@ -9,7 +9,7 @@ namespace Garnet.server public unsafe partial class BitmapManager { [MethodImpl(MethodImplOptions.AggressiveInlining)] - private static byte GetBitFieldSecondaryOp(byte* input) => (*(BitFieldCmdArgs*)(input)).secondaryOpCode; + private static RespCommand GetBitFieldSecondaryOp(byte* input) => (*(BitFieldCmdArgs*)(input)).secondaryCommand; [MethodImpl(MethodImplOptions.AggressiveInlining)] private static byte GetBitFieldType(byte* input) => (*(BitFieldCmdArgs*)(input)).typeInfo; @@ -459,13 +459,13 @@ public static (long, bool) BitFieldExecute(BitFieldCmdArgs args, byte* value, in { var bitCount = (byte)(args.typeInfo & 0x7F); - switch (args.secondaryOpCode) + switch (args.secondaryCommand) { - case (byte)RespCommand.SET: + case RespCommand.SET: return SetBitfieldValue(value, valLen, args.offset, bitCount, args.typeInfo, args.value, args.overflowType); - case (byte)RespCommand.INCRBY: + case RespCommand.INCRBY: return IncrByBitfieldValue(value, valLen, args.offset, bitCount, args.typeInfo, args.value, args.overflowType); - case (byte)RespCommand.GET: + case RespCommand.GET: return (GetBitfieldValue(value, valLen, args.offset, bitCount, args.typeInfo), false); default: throw new GarnetException("BITFIELD secondary op not supported"); diff --git a/libs/server/Resp/Parser/RespCommand.cs b/libs/server/Resp/Parser/RespCommand.cs index a6b94e6b349..7cd956ebb36 100644 --- a/libs/server/Resp/Parser/RespCommand.cs +++ b/libs/server/Resp/Parser/RespCommand.cs @@ -3,6 +3,7 @@ using System; using System.Diagnostics; +using System.Linq; using System.Runtime.CompilerServices; using System.Runtime.InteropServices; using System.Text; @@ -14,7 +15,7 @@ namespace Garnet.server /// /// Basic RESP command enum /// - public enum RespCommand : byte + public enum RespCommand : ushort { NONE = 0x00, @@ -302,7 +303,10 @@ public enum RespCommand : byte HELLO, QUIT, // Note: Update IsNoAuth if adding new no-auth commands after this - INVALID = 0xFF, + // Max value of this enum (not including INVALID) will determine the size of RespCommand.AofIndependentBitLookup and CommandPermissionSet._commandList, + // so avoid manually setting high values unless necessary + + INVALID = 0xFFFF, } /// @@ -357,15 +361,19 @@ public static class RespCommandExtensions RespCommand.MULTI, ]; - // long is 64 bits, 4 longs accomodate 256 resp commands which is more than enough to provide a lookup for each resp command - private static readonly ulong[] AofIndepenedentBitLookup = [0, 0, 0, 0]; + private static readonly ulong[] AofIndependentBitLookup; private const int sizeOfLong = 64; // The static ctor maybe expensive but it is only ever run once, and doesn't interfere with common path static RespCommandExtensions() { - foreach (RespCommand cmd in Enum.GetValues(typeof(RespCommand))) + // # of bits needed to represent all valid commands + var maxBitsNeeded = (ushort)LastValidCommand + 1; + var lookupTableSize = (maxBitsNeeded / 64) + (maxBitsNeeded % 64 == 0 ? 0 : 1); + AofIndependentBitLookup = new ulong[lookupTableSize]; + + foreach (var cmd in Enum.GetValues()) { if (Array.IndexOf(AofIndependentCommands, cmd) == -1) continue; @@ -375,7 +383,7 @@ static RespCommandExtensions() // set the respCommand's bit to indicate int bitIdxOffset = (int)cmd % sizeOfLong; ulong bitmask = 1UL << bitIdxOffset; - AofIndepenedentBitLookup[bitIdxToUse] |= bitmask; + AofIndependentBitLookup[bitIdxToUse] |= bitmask; } } @@ -385,11 +393,13 @@ static RespCommandExtensions() [MethodImpl(MethodImplOptions.AggressiveInlining)] public static bool IsAofIndependent(this RespCommand cmd) { + if (cmd > LastValidCommand) return false; + // check if cmd maps to a bit vec that was set back when static ctor was run int bitIdxToUse = (int)cmd / sizeOfLong; int bitIdxOffset = (int)cmd % sizeOfLong; ulong bitmask = 1UL << bitIdxOffset; - return (AofIndepenedentBitLookup[bitIdxToUse] & bitmask) != 0; + return (AofIndependentBitLookup[bitIdxToUse] & bitmask) != 0; } /// @@ -435,6 +445,11 @@ public static ReadOnlySpan ExpandForACLs(this RespCommand cmd) internal const RespCommand LastDataCommand = RespCommand.EVALSHA; + /// + /// Last valid command (i.e. RespCommand with the largest value excluding INVALID). + /// + public static RespCommand LastValidCommand { get; } = Enum.GetValues().Where(cmd => cmd != RespCommand.INVALID).Max(); + [MethodImpl(MethodImplOptions.AggressiveInlining)] public static bool IsReadOnly(this RespCommand cmd) => cmd <= LastReadCommand; diff --git a/libs/server/Resp/RespCommandAccessor.cs b/libs/server/Resp/RespCommandAccessor.cs index 0c3b2521ee8..b5a3ac927fb 100644 --- a/libs/server/Resp/RespCommandAccessor.cs +++ b/libs/server/Resp/RespCommandAccessor.cs @@ -11,6 +11,6 @@ public static class RespCommandAccessor /// /// MIGRATE /// - public static byte MIGRATE => (byte)RespCommand.MIGRATE; + public static ushort MIGRATE => (ushort)RespCommand.MIGRATE; } } \ No newline at end of file diff --git a/libs/server/Resp/RespServerSession.cs b/libs/server/Resp/RespServerSession.cs index 91332dbdfd5..ef7d28d8aca 100644 --- a/libs/server/Resp/RespServerSession.cs +++ b/libs/server/Resp/RespServerSession.cs @@ -807,7 +807,7 @@ bool NetworkCustomObjCmd(ref TGarnetApi storageApi) } // Perform the operation - TryCustomObjectCommand(currentCustomObjectCommand.GetRespCommand(), currentCustomObjectCommand.subid, + TryCustomObjectCommand(currentCustomObjectCommand.GetObjectType(), currentCustomObjectCommand.subid, currentCustomObjectCommand.type, ref storageApi); currentCustomObjectCommand = null; return true; diff --git a/libs/server/Storage/Functions/MainStore/PrivateMethods.cs b/libs/server/Storage/Functions/MainStore/PrivateMethods.cs index 75b9a988035..50bf90392a4 100644 --- a/libs/server/Storage/Functions/MainStore/PrivateMethods.cs +++ b/libs/server/Storage/Functions/MainStore/PrivateMethods.cs @@ -713,12 +713,12 @@ void WriteLogDelete(ref SpanByte key, long version, int sessionID) BitFieldCmdArgs GetBitFieldArguments(ref RawStringInput input) { var currTokenIdx = input.parseStateFirstArgIdx; - var opCode = (byte)input.parseState.GetEnum(currTokenIdx++, true); + var cmd = input.parseState.GetEnum(currTokenIdx++, true); var encodingArg = input.parseState.GetString(currTokenIdx++); var offsetArg = input.parseState.GetString(currTokenIdx++); long value = default; - if (opCode == (byte)RespCommand.SET || opCode == (byte)RespCommand.INCRBY) + if (cmd == RespCommand.SET || cmd == RespCommand.INCRBY) { value = input.parseState.GetLong(currTokenIdx++); } @@ -738,7 +738,7 @@ BitFieldCmdArgs GetBitFieldArguments(ref RawStringInput input) // Calculate number offset from bitCount if offsetArg starts with # var offset = offsetArg[0] == '#' ? long.Parse(offsetArg.AsSpan(1)) * bitCount : long.Parse(offsetArg); - return new BitFieldCmdArgs(opCode, typeInfo, offset, value, overflowType); + return new BitFieldCmdArgs(cmd, typeInfo, offset, value, overflowType); } } } \ No newline at end of file diff --git a/libs/server/Storage/Functions/MainStore/RMWMethods.cs b/libs/server/Storage/Functions/MainStore/RMWMethods.cs index f81c258629a..58a9c5a6acc 100644 --- a/libs/server/Storage/Functions/MainStore/RMWMethods.cs +++ b/libs/server/Storage/Functions/MainStore/RMWMethods.cs @@ -30,11 +30,11 @@ public bool NeedInitialUpdate(ref SpanByte key, ref RawStringInput input, ref Sp case RespCommand.GETEX: return false; default: - if ((byte)input.header.cmd >= CustomCommandManager.StartOffset) + if ((ushort)input.header.cmd >= CustomCommandManager.StartOffset) { (IMemoryOwner Memory, int Length) outp = (output.Memory, 0); var ret = functionsState - .customCommands[(byte)input.header.cmd - CustomCommandManager.StartOffset].functions + .customCommands[(ushort)input.header.cmd - CustomCommandManager.StartOffset].functions .NeedInitialUpdate(key.AsReadOnlySpan(), ref input, ref outp); output.Memory = outp.Memory; output.Length = outp.Length; @@ -179,9 +179,9 @@ public bool InitialUpdater(ref SpanByte key, ref RawStringInput input, ref SpanB default: value.UnmarkExtraMetadata(); - if ((byte)input.header.cmd >= CustomCommandManager.StartOffset) + if ((ushort)input.header.cmd >= CustomCommandManager.StartOffset) { - var functions = functionsState.customCommands[(byte)input.header.cmd - CustomCommandManager.StartOffset].functions; + var functions = functionsState.customCommands[(ushort)input.header.cmd - CustomCommandManager.StartOffset].functions; // compute metadata size for result var expiration = input.arg1; metadataSize = expiration switch @@ -507,7 +507,7 @@ private bool InPlaceUpdaterWorker(ref SpanByte key, ref RawStringInput input, re return false; default: - var cmd = (byte)input.header.cmd; + var cmd = (ushort)input.header.cmd; if (cmd >= CustomCommandManager.StartOffset) { var functions = functionsState.customCommands[cmd - CustomCommandManager.StartOffset].functions; @@ -576,10 +576,10 @@ public bool NeedCopyUpdate(ref SpanByte key, ref RawStringInput input, ref SpanB } return false; default: - if ((byte)input.header.cmd >= CustomCommandManager.StartOffset) + if ((ushort)input.header.cmd >= CustomCommandManager.StartOffset) { (IMemoryOwner Memory, int Length) outp = (output.Memory, 0); - var ret = functionsState.customCommands[(byte)input.header.cmd - CustomCommandManager.StartOffset].functions + var ret = functionsState.customCommands[(ushort)input.header.cmd - CustomCommandManager.StartOffset].functions .NeedCopyUpdate(key.AsReadOnlySpan(), ref input, oldValue.AsReadOnlySpan(), ref outp); output.Memory = outp.Memory; output.Length = outp.Length; @@ -812,7 +812,7 @@ public bool CopyUpdater(ref SpanByte key, ref RawStringInput input, ref SpanByte break; default: - if ((byte)input.header.cmd >= CustomCommandManager.StartOffset) + if ((ushort)input.header.cmd >= CustomCommandManager.StartOffset) { var functions = functionsState.customCommands[(byte)input.header.cmd - CustomCommandManager.StartOffset].functions; var expiration = input.arg1; diff --git a/libs/server/Storage/Functions/MainStore/ReadMethods.cs b/libs/server/Storage/Functions/MainStore/ReadMethods.cs index 798d3372e53..cd0a0be7853 100644 --- a/libs/server/Storage/Functions/MainStore/ReadMethods.cs +++ b/libs/server/Storage/Functions/MainStore/ReadMethods.cs @@ -19,11 +19,11 @@ public bool SingleReader(ref SpanByte key, ref RawStringInput input, ref SpanByt return false; var cmd = input.header.cmd; - if ((byte)cmd >= CustomCommandManager.StartOffset) + if ((ushort)cmd >= CustomCommandManager.StartOffset) { var valueLength = value.LengthWithoutMetadata; (IMemoryOwner Memory, int Length) output = (dst.Memory, 0); - var ret = functionsState.customCommands[(byte)cmd - CustomCommandManager.StartOffset].functions + var ret = functionsState.customCommands[(ushort)cmd - CustomCommandManager.StartOffset].functions .Reader(key.AsReadOnlySpan(), ref input, value.AsReadOnlySpan(), ref output, ref readInfo); Debug.Assert(valueLength <= value.LengthWithoutMetadata); dst.Memory = output.Memory; @@ -50,11 +50,11 @@ public bool ConcurrentReader(ref SpanByte key, ref RawStringInput input, ref Spa } var cmd = input.header.cmd; - if ((byte)cmd >= CustomCommandManager.StartOffset) + if ((ushort)cmd >= CustomCommandManager.StartOffset) { var valueLength = value.LengthWithoutMetadata; (IMemoryOwner Memory, int Length) output = (dst.Memory, 0); - var ret = functionsState.customCommands[(byte)cmd - CustomCommandManager.StartOffset].functions + var ret = functionsState.customCommands[(ushort)cmd - CustomCommandManager.StartOffset].functions .Reader(key.AsReadOnlySpan(), ref input, value.AsReadOnlySpan(), ref output, ref readInfo); Debug.Assert(valueLength <= value.LengthWithoutMetadata); dst.Memory = output.Memory; diff --git a/libs/server/Storage/Functions/MainStore/VarLenInputMethods.cs b/libs/server/Storage/Functions/MainStore/VarLenInputMethods.cs index 8d1d0db833c..e17917a546e 100644 --- a/libs/server/Storage/Functions/MainStore/VarLenInputMethods.cs +++ b/libs/server/Storage/Functions/MainStore/VarLenInputMethods.cs @@ -119,9 +119,9 @@ public int GetRMWInitialValueLength(ref RawStringInput input) return sizeof(int) + ndigits; default: - if ((byte)cmd >= CustomCommandManager.StartOffset) + if ((ushort)cmd >= CustomCommandManager.StartOffset) { - var functions = functionsState.customCommands[(byte)cmd - CustomCommandManager.StartOffset].functions; + var functions = functionsState.customCommands[(ushort)cmd - CustomCommandManager.StartOffset].functions; // Compute metadata size for result int metadataSize = input.arg1 switch { @@ -236,9 +236,9 @@ public int GetRMWModifiedValueLength(ref SpanByte t, ref RawStringInput input) return sizeof(int) + t.Length + valueLength; default: - if ((byte)cmd >= CustomCommandManager.StartOffset) + if ((ushort)cmd >= CustomCommandManager.StartOffset) { - var functions = functionsState.customCommands[(byte)cmd - CustomCommandManager.StartOffset].functions; + var functions = functionsState.customCommands[(ushort)cmd - CustomCommandManager.StartOffset].functions; // compute metadata for result var metadataSize = input.arg1 switch { diff --git a/libs/server/Storage/Functions/ObjectStore/PrivateMethods.cs b/libs/server/Storage/Functions/ObjectStore/PrivateMethods.cs index 1be16e51a31..b175f7a5096 100644 --- a/libs/server/Storage/Functions/ObjectStore/PrivateMethods.cs +++ b/libs/server/Storage/Functions/ObjectStore/PrivateMethods.cs @@ -184,7 +184,7 @@ static bool EvaluateObjectExpireInPlace(ExpireOption optionType, bool expiryExis [MethodImpl(MethodImplOptions.AggressiveInlining)] private CustomObjectFunctions GetCustomObjectCommand(ref ObjectInput input, GarnetObjectType type) { - var objectId = (byte)((byte)type - CustomCommandManager.StartOffset); + var objectId = (byte)((byte)type - CustomCommandManager.TypeIdStartOffset); var cmdId = input.header.SubId; var customObjectCommand = functionsState.customObjectCommands[objectId].commandMap[cmdId].functions; return customObjectCommand; diff --git a/libs/server/Storage/Functions/ObjectStore/RMWMethods.cs b/libs/server/Storage/Functions/ObjectStore/RMWMethods.cs index bec97eb922f..9186d19b35f 100644 --- a/libs/server/Storage/Functions/ObjectStore/RMWMethods.cs +++ b/libs/server/Storage/Functions/ObjectStore/RMWMethods.cs @@ -26,7 +26,7 @@ public bool NeedInitialUpdate(ref byte[] key, ref ObjectInput input, ref GarnetO case GarnetObjectType.Persist: return false; default: - if ((byte)type < CustomCommandManager.StartOffset) + if ((byte)type < CustomCommandManager.TypeIdStartOffset) return GarnetObject.NeedToCreate(input.header); else { @@ -44,7 +44,7 @@ public bool NeedInitialUpdate(ref byte[] key, ref ObjectInput input, ref GarnetO public bool InitialUpdater(ref byte[] key, ref ObjectInput input, ref IGarnetObject value, ref GarnetObjectStoreOutput output, ref RMWInfo rmwInfo, ref RecordInfo recordInfo) { var type = input.header.type; - if ((byte)type < CustomCommandManager.StartOffset) + if ((byte)type < CustomCommandManager.TypeIdStartOffset) { value = GarnetObject.Create(type); value.Operate(ref input, ref output.spanByteAndMemory, out _, out _); @@ -55,7 +55,7 @@ public bool InitialUpdater(ref byte[] key, ref ObjectInput input, ref IGarnetObj Debug.Assert(type != GarnetObjectType.Expire && type != GarnetObjectType.PExpire && type != GarnetObjectType.Persist, "Expire and Persist commands should have been handled already by NeedInitialUpdate."); var customObjectCommand = GetCustomObjectCommand(ref input, type); - var objectId = (byte)((byte)type - CustomCommandManager.StartOffset); + var objectId = (byte)((byte)type - CustomCommandManager.TypeIdStartOffset); value = functionsState.customObjectCommands[objectId].factory.Create((byte)type); (IMemoryOwner Memory, int Length) outp = (output.spanByteAndMemory.Memory, 0); @@ -140,7 +140,7 @@ bool InPlaceUpdaterWorker(ref byte[] key, ref ObjectInput input, ref IGarnetObje CopyDefaultResp(CmdStrings.RESP_RETURN_VAL_0, ref output.spanByteAndMemory); return true; default: - if ((byte)input.header.type < CustomCommandManager.StartOffset) + if ((byte)input.header.type < CustomCommandManager.TypeIdStartOffset) { var operateSuccessful = value.Operate(ref input, ref output.spanByteAndMemory, out sizeChange, out var removeKey); @@ -232,7 +232,7 @@ public bool PostCopyUpdater(ref byte[] key, ref ObjectInput input, ref IGarnetOb CopyDefaultResp(CmdStrings.RESP_RETURN_VAL_0, ref output.spanByteAndMemory); break; default: - if ((byte)input.header.type < CustomCommandManager.StartOffset) + if ((byte)input.header.type < CustomCommandManager.TypeIdStartOffset) { value.Operate(ref input, ref output.spanByteAndMemory, out _, out var removeKey); if (removeKey) diff --git a/libs/server/Storage/Functions/ObjectStore/ReadMethods.cs b/libs/server/Storage/Functions/ObjectStore/ReadMethods.cs index ab6c2558afd..1fdb4898252 100644 --- a/libs/server/Storage/Functions/ObjectStore/ReadMethods.cs +++ b/libs/server/Storage/Functions/ObjectStore/ReadMethods.cs @@ -46,7 +46,7 @@ public bool SingleReader(ref byte[] key, ref ObjectInput input, ref IGarnetObjec return true; default: - if ((byte)input.header.type < CustomCommandManager.StartOffset) + if ((byte)input.header.type < CustomCommandManager.TypeIdStartOffset) return value.Operate(ref input, ref dst.spanByteAndMemory, out _, out _); if (IncorrectObjectType(ref input, value, ref dst.spanByteAndMemory)) diff --git a/libs/server/Storage/Session/MainStore/BitmapOps.cs b/libs/server/Storage/Session/MainStore/BitmapOps.cs index 3e7451bd90c..458185191b6 100644 --- a/libs/server/Storage/Session/MainStore/BitmapOps.cs +++ b/libs/server/Storage/Session/MainStore/BitmapOps.cs @@ -280,7 +280,7 @@ public unsafe GarnetStatus StringBitField(ArgSlice key, List 0 ? "i"u8 : "u"u8; var encodingSuffix = commandArguments[i].typeInfo & 0x7F; var encodingSuffixLength = NumUtils.NumDigits(encodingSuffix); @@ -339,13 +339,13 @@ public unsafe GarnetStatus StringBitField(ArgSlice key, List(ref SpanByte key, ref Raw where TContext : ITsavoriteContext => Read_MainStore(ref key, ref input, ref output, ref context); - public unsafe GarnetStatus StringBitField(ref SpanByte key, ref RawStringInput input, byte secondaryCommand, ref SpanByteAndMemory output, ref TContext context) + public unsafe GarnetStatus StringBitField(ref SpanByte key, ref RawStringInput input, RespCommand secondaryCommand, ref SpanByteAndMemory output, ref TContext context) where TContext : ITsavoriteContext { GarnetStatus status; - if (secondaryCommand == (byte)RespCommand.GET) + if (secondaryCommand == RespCommand.GET) status = Read_MainStore(ref key, ref input, ref output, ref context); else status = RMW_MainStore(ref key, ref input, ref output, ref context); return status; } - public unsafe GarnetStatus StringBitFieldReadOnly(ref SpanByte key, ref RawStringInput input, byte secondaryCommand, ref SpanByteAndMemory output, ref TContext context) + public unsafe GarnetStatus StringBitFieldReadOnly(ref SpanByte key, ref RawStringInput input, RespCommand secondaryCommand, ref SpanByteAndMemory output, ref TContext context) where TContext : ITsavoriteContext { GarnetStatus status = GarnetStatus.NOTFOUND; - if (secondaryCommand == (byte)RespCommand.GET) + if (secondaryCommand == RespCommand.GET) status = Read_MainStore(ref key, ref input, ref output, ref context); return status; } diff --git a/test/Garnet.test/RespCommandTests.cs b/test/Garnet.test/RespCommandTests.cs index 7f75319ad1e..339266f60df 100644 --- a/test/Garnet.test/RespCommandTests.cs +++ b/test/Garnet.test/RespCommandTests.cs @@ -383,7 +383,7 @@ public void AofIndependentCommandsTest() RespCommand.MULTI, ]; - foreach (RespCommand cmd in Enum.GetValues(typeof(RespCommand))) + foreach (var cmd in Enum.GetValues().Where(cmd => cmd != RespCommand.INVALID)) { var expectedAofIndependence = Array.IndexOf(aofIndpendentCmds, cmd) != -1; ClassicAssert.AreEqual(expectedAofIndependence, cmd.IsAofIndependent()); diff --git a/test/Garnet.test/TestProcedureBitmap.cs b/test/Garnet.test/TestProcedureBitmap.cs index 9712cb21cc9..b9b3e3d65ea 100644 --- a/test/Garnet.test/TestProcedureBitmap.cs +++ b/test/Garnet.test/TestProcedureBitmap.cs @@ -133,10 +133,10 @@ public override void Main(TGarnetApi api, ref CustomProcedureInput p api.SET(bitmapA, data); var listCommands = new List(); - var bitFieldArguments = new BitFieldCmdArgs((byte)RespCommand.GET, ((byte)BitFieldSign.UNSIGNED | 8), 0, 0, (byte)BitFieldOverflow.WRAP); + var bitFieldArguments = new BitFieldCmdArgs(RespCommand.GET, ((byte)BitFieldSign.UNSIGNED | 8), 0, 0, (byte)BitFieldOverflow.WRAP); listCommands.Add(bitFieldArguments); - bitFieldArguments = new BitFieldCmdArgs((byte)RespCommand.INCRBY, ((byte)BitFieldSign.UNSIGNED | 4), 4, 1, (byte)BitFieldOverflow.WRAP); + bitFieldArguments = new BitFieldCmdArgs(RespCommand.INCRBY, ((byte)BitFieldSign.UNSIGNED | 4), 4, 1, (byte)BitFieldOverflow.WRAP); listCommands.Add(bitFieldArguments); api.StringBitField(bitmapA, listCommands, out var resultBitField); From 512c548d003785014db921b07002cfbeeec2c74f Mon Sep 17 00:00:00 2001 From: Yoganand Rajasekaran Date: Fri, 16 Aug 2024 14:39:31 -0700 Subject: [PATCH 02/12] Invoke custom raw string cmd from custom proc/txn --- libs/server/API/GarnetApi.cs | 3 +++ libs/server/API/GarnetWatchApi.cs | 6 +++++ libs/server/API/IGarnetApi.cs | 1 + libs/server/ArgSlice/ArgSlice.cs | 4 ++-- .../Storage/Session/MainStore/MainStoreOps.cs | 19 +++++++++++++++ main/GarnetServer/Extensions/ProcCustomCmd.cs | 24 +++++++++++++++++++ main/GarnetServer/Program.cs | 2 ++ 7 files changed, 57 insertions(+), 2 deletions(-) create mode 100644 main/GarnetServer/Extensions/ProcCustomCmd.cs diff --git a/libs/server/API/GarnetApi.cs b/libs/server/API/GarnetApi.cs index 34c91d00a5f..282f7d36045 100644 --- a/libs/server/API/GarnetApi.cs +++ b/libs/server/API/GarnetApi.cs @@ -416,5 +416,8 @@ public int GetScratchBufferOffset() public bool ResetScratchBuffer(int offset) => storageSession.scratchBufferManager.ResetScratchBuffer(offset); #endregion + + public unsafe GarnetStatus CustomCommand(byte id, ArgSlice key, ArgSlice input, ref SpanByteAndMemory output) + => storageSession.CustomCommand(id, key, input, ref output, ref context); } } \ No newline at end of file diff --git a/libs/server/API/GarnetWatchApi.cs b/libs/server/API/GarnetWatchApi.cs index ef6d82266f3..a23df550c3b 100644 --- a/libs/server/API/GarnetWatchApi.cs +++ b/libs/server/API/GarnetWatchApi.cs @@ -563,6 +563,12 @@ public int GetScratchBufferOffset() public bool ResetScratchBuffer(int offset) => garnetApi.ResetScratchBuffer(offset); + public GarnetStatus CustomCommand(byte id, ArgSlice key, ArgSlice input, ref SpanByteAndMemory output) + { + garnetApi.WATCH(key, StoreType.Main); + return garnetApi.CustomCommand(id, key, input, ref output); + } + #endregion } } \ No newline at end of file diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs index e271407c63f..a0fc9ba5c99 100644 --- a/libs/server/API/IGarnetApi.cs +++ b/libs/server/API/IGarnetApi.cs @@ -1734,6 +1734,7 @@ public bool IterateObjectStore(ref TScanFunctions scanFunctions, #endregion + public GarnetStatus CustomCommand(byte id, ArgSlice key, ArgSlice input, ref SpanByteAndMemory output); } /// diff --git a/libs/server/ArgSlice/ArgSlice.cs b/libs/server/ArgSlice/ArgSlice.cs index 8ea0a889306..984d39b57bc 100644 --- a/libs/server/ArgSlice/ArgSlice.cs +++ b/libs/server/ArgSlice/ArgSlice.cs @@ -21,10 +21,10 @@ public unsafe struct ArgSlice public const int Size = 12; [FieldOffset(0)] - internal byte* ptr; + public byte* ptr; [FieldOffset(8)] - internal int length; + public int length; /// /// Create new ArgSlice from given pointer and length diff --git a/libs/server/Storage/Session/MainStore/MainStoreOps.cs b/libs/server/Storage/Session/MainStore/MainStoreOps.cs index fceaf24d88e..bb1588ee12e 100644 --- a/libs/server/Storage/Session/MainStore/MainStoreOps.cs +++ b/libs/server/Storage/Session/MainStore/MainStoreOps.cs @@ -1124,6 +1124,25 @@ public unsafe GarnetStatus SCAN(long cursor, ArgSlice match, long coun return GarnetStatus.OK; } + public unsafe GarnetStatus CustomCommand(byte id, ArgSlice key, ArgSlice input, ref SpanByteAndMemory output, ref TContext context) + where TContext : ITsavoriteContext + { + var sbKey = key.SpanByte; + + int inputSize = sizeof(int) + RespInputHeader.Size + input.Length; + byte* pbCmdInput = stackalloc byte[inputSize]; + + byte* pcurr = pbCmdInput; + *(int*)pcurr = inputSize - sizeof(int); + pcurr += sizeof(int); + (*(RespInputHeader*)(pcurr)).cmd = (RespCommand)(id + CustomCommandManager.StartOffset); + (*(RespInputHeader*)(pcurr)).flags = 0; + pcurr += RespInputHeader.Size; + Buffer.MemoryCopy(input.ptr, pcurr, input.Length, input.Length); + + return RMW_MainStore(ref sbKey, ref Unsafe.AsRef(pbCmdInput), ref output, ref context); + } + public GarnetStatus GetKeyType(ArgSlice key, out string keyType, ref TContext context, ref TObjectContext objectContext) where TContext : ITsavoriteContext where TObjectContext : ITsavoriteContext diff --git a/main/GarnetServer/Extensions/ProcCustomCmd.cs b/main/GarnetServer/Extensions/ProcCustomCmd.cs new file mode 100644 index 00000000000..468ff850bfe --- /dev/null +++ b/main/GarnetServer/Extensions/ProcCustomCmd.cs @@ -0,0 +1,24 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +using Garnet.common; +using Garnet.server; +using Tsavorite.core; + +namespace Garnet +{ + class ProcCustomCmd : CustomProcedure + { + public override unsafe bool Execute(IGarnetApi garnetApi, ArgSlice input, ref MemoryResult output) + { + var offset = 0; + ArgSlice key = GetNextArg(input, ref offset); + var cmdOutput = new SpanByteAndMemory(null); + + // id from registration of custom raw string cmd + garnetApi.CustomCommand(0, key, new ArgSlice(input.ptr + offset, input.length - offset), ref cmdOutput); + + return true; + } + } +} \ No newline at end of file diff --git a/main/GarnetServer/Program.cs b/main/GarnetServer/Program.cs index 0b5e165d31b..fd5f03126c6 100644 --- a/main/GarnetServer/Program.cs +++ b/main/GarnetServer/Program.cs @@ -94,6 +94,8 @@ static void RegisterExtensions(GarnetServer server) server.Register.NewProcedure("SUM", new Sum()); server.Register.NewProcedure("SETMAINANDOBJECT", new SetStringAndList()); + + server.Register.NewProcedure("PROCCMD", new ProcCustomCmd()); } } } \ No newline at end of file From 020db536260e72183878cad04f68787421491965 Mon Sep 17 00:00:00 2001 From: Yoganand Rajasekaran Date: Sun, 20 Oct 2024 19:50:44 -0700 Subject: [PATCH 03/12] Support cmd name matching --- libs/server/Custom/CustomCommandManager.cs | 2 +- libs/server/Custom/CustomFunctions.cs | 23 ++++ libs/server/Custom/CustomProcedureWrapper.cs | 58 ++++++++- libs/server/Storage/Session/CustomOps.cs | 121 ++++++++++++++++++ .../Storage/Session/MainStore/MainStoreOps.cs | 4 +- main/GarnetServer/Extensions/ProcCustomCmd.cs | 3 +- 6 files changed, 203 insertions(+), 8 deletions(-) create mode 100644 libs/server/Storage/Session/CustomOps.cs diff --git a/libs/server/Custom/CustomCommandManager.cs b/libs/server/Custom/CustomCommandManager.cs index 409b5ba3444..8ad42bea522 100644 --- a/libs/server/Custom/CustomCommandManager.cs +++ b/libs/server/Custom/CustomCommandManager.cs @@ -170,7 +170,7 @@ internal int Register(string name, CustomProcedure customProcedure, RespCommands if (id >= MaxRegistrations) throw new Exception("Out of registration space"); - customProcedureMap[id] = new CustomProcedureWrapper(name, (byte)id, customProcedure); + customProcedureMap[id] = new CustomProcedureWrapper(name, (byte)id, customProcedure, this); if (commandInfo != null) CustomCommandsInfo.Add(name, commandInfo); if (commandDocs != null) CustomCommandsDocs.Add(name, commandDocs); return id; diff --git a/libs/server/Custom/CustomFunctions.cs b/libs/server/Custom/CustomFunctions.cs index a87592e1b5d..45081e075e2 100644 --- a/libs/server/Custom/CustomFunctions.cs +++ b/libs/server/Custom/CustomFunctions.cs @@ -5,7 +5,9 @@ using System.Buffers; using System.Collections.Generic; using System.Diagnostics; +using System.Text; using Garnet.common; +using Tsavorite.core; namespace Garnet.server { @@ -19,6 +21,8 @@ public abstract class CustomFunctions /// protected static MemoryPool MemoryPool => MemoryPool.Shared; + public CustomCommandManager customCommandManager; + /// /// Create output as simple string, from given string /// @@ -208,5 +212,24 @@ protected static unsafe ArgSlice GetNextArg(ref CustomProcedureInput procInput, { return GetNextArg(ref procInput.parseState, procInput.parseStateFirstArgIdx, ref offset); } + + protected void InvokeCustomRawStringCommand(TGarnetApi garnetApi, string cmd, ArgSlice key, ArgSlice input) + where TGarnetApi : IGarnetApi + { + + //var request = scratchBufferManager.FormatCommandAsResp(cmd, new[] { key, input }, null); + //_ = respServerSession.TryConsumeMessages(request.ptr, request.length); + //var response = scratchBufferNetworkSender.GetResponse(); + //var result = ProcessResponse(response.ptr, response.length); + //scratchBufferNetworkSender.Reset(); + //output = result; + + //output = + var output = new SpanByteAndMemory(null); + if (customCommandManager.Match(new ReadOnlySpan(Encoding.UTF8.GetBytes(cmd)), out CustomRawStringCommand customCommand)) + { + garnetApi.CustomCommand(customCommand.id, key, input, ref output); + } + } } } \ No newline at end of file diff --git a/libs/server/Custom/CustomProcedureWrapper.cs b/libs/server/Custom/CustomProcedureWrapper.cs index a6e21736bdd..07da96a371b 100644 --- a/libs/server/Custom/CustomProcedureWrapper.cs +++ b/libs/server/Custom/CustomProcedureWrapper.cs @@ -2,6 +2,7 @@ // Licensed under the MIT license. using System; +using System.Diagnostics; using Garnet.common; namespace Garnet.server @@ -11,6 +12,52 @@ namespace Garnet.server /// public abstract class CustomProcedure : CustomFunctions { + internal ScratchBufferManager scratchBufferManager; + + //unsafe object CustomOperation(TGarnetApi api, string cmd, params object[] args) + // where TGarnetApi : IGarnetApi + //unsafe object CustomOperation(TGarnetApi api, string cmd, params object[] args) + // where TGarnetApi : IGarnetApi + //{ + //switch (cmd) + //{ + // // We special-case a few performance-sensitive operations to directly invoke via the storage API + // case "SET" when args.Length == 2: + // case "set" when args.Length == 2: + // { + // if (!respServerSession.CheckACLPermissions(RespCommand.SET)) + // return Encoding.ASCII.GetString(CmdStrings.RESP_ERR_NOAUTH); + // var key = scratchBufferManager.CreateArgSlice(Convert.ToString(args[0])); + // var value = scratchBufferManager.CreateArgSlice(Convert.ToString(args[1])); + // _ = api.SET(key, value); + // return "OK"; + // } + // case "GET": + // case "get": + // { + // if (!respServerSession.CheckACLPermissions(RespCommand.GET)) + // throw new Exception(Encoding.ASCII.GetString(CmdStrings.RESP_ERR_NOAUTH)); + // var key = scratchBufferManager.CreateArgSlice(Convert.ToString(args[0])); + // var status = api.GET(key, out var value); + // if (status == GarnetStatus.OK) + // return value.ToString(); + // return null; + // } + // // As fallback, we use RespServerSession with a RESP-formatted input. This could be optimized + // // in future to provide parse state directly. + // default: + // { + //var request = scratchBufferManager.FormatCommandAsResp(cmd, args, null); + //_ = respServerSession.TryConsumeMessages(request.ptr, request.length); + //var response = scratchBufferNetworkSender.GetResponse(); + //var result = ProcessResponse(response.ptr, response.length); + //scratchBufferNetworkSender.Reset(); + //return result; + // } + //} + //} + + /// /// Custom command implementation /// @@ -28,18 +75,21 @@ class CustomProcedureWrapper public readonly byte Id; public readonly CustomProcedure CustomProcedureImpl; - internal CustomProcedureWrapper(string name, byte id, CustomProcedure customScriptProc) + internal CustomProcedureWrapper(string name, byte id, CustomProcedure customProcedure, CustomCommandManager customCommandManager) { if (string.IsNullOrEmpty(name)) throw new ArgumentNullException(nameof(name)); - if (customScriptProc == null) - throw new ArgumentNullException(nameof(customScriptProc)); + if (customProcedure == null) + throw new ArgumentNullException(nameof(customProcedure)); + + Debug.Assert(customCommandManager != null); NameStr = name.ToUpperInvariant(); Name = System.Text.Encoding.ASCII.GetBytes(NameStr); Id = id; - CustomProcedureImpl = customScriptProc; + CustomProcedureImpl = customProcedure; + CustomProcedureImpl.customCommandManager = customCommandManager; } } } \ No newline at end of file diff --git a/libs/server/Storage/Session/CustomOps.cs b/libs/server/Storage/Session/CustomOps.cs new file mode 100644 index 00000000000..0c99fee326c --- /dev/null +++ b/libs/server/Storage/Session/CustomOps.cs @@ -0,0 +1,121 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//using System; +//using System.Diagnostics; +//using Garnet.common; +//using Tsavorite.core; + +//namespace Garnet.server +//{ +// using ObjectStoreAllocator = GenericAllocator>>; +// using ObjectStoreFunctions = StoreFunctions>; + +// sealed partial class StorageSession : IDisposable +// { +// private bool TryCustomObjectCommand(byte* ptr, byte* end, RespCommand cmd, byte subid, CommandType type, ref TGarnetApi storageApi) +// where TGarnetApi : IGarnetAdvancedApi +// { +// var keyBytes = parseState.GetArgSliceByRef(0).SpanByte.ToByteArray(); + +// // Prepare input +// var input = new ObjectInput +// { +// header = new RespInputHeader +// { +// cmd = cmd, +// SubId = subid +// }, +// parseState = parseState, +// parseStateStartIdx = 1 +// }; + +// var output = new GarnetObjectStoreOutput { spanByteAndMemory = new SpanByteAndMemory(null) }; + +// GarnetStatus status; + +// if (type == CommandType.ReadModifyWrite) +// { +// status = storageApi.RMW_ObjectStore(ref keyBytes, ref input, ref output); +// Debug.Assert(!output.spanByteAndMemory.IsSpanByte); + +// switch (status) +// { +// case GarnetStatus.WRONGTYPE: +// while (!RespWriteUtils.WriteError(CmdStrings.RESP_ERR_WRONG_TYPE, ref dcurr, dend)) +// SendAndReset(); +// break; +// default: +// if (output.spanByteAndMemory.Memory != null) +// SendAndReset(output.spanByteAndMemory.Memory, output.spanByteAndMemory.Length); +// else +// while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_OK, ref dcurr, dend)) +// SendAndReset(); +// break; +// } +// } +// else +// { +// status = storageApi.Read_ObjectStore(ref keyBytes, ref input, ref output); +// Debug.Assert(!output.spanByteAndMemory.IsSpanByte); + +// switch (status) +// { +// case GarnetStatus.OK: +// if (output.spanByteAndMemory.Memory != null) +// SendAndReset(output.spanByteAndMemory.Memory, output.spanByteAndMemory.Length); +// else +// while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_OK, ref dcurr, dend)) +// SendAndReset(); +// break; +// case GarnetStatus.NOTFOUND: +// Debug.Assert(output.spanByteAndMemory.Memory == null); +// while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_ERRNOTFOUND, ref dcurr, dend)) +// SendAndReset(); +// break; +// case GarnetStatus.WRONGTYPE: +// while (!RespWriteUtils.WriteError(CmdStrings.RESP_ERR_WRONG_TYPE, ref dcurr, dend)) +// SendAndReset(); +// break; +// } +// } + +// return true; +// } + + +// public unsafe GarnetStatus CustomObjectCommand(ArgSlice key, ArgSlice[] elements, ListOperation lop, out int itemsDoneCount, ref TObjectContext objectStoreContext) +// where TObjectContext : ITsavoriteContext +// { +// itemsDoneCount = 0; + +// if (key.Length == 0 || elements.Length == 0) +// return GarnetStatus.OK; + +// // Prepare the parse state +// var parseState = new SessionParseState(); +// ArgSlice[] parseStateBuffer = default; +// parseState.InitializeWithArguments(ref parseStateBuffer, elements); + +// // Prepare the input +// var input = new ObjectInput +// { +// header = new RespInputHeader +// { +// type = GarnetObjectType.List, +// ListOp = lop, +// }, +// parseState = parseState, +// parseStateStartIdx = 0, +// }; + +// var arrKey = key.ToArray(); +// var status = RMWObjectStoreOperation(arrKey, ref input, out var output, ref objectStoreContext); + +// itemsDoneCount = output.result1; +// itemBroker.HandleCollectionUpdate(arrKey); +// return status; +// } + +// } +//} diff --git a/libs/server/Storage/Session/MainStore/MainStoreOps.cs b/libs/server/Storage/Session/MainStore/MainStoreOps.cs index bb1588ee12e..991be07e805 100644 --- a/libs/server/Storage/Session/MainStore/MainStoreOps.cs +++ b/libs/server/Storage/Session/MainStore/MainStoreOps.cs @@ -1135,8 +1135,8 @@ public unsafe GarnetStatus CustomCommand(byte id, ArgSlice key, ArgSli byte* pcurr = pbCmdInput; *(int*)pcurr = inputSize - sizeof(int); pcurr += sizeof(int); - (*(RespInputHeader*)(pcurr)).cmd = (RespCommand)(id + CustomCommandManager.StartOffset); - (*(RespInputHeader*)(pcurr)).flags = 0; + (*(RespInputHeader*)pcurr).cmd = (RespCommand)(id + CustomCommandManager.StartOffset); + (*(RespInputHeader*)pcurr).flags = 0; pcurr += RespInputHeader.Size; Buffer.MemoryCopy(input.ptr, pcurr, input.Length, input.Length); diff --git a/main/GarnetServer/Extensions/ProcCustomCmd.cs b/main/GarnetServer/Extensions/ProcCustomCmd.cs index 468ff850bfe..ad22ff206b6 100644 --- a/main/GarnetServer/Extensions/ProcCustomCmd.cs +++ b/main/GarnetServer/Extensions/ProcCustomCmd.cs @@ -16,7 +16,8 @@ public override unsafe bool Execute(IGarnetApi garnetApi, ArgSlice input, ref Me var cmdOutput = new SpanByteAndMemory(null); // id from registration of custom raw string cmd - garnetApi.CustomCommand(0, key, new ArgSlice(input.ptr + offset, input.length - offset), ref cmdOutput); + //garnetApi.CustomCommand(0, key, new ArgSlice(input.ptr + offset, input.length - offset), ref cmdOutput); + InvokeCustomRawStringCommand(garnetApi, "SETIFPM", key, new ArgSlice(input.ptr + offset, input.length - offset)); return true; } From 2f2715cc3313c6ba12a9fc3f4101d270ca024272 Mon Sep 17 00:00:00 2001 From: Yoganand Rajasekaran Date: Mon, 21 Oct 2024 11:59:14 -0700 Subject: [PATCH 04/12] API changes. --- libs/server/API/GarnetApi.cs | 4 +- libs/server/API/GarnetWatchApi.cs | 4 +- libs/server/API/IGarnetApi.cs | 2 +- libs/server/Custom/CustomFunctions.cs | 16 +---- libs/server/Custom/CustomRespCommands.cs | 60 +++++++++++++++++- .../Storage/Session/MainStore/MainStoreOps.cs | 63 +++++++++++++++---- main/GarnetServer/Extensions/ProcCustomCmd.cs | 12 +++- 7 files changed, 125 insertions(+), 36 deletions(-) diff --git a/libs/server/API/GarnetApi.cs b/libs/server/API/GarnetApi.cs index 282f7d36045..03b414eacd1 100644 --- a/libs/server/API/GarnetApi.cs +++ b/libs/server/API/GarnetApi.cs @@ -417,7 +417,7 @@ public bool ResetScratchBuffer(int offset) => storageSession.scratchBufferManager.ResetScratchBuffer(offset); #endregion - public unsafe GarnetStatus CustomCommand(byte id, ArgSlice key, ArgSlice input, ref SpanByteAndMemory output) - => storageSession.CustomCommand(id, key, input, ref output, ref context); + public unsafe GarnetStatus CustomCommand(RespCommand cmd, ArgSlice key, ArgSlice[] args, CommandType cmdType, out ArgSlice output) + => storageSession.CustomCommand(cmd, key, args, cmdType, out output, ref context); } } \ No newline at end of file diff --git a/libs/server/API/GarnetWatchApi.cs b/libs/server/API/GarnetWatchApi.cs index a23df550c3b..7cd3d12a980 100644 --- a/libs/server/API/GarnetWatchApi.cs +++ b/libs/server/API/GarnetWatchApi.cs @@ -563,10 +563,10 @@ public int GetScratchBufferOffset() public bool ResetScratchBuffer(int offset) => garnetApi.ResetScratchBuffer(offset); - public GarnetStatus CustomCommand(byte id, ArgSlice key, ArgSlice input, ref SpanByteAndMemory output) + public GarnetStatus CustomCommand(RespCommand cmd, ArgSlice key, ArgSlice[] input, CommandType cmdType, out ArgSlice output) { garnetApi.WATCH(key, StoreType.Main); - return garnetApi.CustomCommand(id, key, input, ref output); + return garnetApi.CustomCommand(cmd, key, input, cmdType, out output); } #endregion diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs index a0fc9ba5c99..071e9cebb51 100644 --- a/libs/server/API/IGarnetApi.cs +++ b/libs/server/API/IGarnetApi.cs @@ -1734,7 +1734,7 @@ public bool IterateObjectStore(ref TScanFunctions scanFunctions, #endregion - public GarnetStatus CustomCommand(byte id, ArgSlice key, ArgSlice input, ref SpanByteAndMemory output); + public GarnetStatus CustomCommand(RespCommand cmd, ArgSlice key, ArgSlice[] input, CommandType cmdType, out ArgSlice output); } /// diff --git a/libs/server/Custom/CustomFunctions.cs b/libs/server/Custom/CustomFunctions.cs index 45081e075e2..3fb2ed9e047 100644 --- a/libs/server/Custom/CustomFunctions.cs +++ b/libs/server/Custom/CustomFunctions.cs @@ -7,7 +7,6 @@ using System.Diagnostics; using System.Text; using Garnet.common; -using Tsavorite.core; namespace Garnet.server { @@ -213,22 +212,13 @@ protected static unsafe ArgSlice GetNextArg(ref CustomProcedureInput procInput, return GetNextArg(ref procInput.parseState, procInput.parseStateFirstArgIdx, ref offset); } - protected void InvokeCustomRawStringCommand(TGarnetApi garnetApi, string cmd, ArgSlice key, ArgSlice input) + protected void InvokeCustomRawStringCommand(TGarnetApi garnetApi, string cmd, ArgSlice key, ArgSlice[] input, out ArgSlice output) where TGarnetApi : IGarnetApi { - - //var request = scratchBufferManager.FormatCommandAsResp(cmd, new[] { key, input }, null); - //_ = respServerSession.TryConsumeMessages(request.ptr, request.length); - //var response = scratchBufferNetworkSender.GetResponse(); - //var result = ProcessResponse(response.ptr, response.length); - //scratchBufferNetworkSender.Reset(); - //output = result; - - //output = - var output = new SpanByteAndMemory(null); if (customCommandManager.Match(new ReadOnlySpan(Encoding.UTF8.GetBytes(cmd)), out CustomRawStringCommand customCommand)) { - garnetApi.CustomCommand(customCommand.id, key, input, ref output); + //garnetApi.CustomCommand(customCommand.GetRespCommand(), key, input, customCommand.type, out output); + } } } diff --git a/libs/server/Custom/CustomRespCommands.cs b/libs/server/Custom/CustomRespCommands.cs index d625f667338..ed76bde3be5 100644 --- a/libs/server/Custom/CustomRespCommands.cs +++ b/libs/server/Custom/CustomRespCommands.cs @@ -16,7 +16,7 @@ internal sealed unsafe partial class RespServerSession : ServerSessionBase { private bool TryTransactionProc(byte id, CustomTransactionProcedure proc, int parseStateFirstArgIdx = 0, int parseStateLastArgIdx = -1) { - // Define output + // Define _output var output = new MemoryResult(null, 0); // Run procedure @@ -27,7 +27,7 @@ private bool TryTransactionProc(byte id, CustomTransactionProcedure proc, int pa var procInput = new CustomProcedureInput(ref parseState, parseStateFirstArgIdx, parseStateLastArgIdx); if (txnManager.RunTransactionProc(id, ref procInput, proc, ref output)) { - // Write output to wire + // Write _output to wire if (output.MemoryOwner != null) SendAndReset(output.MemoryOwner, output.Length); else @@ -36,7 +36,7 @@ private bool TryTransactionProc(byte id, CustomTransactionProcedure proc, int pa } else { - // Write output to wire + // Write _output to wire if (output.MemoryOwner != null) SendAndReset(output.MemoryOwner, output.Length); else @@ -81,6 +81,60 @@ private void TryCustomProcedure(CustomProcedure proc, int parseStateFirstArgIdx } } + public bool TryCustomRawStringCommandImpl(RespCommand cmd, long expirationTicks, CommandType type, ref TGarnetApi storageApi, ArgSlice key, ArgSlice[] args, CommandType cmdType, out ArgSlice output) + where TGarnetApi : IGarnetAdvancedApi + { + var sbKey = key.SpanByte; + + var inputArg = expirationTicks > 0 ? DateTimeOffset.UtcNow.Ticks + expirationTicks : expirationTicks; + var sessionParseState = new SessionParseState(); + sessionParseState.InitializeWithArguments(args); + var rawStringInput = new RawStringInput(cmd, ref sessionParseState, 1, -1, inputArg); + + var _output = new SpanByteAndMemory(null); + GarnetStatus status; + if (type == CommandType.ReadModifyWrite) + { + status = storageApi.RMW_MainStore(ref sbKey, ref rawStringInput, ref _output); + Debug.Assert(!_output.IsSpanByte); + + if (_output.Memory != null) + { + output = scratchBufferManager.FormatScratch(0, _output.AsReadOnlySpan()); + _output.Memory.Dispose(); + } + else + { + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_OK); + } + } + else + { + status = storageApi.Read_MainStore(ref sbKey, ref rawStringInput, ref _output); + Debug.Assert(!_output.IsSpanByte); + + if (status == GarnetStatus.OK) + { + if (_output.Memory != null) + { + output = scratchBufferManager.FormatScratch(0, _output.AsReadOnlySpan()); + _output.Memory.Dispose(); + } + else + { + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_OK); + } + } + else + { + Debug.Assert(_output.Memory == null); + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERRNOTFOUND); + } + } + + return true; + } + /// /// Custom command /// diff --git a/libs/server/Storage/Session/MainStore/MainStoreOps.cs b/libs/server/Storage/Session/MainStore/MainStoreOps.cs index 991be07e805..e5e5ef8f78f 100644 --- a/libs/server/Storage/Session/MainStore/MainStoreOps.cs +++ b/libs/server/Storage/Session/MainStore/MainStoreOps.cs @@ -1124,23 +1124,62 @@ public unsafe GarnetStatus SCAN(long cursor, ArgSlice match, long coun return GarnetStatus.OK; } - public unsafe GarnetStatus CustomCommand(byte id, ArgSlice key, ArgSlice input, ref SpanByteAndMemory output, ref TContext context) - where TContext : ITsavoriteContext + public unsafe GarnetStatus CustomCommand(RespCommand cmd, ArgSlice key, ArgSlice[] args, CommandType type, out ArgSlice value, ref TContext context) + where TContext : ITsavoriteContext { var sbKey = key.SpanByte; - int inputSize = sizeof(int) + RespInputHeader.Size + input.Length; - byte* pbCmdInput = stackalloc byte[inputSize]; + var sessionParseState = new SessionParseState(); + sessionParseState.InitializeWithArguments(args); + + var rawStringInput = new RawStringInput(cmd, ref sessionParseState); - byte* pcurr = pbCmdInput; - *(int*)pcurr = inputSize - sizeof(int); - pcurr += sizeof(int); - (*(RespInputHeader*)pcurr).cmd = (RespCommand)(id + CustomCommandManager.StartOffset); - (*(RespInputHeader*)pcurr).flags = 0; - pcurr += RespInputHeader.Size; - Buffer.MemoryCopy(input.ptr, pcurr, input.Length, input.Length); + var _output = new SpanByteAndMemory { SpanByte = scratchBufferManager.ViewRemainingArgSlice().SpanByte }; + + GarnetStatus status; + if (type == CommandType.ReadModifyWrite) + { + status = RMW_MainStore(ref sbKey, ref rawStringInput, ref _output, ref context); + //Debug.Assert(!output.IsSpanByte); + + if (output.Memory != null) + SendAndReset(output.Memory, output.Length); + else + while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_OK, ref dcurr, dend)) + SendAndReset(); + } + else + { + status = Read_MainStore(ref sbKey, ref rawStringInput, ref _output, ref context); + //Debug.Assert(!_output.IsSpanByte); // why? + + if (status == GarnetStatus.OK) + if (!_output.IsSpanByte) + { + value = scratchBufferManager.FormatScratch(0, _output.AsReadOnlySpan()); + _output.Memory.Dispose(); + } + else + { + value = scratchBufferManager.CreateArgSlice(_output.Length); + } + } + + + if (ret == GarnetStatus.OK) + { + if (!_output.IsSpanByte) + { + value = scratchBufferManager.FormatScratch(0, _output.AsReadOnlySpan()); + _output.Memory.Dispose(); + } + else + { + value = scratchBufferManager.CreateArgSlice(_output.Length); + } + } - return RMW_MainStore(ref sbKey, ref Unsafe.AsRef(pbCmdInput), ref output, ref context); + return RMW_MainStore(ref sbKey, ref rawStringInput, ref output, ref context); } public GarnetStatus GetKeyType(ArgSlice key, out string keyType, ref TContext context, ref TObjectContext objectContext) diff --git a/main/GarnetServer/Extensions/ProcCustomCmd.cs b/main/GarnetServer/Extensions/ProcCustomCmd.cs index ad22ff206b6..3ade4e8f5b2 100644 --- a/main/GarnetServer/Extensions/ProcCustomCmd.cs +++ b/main/GarnetServer/Extensions/ProcCustomCmd.cs @@ -9,15 +9,21 @@ namespace Garnet { class ProcCustomCmd : CustomProcedure { - public override unsafe bool Execute(IGarnetApi garnetApi, ArgSlice input, ref MemoryResult output) + public override unsafe bool Execute(TGarnetApi garnetApi, ref CustomProcedureInput procInput, ref MemoryResult output) { var offset = 0; - ArgSlice key = GetNextArg(input, ref offset); + ArgSlice key = GetNextArg(ref procInput, ref offset); + var cmdOutput = new SpanByteAndMemory(null); + ArgSlice[] args = new ArgSlice[procInput.parseState.Count - 1]; + for (int i = 0; i < procInput.parseState.Count - 1; i++) + { + args[i] = GetNextArg(ref procInput, ref offset); + } // id from registration of custom raw string cmd //garnetApi.CustomCommand(0, key, new ArgSlice(input.ptr + offset, input.length - offset), ref cmdOutput); - InvokeCustomRawStringCommand(garnetApi, "SETIFPM", key, new ArgSlice(input.ptr + offset, input.length - offset)); + InvokeCustomRawStringCommand(garnetApi, "SETIFPM", key, args); return true; } From 57c55d80104a21bfe6d6fadc975705ba4a0ab9f5 Mon Sep 17 00:00:00 2001 From: Yoganand Rajasekaran Date: Mon, 21 Oct 2024 21:25:30 -0700 Subject: [PATCH 05/12] Added custom object command support. --- libs/server/Custom/CustomCommandManager.cs | 2 +- .../Custom/CustomCommandManagerSession.cs | 22 ++++- libs/server/Custom/CustomFunctions.cs | 17 ++-- libs/server/Custom/CustomProcedureWrapper.cs | 7 +- libs/server/Custom/CustomRespCommands.cs | 75 ++++++++++++++- libs/server/Module/ModuleRegistrar.cs | 2 +- libs/server/Resp/RespServerSession.cs | 4 +- libs/server/Servers/RegisterApi.cs | 2 +- .../Storage/Session/MainStore/MainStoreOps.cs | 91 ++++++++++--------- libs/server/Transaction/TxnRespCommands.cs | 2 +- main/GarnetServer/Extensions/ProcCustomCmd.cs | 6 +- main/GarnetServer/Extensions/TxnCustomCmd.cs | 38 ++++++++ main/GarnetServer/Program.cs | 7 +- playground/SampleModule/SampleModule.cs | 2 +- test/Garnet.test/Resp/ACL/RespCommandTests.cs | 2 +- test/Garnet.test/RespAdminCommandsTests.cs | 4 +- test/Garnet.test/RespAofTests.cs | 4 +- test/Garnet.test/RespCustomCommandTests.cs | 6 +- 18 files changed, 208 insertions(+), 85 deletions(-) create mode 100644 main/GarnetServer/Extensions/TxnCustomCmd.cs diff --git a/libs/server/Custom/CustomCommandManager.cs b/libs/server/Custom/CustomCommandManager.cs index 8ad42bea522..6cdb1504fd9 100644 --- a/libs/server/Custom/CustomCommandManager.cs +++ b/libs/server/Custom/CustomCommandManager.cs @@ -164,7 +164,7 @@ internal void RegisterType(int objectTypeId, CustomObjectFactory factory) /// /// /// - internal int Register(string name, CustomProcedure customProcedure, RespCommandsInfo commandInfo = null, RespCommandDocs commandDocs = null) + internal int Register(string name, Func customProcedure, RespCommandsInfo commandInfo = null, RespCommandDocs commandDocs = null) { int id = Interlocked.Increment(ref CustomProcedureId) - 1; if (id >= MaxRegistrations) diff --git a/libs/server/Custom/CustomCommandManagerSession.cs b/libs/server/Custom/CustomCommandManagerSession.cs index 384cf75125b..342111e66a8 100644 --- a/libs/server/Custom/CustomCommandManagerSession.cs +++ b/libs/server/Custom/CustomCommandManagerSession.cs @@ -12,25 +12,40 @@ internal sealed class CustomCommandManagerSession { readonly CustomCommandManager customCommandManager; public readonly (CustomTransactionProcedure, int)[] sessionTransactionProcMap; + public readonly CustomProcedure[] sessionCustomProcMap; + public CustomCommandManagerSession(CustomCommandManager customCommandManager) { this.customCommandManager = customCommandManager; sessionTransactionProcMap = new (CustomTransactionProcedure, int)[CustomCommandManager.MaxRegistrations]; + sessionCustomProcMap = new CustomProcedure[CustomCommandManager.MaxRegistrations]; + } + + public CustomProcedure GetCustomProcedure(int id, RespServerSession respServerSession) + { + if (sessionCustomProcMap[id] == null) + { + var entry = customCommandManager.customProcedureMap[id] ?? throw new GarnetException($"Custom procedure {id} not found"); + sessionCustomProcMap[id] = entry.CustomProcedure(); + sessionCustomProcMap[id].respServerSession = respServerSession; + } + + return sessionCustomProcMap[id]; } - public (CustomTransactionProcedure, int) GetCustomTransactionProcedure(int id, TransactionManager txnManager, ScratchBufferManager scratchBufferManager) + public (CustomTransactionProcedure, int) GetCustomTransactionProcedure(int id, RespServerSession respServerSession, TransactionManager txnManager, ScratchBufferManager scratchBufferManager) { if (sessionTransactionProcMap[id].Item1 == null) { var entry = customCommandManager.transactionProcMap[id] ?? throw new GarnetException($"Transaction procedure {id} not found"); _ = customCommandManager.CustomCommandsInfo.TryGetValue(entry.NameStr, out var cmdInfo); - return GetCustomTransactionProcedure(entry, txnManager, scratchBufferManager, cmdInfo?.Arity ?? 0); + return GetCustomTransactionProcedure(entry, respServerSession, txnManager, scratchBufferManager, cmdInfo?.Arity ?? 0); } return sessionTransactionProcMap[id]; } - public (CustomTransactionProcedure, int) GetCustomTransactionProcedure(CustomTransaction entry, TransactionManager txnManager, ScratchBufferManager scratchBufferManager, int arity) + public (CustomTransactionProcedure, int) GetCustomTransactionProcedure(CustomTransaction entry, RespServerSession respServerSession, TransactionManager txnManager, ScratchBufferManager scratchBufferManager, int arity) { int id = entry.id; if (sessionTransactionProcMap[id].Item1 == null) @@ -40,6 +55,7 @@ public CustomCommandManagerSession(CustomCommandManager customCommandManager) sessionTransactionProcMap[id].Item1.txnManager = txnManager; sessionTransactionProcMap[id].Item1.scratchBufferManager = scratchBufferManager; + sessionTransactionProcMap[id].Item1.respServerSession = respServerSession; } return sessionTransactionProcMap[id]; } diff --git a/libs/server/Custom/CustomFunctions.cs b/libs/server/Custom/CustomFunctions.cs index 3fb2ed9e047..0b394f7b408 100644 --- a/libs/server/Custom/CustomFunctions.cs +++ b/libs/server/Custom/CustomFunctions.cs @@ -5,7 +5,6 @@ using System.Buffers; using System.Collections.Generic; using System.Diagnostics; -using System.Text; using Garnet.common; namespace Garnet.server @@ -20,7 +19,9 @@ public abstract class CustomFunctions /// protected static MemoryPool MemoryPool => MemoryPool.Shared; - public CustomCommandManager customCommandManager; + //public CustomCommandManager customCommandManager; + + internal RespServerSession respServerSession; /// /// Create output as simple string, from given string @@ -212,14 +213,16 @@ protected static unsafe ArgSlice GetNextArg(ref CustomProcedureInput procInput, return GetNextArg(ref procInput.parseState, procInput.parseStateFirstArgIdx, ref offset); } - protected void InvokeCustomRawStringCommand(TGarnetApi garnetApi, string cmd, ArgSlice key, ArgSlice[] input, out ArgSlice output) + protected void ExecuteCustomRawStringCommand(TGarnetApi garnetApi, string cmd, ArgSlice key, ArgSlice[] input, out ArgSlice output) where TGarnetApi : IGarnetApi { - if (customCommandManager.Match(new ReadOnlySpan(Encoding.UTF8.GetBytes(cmd)), out CustomRawStringCommand customCommand)) - { - //garnetApi.CustomCommand(customCommand.GetRespCommand(), key, input, customCommand.type, out output); + respServerSession.InvokeCustomRawStringCommand(ref garnetApi, cmd, key, input, out output); + } - } + protected void ExecuteCustomObjectCommand(TGarnetApi garnetApi, string cmd, ArgSlice key, ArgSlice[] input, out ArgSlice output) + where TGarnetApi : IGarnetApi + { + respServerSession.InvokeCustomObjectCommand(ref garnetApi, cmd, key, input, out output); } } } \ No newline at end of file diff --git a/libs/server/Custom/CustomProcedureWrapper.cs b/libs/server/Custom/CustomProcedureWrapper.cs index 07da96a371b..9b97d33e963 100644 --- a/libs/server/Custom/CustomProcedureWrapper.cs +++ b/libs/server/Custom/CustomProcedureWrapper.cs @@ -73,9 +73,9 @@ class CustomProcedureWrapper public readonly string NameStr; public readonly byte[] Name; public readonly byte Id; - public readonly CustomProcedure CustomProcedureImpl; + public readonly Func CustomProcedure; - internal CustomProcedureWrapper(string name, byte id, CustomProcedure customProcedure, CustomCommandManager customCommandManager) + internal CustomProcedureWrapper(string name, byte id, Func customProcedure, CustomCommandManager customCommandManager) { if (string.IsNullOrEmpty(name)) throw new ArgumentNullException(nameof(name)); @@ -88,8 +88,7 @@ internal CustomProcedureWrapper(string name, byte id, CustomProcedure customProc NameStr = name.ToUpperInvariant(); Name = System.Text.Encoding.ASCII.GetBytes(NameStr); Id = id; - CustomProcedureImpl = customProcedure; - CustomProcedureImpl.customCommandManager = customCommandManager; + CustomProcedure = customProcedure; } } } \ No newline at end of file diff --git a/libs/server/Custom/CustomRespCommands.cs b/libs/server/Custom/CustomRespCommands.cs index ed76bde3be5..f038df300c4 100644 --- a/libs/server/Custom/CustomRespCommands.cs +++ b/libs/server/Custom/CustomRespCommands.cs @@ -3,6 +3,7 @@ using System; using System.Diagnostics; +using System.Text; using Garnet.common; using Tsavorite.core; @@ -51,7 +52,7 @@ private bool TryTransactionProc(byte id, CustomTransactionProcedure proc, int pa public bool RunTransactionProc(byte id, ref CustomProcedureInput procInput, ref MemoryResult output) { var proc = customCommandManagerSession - .GetCustomTransactionProcedure(id, txnManager, scratchBufferManager).Item1; + .GetCustomTransactionProcedure(id, this, txnManager, scratchBufferManager).Item1; return txnManager.RunTransactionProc(id, ref procInput, proc, ref output); } @@ -81,19 +82,22 @@ private void TryCustomProcedure(CustomProcedure proc, int parseStateFirstArgIdx } } - public bool TryCustomRawStringCommandImpl(RespCommand cmd, long expirationTicks, CommandType type, ref TGarnetApi storageApi, ArgSlice key, ArgSlice[] args, CommandType cmdType, out ArgSlice output) + public bool InvokeCustomRawStringCommand(ref TGarnetApi storageApi, string cmd, ArgSlice key, ArgSlice[] args, out ArgSlice output) where TGarnetApi : IGarnetAdvancedApi { + // TODO: check if cmd exists + storeWrapper.customCommandManager.Match(new ReadOnlySpan(Encoding.UTF8.GetBytes(cmd)), out CustomRawStringCommand customCommand); + var sbKey = key.SpanByte; - var inputArg = expirationTicks > 0 ? DateTimeOffset.UtcNow.Ticks + expirationTicks : expirationTicks; + var inputArg = customCommand.expirationTicks > 0 ? DateTimeOffset.UtcNow.Ticks + customCommand.expirationTicks : customCommand.expirationTicks; var sessionParseState = new SessionParseState(); sessionParseState.InitializeWithArguments(args); - var rawStringInput = new RawStringInput(cmd, ref sessionParseState, 1, -1, inputArg); + var rawStringInput = new RawStringInput(customCommand.GetRespCommand(), ref sessionParseState, arg1: inputArg); var _output = new SpanByteAndMemory(null); GarnetStatus status; - if (type == CommandType.ReadModifyWrite) + if (customCommand.type == CommandType.ReadModifyWrite) { status = storageApi.RMW_MainStore(ref sbKey, ref rawStringInput, ref _output); Debug.Assert(!_output.IsSpanByte); @@ -183,6 +187,67 @@ private bool TryCustomRawStringCommand(RespCommand cmd, long expirat return true; } + public bool InvokeCustomObjectCommand(ref TGarnetApi storageApi, string cmd, ArgSlice key, ArgSlice[] args, out ArgSlice output) + where TGarnetApi : IGarnetAdvancedApi + { + output = default; + // TODO: check if cmd exists + storeWrapper.customCommandManager.Match(new ReadOnlySpan(Encoding.UTF8.GetBytes(cmd)), out CustomObjectCommand customObjCommand); + + var keyBytes = key.ToArray(); + + // Prepare input + var header = new RespInputHeader(customObjCommand.GetRespCommand()) { SubId = customObjCommand.subid }; + var sessionParseState = new SessionParseState(); + sessionParseState.InitializeWithArguments(args); + var input = new ObjectInput(header, ref sessionParseState); + + var _output = new GarnetObjectStoreOutput { spanByteAndMemory = new SpanByteAndMemory(null) }; + GarnetStatus status; + if (customObjCommand.type == CommandType.ReadModifyWrite) + { + status = storageApi.RMW_ObjectStore(ref keyBytes, ref input, ref _output); + Debug.Assert(!_output.spanByteAndMemory.IsSpanByte); + + switch (status) + { + case GarnetStatus.WRONGTYPE: + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERR_WRONG_TYPE); + break; + default: + if (_output.spanByteAndMemory.Memory != null) + output = scratchBufferManager.FormatScratch(0, _output.spanByteAndMemory.AsReadOnlySpan()); + else + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_OK); + break; + } + } + else + { + status = storageApi.Read_ObjectStore(ref keyBytes, ref input, ref _output); + Debug.Assert(!_output.spanByteAndMemory.IsSpanByte); + + switch (status) + { + case GarnetStatus.OK: + if (_output.spanByteAndMemory.Memory != null) + output = scratchBufferManager.FormatScratch(0, _output.spanByteAndMemory.AsReadOnlySpan()); + else + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_OK); + break; + case GarnetStatus.NOTFOUND: + Debug.Assert(_output.spanByteAndMemory.Memory == null); + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERRNOTFOUND); + break; + case GarnetStatus.WRONGTYPE: + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERR_WRONG_TYPE); + break; + } + } + + return true; + } + /// /// Custom object command /// diff --git a/libs/server/Module/ModuleRegistrar.cs b/libs/server/Module/ModuleRegistrar.cs index 90dfa9e3de9..891b2df33d9 100644 --- a/libs/server/Module/ModuleRegistrar.cs +++ b/libs/server/Module/ModuleRegistrar.cs @@ -165,7 +165,7 @@ public ModuleActionStatus RegisterCommand(string name, CustomObjectFactory facto /// Command info /// Command docs /// Registration status - public ModuleActionStatus RegisterProcedure(string name, CustomProcedure customScriptProc, RespCommandsInfo commandInfo = null, RespCommandDocs commandDocs = null) + public ModuleActionStatus RegisterProcedure(string name, Func customScriptProc, RespCommandsInfo commandInfo = null, RespCommandDocs commandDocs = null) { if (string.IsNullOrEmpty(name) || customScriptProc == null) return ModuleActionStatus.InvalidRegistrationInfo; diff --git a/libs/server/Resp/RespServerSession.cs b/libs/server/Resp/RespServerSession.cs index ef7d28d8aca..0320311e7fd 100644 --- a/libs/server/Resp/RespServerSession.cs +++ b/libs/server/Resp/RespServerSession.cs @@ -751,7 +751,7 @@ bool NetworkCustomTxn() // Perform the operation TryTransactionProc(currentCustomTransaction.id, customCommandManagerSession - .GetCustomTransactionProcedure(currentCustomTransaction.id, txnManager, scratchBufferManager) + .GetCustomTransactionProcedure(currentCustomTransaction.id, this, txnManager, scratchBufferManager) .Item1); currentCustomTransaction = null; return true; @@ -765,7 +765,7 @@ bool NetworkCustomProcedure() return true; } - TryCustomProcedure(currentCustomProcedure.CustomProcedureImpl); + TryCustomProcedure(customCommandManagerSession.GetCustomProcedure(currentCustomProcedure.Id, this)); currentCustomProcedure = null; return true; diff --git a/libs/server/Servers/RegisterApi.cs b/libs/server/Servers/RegisterApi.cs index 86437d5c2bc..ed2995280bb 100644 --- a/libs/server/Servers/RegisterApi.cs +++ b/libs/server/Servers/RegisterApi.cs @@ -86,7 +86,7 @@ public void NewType(int type, CustomObjectFactory factory) /// /// /// - public int NewProcedure(string name, CustomProcedure customProcedure, RespCommandsInfo commandInfo = null, RespCommandDocs commandDocs = null) + public int NewProcedure(string name, Func customProcedure, RespCommandsInfo commandInfo = null, RespCommandDocs commandDocs = null) => provider.StoreWrapper.customCommandManager.Register(name, customProcedure, commandInfo, commandDocs); } } \ No newline at end of file diff --git a/libs/server/Storage/Session/MainStore/MainStoreOps.cs b/libs/server/Storage/Session/MainStore/MainStoreOps.cs index e5e5ef8f78f..f867807e4dd 100644 --- a/libs/server/Storage/Session/MainStore/MainStoreOps.cs +++ b/libs/server/Storage/Session/MainStore/MainStoreOps.cs @@ -1128,6 +1128,7 @@ public unsafe GarnetStatus CustomCommand(RespCommand cmd, ArgSlice key where TContext : ITsavoriteContext { var sbKey = key.SpanByte; + value = default; var sessionParseState = new SessionParseState(); sessionParseState.InitializeWithArguments(args); @@ -1135,51 +1136,51 @@ public unsafe GarnetStatus CustomCommand(RespCommand cmd, ArgSlice key var rawStringInput = new RawStringInput(cmd, ref sessionParseState); var _output = new SpanByteAndMemory { SpanByte = scratchBufferManager.ViewRemainingArgSlice().SpanByte }; - - GarnetStatus status; - if (type == CommandType.ReadModifyWrite) - { - status = RMW_MainStore(ref sbKey, ref rawStringInput, ref _output, ref context); - //Debug.Assert(!output.IsSpanByte); - - if (output.Memory != null) - SendAndReset(output.Memory, output.Length); - else - while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_OK, ref dcurr, dend)) - SendAndReset(); - } - else - { - status = Read_MainStore(ref sbKey, ref rawStringInput, ref _output, ref context); - //Debug.Assert(!_output.IsSpanByte); // why? - - if (status == GarnetStatus.OK) - if (!_output.IsSpanByte) - { - value = scratchBufferManager.FormatScratch(0, _output.AsReadOnlySpan()); - _output.Memory.Dispose(); - } - else - { - value = scratchBufferManager.CreateArgSlice(_output.Length); - } - } - - - if (ret == GarnetStatus.OK) - { - if (!_output.IsSpanByte) - { - value = scratchBufferManager.FormatScratch(0, _output.AsReadOnlySpan()); - _output.Memory.Dispose(); - } - else - { - value = scratchBufferManager.CreateArgSlice(_output.Length); - } - } - - return RMW_MainStore(ref sbKey, ref rawStringInput, ref output, ref context); + return GarnetStatus.OK; + //GarnetStatus status; + //if (type == CommandType.ReadModifyWrite) + //{ + // status = RMW_MainStore(ref sbKey, ref rawStringInput, ref _output, ref context); + // //Debug.Assert(!output.IsSpanByte); + + // if (output.Memory != null) + // SendAndReset(output.Memory, output.Length); + // else + // while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_OK, ref dcurr, dend)) + // SendAndReset(); + //} + //else + //{ + // status = Read_MainStore(ref sbKey, ref rawStringInput, ref _output, ref context); + // //Debug.Assert(!_output.IsSpanByte); // why? + + // if (status == GarnetStatus.OK) + // if (!_output.IsSpanByte) + // { + // value = scratchBufferManager.FormatScratch(0, _output.AsReadOnlySpan()); + // _output.Memory.Dispose(); + // } + // else + // { + // value = scratchBufferManager.CreateArgSlice(_output.Length); + // } + //} + + + //if (ret == GarnetStatus.OK) + //{ + // if (!_output.IsSpanByte) + // { + // value = scratchBufferManager.FormatScratch(0, _output.AsReadOnlySpan()); + // _output.Memory.Dispose(); + // } + // else + // { + // value = scratchBufferManager.CreateArgSlice(_output.Length); + // } + //} + + //return RMW_MainStore(ref sbKey, ref rawStringInput, ref output, ref context); } public GarnetStatus GetKeyType(ArgSlice key, out string keyType, ref TContext context, ref TObjectContext objectContext) diff --git a/libs/server/Transaction/TxnRespCommands.cs b/libs/server/Transaction/TxnRespCommands.cs index 3841c09ea31..14186a0537d 100644 --- a/libs/server/Transaction/TxnRespCommands.cs +++ b/libs/server/Transaction/TxnRespCommands.cs @@ -266,7 +266,7 @@ private bool NetworkRUNTXP() try { - (proc, arity) = customCommandManagerSession.GetCustomTransactionProcedure(txId, txnManager, scratchBufferManager); + (proc, arity) = customCommandManagerSession.GetCustomTransactionProcedure(txId, this, txnManager, scratchBufferManager); } catch (Exception e) { diff --git a/main/GarnetServer/Extensions/ProcCustomCmd.cs b/main/GarnetServer/Extensions/ProcCustomCmd.cs index 3ade4e8f5b2..ca1530862b3 100644 --- a/main/GarnetServer/Extensions/ProcCustomCmd.cs +++ b/main/GarnetServer/Extensions/ProcCustomCmd.cs @@ -3,7 +3,6 @@ using Garnet.common; using Garnet.server; -using Tsavorite.core; namespace Garnet { @@ -14,7 +13,7 @@ public override unsafe bool Execute(TGarnetApi garnetApi, ref Custom var offset = 0; ArgSlice key = GetNextArg(ref procInput, ref offset); - var cmdOutput = new SpanByteAndMemory(null); + //var cmdOutput = new SpanByteAndMemory(null); ArgSlice[] args = new ArgSlice[procInput.parseState.Count - 1]; for (int i = 0; i < procInput.parseState.Count - 1; i++) @@ -23,8 +22,9 @@ public override unsafe bool Execute(TGarnetApi garnetApi, ref Custom } // id from registration of custom raw string cmd //garnetApi.CustomCommand(0, key, new ArgSlice(input.ptr + offset, input.length - offset), ref cmdOutput); - InvokeCustomRawStringCommand(garnetApi, "SETIFPM", key, args); + //InvokeCustomRawStringCommand(garnetApi, "SETIFPM", key, args); + ExecuteCustomRawStringCommand(garnetApi, "SETIFPM", key, args, out var _output); return true; } } diff --git a/main/GarnetServer/Extensions/TxnCustomCmd.cs b/main/GarnetServer/Extensions/TxnCustomCmd.cs new file mode 100644 index 00000000000..31c004c42d0 --- /dev/null +++ b/main/GarnetServer/Extensions/TxnCustomCmd.cs @@ -0,0 +1,38 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +using Garnet.common; +using Garnet.server; +using Tsavorite.core; + +namespace Garnet +{ + class TxnCustomCmd : CustomTransactionProcedure + { + public override void Main(TGarnetApi api, ref CustomProcedureInput procInput, ref MemoryResult output) => throw new System.NotImplementedException(); + public override bool Prepare(TGarnetReadApi api, ref CustomProcedureInput procInput) => false; + + public override void Finalize(TGarnetApi api, ref CustomProcedureInput procInput, ref MemoryResult output) + { + var offset = 0; + ArgSlice key = GetNextArg(ref procInput, ref offset); + + var cmdOutput = new SpanByteAndMemory(null); + + ArgSlice[] args = new ArgSlice[procInput.parseState.Count - 1]; + for (int i = 0; i < procInput.parseState.Count - 1; i++) + { + args[i] = GetNextArg(ref procInput, ref offset); + } + // id from registration of custom raw string cmd + //garnetApi.CustomCommand(0, key, new ArgSlice(input.ptr + offset, input.length - offset), ref cmdOutput); + //InvokeCustomRawStringCommand(garnetApi, "SETIFPM", key, args); + + //ExecuteCustomRawStringCommand(api, "SETIFPM", key, args, out var _output); + + ExecuteCustomObjectCommand(api, "MYDICTSET", key, args, out var _output); + + WriteSimpleString(ref output, "OK"); + } + } +} diff --git a/main/GarnetServer/Program.cs b/main/GarnetServer/Program.cs index fd5f03126c6..882902b47f5 100644 --- a/main/GarnetServer/Program.cs +++ b/main/GarnetServer/Program.cs @@ -92,10 +92,11 @@ static void RegisterExtensions(GarnetServer server) server.Register.NewTransactionProc("SAMPLEUPDATETX", () => new SampleUpdateTxn(), new RespCommandsInfo { Arity = 9 }); server.Register.NewTransactionProc("SAMPLEDELETETX", () => new SampleDeleteTxn(), new RespCommandsInfo { Arity = 6 }); - server.Register.NewProcedure("SUM", new Sum()); - server.Register.NewProcedure("SETMAINANDOBJECT", new SetStringAndList()); + server.Register.NewProcedure("SUM", () => new Sum()); + server.Register.NewProcedure("SETMAINANDOBJECT", () => new SetStringAndList()); - server.Register.NewProcedure("PROCCMD", new ProcCustomCmd()); + server.Register.NewProcedure("PROCCMD", () => new ProcCustomCmd()); + server.Register.NewTransactionProc("TXNCMD", () => new TxnCustomCmd()); } } } \ No newline at end of file diff --git a/playground/SampleModule/SampleModule.cs b/playground/SampleModule/SampleModule.cs index 7463d259339..96cea06277f 100644 --- a/playground/SampleModule/SampleModule.cs +++ b/playground/SampleModule/SampleModule.cs @@ -28,7 +28,7 @@ public override void OnLoad(ModuleLoadContext context, string[] args) context.RegisterCommand("SampleModule.MYDICTSET", factory, new MyDictSet()); context.RegisterCommand("SampleModule.MYDICTGET", factory, new MyDictGet(), CommandType.Read); - context.RegisterProcedure("SampleModule.SUM", new Sum()); + context.RegisterProcedure("SampleModule.SUM", () => new Sum()); } } } \ No newline at end of file diff --git a/test/Garnet.test/Resp/ACL/RespCommandTests.cs b/test/Garnet.test/Resp/ACL/RespCommandTests.cs index c71932497f1..5667f65b71b 100644 --- a/test/Garnet.test/Resp/ACL/RespCommandTests.cs +++ b/test/Garnet.test/Resp/ACL/RespCommandTests.cs @@ -38,7 +38,7 @@ public void Setup() server.Register.NewCommand("SETWPIFPGT", CommandType.ReadModifyWrite, new SetWPIFPGTCustomCommand(), respCustomCommandsInfo["SETWPIFPGT"]); server.Register.NewCommand("MYDICTGET", CommandType.Read, new MyDictFactory(), new MyDictGet(), respCustomCommandsInfo["MYDICTGET"]); server.Register.NewTransactionProc("READWRITETX", () => new ReadWriteTxn(), new RespCommandsInfo { Arity = 4 }); - server.Register.NewProcedure("SUM", new Sum()); + server.Register.NewProcedure("SUM", () => new Sum()); server.Start(); } diff --git a/test/Garnet.test/RespAdminCommandsTests.cs b/test/Garnet.test/RespAdminCommandsTests.cs index 19963ba9b3b..82249c5a7ac 100644 --- a/test/Garnet.test/RespAdminCommandsTests.cs +++ b/test/Garnet.test/RespAdminCommandsTests.cs @@ -279,7 +279,7 @@ static void ValidateServerData(IDatabase db, string strKey, string strValue, str var listValue = "ListValue"; // Register sample custom script that updates both main store and object store keys - server.Register.NewProcedure("SETMAINANDOBJECT", new SetStringAndList()); + server.Register.NewProcedure("SETMAINANDOBJECT", () => new SetStringAndList()); using (var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig(allowAdmin: true))) { @@ -295,7 +295,7 @@ static void ValidateServerData(IDatabase db, string strKey, string strValue, str server.Dispose(false); server = TestUtils.CreateGarnetServer(TestUtils.MethodTestDir, tryRecover: true); - server.Register.NewProcedure("SETMAINANDOBJECT", new SetStringAndList()); + server.Register.NewProcedure("SETMAINANDOBJECT", () => new SetStringAndList()); server.Start(); using (var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig(allowAdmin: true))) diff --git a/test/Garnet.test/RespAofTests.cs b/test/Garnet.test/RespAofTests.cs index 76271caea14..f5193b10b4f 100644 --- a/test/Garnet.test/RespAofTests.cs +++ b/test/Garnet.test/RespAofTests.cs @@ -539,7 +539,7 @@ static void ValidateServerData(IDatabase db, string strKey, string strValue, str server.Dispose(false); server = TestUtils.CreateGarnetServer(TestUtils.MethodTestDir, enableAOF: true); - server.Register.NewProcedure("SETMAINANDOBJECT", new SetStringAndList()); + server.Register.NewProcedure("SETMAINANDOBJECT", () => new SetStringAndList()); server.Start(); var strKey = "StrKey"; @@ -557,7 +557,7 @@ static void ValidateServerData(IDatabase db, string strKey, string strValue, str server.Store.CommitAOF(true); server.Dispose(false); server = TestUtils.CreateGarnetServer(TestUtils.MethodTestDir, tryRecover: true, enableAOF: true); - server.Register.NewProcedure("SETMAINANDOBJECT", new SetStringAndList()); + server.Register.NewProcedure("SETMAINANDOBJECT", () => new SetStringAndList()); server.Start(); using (var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig())) diff --git a/test/Garnet.test/RespCustomCommandTests.cs b/test/Garnet.test/RespCustomCommandTests.cs index 336897bc094..01b9cc17343 100644 --- a/test/Garnet.test/RespCustomCommandTests.cs +++ b/test/Garnet.test/RespCustomCommandTests.cs @@ -636,7 +636,7 @@ public async Task CustomCommandSetWithCustomExpirationTestAsync() [Test] public void CustomCommandRegistrationTest() { - server.Register.NewProcedure("SUM", new Sum()); + server.Register.NewProcedure("SUM", () => new Sum()); using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig()); var db = redis.GetDatabase(0); @@ -653,7 +653,7 @@ public void CustomCommandRegistrationTest() [Test] public void CustomProcedureFreeBufferTest() { - server.Register.NewProcedure("LARGEGET", new LargeGet()); + server.Register.NewProcedure("LARGEGET", () => new LargeGet()); var key = "key"; var hashKey = "hashKey"; var hashField = "field"; @@ -701,7 +701,7 @@ public void CustomTxnFreeBufferTest() [Test] public void CustomProcedureOutOfOrderFreeBufferTest() { - server.Register.NewProcedure("OUTOFORDERFREE", new OutOfOrderFreeBuffer()); + server.Register.NewProcedure("OUTOFORDERFREE", () => new OutOfOrderFreeBuffer()); var key = "key"; using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig()); var db = redis.GetDatabase(0); From 5175423265559817a62a2c0520883c746eff386d Mon Sep 17 00:00:00 2001 From: Yoganand Rajasekaran Date: Mon, 21 Oct 2024 21:35:15 -0700 Subject: [PATCH 06/12] Cleanup. --- libs/server/API/GarnetApi.cs | 3 - libs/server/API/GarnetWatchApi.cs | 6 - libs/server/API/IGarnetApi.cs | 1 - libs/server/ArgSlice/ArgSlice.cs | 4 +- libs/server/Custom/CustomFunctions.cs | 2 - libs/server/Custom/CustomProcedureWrapper.cs | 46 ------- libs/server/Custom/CustomRespCommands.cs | 6 +- libs/server/Storage/Session/CustomOps.cs | 121 ------------------ .../Storage/Session/MainStore/MainStoreOps.cs | 59 --------- main/GarnetServer/Extensions/ProcCustomCmd.cs | 3 - main/GarnetServer/Extensions/TxnCustomCmd.cs | 3 - 11 files changed, 5 insertions(+), 249 deletions(-) delete mode 100644 libs/server/Storage/Session/CustomOps.cs diff --git a/libs/server/API/GarnetApi.cs b/libs/server/API/GarnetApi.cs index 03b414eacd1..34c91d00a5f 100644 --- a/libs/server/API/GarnetApi.cs +++ b/libs/server/API/GarnetApi.cs @@ -416,8 +416,5 @@ public int GetScratchBufferOffset() public bool ResetScratchBuffer(int offset) => storageSession.scratchBufferManager.ResetScratchBuffer(offset); #endregion - - public unsafe GarnetStatus CustomCommand(RespCommand cmd, ArgSlice key, ArgSlice[] args, CommandType cmdType, out ArgSlice output) - => storageSession.CustomCommand(cmd, key, args, cmdType, out output, ref context); } } \ No newline at end of file diff --git a/libs/server/API/GarnetWatchApi.cs b/libs/server/API/GarnetWatchApi.cs index 7cd3d12a980..ef6d82266f3 100644 --- a/libs/server/API/GarnetWatchApi.cs +++ b/libs/server/API/GarnetWatchApi.cs @@ -563,12 +563,6 @@ public int GetScratchBufferOffset() public bool ResetScratchBuffer(int offset) => garnetApi.ResetScratchBuffer(offset); - public GarnetStatus CustomCommand(RespCommand cmd, ArgSlice key, ArgSlice[] input, CommandType cmdType, out ArgSlice output) - { - garnetApi.WATCH(key, StoreType.Main); - return garnetApi.CustomCommand(cmd, key, input, cmdType, out output); - } - #endregion } } \ No newline at end of file diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs index 071e9cebb51..e271407c63f 100644 --- a/libs/server/API/IGarnetApi.cs +++ b/libs/server/API/IGarnetApi.cs @@ -1734,7 +1734,6 @@ public bool IterateObjectStore(ref TScanFunctions scanFunctions, #endregion - public GarnetStatus CustomCommand(RespCommand cmd, ArgSlice key, ArgSlice[] input, CommandType cmdType, out ArgSlice output); } /// diff --git a/libs/server/ArgSlice/ArgSlice.cs b/libs/server/ArgSlice/ArgSlice.cs index 984d39b57bc..8ea0a889306 100644 --- a/libs/server/ArgSlice/ArgSlice.cs +++ b/libs/server/ArgSlice/ArgSlice.cs @@ -21,10 +21,10 @@ public unsafe struct ArgSlice public const int Size = 12; [FieldOffset(0)] - public byte* ptr; + internal byte* ptr; [FieldOffset(8)] - public int length; + internal int length; /// /// Create new ArgSlice from given pointer and length diff --git a/libs/server/Custom/CustomFunctions.cs b/libs/server/Custom/CustomFunctions.cs index 0b394f7b408..5504185146a 100644 --- a/libs/server/Custom/CustomFunctions.cs +++ b/libs/server/Custom/CustomFunctions.cs @@ -19,8 +19,6 @@ public abstract class CustomFunctions /// protected static MemoryPool MemoryPool => MemoryPool.Shared; - //public CustomCommandManager customCommandManager; - internal RespServerSession respServerSession; /// diff --git a/libs/server/Custom/CustomProcedureWrapper.cs b/libs/server/Custom/CustomProcedureWrapper.cs index 9b97d33e963..899349b135f 100644 --- a/libs/server/Custom/CustomProcedureWrapper.cs +++ b/libs/server/Custom/CustomProcedureWrapper.cs @@ -12,52 +12,6 @@ namespace Garnet.server /// public abstract class CustomProcedure : CustomFunctions { - internal ScratchBufferManager scratchBufferManager; - - //unsafe object CustomOperation(TGarnetApi api, string cmd, params object[] args) - // where TGarnetApi : IGarnetApi - //unsafe object CustomOperation(TGarnetApi api, string cmd, params object[] args) - // where TGarnetApi : IGarnetApi - //{ - //switch (cmd) - //{ - // // We special-case a few performance-sensitive operations to directly invoke via the storage API - // case "SET" when args.Length == 2: - // case "set" when args.Length == 2: - // { - // if (!respServerSession.CheckACLPermissions(RespCommand.SET)) - // return Encoding.ASCII.GetString(CmdStrings.RESP_ERR_NOAUTH); - // var key = scratchBufferManager.CreateArgSlice(Convert.ToString(args[0])); - // var value = scratchBufferManager.CreateArgSlice(Convert.ToString(args[1])); - // _ = api.SET(key, value); - // return "OK"; - // } - // case "GET": - // case "get": - // { - // if (!respServerSession.CheckACLPermissions(RespCommand.GET)) - // throw new Exception(Encoding.ASCII.GetString(CmdStrings.RESP_ERR_NOAUTH)); - // var key = scratchBufferManager.CreateArgSlice(Convert.ToString(args[0])); - // var status = api.GET(key, out var value); - // if (status == GarnetStatus.OK) - // return value.ToString(); - // return null; - // } - // // As fallback, we use RespServerSession with a RESP-formatted input. This could be optimized - // // in future to provide parse state directly. - // default: - // { - //var request = scratchBufferManager.FormatCommandAsResp(cmd, args, null); - //_ = respServerSession.TryConsumeMessages(request.ptr, request.length); - //var response = scratchBufferNetworkSender.GetResponse(); - //var result = ProcessResponse(response.ptr, response.length); - //scratchBufferNetworkSender.Reset(); - //return result; - // } - //} - //} - - /// /// Custom command implementation /// diff --git a/libs/server/Custom/CustomRespCommands.cs b/libs/server/Custom/CustomRespCommands.cs index f038df300c4..cb46fe3f27d 100644 --- a/libs/server/Custom/CustomRespCommands.cs +++ b/libs/server/Custom/CustomRespCommands.cs @@ -17,7 +17,7 @@ internal sealed unsafe partial class RespServerSession : ServerSessionBase { private bool TryTransactionProc(byte id, CustomTransactionProcedure proc, int parseStateFirstArgIdx = 0, int parseStateLastArgIdx = -1) { - // Define _output + // Define output var output = new MemoryResult(null, 0); // Run procedure @@ -28,7 +28,7 @@ private bool TryTransactionProc(byte id, CustomTransactionProcedure proc, int pa var procInput = new CustomProcedureInput(ref parseState, parseStateFirstArgIdx, parseStateLastArgIdx); if (txnManager.RunTransactionProc(id, ref procInput, proc, ref output)) { - // Write _output to wire + // Write output to wire if (output.MemoryOwner != null) SendAndReset(output.MemoryOwner, output.Length); else @@ -37,7 +37,7 @@ private bool TryTransactionProc(byte id, CustomTransactionProcedure proc, int pa } else { - // Write _output to wire + // Write output to wire if (output.MemoryOwner != null) SendAndReset(output.MemoryOwner, output.Length); else diff --git a/libs/server/Storage/Session/CustomOps.cs b/libs/server/Storage/Session/CustomOps.cs deleted file mode 100644 index 0c99fee326c..00000000000 --- a/libs/server/Storage/Session/CustomOps.cs +++ /dev/null @@ -1,121 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -//using System; -//using System.Diagnostics; -//using Garnet.common; -//using Tsavorite.core; - -//namespace Garnet.server -//{ -// using ObjectStoreAllocator = GenericAllocator>>; -// using ObjectStoreFunctions = StoreFunctions>; - -// sealed partial class StorageSession : IDisposable -// { -// private bool TryCustomObjectCommand(byte* ptr, byte* end, RespCommand cmd, byte subid, CommandType type, ref TGarnetApi storageApi) -// where TGarnetApi : IGarnetAdvancedApi -// { -// var keyBytes = parseState.GetArgSliceByRef(0).SpanByte.ToByteArray(); - -// // Prepare input -// var input = new ObjectInput -// { -// header = new RespInputHeader -// { -// cmd = cmd, -// SubId = subid -// }, -// parseState = parseState, -// parseStateStartIdx = 1 -// }; - -// var output = new GarnetObjectStoreOutput { spanByteAndMemory = new SpanByteAndMemory(null) }; - -// GarnetStatus status; - -// if (type == CommandType.ReadModifyWrite) -// { -// status = storageApi.RMW_ObjectStore(ref keyBytes, ref input, ref output); -// Debug.Assert(!output.spanByteAndMemory.IsSpanByte); - -// switch (status) -// { -// case GarnetStatus.WRONGTYPE: -// while (!RespWriteUtils.WriteError(CmdStrings.RESP_ERR_WRONG_TYPE, ref dcurr, dend)) -// SendAndReset(); -// break; -// default: -// if (output.spanByteAndMemory.Memory != null) -// SendAndReset(output.spanByteAndMemory.Memory, output.spanByteAndMemory.Length); -// else -// while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_OK, ref dcurr, dend)) -// SendAndReset(); -// break; -// } -// } -// else -// { -// status = storageApi.Read_ObjectStore(ref keyBytes, ref input, ref output); -// Debug.Assert(!output.spanByteAndMemory.IsSpanByte); - -// switch (status) -// { -// case GarnetStatus.OK: -// if (output.spanByteAndMemory.Memory != null) -// SendAndReset(output.spanByteAndMemory.Memory, output.spanByteAndMemory.Length); -// else -// while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_OK, ref dcurr, dend)) -// SendAndReset(); -// break; -// case GarnetStatus.NOTFOUND: -// Debug.Assert(output.spanByteAndMemory.Memory == null); -// while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_ERRNOTFOUND, ref dcurr, dend)) -// SendAndReset(); -// break; -// case GarnetStatus.WRONGTYPE: -// while (!RespWriteUtils.WriteError(CmdStrings.RESP_ERR_WRONG_TYPE, ref dcurr, dend)) -// SendAndReset(); -// break; -// } -// } - -// return true; -// } - - -// public unsafe GarnetStatus CustomObjectCommand(ArgSlice key, ArgSlice[] elements, ListOperation lop, out int itemsDoneCount, ref TObjectContext objectStoreContext) -// where TObjectContext : ITsavoriteContext -// { -// itemsDoneCount = 0; - -// if (key.Length == 0 || elements.Length == 0) -// return GarnetStatus.OK; - -// // Prepare the parse state -// var parseState = new SessionParseState(); -// ArgSlice[] parseStateBuffer = default; -// parseState.InitializeWithArguments(ref parseStateBuffer, elements); - -// // Prepare the input -// var input = new ObjectInput -// { -// header = new RespInputHeader -// { -// type = GarnetObjectType.List, -// ListOp = lop, -// }, -// parseState = parseState, -// parseStateStartIdx = 0, -// }; - -// var arrKey = key.ToArray(); -// var status = RMWObjectStoreOperation(arrKey, ref input, out var output, ref objectStoreContext); - -// itemsDoneCount = output.result1; -// itemBroker.HandleCollectionUpdate(arrKey); -// return status; -// } - -// } -//} diff --git a/libs/server/Storage/Session/MainStore/MainStoreOps.cs b/libs/server/Storage/Session/MainStore/MainStoreOps.cs index f867807e4dd..fceaf24d88e 100644 --- a/libs/server/Storage/Session/MainStore/MainStoreOps.cs +++ b/libs/server/Storage/Session/MainStore/MainStoreOps.cs @@ -1124,65 +1124,6 @@ public unsafe GarnetStatus SCAN(long cursor, ArgSlice match, long coun return GarnetStatus.OK; } - public unsafe GarnetStatus CustomCommand(RespCommand cmd, ArgSlice key, ArgSlice[] args, CommandType type, out ArgSlice value, ref TContext context) - where TContext : ITsavoriteContext - { - var sbKey = key.SpanByte; - value = default; - - var sessionParseState = new SessionParseState(); - sessionParseState.InitializeWithArguments(args); - - var rawStringInput = new RawStringInput(cmd, ref sessionParseState); - - var _output = new SpanByteAndMemory { SpanByte = scratchBufferManager.ViewRemainingArgSlice().SpanByte }; - return GarnetStatus.OK; - //GarnetStatus status; - //if (type == CommandType.ReadModifyWrite) - //{ - // status = RMW_MainStore(ref sbKey, ref rawStringInput, ref _output, ref context); - // //Debug.Assert(!output.IsSpanByte); - - // if (output.Memory != null) - // SendAndReset(output.Memory, output.Length); - // else - // while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_OK, ref dcurr, dend)) - // SendAndReset(); - //} - //else - //{ - // status = Read_MainStore(ref sbKey, ref rawStringInput, ref _output, ref context); - // //Debug.Assert(!_output.IsSpanByte); // why? - - // if (status == GarnetStatus.OK) - // if (!_output.IsSpanByte) - // { - // value = scratchBufferManager.FormatScratch(0, _output.AsReadOnlySpan()); - // _output.Memory.Dispose(); - // } - // else - // { - // value = scratchBufferManager.CreateArgSlice(_output.Length); - // } - //} - - - //if (ret == GarnetStatus.OK) - //{ - // if (!_output.IsSpanByte) - // { - // value = scratchBufferManager.FormatScratch(0, _output.AsReadOnlySpan()); - // _output.Memory.Dispose(); - // } - // else - // { - // value = scratchBufferManager.CreateArgSlice(_output.Length); - // } - //} - - //return RMW_MainStore(ref sbKey, ref rawStringInput, ref output, ref context); - } - public GarnetStatus GetKeyType(ArgSlice key, out string keyType, ref TContext context, ref TObjectContext objectContext) where TContext : ITsavoriteContext where TObjectContext : ITsavoriteContext diff --git a/main/GarnetServer/Extensions/ProcCustomCmd.cs b/main/GarnetServer/Extensions/ProcCustomCmd.cs index ca1530862b3..fa21144ada1 100644 --- a/main/GarnetServer/Extensions/ProcCustomCmd.cs +++ b/main/GarnetServer/Extensions/ProcCustomCmd.cs @@ -20,9 +20,6 @@ public override unsafe bool Execute(TGarnetApi garnetApi, ref Custom { args[i] = GetNextArg(ref procInput, ref offset); } - // id from registration of custom raw string cmd - //garnetApi.CustomCommand(0, key, new ArgSlice(input.ptr + offset, input.length - offset), ref cmdOutput); - //InvokeCustomRawStringCommand(garnetApi, "SETIFPM", key, args); ExecuteCustomRawStringCommand(garnetApi, "SETIFPM", key, args, out var _output); return true; diff --git a/main/GarnetServer/Extensions/TxnCustomCmd.cs b/main/GarnetServer/Extensions/TxnCustomCmd.cs index 31c004c42d0..af56e943217 100644 --- a/main/GarnetServer/Extensions/TxnCustomCmd.cs +++ b/main/GarnetServer/Extensions/TxnCustomCmd.cs @@ -24,9 +24,6 @@ public override void Finalize(TGarnetApi api, ref CustomProcedureInp { args[i] = GetNextArg(ref procInput, ref offset); } - // id from registration of custom raw string cmd - //garnetApi.CustomCommand(0, key, new ArgSlice(input.ptr + offset, input.length - offset), ref cmdOutput); - //InvokeCustomRawStringCommand(garnetApi, "SETIFPM", key, args); //ExecuteCustomRawStringCommand(api, "SETIFPM", key, args, out var _output); From 31d84f043b453061b7ffc8093c4af15adbff4ed3 Mon Sep 17 00:00:00 2001 From: Yoganand Rajasekaran Date: Mon, 21 Oct 2024 22:25:33 -0700 Subject: [PATCH 07/12] Updated tests. --- test/Garnet.test/RespModuleTests.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/Garnet.test/RespModuleTests.cs b/test/Garnet.test/RespModuleTests.cs index c8bf7f46d14..c4cca16109d 100644 --- a/test/Garnet.test/RespModuleTests.cs +++ b/test/Garnet.test/RespModuleTests.cs @@ -116,7 +116,7 @@ public void TestModuleLoad() new RespCommandsInfo { Name = ""TestModule.MYDICTGET"", Arity = 3, FirstKey = 1, LastKey = 1, Step = 1, Flags = RespCommandFlags.ReadOnly, AclCategories = RespAclCategories.Read }); - context.RegisterProcedure(""TestModule.SUM"", new Sum());"; + context.RegisterProcedure(""TestModule.SUM"", () => new Sum());"; var modulePath = CreateTestModule(onLoad); @@ -186,7 +186,7 @@ public void TestModuleLoadUsingGarnetOptions() var onLoad2 = @"context.Initialize(""TestModule2"", 1); - context.RegisterProcedure(""TestModule2.SUM"", new Sum());"; + context.RegisterProcedure(""TestModule2.SUM"", () => new Sum());"; var module1Path = CreateTestModule(onLoad, "TestModule1.dll"); var module2Path = CreateTestModule(onLoad2, "TestModule2.dll"); From e6e3c0eaeca6d1943874f7de4d4c53c358b4be51 Mon Sep 17 00:00:00 2001 From: Yoganand Rajasekaran Date: Mon, 21 Oct 2024 22:29:35 -0700 Subject: [PATCH 08/12] Format fix. --- main/GarnetServer/Extensions/TxnCustomCmd.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/main/GarnetServer/Extensions/TxnCustomCmd.cs b/main/GarnetServer/Extensions/TxnCustomCmd.cs index af56e943217..a0d3a4419df 100644 --- a/main/GarnetServer/Extensions/TxnCustomCmd.cs +++ b/main/GarnetServer/Extensions/TxnCustomCmd.cs @@ -32,4 +32,4 @@ public override void Finalize(TGarnetApi api, ref CustomProcedureInp WriteSimpleString(ref output, "OK"); } } -} +} \ No newline at end of file From 7972d75761d9cfc79aadbeec40ad4f2cb088b8b5 Mon Sep 17 00:00:00 2001 From: Yoganand Rajasekaran Date: Tue, 22 Oct 2024 09:45:36 -0700 Subject: [PATCH 09/12] Added tests. --- libs/server/Custom/CustomRespCommands.cs | 14 +++-- main/GarnetServer/Extensions/ProcCustomCmd.cs | 12 ++-- main/GarnetServer/Extensions/TxnCustomCmd.cs | 38 ++++++++---- test/Garnet.test/Garnet.test.csproj | 4 +- test/Garnet.test/RespCustomCommandTests.cs | 58 +++++++++++++++++++ 5 files changed, 101 insertions(+), 25 deletions(-) diff --git a/libs/server/Custom/CustomRespCommands.cs b/libs/server/Custom/CustomRespCommands.cs index cb46fe3f27d..c04223f930e 100644 --- a/libs/server/Custom/CustomRespCommands.cs +++ b/libs/server/Custom/CustomRespCommands.cs @@ -85,8 +85,11 @@ private void TryCustomProcedure(CustomProcedure proc, int parseStateFirstArgIdx public bool InvokeCustomRawStringCommand(ref TGarnetApi storageApi, string cmd, ArgSlice key, ArgSlice[] args, out ArgSlice output) where TGarnetApi : IGarnetAdvancedApi { - // TODO: check if cmd exists - storeWrapper.customCommandManager.Match(new ReadOnlySpan(Encoding.UTF8.GetBytes(cmd)), out CustomRawStringCommand customCommand); + if (!storeWrapper.customCommandManager.Match(new ReadOnlySpan(Encoding.UTF8.GetBytes(cmd)), out CustomRawStringCommand customCommand)) + { + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERR_GENERIC_UNK_CMD); + return false; + } var sbKey = key.SpanByte; @@ -191,8 +194,11 @@ public bool InvokeCustomObjectCommand(ref TGarnetApi storageApi, str where TGarnetApi : IGarnetAdvancedApi { output = default; - // TODO: check if cmd exists - storeWrapper.customCommandManager.Match(new ReadOnlySpan(Encoding.UTF8.GetBytes(cmd)), out CustomObjectCommand customObjCommand); + if (!storeWrapper.customCommandManager.Match(new ReadOnlySpan(Encoding.UTF8.GetBytes(cmd)), out CustomObjectCommand customObjCommand)) + { + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERR_GENERIC_UNK_CMD); + return false; + } var keyBytes = key.ToArray(); diff --git a/main/GarnetServer/Extensions/ProcCustomCmd.cs b/main/GarnetServer/Extensions/ProcCustomCmd.cs index fa21144ada1..a063b35b05a 100644 --- a/main/GarnetServer/Extensions/ProcCustomCmd.cs +++ b/main/GarnetServer/Extensions/ProcCustomCmd.cs @@ -11,15 +11,11 @@ class ProcCustomCmd : CustomProcedure public override unsafe bool Execute(TGarnetApi garnetApi, ref CustomProcedureInput procInput, ref MemoryResult output) { var offset = 0; - ArgSlice key = GetNextArg(ref procInput, ref offset); + var key = GetNextArg(ref procInput, ref offset); - //var cmdOutput = new SpanByteAndMemory(null); - - ArgSlice[] args = new ArgSlice[procInput.parseState.Count - 1]; - for (int i = 0; i < procInput.parseState.Count - 1; i++) - { - args[i] = GetNextArg(ref procInput, ref offset); - } + var args = new ArgSlice[2]; + args[0] = GetNextArg(ref procInput, ref offset); // value to set + args[1] = GetNextArg(ref procInput, ref offset); // prefix to match ExecuteCustomRawStringCommand(garnetApi, "SETIFPM", key, args, out var _output); return true; diff --git a/main/GarnetServer/Extensions/TxnCustomCmd.cs b/main/GarnetServer/Extensions/TxnCustomCmd.cs index a0d3a4419df..563911162b0 100644 --- a/main/GarnetServer/Extensions/TxnCustomCmd.cs +++ b/main/GarnetServer/Extensions/TxnCustomCmd.cs @@ -9,25 +9,39 @@ namespace Garnet { class TxnCustomCmd : CustomTransactionProcedure { - public override void Main(TGarnetApi api, ref CustomProcedureInput procInput, ref MemoryResult output) => throw new System.NotImplementedException(); - public override bool Prepare(TGarnetReadApi api, ref CustomProcedureInput procInput) => false; + public override bool Prepare(TGarnetReadApi api, ref CustomProcedureInput procInput) + { + var offset = 0; + + var mainStoreKey = GetNextArg(ref procInput, ref offset); + _ = GetNextArg(ref procInput, ref offset); // mainStoreValue + + AddKey(mainStoreKey, LockType.Exclusive, false); - public override void Finalize(TGarnetApi api, ref CustomProcedureInput procInput, ref MemoryResult output) + var myDictKey = GetNextArg(ref procInput, ref offset); + AddKey(myDictKey, LockType.Exclusive, true); + + return true; + } + + public override void Main(TGarnetApi api, ref CustomProcedureInput procInput, ref MemoryResult output) { var offset = 0; - ArgSlice key = GetNextArg(ref procInput, ref offset); - var cmdOutput = new SpanByteAndMemory(null); + var mainStoreKey = GetNextArg(ref procInput, ref offset); + var mainStoreValue = GetNextArg(ref procInput, ref offset); + + api.SET(mainStoreKey, mainStoreValue); - ArgSlice[] args = new ArgSlice[procInput.parseState.Count - 1]; - for (int i = 0; i < procInput.parseState.Count - 1; i++) - { - args[i] = GetNextArg(ref procInput, ref offset); - } + var myDictKey = GetNextArg(ref procInput, ref offset); + var myDictField = GetNextArg(ref procInput, ref offset); + var myDictValue = GetNextArg(ref procInput, ref offset); - //ExecuteCustomRawStringCommand(api, "SETIFPM", key, args, out var _output); + var args = new ArgSlice[2]; + args[0] = myDictField; + args[1] = myDictValue; - ExecuteCustomObjectCommand(api, "MYDICTSET", key, args, out var _output); + ExecuteCustomObjectCommand(api, "MYDICTSET", myDictKey, args, out var _output); WriteSimpleString(ref output, "OK"); } diff --git a/test/Garnet.test/Garnet.test.csproj b/test/Garnet.test/Garnet.test.csproj index f60804d6c92..6b180dfb7bf 100644 --- a/test/Garnet.test/Garnet.test.csproj +++ b/test/Garnet.test/Garnet.test.csproj @@ -26,7 +26,9 @@ - + + + diff --git a/test/Garnet.test/RespCustomCommandTests.cs b/test/Garnet.test/RespCustomCommandTests.cs index 01b9cc17343..692d0874868 100644 --- a/test/Garnet.test/RespCustomCommandTests.cs +++ b/test/Garnet.test/RespCustomCommandTests.cs @@ -962,5 +962,63 @@ .. libraryPaths.Skip(1), } ClassicAssert.IsNull(resp); } + + [Test] + public void CustomProcInvokingCustomCmdTest() + { + using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig()); + var db = redis.GetDatabase(0); + + server.Register.NewCommand("SETIFPM", CommandType.ReadModifyWrite, new SetIfPMCustomCommand(), new RespCommandsInfo { Arity = 4 }); + server.Register.NewProcedure("PROCCMD", () => new ProcCustomCmd()); + + var key = "mainKey"; + var value = "foovalue0"; + db.StringSet(key, value); + + var newValue1 = "foovalue1"; + var newValue2 = "foovalue2"; + + // This conditional set should pass (prefix matches) + var result = db.Execute("SETIFPM", key, newValue1, "foo"); + ClassicAssert.AreEqual("OK", (string)result); + + var retValue = db.StringGet(key); + ClassicAssert.AreEqual(newValue1, retValue.ToString()); + + // This conditional set should fail (prefix does not match) + result = db.Execute("SETIFPM", key, newValue2, "bar"); + ClassicAssert.AreEqual("OK", (string)result); + + retValue = db.StringGet(key); + ClassicAssert.AreEqual(newValue1, retValue.ToString()); + } + + [Test] + public void CustomTransactionInvokingCustomCmdTest() + { + var factory = new MyDictFactory(); + server.Register.NewCommand("MYDICTSET", CommandType.ReadModifyWrite, factory, new MyDictSet(), new RespCommandsInfo { Arity = 4 }); + server.Register.NewCommand("MYDICTGET", CommandType.Read, factory, new MyDictGet(), new RespCommandsInfo { Arity = 3 }); + server.Register.NewTransactionProc("TXNCMD", () => new TxnCustomCmd()); + + using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig()); + var db = redis.GetDatabase(0); + + var mainKey = "mainKey"; + var mainValue = "foovalue0"; + var dictKey = "dictKey"; + var dictField = "dictField"; + var dictValue = "dictValue"; + + var result = db.Execute("TXNCMD", mainKey, mainValue, dictKey, dictField, dictValue); + ClassicAssert.AreEqual("OK", (string)result); + + result = db.Execute("MYDICTGET", dictKey, dictField); + ClassicAssert.AreEqual(dictValue, (string)result); + + var retValue = db.StringGet(mainKey); + ClassicAssert.AreEqual(mainValue, (string)retValue); + } } } \ No newline at end of file From bc6bd523a2060b6e5f736c69896be1d1f944e1da Mon Sep 17 00:00:00 2001 From: Yoganand Rajasekaran Date: Tue, 22 Oct 2024 12:30:30 -0700 Subject: [PATCH 10/12] Unified API. Added test for invalid command. --- libs/server/Custom/CustomFunctions.cs | 19 +- libs/server/Custom/CustomRespCommands.cs | 245 +++++++++--------- main/GarnetServer/Extensions/ProcCustomCmd.cs | 2 +- main/GarnetServer/Extensions/TxnCustomCmd.cs | 2 +- test/Garnet.test/RespCustomCommandTests.cs | 36 ++- 5 files changed, 173 insertions(+), 131 deletions(-) diff --git a/libs/server/Custom/CustomFunctions.cs b/libs/server/Custom/CustomFunctions.cs index 5504185146a..7f968a79d58 100644 --- a/libs/server/Custom/CustomFunctions.cs +++ b/libs/server/Custom/CustomFunctions.cs @@ -211,16 +211,19 @@ protected static unsafe ArgSlice GetNextArg(ref CustomProcedureInput procInput, return GetNextArg(ref procInput.parseState, procInput.parseStateFirstArgIdx, ref offset); } - protected void ExecuteCustomRawStringCommand(TGarnetApi garnetApi, string cmd, ArgSlice key, ArgSlice[] input, out ArgSlice output) - where TGarnetApi : IGarnetApi - { - respServerSession.InvokeCustomRawStringCommand(ref garnetApi, cmd, key, input, out output); - } - - protected void ExecuteCustomObjectCommand(TGarnetApi garnetApi, string cmd, ArgSlice key, ArgSlice[] input, out ArgSlice output) + /// + /// Execute custom raw string or object command + /// + /// + /// Command to execute + /// Key argument + /// Array of arguments to command + /// Output from command + /// True if command found, else false + protected bool ExecuteCustomCommand(TGarnetApi garnetApi, string cmd, ArgSlice key, ArgSlice[] input, out ArgSlice output) where TGarnetApi : IGarnetApi { - respServerSession.InvokeCustomObjectCommand(ref garnetApi, cmd, key, input, out output); + return respServerSession.InvokeCustomCommand(ref garnetApi, cmd, key, input, out output); } } } \ No newline at end of file diff --git a/libs/server/Custom/CustomRespCommands.cs b/libs/server/Custom/CustomRespCommands.cs index c04223f930e..010ff3b1f70 100644 --- a/libs/server/Custom/CustomRespCommands.cs +++ b/libs/server/Custom/CustomRespCommands.cs @@ -82,66 +82,6 @@ private void TryCustomProcedure(CustomProcedure proc, int parseStateFirstArgIdx } } - public bool InvokeCustomRawStringCommand(ref TGarnetApi storageApi, string cmd, ArgSlice key, ArgSlice[] args, out ArgSlice output) - where TGarnetApi : IGarnetAdvancedApi - { - if (!storeWrapper.customCommandManager.Match(new ReadOnlySpan(Encoding.UTF8.GetBytes(cmd)), out CustomRawStringCommand customCommand)) - { - output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERR_GENERIC_UNK_CMD); - return false; - } - - var sbKey = key.SpanByte; - - var inputArg = customCommand.expirationTicks > 0 ? DateTimeOffset.UtcNow.Ticks + customCommand.expirationTicks : customCommand.expirationTicks; - var sessionParseState = new SessionParseState(); - sessionParseState.InitializeWithArguments(args); - var rawStringInput = new RawStringInput(customCommand.GetRespCommand(), ref sessionParseState, arg1: inputArg); - - var _output = new SpanByteAndMemory(null); - GarnetStatus status; - if (customCommand.type == CommandType.ReadModifyWrite) - { - status = storageApi.RMW_MainStore(ref sbKey, ref rawStringInput, ref _output); - Debug.Assert(!_output.IsSpanByte); - - if (_output.Memory != null) - { - output = scratchBufferManager.FormatScratch(0, _output.AsReadOnlySpan()); - _output.Memory.Dispose(); - } - else - { - output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_OK); - } - } - else - { - status = storageApi.Read_MainStore(ref sbKey, ref rawStringInput, ref _output); - Debug.Assert(!_output.IsSpanByte); - - if (status == GarnetStatus.OK) - { - if (_output.Memory != null) - { - output = scratchBufferManager.FormatScratch(0, _output.AsReadOnlySpan()); - _output.Memory.Dispose(); - } - else - { - output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_OK); - } - } - else - { - Debug.Assert(_output.Memory == null); - output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERRNOTFOUND); - } - } - - return true; - } - /// /// Custom command /// @@ -190,63 +130,65 @@ private bool TryCustomRawStringCommand(RespCommand cmd, long expirat return true; } - public bool InvokeCustomObjectCommand(ref TGarnetApi storageApi, string cmd, ArgSlice key, ArgSlice[] args, out ArgSlice output) + /// + /// Custom object command + /// + private bool TryCustomObjectCommand(GarnetObjectType objType, byte subid, CommandType type, ref TGarnetApi storageApi) where TGarnetApi : IGarnetAdvancedApi { - output = default; - if (!storeWrapper.customCommandManager.Match(new ReadOnlySpan(Encoding.UTF8.GetBytes(cmd)), out CustomObjectCommand customObjCommand)) - { - output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERR_GENERIC_UNK_CMD); - return false; - } - - var keyBytes = key.ToArray(); + var keyBytes = parseState.GetArgSliceByRef(0).SpanByte.ToByteArray(); // Prepare input - var header = new RespInputHeader(customObjCommand.GetRespCommand()) { SubId = customObjCommand.subid }; - var sessionParseState = new SessionParseState(); - sessionParseState.InitializeWithArguments(args); - var input = new ObjectInput(header, ref sessionParseState); - var _output = new GarnetObjectStoreOutput { spanByteAndMemory = new SpanByteAndMemory(null) }; + var header = new RespInputHeader(objType) { SubId = subid }; + var input = new ObjectInput(header, ref parseState, 1); + + var output = new GarnetObjectStoreOutput { spanByteAndMemory = new SpanByteAndMemory(null) }; + GarnetStatus status; - if (customObjCommand.type == CommandType.ReadModifyWrite) + + if (type == CommandType.ReadModifyWrite) { - status = storageApi.RMW_ObjectStore(ref keyBytes, ref input, ref _output); - Debug.Assert(!_output.spanByteAndMemory.IsSpanByte); + status = storageApi.RMW_ObjectStore(ref keyBytes, ref input, ref output); + Debug.Assert(!output.spanByteAndMemory.IsSpanByte); switch (status) { case GarnetStatus.WRONGTYPE: - output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERR_WRONG_TYPE); + while (!RespWriteUtils.WriteError(CmdStrings.RESP_ERR_WRONG_TYPE, ref dcurr, dend)) + SendAndReset(); break; default: - if (_output.spanByteAndMemory.Memory != null) - output = scratchBufferManager.FormatScratch(0, _output.spanByteAndMemory.AsReadOnlySpan()); + if (output.spanByteAndMemory.Memory != null) + SendAndReset(output.spanByteAndMemory.Memory, output.spanByteAndMemory.Length); else - output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_OK); + while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_OK, ref dcurr, dend)) + SendAndReset(); break; } } else { - status = storageApi.Read_ObjectStore(ref keyBytes, ref input, ref _output); - Debug.Assert(!_output.spanByteAndMemory.IsSpanByte); + status = storageApi.Read_ObjectStore(ref keyBytes, ref input, ref output); + Debug.Assert(!output.spanByteAndMemory.IsSpanByte); switch (status) { case GarnetStatus.OK: - if (_output.spanByteAndMemory.Memory != null) - output = scratchBufferManager.FormatScratch(0, _output.spanByteAndMemory.AsReadOnlySpan()); + if (output.spanByteAndMemory.Memory != null) + SendAndReset(output.spanByteAndMemory.Memory, output.spanByteAndMemory.Length); else - output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_OK); + while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_OK, ref dcurr, dend)) + SendAndReset(); break; case GarnetStatus.NOTFOUND: - Debug.Assert(_output.spanByteAndMemory.Memory == null); - output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERRNOTFOUND); + Debug.Assert(output.spanByteAndMemory.Memory == null); + while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_ERRNOTFOUND, ref dcurr, dend)) + SendAndReset(); break; case GarnetStatus.WRONGTYPE: - output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERR_WRONG_TYPE); + while (!RespWriteUtils.WriteError(CmdStrings.RESP_ERR_WRONG_TYPE, ref dcurr, dend)) + SendAndReset(); break; } } @@ -254,65 +196,130 @@ public bool InvokeCustomObjectCommand(ref TGarnetApi storageApi, str return true; } - /// - /// Custom object command - /// - private bool TryCustomObjectCommand(GarnetObjectType objType, byte subid, CommandType type, ref TGarnetApi storageApi) + public bool InvokeCustomCommand(ref TGarnetApi storageApi, string cmd, ArgSlice key, ArgSlice[] args, out ArgSlice output) where TGarnetApi : IGarnetAdvancedApi { - var keyBytes = parseState.GetArgSliceByRef(0).SpanByte.ToByteArray(); - - // Prepare input + if (storeWrapper.customCommandManager.Match(new ReadOnlySpan(Encoding.UTF8.GetBytes(cmd)), out CustomRawStringCommand customCommand)) + { + return InvokeCustomRawStringCommand(ref storageApi, customCommand, key, args, out output); + } + else if (storeWrapper.customCommandManager.Match(new ReadOnlySpan(Encoding.UTF8.GetBytes(cmd)), out CustomObjectCommand customObjCommand)) + { + return InvokeCustomObjectCommand(ref storageApi, customObjCommand, key, args, out output); + } + else + { + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERR_GENERIC_UNK_CMD); + return false; + } + } - var header = new RespInputHeader(objType) { SubId = subid }; - var input = new ObjectInput(header, ref parseState, 1); + private bool InvokeCustomRawStringCommand(ref TGarnetApi storageApi, CustomRawStringCommand customCommand, ArgSlice key, ArgSlice[] args, out ArgSlice output) + where TGarnetApi : IGarnetAdvancedApi + { + var sbKey = key.SpanByte; - var output = new GarnetObjectStoreOutput { spanByteAndMemory = new SpanByteAndMemory(null) }; + var inputArg = customCommand.expirationTicks > 0 ? DateTimeOffset.UtcNow.Ticks + customCommand.expirationTicks : customCommand.expirationTicks; + var sessionParseState = new SessionParseState(); + sessionParseState.InitializeWithArguments(args); + var rawStringInput = new RawStringInput(customCommand.GetRespCommand(), ref sessionParseState, arg1: inputArg); + var _output = new SpanByteAndMemory(null); GarnetStatus status; + if (customCommand.type == CommandType.ReadModifyWrite) + { + status = storageApi.RMW_MainStore(ref sbKey, ref rawStringInput, ref _output); + Debug.Assert(!_output.IsSpanByte); - if (type == CommandType.ReadModifyWrite) + if (_output.Memory != null) + { + output = scratchBufferManager.FormatScratch(0, _output.AsReadOnlySpan()); + _output.Memory.Dispose(); + } + else + { + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_OK); + } + } + else { - status = storageApi.RMW_ObjectStore(ref keyBytes, ref input, ref output); - Debug.Assert(!output.spanByteAndMemory.IsSpanByte); + status = storageApi.Read_MainStore(ref sbKey, ref rawStringInput, ref _output); + Debug.Assert(!_output.IsSpanByte); + + if (status == GarnetStatus.OK) + { + if (_output.Memory != null) + { + output = scratchBufferManager.FormatScratch(0, _output.AsReadOnlySpan()); + _output.Memory.Dispose(); + } + else + { + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_OK); + } + } + else + { + Debug.Assert(_output.Memory == null); + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERRNOTFOUND); + } + } + + return true; + } + + private bool InvokeCustomObjectCommand(ref TGarnetApi storageApi, CustomObjectCommand customObjCommand, ArgSlice key, ArgSlice[] args, out ArgSlice output) + where TGarnetApi : IGarnetAdvancedApi + { + output = default; + + var keyBytes = key.ToArray(); + + // Prepare input + var header = new RespInputHeader(customObjCommand.GetRespCommand()) { SubId = customObjCommand.subid }; + var sessionParseState = new SessionParseState(); + sessionParseState.InitializeWithArguments(args); + var input = new ObjectInput(header, ref sessionParseState); + + var _output = new GarnetObjectStoreOutput { spanByteAndMemory = new SpanByteAndMemory(null) }; + GarnetStatus status; + if (customObjCommand.type == CommandType.ReadModifyWrite) + { + status = storageApi.RMW_ObjectStore(ref keyBytes, ref input, ref _output); + Debug.Assert(!_output.spanByteAndMemory.IsSpanByte); switch (status) { case GarnetStatus.WRONGTYPE: - while (!RespWriteUtils.WriteError(CmdStrings.RESP_ERR_WRONG_TYPE, ref dcurr, dend)) - SendAndReset(); + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERR_WRONG_TYPE); break; default: - if (output.spanByteAndMemory.Memory != null) - SendAndReset(output.spanByteAndMemory.Memory, output.spanByteAndMemory.Length); + if (_output.spanByteAndMemory.Memory != null) + output = scratchBufferManager.FormatScratch(0, _output.spanByteAndMemory.AsReadOnlySpan()); else - while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_OK, ref dcurr, dend)) - SendAndReset(); + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_OK); break; } } else { - status = storageApi.Read_ObjectStore(ref keyBytes, ref input, ref output); - Debug.Assert(!output.spanByteAndMemory.IsSpanByte); + status = storageApi.Read_ObjectStore(ref keyBytes, ref input, ref _output); + Debug.Assert(!_output.spanByteAndMemory.IsSpanByte); switch (status) { case GarnetStatus.OK: - if (output.spanByteAndMemory.Memory != null) - SendAndReset(output.spanByteAndMemory.Memory, output.spanByteAndMemory.Length); + if (_output.spanByteAndMemory.Memory != null) + output = scratchBufferManager.FormatScratch(0, _output.spanByteAndMemory.AsReadOnlySpan()); else - while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_OK, ref dcurr, dend)) - SendAndReset(); + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_OK); break; case GarnetStatus.NOTFOUND: - Debug.Assert(output.spanByteAndMemory.Memory == null); - while (!RespWriteUtils.WriteDirect(CmdStrings.RESP_ERRNOTFOUND, ref dcurr, dend)) - SendAndReset(); + Debug.Assert(_output.spanByteAndMemory.Memory == null); + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERRNOTFOUND); break; case GarnetStatus.WRONGTYPE: - while (!RespWriteUtils.WriteError(CmdStrings.RESP_ERR_WRONG_TYPE, ref dcurr, dend)) - SendAndReset(); + output = scratchBufferManager.CreateArgSlice(CmdStrings.RESP_ERR_WRONG_TYPE); break; } } diff --git a/main/GarnetServer/Extensions/ProcCustomCmd.cs b/main/GarnetServer/Extensions/ProcCustomCmd.cs index a063b35b05a..1e87b1d17e1 100644 --- a/main/GarnetServer/Extensions/ProcCustomCmd.cs +++ b/main/GarnetServer/Extensions/ProcCustomCmd.cs @@ -17,7 +17,7 @@ public override unsafe bool Execute(TGarnetApi garnetApi, ref Custom args[0] = GetNextArg(ref procInput, ref offset); // value to set args[1] = GetNextArg(ref procInput, ref offset); // prefix to match - ExecuteCustomRawStringCommand(garnetApi, "SETIFPM", key, args, out var _output); + ExecuteCustomCommand(garnetApi, "SETIFPM", key, args, out var _output); return true; } } diff --git a/main/GarnetServer/Extensions/TxnCustomCmd.cs b/main/GarnetServer/Extensions/TxnCustomCmd.cs index 563911162b0..ba26c2270c8 100644 --- a/main/GarnetServer/Extensions/TxnCustomCmd.cs +++ b/main/GarnetServer/Extensions/TxnCustomCmd.cs @@ -41,7 +41,7 @@ public override void Main(TGarnetApi api, ref CustomProcedureInput p args[0] = myDictField; args[1] = myDictValue; - ExecuteCustomObjectCommand(api, "MYDICTSET", myDictKey, args, out var _output); + ExecuteCustomCommand(api, "MYDICTSET", myDictKey, args, out var _output); WriteSimpleString(ref output, "OK"); } diff --git a/test/Garnet.test/RespCustomCommandTests.cs b/test/Garnet.test/RespCustomCommandTests.cs index 692d0874868..5f3ceac6643 100644 --- a/test/Garnet.test/RespCustomCommandTests.cs +++ b/test/Garnet.test/RespCustomCommandTests.cs @@ -120,6 +120,26 @@ public override bool Execute(TGarnetApi garnetApi, ref CustomProcedu } } + public class InvalidCommandProc : CustomProcedure + { + public override bool Execute(TGarnetApi garnetApi, ref CustomProcedureInput procInput, ref MemoryResult output) + { + var offset = 0; + var key = GetNextArg(ref procInput, ref offset); + + if (ExecuteCustomCommand(garnetApi, "INVALIDCMD", key, null, out var _output)) + { + WriteError(ref output, "ERR ExecuteCustomCommand should have failed"); + } + else + { + WriteSimpleString(ref output, "OK"); + } + + return true; + } + } + [TestFixture] public class RespCustomCommandTests { @@ -980,14 +1000,14 @@ public void CustomProcInvokingCustomCmdTest() var newValue2 = "foovalue2"; // This conditional set should pass (prefix matches) - var result = db.Execute("SETIFPM", key, newValue1, "foo"); + var result = db.Execute("PROCCMD", key, newValue1, "foo"); ClassicAssert.AreEqual("OK", (string)result); var retValue = db.StringGet(key); ClassicAssert.AreEqual(newValue1, retValue.ToString()); // This conditional set should fail (prefix does not match) - result = db.Execute("SETIFPM", key, newValue2, "bar"); + result = db.Execute("PROCCMD", key, newValue2, "bar"); ClassicAssert.AreEqual("OK", (string)result); retValue = db.StringGet(key); @@ -1020,5 +1040,17 @@ public void CustomTransactionInvokingCustomCmdTest() var retValue = db.StringGet(mainKey); ClassicAssert.AreEqual(mainValue, (string)retValue); } + + [Test] + public void CustomProcedureInvokingInvalidCommandTest() + { + server.Register.NewProcedure("PROCINVALIDCMD", () => new InvalidCommandProc()); + + using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig()); + var db = redis.GetDatabase(0); + + var result = db.Execute("PROCINVALIDCMD", "key"); + ClassicAssert.AreEqual("OK", (string)result); + } } } \ No newline at end of file From 3890dc57969b4e2c1950298983221e6f2cb5adeb Mon Sep 17 00:00:00 2001 From: Yoganand Rajasekaran Date: Tue, 22 Oct 2024 12:34:48 -0700 Subject: [PATCH 11/12] Fixed formatting. --- test/Garnet.test/Garnet.test.csproj | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/Garnet.test/Garnet.test.csproj b/test/Garnet.test/Garnet.test.csproj index 6b180dfb7bf..e200ab39f95 100644 --- a/test/Garnet.test/Garnet.test.csproj +++ b/test/Garnet.test/Garnet.test.csproj @@ -26,9 +26,9 @@ - - - + + + From ce69dfe8d8d6523175ddfd2828d7eb29e0d67f2f Mon Sep 17 00:00:00 2001 From: Yoganand Rajasekaran Date: Tue, 22 Oct 2024 20:11:17 -0700 Subject: [PATCH 12/12] Merged latest changes. Moved custom procs, txn to test project. --- libs/server/Custom/CustomRespCommands.cs | 2 +- main/GarnetServer/Program.cs | 3 --- .../Garnet.test}/Extensions/ProcCustomCmd.cs | 0 .../Garnet.test}/Extensions/TxnCustomCmd.cs | 0 test/Garnet.test/Garnet.test.csproj | 2 -- 5 files changed, 1 insertion(+), 6 deletions(-) rename {main/GarnetServer => test/Garnet.test}/Extensions/ProcCustomCmd.cs (100%) rename {main/GarnetServer => test/Garnet.test}/Extensions/TxnCustomCmd.cs (100%) diff --git a/libs/server/Custom/CustomRespCommands.cs b/libs/server/Custom/CustomRespCommands.cs index 010ff3b1f70..71cb54a60f0 100644 --- a/libs/server/Custom/CustomRespCommands.cs +++ b/libs/server/Custom/CustomRespCommands.cs @@ -276,7 +276,7 @@ private bool InvokeCustomObjectCommand(ref TGarnetApi storageApi, Cu var keyBytes = key.ToArray(); // Prepare input - var header = new RespInputHeader(customObjCommand.GetRespCommand()) { SubId = customObjCommand.subid }; + var header = new RespInputHeader(customObjCommand.GetObjectType()) { SubId = customObjCommand.subid }; var sessionParseState = new SessionParseState(); sessionParseState.InitializeWithArguments(args); var input = new ObjectInput(header, ref sessionParseState); diff --git a/main/GarnetServer/Program.cs b/main/GarnetServer/Program.cs index 882902b47f5..314a6a39baf 100644 --- a/main/GarnetServer/Program.cs +++ b/main/GarnetServer/Program.cs @@ -94,9 +94,6 @@ static void RegisterExtensions(GarnetServer server) server.Register.NewProcedure("SUM", () => new Sum()); server.Register.NewProcedure("SETMAINANDOBJECT", () => new SetStringAndList()); - - server.Register.NewProcedure("PROCCMD", () => new ProcCustomCmd()); - server.Register.NewTransactionProc("TXNCMD", () => new TxnCustomCmd()); } } } \ No newline at end of file diff --git a/main/GarnetServer/Extensions/ProcCustomCmd.cs b/test/Garnet.test/Extensions/ProcCustomCmd.cs similarity index 100% rename from main/GarnetServer/Extensions/ProcCustomCmd.cs rename to test/Garnet.test/Extensions/ProcCustomCmd.cs diff --git a/main/GarnetServer/Extensions/TxnCustomCmd.cs b/test/Garnet.test/Extensions/TxnCustomCmd.cs similarity index 100% rename from main/GarnetServer/Extensions/TxnCustomCmd.cs rename to test/Garnet.test/Extensions/TxnCustomCmd.cs diff --git a/test/Garnet.test/Garnet.test.csproj b/test/Garnet.test/Garnet.test.csproj index e200ab39f95..f60804d6c92 100644 --- a/test/Garnet.test/Garnet.test.csproj +++ b/test/Garnet.test/Garnet.test.csproj @@ -26,8 +26,6 @@ - -