Imagine that you have a Task handed to you by a third party, and that you would like to force this Task to complete within a specified time period. However, you cannot alter the “natural” completion path and completion state of the Task, as that may cause problems with other consumers of the Task. So you need a way to obtain a copy or “proxy” of the Task that will either (A) complete within the specified time period, or (B) will complete with an indication that it had timed out.

In this blog post, I will show how one might go about implementing a Task.TimeoutAfter method to support this scenario. The signature of TimeoutAfter would look like this:

public static Task TimeoutAfter(this Task task, int millisecondsTimeout)

The returned proxy Task can complete in one of two ways:

1. If task completes before the specified timeout period has elapsed, then the proxy Task finishes when task finishes, with task’s completion status being copied to the proxy.
2. If task fails to complete before the specified timeout period has elapsed, then the proxy Task finishes when the timeout period expires, in Faulted state with a TimeoutException.

In addition to showing how to implement Task.TimeoutAfter, this post will also shed some light on the general thought process that should go into implementing such a feature.

A First Try

Here’s some code that will do the trick:

internal struct VoidTypeStruct { }  // See Footnote #1

static class TaskExtensions
{
    public static Task TimeoutAfter(this Task task, int millisecondsTimeout)
    {
        // tcs.Task will be returned as a proxy to the caller
        TaskCompletionSource<VoidTypeStruct> tcs = 
            new TaskCompletionSource<VoidTypeStruct>();

        // Set up a timer to complete after the specified timeout period
        Timer timer = new Timer(_ => 
        {
            // Fault our proxy Task with a TimeoutException
            tcs.TrySetException(new TimeoutException()); 
        }, null, millisecondsTimeout, Timeout.Infinite);

        // Wire up the logic for what happens when source task completes
        task.ContinueWith(antecedent =>
        {
            timer.Dispose(); // Cancel the timer
            MarshalTaskResults(antecedent, tcs); // Marshal results to proxy
        }, CancellationToken.None, 
            TaskContinuationOptions.ExecuteSynchronously, 
            TaskScheduler.Default);

        return tcs.Task;
    }
}

Simple enough, right? You start a Timer job that faults the proxy Task, and also add a continuation off of the source Task that transfers the completion state of the source to the proxy. The final state of the proxy will therefore depend on which completes first, the Timer job or the source Task.

And by the way, MarshalTaskResults is implemented like this:

internal static void MarshalTaskResults<TResult>(
    Task source, TaskCompletionSource<TResult> proxy)
{
    switch (source.Status)
    {
        case TaskStatus.Faulted:
            proxy.TrySetException(source.Exception);
            break;
        case TaskStatus.Canceled:
            proxy.TrySetCanceled();
            break;
        case TaskStatus.RanToCompletion:
            Task<TResult> castedSource = source as Task<TResult>;
            proxy.TrySetResult(
                castedSource == null ? default(TResult) : // source is a Task
                    castedSource.Result); // source is a Task<TResult>
            break;
    }
}

he “RanToCompletion” handling might seem a little more complicated than it needs to be, but it will allow us to handle Task<TResult> objects correctly (discussed briefly below).

Can We Do Better?

While our first stab at a TimeoutAfter method is functionally correct, we could streamline it and improve its performance. Specifically, notice that our Timer and continuation delegates “capture” variables; this will cause the compiler to allocate special “closure” classes for these delegates behind the scenes, which will slow down our method. To eliminate the need for the closure class allocations, we can pass all “captured” variables in through state variables for those respective calls, like this:

public static Task TimeoutAfter(this Task task, int millisecondsTimeout)
{
    // tcs.Task will be returned as a proxy to the caller
    TaskCompletionSource<VoidTypeStruct> tcs = 
        new TaskCompletionSource<VoidTypeStruct>();

    // Set up a timer to complete after the specified timeout period
    Timer timer = new Timer(state => 
    { 
        // Recover our state data
        var myTcs = (TaskCompletionSource<VoidTypeStruct>)state;

        // Fault our proxy Task with a TimeoutException
        myTcs.TrySetException(new TimeoutException()); 
    }, tcs, millisecondsTimeout, Timeout.Infinite);

    // Wire up the logic for what happens when source task completes
    task.ContinueWith((antecedent,state) =>
    {
        // Recover our state data
        var tuple = 
            (Tuple<Timer, TaskCompletionSource<VoidTypeStruct>>)state;

        // Cancel the timer
        tuple.Item1.Dispose();
        // Marshal results to proxy
        MarshalTaskResults(antecedent, tuple.Item2);    
    }, 
    Tuple.Create(timer,tcs),  // See Footnote #2
    CancellationToken.None,
    TaskContinuationOptions.ExecuteSynchronously,
    TaskScheduler.Default);

    return tcs.Task;
}

