From d8e2d11a1e6891939fc9e124a51401b3e4b246d3 Mon Sep 17 00:00:00 2001
From: Koundinya Veluri <kouvel@users.noreply.github.com>
Date: Wed, 29 Sep 2021 21:36:10 -0700
Subject: [PATCH] Fix an occasional small perf issue in `ThreadLocal`
 constructor/dispose (#59300)

- Replaced the linear search for a free ID with a pair of collections that operate in O(1) time for insertion and removal
- See https://github.com/dotnet/runtime/issues/59145#issuecomment-922162343 for more information
- Fixes https://github.com/dotnet/runtime/issues/59145
---
 .../src/System/Threading/ThreadLocal.cs       | 49 +++++++++++--------
 1 file changed, 28 insertions(+), 21 deletions(-)

diff --git a/src/libraries/System.Private.CoreLib/src/System/Threading/ThreadLocal.cs b/src/libraries/System.Private.CoreLib/src/System/Threading/ThreadLocal.cs
index d994a9db8c4..cf05a40a68e 100644
--- a/src/libraries/System.Private.CoreLib/src/System/Threading/ThreadLocal.cs
+++ b/src/libraries/System.Private.CoreLib/src/System/Threading/ThreadLocal.cs
@@ -652,39 +652,46 @@ private sealed class IdManager
             // Keep track of the count of non-TrackAllValues ids in use. A count of 0 leads to more efficient thread cleanup
             private volatile int _idsThatDoNotTrackAllValues;
 
-            private const byte IdFree = 0;
-            private const byte TrackAllValuesAllocated = 1;
-            private const byte DoNotTrackAllValuesAllocated = 2;
+            // Stores IDs that are used, and if each ID tracksAllValues or not.
+            private readonly Dictionary<int, bool> _usedIdToTracksAllValuesMap = new Dictionary<int, bool>();
 
-            // Stores whether each ID is free or not, and if it tracksAllValues or not. Additionally, the object is also used as a lock for the IdManager.
-            private readonly List<byte> _ids = new List<byte>();
+            // Stores IDs that were previously used and are now free to reuse. Additionally, the object is also used as a lock
+            // for the IdManager.
+            private readonly List<int> _freeIds = new List<int>();
 
             internal int GetId(bool trackAllValues)
             {
-                lock (_ids)
+                lock (_freeIds)
                 {
-                    int availableId = _nextIdToTry;
-                    while (availableId < _ids.Count)
+                    int availableId;
+                    int freeIdCount = _freeIds.Count;
+                    if (freeIdCount > 0)
                     {
-                        if (_ids[availableId] == IdFree) { break; }
-                        availableId++;
+                        availableId = _freeIds[freeIdCount - 1];
+                    }
+                    else
+                    {
+                        availableId = _nextIdToTry;
                     }
 
-                    byte allocatedFlag = trackAllValues ? TrackAllValuesAllocated : DoNotTrackAllValuesAllocated;
-                    if (availableId == _ids.Count)
+                    // Ensure that all of the IDs that will be used can be freed without throwing due to OOM when disposing or
+                    // finalizing
+                    _freeIds.EnsureCapacity(_usedIdToTracksAllValuesMap.Count + 1);
+
+                    _usedIdToTracksAllValuesMap.Add(availableId, trackAllValues);
+
+                    if (freeIdCount > 0)
                     {
-                        _ids.Add(allocatedFlag);
+                        _freeIds.RemoveAt(freeIdCount - 1);
                     }
                     else
                     {
-                        _ids[availableId] = allocatedFlag;
+                        _nextIdToTry = availableId + 1;
                     }
 
                     if (!trackAllValues)
                         _idsThatDoNotTrackAllValues++;
 
-                    _nextIdToTry = availableId + 1;
-
                     return availableId;
                 }
             }
@@ -692,9 +699,9 @@ internal int GetId(bool trackAllValues)
             // Identify if an allocated id tracks all values or not
             internal bool IdTracksAllValues(int id)
             {
-                lock (_ids)
+                lock (_freeIds)
                 {
-                    return _ids[id] == TrackAllValuesAllocated;
+                    return _usedIdToTracksAllValuesMap.TryGetValue(id, out bool tracksAllValues) && tracksAllValues;
                 }
             }
 
@@ -703,13 +710,13 @@ internal bool IdTracksAllValues(int id)
             // Return an ID to the pool
             internal void ReturnId(int id, bool idTracksAllValues)
             {
-                lock (_ids)
+                lock (_freeIds)
                 {
                     if (!idTracksAllValues)
                         _idsThatDoNotTrackAllValues--;
 
-                    _ids[id] = IdFree;
-                    if (id < _nextIdToTry) _nextIdToTry = id;
+                    _usedIdToTracksAllValuesMap.Remove(id);
+                    _freeIds.Add(id); // does not throw because the capacity is ensured in GetId()
                 }
             }
         }
-- 
GitLab