optimize CommandMap by pre-generating all the RESP chunks

Author: mgravell

src/RESPite/PublicAPI/PublicAPI.Unshipped.txt

@@ -1,2 +1,6 @@
 #nullable enable
 [SER004]RESPite.Messages.RespReader.TryParseScalar<T>(RESPite.Messages.RespReader.ScalarParser<char, T>! parser, out T value) -> bool
+[SER004]static RESPite.AsciiHash.EqualsCI(System.ReadOnlySpan<byte> first, System.ReadOnlySpan<char> second) -> bool
+[SER004]static RESPite.AsciiHash.EqualsCI(System.ReadOnlySpan<char> first, System.ReadOnlySpan<byte> second) -> bool
+[SER004]static RESPite.AsciiHash.SequenceEqualsCI(System.ReadOnlySpan<byte> first, System.ReadOnlySpan<char> second) -> bool
+[SER004]static RESPite.AsciiHash.SequenceEqualsCI(System.ReadOnlySpan<char> first, System.ReadOnlySpan<byte> second) -> bool

src/RESPite/Shared/AsciiHash.cs

@@ -82,6 +82,9 @@ public static bool EqualsCI(ReadOnlySpan<byte> first, ReadOnlySpan<byte> second)
         return len <= MaxBytesHashed ? HashUC(first) == HashUC(second) : SequenceEqualsCI(first, second);
     }
 
+    public static bool EqualsCI(ReadOnlySpan<byte> first, ReadOnlySpan<char> second)
+        => EqualsCI(second, first);
+
     public static unsafe bool SequenceEqualsCI(ReadOnlySpan<byte> first, ReadOnlySpan<byte> second)
     {
         var len = first.Length;
@@ -117,6 +120,9 @@ public static unsafe bool SequenceEqualsCI(ReadOnlySpan<byte> first, ReadOnlySpa
         }
     }
 
+    public static bool SequenceEqualsCI(ReadOnlySpan<byte> first, ReadOnlySpan<char> second)
+        => SequenceEqualsCI(second, first);
+
     public static bool EqualsCS(ReadOnlySpan<char> first, ReadOnlySpan<char> second)
     {
         var len = first.Length;
@@ -136,6 +142,14 @@ public static bool EqualsCI(ReadOnlySpan<char> first, ReadOnlySpan<char> second)
         return len <= MaxBytesHashed ? HashUC(first) == HashUC(second) : SequenceEqualsCI(first, second);
     }
 
+    public static bool EqualsCI(ReadOnlySpan<char> first, ReadOnlySpan<byte> second)
+    {
+        var len = first.Length;
+        if (len != second.Length) return false;
+        // for very short values, the UC hash performs CI equality
+        return len <= MaxBytesHashed ? HashUC(first) == HashUC(second) : SequenceEqualsCI(first, second);
+    }
+
     public static unsafe bool SequenceEqualsCI(ReadOnlySpan<char> first, ReadOnlySpan<char> second)
     {
         var len = first.Length;
@@ -171,6 +185,41 @@ public static unsafe bool SequenceEqualsCI(ReadOnlySpan<char> first, ReadOnlySpa
         }
     }
 
