A Custom TaskScheduler for STA-Threads

Setting the apartment state of a thread to STA is easy when you create it yourself. Using TPL-Tasks, a custom TaskScheduler is a nice solution.

Apartments (not suitable for Threads) near San Francisco
Apartments (not suitable for Threads) near San Francisco

What is the apartment state of a thread and why should I care?

When dealing with COM-components in multi threaded environments, sometimes strange exceptions occur. Often enough, a simple fix lies in changing the so-called “Apartment State” of the thread which is using the component. Most of the times, you will need to change it to ApartmentState.STA indicating a single-threaded apartment. For more information on apartment states check out Could you explain STA and MTA?, ApartmentState for dummies or for more details Understanding The COM Single-Threaded Apartment Part 1.

Setting the apartment state of a thread

Usually, the change of the apartment state can be achieved by adding a STAThreadAttribute to the main method of your application. However, this doesn’t work in multi-threaded environments as Threads are created by the ThreadPool or by the TPL (Task Parallel Library). Manually creating the Thread is always an option as you have the possibility to set the apartment state of a Thread during its creation time:

var thread = new Thread(this.ThreadStart);
thread.TrySetApartmentState(ApartmentState.STA);
thread.Start();

Sticking with the comfort of the TPL, using the TaskScheduler.FromCurrentSynchronizationContex method can be a simple solution to execute the code in the UI thread instead which always has to be a STA-Thread. Nevertheless, there is one major downside to this approach: Especially with long-running operations, you can easily block the UI thread. Thus, the wish arises to keep the application responsive without sacrificing the comfort of using System.Threading.Tasks.Tasks.

The SingleThreadTaskScheduler

A very nice solution is writing a custom TaskScheduler automatically scheduling Tasks in the same Thread with a defined apartment state. What might sound like a difficult task at first turned out to be quite easy in the end. The result is the SingleThreadTaskScheduler class.

When implementing a custom task scheduler, we have to inherit from the abstract class TaskScheduler and overwrite three methods:

  • void QueueTask(Task task) which is called whenever a Task should be scheduled on the scheduler.
  • bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued) to support “inlining” of Tasks meaning to execute them without queuing. This is optional and the method can simply return false .
  • IEnumerable<Task> GetScheduledTasks() which is used for debugging purposes and should return a list of all scheduled tasks

In order to execute a Task, it can be passed to the TryExecuteTask(task); method provided by TaskScheduler .

SingleThreadTaskScheduler implements those methods by putting all scheduled tasks into a thread-safe System.Collections.Concurrent.BlockingCollection<Task> and starting a single, dedicated Thread to process them:

public SingleThreadTaskScheduler(Action initAction, ApartmentState apartmentState = ApartmentState.STA)
{
    //...
    this._tasks = new BlockingCollection<Task>();

    this._thread = new Thread(this.ThreadStart);
    this._thread.IsBackground = true;
    this._thread.TrySetApartmentState(apartmentState);
    this._thread.Start();
}


protected override void QueueTask(Task task)
{
    this.VerifyNotDisposed();

    this._tasks.Add(task, this._cancellationToken.Token);
}
private void ThreadStart()
{
    try
    {
        var token = this._cancellationToken.Token;

        this._initAction();

        foreach (var task in this._tasks.GetConsumingEnumerable(token))
            this.TryExecuteTask(task);
    }
    finally
    {
        this._tasks.Dispose();
    }
}

In addition, inlining is done when the Thread trying to add a Task equals the Thread which is used to process the Tasks. As our task scheduler only uses one Thread, providing this capability is very important. Otherwise, a deadlock could be caused by submitting a Task from within a currently executing Task and waiting for the result (causing the Thread to be blocked):

protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued)
{
    this.VerifyNotDisposed();

    if (this._thread != Thread.CurrentThread)
        return false;
    if (this._cancellationToken.Token.IsCancellationRequested)
        return false;

    this.TryExecuteTask(task);
    return true;
}

To verify the correctness of the implementation, eighteen unit tests are provided which test various aspects of the class.

Usage

The SingleThreadTaskScheduler can be used like any other task scheduler: When starting a new Task using Task.Factory.StartNew(…) or using task.ContinueWith(…), the task scheduler to use is one of the possible arguments. For maximum convenience, the library also includes two extension methods which allow you to write:

var taskScheduler = //...
Task.Factory.StartNew(() => /*Do something*/, taskScheduler);
//respectively:
var result = await Task.Factory.StartNew(() => return /*Some value*/, taskScheduler);

Note the usage of await (only possible within methods marked as async ) do retrieve the result without blocking the calling thread. For example, this can allow you to run all interaction with a specific COM component on a dedicated Thread (provided by the SingleThreadTaskScheduler) to avoid blocking of the UI while reducing the additional clutter in the code.

Download

The source code including unit tests is available at GitHub (direct link to the source file).
You can also add the package SingleThreadScheduler from my NuGet feed or simply download the binaries.

  • Martin M.

    Vielen Dank für den Code. You safed my day!

  • VURSO

    Nice work man saved me as well. I used your classes (from GitHub) to help me run some TPL code that was starting a web browser object programmatically to read in a site from a url, convert it to an image and back out in byte array format but I could get it to work in a Web API 2 project.

    Using the SingleThreadScheduler I was able to pass this into the TPL setup and get it working!