Returning "Friendly IDs" for Custom Commands / Custom Types (#906)

* Converting custom command IDs to "friendly IDs" when returning to client

* Some simplifications

* Simplified ExpandableMap to support only ascending IDs + Added test for multiple object registrations

* small format

libs/server/Custom/CustomCommandManager.cs

@@ -12,14 +12,21 @@ namespace Garnet.server
     /// </summary>
     public class CustomCommandManager
     {
-        internal static readonly int MinMapSize = 8;
-        internal static readonly byte TypeIdStartOffset = byte.MaxValue - (byte)GarnetObjectTypeExtensions.FirstSpecialObjectType;
+        private static readonly int MaxCustomRawStringCommands = 256;
+
+        private static readonly int CustomRawStringCommandMinId = (ushort)RespCommand.INVALID - MaxCustomRawStringCommands;
+        private static readonly int CustomRawStringCommandMaxId = (ushort)RespCommand.INVALID - 1;
+        private static readonly int CustomObjectTypeMinId = (byte)GarnetObjectTypeExtensions.LastObjectType + 1;
+        private static readonly int CustomObjectTypeMaxId = (byte)GarnetObjectTypeExtensions.FirstSpecialObjectType - 1;
 
         private ConcurrentExpandableMap<CustomRawStringCommand> rawStringCommandMap;
         private ConcurrentExpandableMap<CustomObjectCommandWrapper> objectCommandMap;
         private ConcurrentExpandableMap<CustomTransaction> transactionProcMap;
         private ConcurrentExpandableMap<CustomProcedureWrapper> customProcedureMap;
 
+        internal static readonly int MinMapSize = 8;
+        internal static readonly byte CustomTypeIdStartOffset = (byte)CustomObjectTypeMinId;
+
         internal int CustomCommandsInfoCount => customCommandsInfo.Count;
         internal readonly ConcurrentDictionary<string, RespCommandsInfo> customCommandsInfo = new(StringComparer.OrdinalIgnoreCase);
         internal readonly ConcurrentDictionary<string, RespCommandDocs> customCommandsDocs = new(StringComparer.OrdinalIgnoreCase);
@@ -29,30 +36,35 @@ public class CustomCommandManager
         /// </summary>
         public CustomCommandManager()
         {
+            Debug.Assert(CustomRawStringCommandMinId > (ushort)RespCommandExtensions.LastValidCommand);
+
             rawStringCommandMap = new ConcurrentExpandableMap<CustomRawStringCommand>(MinMapSize,
-                (ushort)RespCommand.INVALID - 1,
-                (ushort)RespCommandExtensions.LastValidCommand + 1);
+                CustomRawStringCommandMinId, CustomRawStringCommandMaxId);
+
             objectCommandMap = new ConcurrentExpandableMap<CustomObjectCommandWrapper>(MinMapSize,
-                (byte)GarnetObjectTypeExtensions.FirstSpecialObjectType - 1,
-                (byte)GarnetObjectTypeExtensions.LastObjectType + 1);
+                CustomObjectTypeMinId, CustomObjectTypeMaxId);
+
             transactionProcMap = new ConcurrentExpandableMap<CustomTransaction>(MinMapSize, 0, byte.MaxValue);
             customProcedureMap = new ConcurrentExpandableMap<CustomProcedureWrapper>(MinMapSize, 0, byte.MaxValue);
         }
 
-        internal int Register(string name, CommandType type, CustomRawStringFunctions customFunctions, RespCommandsInfo commandInfo, RespCommandDocs commandDocs, long expirationTicks)
+        internal int Register(string name, CommandType type, CustomRawStringFunctions customFunctions,
+            RespCommandsInfo commandInfo, RespCommandDocs commandDocs, long expirationTicks)
         {
             if (!rawStringCommandMap.TryGetNextId(out var cmdId))
                 throw new Exception("Out of registration space");
             Debug.Assert(cmdId <= ushort.MaxValue);
-            var newCmd = new CustomRawStringCommand(name, (ushort)cmdId, type, customFunctions, expirationTicks);
+            var extId = cmdId - CustomRawStringCommandMinId;
+            var newCmd = new CustomRawStringCommand(name, (ushort)extId, type, customFunctions, expirationTicks);
             var setSuccessful = rawStringCommandMap.TrySetValue(cmdId, ref newCmd);
             Debug.Assert(setSuccessful);
             if (commandInfo != null) customCommandsInfo.AddOrUpdate(name, commandInfo, (_, _) => commandInfo);
             if (commandDocs != null) customCommandsDocs.AddOrUpdate(name, commandDocs, (_, _) => commandDocs);
-            return cmdId;
+            return extId;
         }
 
