Stored the initial capacity of the ConcurrentDictionary for correctly sizing the backing array after clearing the collection. (#108065)

* Stored the initial capacity of the ConcurrentDictionary for correctly sizing the backing array after clearing the collection.

* Stored the capacity in the ctor.
* Used the stored capacity in Clear().

Fixes #107016

* Added a test to check the capacity logic of the ConcurrentDictionary.

Author: koenigst

src/libraries/System.Collections.Concurrent/src/System/Collections/Concurrent/ConcurrentDictionary.cs

@@ -42,6 +42,11 @@ public class ConcurrentDictionary<TKey, TValue> : IDictionary<TKey, TValue>, IDi
         /// extra branch when using a custom comparer with a reference type key.
         /// </remarks>
         private readonly bool _comparerIsDefaultForClasses;
+        /// <summary>The initial size of the _buckets array.</summary>
+        /// <remarks>
+        /// We store this to retain the initially specified growing behavior of the _buckets array even after clearing the collection.
+        /// </remarks>
+        private readonly int _initialCapacity;
 
         /// <summary>The default capacity, i.e. the initial # of buckets.</summary>
         /// <remarks>
@@ -220,6 +225,7 @@ internal ConcurrentDictionary(int concurrencyLevel, int capacity, bool growLockA
 
             _tables = new Tables(buckets, locks, countPerLock, comparer);
             _growLockArray = growLockArray;
+            _initialCapacity = capacity;
             _budget = buckets.Length / locks.Length;
         }
 
@@ -716,7 +722,7 @@ public void Clear()
                 }
 
                 Tables tables = _tables;
-                var newTables = new Tables(new VolatileNode[HashHelpers.GetPrime(DefaultCapacity)], tables._locks, new int[tables._countPerLock.Length], tables._comparer);
+                var newTables = new Tables(new VolatileNode[HashHelpers.GetPrime(_initialCapacity)], tables._locks, new int[tables._countPerLock.Length], tables._comparer);
                 _tables = newTables;
                 _budget = Math.Max(1, newTables._buckets.Length / newTables._locks.Length);
             }

src/libraries/System.Collections.Concurrent/tests/ConcurrentDictionary/ConcurrentDictionaryTests.cs

@@ -954,6 +954,25 @@ public static void TestClear()
             Assert.True(dictionary.IsEmpty, "TestClear: FAILED.  IsEmpty returned false after Clear");
         }
 
+        [Theory]
+        [InlineData(3)]
+        [InlineData(1_162_687)]
+        public static void TestCapacity(int capacity)
+        {
+            int itemsCount = capacity + 100;
+            var dictionary = new ConcurrentDictionary<int, int>(1, capacity);
+            Assert.Equal(capacity, GetCapacity(dictionary));
+
+            for (int i = 0; i < itemsCount; i++)
+                dictionary.TryAdd(i, i);
+
+            Assert.Equal(itemsCount, dictionary.Count);
+            Assert.InRange(GetCapacity(dictionary), capacity + 1, int.MaxValue);
+
+            dictionary.Clear();
+            Assert.Equal(capacity, GetCapacity(dictionary));
+        }
+
         [ConditionalFact(typeof(PlatformDetection), nameof(PlatformDetection.IsThreadingSupported))]
         public static void TestTryUpdate()
         {
@@ -1231,6 +1250,20 @@ public void Dictionary_NotCorruptedByNullReturningComparer()
         }
 
         #region Helper Classes and Methods
+
+        private static int GetCapacity(ConcurrentDictionary<int, int> dictionary)
+        {
+            var tables = typeof(ConcurrentDictionary<int, int>)
+                .GetField("_tables", BindingFlags.Instance | BindingFlags.NonPublic)!
+                .GetValue(dictionary);
+
+            var buckets = (ICollection)tables.GetType()
+                .GetField("_buckets", BindingFlags.Instance | BindingFlags.NonPublic)!
+                .GetValue(tables);
+
+            return buckets.Count;
+        }
+
         private sealed class CreateThrowsComparer : IEqualityComparer<string>, IAlternateEqualityComparer<ReadOnlySpan<char>, string>
         {
             public bool Equals(string? x, string? y) => EqualityComparer<string>.Default.Equals(x, y);