Testing asynchronous background workers in .NET

Posted on Tue 08 May 2012 in Coding

When you build a GUI, all lengthy operations that can be triggered by the user should take place on a background thread so that the GUI doesn’t become unresponsive. Why would it? Well, it’s because all GUI operations take place on a single thread - the GUI thread (the Event-Dispatch Thread in Java). At the core of the GUI thread is a message pump that dispatches messages such as mouse and keyboard input to various controls. If you perform a lengthy operation on the GUI thread in response to a click message, for example, the message pump will “stand still” until the operation is complete. During this time, other messages just sit in the queue, which effectively renders the GUI inoperable. Instead, you’d spawn off a background thread to do the work, and let that thread post back the result to the GUI thread when done.

Using the Task Parallel Library in .NET, creating a background thread to do some work is very simple:

private void DoWorkInBackground() {
    Task<string>.Factory.StartNew(DoWork);
}

private string DoWork() {
    // perform some lengthy operation, then return a result
}

But how do we get to the result? The StartNew method returns a Task instance, and we can query its Result property to get the result. The problem is that the calling thread blocks until the result is available, and blocking the GUI thread was exactly what we wanted to avoid. So let’s try a different approach, utilizing the facts that the GUI thread has its own synchronization context, that a task can have a continuation task, and that we can run such a continuation task with a specific task scheduler:

private void DoWorkInBackground(Action<string> callback) {
    var taskScheduler = TaskScheduler.FromCurrentSynchronizationContext();
    var task = new Task<string>(DoWork);
    task.ContinueWith(t => callback(t.Result), taskScheduler);
    task.Start();
}

private string DoWork() {
    // perform some lengthy operation, then return a result
}

As you can see here, the continuation task runs using a task scheduler created from the current synchronization context. This works because the synchronization context for the GUI thread (this or this) posts to the message pump a special message that, when dispatched, runs the task code.

What I’ve just described is a very common way of doing background work. So common that Microsoft has chosen to simplify it a lot with the new Async API (which will be a part of .NET 4.5). With the new async and await keywords, we can just do:

private async void DoWorkInBackground(Action<string> callback) {
    var task = Task<string>.Factory.StartNew(DoWork);
    callback(await task);
}

private string DoWork() {
    // perform some lengthy operation, then return a result
}

This is especially useful if you do multiple asynchronous calls with some GUI operations in-between (e.g., pausing the background work to ask for some user input).

Alright, so how do we test asynchronous code like the one above? Well, we could (and should) break out the “work” to a separate entity and test that entity synchronously. While that’s a good start, perhaps we still want to test the entity that chooses to run code in the background in the first place - most likely an MVC controller class. The answer is that the test code has to wait until the asynchronous code has finished running. But what should it wait for? Under our background worker model, there’s some continuation code that runs on the GUI thread by means of its message pump. But during a test run, the original thread is the test runner thread, which will be stuck waiting for the background worker to complete. Can we run a message pump in the test runner thread? We probably could, but luckily we don’t have to! All we need is a simulation of one, using a custom synchronization context! Next, I’ll show you how to write one.

Let’s start by studying how the custom synchronization context can be used. To be able to do TDD, we do need to specify a timeout for the wait. While we could put the timeout on the test method itself, I think specifying the timeout when the context is created is slightly more readable.

[TestMethod]
public void TestTheOutcomeOfSomeBackgroundOperation() {
    var sut = ...;
    var tester = AsyncTester.WithTimeout(200);
    sut.SomeAction();
    tester.Uninstall();
    Assert.AreEqual(...);
}

The factory method WithTimeout creates and automatically installs the context. After having triggered the background operation, we uninstall the context. It’s a bit annoying that uninstallation is manual, and it also opens up for some nasty side-effects if the action fails with an exception and the context never gets uninstalled. A using block is much nicer:

[TestMethod]
public void TestTheOutcomeOfSomeBackgroundOperation() {
    var sut = ...;
    using (AsyncTester.WithTimeout(200)) {
        sut.SomeAction();
    }
    Assert.AreEqual(...);
}

