TDD + async: Introducing MemoryChannel

Can you use test-driven development and unit testing with asynchronous code? Yes, with a few caveats. Above all, unit tests should be deterministic and fast. This means you should stick to single-threaded workflows as much as possible and never sleep.

To demonstrate these principles along with typical .NET 4.5 asynchronous patterns, I came up with a useful and complete example — MemoryChannel. MemoryChannel is a simple one-way communication channel which allows a sender to pass raw byte buffers to an asynchronous receiver. With some appropriate adapters, it could replace a named pipe or TCP socket (thus helping to eliminating external dependencies, for example, in a set of unit tests).

Before writing the code, I sketched a crude FSM diagram on a piece of paper. Here is a much more refined version of that diagram (courtesy of PowerPoint, my diagramming tool of choice):

There are only four states (no data, data available, receive pending, and disposed) and three operations (send, receive, dispose). The only real complications here are the conditions relating to the sent and received data sizes. For example, if there is some data in the channel and a receiver requests more data than is available, the request should still complete but notify the receiver that only S bytes were transferred. However, if less data is requested (R), the excess (S - R) should remain available for later receive attempts.

In true TDD style, I started with the minimum possible interface and decided to focus on the receive side of the state machine.

public class MemoryChannel
{
    public MemoryChannel()
    {
    }

    public Task ReceiveAsync(byte[] buffer)
    {
        throw new NotImplementedException();
    }

    public void Send(byte[] buffer)
    {
        throw new NotImplementedException();
    }
}

Here is the first unit test, covering just the highlighted part of the state machine:

public void Pending_receive_completes_after_send_with_same_data_size()
{
    MemoryChannel channel = new MemoryChannel();

    byte[] receiveBuffer = new byte[3];
    Task<int> receiveTask = channel.ReceiveAsync(receiveBuffer);

    Assert.False(receiveTask.IsCompleted);
    Assert.False(receiveTask.IsFaulted);

    byte[] sendBuffer = new byte[] { 1, 2, 3 };
    channel.Send(sendBuffer);

    Assert.Equal(TaskStatus.RanToCompletion, receiveTask.Status);
    Assert.Equal(3, receiveTask.Result);
    Assert.Equal(new byte[] { 1, 2, 3 }, receiveBuffer);
}

To fulfill the conditions of this test, we need to return a Task<int> from ReceiveAsync which has not yet completed and only does so when a matching Send call is made. Additionally, to satisfy the unit testing principles above we need to do this without external threads. How can we do this?

Enter TaskCompletionSource — one of the most useful library classes in the async .NET toolbox. Essentially, TCS lets us hand out a Task representing a custom asynchronous operation while giving us the ability to explicitly control the lifetime of this operation.

To make the test above pass, we can use TCS to represent the pending receive operation like so, initializing it on start of receive and completing it (via SetResult) when Send is invoked:

public class MemoryChannel
{
    private TaskCompletionSource<int> pendingReceive;
    private byte[] pendingReceiveBuffer;

    public MemoryChannel()
    { 
    }

    public Task<int> ReceiveAsync(byte[] buffer)
    {
        this.pendingReceive = new TaskCompletionSource<int>();
        this.pendingReceiveBuffer = buffer;
        return this.pendingReceive.Task;
    }
 
    public void Send(byte[] buffer)
    {
        buffer.CopyTo(this.pendingReceiveBuffer, 0);
        this.pendingReceive.SetResult(buffer.Length);
    }
}

Note that SetResult will transition the task to a completed state as well run any synchronous continuations. This can be undesirable in some situations which I will get into in a later post.

You can track the evolution of MemoryChannel into its full implementation in my CommSample project. The commit history shows the changes following every new unit test or refactoring.