Per Script Invocation Lua Memory Limits (#903)

* Lua allocations go through .NET; really crummy allocator on the POH just to prove it's possible; squashing to clean up _a lot_ of experimentation commits

* punch a LuaOptions into settings

* wire up the Lua options

* prep for other allocators

* Implement more allocators and expand testing

* Knock out a number of todos, consider allocator in benchmarks, additional validation

* add a test for OOMs

* formatting

* convert ScriptOperations to explore different allocators

* cleanup Lua error messgages; this revealed a bug in buffer management for LuaScripts, fixes that

* Make managed allocator less naive, and benchmark on par.
It will probably be unusual to use this allocator, but it shouldn't be _bad_ either.

* formatting

* address nit

* address feedback; only copy relevant bits, not whole buffer

* set lua options in OperationsBase, fixing benchmarks

* BDN Updates:
1) Added a check for NA in results which is an indication that the BDN test failed at run time
2) Added 'Lua.LuaScriptCacheOperations','Lua.LuaRunnerOperations' to BDN Github Action
3) Updated Expected values for the new Lua BDN tests

* Updated some of the LuaScriptCacheOperations expected values

* [Compatibility] Added LCS command (#843)

* Added LCS command

* Format fix

* Reverted CommandDocsUpdater.cs

* Fix cluster test

* Fixed wrong change

* Moved to constant

* Review command fixes

* Fixed review comment

* Fixed test issue

---------

Co-authored-by: Vasileios Zois <[email protected]>
Co-authored-by: Tal Zaccai <[email protected]>

* Configure min and max IO completion threads (#904)

* Configure min and max IO completion threads separately from min and max threads (in the ThreadPool). This is needed as some scenarios may limit number of thread pool threads but require a larger number of IO completion threads.

* nit

* address nit

* address feedback; only copy relevant bits, not whole buffer

* set lua options in OperationsBase, fixing benchmarks

* Yak shave a bunch of cleanup, looking to reduce amount of work done generally in Lua to get some perf back.

* formatting

* fix LuaRunnerOperations ; session wouldn't always be initialized before, which would cause SendAndReset() to fail thinking the message was too large

* Updated Expected values for Allocated in the BDN perf tests to match current Lua changes.

* Missed one expected value. Fixed that so should be ok.

* Seen some instances where allocated bytes are 1024 on one run and then run it a second time and see 1312 without any code changes. Just a small variance between runs so decided to just make sure doesn't go above 1312 for any that were set to 1024.  The charts and other things will show the small nuances if needed.

* Bumped expected values up a bit to handle variance

---------

Co-authored-by: darrenge <[email protected]>
Co-authored-by: Vijay Nirmal <[email protected]>
Co-authored-by: Vasileios Zois <[email protected]>
Co-authored-by: Tal Zaccai <[email protected]>
Co-authored-by: Badrish Chandramouli <[email protected]>

.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,21 @@ 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");
 
-            smallCompileRunner = new LuaRunner(SmallScript);
-            largeCompileRunner = new LuaRunner(LargeScript);
+            paramsRunner = new LuaRunner(opts, "return nil");
+
+            smallCompileRunner = new LuaRunner(opts, SmallScript);
+            largeCompileRunner = new LuaRunner(opts, LargeScript);
         }
 
         [GlobalCleanup]
@@ -171,6 +180,18 @@ public void GlobalCleanup()
             paramsRunner.Dispose();
         }
 
+        [IterationSetup]
+        public void IterationSetup()
+        {
+            session.EnterAndGetResponseObject();
+        }
+
+        [IterationCleanup]
+        public void IterationCleanup()
+        {
+            session.ExitAndReturnResponseObject();
+        }
+
         [Benchmark]
         public void ResetParametersSmall()
         {
@@ -194,13 +215,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,
             };
         }