Perf Improvements to Custom Commands' Execution Path

libs/server/Custom/CustomCommandManagerSession.cs

@@ -1,100 +1,190 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-using System.Diagnostics;
+using System;
+using System.Collections.Generic;
 using Garnet.common;
 
 namespace Garnet.server
 {
     /// <summary>
-    /// Server session for RESP protocol - basic commands are in this file
+    /// Session-specific access to custom command data held by CustomCommandManager
+    /// For both custom procedures and custom transactions, this class maintains a cached instance of the transaction / procedure that can be called directly per-session
+    /// This class also maintains a cache of custom command info and docs to avoid the cost of synchronization when calling CustomCommandManager
     /// </summary>
     internal sealed class CustomCommandManagerSession
     {
+        // Initial size of expandable maps
+        private static readonly int MinMapSize = 8;
+
+        // The instance of the CustomCommandManager that this session is associated with
         readonly CustomCommandManager customCommandManager;
 
-        // These session specific arrays are indexed by the same ID as the arrays in CustomCommandManager
-        ExpandableMap<CustomTransactionProcedureWithArity> sessionTransactionProcMap;
+        // These session specific maps are indexed by the same ID as the arrays in CustomCommandManager
+        // Maps between the transaction ID and a tuple of the per-session transaction procedure and its arity
+        ExpandableMap<(CustomTransactionProcedure, int)> sessionTransactionProcMap;
+        // Maps between the custom procedure ID and the per-session custom procedure instance
         ExpandableMap<CustomProcedure> sessionCustomProcMap;
 
         public CustomCommandManagerSession(CustomCommandManager customCommandManager)
         {
             this.customCommandManager = customCommandManager;
-            sessionTransactionProcMap = new ExpandableMap<CustomTransactionProcedureWithArity>(CustomCommandManager.MinMapSize, 0, byte.MaxValue);
-            sessionCustomProcMap = new ExpandableMap<CustomProcedure>(CustomCommandManager.MinMapSize, 0, byte.MaxValue);
+
+            sessionTransactionProcMap = new ExpandableMap<(CustomTransactionProcedure, int)>(MinMapSize, 0, byte.MaxValue);
+            sessionCustomProcMap = new ExpandableMap<CustomProcedure>(MinMapSize, 0, byte.MaxValue);
         }
 
+        /// <summary>
+        /// Get a custom procedure by ID
+        /// </summary>
+        /// <param name="id">The procedure ID</param>
+        /// <param name="respServerSession">The current session</param>
+        /// <returns>The per-session instance of the procedure</returns>
+        /// <exception cref="GarnetException"></exception>
         public CustomProcedure GetCustomProcedure(int id, RespServerSession respServerSession)
         {
+            // Check if we already have a cached entry of the per-session custom procedure instance
             if (!sessionCustomProcMap.TryGetValue(id, out var customProc))
             {
+                // If not, get the custom procedure from the CustomCommandManager
                 if (!customCommandManager.TryGetCustomProcedure(id, out var entry))
                     throw new GarnetException($"Custom procedure {id} not found");
 
+                // Create the session-specific instance and add it to the cache
                 customProc = entry.CustomProcedureFactory();
                 customProc.respServerSession = respServerSession;
-                var setSuccessful = sessionCustomProcMap.TrySetValue(id, ref customProc);
-                Debug.Assert(setSuccessful);
+                sessionCustomProcMap.TrySetValue(id, ref customProc);
             }
 
             return customProc;
         }
 
+        /// <summary>
+        /// Get a custom transaction procedure by ID
+        /// </summary>
+        /// <param name="id">The transaction ID</param>
+        /// <param name="respServerSession">The current session</param>
+        /// <param name="txnManager">txnManager</param>
+        /// <param name="scratchBufferManager">scratchBufferManager</param>
+        /// <param name="arity">The arity of the transaction</param>
+        /// <returns>The per-session instance of the transaction</returns>
+        /// <exception cref="GarnetException"></exception>
         public CustomTransactionProcedure GetCustomTransactionProcedure(int id, RespServerSession respServerSession, TransactionManager txnManager, ScratchBufferManager scratchBufferManager, out int arity)
         {
-            if (sessionTransactionProcMap.Exists(id))
+            // Check if we already have a cached entry of the per-session custom transaction instance
+            if (sessionTransactionProcMap.TryGetValue(id, out var tranToArity))
             {
-                ref var customTranProc = ref sessionTransactionProcMap.GetValueByRef(id);
-                if (customTranProc.Procedure != null)
+                if (tranToArity.Item1 != null)
                 {
-                    arity = customTranProc.Arity;
-                    return customTranProc.Procedure;
+                    arity = tranToArity.Item2;
+                    return tranToArity.Item1;
                 }
             }
 
             if (!customCommandManager.TryGetCustomTransactionProcedure(id, out var entry))
                 throw new GarnetException($"Transaction procedure {id} not found");
-            _ = customCommandManager.customCommandsInfo.TryGetValue(entry.NameStr, out var cmdInfo);
-            arity = cmdInfo?.Arity ?? 0;
-            return GetCustomTransactionProcedureAndSetArity(entry, respServerSession, txnManager, scratchBufferManager, cmdInfo?.Arity ?? 0);
+
+            // Create the session-specific instance and add it to the cache
+            var customTranProc = entry.proc();
+            customTranProc.txnManager = txnManager;
+            customTranProc.scratchBufferManager = scratchBufferManager;
+            customTranProc.respServerSession = respServerSession;
+
+            arity = entry.arity;
+            tranToArity = new ValueTuple<CustomTransactionProcedure, int>(customTranProc, arity);
+            sessionTransactionProcMap.TrySetValue(id, ref tranToArity);
+
+            return customTranProc;
         }
 
+        /// <summary>
+        /// Get a custom raw-string command by name
+        /// </summary>
+        /// <param name="command">The command name to match</param>
+        /// <param name="cmd">The matching command</param>
+        /// <returns>True if command name matched an existing command</returns>
+        public bool Match(ReadOnlySpan<byte> command, out CustomRawStringCommand cmd)
+            => customCommandManager.Match(command, out cmd);
+
+        /// <summary>
+        /// Get a custom transaction by name
+        /// </summary>
+        /// <param name="command">The transaction name to match</param>
+        /// <param name="cmd">The matching transaction</param>
+        /// <returns>True if transaction name matched an existing transaction</returns>
+        public bool Match(ReadOnlySpan<byte> command, out CustomTransaction cmd)
+            => customCommandManager.Match(command, out cmd);
+
+        /// <summary>
+        /// Get a custom object command by name
+        /// </summary>
+        /// <param name="command">The command name to match</param>
+        /// <param name="cmd">The matching command</param>
+        /// <returns>True if command name matched an existing command</returns>
+        public bool Match(ReadOnlySpan<byte> command, out CustomObjectCommand cmd)
+            => customCommandManager.Match(command, out cmd);
+
+        /// <summary>
+        /// Get a custom procedure by name
+        /// </summary>
+        /// <param name="command">The procedure name to match</param>
+        /// <param name="cmd">The matching procedure</param>
+        /// <returns>True if procedure name matched an existing procedure</returns>
+        public bool Match(ReadOnlySpan<byte> command, out CustomProcedureWrapper cmd)
+            => customCommandManager.Match(command, out cmd);
+
+        /// <summary>
+        /// Get custom command info by name
+        /// </summary>
+        /// <param name="cmdName">The command name</param>
+        /// <param name="respCommandsInfo">The matching command info</param>
+        /// <returns>True if command info was found</returns>
+        public bool TryGetCustomCommandInfo(string cmdName, out RespCommandsInfo respCommandsInfo)
+            => customCommandManager.TryGetCustomCommandInfo(cmdName, out respCommandsInfo);
+
+        /// <summary>
+        /// Get custom command docs by name
+        /// </summary>
+        /// <param name="cmdName">The command name</param>
+        /// <param name="respCommandsDocs">The matching command docs</param>
+        /// <returns>True if command docs was found</returns>
+        public bool TryGetCustomCommandDocs(string cmdName, out RespCommandDocs respCommandsDocs)
+            => customCommandManager.TryGetCustomCommandDocs(cmdName, out respCommandsDocs);
+
+        /// <summary>
+        /// Get all custom command infos
+        /// </summary>
+        /// <returns>Map between custom command name and custom command info</returns>
+        internal IReadOnlyDictionary<string, RespCommandsInfo> GetAllCustomCommandsInfos()
+            => customCommandManager.GetAllCustomCommandsInfos();
+
+        /// <summary>
+        /// Get all custom command docs
+        /// </summary>
+        /// <returns>Map between custom command name and custom command docs</returns>
+        internal IReadOnlyDictionary<string, RespCommandDocs> GetAllCustomCommandsDocs()
+            => customCommandManager.GetAllCustomCommandsDocs();
+
+        /// <summary>
+        /// Get count of all custom command infos
+        /// </summary>
+        /// <returns>Count</returns>
+        internal int GetCustomCommandInfoCount() => customCommandManager.GetCustomCommandInfoCount();
+
+        /// <summary>
+        /// Get RespCommand enum by command ID
+        /// </summary>
+        /// <param name="id">Command ID</param>
+        /// <returns>Matching RespCommand</returns>
         public RespCommand GetCustomRespCommand(int id)
             => customCommandManager.GetCustomRespCommand(id);
 
+        /// <summary>
+        /// Get GarnetObjectType enum by object type ID
+        /// </summary>
+        /// <param name="id">Object type ID</param>
+        /// <returns>Matching GarnetObjectType</returns>
         public GarnetObjectType GetCustomGarnetObjectType(int id)
             => customCommandManager.GetCustomGarnetObjectType(id);
-
-        private CustomTransactionProcedure GetCustomTransactionProcedureAndSetArity(CustomTransaction entry, RespServerSession respServerSession, TransactionManager txnManager, ScratchBufferManager scratchBufferManager, int arity)
-        {
-            int id = entry.id;
-
-            var customTranProc = new CustomTransactionProcedureWithArity(entry.proc(), arity)
-            {
-                Procedure =
-                {
-                    txnManager = txnManager,
-                    scratchBufferManager = scratchBufferManager,
-                    respServerSession = respServerSession
-                }
-            };
-            var setSuccessful = sessionTransactionProcMap.TrySetValue(id, ref customTranProc);
-            Debug.Assert(setSuccessful);
-
-            return customTranProc.Procedure;
-        }
-
-        private struct CustomTransactionProcedureWithArity
-        {
-            public CustomTransactionProcedure Procedure { get; }
-
-            public int Arity { get; }
-
-            public CustomTransactionProcedureWithArity(CustomTransactionProcedure procedure, int arity)
-            {
-                this.Procedure = procedure;
-                this.Arity = arity;
-            }
-        }
     }
 }

