From fc96527a366f263df1d77f364c4d1a2cff3dde72 Mon Sep 17 00:00:00 2001
From: Konst Khurin <kkhurin@microsoft.com>
Date: Wed, 1 Mar 2023 00:15:43 -0800
Subject: [PATCH] Fixes for ArrayDictionary and extra memory allocations
 (#1764)

* Fixing ArrayDictionary namespace and further reducing allocations

* updated fix
---
 .../ArrayDictionary.cs                        | 13 +++--
 .../ServiceProtocol.cs                        | 47 ++++++++++++++++---
 2 files changed, 47 insertions(+), 13 deletions(-)

diff --git a/src/Microsoft.Azure.SignalR.Protocols/ArrayDictionary.cs b/src/Microsoft.Azure.SignalR.Protocols/ArrayDictionary.cs
index c8126458e..8c8afc2ed 100644
--- a/src/Microsoft.Azure.SignalR.Protocols/ArrayDictionary.cs
+++ b/src/Microsoft.Azure.SignalR.Protocols/ArrayDictionary.cs
@@ -5,8 +5,9 @@
 
 #nullable enable
 
-namespace Microsoft.Azure.SignalR
+namespace Microsoft.Azure.SignalR.Protocol
 {
+
     /// <summary>
     /// Lightweight, read-only IDictionary implementation using two arrays
     /// and O(n) lookup.
@@ -19,8 +20,9 @@ namespace Microsoft.Azure.SignalR
     public class ArrayDictionary<TKey, TValue>
         : IDictionary<TKey, TValue>, IReadOnlyDictionary<TKey, TValue>
         where TKey : notnull
-        where TValue : struct
     {
+        public static readonly ArrayDictionary<TKey, TValue> Empty = new(0);
+
         private readonly IEqualityComparer<TKey> _comparer;
         private readonly TKey[] _keys;
         private readonly TValue[] _values;
@@ -34,9 +36,6 @@ public ArrayDictionary(int capacity, IEqualityComparer<TKey>? comparer = null)
             _comparer = comparer ?? EqualityComparer<TKey>.Default;
         }
 
-        public static IDictionary<TKey, TValue> Create(int capacity) =>
-            new ArrayDictionary<TKey, TValue>(capacity);
-
         public TValue this[TKey key]
         {
             get
@@ -150,7 +149,7 @@ public bool TryGetValue(TKey key, out TValue value)
                     return true;
                 }
             }
-            value = default;
+            value = default!;
             return false;
         }
 
@@ -188,4 +187,4 @@ public void Reset()
             }
         }
     }
-}
+}
\ No newline at end of file
diff --git a/src/Microsoft.Azure.SignalR.Protocols/ServiceProtocol.cs b/src/Microsoft.Azure.SignalR.Protocols/ServiceProtocol.cs
index 47d7ecba3..3a218ec76 100644
--- a/src/Microsoft.Azure.SignalR.Protocols/ServiceProtocol.cs
+++ b/src/Microsoft.Azure.SignalR.Protocols/ServiceProtocol.cs
@@ -811,7 +811,7 @@ private static PingMessage CreatePingMessage(ref MessagePackReader reader, int a
                 var values = new string[length];
                 for (int i = 0; i < length; i++)
                 {
-                    values[i] = ReadString(ref reader, $"messages[{i}]");
+                    values[i] = ReadString(ref reader, "messages[{0}]", i);
                 }
 
                 return new PingMessage { Messages = values };
@@ -1266,8 +1266,8 @@ private static Claim[] ReadClaims(ref MessagePackReader reader)
 
                 for (var i = 0; i < claimCount; i++)
                 {
-                    var type = ReadString(ref reader, $"claims[{i}].Type");
-                    var value = ReadString(ref reader, $"claims[{i}].Value");
+                    var type = ReadString(ref reader, "claims[{0}].Type", i);
+                    var value = ReadString(ref reader, "claims[{0}].Value", i);
                     claims[i] = new Claim(type, value);
                 }
 
@@ -1285,8 +1285,8 @@ private static IDictionary<string, ReadOnlyMemory<byte>> ReadPayloads(ref Messag
                 var payloads = new ArrayDictionary<string, ReadOnlyMemory<byte>>((int)payloadCount, StringComparer.OrdinalIgnoreCase);
                 for (var i = 0; i < payloadCount; i++)
                 {
-                    var key = ReadString(ref reader, $"payloads[{i}].key");
-                    var value = ReadBytes(ref reader, $"payloads[{i}].value");
+                    var key = ReadString(ref reader, "payloads[{0}].key", i);
+                    var value = ReadBytes(ref reader, "payloads[{0}].value", i);
                     payloads.Add(key, value);
                 }
 
@@ -1345,6 +1345,30 @@ private static string ReadString(ref MessagePackReader reader, string field)
             }
         }
 
+        private static string ReadString(ref MessagePackReader reader, string formatField, int param)
+        {
+            try
+            {
+                return reader.ReadString();
+            }
+            catch (Exception ex)
+            {
+                throw new InvalidDataException($"Reading '{string.Format(formatField, param)}' as String failed.", ex);
+            }
+        }
+
+        private static string ReadString(ref MessagePackReader reader, string formatField, string param1, int param2)
+        {
+            try
+            {
+                return reader.ReadString();
+            }
+            catch (Exception ex)
+            {
+                throw new InvalidDataException($"Reading '{string.Format(formatField, param1, param2)}' as String failed.", ex);
+            }
+        }
+
         private static string[] ReadStringArray(ref MessagePackReader reader, string field)
         {
             var arrayLength = ReadArrayLength(ref reader, field);
@@ -1353,7 +1377,7 @@ private static string[] ReadStringArray(ref MessagePackReader reader, string fie
                 var array = new string[arrayLength];
                 for (int i = 0; i < arrayLength; i++)
                 {
-                    array[i] = ReadString(ref reader, $"{field}[{i}]");
+                    array[i] = ReadString(ref reader, "{0}[{1}]", field, i);
                 }
 
                 return array;
@@ -1372,7 +1396,18 @@ private static byte[] ReadBytes(ref MessagePackReader reader, string field)
             {
                 throw new InvalidDataException($"Reading '{field}' as Byte[] failed.", ex);
             }
+        }
 
+        private static byte[] ReadBytes(ref MessagePackReader reader, string formatField, int param)
+        {
+            try
+            {
+                return reader.ReadBytes()?.ToArray() ?? Array.Empty<byte>();
+            }
+            catch (Exception ex)
+            {
+                throw new InvalidDataException($"Reading '{string.Format(formatField, param)}' as Byte[] failed.", ex);
+            }
         }
 
         private static long ReadMapLength(ref MessagePackReader reader, string field)