-        internal int Register(string name, Func<CustomTransactionProcedure> proc, RespCommandsInfo commandInfo = null, RespCommandDocs commandDocs = null)
+        internal int Register(string name, Func<CustomTransactionProcedure> proc, RespCommandsInfo commandInfo = null,
+            RespCommandDocs commandDocs = null)
         {
             if (!transactionProcMap.TryGetNextId(out var cmdId))
                 throw new Exception("Out of registration space");
@@ -71,45 +83,33 @@ internal int RegisterType(CustomObjectFactory factory)
             if (objectCommandMap.TryGetFirstId(c => c.factory == factory, out var dupRegistrationId))
                 throw new Exception($"Type already registered with ID {dupRegistrationId}");
 
-            if (!objectCommandMap.TryGetNextId(out var cmdId))
-                throw new Exception("Out of registration space");
-            Debug.Assert(cmdId <= byte.MaxValue);
-
-            var newCmd = new CustomObjectCommandWrapper((byte)cmdId, factory);
-            var setSuccessful = objectCommandMap.TrySetValue(cmdId, ref newCmd);
-            Debug.Assert(setSuccessful);
-
-            return cmdId;
+            var typeId = RegisterNewType(factory);
+            return typeId - CustomObjectTypeMinId;
         }
 
         internal (int objectTypeId, int subCommand) Register(string name, CommandType commandType, CustomObjectFactory factory, RespCommandsInfo commandInfo, RespCommandDocs commandDocs, CustomObjectFunctions customObjectFunctions = null)
         {
             if (!objectCommandMap.TryGetFirstId(c => c.factory == factory, out var typeId))
             {
-                if (!objectCommandMap.TryGetNextId(out typeId))
-                    throw new Exception("Out of registration space");
-
-                Debug.Assert(typeId <= byte.MaxValue);
-
-                var newCmd = new CustomObjectCommandWrapper((byte)typeId, factory);
-                var setSuccessful = objectCommandMap.TrySetValue(typeId, ref newCmd);
-                Debug.Assert(setSuccessful);
+                typeId = RegisterNewType(factory);
             }
 
+            var extId = typeId - CustomObjectTypeMinId;
+
             objectCommandMap.TryGetValue(typeId, out var wrapper);
             if (!wrapper.commandMap.TryGetNextId(out var scId))
                 throw new Exception("Out of registration space");
 
             Debug.Assert(scId <= byte.MaxValue);
-            var newSubCmd = new CustomObjectCommand(name, (byte)typeId, (byte)scId, commandType, wrapper.factory,
+            var newSubCmd = new CustomObjectCommand(name, (byte)extId, (byte)scId, commandType, wrapper.factory,
                 customObjectFunctions);
             var scSetSuccessful = wrapper.commandMap.TrySetValue(scId, ref newSubCmd);
             Debug.Assert(scSetSuccessful);
 
             if (commandInfo != null) customCommandsInfo.AddOrUpdate(name, commandInfo, (_, _) => commandInfo);
             if (commandDocs != null) customCommandsDocs.AddOrUpdate(name, commandDocs, (_, _) => commandDocs);
 
-            return (typeId, scId);
+            return (extId, scId);
         }
 
         /// <summary>
@@ -177,5 +177,25 @@ internal bool TryGetCustomCommandDocs(string cmdName, out RespCommandDocs respCo
         {
             return this.customCommandsDocs.TryGetValue(cmdName, out respCommandsDocs);
         }
+
+        internal RespCommand GetCustomRespCommand(int id)
+            => (RespCommand)(CustomRawStringCommandMinId + id);
+
+        internal GarnetObjectType GetCustomGarnetObjectType(int id)
+            => (GarnetObjectType)(CustomObjectTypeMinId + id);
+
+        private int RegisterNewType(CustomObjectFactory factory)
+        {
+            if (!objectCommandMap.TryGetNextId(out var typeId))
+                throw new Exception("Out of registration space");
+            Debug.Assert(typeId <= byte.MaxValue);
+
+            var extId = typeId - CustomObjectTypeMinId;
+            var newCmd = new CustomObjectCommandWrapper((byte)extId, factory);
+            var setSuccessful = objectCommandMap.TrySetValue(typeId, ref newCmd);
+            Debug.Assert(setSuccessful);
+
+            return typeId;
+        }
     }
 }

libs/server/Custom/CustomCommandManagerSession.cs

@@ -59,6 +59,12 @@ public CustomTransactionProcedure GetCustomTransactionProcedure(int id, RespServ
             return GetCustomTransactionProcedureAndSetArity(entry, respServerSession, txnManager, scratchBufferManager, cmdInfo?.Arity ?? 0);
         }
 
+        public RespCommand GetCustomRespCommand(int id)
+            => customCommandManager.GetCustomRespCommand(id);
+
+        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;