libs/server/Custom/CustomObjectCommand.cs

@@ -11,16 +11,18 @@ public class CustomObjectCommand : ICustomCommand
         public readonly byte id;
         public readonly byte subid;
         public readonly CommandType type;
+        public readonly int arity;
         public readonly CustomObjectFactory factory;
         public readonly CustomObjectFunctions functions;
 
-        internal CustomObjectCommand(string name, byte id, byte subid, CommandType type, CustomObjectFactory factory, CustomObjectFunctions functions = null)
+        internal CustomObjectCommand(string name, byte id, byte subid, CommandType type, int arity, CustomObjectFactory factory, CustomObjectFunctions functions = null)
         {
             NameStr = name.ToUpperInvariant();
             this.Name = System.Text.Encoding.ASCII.GetBytes(NameStr);
             this.id = id;
             this.subid = subid;
             this.type = type;
+            this.arity = arity;
             this.factory = factory;
             this.functions = functions;
         }

libs/server/Custom/CustomObjectCommandWrapper.cs

@@ -13,13 +13,13 @@ class CustomObjectCommandWrapper
 
         public readonly byte id;
         public readonly CustomObjectFactory factory;
-        public ConcurrentExpandableMap<CustomObjectCommand> commandMap;
+        public ExpandableMap<CustomObjectCommand> commandMap;
 
         public CustomObjectCommandWrapper(byte id, CustomObjectFactory functions)
         {
             this.id = id;
             this.factory = functions;
-            this.commandMap = new ConcurrentExpandableMap<CustomObjectCommand>(MinMapSize, 0, MaxSubId);
+            this.commandMap = new ExpandableMap<CustomObjectCommand>(MinMapSize, 0, MaxSubId);
         }
     }
 }

