Per Script Invocation Lua Memory Limits

.github/workflows/ci-bdnbenchmark.yml

@@ -42,7 +42,7 @@ jobs:
         os: [ ubuntu-latest, windows-latest ]
         framework: [ 'net8.0' ]
         configuration: [ 'Release' ]
-        test: [ 'Operations.BasicOperations', 'Operations.ObjectOperations', 'Operations.HashObjectOperations', 'Cluster.ClusterMigrate', 'Cluster.ClusterOperations', 'Lua.LuaScripts', 'Operations.CustomOperations', 'Operations.RawStringOperations', 'Operations.ScriptOperations','Network.BasicOperations', 'Network.RawStringOperations' ]
+        test: [ 'Operations.BasicOperations', 'Operations.ObjectOperations', 'Operations.HashObjectOperations', 'Cluster.ClusterMigrate', 'Cluster.ClusterOperations', 'Lua.LuaScripts', 'Lua.LuaScriptCacheOperations','Lua.LuaRunnerOperations','Operations.CustomOperations', 'Operations.RawStringOperations', 'Operations.ScriptOperations','Network.BasicOperations', 'Network.RawStringOperations' ]
     steps:
       - name: Check out code
         uses: actions/checkout@v4

benchmark/BDN.benchmark/Lua/LuaParams.cs

@@ -1,26 +1,38 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
+using BenchmarkDotNet.Code;
+using Garnet.server;
+
 namespace BDN.benchmark.Lua
 {
     /// <summary>
-    /// Cluster parameters
+    /// Lua parameters
     /// </summary>
-    public struct LuaParams
+    public readonly struct LuaParams
     {
+        public readonly LuaMemoryManagementMode Mode { get; }
+        public readonly bool MemoryLimit { get; }
+
         /// <summary>
         /// Constructor
         /// </summary>
-        public LuaParams()
+        public LuaParams(LuaMemoryManagementMode mode, bool memoryLimit)
         {
+            Mode = mode;
+            MemoryLimit = memoryLimit;
         }
 
+        /// <summary>
+        /// Get the equivalent <see cref="LuaOptions"/>.
+        /// </summary>
+        public LuaOptions CreateOptions()
+        => new(Mode, MemoryLimit ? "2m" : "");
+
         /// <summary>
         /// String representation
         /// </summary>
         public override string ToString()
-        {
-            return "None";
-        }
+        => $"{Mode},{(MemoryLimit ? "Limit" : "None")}";
     }
 }

benchmark/BDN.benchmark/Lua/LuaRunnerOperations.cs

@@ -139,9 +139,13 @@ public unsafe class LuaRunnerOperations
         /// Lua parameters provider
         /// </summary>
         public IEnumerable<LuaParams> LuaParamsProvider()
-        {
-            yield return new();
-        }
+        => [
+            new(LuaMemoryManagementMode.Native, false),
+            new(LuaMemoryManagementMode.Tracked, false),
+            new(LuaMemoryManagementMode.Tracked, true),
+            new(LuaMemoryManagementMode.Managed, false),
+            new(LuaMemoryManagementMode.Managed, true),
+        ];
 
         private EmbeddedRespServer server;
         private RespServerSession session;
@@ -151,16 +155,20 @@ public IEnumerable<LuaParams> LuaParamsProvider()
         private LuaRunner smallCompileRunner;
         private LuaRunner largeCompileRunner;
 
+        private LuaOptions opts;
+
         [GlobalSetup]
         public void GlobalSetup()
         {
-            server = new EmbeddedRespServer(new GarnetServerOptions() { EnableLua = true, QuietMode = true });
+            opts = Params.CreateOptions();
+
+            server = new EmbeddedRespServer(new GarnetServerOptions() { EnableLua = true, QuietMode = true, LuaOptions = opts });
 
             session = server.GetRespSession();
-            paramsRunner = new LuaRunner("return nil");
+            paramsRunner = new LuaRunner(opts, "return nil");
 
-            smallCompileRunner = new LuaRunner(SmallScript);
-            largeCompileRunner = new LuaRunner(LargeScript);
+            smallCompileRunner = new LuaRunner(opts, SmallScript);
+            largeCompileRunner = new LuaRunner(opts, LargeScript);
         }
 
         [GlobalCleanup]
