Lua Scripting Allocation, Performance, and Correctness Improvements

Directory.Packages.props

@@ -8,7 +8,7 @@
     <PackageVersion Include="BenchmarkDotNet.Diagnostics.Windows" Version="0.13.12" />
     <PackageVersion Include="CommandLineParser" Version="2.9.1" />
     <PackageVersion Include="JsonPath.Net" Version="1.1.6" />
-    <PackageVersion Include="NLua" Version="1.7.3" />
+    <PackageVersion Include="KeraLua" Version="1.4.1" />
     <PackageVersion Include="NUnit" Version="4.1.0" />
     <PackageVersion Include="NUnit3TestAdapter" Version="4.6.0" />
     <PackageVersion Include="Microsoft.CodeAnalysis" Version="4.11.0" />

benchmark/BDN.benchmark/BDN.benchmark.csproj

@@ -10,7 +10,7 @@
 
   <ItemGroup>
     <PackageReference Include="BenchmarkDotNet" />
-		<PackageReference Include="NLua" />
+		<PackageReference Include="KeraLua" />
 	</ItemGroup>
 
   <ItemGroup>

benchmark/BDN.benchmark/Lua/LuaRunnerOperations.cs

@@ -0,0 +1,220 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+using BenchmarkDotNet.Attributes;
+using Embedded.perftest;
+using Garnet.server;
+
+namespace BDN.benchmark.Lua
+{
+    /// <summary>
+    /// Benchmark for non-script running operations in LuaRunner
+    /// </summary>
+    [MemoryDiagnoser]
+    public unsafe class LuaRunnerOperations
+    {
+        private const string SmallScript = "return nil";
+
+        private const string LargeScript = @"
+-- based on a real UDF, with some extraneous ops removed
+
+local userKey = KEYS[1]
+local identifier = ARGV[1]
+local currentTime = ARGV[2]
+local newSession = ARGV[3]
+
+local userKeyValue = redis.call(""GET"", userKey)
+
+local updatedValue = nil
+local returnValue = -1
+
+if userKeyValue then
+-- needs to be updated
+
+	local oldestEntry = nil
+	local oldestEntryUpdateTime = nil
+
+	local match = nil
+
+	local entryCount = 0
+
+-- loop over each entry, looking for one to update
+	for entry in string.gmatch(userKeyValue, ""([^%|]+)"") do
+		entryCount = entryCount + 1
+
+		local entryIdentifier = nil
+		local entrySessionNumber = -1
+		local entryRequestCount = -1
+		local entryLastSessionUpdateTime = -1
+
+		local ix = 0
+		for part in string.gmatch(entry, ""([^:]+)"") do 
+			if ix == 0 then
+				entryIdentifier = part
+			elseif ix == 1 then
+				entrySessionNumber = tonumber(part)
+			elseif ix == 2 then
+				entryRequestCount = tonumber(part)
+			elseif ix == 3 then
+				entryLastSessionUpdateTime = tonumber(part)
+			else
+-- malformed, too many parts
+				return -1
+			end
+
+			ix = ix + 1
+		end
+
+		if ix ~= 4 then
+-- malformed, too few parts
+			return -2
+		end
+
+		if entryIdentifier == identifier then
+-- found the one to update
+			local updatedEntry = nil
+
+			if tonumber(newSession) == 1 then
+				local updatedSessionNumber = entrySessionNumber + 1
+				updatedEntry = entryIdentifier .. "":"" .. tostring(updatedSessionNumber) .. "":1:"" .. tostring(currentTime)
+				returnValue = 3
+			else
+				local updatedRequestCount = entryRequestCount + 1
+				updatedEntry = entryIdentifier .. "":"" .. tostring(entrySessionNumber) .. "":"" .. tostring(updatedRequestCount) .. "":"" .. tostring(currentTime)
+				returnValue = 2
+			end
+
+-- have to escape the replacement, since Lua doesn't have a literal replace :/
+			local escapedEntry = string.gsub(entry, ""%p"", ""%%%1"")
+			updatedValue = string.gsub(userKeyValue, escapedEntry, updatedEntry)
+
+			break
+		end
+
+		if oldestEntryUpdateTime == nil or oldestEntryUpdateTime > entryLastSessionUpdateTime then
+-- remember the oldest entry, so we can replace it if needed
+			oldestEntry = entry
+			oldestEntryUpdateTime = entryLastSessionUpdateTime
+		end
+	end
+
+	if updatedValue == nil then
+-- we didn't update an existing value, so we need to add it
+
+		local newEntry = identifier .. "":1:1:"" .. tostring(currentTime)
+
+		if entryCount < 20 then
+-- there's room, just append it
+			updatedValue = userKeyValue .. ""|"" .. newEntry
+			returnValue = 4
+		else
+-- there isn't room, replace the LRU entry
+
+-- have to escape the replacement, since Lua doesn't have a literal replace :/
+			local escapedOldestEntry = string.gsub(oldestEntry, ""%p"", ""%%%1"")
+
+			updatedValue = string.gsub(userKeyValue, escapedOldestEntry, newEntry)
+			returnValue = 5
+		end
+	end
+else
+-- needs to be created
+	updatedValue = identifier .. "":1:1:"" .. tostring(currentTime)
+
+	returnValue = 1
+end
+
+redis.call(""SET"", userKey, updatedValue)
+
+return returnValue
+";
+
+        /// <summary>
+        /// Lua parameters
+        /// </summary>
+        [ParamsSource(nameof(LuaParamsProvider))]
+        public LuaParams Params { get; set; }
+
+        /// <summary>
+        /// Lua parameters provider
+        /// </summary>
+        public IEnumerable<LuaParams> LuaParamsProvider()
+        {
+            yield return new();
+        }
+
+        private EmbeddedRespServer server;
+        private RespServerSession session;
+
+        private LuaRunner paramsRunner;
+
+        private LuaRunner smallCompileRunner;
+        private LuaRunner largeCompileRunner;
+
+        [GlobalSetup]
+        public void GlobalSetup()
+        {
+            server = new EmbeddedRespServer(new GarnetServerOptions() { EnableLua = true, QuietMode = true });
+
+            session = server.GetRespSession();
+            paramsRunner = new LuaRunner("return nil");
+
+            smallCompileRunner = new LuaRunner(SmallScript);
+            largeCompileRunner = new LuaRunner(LargeScript);
+        }
+
+        [GlobalCleanup]
+        public void GlobalCleanup()
+        {
+            session.Dispose();
+            server.Dispose();
+            paramsRunner.Dispose();
+        }
+
+        [Benchmark]
+        public void ResetParametersSmall()
+        {
+            // First force up
+            paramsRunner.ResetParameters(1, 1);
+
+            // Then require a small amount of clearing (1 key, 1 arg)
+            paramsRunner.ResetParameters(0, 0);
+        }
+
+        [Benchmark]
+        public void ResetParametersLarge()
+        {
+            // First force up
+            paramsRunner.ResetParameters(10, 10);
+
+            // Then require a large amount of clearing (10 keys, 10 args)
+            paramsRunner.ResetParameters(0, 0);
+        }
+
+        [Benchmark]
+        public void ConstructSmall()
+        {
+            using var runner = new LuaRunner(SmallScript);
+        }
+
+        [Benchmark]
+        public void ConstructLarge()
+        {
+            using var runner = new LuaRunner(LargeScript);
+        }
+
+        [Benchmark]
+        public void CompileForSessionSmall()
+        {
+            smallCompileRunner.ResetCompilation();
+            smallCompileRunner.CompileForSession(session);
+        }
+
+        [Benchmark]
+        public void CompileForSessionLarge()
+        {
+            largeCompileRunner.ResetCompilation();
+            largeCompileRunner.CompileForSession(session);
+        }
+    }
+}

