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
DoStuffInnerAsyncandTimeoutAfter, 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 toct.Register(...)might actually make my application crash ifctwas canceled outside of theTimeOutAfter-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!
Leave a Reply