What is the correct way to cancel an async operation that doesn't accept a CancellationToken?

Question

What is the correct way to cancel the following?

var tcpListener = new TcpListener(connection);
tcpListener.Start();
var client = await tcpListener.AcceptTcpClientAsync();

Simply calling tcpListener.Stop() seems to result in an ObjectDisposedException and the AcceptTcpClientAsync method doesn't accept a CancellationToken structure.

Am I totally missing something obvious?

Solution 1

Assuming that you don't want to call the Stop method on the TcpListener class, there's no perfect solution here.

If you're alright with being notified when the operation doesn't complete within a certain time frame, but allowing the original operation to complete, then you can create an extension method, like so:

public static async Task<T> WithWaitCancellation<T>( 
    this Task<T> task, CancellationToken cancellationToken) 
{
    // The tasck completion source. 
    var tcs = new TaskCompletionSource<bool>(); 

    // Register with the cancellation token.
    using(cancellationToken.Register( s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs) ) 
    {
        // If the task waited on is the cancellation token...
        if (task != await Task.WhenAny(task, tcs.Task)) 
            throw new OperationCanceledException(cancellationToken); 
    }

    // Wait for one or the other to complete.
    return await task; 
}

The above is from Stephen Toub's blog post "How do I cancel non-cancelable async operations?".

The caveat here bears repeating, this doesn't actually cancel the operation, because there is not an overload of the AcceptTcpClientAsync method that takes a CancellationToken, it's not able to be cancelled.

That means that if the extension method indicates that a cancellation did happen, you are cancelling the wait on the callback of the original Task, not cancelling the operation itself.

To that end, that is why I've renamed the method from WithCancellation to WithWaitCancellation to indicate that you are cancelling the wait, not the actual action.

From there, it's easy to use in your code:

// Create the listener.
var tcpListener = new TcpListener(connection);

// Start.
tcpListener.Start();

// The CancellationToken.
var cancellationToken = ...;

// Have to wait on an OperationCanceledException
// to see if it was cancelled.
try
{
    // Wait for the client, with the ability to cancel
    // the *wait*.
    var client = await tcpListener.AcceptTcpClientAsync().
        WithWaitCancellation(cancellationToken);
}
catch (AggregateException ae)
{
    // Async exceptions are wrapped in
    // an AggregateException, so you have to
    // look here as well.
}
catch (OperationCancelledException oce)
{
    // The operation was cancelled, branch
    // code here.
}

Note that you'll have to wrap the call for your client to capture the OperationCanceledException instance thrown if the wait is cancelled.

I've also thrown in an AggregateException catch as exceptions are wrapped when thrown from asynchronous operations (you should test for yourself in this case).

That leaves the question of which approach is a better approach in the face of having a method like the Stop method (basically, anything which violently tears everything down, regardless of what is going on), which of course, depends on your circumstances.

If you are not sharing the resource that you're waiting on (in this case, the TcpListener), then it would probably be a better use of resources to call the abort method and swallow any exceptions that come from operations you're waiting on (you'll have to flip a bit when you call stop and monitor that bit in the other areas you're waiting on an operation). This adds some complexity to the code but if you're concerned about resource utilization and cleaning up as soon as possible, and this choice is available to you, then this is the way to go.

If resource utilization is not an issue and you're comfortable with a more cooperative mechanism, and you're not sharing the resource, then using the WithWaitCancellation method is fine. The pros here are that it's cleaner code, and easier to maintain.

Solution 2

While casperOne's answer is correct, there's a cleaner potential implementation to the WithCancellation (or WithWaitCancellation) extension method that achieves the same goals:

static Task<T> WithCancellation<T>(this Task<T> task, CancellationToken cancellationToken)
{
    return task.IsCompleted
        ? task
        : task.ContinueWith(
            completedTask => completedTask.GetAwaiter().GetResult(),
            cancellationToken,
            TaskContinuationOptions.ExecuteSynchronously,
            TaskScheduler.Default);
}

• First we have a fast-path optimization by checking whether the task has already completed.
• Then we simply register a continuation to the original task and pass on the CancellationToken parameter.
• The continuation extracts the original task's result (or exception if there is one) synchronously if possible (TaskContinuationOptions.ExecuteSynchronously) and using a ThreadPool thread if not (TaskScheduler.Default) while observing the CancellationToken for cancellation.

If the original task completes before the CancellationToken is canceled then the returned task stores the result, otherwise the task is canceled and will throw a TaskCancelledException when awaited.

Comments

[Glenn Slayden]

Why do you use completedTask.GetAwaiter().GetResult() instead of just completedTask.Result?

[i3arnon]

@GlennSlayden it throws the first (i.e. "real") exception if there is one as opposed to Task.Result that throws the AggregateException

[Glenn Slayden]

Ok, great; thanks. Do you also want to add cancellationToken.ThrowIfCancellationRequested(); as the first line of your helpful extension method?

[i3arnon]

@GlennSlayden you can...though it's not really needed. The continuation would be cancelled immediately. Also, consumers may be surprised by an exception thrown synchronously and not from awaiting the task.

[Jack Ukleja]

@i3arnon The cancellationToken parameter for ContinueWith is described as "The CancellationToken that will be assigned to the new continuation task." Are you sure this will actually cancel "early" before the original task completes? Documentation indicates otherwise

[i3arnon]

@Schneider it doesn't. As you quoted, the new continuation will be cancelled before the original task completes.

[Alonzzo2]

what's the difference between this line: using (cancellationToken.Register(s => ((TaskCompletionSource<bool>)s).TrySetResult(true) , tcs)) And the way I wrote it (cause I didn't understand the previous one): using (cancellationToken.Register(() => tcs.TrySetResult(true))) ?