Fixing module registration logic + Implementing no-op & JSON module benchmarking

Garnet.sln

@@ -109,6 +109,8 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Garnet.worker", "hosting\Wi
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Garnet.resources", "libs\resources\Garnet.resources.csproj", "{A48412B4-FD60-467E-A5D9-F155CAB4F907}"
 EndProject
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "NoOpModule", "playground\NoOpModule\NoOpModule.csproj", "{D4C9A1A0-7053-F072-21F5-4E0C5827136D}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -325,6 +327,14 @@ Global
 		{A48412B4-FD60-467E-A5D9-F155CAB4F907}.Release|Any CPU.Build.0 = Release|Any CPU
 		{A48412B4-FD60-467E-A5D9-F155CAB4F907}.Release|x64.ActiveCfg = Release|Any CPU
 		{A48412B4-FD60-467E-A5D9-F155CAB4F907}.Release|x64.Build.0 = Release|Any CPU
+		{D4C9A1A0-7053-F072-21F5-4E0C5827136D}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{D4C9A1A0-7053-F072-21F5-4E0C5827136D}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{D4C9A1A0-7053-F072-21F5-4E0C5827136D}.Debug|x64.ActiveCfg = Debug|Any CPU
+		{D4C9A1A0-7053-F072-21F5-4E0C5827136D}.Debug|x64.Build.0 = Debug|Any CPU
+		{D4C9A1A0-7053-F072-21F5-4E0C5827136D}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{D4C9A1A0-7053-F072-21F5-4E0C5827136D}.Release|Any CPU.Build.0 = Release|Any CPU
+		{D4C9A1A0-7053-F072-21F5-4E0C5827136D}.Release|x64.ActiveCfg = Release|Any CPU
+		{D4C9A1A0-7053-F072-21F5-4E0C5827136D}.Release|x64.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE
@@ -359,6 +369,7 @@ Global
 		{697766CD-2046-46D9-958A-0FD3B46C98D4} = {01823EA4-4446-4D66-B268-DFEE55951964}
 		{DF2DD03E-87EE-482A-9FBA-6C8FBC23BDC5} = {697766CD-2046-46D9-958A-0FD3B46C98D4}
 		{A48412B4-FD60-467E-A5D9-F155CAB4F907} = {147FCE31-EC09-4C90-8E4D-37CA87ED18C3}
+		{D4C9A1A0-7053-F072-21F5-4E0C5827136D} = {69A71E2C-00E3-42F3-854E-BE157A24834E}
 	EndGlobalSection
 	GlobalSection(ExtensibilityGlobals) = postSolution
 		SolutionGuid = {2C02C405-4798-41CA-AF98-61EDFEF6772E}

benchmark/BDN.benchmark/BDN.benchmark.csproj

@@ -18,6 +18,8 @@
     <ProjectReference Include="..\..\libs\host\Garnet.host.csproj" />
     <ProjectReference Include="..\..\libs\common\Garnet.common.csproj" />
     <ProjectReference Include="..\..\libs\server\Garnet.server.csproj" />
+    <ProjectReference Include="..\..\playground\GarnetJSON\GarnetJSON.csproj" />
+    <ProjectReference Include="..\..\playground\NoOpModule\NoOpModule.csproj" />
   </ItemGroup>
 
   <ItemGroup>

benchmark/BDN.benchmark/Operations/ModuleOperations.cs