@@ -194,13 +202,13 @@ public void ResetParametersLarge()
         [Benchmark]
         public void ConstructSmall()
         {
-            using var runner = new LuaRunner(SmallScript);
+            using var runner = new LuaRunner(opts, SmallScript);
         }
 
         [Benchmark]
         public void ConstructLarge()
         {
-            using var runner = new LuaRunner(LargeScript);
+            using var runner = new LuaRunner(opts, LargeScript);
         }
 
         [Benchmark]

benchmark/BDN.benchmark/Lua/LuaScriptCacheOperations.cs

@@ -22,9 +22,13 @@ public class LuaScriptCacheOperations
         /// Lua parameters provider
         /// </summary>
         public IEnumerable<LuaParams> LuaParamsProvider()
-        {
-            yield return new();
-        }
+        => [
+            new(LuaMemoryManagementMode.Native, false),
+            new(LuaMemoryManagementMode.Tracked, false),
+            new(LuaMemoryManagementMode.Tracked, true),
+            new(LuaMemoryManagementMode.Managed, false),
+            new(LuaMemoryManagementMode.Managed, true),
+        ];
 
         private EmbeddedRespServer server;
         private StoreWrapper storeWrapper;
@@ -38,7 +42,9 @@ public IEnumerable<LuaParams> LuaParamsProvider()
         [GlobalSetup]
         public void GlobalSetup()
         {
-            server = new EmbeddedRespServer(new GarnetServerOptions() { EnableLua = true, QuietMode = true });
+            var options = Params.CreateOptions();
+
+            server = new EmbeddedRespServer(new GarnetServerOptions() { EnableLua = true, QuietMode = true, LuaOptions = options });
             storeWrapper = server.StoreWrapper;
             sessionScriptCache = new SessionScriptCache(storeWrapper, new GarnetNoAuthAuthenticator());
             session = server.GetRespSession();

benchmark/BDN.benchmark/Lua/LuaScripts.cs

@@ -22,23 +22,29 @@ public unsafe class LuaScripts
         /// Lua parameters provider
         /// </summary>
         public IEnumerable<LuaParams> LuaParamsProvider()
-        {
-            yield return new();
-        }
+        => [
+            new(LuaMemoryManagementMode.Native, false),
+            new(LuaMemoryManagementMode.Tracked, false),
+            new(LuaMemoryManagementMode.Tracked, true),
+            new(LuaMemoryManagementMode.Managed, false),
+            new(LuaMemoryManagementMode.Managed, true),
+        ];
 
         LuaRunner r1, r2, r3, r4;
         readonly string[] keys = ["key1"];
 
         [GlobalSetup]
         public void GlobalSetup()
         {
-            r1 = new LuaRunner("return");
+            var options = Params.CreateOptions();
+
+            r1 = new LuaRunner(options, "return");
             r1.CompileForRunner();
-            r2 = new LuaRunner("return 1 + 1");
+            r2 = new LuaRunner(options, "return 1 + 1");
             r2.CompileForRunner();
-            r3 = new LuaRunner("return KEYS[1]");
+            r3 = new LuaRunner(options, "return KEYS[1]");
             r3.CompileForRunner();
-            r4 = new LuaRunner("return redis.call(KEYS[1])");
+            r4 = new LuaRunner(options, "return redis.call(KEYS[1])");
             r4.CompileForRunner();
         }
 

benchmark/BDN.benchmark/Operations/OperationsBase.cs

@@ -55,6 +55,7 @@ public virtual void GlobalSetup()
                 QuietMode = true,
                 EnableLua = true,
                 DisablePubSub = true,
+                LuaOptions = new(LuaMemoryManagementMode.Native, ""),
             };
             if (Params.useAof)
             {

benchmark/BDN.benchmark/Operations/ScriptOperations.cs

@@ -4,16 +4,18 @@
 using System.Runtime.CompilerServices;
 using System.Security.Cryptography;
 using System.Text;
+using BDN.benchmark.Lua;
 using BenchmarkDotNet.Attributes;
 using Embedded.server;
+using Garnet.server;
 
 namespace BDN.benchmark.Operations
 {
     /// <summary>
     /// Benchmark for SCRIPT LOAD, SCRIPT EXISTS, EVAL, and EVALSHA
     /// </summary>
     [MemoryDiagnoser]
-    public unsafe class ScriptOperations : OperationsBase
+    public unsafe class ScriptOperations
     {
         // Small script that does 1 operation and no logic
         const string SmallScriptText = @"return redis.call('GET', KEYS[1]);";
@@ -156,9 +158,54 @@ public unsafe class ScriptOperations : OperationsBase
         static ReadOnlySpan<byte> ARRAY_RETURN => "*3\r\n$4\r\nEVAL\r\n$22\r\nreturn {1, 2, 3, 4, 5}\r\n$1\r\n0\r\n"u8;
         Request arrayReturn;
 
-        public override void GlobalSetup()
+
+        /// <summary>
+        /// Lua parameters
+        /// </summary>
+        [ParamsSource(nameof(LuaParamsProvider))]
+        public LuaParams Params { get; set; }
+
+        /// <summary>
+        /// Lua parameters provider
+        /// </summary>
+        public static IEnumerable<LuaParams> LuaParamsProvider()
+        => [
+            new(LuaMemoryManagementMode.Native, false),
+            new(LuaMemoryManagementMode.Tracked, false),
+            new(LuaMemoryManagementMode.Tracked, true),
+            new(LuaMemoryManagementMode.Managed, false),
+            new(LuaMemoryManagementMode.Managed, true),
+        ];
+
+        /// <summary>
+        /// Batch size per method invocation
+        /// With a batchSize of 100, we have a convenient conversion of latency to throughput:
+        ///   5 us = 20 Mops/sec
+        ///  10 us = 10 Mops/sec
+        ///  20 us =  5 Mops/sec
+        ///  25 us =  4 Mops/sec
+        /// 100 us =  1 Mops/sec
+        /// </summary>
+        internal const int batchSize = 100;
+        internal EmbeddedRespServer server;
+        internal RespServerSession session;
+
+        /// <summary>
+        /// Setup
+        /// </summary>
+        [GlobalSetup]
+        public void GlobalSetup()
         {
-            base.GlobalSetup();
+            var opts = new GarnetServerOptions
+            {
+                QuietMode = true,
+                EnableLua = true,
+                LuaOptions = Params.CreateOptions(),
+            };
+
+            server = new EmbeddedRespServer(opts);
+
+            session = server.GetRespSession();
 
             SetupOperation(ref scriptLoad, SCRIPT_LOAD);
 
@@ -216,6 +263,16 @@ public override void GlobalSetup()
             SetupOperation(ref evalShaLargeScript, largeScriptEvals);
         }
 
+        /// <summary>
+        /// Cleanup
+        /// </summary>
+        [GlobalCleanup]
+        public virtual void GlobalCleanup()
+        {
+            session.Dispose();
+            server.Dispose();
+        }
+
         [Benchmark]
         public void ScriptLoad()
         {
@@ -263,5 +320,36 @@ public void ArrayReturn()
         {
             Send(arrayReturn);
         }
+
+        private void Send(Request request)
+        {
+            _ = session.TryConsumeMessages(request.bufferPtr, request.buffer.Length);
+        }
+
+        private unsafe void SetupOperation(ref Request request, ReadOnlySpan<byte> operation, int batchSize = batchSize)
+        {
+            request.buffer = GC.AllocateArray<byte>(operation.Length * batchSize, pinned: true);
+            request.bufferPtr = (byte*)Unsafe.AsPointer(ref request.buffer[0]);
+            for (int i = 0; i < batchSize; i++)
+                operation.CopyTo(new Span<byte>(request.buffer).Slice(i * operation.Length));
+        }
+
+        private unsafe void SetupOperation(ref Request request, string operation, int batchSize = batchSize)
+        {
+            request.buffer = GC.AllocateUninitializedArray<byte>(operation.Length * batchSize, pinned: true);
+            for (var i = 0; i < batchSize; i++)
+            {
+                var start = i * operation.Length;
+                Encoding.UTF8.GetBytes(operation, request.buffer.AsSpan().Slice(start, operation.Length));
+            }
+            request.bufferPtr = (byte*)Unsafe.AsPointer(ref request.buffer[0]);
+        }
+
+        private unsafe void SetupOperation(ref Request request, List<byte> operationBytes)
+        {
+            request.buffer = GC.AllocateUninitializedArray<byte>(operationBytes.Count, pinned: true);
+            operationBytes.CopyTo(request.buffer);
+            request.bufferPtr = (byte*)Unsafe.AsPointer(ref request.buffer[0]);
+        }
     }
 }

libs/host/Configuration/Options.cs

@@ -8,6 +8,7 @@
 using System.Linq;
 using System.Net;
 using System.Reflection;
+using System.Runtime.InteropServices;
 using System.Security.Cryptography.X509Certificates;
 using CommandLine;
 using Garnet.server;
@@ -508,6 +509,14 @@ internal sealed class Options
         [Option("fail-on-recovery-error", Default = false, Required = false, HelpText = "Server bootup should fail if errors happen during bootup of AOF and checkpointing")]
         public bool? FailOnRecoveryError { get; set; }
 
+        [Option("lua-memory-management-mode", Default = LuaMemoryManagementMode.Native, Required = false, HelpText = "Memory management mode for Lua scripts, must be set to LimittedNative or Managed to impose script limits")]
+        public LuaMemoryManagementMode LuaMemoryManagementMode { get; set; }
+
+        [MemorySizeValidation(false)]
+        [ForbiddenWithOption(nameof(LuaMemoryManagementMode), nameof(LuaMemoryManagementMode.Native))]
+        [Option("lua-script-memory-limit", Default = null, HelpText = "Memory limit for a Lua instances while running a script, lua-memory-management-mode must be set to something other than Native to use this flag")]
+        public string LuaScriptMemoryLimit { get; set; }
+
         /// <summary>
         /// This property contains all arguments that were not parsed by the command line argument parser
         /// </summary>
@@ -718,7 +727,8 @@ public GarnetServerOptions GetServerOptions(ILogger logger = null)
                 IndexResizeFrequencySecs = IndexResizeFrequencySecs,
                 IndexResizeThreshold = IndexResizeThreshold,
                 LoadModuleCS = LoadModuleCS,
-                FailOnRecoveryError = FailOnRecoveryError.GetValueOrDefault()
+                FailOnRecoveryError = FailOnRecoveryError.GetValueOrDefault(),
+                LuaOptions = EnableLua.GetValueOrDefault() ? new LuaOptions(LuaMemoryManagementMode, LuaScriptMemoryLimit, logger) : null,
             };
         }
 

