TaskExtensions For Dennis

TaskExtensions.cs

public static class TaskExtensions
    {
        public static Task<TResult> WaitWithCancellation<TResult>(this Task<TResult> task, CancellationToken token = default(CancellationToken))
        {
            var tcs = new TaskCompletionSource<TResult>();
            var registration = token.Register(s =>
            {
                var source = (TaskCompletionSource<TResult>) s;
                source.TrySetCanceled();
            }, tcs);

            task.ContinueWith((t, s) =>
            {
                var tcsAndRegistration = (Tuple<TaskCompletionSource<TResult>, CancellationTokenRegistration>) s;

                if (t.IsFaulted && t.Exception!= null)
                {
                    tcsAndRegistration.Item1.TrySetException(t.Exception.InnerException);
                }

                if (t.IsCanceled)
                {
                    tcsAndRegistration.Item1.TrySetCanceled();
                }

                if (t.IsCompleted)
                {
                    tcsAndRegistration.Item1.TrySetResult(t.Result);
                }

                tcsAndRegistration.Item2.Dispose();
            }, Tuple.Create(tcs, registration), CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);

            return tcs.Task;
        }
    }

Test.cs

 [Test]
        public async Task Completed()
        {
            var cts = new CancellationTokenSource();
            var t = DoSomething();

            var result = await t.WaitWithCancellation(cts.Token);

            Assert.AreEqual("Hello", result);
        }

        [Test]
        public async Task Cancelled()
        {
            var cts = new CancellationTokenSource(TimeSpan.FromSeconds(1));
            var t = DoSomething();

            Assert.ThrowsAsync<TaskCanceledException>(async () => await t.WaitWithCancellation(cts.Token));
        }

        [Test]
        public async Task Throws()
        {
            var cts = new CancellationTokenSource();
            var t = DoSomethingWithThrows();

            Assert.ThrowsAsync<InvalidOperationException>(async () => await t.WaitWithCancellation(cts.Token));
        }

        static async Task<string> DoSomething()
        {
            await Task.Delay(5000);
            return "Hello";
        }

        static async Task<string> DoSomethingWithThrows()
        {
            await Task.Delay(5000);
            throw new InvalidOperationException();
        }