Sourcetrail/ide_plugins/vs/vs2015/SourcetrailPlugin/SourcetrailPlugin/Multitasking/LimitedThreadsTaskScheduler.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

// based on: https://msdn.microsoft.com/en-us/library/system.threading.tasks.taskscheduler(v=vs.110).aspx

namespace CoatiSoftware.SourcetrailPlugin.Multitasking
{
	class LimitedThreadsTaskScheduler : TaskScheduler
	{
		[ThreadStatic]
		private static bool _currentThreadIsProcessingItems;

		private readonly LinkedList<Task> _tasks = new LinkedList<Task>();

		private readonly int _maxNumberOfRunningThreads = 0;

		private int _delegatesQueuedOrRunning = 0;

		public LimitedThreadsTaskScheduler(int maxNumberOfRunningThreads)
		{
			if (maxNumberOfRunningThreads < 1)
			{
				maxNumberOfRunningThreads = 1;
			}

			_maxNumberOfRunningThreads = maxNumberOfRunningThreads;
		}

		protected override IEnumerable<Task> GetScheduledTasks()
		{
			bool lockTaken = false;

			try
			{
				Monitor.TryEnter(_tasks, ref lockTaken);
				if(lockTaken)
				{
					return _tasks;
				}
				else
				{
					throw new NotSupportedException();
				}
			}
			finally
			{
				if (lockTaken)
				{
					Monitor.Exit(_tasks);
				}
			}
		}

		protected override void QueueTask(Task task)
		{
			lock(_tasks)
			{
				_tasks.AddLast(task);
				
				if(_delegatesQueuedOrRunning < _maxNumberOfRunningThreads)
				{
					++_delegatesQueuedOrRunning;
					NotifyThreadPoolOfPendingWork();
				}
			}
		}

		protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued)
		{
			if(!_currentThreadIsProcessingItems)
			{
				return false;
			}

			if(taskWasPreviouslyQueued)
			{
				if(TryDequeue(task))
				{
					return base.TryExecuteTask(task);
				}
				else
				{
					return false;
				}
			}
			else
			{
				return base.TryExecuteTask(task);
			}
		}

		private void NotifyThreadPoolOfPendingWork()
		{
			ThreadPool.UnsafeQueueUserWorkItem(_ =>
			{
				_currentThreadIsProcessingItems = true;
				try
				{
					while(true)
					{
						Task item;

						lock(_tasks)
						{
							if(_tasks.Count <= 0)
							{
								--_delegatesQueuedOrRunning;
								break;
							}

							item = _tasks.First.Value;
							_tasks.RemoveFirst();
						}

						base.TryExecuteTask(item);
					}
				}
				finally
				{
					_currentThreadIsProcessingItems = false;
				}

			}, null);
		}

		protected sealed override bool TryDequeue(Task task)
		{
			lock(_tasks)
			{
				return _tasks.Remove(task);
			}
		}
	}
}