using System; using System.Collections.Generic; using System.Linq; using System.Threading; using System.Threading.Tasks; namespace NzbDrone.Common.TPL { public class LimitedConcurrencyLevelTaskScheduler : TaskScheduler { /// Whether the current thread is processing work items. [ThreadStatic] private static bool _currentThreadIsProcessingItems; /// The list of tasks to be executed. private readonly LinkedList _tasks = new LinkedList(); /// The maximum concurrency level allowed by this scheduler. private readonly int _maxDegreeOfParallelism; /// Whether the scheduler is currently processing work items. private int _delegatesQueuedOrRunning; /// /// Initializes an instance of the LimitedConcurrencyLevelTaskScheduler class with the /// specified degree of parallelism. /// /// The maximum degree of parallelism provided by this scheduler. public LimitedConcurrencyLevelTaskScheduler(int maxDegreeOfParallelism) { if (maxDegreeOfParallelism < 1) { throw new ArgumentOutOfRangeException("maxDegreeOfParallelism"); } _maxDegreeOfParallelism = maxDegreeOfParallelism; } /// Queues a task to the scheduler. /// The task to be queued. protected sealed override void QueueTask(Task task) { // Add the task to the list of tasks to be processed. If there aren't enough // delegates currently queued or running to process tasks, schedule another. lock (_tasks) { _tasks.AddLast(task); if (_delegatesQueuedOrRunning < _maxDegreeOfParallelism) { ++_delegatesQueuedOrRunning; NotifyThreadPoolOfPendingWork(); } } } /// /// Informs the ThreadPool that there's work to be executed for this scheduler. /// private void NotifyThreadPoolOfPendingWork() { ThreadPool.UnsafeQueueUserWorkItem(_ => { // Note that the current thread is now processing work items. // This is necessary to enable inlining of tasks into this thread. _currentThreadIsProcessingItems = true; try { // Process all available items in the queue. while (true) { Task item; lock (_tasks) { // When there are no more items to be processed, // note that we're done processing, and get out. if (_tasks.Count == 0) { --_delegatesQueuedOrRunning; break; } // Get the next item from the queue item = _tasks.First.Value; _tasks.RemoveFirst(); } // Execute the task we pulled out of the queue TryExecuteTask(item); } } // We're done processing items on the current thread finally { _currentThreadIsProcessingItems = false; } }, null); } /// Attempts to execute the specified task on the current thread. /// The task to be executed. /// /// Whether the task could be executed on the current thread. protected sealed override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued) { // If this thread isn't already processing a task, we don't support inlining if (!_currentThreadIsProcessingItems) { return false; } // If the task was previously queued, remove it from the queue if (taskWasPreviouslyQueued) { TryDequeue(task); } // Try to run the task. return TryExecuteTask(task); } /// Attempts to remove a previously scheduled task from the scheduler. /// The task to be removed. /// Whether the task could be found and removed. protected sealed override bool TryDequeue(Task task) { lock (_tasks) { return _tasks.Remove(task); } } /// Gets the maximum concurrency level supported by this scheduler. public sealed override int MaximumConcurrencyLevel => _maxDegreeOfParallelism; /// Gets an enumerable of the tasks currently scheduled on this scheduler. /// An enumerable of the tasks currently scheduled. protected sealed override IEnumerable GetScheduledTasks() { var lockTaken = false; try { Monitor.TryEnter(_tasks, ref lockTaken); if (lockTaken) { return _tasks.ToArray(); } else { throw new NotSupportedException(); } } finally { if (lockTaken) { Monitor.Exit(_tasks); } } } } }