libs/server/Custom/CustomProcedureWrapper.cs

@@ -28,9 +28,10 @@ class CustomProcedureWrapper : ICustomCommand
 
         public readonly string NameStr;
         public readonly byte Id;
+        public readonly int Arity;
         public readonly Func<CustomProcedure> CustomProcedureFactory;
 
-        internal CustomProcedureWrapper(string name, byte id, Func<CustomProcedure> customProcedureFactory, CustomCommandManager customCommandManager)
+        internal CustomProcedureWrapper(string name, byte id, int arity, Func<CustomProcedure> customProcedureFactory, CustomCommandManager customCommandManager)
         {
             if (string.IsNullOrEmpty(name))
                 throw new ArgumentNullException(nameof(name));
@@ -43,6 +44,7 @@ internal CustomProcedureWrapper(string name, byte id, Func<CustomProcedure> cust
             NameStr = name.ToUpperInvariant();
             Name = System.Text.Encoding.ASCII.GetBytes(NameStr);
             Id = id;
+            Arity = arity;
             CustomProcedureFactory = customProcedureFactory;
         }
     }

libs/server/Custom/CustomRawStringCommand.cs

@@ -10,15 +10,17 @@ public class CustomRawStringCommand : ICustomCommand
         public readonly string NameStr;
         public readonly ushort id;
         public readonly CommandType type;