@@ -0,0 +1,115 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+using System.Reflection;
+using System.Text;
+using BenchmarkDotNet.Attributes;
+using Embedded.server;
+
+namespace BDN.benchmark.Operations
+{
+    /// <summary>
+    /// Benchmark for ModuleOperations
+    /// </summary>
+    [MemoryDiagnoser]
+    public class ModuleOperations : OperationsBase
+    {
+        static ReadOnlySpan<byte> NOOPCMDREAD => "*2\r\n$22\r\nNoOpModule.NOOPCMDREAD\r\n$2\r\nk1\r\n"u8;
+        Request noOpCmdRead;
+
+        static ReadOnlySpan<byte> NOOPCMDRMW => "*2\r\n$21\r\nNoOpModule.NOOPCMDRMW\r\n$2\r\nk1\r\n"u8;
+        Request noOpCmdRmw;
+
+        static ReadOnlySpan<byte> NOOPOBJRMW => "*2\r\n$21\r\nNoOpModule.NOOPOBJRMW\r\n$2\r\nk2\r\n"u8;
+        Request noOpObjRmw;
+
+        static ReadOnlySpan<byte> NOOPOBJREAD => "*2\r\n$22\r\nNoOpModule.NOOPOBJREAD\r\n$2\r\nk2\r\n"u8;
+        Request noOpObjRead;
+
+        static ReadOnlySpan<byte> NOOPPROC => "*1\r\n$19\r\nNoOpModule.NOOPPROC\r\n"u8;
+        Request noOpProc;
+
+        static ReadOnlySpan<byte> NOOPTXN => "*1\r\n$18\r\nNoOpModule.NOOPTXN\r\n"u8;
+        Request noOpTxn;
+
+        static ReadOnlySpan<byte> JSONGETCMD => "*3\r\n$8\r\nJSON.GET\r\n$2\r\nk3\r\n$1\r\n$\r\n"u8;
+        Request jsonGetCmd;
+
+        static ReadOnlySpan<byte> JSONSETCMD => "*4\r\n$8\r\nJSON.SET\r\n$2\r\nk3\r\n$4\r\n$.f2\r\n$1\r\n2\r\n"u8;
+        Request jsonSetCmd;
+
+        private void RegisterModules()
+        {
+            server.Register.NewModule(new NoOpModule.NoOpModule(), [], out _);
+            server.Register.NewModule(new GarnetJSON.Module(), [], out _);
+        }
+
+        public override void GlobalSetup()
+        {
+            base.GlobalSetup();
+            RegisterModules();
+
+            SetupOperation(ref noOpCmdRead, NOOPCMDREAD);
+            SetupOperation(ref noOpCmdRmw, NOOPCMDRMW);
+            SetupOperation(ref noOpObjRead, NOOPOBJREAD);
+            SetupOperation(ref noOpObjRmw, NOOPOBJRMW);
+            SetupOperation(ref noOpProc, NOOPPROC);
+            SetupOperation(ref noOpTxn, NOOPTXN);
+
+            SetupOperation(ref jsonGetCmd, JSONGETCMD);
+            SetupOperation(ref jsonSetCmd, JSONSETCMD);
+
+            SlowConsumeMessage("*3\r\n$3\r\nSET\r\n$2\r\nk1\r\n$1\r\nc\r\n"u8);
+            SlowConsumeMessage(NOOPOBJRMW);
+            SlowConsumeMessage("*4\r\n$8\r\nJSON.SET\r\n$2\r\nk3\r\n$1\r\n$\r\n$14\r\n{\"f1\":{\"a\":1}}\r\n"u8);
+        }
+
+        [Benchmark]
+        public void ModuleNoOpRawStringReadCommand()
+        {
+            Send(noOpCmdRead);
+        }
+
+        [Benchmark]
+        public void ModuleNoOpRawStringRmwCommand()
+        {
+            Send(noOpCmdRmw);
+        }
+
+        [Benchmark]
+        public void ModuleNoOpObjRmwCommand()
+        {
+            Send(noOpObjRmw);
+        }
+
+        [Benchmark]
+        public void ModuleNoOpObjReadCommand()
+        {
+            Send(noOpObjRead);
+        }
+
+        [Benchmark]
+        public void ModuleNoOpProc()
+        {
+            Send(noOpProc);
+        }
+
+        [Benchmark]
+        public void ModuleNoOpTxn()
+        {
+            Send(noOpTxn);
+        }
+
+        [Benchmark]
+        public void ModuleJsonGetCommand()
+        {
+            Send(jsonGetCmd);
+        }
+
+        [Benchmark]
+        public void ModuleJsonSetCommand()
+        {
+            Send(jsonSetCmd);
+        }
+    }
+}

libs/server/Custom/CustomCommandManager.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Concurrent;
 using System.Diagnostics;
