// @nuget: MSTest.TestFramework

​

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())

        {

            var delayTask = Task.Delay(millisecondsTimeout, timeoutCancellation.Token);

            var completedTask = await Task.WhenAny(task, delayTask).ConfigureAwait(false);

​

            if (completedTask == task)

            {

                // original task completed

                timeoutCancellation.Cancel();

                return await task;

            }

​

            // timeout

            return onTimeout != null ? onTimeout() : throw new TimeoutException();

        }

    }

}