Stop tracking double disposal in TaskThreadPool

PerfectNumbers.Core/TaskThreadPool.cs

@@ -19,9 +19,16 @@ internal sealed class TaskThreadPool
     }
 
     private readonly ConcurrentQueue<WorkItem> _pendingTaskQueue = new();
-    private readonly BlockingCollection<WorkItem> _pendingTasks;
+    private readonly HashSet<Thread> _runningThreads = new();
-    private readonly Thread[] _threads;
+    private readonly object _lock = new();
     private readonly Action<Task> _taskExecutor;
+    private readonly int _maximumConcurrency;
+    private readonly SemaphoreSlim _workAvailable = new(0);
+    private readonly CancellationTokenSource _disposeCancellation = new();
+    private int _pendingTaskCount;
+    private int _currentThreads;
+    private int _threadId;
+
     public TaskThreadPool(int minimumThreads, Action<Task> taskExecutor)
     {
         if (minimumThreads <= 0)
@@ -30,23 +37,28 @@ internal sealed class TaskThreadPool
         }
 
         _taskExecutor = taskExecutor ?? throw new ArgumentNullException(nameof(taskExecutor));
-        _pendingTasks = new BlockingCollection<WorkItem>(_pendingTaskQueue);
+        _maximumConcurrency = minimumThreads;
-        _threads = new Thread[minimumThreads];
 
-        for (int i = 0; i < minimumThreads; i++)
+        List<Thread> threadsToStart = new(minimumThreads);
+
+        lock (_lock)
         {
-            Thread thread = new(WorkerLoop)
+            for (int i = 0; i < minimumThreads; i++)
             {
-                IsBackground = true,
+                Thread thread = CreateWorkerThread();
-                Name = $"{nameof(TaskThreadPool)}-{i + 1}"
+                _runningThreads.Add(thread);
-            };
+                _currentThreads++;
+                threadsToStart.Add(thread);
+            }
+        }
 
-            _threads[i] = thread;
+        foreach (Thread thread in threadsToStart)
+        {
             thread.Start();
         }
     }
 
-    public int ThreadCount => _threads.Length;
+    public int ThreadCount => _maximumConcurrency;
 
     public void Queue(Task task)
     {
@@ -55,18 +67,29 @@ internal sealed class TaskThreadPool
             throw new ArgumentNullException(nameof(task));
         }
 
-        try
+        List<Thread>? threadsToStart;
+
+        lock (_lock)
         {
-            _pendingTasks.Add(new WorkItem(task));
+            if (_disposeCancellation.IsCancellationRequested)
-        }
-        catch (Exception ex)
-        {
-            if (ex is InvalidOperationException || ex is ObjectDisposedException)
             {
-                throw new ObjectDisposedException(nameof(TaskThreadPool), ex);
+                throw new ObjectDisposedException(nameof(TaskThreadPool));
             }
 
-            throw;
+            _pendingTaskQueue.Enqueue(new WorkItem(task));
+            _pendingTaskCount++;
+
+            threadsToStart = EnsureWorkerCapacityLocked(null);
+        }
+
+        _workAvailable.Release();
+
+        if (threadsToStart is not null)
+        {
+            foreach (Thread thread in threadsToStart)
+            {
+                thread.Start();
+            }
         }
     }
 
@@ -91,28 +114,132 @@ internal sealed class TaskThreadPool
 
     public void Dispose()
     {
-        _pendingTasks.CompleteAdding();
+        Thread[] threadsToJoin;
 
-        foreach (Thread thread in _threads)
+        lock (_lock)
+        {
+            _disposeCancellation.Cancel();
+
+            threadsToJoin = new Thread[_runningThreads.Count];
+            _runningThreads.CopyTo(threadsToJoin);
+        }
+
+        if (threadsToJoin.Length > 0)
+        {
+            _workAvailable.Release(threadsToJoin.Length);
+        }
+
+        foreach (Thread thread in threadsToJoin)
         {
             thread.Join();
         }
 
-        _pendingTasks.Dispose();
+        _workAvailable.Dispose();
+        _disposeCancellation.Dispose();
+
+        lock (_lock)
+        {
+            _runningThreads.Clear();
+
+            while (_pendingTaskQueue.TryDequeue(out _))
+            {
+            }
+
+            _pendingTaskCount = 0;
+            _currentThreads = 0;
+        }
+    }
+
+    private bool TryDequeueWork(out WorkItem workItem)
+    {
+        CancellationToken cancellationToken = _disposeCancellation.Token;
+
+        while (true)
+        {
+            try
+            {
+                _workAvailable.Wait(cancellationToken);
+            }
+            catch (OperationCanceledException)
+            {
+                workItem = default;
+                return false;
+            }
+
+            if (_pendingTaskQueue.TryDequeue(out workItem))
+            {
+                Interlocked.Decrement(ref _pendingTaskCount);
+                return true;
+            }
+        }
     }
 
     private void WorkerLoop()
     {
         try
         {
-            foreach (WorkItem workItem in _pendingTasks.GetConsumingEnumerable())
+            if (!TryDequeueWork(out WorkItem workItem))
             {
-                _taskExecutor(workItem.Task);
+                return;
             }
+
+            _taskExecutor(workItem.Task);
         }
-        catch (ObjectDisposedException)
+        finally
         {
-            // The collection was disposed while the worker was waiting. Exit gracefully.
+            CompleteWorker();
         }
     }
+
+    private void CompleteWorker()
+    {
+        List<Thread>? threadsToStart = null;
+
+        lock (_lock)
+        {
+            _runningThreads.Remove(Thread.CurrentThread);
+            _currentThreads--;
+
+            if (!_disposeCancellation.IsCancellationRequested)
+            {
+                threadsToStart = EnsureWorkerCapacityLocked(null);
+            }
+        }
+
+        if (threadsToStart is not null)
+        {
+            foreach (Thread thread in threadsToStart)
+            {
+                thread.Start();
+            }
+        }
+    }
+
+    private List<Thread>? EnsureWorkerCapacityLocked(List<Thread>? threadsToStart)
+    {
+        int pendingTasks = Math.Min(_maximumConcurrency, Volatile.Read(ref _pendingTaskCount));
+
+        while (!_disposeCancellation.IsCancellationRequested && _currentThreads < pendingTasks)
+        {
+            Thread thread = CreateWorkerThread();
+            _runningThreads.Add(thread);
+            _currentThreads++;
+
+            threadsToStart ??= new List<Thread>();
+            threadsToStart.Add(thread);
+        }
+
+        return threadsToStart;
+    }
+
+    private Thread CreateWorkerThread()
+    {
+        int threadIndex = Interlocked.Increment(ref _threadId);
+
+        return new Thread(WorkerLoop)
+        {
+            IsBackground = true,
+            Name = $"{nameof(TaskThreadPool)}-{threadIndex}"
+        };
+    }
 }