Some ad-hoc performance tests show that this little optimization shaves about 12% off of the overhead from the TimeoutAfter method.

What about Edge Cases?

What do we do when the caller specifies a zero timeout, or an infinite timeout? What if the source Task has already completed by the time that we enter the TimeoutAfter method? We can address these edge cases in the TimeoutAfter implementation as follows:

public static Task TimeoutAfter(this Task task, int millisecondsTimeout)
{
    // Short-circuit #1: infinite timeout or task already completed
    if (task.IsCompleted || (millisecondsTimeout == Timeout.Infinite))
    {
        // Either the task has already completed or timeout will never occur.
        // No proxy necessary.
        return task;
    }

    // tcs.Task will be returned as a proxy to the caller
    TaskCompletionSource<VoidTypeStruct> tcs = 
        new TaskCompletionSource<VoidTypeStruct>();

    // Short-circuit #2: zero timeout
    if (millisecondsTimeout == 0)
    {
        // We've already timed out.
        tcs.SetException(new TimeoutException());
        return tcs.Task;
    }

    // Set up a timer to complete after the specified timeout period
    Timer timer = new Timer(state => 
    {
        // Recover your state information
        var myTcs = (TaskCompletionSource<VoidTypeStruct>)state;

        // Fault our proxy with a TimeoutException
        myTcs.TrySetException(new TimeoutException()); 
    }, tcs, millisecondsTimeout, Timeout.Infinite);

    // Wire up the logic for what happens when source task completes
    task.ContinueWith((antecedent, state) =>
    {
        // Recover our state data
        var tuple = 
            (Tuple<Timer, TaskCompletionSource<VoidTypeStruct>>)state;

        // Cancel the Timer
        tuple.Item1.Dispose();

        // Marshal results to proxy
        MarshalTaskResults(antecedent, tuple.Item2);
    }, 
    Tuple.Create(timer, tcs),
    CancellationToken.None,
    TaskContinuationOptions.ExecuteSynchronously,
    TaskScheduler.Default);

    return tcs.Task;
}

Such changes ensure efficient handling of the edge cases associated with TimeoutAfter.

A Different Approach

My colleague Stephen Toub informed me of another potential implementation of Task.TimeoutAfter:

public static async Task TimeoutAfter(this Task task, int millisecondsTimeout)
{
    if (task == await Task.WhenAny(task, Task.Delay(millisecondsTimeout))) 
        await task;
    else
        throw new TimeoutException();
}

The implementation above takes advantage of the new async/await support in .NET 4.5, and is pleasingly concise. However, it does lack some optimizations:

1. The edge cases described previously are not handled well. (But that could probably be fixed.)
2. A Task is created via Task.Delay, instead of just a simple timer job.
3. In the cases where the source Task (task) completes before the timeout expires, no effort is made to cancel the internal timer job that was launched in the Task.Delay call. If the number of “zombie” timer jobs starts becoming significant, performance could suffer.

Nevertheless, it is good to consider the use of async/await support in implementing features like this. Often await will be optimized in ways that simple continuations are not.

What about TimeoutAfter<TResult>?

Suppose that we want to implement the generic version of TimeoutAfter?

public static Task TimeoutAfter<TResult>(
    this Task<TResult> task, 
    int millisecondsTimeout)

It turns out that the implementation of the above would be nearly identical to the non-generic version, except that a TaskCompletionSource<TResult> would be used instead of a TaskCompletionSource<VoidTypeStruct>. The MarshalTaskResults method was already written to correctly handle the marshaling of the results of generic Tasks.

[1] There is no non-generic version of TaskCompletionSource<TResult>. So, if you want a completion source for a Task (as opposed to a Task<TResult>), you still need to provide some throwaway TResult type to TaskCompletionSource. For this example, we’ve created a dummy type (VoidTypeStruct), and we create a TaskCompletionSource<VoidTypeStruct>.

[2] So does it really buy you anything to replace a closure allocation with a tuple allocation? The answer is “yes”. If you were to examine the IL produced from the original code, you would see that both a closure object and a delegate need to be allocated for this call. Eliminating variable capture in the delegate typically allows the compiler to cache the delegate, so in effect two allocations are saved by eliminating variable capture. Thus in this code we’ve traded closure and delegate allocations for a Tuple allocation, so we still come out ahead.