Lua Scripting Allocation, Performance, and Correctness Improvements (#…

…882)

* NLua -> KeraLua; literally nothing compiles

* blind idiot translation into KeraLua; nothing works

* very basic functionality restored

* more complicated functionality sketched out; far from ideal, will need to upstream some KeraLua additions; also horribly broken

* handle errors during script run (compilation is still unguarded, as we can't really do much there)

* faster (and easier to understand, frankly) way to set and clear KEYS and ARGV at start

* a little more functional; moving OK/ERR checking into the loader script may not be viable... but was worth experimenting with

* redis.call should not accept table arguments, only strings and numbers; fix some ZADD scripting tests; special err response stuff can't be done on the Lua side, remove it

* scripting appears to be at parity with NLua implementation

* fix error propogation; implement needed functions for direct script running; all tests passing again

* start on removing allocations and fixing encodings

* remove more allocs

* add a test (and fixes) confirming that redis.call errors match Redis behavior

* add a test (and fixes) confirming that redis.call errors match Redis behavior

* remove more allocations

* add a test for weird binary values, note this fails in main today

* knock out some alloc todos; do some cleanup

* remove some more allocations

* response processesing is now allocation free

* kill a hand full of additional allocations

* DRY up some repeated checks; assert assumptions about Lua stack in DEBUG builds

* adjust some names, these keep confusing me; preparing for writing response directly to network stream

* first whack at getting script results directly into network stream; lots broken right now

* first pass of tables directly into network stream

* complex data types now written out directly

* Runner (ie. mapping Resp <-> .NET) restored to functionality; bits of cleanup

* knock out some todo; pulling Lua constants in CmdStrings like everything else; avoid copying regular used strings into Lua each time they're needed

* Benchmark depends on missing sessions always causing nil responses, which is odd, but easy enough to restore

* remove one bespoke array rental

* Remvoe another bespoke array rental

* Yet another bespoke array rental

* Cleanup

* more benchmarks for LuaRunner; test that parameter resets work

* most overhead is in pinvoke, so start moving some stuff over

* most overhead is in pinvoke, so start moving some stuff over

* where we've already proven a type is a number or string, skip a pinvoke to double check the type

* avoid checkstack calls by tracking stack capacity on our side

* remove more allocs

* remove more allocs

* add a benchmark for script operations

* script lookup is in the hot path, so optimize the key type we're using a bit

* expand ScriptOperations benchmark to actually invoke some functions and do some logic

* huge cleanup; move all the Lua interop into a dedicated class, normalize stack checking, normalize assertions

* switch to LibraryImport since we're on .NET 8

* do a big audit of Lua invokes and mark where GC transition can be suppressed - it's unclear if .NET is actually doing that today, but it's safe

* add a benchmark that returns an array, as there's an outstanding TODO to look at removing some p/invokes

* nope

* add a test for metatable behavior matching Redis; switch to Raw operations where now allowed; more closely sync methods exposed in Lua to those provided by Redis

* todone

* formatting

* address feedback; spelling nits

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 _);
+                    }
+                }
+            }
+        }
+    }
+}