Wake up!

In Effective Unit Testing: A guide for Java developers, author Lasse Koskela writes about the scourge of the sleeping snail — slow tests littered with calls to Thread.sleep. His suggested alternative is perfectly sound and practical: turn nondeterministic waiting into explicit synchronization with help from primitives like a CountDownLatch (that’s CountdownEvent for us .NET programmers).

There is however another way that should be used when possible: avoid background threads completely and bring the action right to you. Wake up and take part!

Consider an example of an “expiring cache.” Items in the cache have a certain time to live after which they are evicted. A naïve way to implement this would involve direct use of a Timer and some “sleepy” unit tests to verify the implementation, like so:

public sealed class ExpiringCache<TKey, TValue>
{
    private readonly ConcurrentDictionary<TKey, ExpiringValue> items;

    public ExpiringCache()
    {
        this.items = new ConcurrentDictionary<TKey, ExpiringValue>();
    }

    public void Add(TKey key, TValue value, TimeSpan timeToLive)
    {
        this.items.AddOrUpdate(key, k => new ExpiringValue(this, k, value, timeToLive), (k, v) => v);
    }

    public bool TryGet(TKey key, out TValue value)
    {
        ExpiringValue expiringValue;
        if (!this.items.TryGetValue(key, out expiringValue))
        {
            value = default(TValue);
            return false;
        }

        value = expiringValue.Value;
        return true;
    }

    private void Remove(TKey key)
    {
        ExpiringValue value;
        this.items.TryRemove(key, out value);
    }

    private sealed class ExpiringValue
    {
        private readonly ExpiringCache<TKey, TValue> parent;
        private readonly TKey key;
        private readonly Timer timer;

        public ExpiringValue(ExpiringCache<TKey, TValue> parent, TKey key, TValue value, TimeSpan timeToLive)
        {
            this.parent = parent;
            this.key = key;
            this.Value = value;
            this.timer = new Timer(Expire, this, Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan);
            this.timer.Change(timeToLive, Timeout.InfiniteTimeSpan);
        }

        public TValue Value { get; private set; }

        private static void Expire(object thisObj)
        {
            ((ExpiringValue)thisObj).Expire();
        }

        private void Expire()
        {
            this.parent.Remove(this.key);
            this.timer.Dispose();
        }
    }
}

[TestClass]
public class ExpiringCacheTest
{
    [TestMethod]
    public void AddedItemShouldStillExistBeforeExpirationTime()
    {
        ExpiringCache<int, string> cache = new ExpiringCache<int, string>();
        cache.Add(1, "two", TimeSpan.FromSeconds(1.0d));

        string value;
        Assert.IsTrue(cache.TryGet(1, out value));
        Assert.AreEqual("two", value);
    }

    [TestMethod]
    public void AddedItemShouldNotExistAfterExpirationTime()
    {
        ExpiringCache<int, string> cache = new ExpiringCache<int, string>();
        cache.Add(1, "two", TimeSpan.FromSeconds(0.5d));

        Thread.Sleep(1000);

        string value;
        Assert.IsFalse(cache.TryGet(1, out value));
    }

    [TestMethod]
    public void UpdatedItemShouldNotUpdateValue()
    {
        ExpiringCache<int, string> cache = new ExpiringCache<int, string>();
        cache.Add(1, "two", TimeSpan.FromSeconds(0.5d));
        cache.Add(1, "three", TimeSpan.FromSeconds(0.5d));

        string value;
        Assert.IsTrue(cache.TryGet(1, out value));
        Assert.AreEqual("two", value);
    }

    [TestMethod]
    public void UpdatedItemShouldNotUpdateExpirationTime()
    {
        ExpiringCache<int, string> cache = new ExpiringCache<int, string>();
        cache.Add(1, "two", TimeSpan.FromSeconds(0.5d));
        cache.Add(1, "three", TimeSpan.FromSeconds(4.0d));

        Thread.Sleep(1000);

        string value;
        Assert.IsFalse(cache.TryGet(1, out value));
    }
}

This works (well… most of the time), but it is enormously slow in unit testing terms. The two seconds spent doing absolutely nothing is enough time to run thousands of other fast unit tests. If we instead take a page from the async playbook we can gain control of the fabric of time, so to speak. How?

