using System;
using System.Threading;
using System.Threading.Tasks;

namespace ExtensionMethods
{
    public static class ExtensionMethods
    {
        public static async Task WithCancellation(this Task task,
            CancellationToken cancellationToken)
        {
            var completionSource = new TaskCompletionSource<bool>();
            using (cancellationToken.Register(
                s => ((TaskCompletionSource<bool>)s).TrySetResult(true),
                completionSource))
            {
                if (task != await Task.WhenAny(task, completionSource.Task).ConfigureAwait(false))
                    throw new OperationCanceledException(cancellationToken);
            }
        }
    }
}