+        public readonly int arity;
         public readonly CustomRawStringFunctions functions;
         public long expirationTicks;
 
-        internal CustomRawStringCommand(string name, ushort id, CommandType type, CustomRawStringFunctions functions, long expirationTicks)
+        internal CustomRawStringCommand(string name, ushort id, CommandType type, int arity, CustomRawStringFunctions functions, long expirationTicks)
         {
             NameStr = name.ToUpperInvariant();
             this.Name = System.Text.Encoding.ASCII.GetBytes(NameStr);
             this.id = id;
             this.type = type;
+            this.arity = arity;
             this.functions = functions;
             this.expirationTicks = expirationTicks;
         }

libs/server/Custom/CustomRespCommands.cs

@@ -200,14 +200,14 @@ private bool TryCustomObjectCommand<TGarnetApi>(GarnetObjectType objType, byte s
         /// <param name="customCommand">Parsed raw string command</param>
         /// <returns>True if command found, false otherwise</returns>
         public bool ParseCustomRawStringCommand(string cmd, out CustomRawStringCommand customCommand) =>
-            storeWrapper.customCommandManager.Match(new ReadOnlySpan<byte>(Encoding.UTF8.GetBytes(cmd)), out customCommand);
+            customCommandManagerSession.Match(new ReadOnlySpan<byte>(Encoding.UTF8.GetBytes(cmd)), out customCommand);
 
         /// <summary>Parse custom object command</summary>
         /// <param name="cmd">Command name</param>
         /// <param name="customObjCommand">Parsed object command</param>
         /// <returns>True if command found, false othrewise</returns>
         public bool ParseCustomObjectCommand(string cmd, out CustomObjectCommand customObjCommand) =>
-            storeWrapper.customCommandManager.Match(new ReadOnlySpan<byte>(Encoding.UTF8.GetBytes(cmd)), out customObjCommand);
+            customCommandManagerSession.Match(new ReadOnlySpan<byte>(Encoding.UTF8.GetBytes(cmd)), out customObjCommand);
 
         /// <summary>Execute a specific custom raw string command</summary>
         /// <typeparam name="TGarnetApi"></typeparam>

libs/server/Custom/CustomTransaction.cs

@@ -12,15 +12,17 @@ class CustomTransaction : ICustomCommand
 
         public readonly string NameStr;
         public readonly byte id;
+        public readonly int arity;
         public readonly Func<CustomTransactionProcedure> proc;
 
-        internal CustomTransaction(string name, byte id, Func<CustomTransactionProcedure> proc)
+        internal CustomTransaction(string name, byte id, int arity, Func<CustomTransactionProcedure> proc)
         {
             if (name == null)
                 throw new GarnetException("CustomTransaction name is null");
             NameStr = name.ToUpperInvariant();
             this.Name = System.Text.Encoding.ASCII.GetBytes(NameStr);
             this.id = id;
+            this.arity = arity;
             this.proc = proc ?? throw new GarnetException("CustomTransactionProcedure is null");
         }
     }