+    public static unsafe bool SequenceEqualsCI(ReadOnlySpan<char> first, ReadOnlySpan<byte> second)
+    {
+        var len = first.Length;
+        if (len != second.Length) return false;
+
+        // OK, don't be clever (SIMD, etc); the purpose of FashHash is to compare RESP key tokens, which are
+        // typically relatively short, think 3-20 bytes. That wouldn't even touch a SIMD vector, so:
+        // just loop (the exact thing we'd need to do *anyway* in a SIMD implementation, to mop up the non-SIMD
+        // trailing bytes).
+        fixed (char* firstPtr = &MemoryMarshal.GetReference(first))
+        {
+            fixed (byte* secondPtr = &MemoryMarshal.GetReference(second))
+            {
+                const int CS_MASK = 0b0101_1111;
+                for (int i = 0; i < len; i++)
+                {
+                    int x = (byte)firstPtr[i];
+                    var xCI = x & CS_MASK;
+                    if (xCI >= 'A' & xCI <= 'Z')
+                    {
+                        // alpha mismatch
+                        if (xCI != (secondPtr[i] & CS_MASK)) return false;
+                    }
+                    else if (x != secondPtr[i])
+                    {
+                        // non-alpha mismatch
+                        return false;
+                    }
+                }
+
+                return true;
+            }
+        }
+    }
+
     public static void Hash(scoped ReadOnlySpan<byte> value, out long cs, out long uc)
     {
         cs = HashCS(value);

src/StackExchange.Redis/CommandMap.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Runtime.CompilerServices;
 using System.Text;
 using RESPite;
 
@@ -10,9 +11,14 @@ namespace StackExchange.Redis
     /// </summary>
     public sealed class CommandMap
     {
-        private readonly AsciiHash[] map;
+        private readonly CommandBytes[] map;
+        private readonly byte[] bytes;
 
-        internal CommandMap(AsciiHash[] map) => this.map = map;
+        private CommandMap(CommandBytes[] map, byte[] bytes)
+        {
+            this.map = map;
+            this.bytes = bytes;
+        }
 
         /// <summary>
         /// The default commands specified by redis.
@@ -184,22 +190,31 @@ internal void AppendDeltas(StringBuilder sb)
                 var knownCmd = all[i];
                 if (knownCmd is RedisCommand.UNKNOWN) continue;
                 var keyString = knownCmd.ToString();
-                var keyBytes = new AsciiHash(keyString);
                 var value = map[i];
-                if (!keyBytes.Equals(value))
+                var valueBytes = value.GetCommandBytes(bytes);
+                if (!AsciiHash.EqualsCI(keyString, valueBytes))
                 {
                     if (sb.Length != 0) sb.Append(',');
-                    sb.Append('$').Append(keyString).Append('=').Append(value);
+                    sb.Append('$').Append(keyString).Append('=');
+                    if (!valueBytes.IsEmpty)
+                    {
+                        sb.Append(Encoding.ASCII.GetString(valueBytes));
+                    }
                 }
             }
         }
 
         internal void AssertAvailable(RedisCommand command)
         {
-            if (map[(int)command].IsEmpty) throw ExceptionFactory.CommandDisabled(command);
+            if (map[(int)command].IsEmpty) ThrowCommandDisabled(command);
+
+            [MethodImpl(MethodImplOptions.NoInlining)]
+            static void ThrowCommandDisabled(RedisCommand command) => throw ExceptionFactory.CommandDisabled(command);
         }
 
-        internal AsciiHash GetBytes(RedisCommand command) => map[(int)command];
+        internal ReadOnlySpan<byte> GetCommandBytes(RedisCommand command) => map[(int)command].GetCommandBytes(bytes);
+
+        internal ReadOnlySpan<byte> GetResp(RedisCommand command) => map[(int)command].GetResp(bytes);
 
         internal bool IsAvailable(RedisCommand command) => !map[(int)command].IsEmpty;
 
