Avoid using ConcurrentDictionary for channels with few methods (#2597)

Co-authored-by: Günther Foidl <[email protected]>

Author: JamesNK

.github/workflows/test.yml

@@ -15,7 +15,7 @@ jobs:
 
   build:
     name: Basic Tests
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     steps:
 
     - name: Check out code
@@ -36,7 +36,7 @@ jobs:
 
   grpc_web:
     name: gRPC-Web Tests
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     steps:
 
     - name: Check out code

src/Grpc.Net.Client/GrpcChannel.cs

@@ -16,7 +16,6 @@
 
 #endregion
 
-using System.Collections.Concurrent;
 using System.Diagnostics;
 using Grpc.Core;
 #if SUPPORT_LOAD_BALANCING
@@ -51,7 +50,7 @@ public sealed partial class GrpcChannel : ChannelBase, IDisposable
     internal const long DefaultMaxRetryBufferPerCallSize = 1024 * 1024; // 1 MB
 
     private readonly object _lock;
-    private readonly ConcurrentDictionary<IMethod, GrpcMethodInfo> _methodInfoCache;
+    private readonly ThreadSafeLookup<IMethod, GrpcMethodInfo> _methodInfoCache;
     private readonly Func<IMethod, GrpcMethodInfo> _createMethodInfoFunc;
     private readonly Dictionary<MethodKey, MethodConfig>? _serviceConfigMethods;
     private readonly bool _isSecure;
@@ -109,7 +108,7 @@ public sealed partial class GrpcChannel : ChannelBase, IDisposable
     internal GrpcChannel(Uri address, GrpcChannelOptions channelOptions) : base(address.Authority)
     {
         _lock = new object();
-        _methodInfoCache = new ConcurrentDictionary<IMethod, GrpcMethodInfo>();
+        _methodInfoCache = new ThreadSafeLookup<IMethod, GrpcMethodInfo>();
 
         // Dispose the HTTP client/handler if...
         //   1. No client/handler was specified and so the channel created the client itself

src/Grpc.Net.Client/Internal/ThreadSafeLookup.cs

@@ -0,0 +1,104 @@
+#region Copyright notice and license
+
+// Copyright 2019 The gRPC Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#endregion
+
+using System.Collections.Concurrent;
+
+internal sealed class ThreadSafeLookup<TKey, TValue> where TKey : notnull
+{
+    // Avoid allocating ConcurrentDictionary until the threshold is reached.
+    // Looking up a key in an array is as fast as a dictionary for small collections and uses much less memory.
+    internal const int Threshold = 10;
+
+    private KeyValuePair<TKey, TValue>[] _array = Array.Empty<KeyValuePair<TKey, TValue>>();
+    private ConcurrentDictionary<TKey, TValue>? _dictionary;
+
+    /// <summary>
+    /// Gets the value for the key if it exists. If the key does not exist then the value is created using the valueFactory.
+    /// The value is created outside of a lock and there is no guarentee which value will be stored or returned.
+    /// </summary>
+    public TValue GetOrAdd(TKey key, Func<TKey, TValue> valueFactory)
+    {
+        if (_dictionary != null)
+        {
+            return _dictionary.GetOrAdd(key, valueFactory);
+        }
+
+        if (TryGetValue(_array, key, out var value))
+        {
+            return value;
+        }
+
+        var newValue = valueFactory(key);
+
+        lock (this)
+        {
+            if (_dictionary != null)
+            {
+                _dictionary.TryAdd(key, newValue);
+            }
+            else
+            {
+                // Double check inside lock if the key was added to the array by another thread.
+                if (TryGetValue(_array, key, out value))
+                {
+                    return value;
+                }
+
+                if (_array.Length > Threshold - 1)
+                {
+                    // Array length exceeds threshold so switch to dictionary.
+                    var newDict = new ConcurrentDictionary<TKey, TValue>();
+                    foreach (var kvp in _array)
+                    {
+                        newDict.TryAdd(kvp.Key, kvp.Value);
+                    }
+                    newDict.TryAdd(key, newValue);
+
+                    _dictionary = newDict;
+                    _array = Array.Empty<KeyValuePair<TKey, TValue>>();
+                }
+                else
+                {
+                    // Add new value by creating a new array with old plus new value.
+                    var newArray = new KeyValuePair<TKey, TValue>[_array.Length + 1];
+                    Array.Copy(_array, newArray, _array.Length);
+                    newArray[newArray.Length - 1] = new KeyValuePair<TKey, TValue>(key, newValue);
+
+                    _array = newArray;
+                }
+            }
+        }
+
+        return newValue;
+    }
+
+    private static bool TryGetValue(KeyValuePair<TKey, TValue>[] array, TKey key, out TValue value)
+    {
+        foreach (var kvp in array)
+        {
+            if (EqualityComparer<TKey>.Default.Equals(kvp.Key, key))
+            {
+                value = kvp.Value;
+                return true;
+            }
+        }
+
+        value = default!;
+        return false;
+    }
+}

test/Grpc.Net.Client.Tests/ThreadSafeLookupTests.cs

@@ -0,0 +1,69 @@
+#region Copyright notice and license
+
+// Copyright 2019 The gRPC Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#endregion
+
+namespace Grpc.Net.Client.Tests;
+
+[TestFixture]
+public class ThreadSafeLookupTests
+{
+    [Test]
+    public void GetOrAdd_ReturnsCorrectValueForNewKey()
+    {
+        var lookup = new ThreadSafeLookup<int, string>();
+        var result = lookup.GetOrAdd(1, k => "Value-1");
+
+        Assert.AreEqual("Value-1", result);
+    }
+
+    [Test]
+    public void GetOrAdd_ReturnsExistingValueForExistingKey()
+    {
+        var lookup = new ThreadSafeLookup<int, string>();
+        lookup.GetOrAdd(1, k => "InitialValue");
+        var result = lookup.GetOrAdd(1, k => "NewValue");
+
+        Assert.AreEqual("InitialValue", result);
+    }
+
+    [Test]
+    public void GetOrAdd_SwitchesToDictionaryAfterThreshold()
+    {
+        var addCount = (ThreadSafeLookup<int, string>.Threshold * 2);
+        var lookup = new ThreadSafeLookup<int, string>();
+
+        for (var i = 0; i <= addCount; i++)
+        {
+            lookup.GetOrAdd(i, k => $"Value-{k}");
+        }
+
+        var result = lookup.GetOrAdd(addCount, k => $"NewValue-{addCount}");
+
+        Assert.AreEqual($"Value-{addCount}", result);
+    }
+
+    [Test]
+    public void GetOrAdd_HandlesConcurrentAccess()
+    {
+        var lookup = new ThreadSafeLookup<int, string>();
+        Parallel.For(0, 1000, i =>
+        {
+            var value = lookup.GetOrAdd(i, k => $"Value-{k}");
+            Assert.AreEqual($"Value-{i}", value);
+        });
+    }
+}