Implement Timeouts For Lua Scripts

benchmark/BDN.benchmark/Lua/LuaParams.cs

@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-using BenchmarkDotNet.Code;
 using Garnet.server;
 
 namespace BDN.benchmark.Lua
@@ -14,25 +13,38 @@ public readonly struct LuaParams
         public readonly LuaMemoryManagementMode Mode { get; }
         public readonly bool MemoryLimit { get; }
 
+        public readonly TimeSpan? Timeout { get; }
+
         /// <summary>
         /// Constructor
         /// </summary>
-        public LuaParams(LuaMemoryManagementMode mode, bool memoryLimit)
+        public LuaParams(LuaMemoryManagementMode mode, bool memoryLimit, TimeSpan? timeout = null)
         {
             Mode = mode;
             MemoryLimit = memoryLimit;
+            Timeout = timeout;
         }
 
         /// <summary>
         /// Get the equivalent <see cref="LuaOptions"/>.
         /// </summary>
         public LuaOptions CreateOptions()
-        => new(Mode, MemoryLimit ? "2m" : "");
+        => new(Mode, MemoryLimit ? "2m" : "", Timeout ?? System.Threading.Timeout.InfiniteTimeSpan);
 
         /// <summary>
         /// String representation
         /// </summary>
         public override string ToString()
-        => $"{Mode},{(MemoryLimit ? "Limit" : "None")}";
+        {
+            if (Timeout != null)
+            {
+                return $"{Mode},{(MemoryLimit ? "Limit" : "None")},{(Timeout == System.Threading.Timeout.InfiniteTimeSpan ? "-" : Timeout.ToString())}";
+            }
+            else
+            {
+                // Keep old format for benchmarks that don't care about timeouts
+                return $"{Mode},{(MemoryLimit ? "Limit" : "None")}";
+            }
+        }
     }
 }

benchmark/BDN.benchmark/Lua/LuaScriptCacheOperations.cs

@@ -46,7 +46,7 @@ public void GlobalSetup()
 
             server = new EmbeddedRespServer(new GarnetServerOptions() { EnableLua = true, QuietMode = true, LuaOptions = options });
             storeWrapper = server.StoreWrapper;
-            sessionScriptCache = new SessionScriptCache(storeWrapper, new GarnetNoAuthAuthenticator());
+            sessionScriptCache = new SessionScriptCache(storeWrapper, new GarnetNoAuthAuthenticator(), null);
             session = server.GetRespSession();
 
             outerHitDigest = GC.AllocateUninitializedArray<byte>(SessionScriptCache.SHA1Len, pinned: true);

benchmark/BDN.benchmark/Lua/LuaTimeouts.cs

@@ -0,0 +1,89 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+using System.Security.Cryptography;
+using System.Text;
+using BDN.benchmark.Operations;
+using BenchmarkDotNet.Attributes;
+using Embedded.server;
+using Garnet.server;
+
+namespace BDN.benchmark.Lua
+{
+    /// <summary>
+    /// Benchmark the overhead of enabling timeouts on Lua scripting objects.
+    /// </summary>
+    public class LuaTimeouts
+    {
+        private const string Script = @"
+local counter = 7
+for i = 1, 5 do
+    counter = counter + 1
+end
+
+return counter";
+
+        /// <summary>
+        /// Lua parameters
+        /// </summary>
+        [ParamsSource(nameof(LuaParamsProvider))]
+        public LuaParams Params { get; set; }
+
+        internal EmbeddedRespServer server;
+        internal RespServerSession session;
+
+        /// <summary>
+        /// Lua parameters provider
+        /// </summary>
+        public IEnumerable<LuaParams> LuaParamsProvider()
+        => [
+            // We don't expect this to vary by allocator
+            new(LuaMemoryManagementMode.Native, false, Timeout.InfiniteTimeSpan),
+            new(LuaMemoryManagementMode.Native, false, TimeSpan.FromSeconds(1)),
+            new(LuaMemoryManagementMode.Native, false, TimeSpan.FromMinutes(1)),
+        ];
+
+        private Request evalShaRequest;
+
+        [GlobalSetup]
+        public void GlobalSetup()
+        {
+            var opts = new GarnetServerOptions
+            {
+                QuietMode = true,
+                EnableLua = true,
+                LuaOptions = Params.CreateOptions(),
+            };
+
+            server = new EmbeddedRespServer(opts);
+
+            session = server.GetRespSession();
+
+            var scriptLoadStr = $"*3\r\n$6\r\nSCRIPT\r\n$4\r\nLOAD\r\n${Script.Length}\r\n{Script}\r\n";
+            Request scriptLoad = default;
+
+            ScriptOperations.SetupOperation(ref scriptLoad, scriptLoadStr, batchSize: 1);
+            ScriptOperations.Send(session, scriptLoad);
+
+            var scriptHash = string.Join("", SHA1.HashData(Encoding.UTF8.GetBytes(Script)).Select(static x => x.ToString("x2")));
+
+            var evalShaStr = $"*3\r\n$7\r\nEVALSHA\r\n$40\r\n{scriptHash}\r\n$1\r\n0\r\n";
+
+            // Use a batchSize that gets us up around ~100ms
+            ScriptOperations.SetupOperation(ref evalShaRequest, evalShaStr, batchSize: 500 * 1_000);
+        }
+
+        [GlobalCleanup]
+        public void GlobalCleanup()
+        {
+            session.Dispose();
+            server.Dispose();
+        }
+
+        [Benchmark]
+        public void RunScript()
+        {
+            ScriptOperations.Send(session, evalShaRequest);
+        }
+    }
+}

