Wait for a canceled task to complete before continuing execution

Question

I have the following code, where a Task can be canceled, but I basically need to wait for it to complete (to ensure integrity) before throwing the OperationCanceledException to the caller.

public static void TaskCancellationTest() {
    try {
        Console.WriteLine("TaskCancellationTest started.");
        var cts = new CancellationTokenSource();
        var t = Task.Run(() => {
            if (cts.Token.IsCancellationRequested) return;
            Console.WriteLine("1");
            Task.Delay(2000).Wait();
            Console.WriteLine("2");
        }).ContinueWith(task => {
            if (cts.Token.IsCancellationRequested) return;
            Console.WriteLine("3");
            Task.Delay(2000).Wait();
            Console.WriteLine("4");
        });
        Task.Run(() => {
            Task.Delay(1000).Wait();
            Console.WriteLine("Cancelling...");
            cts.Cancel();
        });
        t.Wait();
        try {
            cts.Token.ThrowIfCancellationRequested();
        }
        catch (OperationCanceledException) {
            Console.WriteLine("Gracefully canceled.");
        }
        Console.WriteLine("TaskCancellationTest completed.");
    }
    catch (Exception ex) {
        Console.WriteLine("TaskCancellationTest... Failure: " + ex);
    }
}

The result, as expected, is:

1
Cancelling...
2
Gracefully canceled.

It works, but I would prefer to pass the CancellationToken to the methods as I understand this is a better pattern. I would also like to be able to observe the token inside the method body and to call ThrowIfCancellationRequested() to abort without having to wait for the next ContinueWith().

I was playing with the following alternative code, which also works, but is there any way to have an OperationCanceledException raised instead of an AggregateException?

If I pass the cancellationToken to the WaitAll() method, the problem is that it will throw an OperationCanceledException immediately upon cancellation of the token, rather than waiting for the tasks t1 and t2 to actually complete (they will continue running in the background) and then only throwing the exception.

public static void TaskCancellationTest2() {
    try {
        Console.WriteLine("TaskCancellationTest2 started.");
        var cts = new CancellationTokenSource();
        var t1 = Task.Run(() => {
            Console.WriteLine("1");
            Task.Delay(2000).Wait();
            Console.WriteLine("2");
        }, cts.Token);
        var t2 = t1.ContinueWith(task => {
            Console.WriteLine("3");
            Task.Delay(2000).Wait();
            cts.Token.ThrowIfCancellationRequested();
            Console.WriteLine("4");
        }, cts.Token);
        Task.Run(() => {
            Task.Delay(1000).Wait();
            Console.WriteLine("Cancelling...");
            cts.Cancel();
        });
        try {
            try {
                Task.WaitAll(t1, t2);
            }
            catch (AggregateException ae) {
                if (ae.InnerExceptions.Count == 1 && ae.InnerExceptions.Single() is OperationCanceledException) {
                    throw ae.InnerExceptions.Single();
                }
                throw;
            }
        }
        catch (OperationCanceledException) {
            Console.WriteLine("Gracefully canceled.");
        }
        Console.WriteLine("TaskCancellationTest2 completed.");
    }
    catch (Exception ex) {
        Console.WriteLine("TaskCancellationTest2... Failure: " + ex);
    }
}

I have prepared a fiddle here.

This question's title is very similar to mine, but the accepted answer is unfortunately not relevant to my case.

Do you know of any way to achieve what I would like, that makes as good use of CancellationToken as possible?

Answer 1

I think the TPL is designed to eagerly complete tasks if the CancellationToken is set. Part of the reason you are seeing this behavior is because you are calling t.Wait(cts.Token). The overload that takes a CancellationToken will stop waiting if the token is set even if the task hasn't ran to completion.

It's the same with ContinueWith if you pass in a CancellationToken the task can complete as soon as that token is set.

Change your code to call t.Wait() and ContinueWith without a token and you'll get the behavior you want.

    public static void TaskCancellationTestNotWorking1()
    {
        try
        {
            Console.WriteLine("TaskCancellationTestNotWorking started.");
            var cts = new CancellationTokenSource();
            var t = Task.Run(() =>
            {
                Console.WriteLine("1");
                Thread.Sleep(2000);
                Console.WriteLine("2");
            }, cts.Token).ContinueWith(task =>
            {
                Console.WriteLine("3");
                Thread.Sleep(2000);
                cts.Token.ThrowIfCancellationRequested();
                Console.WriteLine("4");
            });

            Task.Run(() =>
            {
                Thread.Sleep(1000);
                Console.WriteLine("Cancelling...");
                cts.Cancel();
            }, cts.Token);

            try
            {
                t.Wait();
            }
            catch (OperationCanceledException)
            {
                Console.WriteLine("IsCanceled " + t.IsCanceled);
                Console.WriteLine("IsCompleted " + t.IsCompleted);                    

                Console.WriteLine("Gracefully canceled.");
            }
            catch (AggregateException)
            {
                Console.WriteLine("IsCanceled " + t.IsCanceled);
                Console.WriteLine("IsCompleted " + t.IsCompleted);

                Console.WriteLine("Gracefully canceled 1.");
            }

            Console.WriteLine("TaskCancellationTestNotWorking completed.");
        }
        catch (Exception ex)
        {
            Console.WriteLine("TaskCancellationTestNotWorking... Failure: " + ex);
        }
    }

You might find this article useful How do I cancel non-cancelable async operations?