From 30a2709528ce52c9576513c43d5370def593171b Mon Sep 17 00:00:00 2001
From: Koundinya Veluri <kouvel@users.noreply.github.com>
Date: Fri, 13 Aug 2021 21:52:46 -0700
Subject: [PATCH] Fix thread pool completed work item counter to remove
 finalized nodes (#57305)

* Fix thread pool completed work item counter to remove finalized nodes

- New nodes are being added to the set but finalized ones were not being removed
- This would result in double-counting the completed work item count for a thread that exits, and memory would grow over time

* Add a separate finalization helper for each node to track thread-local lifetimes

* Remove unnecessary IDisposable

* Fix nullability issues

* Update initialization of thread-local as suggested

Co-authored-by: Stephen Toub <stoub@microsoft.com>

Co-authored-by: Stephen Toub <stoub@microsoft.com>
---
 .../Threading/ThreadInt64PersistentCounter.cs      | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/src/libraries/System.Private.CoreLib/src/System/Threading/ThreadInt64PersistentCounter.cs b/src/libraries/System.Private.CoreLib/src/System/Threading/ThreadInt64PersistentCounter.cs
index 61fe075..c152087 100644
--- a/src/libraries/System.Private.CoreLib/src/System/Threading/ThreadInt64PersistentCounter.cs
+++ b/src/libraries/System.Private.CoreLib/src/System/Threading/ThreadInt64PersistentCounter.cs
@@ -12,6 +12,9 @@ namespace System.Threading
     {
         private static readonly LowLevelLock s_lock = new LowLevelLock();
 
+        [ThreadStatic]
+        private static List<ThreadLocalNodeFinalizationHelper>? t_nodeFinalizationHelpers;
+
         private long _overflowCount;
         private HashSet<ThreadLocalNode> _nodes = new HashSet<ThreadLocalNode>();
 
@@ -26,6 +29,9 @@ namespace System.Threading
         {
             var node = new ThreadLocalNode(this);
 
+            List<ThreadLocalNodeFinalizationHelper>? nodeFinalizationHelpers = t_nodeFinalizationHelpers ??= new List<ThreadLocalNodeFinalizationHelper>(1);
+            nodeFinalizationHelpers.Add(new ThreadLocalNodeFinalizationHelper(node));
+
             s_lock.Acquire();
             try
             {
@@ -76,13 +82,14 @@ namespace System.Threading
                 _counter = counter;
             }
 
-            ~ThreadLocalNode()
+            public void Dispose()
             {
                 ThreadInt64PersistentCounter counter = _counter;
                 s_lock.Acquire();
                 try
                 {
                     counter._overflowCount += _count;
+                    counter._nodes.Remove(this);
                 }
                 finally
                 {
@@ -125,5 +132,18 @@ namespace System.Threading
                 }
             }
         }
+
+        private sealed class ThreadLocalNodeFinalizationHelper
+        {
+            private readonly ThreadLocalNode _node;
+
+            public ThreadLocalNodeFinalizationHelper(ThreadLocalNode node)
+            {
+                Debug.Assert(node != null);
+                _node = node;
+            }
+
+            ~ThreadLocalNodeFinalizationHelper() => _node.Dispose();
+        }
     }
 }
-- 
2.7.4