Crafting a Task.TimeoutAfter Method

Crafting a Task.TimeoutAfter Method

  • Comments 18

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;
    }
}

The “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.

Leave a Comment
  • Please add 4 and 3 and type the answer here:
  • Post
  • Another potential implementation of Task.TimeoutAfter with Rx:

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

    {

       return task.ToObservable().Timeout(TimeSpan.FromMilliseconds(millisecondsTimeout)).ToTask();

    }

  • Another approach:

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

           {

               var timedTask = Task.Factory.StartNew(() =>

               {

                   var completed = task.Wait(millisecondsTimeout);

                   if (!completed)

                   {

                       throw new TimeoutException();

                   }

                   return task.Result;

               });

               return timedTask;

           }

    In this case the waiting will stop either by completion, exception, cancellation, or timeout.

  • Omer, while that will provide expected behavior for the returned task, it will also block a thread for the duration of the timeout.  In contrast, the solutions provided in this post will not block threads unnecessarily.  Especially for timer-related functionality, this becomes very important, as it's reasonable to expect a process to have thousands of these active at a time.

  • You are right of course. Don't know how I've missed it. In that case - I'd go with the Rx option that Steffen mentioned. It's concise, readable, and I love Rx :-)

  • I attempted to implement this code, however, I found that the ContinueWith line errors with the following:

    "No overload for method 'ContinueWith' takes 5 arguments". Is there something else that needs to be added or changed to get this to work.

  • @Michael Rivera: Are you using .NET4 or .NET4.5?  I believe that the overload in question only exists in the latter.

  • After some research I found that this is a .net 4.5 framework feature. Do we know if this can work in a 4.0 framework way. So far my attempts have been unsuccessful. The ContinueWith is never called. Any assistance you can provide would be helpful. Thank you.

  • @thxmike: Can you simply revert to the ContinueWith that was used in the first implementation given in the post?  That should work on .NET4.

  • @Joe Hoag: I am attempting to use the ContinueWith indicated in the first post with .Net 4. The execution gets to the timers timeout, sets the exception,  but then never calls ContinueWith after. I think I am just missing something but not sure what.

  • @thxmike: The ContinueWith will only be called if the original task (the "task" parameter to TimeoutAfter) completes.  When the timer expires, it will cause the task *returned* by TimeoutAfter to complete.  So:

       Task taskA = Task.Factory.StartNew(...);

       taskB = taskA.TimeoutAfter(1000); // 1 second

    If the timeout expires, taskB will complete with a TimeoutException.  taskA may well continue running.

    Does that help?

  • @Joe Hoag: I get it now. Very Interesting.  Thanks for you patience. This next question may be out of scope but what happens to the orphaned taskA, lets say it never stops executing ?(i.e. some run away code, or some type of blocking code which never lets it complete).

  • @thxmike: The ultimate disposition of taskA is the concern of the person who creates/launches taskA.  The purpose of Task.TimeoutAfter is to return a *proxy* (taskB in the example above) for taskA that will either (1) be completed with a TimeoutException if taskA fails to complete within the specified time span, or (2) be completed with taskA's completion information if taskA completes within the specified time span.

    This proxy can be used to effectively "time out" a task that was handed to you, which you did not create, without actually perturbing that original task.  Consider something like this:

       Task taskA; // handed to you

       var proxy = taskA.TimeoutAfter(1000); // times out after one second

       proxy.ContinueWith( antecedent => {DoSomething(antecedent);});

    That continuation can now fire, one way or another, whether or not taskA completes.  The DoSomething() method can check to see whether or not "antecedent" completed or timed out, and act accordingly.

  • Is this thread safe? Timer delegate and task.ContinueWith can execute at same time.

  • @Vamsi: Which example, and what race specifically are you concerned about?  In these examples, the timer just serves to complete the task, so the timer's callback won't be doing anything meaningful by the time the continuation is running.

  • Can anyone provide Steffen's implementation but for Task<T> which returns a result?

Page 1 of 2 (18 items) 12