benchmark/BDN.benchmark/Lua/LuaScriptCacheOperations.cs

@@ -0,0 +1,153 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+using BenchmarkDotNet.Attributes;
+using Embedded.perftest;
+using Garnet.common;
+using Garnet.server;
+using Garnet.server.Auth;
+
+namespace BDN.benchmark.Lua
+{
+    [MemoryDiagnoser]
+    public class LuaScriptCacheOperations
+    {
+        /// <summary>
+        /// Lua parameters
+        /// </summary>
+        [ParamsSource(nameof(LuaParamsProvider))]
+        public LuaParams Params { get; set; }
+
+        /// <summary>
+        /// Lua parameters provider
+        /// </summary>
+        public IEnumerable<LuaParams> LuaParamsProvider()
+        {
+            yield return new();
+        }
+
+        private EmbeddedRespServer server;
+        private StoreWrapper storeWrapper;
+        private SessionScriptCache sessionScriptCache;
+        private RespServerSession session;
+
+        private byte[] outerHitDigest;
+        private byte[] innerHitDigest;
+        private byte[] missDigest;
+
+        [GlobalSetup]
+        public void GlobalSetup()
+        {
+            server = new EmbeddedRespServer(new GarnetServerOptions() { EnableLua = true, QuietMode = true });
+            storeWrapper = server.StoreWrapper;
+            sessionScriptCache = new SessionScriptCache(storeWrapper, new GarnetNoAuthAuthenticator());
+            session = server.GetRespSession();
+
+            outerHitDigest = GC.AllocateUninitializedArray<byte>(SessionScriptCache.SHA1Len, pinned: true);
+            sessionScriptCache.GetScriptDigest("return 1"u8, outerHitDigest);
+            if (!storeWrapper.storeScriptCache.TryAdd(new(outerHitDigest), "return 1"u8.ToArray()))
+            {
+                throw new InvalidOperationException("Should have been able to load into global cache");
+            }
+
+            innerHitDigest = GC.AllocateUninitializedArray<byte>(SessionScriptCache.SHA1Len, pinned: true);
+            sessionScriptCache.GetScriptDigest("return 1 + 1"u8, innerHitDigest);
+            if (!storeWrapper.storeScriptCache.TryAdd(new(innerHitDigest), "return 1 + 1"u8.ToArray()))
+            {
+                throw new InvalidOperationException("Should have been able to load into global cache");
+            }
+
+            missDigest = GC.AllocateUninitializedArray<byte>(SessionScriptCache.SHA1Len, pinned: true);
+            sessionScriptCache.GetScriptDigest("foobar"u8, missDigest);
+        }
+
+        [GlobalCleanup]
+        public void GlobalCleanup()
+        {
+            session?.Dispose();
+            server?.Dispose();
+        }
+
+        [IterationSetup]
+        public void IterationSetup()
+        {
+            // Force lookup to do work
+            sessionScriptCache.Clear();
+
+            // Make outer hit available for every iteration
+            if (!sessionScriptCache.TryLoad(session, "return 1"u8, new(outerHitDigest), out _, out _, out var error))
+            {
+                throw new InvalidOperationException($"Should have been able to load: {error}");
+            }
+        }
+
+        [Benchmark]
+        public void LookupHit()
+        {
+            _ = sessionScriptCache.TryGetFromDigest(new(outerHitDigest), out _);
+        }
+
+        [Benchmark]
+        public void LookupMiss()
+        {
+            _ = sessionScriptCache.TryGetFromDigest(new(missDigest), out _);
+        }
+
+        [Benchmark]
+        public void LoadOuterHit()
+        {
+            // First if returns true
+            //
+            // This is the common case
+            LoadScript(outerHitDigest);
+        }
+
+        [Benchmark]
+        public void LoadInnerHit()
+        {
+            // First if returns false, second if returns true
+            //
+            // This is expected, but rare
+            LoadScript(innerHitDigest);
+        }
+
+        [Benchmark]
+        public void LoadMiss()
+        {
+            // First if returns false, second if returns false
+            //
+            // This is extremely unlikely, basically implies an error on the client
+            LoadScript(missDigest);
+        }
+
+        [Benchmark]
+        public void Digest()
+        {
+            Span<byte> digest = stackalloc byte[SessionScriptCache.SHA1Len];
+            sessionScriptCache.GetScriptDigest("return 1 + redis.call('GET', KEYS[1])"u8, digest);
+        }
+
+        /// <summary>
+        /// The moral equivalent to our cache load operation.
+        /// </summary>
+        private void LoadScript(Span<byte> digest)
+        {
+            AsciiUtils.ToLowerInPlace(digest);
+
+            var digestKey = new ScriptHashKey(digest);
+
+            if (!sessionScriptCache.TryGetFromDigest(digestKey, out var runner))
+            {
+                if (storeWrapper.storeScriptCache.TryGetValue(digestKey, out var source))
+                {
+                    if (!sessionScriptCache.TryLoad(session, source, digestKey, out runner, out _, out var error))
+                    {
+                        // TryLoad will have written an error out, it any
+
+                        _ = storeWrapper.storeScriptCache.TryRemove(digestKey, out _);
+                    }
+                }
+            }
+        }
+    }
+}