[9.0] Make counting of IO completion work items more precise on Windows (#112794)

* Stop counting work items from ThreadPoolTypedWorkItemQueue for ThreadPool.CompletedWorkItemCount (#106854)

* Stop counting work items from ThreadPoolTypedWorkItemQueue as completed work items

---------

Co-authored-by: Eduardo Manuel Velarde Polar <[email protected]>
Co-authored-by: Koundinya Veluri <[email protected]>

* Make counting of IO completion work items more precise on Windows

- Follow-up to #106854. Issue: #104284.
- Before the change, the modified test case often yields 5 or 6 completed work items, due to queue-processing work items that happen to not process any user work items. After the change, it always yields 4.
- Looks like it doesn't hurt to have more-precise counting, and there was a request to backport a fix to .NET 8, where it's more necessary to fix the issue

---------

Co-authored-by: Eduardo Velarde <[email protected]>
Co-authored-by: Eduardo Manuel Velarde Polar <[email protected]>

Author: 3 people

src/libraries/System.Private.CoreLib/src/System/Threading/ThreadInt64PersistentCounter.cs

@@ -15,6 +15,7 @@ internal sealed class ThreadInt64PersistentCounter
         private static List<ThreadLocalNodeFinalizationHelper>? t_nodeFinalizationHelpers;
 
         private long _overflowCount;
+        private long _lastReturnedCount;
 
         // dummy node serving as a start and end of the ring list
         private readonly ThreadLocalNode _nodes;
@@ -31,6 +32,13 @@ public static void Increment(object threadLocalCountObject)
             Unsafe.As<ThreadLocalNode>(threadLocalCountObject).Increment();
         }
 
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static void Decrement(object threadLocalCountObject)
+        {
+            Debug.Assert(threadLocalCountObject is ThreadLocalNode);
+            Unsafe.As<ThreadLocalNode>(threadLocalCountObject).Decrement();
+        }
+
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static void Add(object threadLocalCountObject, uint count)
         {
@@ -76,6 +84,17 @@ public long Count
                         count += node.Count;
                         node = node._next;
                     }
+
+                    // Ensure that the returned value is monotonically increasing
+                    long lastReturnedCount = _lastReturnedCount;
+                    if (count > lastReturnedCount)
+                    {
+                        _lastReturnedCount = count;
+                    }
+                    else
+                    {
+                        count = lastReturnedCount;
+                    }
                 }
                 finally
                 {
@@ -134,6 +153,18 @@ public void Increment()
                 OnAddOverflow(1);
             }
 
+            [MethodImpl(MethodImplOptions.AggressiveInlining)]
+            public void Decrement()
+            {
+                if (_count != 0)
+                {
+                    _count--;
+                    return;
+                }
+
+                OnAddOverflow(-1);
+            }
+
             public void Add(uint count)
             {
                 Debug.Assert(count != 0);
@@ -149,7 +180,7 @@ public void Add(uint count)
             }
 
             [MethodImpl(MethodImplOptions.NoInlining)]
-            private void OnAddOverflow(uint count)
+            private void OnAddOverflow(long count)
             {
                 Debug.Assert(count != 0);
 
@@ -161,7 +192,7 @@ private void OnAddOverflow(uint count)
                 counter._lock.Acquire();
                 try
                 {
-                    counter._overflowCount += (long)_count + count;
+                    counter._overflowCount += _count + count;
                     _count = 0;
                 }
                 finally

src/libraries/System.Private.CoreLib/src/System/Threading/ThreadPoolWorkQueue.cs

@@ -1375,6 +1375,9 @@ void IThreadPoolWorkItem.Execute()
                 Debug.Assert(stageBeforeUpdate != QueueProcessingStage.NotScheduled);
                 if (stageBeforeUpdate == QueueProcessingStage.Determining)
                 {
+                    // Discount a work item here to avoid counting this queue processing work item
+                    ThreadInt64PersistentCounter.Decrement(
+                        ThreadPoolWorkQueueThreadLocals.threadLocals!.threadLocalCompletionCountObject!);
                     return;
                 }
             }
@@ -1414,7 +1417,11 @@ void IThreadPoolWorkItem.Execute()
                 currentThread.ResetThreadPoolThread();
             }
 
-            ThreadInt64PersistentCounter.Add(tl.threadLocalCompletionCountObject!, completedCount);
+            // Discount a work item here to avoid counting this queue processing work item
+            if (completedCount > 1)
+            {
+                ThreadInt64PersistentCounter.Add(tl.threadLocalCompletionCountObject!, completedCount - 1);
+            }
         }
     }
 

src/libraries/System.Threading.ThreadPool/tests/ThreadPoolTests.cs

@@ -1462,6 +1462,42 @@ static async Task RunAsyncIOTest()
             }, ioCompletionPortCount.ToString()).Dispose();
         }
 
+
+        [ConditionalFact(nameof(IsThreadingAndRemoteExecutorSupported))]
+        [PlatformSpecific(TestPlatforms.Windows)]
+        public static unsafe void ThreadPoolCompletedWorkItemCountTest()
+        {
+            // Run in a separate process to test in a clean thread pool environment such that we don't count external work items
+            RemoteExecutor.Invoke(() =>
+            {
+                const int WorkItemCount = 4;
+
+                int completedWorkItemCount = 0;
+                using var allWorkItemsCompleted = new AutoResetEvent(false);
+
+                IOCompletionCallback callback =
+                    (errorCode, numBytes, innerNativeOverlapped) =>
+                    {
+                        Overlapped.Free(innerNativeOverlapped);
+                        if (Interlocked.Increment(ref completedWorkItemCount) == WorkItemCount)
+                        {
+                            allWorkItemsCompleted.Set();
+                        }
+                    };
+                for (int i = 0; i < WorkItemCount; i++)
+                {
+                    ThreadPool.UnsafeQueueNativeOverlapped(new Overlapped().Pack(callback, null));
+                }
+
+                allWorkItemsCompleted.CheckedWait();
+
+                // Allow work items to be marked as completed during this time
+                ThreadTestHelpers.WaitForCondition(() => ThreadPool.CompletedWorkItemCount >= WorkItemCount);
+                Thread.Sleep(50);
+                Assert.Equal(WorkItemCount, ThreadPool.CompletedWorkItemCount);
+            }).Dispose();
+        }
+
         public static bool IsThreadingAndRemoteExecutorSupported =>
             PlatformDetection.IsThreadingSupported && RemoteExecutor.IsSupported;