+using Microsoft.Extensions.Logging;
 
 namespace Garnet.server
 {
@@ -24,6 +25,8 @@ public class CustomCommandManager
         private ConcurrentExpandableMap<CustomTransaction> transactionProcMap;
         private ConcurrentExpandableMap<CustomProcedureWrapper> customProcedureMap;
 
+        private readonly ConcurrentDictionary<string, ModuleLoadContext> modules;
+
         internal static readonly int MinMapSize = 8;
         internal static readonly byte CustomTypeIdStartOffset = (byte)CustomObjectTypeMinId;
 
@@ -46,6 +49,8 @@ public CustomCommandManager()
 
             transactionProcMap = new ConcurrentExpandableMap<CustomTransaction>(MinMapSize, 0, byte.MaxValue);
             customProcedureMap = new ConcurrentExpandableMap<CustomProcedureWrapper>(MinMapSize, 0, byte.MaxValue);
+
+            modules = new();
         }
 
         internal int Register(string name, CommandType type, CustomRawStringFunctions customFunctions,
@@ -137,6 +142,45 @@ internal int Register(string name, Func<CustomProcedure> customProcedure, RespCo
             return cmdId;
         }
 
+        /// <summary>
+        /// Register module
+        /// </summary>
+        /// <param name="module">Module to register</param>
+        /// <param name="moduleArgs">Module arguments</param>
+        /// <param name="logger">Logger</param>
+        /// <param name="errorMessage">Error message</param>
+        /// <returns>True if module registered successfully</returns>
+        public bool RegisterModule(ModuleBase module, string[] moduleArgs, ILogger logger,
+            out ReadOnlySpan<byte> errorMessage)
+        {
+            errorMessage = default;
+
+            var moduleLoadContext = new ModuleLoadContext(this, logger);
+            try
+            {
+                module.OnLoad(moduleLoadContext, moduleArgs);
+            }
+            catch (Exception ex)
+            {
+                logger?.LogError(ex, "Error loading module");
+                errorMessage = CmdStrings.RESP_ERR_MODULE_ONLOAD;
+                return false;
+            }
+
+            if (!moduleLoadContext.Initialized)
+            {
+                errorMessage = CmdStrings.RESP_ERR_MODULE_ONLOAD;
+                logger?.LogError("Module {0} failed to initialize", moduleLoadContext.Name);
+                return false;
+            }
+
+            logger?.LogInformation("Module {0} version {1} loaded", moduleLoadContext.Name, moduleLoadContext.Version);
+            return true;
+        }
+
+        internal bool TryAddModule(ModuleLoadContext moduleLoadContext)
+            => modules.TryAdd(moduleLoadContext.Name, moduleLoadContext);
+
         internal bool TryGetCustomProcedure(int id, out CustomProcedureWrapper value)
             => customProcedureMap.TryGetValue(id, out value);
 

libs/server/Module/ModuleRegistrar.cs

@@ -48,16 +48,13 @@ public class ModuleLoadContext
         internal uint Version;
         internal bool Initialized;
 
-        private readonly ModuleRegistrar moduleRegistrar;
         private readonly CustomCommandManager customCommandManager;
 
-        internal ModuleLoadContext(ModuleRegistrar moduleRegistrar, CustomCommandManager customCommandManager, ILogger logger)
+        internal ModuleLoadContext(CustomCommandManager customCommandManager, ILogger logger)
         {
-            Debug.Assert(moduleRegistrar != null);
             Debug.Assert(customCommandManager != null);
 
             Initialized = false;
-            this.moduleRegistrar = moduleRegistrar;
             this.customCommandManager = customCommandManager;
             Logger = logger;
         }
@@ -79,7 +76,7 @@ public ModuleActionStatus Initialize(string name, uint version)
             Name = name;
             Version = version;
 
-            if (!moduleRegistrar.TryAdd(this))
+            if (!customCommandManager.TryAddModule(this))
                 return ModuleActionStatus.AlreadyExists;
 
             Initialized = true;
@@ -177,17 +174,26 @@ public ModuleActionStatus RegisterProcedure(string name, Func<CustomProcedure> c
 
     public sealed class ModuleRegistrar
     {
-        private static readonly Lazy<ModuleRegistrar> lazy = new Lazy<ModuleRegistrar>(() => new ModuleRegistrar());
+        private static readonly Lazy<ModuleRegistrar> lazy = new(() => new ModuleRegistrar());
 
         public static ModuleRegistrar Instance { get { return lazy.Value; } }
 
         private ModuleRegistrar()
         {
-            modules = new();
+            moduleConstructors = new();
         }
 
-        private readonly ConcurrentDictionary<string, ModuleLoadContext> modules;
+        private readonly ConcurrentDictionary<string, ConstructorInfo> moduleConstructors;
 
+        /// <summary>
+        /// Load module from a specified assembly and register it with the custom command managers
+        /// </summary>
+        /// <param name="customCommandManager">Custom command manager instance</param>
+        /// <param name="loadedAssembly">Loaded assembly containing module</param>
+        /// <param name="moduleArgs">Module arguments</param>
+        /// <param name="logger">Logger</param>
+        /// <param name="errorMessage">Error message</param>
+        /// <returns>True if module loaded successfully</returns>
         public bool LoadModule(CustomCommandManager customCommandManager, Assembly loadedAssembly, string[] moduleArgs, ILogger logger, out ReadOnlySpan<byte> errorMessage)
         {
             errorMessage = default;
@@ -210,47 +216,40 @@ public bool LoadModule(CustomCommandManager customCommandManager, Assembly loade
                 return false;
             }
 
-            // Check that type has empty constructor
             var moduleType = loadedTypes[0];
-            if (moduleType.GetConstructor(Type.EmptyTypes) == null)
-            {
-                errorMessage = CmdStrings.RESP_ERR_GENERIC_INSTANTIATING_CLASS;
-                return false;
-            }
 
-            var instance = (ModuleBase)Activator.CreateInstance(moduleType);
-            if (instance == null)
+            // Uniquely identify the module by assembly and type name
+            var moduleAssemblyAnyTypeName = $"{loadedAssembly.FullName}:{moduleType.FullName!}";
+
+            // Check if module was previously loaded
+            if (!moduleConstructors.TryGetValue(moduleAssemblyAnyTypeName, out var moduleConstructor))
             {
-                errorMessage = CmdStrings.RESP_ERR_GENERIC_INSTANTIATING_CLASS;
-                return false;
+                // Check that type has empty constructor
+                moduleConstructor = moduleType.GetConstructor(Type.EmptyTypes);
+
+                if (moduleConstructor == null)
+                {
+                    errorMessage = CmdStrings.RESP_ERR_GENERIC_INSTANTIATING_CLASS;
+                    return false;
+                }
+
+                moduleConstructors.TryAdd(moduleAssemblyAnyTypeName, moduleConstructor);
             }
 
-            var moduleLoadContext = new ModuleLoadContext(this, customCommandManager, logger);
+            // Instantiate module using constructor info
+            ModuleBase module;
             try
             {
-                instance.OnLoad(moduleLoadContext, moduleArgs);
+                module = (ModuleBase)moduleConstructor.Invoke([]);
             }
-            catch (Exception ex)
-            {
-                logger?.LogError(ex, "Error loading module");
-                errorMessage = CmdStrings.RESP_ERR_MODULE_ONLOAD;
-                return false;
-            }
-
-            if (!moduleLoadContext.Initialized)
+            catch (Exception)
             {
-                errorMessage = CmdStrings.RESP_ERR_MODULE_ONLOAD;
-                logger?.LogError("Module {0} failed to initialize", moduleLoadContext.Name);
+                errorMessage = CmdStrings.RESP_ERR_GENERIC_INSTANTIATING_CLASS;
                 return false;
             }
 
-            logger?.LogInformation("Module {0} version {1} loaded", moduleLoadContext.Name, moduleLoadContext.Version);
-            return true;
-        }
-
-        internal bool TryAdd(ModuleLoadContext moduleLoadContext)
-        {
-            return modules.TryAdd(moduleLoadContext.Name, moduleLoadContext);
+            // Register module with custom command manager
+            return customCommandManager.RegisterModule(module, moduleArgs, logger, out errorMessage);
         }
     }
 }

libs/server/Servers/RegisterApi.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license.
 
 using System;
+using Microsoft.Extensions.Logging;
 
 namespace Garnet.server
 {
@@ -80,5 +81,16 @@ public int NewType(CustomObjectFactory factory)
         /// <returns></returns>
         public int NewProcedure(string name, Func<CustomProcedure> customProcedure, RespCommandsInfo commandInfo = null, RespCommandDocs commandDocs = null)
             => provider.StoreWrapper.customCommandManager.Register(name, customProcedure, commandInfo, commandDocs);
+
+        /// <summary>
+        /// Register custom module with Garnet
+        /// </summary>
+        /// <param name="module"></param>
+        /// <param name="moduleArgs"></param>
+        /// <param name="logger"></param>
+        /// <param name="errorMessage"></param>
+        /// <returns></returns>
+        public bool NewModule(ModuleBase module, string[] moduleArgs, out ReadOnlySpan<byte> errorMessage, ILogger logger = null)
+            => provider.StoreWrapper.customCommandManager.RegisterModule(module, moduleArgs, logger, out errorMessage);
     }
 }