libs/host/Configuration/OptionsValidators.cs

@@ -563,4 +563,42 @@ protected override ValidationResult IsValid(object value, ValidationContext vali
             return base.IsValid(value, validationContext);
         }
     }
+
+    /// <summary>
+    /// Forbids a config option from being set if the another option has particular values.
+    /// </summary>
+    [AttributeUsage(AttributeTargets.Property)]
+    internal sealed class ForbiddenWithOptionAttribute : ValidationAttribute
+    {
+        private readonly string otherOptionName;
+        private readonly string[] forbiddenValues;
+
+        internal ForbiddenWithOptionAttribute(string otherOptionName, string forbiddenValue, params string[] otherForbiddenValues)
+        {
+            this.otherOptionName = otherOptionName;
+            forbiddenValues = [forbiddenValue, .. otherForbiddenValues];
+        }
+
+        /// <inheritdoc/>
+        protected override ValidationResult IsValid(object value, ValidationContext validationContext)
+        {
+            var optionIsSet = value != null && !(value is string valueStr && string.IsNullOrEmpty(valueStr));
+            if (optionIsSet)
+            {
+                var propAccessor = validationContext.ObjectInstance?.GetType()?.GetProperty(otherOptionName, BindingFlags.Instance | BindingFlags.Public);
+                if (propAccessor != null)
+                {
+                    var otherOptionValue = propAccessor.GetValue(validationContext.ObjectInstance);
+                    var otherOptionValueAsString = otherOptionValue is string strVal ? strVal : otherOptionValue?.ToString();
+
+                    if (forbiddenValues.Contains(otherOptionValueAsString, StringComparer.OrdinalIgnoreCase))
+                    {
+                        return new ValidationResult($"{nameof(validationContext.DisplayName)} cannot be set with {otherOptionName} has value '{otherOptionValueAsString}'");
+                    }
+                }
+            }
+
+            return ValidationResult.Success;
+        }
+    }
 }