Asynchronously Wait for Task to Complete With Timeout

I was recently working at an async method which could possibly hang if some dependent stuff did not happen. To prevent the method from hanging, I wanted to implement some kind of timeout. Now, how can I make a task abort, given the following method signature?

1
2
3
4
5
6
7
public class SomeClient
{
    public async Task DoStuffAsync(CancellationToken? ct = null)
    {
        ...
    }
}

Take One - Task Extension Method

The plan was to implement som kind of extension method to Task which could add the timeout-functionality.

1
2
3
4
5
6
7
8
9
10
11
12
internal static class TaskExtensions
{
    public static async Task TimeoutAfter(
        this Task task, int millisecondsTimeout, CancellationToken ct)
    {
        var completedTask = await Task.WhenAny(
            task, Task.Delay(millisecondsTimeout, ct));
        if (completedTask == task)
            return;
        throw new TimeoutException();
    }
}

And use it like this.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
public class SomeClient
{
    public async Task DoStuffAsync(
        CancellationToken? ct = null, int millisecondsTimeout = 20_0000)
    {
        ...
        var notNullCt = ct ?? CancellationToken.None;
        await DoStuffInnerAsync(notNullCt).TimeoutAfter(
            millisecondsTimeout, notNullCt);
        ...
    }

    private async Task DoStuffInnerAsync(CancellationToken ct)
    {
        ...
    }
}

This design allowed that the internal Delay-task would be cancelled if the user of my API canceled the method call. Nice! But it also had some mayor disadvantages:

  • No task will be canceled if either of them finish successfully, which leads to having tasks running in the background for no reason, eating system resources.
  • I had to make sure to pass the same cancellation token both to DoStuffInnerAsync and TimeoutAfter, which might be something that could lead to mistakes further down.

Take Two - Expanding the Extension Method

To be able to cancel the TimeoutAfter-task i needed a CancellationTokenSource-instance, and pass its token to the TimeoutAfter-method. And I also wanted the TimeoutAfter-task to cancel if the user canceled the public API call.

This is what I came up with.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
internal static class TaskExtensions
{
    public static async Task TimeoutAfter(
        this Task task, int millisecondsTimeout, CancellationToken ct)
    {
        using (var cts = new CancellationTokenSource())
        {
            using (ct.Register(() => cts.Cancel()))
            {
                var completedTask = await Task.WhenAny(
                    task, Task.Delay(millisecondsTimeout, cts.Token));
                if (completedTask == task)
                {
                    cts.Cancel();
                    return;
                }
                throw new TimeoutException();                    
            }
        }
    }
}

This is some seriously dangerous programming.

  • By subscribing to cancel-events with ct.Register(...) I opened upp the possibility for memory leaks if I do not unsubscribe somehow.
  • Also, using cts (which can be disposed) in the delegate passed to ct.Register(...) might actually make my application crash if ct was canceled outside of the TimeOutAfter-method scope.

Register returns a disposable something, which when disposed will unsubscribe. By adding the inner using-block, I fixed both of these problems.

This made it possible to cancel the Delay-task when the actual task completed, but not the reverse. How should I solve the bigger problem, how to cancel the actual task if it would hang? Eating up system resources indefinitely…

Take Three - Changing the Public API

With much hesitation I finally decided to make a breaking change to the public API by replacing the CancellationToken with a CancellationTokenSource in the DoStuffAsync-method.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
public class SomeClient
{
    public async Task DoStuffAsync(
        CancellationTokenSource cts = null, int millisecondsTimeout = 20_0000)
    {
        ...
        var notNullCts = cts ?? new CancellationTokenSource();
        await DoStuffInnerAsync(notNullCts.Token).TimeoutAfter(
            millisecondsTimeout, notNullCts);
        ...
    }

    private async Task DoStuffInnerAsync(CancellationToken ct)
    {
        ...
    }
}

internal static class TaskExtensions
{
    public static async Task TimeoutAfter(
        this Task task, int millisecondsTimeout, CancellationTokenSource cts)
    {
        var completedTask = await Task.WhenAny(
            task, Task.Delay(millisecondsTimeout, cts.Token));
        if (completedTask == task)
        {
            cts.Cancel();
            return;
        }
        cts.Cancel();
        throw new TimeoutException();                    
    }
}

Nice! But this still did not solve that I had make sure to pass the same cts to both the actual task and the Delay-task.

Final Solution - Doing the Obvious

Most things is really easy when you know the answer. By accident I stumbled upon that CancellationTokenSource has a CancelAfter(...)-method. This solves my problem entirely without the need to update my public API.

1
2
3
4
var client = new SomeClient();
var cts = new CancellationTokenSource();
cts.CancelAfter(20_000);
await client.DoStuffAsync(cts.Token);

Easy peasy. I wish I had known about this earlier!