Useful task extensions

TaskExtensions.cs

namespace System.Threading.Tasks
{
    using System;
    using System.Runtime.CompilerServices;
    using System.Threading;

    /// <summary>
    /// Extensions on Task and Task<T>.
    /// </summary>
    public static class TaskExtensions
    {
        /// <summary>
        /// ConfigureAwait(false)
        /// </summary>
        /// <param name="task">The task to to convert.</param>
        /// <returns>A ConfiguredTaskAwaitable.</returns>
        public static ConfiguredTaskAwaitable NotOnCapturedContext(this Task task)
            => task.ConfigureAwait(false);

        /// <summary>
        /// ConfigureAwait(false)
        /// </summary>
        /// <typeparam name="T">The type result of the task.</typeparam>
        /// <param name="task">The task to to convert.</param>
        /// <returns>A ConfiguredTaskAwaitable.</returns>
        public static ConfiguredTaskAwaitable<T> NotOnCapturedContext<T>(this Task<T> task)
            => task.ConfigureAwait(false);

        /// <summary>
        /// Asynchronously waits for the task to complete, or for the cancellation token to be canceled.
        /// </summary>
        /// <param name="task">The task to wait for. May not be <c>null</c>.</param>
        /// <param name="cancellationToken">The cancellation token that cancels the wait.</param>
        public static Task<T> WithCancellation<T>(this Task<T> task, CancellationToken cancellationToken)
        {
            if (task == null)
            {
                throw new ArgumentNullException(nameof(task));
            }

            if (!cancellationToken.CanBeCanceled)
            {
                return task;
            }

            if (cancellationToken.IsCancellationRequested)
            {
                return Task.FromCanceled<T>(cancellationToken);
            }

            return AwaitWithCancellation(task, cancellationToken);

            async Task<T> AwaitWithCancellation(Task<T> t, CancellationToken ct)
            {
                var tcs = new TaskCompletionSource<T>();
                using (ct.Register(() => tcs.TrySetCanceled(ct), useSynchronizationContext: false))
                {
                    return await await Task.WhenAny(t, tcs.Task).ConfigureAwait(false);
                }
            }
        }

        /// <summary>
        /// Asynchronously waits for any of the source tasks to complete, or for the cancellation token to be canceled.
        /// </summary>
        /// <typeparam name="TResult">The type of the task results.</typeparam>
        /// <param name="task">The tasks to wait for. May not be <c>null</c>.</param>
        /// <param name="cancellationToken">The cancellation token that cancels the wait.</param>
        public static async Task WithCancellation(this Task task, CancellationToken cancellationToken)
        {
            if (task == null)
            {
                throw new ArgumentNullException(nameof(task));
            }

            if (!cancellationToken.CanBeCanceled)
            {
                await task;
            }

            if (cancellationToken.IsCancellationRequested)
            {
                await Task.FromCanceled(cancellationToken);
            }

            await AwaitWithCancellation(task, cancellationToken);

            async Task AwaitWithCancellation(Task t, CancellationToken ct)
            {
                var tcs = new TaskCompletionSource<object>();
                using (ct.Register(() => tcs.TrySetCanceled(ct), useSynchronizationContext: false))
                {
                    await await Task.WhenAny(t, tcs.Task).ConfigureAwait(false);
                }
            }
        }
        
        public static async Task<ValueTuple<T1, T2>> WhenAll<T1, T2>(this ValueTuple<Task<T1>, Task<T2>> tasks)
        {
            await Task.WhenAll(tasks.Item1, tasks.Item2);

            return (tasks.Item1.Result, tasks.Item2.Result);
        }

        public static async Task<ValueTuple<T1, T2, T3>> WhenAll<T1, T2, T3>(this ValueTuple<Task<T1>, Task<T2>, Task<T3>> tasks)
        {
            await Task.WhenAll(tasks.Item1, tasks.Item2, tasks.Item3);

            return (tasks.Item1.Result, tasks.Item2.Result, tasks.Item3.Result);
        }

        public static async Task<ValueTuple<T1, T2, T3, T4>> WhenAll<T1, T2, T3, T4>(this ValueTuple<Task<T1>, Task<T2>, Task<T3>, Task<T4>> tasks)
        {
            await Task.WhenAll(tasks.Item1, tasks.Item2, tasks.Item3, tasks.Item4);

            return (tasks.Item1.Result, tasks.Item2.Result, tasks.Item3.Result, tasks.Item4.Result);
        }
        
        public static Task AsTimedTask(this Task task, Action<TimeSpan> action)
        {
            return TimedExecution();

            async Task TimedExecution()
            {
                var sw = new Stopwatch();
                sw.Start();

                await task;

                sw.Stop();
                action(sw.Elapsed);
            }
        }

        public static Task<T> AsTimedTask<T>(this Task<T> task, Action<TimeSpan> action)
        {
            return TimedExecution();

            async Task<T> TimedExecution()
            {
                var sw = new Stopwatch();
                sw.Start();

                var result = await task;

                sw.Stop();
                action(sw.Elapsed);

                return result;
            }
        }
    }
}

TaskExtensionsTests.cs

namespace System.Threading.Tasks
{
    using Shouldly;
    using System;
    using System.Threading;
    using Xunit;

    public class TaskExtensionsTests
    {
        [Fact]
        public void WithCancellation_ThrowsOriginalException()
        {
            var exception = new Exception();

            Func<Task> act = () => Task.FromException(exception).WithCancellation(CancellationToken.None);
            act.ShouldThrow<Exception>().ShouldBeSameAs(exception);
        }

        [Fact]
        public void WithCancellation_TokenNotCancellable_ReturnsTcsTask()
        {
            var tcs = new TaskCompletionSource<object>();
            var task = tcs.Task.WithCancellation(CancellationToken.None);

            tcs.Task.ShouldBeSameAs(task);
        }

        [Fact]
        public void WithCancellation_AlreadyCanceledToken_ReturnsSynchronouslyCanceledTask()
        {
            var ct = new CancellationToken(true);
            Task task = new TaskCompletionSource<object>().Task.WithCancellation(ct);

            task.IsCanceled.ShouldBeTrue();
        }

        [Fact]
        public async Task WithCancellation_TokenCanceled_ThrowsOperationCanceledException()
        {
            var cts = new CancellationTokenSource();
            Task task = new TaskCompletionSource<object>().Task.WithCancellation(cts.Token);

            cts.Cancel();

            Func<Task> act = () => task;

            await act.ShouldThrowAsync<OperationCanceledException>();
        }

        [Fact]
        public void WithCancellationTResult_TokenNotCancellable_ReturnsTcsTask()
        {
            var tcs = new TaskCompletionSource<object>();
            var task = tcs.Task.WithCancellation(CancellationToken.None);

            tcs.Task.ShouldBeSameAs(task);
        }

        [Fact]
        public void WithCancellationTResultT_AlreadyCanceledToken_ReturnsSynchronouslyCanceledTask()
        {
            var ct = new CancellationToken(true);
            var task = new TaskCompletionSource<object>().Task.WithCancellation(ct);

            task.IsCanceled.ShouldBeTrue();
        }

        [Fact]
        public async Task WithCancellationTResult_TokenCanceled_ThrowsOperationCanceledException()
        {
            var cts = new CancellationTokenSource();
            var task = new TaskCompletionSource<object>().Task.WithCancellation(cts.Token);

            cts.Cancel();

            Func<Task> act = () => task;

            await act.ShouldThrowAsync<OperationCanceledException>();
        }
    }
}