Ok, once we have the usage clear, let’s define the class! Starting with the mechanics that allow installation and uninstallation as we’ve seen above:

public class AsyncTester : SynchronizationContext, IDisposable {
    public static AsyncTester WithTimeout(int timeoutMs) {
        var sc = new AsyncTester(timeoutMs, Current);
        SetSynchronizationContext(sc);
        return sc;
    }

    public void Uninstall() {
        SetSynchronizationContext(_prevSyncContext);
    }

    public void Dispose() {
        Uninstall();
    }

The WithTimeout method creates and sets the synchronization context. The Uninstall method restores the previous context. The Dispose method is required by the IDisposable interface and simply delegates the uninstallation. We need a constructor and some fields to support the above (the timeout comes into play later):

    private readonly int _timeoutMs;
    private readonly SynchronizationContext _prevSyncContext;

    private AsyncTester(int timeoutMs, SynchronizationContext prevContext) {
        _timeoutMs = timeoutMs;
        _prevSyncContext = prevContext;
    }

I’ve chosen to make the constructor private for two reasons. First, it forces the use of the factory method, which IMHO results in more readable code. Second, I don’t want to expose the fact that we keep track of the previous context.

Next, we need to take care of the actual juice of the synchronization context - the code that waits for whatever background operation is going on to complete. Recall that when the background task is finished, there is a continuation task to be run on the original thread. That task is posted by a task scheduler created from the current synchronization context. Posted to what? To the synchronization context, of course! So let’s override the Post method:

    private readonly BlockingCollection<Tuple<SendOrPostCallback, object>> _queue = 
        new BlockingCollection<Tuple<SendOrPostCallback, object>>();

    public override void Post(SendOrPostCallback d, object state) {
        _queue.Add(Tuple.Create(d, state));
    }

I’ve chose to use a BlockingCollection to queue callbacks that correspond to continuation tasks, since it is thread safe and the Post method is called from a background thread. Note that no execution happens here. That is where the next piece of the puzzle comes in, the method that consumes the queue. This is our message pump simulation!

    public void ProcessQueue(int timeoutMs = 0) {
        // Use the constructor-specified timeout if not overridden.
        if (timeoutMs <= 0)
            timeoutMs = _timeoutMs;
        var sw = Stopwatch.StartNew();
        int timeLeft;
        while (true) {
            timeLeft = (int)(timeoutMs - sw.ElapsedMilliseconds);
            if (timeLeft <= 0)
                break; // we're done waiting
            Tuple<SendOrPostCallback, object> tuple;
            if (_queue.TryTake(out tuple, timeLeft)) {
                // Item1 = the callback, Item2 = the state object.
                tuple.Item1(tuple.Item2);
            }
        }
    }

The method is pretty straight-forward; we use a Stopwatch instance to keep track of elapsed time. For each iteration of the loop (there will be many in the case of multiple background workers) we calculate how long time we can try to take a previously added pair of callback and state object. The callback has been created by the task scheduler to run the code in the continuation task.

You may have reacted on the fact that the method is public, and that it takes an optional timeout. I’ll come to these in a moment. First, though, we need to call the method from somewhere. Let’s modify the Uninstall method:

    public void Uninstall() {
        if (this != Current) return;
        ProcessQueue();
        SetSynchronizationContext(_prevSyncContext);
    }

The modification means that when the context is uninstalled, all queued continuation tasks are run. There is also a small safeguard to prevent multiple uninstallations (which is a real problem since the context instance eventually will be disposed, and Dispose calls Uninstall).

What about the public ProcessQueue method and the local timeout? Well, I’ve had some cases where the context setup (the given) involved background work as well. In such cases, we can do:

[TestMethod]
public void SomeTest() {
    // context setup...
    using (var tester = AsyncTester.WithTimeout(300)) {
        // more context setup, including background work...
        tester.ProcessQueue(500);
        // call code that uses a background worker...
    }
    // assert result...
}

The local timeout is unrelated to the the overall timeout. Both affect the running time of the test - 800 milliseconds in the test method above.

That’s it! I hope that you find the class useful! You can find it in its entirety in this gist. Actually, there is one more thing to make it complete: :-)

}