First, let’s replace the timer with a strategically placed Task.Delay:

private sealed class ExpiringValue
{
    private readonly ExpiringCache<TKey, TValue> parent;
    private readonly TKey key;
    private readonly Task expired;

    public ExpiringValue(ExpiringCache<TKey, TValue> parent, TKey key, TValue value, TimeSpan timeToLive)
    {
        this.parent = parent;
        this.key = key;
        this.Value = value;
        this.expired = Task.Delay(timeToLive).ContinueWith(this.Expire, TaskContinuationOptions.ExecuteSynchronously);
    }

    public TValue Value { get; private set; }

    private void Expire(Task task)
    {
        this.parent.Remove(this.key);
    }
}

Nothing much has changed, but this gives us the beginning of an important test seam. It turns out that Task’s role as a promise/future is useful for more than non-blocking I/O calls. We just need to go one step further and return a Task whose lifetime is under our complete control, and in essence we will become the timekeeper with no need for sleep!

Keeping it simple, we can add a pluggable delay delegate with the same signature as Task.Delay and override the implementation for testing purposes:

public sealed class ExpiringCache<TKey, TValue>
{
    private static readonly Func<Task> DefaultDelay = Task.Delay;
    // ...
    public ExpiringCache()
    {
        // ...
        this.Delay = DefaultDelay;
    }

    public Func<TimeSpan, Task> Delay { get; set; }
    // ...
    public ExpiringValue(ExpiringCache<TKey, TValue> parent, TKey key, TValue value, TimeSpan timeToLive)
    {
        // ...
        this.expired = parent.Delay(timeToLive).ContinueWith(this.Expire, TaskContinuationOptions.ExecuteSynchronously);
    }
    // ...
}

With all tests still passing, we can now attack the sleeping giant. First, we’ll define a test class to simulate a timer event that fires at some future time (defined as an inner class of its only user, ExpiringCacheTest):

private sealed class AlarmClock
{
    private TaskCompletionSource<bool> tcs;
    private TimeSpan dueTime;

    public AlarmClock()
    {
    }

    public Task Set(TimeSpan dueTime)
    {
        this.tcs = new TaskCompletionSource<bool>();
        this.dueTime = dueTime;
        return this.tcs.Task;
    }

    public TimeSpan Wake()
    {
        this.tcs.SetResult(false);
        return this.dueTime;
    }
}

Now let’s fix the implementation of our two long running tests:

[TestMethod]
public void AddedItemShouldNotExistAfterExpirationTime()
{
    AlarmClock clock = new AlarmClock(); 
    ExpiringCache<int, string> cache = new ExpiringCache<int, string>()
    {
        Delay = clock.Set
    };

    cache.Add(1, "two", TimeSpan.FromSeconds(0.5d));

    TimeSpan dueTime = clock.Wake();
    Assert.AreEqual(TimeSpan.FromSeconds(0.5d), dueTime);

    string value;
    Assert.IsFalse(cache.TryGet(1, out value));
}
// ...
[TestMethod]
public void UpdatedItemShouldNotUpdateExpirationTime()
{
    AlarmClock clock = new AlarmClock();
    ExpiringCache<int, string> cache = new ExpiringCache<int, string>()
    {
        Delay = clock.Set
    }; 

    cache.Add(1, "two", TimeSpan.FromSeconds(0.5d));
    cache.Add(1, "three", TimeSpan.FromSeconds(4.0d));

    TimeSpan dueTime = clock.Wake();
    Assert.AreEqual(TimeSpan.FromSeconds(0.5d), dueTime);

    string value;
    Assert.IsFalse(cache.TryGet(1, out value));
}

Basically, the AlarmClock keeps track of what time was set and allows manual control of when we awake. Not only have we eliminated the sleeps, we have eliminated background threads entirely. The call to TaskCompletionSource.SetResult runs the continuation which expires the cache item synchronously because we strategically specified TaskContinuationOptions.ExecuteSynchronously above.

This approach is simple and noninvasive enough that it leaves no good excuse for sleeping on the job. Try it out and kick your unit test suites into high gear!