/Code/Source/Teaching.Core/Threading/StaTaskScheduler.cs

using System;

using System.Collections.Concurrent;

using System.Collections.Generic;

using System.Diagnostics.Contracts;

using System.Linq;

using System.Threading;

using System.Threading.Tasks;



namespace Teaching.Core.Threading

{

    public sealed class StaTaskScheduler : TaskScheduler, IDisposable

    {

        private BlockingCollection<Task>    tasks = new BlockingCollection<Task>();

        private readonly List<Thread>       threads;



        public override int MaximumConcurrencyLevel

        {

            get { return threads.Count; }

        }



        protected override void QueueTask(Task task)

        {

            tasks.Add(task);

        }



        protected override IEnumerable<Task> GetScheduledTasks()

        {

            return tasks.ToArray();

        }



        protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued)

        {

            return

                Thread.CurrentThread.GetApartmentState() == ApartmentState.STA &&

                TryExecuteTask(task);

        }



        public void Dispose()

        {

            if (tasks != null)

            {

                tasks.CompleteAdding();



                foreach (var thread in threads)

                {

                    thread.Join();

                }



                tasks.Dispose();

                tasks = null;

            }

        }



        public StaTaskScheduler(int numberOfThreads)

        {

            Contract.Requires<ArgumentOutOfRangeException>(numberOfThreads > 0,"numberOfThreads");



            threads = Enumerable.Range(0, numberOfThreads).Select(i =>

            {

                var thread = new Thread(() =>

                {

                    foreach (var t in tasks.GetConsumingEnumerable())

                    {

                        TryExecuteTask(t);

                    }

                }) {IsBackground = true};

                thread.SetApartmentState(ApartmentState.STA);

                return thread;

            }).ToList();



            threads.ForEach(t => t.Start());

        }

    } 

}