libs/server/Custom/CustomRespCommands.cs

@@ -54,7 +54,6 @@ public bool RunTransactionProc(byte id, ref CustomProcedureInput procInput, ref
             var proc = customCommandManagerSession
                 .GetCustomTransactionProcedure(id, this, txnManager, scratchBufferManager, out _);
             return txnManager.RunTransactionProc(id, ref procInput, proc, ref output);
-
         }
 
         private void TryCustomProcedure(CustomProcedure proc, int startIdx = 0)
@@ -226,7 +225,8 @@ public bool InvokeCustomRawStringCommand<TGarnetApi>(ref TGarnetApi storageApi,
             var sbKey = key.SpanByte;
             var inputArg = customCommand.expirationTicks > 0 ? DateTimeOffset.UtcNow.Ticks + customCommand.expirationTicks : customCommand.expirationTicks;
             customCommandParseState.InitializeWithArguments(args);
-            var rawStringInput = new RawStringInput((RespCommand)customCommand.id, ref customCommandParseState, arg1: inputArg);
+            var cmd = customCommandManagerSession.GetCustomRespCommand(customCommand.id);
+            var rawStringInput = new RawStringInput(cmd, ref customCommandParseState, arg1: inputArg);
 
             var _output = new SpanByteAndMemory(null);
             GarnetStatus status;
@@ -290,7 +290,8 @@ public bool InvokeCustomObjectCommand<TGarnetApi>(ref TGarnetApi storageApi, Cus
             var keyBytes = key.ToArray();
 
             // Prepare input
-            var header = new RespInputHeader((GarnetObjectType)customObjCommand.id) { SubId = customObjCommand.subid };
+            var type = customCommandManagerSession.GetCustomGarnetObjectType(customObjCommand.id);
+            var header = new RespInputHeader(type) { SubId = customObjCommand.subid };
             customCommandParseState.InitializeWithArguments(args);
             var input = new ObjectInput(header, ref customCommandParseState);
 

libs/server/Custom/ExpandableMap.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license.
 
 using System;
+using System.Diagnostics;
 using System.Threading;
 using Garnet.common;
 
@@ -88,8 +89,6 @@ internal struct ExpandableMap<T> : IExpandableMap<T>
         readonly int minId;
         // Value of max item ID
         readonly int maxSize;
-        // True if item IDs are in descending order
-        readonly bool descIds;
 
         /// <summary>
         /// Creates a new instance of ExpandableMap
@@ -99,18 +98,19 @@ internal struct ExpandableMap<T> : IExpandableMap<T>
         /// <param name="maxId">The maximal item ID value (can be smaller than minId for descending order of IDs)</param>
         public ExpandableMap(int minSize, int minId, int maxId)
         {
+            Debug.Assert(maxId > minId);
+
             this.Map = null;
             this.minSize = minSize;
             this.minId = minId;
-            this.maxSize = Math.Abs(maxId - minId) + 1;
-            this.descIds = minId > maxId;
+            this.maxSize = maxId - minId + 1;
         }
 
         /// <inheritdoc />
         public bool TryGetValue(int id, out T value)
         {
             value = default;
-            var idx = GetIndexFromId(id);
+            var idx = id - minId;
             if (idx < 0 || idx >= ActualSize)
                 return false;
 
@@ -121,14 +121,14 @@ public bool TryGetValue(int id, out T value)
         /// <inheritdoc />
         public bool Exists(int id)
         {
-            var idx = GetIndexFromId(id);
+            var idx = id - minId;
             return idx >= 0 && idx < ActualSize;
         }
 
         /// <inheritdoc />
         public ref T GetValueByRef(int id)
         {
-            var idx = GetIndexFromId(id);
+            var idx = id - minId;
             if (idx < 0 || idx >= ActualSize)
                 throw new ArgumentOutOfRangeException(nameof(idx));
 
@@ -147,7 +147,7 @@ public bool TryGetNextId(out int id)
 
             if (nextIdx >= maxSize)
                 return false;
-            id = GetIdFromIndex(nextIdx);
+            id = minId + nextIdx;
 
             return true;
         }
@@ -160,7 +160,7 @@ public bool TryGetFirstId(Func<T, bool> predicate, out int id)
             {
                 if (predicate(Map[i]))
                 {
-                    id = GetIdFromIndex(i);
+                    id = minId + i;
                     return true;
                 }
             }
@@ -180,7 +180,7 @@ public bool TryGetNextIdSafe(out int id)
 
             if (nextIdx >= maxSize)
                 return false;
-            id = GetIdFromIndex(nextIdx);
+            id = minId + nextIdx;
 
             return true;
         }
@@ -193,7 +193,7 @@ public bool TryGetNextIdSafe(out int id)
         /// <returns>True if actual size should be updated (or was updated if noUpdate is false)</returns>
         internal bool TryUpdateSize(int id, bool noUpdate = false)
         {
-            var idx = GetIndexFromId(id);
+            var idx = id - minId;
 
             // Should not update the size if the index is out of bounds
             // or if index is smaller than the current actual size
@@ -215,7 +215,7 @@ internal bool TryUpdateSize(int id, bool noUpdate = false)
         /// <returns>True if assignment succeeded</returns>
         internal bool TrySetValue(int id, ref T value, bool noExpansion, bool updateSize)
         {
-            var idx = GetIndexFromId(id);
+            var idx = id - minId;
             if (idx < 0 || idx >= maxSize) return false;
 
             // If index within array bounds, set item
@@ -247,20 +247,6 @@ internal bool TrySetValue(int id, ref T value, bool noExpansion, bool updateSize
             if (updateSize) TryUpdateSize(id);
             return true;
         }
-
-        /// <summary>
-        /// Maps map index to item ID
-        /// </summary>
-        /// <param name="index">Map index</param>
-        /// <returns>Item ID</returns>
-        private int GetIdFromIndex(int index) => descIds ? minId - index : index;
-
-        /// <summary>
-        /// Maps an item ID to a map index
-        /// </summary>
-        /// <param name="id">Item ID</param>
-        /// <returns>Map index</returns>
-        private int GetIndexFromId(int id) => descIds ? minId - id : id;
     }
 
     /// <summary>

libs/server/Objects/Types/GarnetObjectSerializer.cs

@@ -58,7 +58,7 @@ private IGarnetObject DeserializeInternal(BinaryReader binaryReader)
 
         private IGarnetObject CustomDeserialize(byte type, BinaryReader binaryReader)
         {
-            if (type < CustomCommandManager.TypeIdStartOffset ||
+            if (type < CustomCommandManager.CustomTypeIdStartOffset ||
                 !customCommandManager.TryGetCustomObjectCommand(type, out var cmd)) return null;
             return cmd.factory.Deserialize(type, binaryReader);
         }

libs/server/Resp/RespServerSession.cs

@@ -826,8 +826,8 @@ private bool NetworkCustomRawStringCmd<TGarnetApi>(ref TGarnetApi storageApi)
             }
 
             // Perform the operation
-            TryCustomRawStringCommand((RespCommand)currentCustomRawStringCommand.id,
-                currentCustomRawStringCommand.expirationTicks, currentCustomRawStringCommand.type, ref storageApi);
+            var cmd = customCommandManagerSession.GetCustomRespCommand(currentCustomRawStringCommand.id);
+            TryCustomRawStringCommand(cmd, currentCustomRawStringCommand.expirationTicks, currentCustomRawStringCommand.type, ref storageApi);
             currentCustomRawStringCommand = null;
             return true;
         }
@@ -842,7 +842,8 @@ bool NetworkCustomObjCmd<TGarnetApi>(ref TGarnetApi storageApi)
             }
 
             // Perform the operation
-            TryCustomObjectCommand((GarnetObjectType)currentCustomObjectCommand.id, currentCustomObjectCommand.subid,
+            var type = customCommandManagerSession.GetCustomGarnetObjectType(currentCustomObjectCommand.id);
+            TryCustomObjectCommand(type, currentCustomObjectCommand.subid,
                 currentCustomObjectCommand.type, ref storageApi);
             currentCustomObjectCommand = null;
             return true;

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.TypeIdStartOffset)
+                    if ((byte)type < CustomCommandManager.CustomTypeIdStartOffset)
                         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.TypeIdStartOffset)
+            if ((byte)type < CustomCommandManager.CustomTypeIdStartOffset)
             {
                 value = GarnetObject.Create(type);
                 value.Operate(ref input, ref output.spanByteAndMemory, out _, out _);
@@ -139,7 +139,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.TypeIdStartOffset)
+                    if ((byte)input.header.type < CustomCommandManager.CustomTypeIdStartOffset)
                     {
                         var operateSuccessful = value.Operate(ref input, ref output.spanByteAndMemory, out sizeChange,
                         out var removeKey);
@@ -231,7 +231,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.TypeIdStartOffset)
+                    if ((byte)input.header.type < CustomCommandManager.CustomTypeIdStartOffset)
                     {
                         value.Operate(ref input, ref output.spanByteAndMemory, out _, out var removeKey);
                         if (removeKey)

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.TypeIdStartOffset)
+                        if ((byte)input.header.type < CustomCommandManager.CustomTypeIdStartOffset)
                             return value.Operate(ref input, ref dst.spanByteAndMemory, out _, out _);
 
                         if (IncorrectObjectType(ref input, value, ref dst.spanByteAndMemory))