@@ -211,27 +226,107 @@ private static CommandMap CreateImpl(Dictionary<string, string?>? caseInsensitiv
         {
             var commands = AllCommands;
 
-            // todo: optimize and support ad-hoc overrides/disables, and shared buffer rather than multiple arrays
-            var map = new AsciiHash[commands.Length];
+            int totalLength = 0;
             for (int i = 0; i < commands.Length; i++)
             {
-                int idx = (int)commands[i];
-                string? name = commands[i].ToString(), value = name;
+                var value = GetCommandValue(commands[i], caseInsensitiveOverrides, exclusions);
+                if (string.IsNullOrEmpty(value)) continue;
 
-                if (commands[i] is RedisCommand.UNKNOWN || exclusions?.Contains(commands[i]) == true)
-                {
-                    map[idx] = default;
-                }
-                else
-                {
-                    if (caseInsensitiveOverrides != null && caseInsensitiveOverrides.TryGetValue(name, out string? tmp))
-                    {
-                        value = tmp?.ToUpperInvariant();
-                    }
-                    map[idx] = new AsciiHash(value);
-                }
+                totalLength += GetBulkStringLength(Encoding.ASCII.GetByteCount(value));
+            }
+
+            // Store all mapped command names as RESP bulk-string fragments in one buffer - everything is then
+            // ready to throw directly into the stream.
+            var map = new CommandBytes[commands.Length];
+
+            // Currently (8.8-ish) this is approx 3k; that's very reasonable to avoid a ton of CPU cycles in
+            // the most common write path.
+            var bytes = totalLength == 0 ? Array.Empty<byte>() : new byte[totalLength];
+            int offset = 0;
+            for (int i = 0; i < commands.Length; i++)
+            {
+                var command = commands[i];
+                var value = GetCommandValue(command, caseInsensitiveOverrides, exclusions);
+                if (string.IsNullOrEmpty(value)) continue;
+
+                int payloadLength = Encoding.ASCII.GetByteCount(value);
+                int respLength = GetBulkStringLength(payloadLength);
+                map[(int)command] = new CommandBytes(offset, respLength, GetPayloadOffset(payloadLength));
+
+                var span = bytes.AsSpan(offset, respLength);
+                span[0] = (byte)'$';
+                int payloadOffset = MessageWriter.WriteRaw(span, payloadLength, offset: 1);
+                var payload = span.Slice(payloadOffset, payloadLength);
+                int written = Encoding.ASCII.GetBytes(value.AsSpan(), payload);
+                if (written != payloadLength) ThrowAsciiEncodeLengthCheckFailure();
+                AsciiHash.ToUpper(payload);
+                MessageWriter.WriteCrlf(span, payloadOffset + payloadLength);
+                offset += respLength;
+            }
+            return new CommandMap(map, bytes);
+
+            [MethodImpl(MethodImplOptions.NoInlining)]
+            static void ThrowAsciiEncodeLengthCheckFailure() => throw new InvalidOperationException("ASCII encode length check failure");
+        }
+
+        private static string? GetCommandValue(
+            RedisCommand command,
+            Dictionary<string, string?>? caseInsensitiveOverrides,
+            HashSet<RedisCommand>? exclusions)
+        {
+            if (command is RedisCommand.UNKNOWN || exclusions?.Contains(command) == true) return null;
+
+            var name = command.ToString();
+            if (caseInsensitiveOverrides != null && caseInsensitiveOverrides.TryGetValue(name, out string? value))
+            {
+                return value;
+            }
+
+            return name;
+        }
+
+        // ${N}\r\n{RAW}\r\n
+        private static int GetBulkStringLength(int payloadLength) => 5 + GetDigitCount(payloadLength) + payloadLength;
+
+        // ${N}\r\n
+        private static byte GetPayloadOffset(int payloadLength) => checked((byte)(3 + GetDigitCount(payloadLength)));
+
+        private static int GetDigitCount(int value)
+        {
+            if (value < 10) return 1;
+            if (value < 100) return 2;
+            int digits = 1;
+            while ((value /= 10) != 0)
+            {
+                digits++;
+            }
+            return digits;
+        }
+
+        private readonly struct CommandBytes(int offset, int length, byte payloadOffset)
+        {
+            // Tracks position inside a shared buffer; given
+            // $3\r\nFOO\r\n$3\r\nBAR\r\n we have the positions (for BAR):
+            //              ^ a   ^ b    ^c
+            // We know that the trailer is always exactly 2 bytes, so we don't need to store the
+            // length of the command itself - we can infer from the other values.
+            // offset is a, payloadOffset is a-to-b, length is a-to-c
+            private readonly uint offset = checked((uint)offset);
+            private readonly ushort length = checked((ushort)length);
+
+            public bool IsEmpty => length == 0;
+
+            // this will be fine even for a default instance
+            public ReadOnlySpan<byte> GetResp(byte[] bytes) => new(bytes, (int)offset, length);
+
+            public ReadOnlySpan<byte> GetCommandBytes(byte[] bytes)
+            {
+                if (IsEmpty) return default;
+                return new ReadOnlySpan<byte>(
+                    bytes,
+                    checked((int)offset) + payloadOffset,
+                    length - payloadOffset - 2);
             }
-            return new CommandMap(map);
         }
     }
 }

src/StackExchange.Redis/MessageWriter.cs

@@ -131,7 +131,28 @@ internal void WriteHeader(string command, int arguments)
         }
     }
 
