// @nuget: MSTest.TestFramework

// @nuget: Microsoft.Bcl.AsyncInterfaces

// @nuget: Microsoft.Bcl.TimeProvider

​

using System;

using System.Threading;

using System.Threading.Tasks;

using Microsoft.VisualStudio.TestTools.UnitTesting;

public class Program

{

    private const int RepeatTestCount = 5;

    public static async Task Main()

    {

        await ShouldReturnInTime();

        await ShouldTimeout();

    }

    private static async Task ShouldReturnInTime()

    {

        for (int i = 0; i < RepeatTestCount; i++)

            await TestWithTimeout(10, 25, false);

        Console.WriteLine($"##### {nameof(ShouldReturnInTime)} passed! #####");

    }

    private static async Task ShouldTimeout()

    {

        for (int i = 0; i < RepeatTestCount; i++)

            await TestWithTimeout(25, 10, true);    

        Console.WriteLine($"##### {nameof(ShouldTimeout)} passed! #####");

    }

    private static async Task TestWithTimeout(int longRunningTaskInMs, int waitTimeoutInMs, bool expectedWaitCallbackFired)

    {

        var timeOutCallbackFired = false;

        var actual = await Task.Run(async () =>

        {

            Console.WriteLine($"   {DateTime.Now:O} long running task started...");

            await Task.Delay(longRunningTaskInMs);

            Console.WriteLine($"   {DateTime.Now:O} long running task finished.");

            return "   returned in time";

        }).WithTimeout(waitTimeoutInMs, () =>

        {

            Console.WriteLine("   timeout!");

            timeOutCallbackFired = true;

            return "   fired";

        });

​

        Assert.AreEqual(expectedWaitCallbackFired ? "   fired" : "   returned in time", actual);

        Assert.AreEqual(expectedWaitCallbackFired, timeOutCallbackFired);

    }

}

​

public static class TaskExtensions

{

    public static async Task<T> WithTimeout<T>(this Task<T> task, int millisecondsTimeout, Func<T> onTimeout = null)

    {

        using (var timeoutCancellation = new CancellationTokenSource())

        {

            try

            {

                return await task

                    .WaitAsync(TimeSpan.FromMilliseconds(millisecondsTimeout), TimeProvider.System, timeoutCancellation.Token)

                    .ConfigureAwait(false);

            }

            catch (TimeoutException timeoutException)

            {

                return onTimeout != null ? onTimeout() : throw timeoutException;

            }

        }

    }

}