benchmark/BDN.benchmark/Operations/OperationsBase.cs

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

benchmark/BDN.benchmark/Operations/ScriptOperations.cs

@@ -322,19 +322,20 @@ public void ArrayReturn()
         }
 
         private void Send(Request request)
-        {
-            _ = session.TryConsumeMessages(request.bufferPtr, request.buffer.Length);
-        }
+        => Send(session, request);
+
+        internal static void Send(RespServerSession session, Request request)
+        => session.TryConsumeMessages(request.bufferPtr, request.buffer.Length);
 
-        private unsafe void SetupOperation(ref Request request, ReadOnlySpan<byte> operation, int batchSize = batchSize)
+        internal static 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)
+        internal static 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++)
@@ -345,7 +346,7 @@ private unsafe void SetupOperation(ref Request request, string operation, int ba
             request.bufferPtr = (byte*)Unsafe.AsPointer(ref request.buffer[0]);
         }
 
-        private unsafe void SetupOperation(ref Request request, List<byte> operationBytes)
+        internal static unsafe void SetupOperation(ref Request request, List<byte> operationBytes)
         {
             request.buffer = GC.AllocateUninitializedArray<byte>(operationBytes.Count, pinned: true);
             operationBytes.CopyTo(request.buffer);

libs/host/Configuration/Options.cs

@@ -11,6 +11,7 @@
 using System.Reflection;
 using System.Runtime.InteropServices;
 using System.Security.Cryptography.X509Certificates;
+using System.Threading;
 using CommandLine;
 using Garnet.server;
 using Garnet.server.Auth.Aad;
@@ -534,6 +535,10 @@ internal sealed class Options
         [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; }
 
+        [TimeSpanValidation(allowZero: false, allowNegative: false)]
+        [Option("lua-script-timeout", Default = null, Required = false, HelpText = "Timeout for a Lua instance while running a script, specified as a TimeSpan ('c' format, ie. d.hh:mm:ss.ffff)")]
+        public string LuaScriptTimeout { get; set; }
+
         [FilePathValidation(false, true, false)]
         [Option("unixsocket", Required = false, HelpText = "Unix socket address path to bind server to")]
         public string UnixSocketPath { get; set; }
@@ -776,7 +781,7 @@ public GarnetServerOptions GetServerOptions(ILogger logger = null)
                 LoadModuleCS = LoadModuleCS,
                 FailOnRecoveryError = FailOnRecoveryError.GetValueOrDefault(),
                 SkipRDBRestoreChecksumValidation = SkipRDBRestoreChecksumValidation.GetValueOrDefault(),
-                LuaOptions = EnableLua.GetValueOrDefault() ? new LuaOptions(LuaMemoryManagementMode, LuaScriptMemoryLimit, logger) : null,
+                LuaOptions = EnableLua.GetValueOrDefault() ? new LuaOptions(LuaMemoryManagementMode, LuaScriptMemoryLimit, string.IsNullOrEmpty(LuaScriptTimeout) ? Timeout.InfiniteTimeSpan : TimeSpan.ParseExact(LuaScriptTimeout, "c", null), logger) : null,
                 UnixSocketPath = UnixSocketPath,
                 UnixSocketPermission = unixSocketPermissions
             };

libs/host/Configuration/OptionsValidators.cs

@@ -605,7 +605,7 @@ protected override ValidationResult IsValid(object value, ValidationContext vali
 
                     if (forbiddenValues.Contains(otherOptionValueAsString, StringComparer.OrdinalIgnoreCase))
                     {
-                        return new ValidationResult($"{nameof(validationContext.DisplayName)} cannot be set with {otherOptionName} has value '{otherOptionValueAsString}'");
+                        return new ValidationResult($"{validationContext.DisplayName} cannot be set with {otherOptionName} has value '{otherOptionValueAsString}'");
                     }
                 }
             }
@@ -614,6 +614,68 @@ protected override ValidationResult IsValid(object value, ValidationContext vali
         }
     }
 