-    internal void WriteHeader(RedisCommand command, int arguments) => WriteHeader(command, arguments, _map.GetBytes(command).Span);
+    internal void WriteHeader(RedisCommand command, int arguments)
+    {
+        // using >= here because we will be adding 1 for the command itself (which is an arg for the purposes of the multi-bulk protocol)
+        if (arguments >= REDIS_MAX_ARGS) throw ExceptionFactory.TooManyArgs(command.ToString(), arguments);
+
+        // in theory we should never see this; CheckMessage dealt with "regular" messages, and
+        // ExecuteMessage should have dealt with everything else
+        var commandBytes = _map.GetResp(command);
+        if (commandBytes.IsEmpty) throw ExceptionFactory.CommandDisabled(command);
+
+        // *{argCount}\r\n      = 3 + MaxInt32TextLen
+        // ${cmd-len}\r\n       = precomputed
+        // {cmd}\r\n            = precomputed
+        var span = _writer.GetSpan(commandBytes.Length + 3 + Format.MaxInt32TextLen);
+        span[0] = (byte)'*';
+
+        int offset = WriteRaw(span, arguments + 1, offset: 1);
+        commandBytes.CopyTo(span.Slice(offset));
+        offset += commandBytes.Length;
+
+        _writer.Advance(offset);
+    }
 
     internal void WriteHeader(RedisCommand command, int arguments, ReadOnlySpan<byte> commandBytes)
     {

tests/StackExchange.Redis.Tests/AsciiHashUnitTests.cs

@@ -150,6 +150,33 @@ public void CaseInsensitiveEquality(string text)
         Assert.Equal(hashLowerUC, hashUpperUC);
     }
 
+    [Theory]
+    [InlineData("a")] // length 1
+    [InlineData("ab")] // length 2
+    [InlineData("abc")] // length 3
+    [InlineData("abcd")] // length 4
+    [InlineData("abcde")] // length 5
+    [InlineData("abcdef")] // length 6
+    [InlineData("abcdefg")] // length 7
+    [InlineData("abcdefgh")] // length 8
+    [InlineData("abcdefghi")] // length 9
+    [InlineData("abcdefghij")] // length 10
+    [InlineData("abcdefghijklmnop")] // length 16
+    [InlineData("abcdefghijklmnopqrst")] // length 20
+    public void CaseInsensitiveEquality_MixedBytesAndChars(string text)
+    {
+        var lowerChars = text.AsSpan();
+        var upperBytes = Encoding.UTF8.GetBytes(text.ToUpperInvariant());
+
+        Assert.True(AsciiHash.EqualsCI(lowerChars, upperBytes), "CI: chars lower == bytes upper");
+        Assert.True(AsciiHash.EqualsCI(upperBytes, lowerChars), "CI: bytes upper == chars lower");
+
+        Assert.True(AsciiHash.SequenceEqualsCI(lowerChars, upperBytes), "CI sequence: chars lower == bytes upper");
+        Assert.True(AsciiHash.SequenceEqualsCI(upperBytes, lowerChars), "CI sequence: bytes upper == chars lower");
+
+        Assert.False(AsciiHash.EqualsCI((text + "x").AsSpan(), upperBytes), "CI: length mismatch");
+    }
+
     [Theory]
     [InlineData("a")] // length 1
     [InlineData("ab")] // length 2

tests/StackExchange.Redis.Tests/TestHarnessTests.cs

@@ -53,6 +53,15 @@ public void WithKeyPrefix()
         resp.ValidateResp("*3\r\n$3\r\nPUT\r\n$9\r\n123/mykey\r\n$2\r\n42\r\n", "set", args);
     }
 
+    [Fact]
+    public void CommandMapDeltas()
+    {
+        Assert.Equal("", CommandMap.Default.ToString());
+        Assert.Equal("", CommandMap.Create(new() { ["sEt"] = "set" }).ToString());
+        Assert.Equal("$SET=PUT", CommandMap.Create(new() { ["sEt"] = "put" }).ToString());
+        Assert.Equal("$ECHO=", CommandMap.Create(new() { "echo" }, available: false).ToString());
+    }
+
     [Fact]
     public void WithKeyPrefix_DetectIncorrectUsage()
     {