+    /// <summary>
+    /// Validates TimeSpan options which are typed as strings in the options object.
+    /// 
+    /// This is a little awkward because attributes can't taken TimeSpan objects.
+    /// </summary>
+    [AttributeUsage(AttributeTargets.Property)]
+    internal sealed class TimeSpanValidation : ValidationAttribute
+    {
+        private readonly bool allowZero;
+        private readonly bool allowPositive;
+        private readonly bool allowNegative;
+
+        internal TimeSpanValidation(bool allowZero = true, bool allowPositive = true, bool allowNegative = true)
+        {
+            this.allowNegative = allowNegative;
+            this.allowPositive = allowPositive;
+            this.allowZero = allowZero;
+
+        }
+
+        /// <inheritdoc/>
+        protected override ValidationResult IsValid(object value, ValidationContext validationContext)
+        {
+            if (value is string valueStr)
+            {
+                if (string.IsNullOrEmpty(valueStr))
+                {
+                    return ValidationResult.Success;
+                }
+
+                if (!TimeSpan.TryParseExact(valueStr, "c", null, out var parsed))
+                {
+                    return new ValidationResult($"{validationContext.DisplayName} value of \"{valueStr}\" could not be parsed as a TimeSpan");
+                }
+
+                if (!allowZero && parsed == TimeSpan.Zero)
+                {
+                    return new ValidationResult($"{validationContext.DisplayName} cannot be 0");
+                }
+
+                if (!allowPositive && parsed > TimeSpan.Zero)
+                {
+                    return new ValidationResult($"{validationContext.DisplayName} cannot be positive");
+                }
+
+                if (!allowNegative && parsed < TimeSpan.Zero)
+                {
+                    return new ValidationResult($"{validationContext.DisplayName} cannot be negative");
+                }
+
+                return ValidationResult.Success;
+            }
+
+            if (value is not null)
+            {
+                return new ValidationResult($"{validationContext.DisplayName} set with non-string value \"{value}\"");
+            }
+
+            return ValidationResult.Success;
+        }
+    }
+
     /// <summary>
     /// Forbids a config option from being set if the current OS platform is not supported.
     /// </summary>
@@ -639,7 +701,7 @@ protected override ValidationResult IsValid(object value, ValidationContext vali
                     return ValidationResult.Success;
             }
 
-            return new ValidationResult($"{nameof(validationContext.DisplayName)} can only bet set on following platforms: {string.Join(',', supportedPlatforms)}");
+            return new ValidationResult($"{validationContext.DisplayName} can only bet set on following platforms: {string.Join(',', supportedPlatforms)}");
         }
     }
 }

libs/host/defaults.conf

@@ -359,9 +359,12 @@
 	/* Lua uses the default, unmanaged and untracked, allocator */
 	"LuaMemoryManagementMode": "Native",
 
-	/* Lua limits are ignored for Native, but can be set for other modes */
+	/* Lua memory limits are ignored for Native, but can be set for other modes */
 	"LuaScriptMemoryLimit": null,
 
+	/* Impose no timeout on Lua scripts by default */
+	"LuaScriptTimeout": null,
+
 	/* Unix socket address path to bind the server to */
 	"UnixSocketPath": null,
 

libs/server/Lua/LuaCommands.cs

@@ -30,13 +30,13 @@ private unsafe bool TryEVALSHA()
 
             ref var digest = ref parseState.GetArgSliceByRef(0);
 
+            var convertedToLower = false;
             LuaRunner runner = null;
 
             // Length check is mandatory, as ScriptHashKey assumes correct length
             if (digest.length == SessionScriptCache.SHA1Len)
             {
-                AsciiUtils.ToLowerInPlace(digest.Span);
-
+            tryAgain:
                 var scriptKey = new ScriptHashKey(digest.Span);
 
                 if (!sessionScriptCache.TryGetFromDigest(scriptKey, out runner))
@@ -51,6 +51,17 @@ private unsafe bool TryEVALSHA()
                             return true;
                         }
                     }
+                    else if (!convertedToLower)
+                    {
+                        // On a miss (which should be rare) make sure the hash is lower case and try again.
+                        //
+                        // We assume that hashes will be sent in the same format as we return them (lower)
+                        // most of the time, so optimize for that.
+
+                        AsciiUtils.ToLowerInPlace(digest.Span);
+                        convertedToLower = true;
+                        goto tryAgain;
+                    }
                 }
             }
 
@@ -61,7 +72,10 @@ private unsafe bool TryEVALSHA()
             }
             else
             {
+                // We assume here that ExecuteScript does not raise exceptions
+                sessionScriptCache.StartRunningScript(runner);
                 ExecuteScript(count - 1, runner);
+                sessionScriptCache.StopRunningScript();
             }
 
             return true;
@@ -104,7 +118,10 @@ private unsafe bool TryEVAL()
             }
             else
             {
+                // We assume here that ExecuteScript does not raise exceptions
+                sessionScriptCache.StartRunningScript(runner);
                 ExecuteScript(count - 1, runner);
+                sessionScriptCache.StopRunningScript();
             }
 
             return true;
@@ -168,7 +185,7 @@ private bool NetworkScriptFlush()
             }
             else if (parseState.Count == 1)
             {
-                // we ignore this, but should validate it
+                // We ignore this, but should validate it
                 ref var arg = ref parseState.GetArgSliceByRef(0);
 
                 AsciiUtils.ToUpperInPlace(arg.Span);