// @nuget: NUnit

// @nuget: Microsoft.Bcl.AsyncInterfaces

// @nuget: Microsoft.Bcl.TimeProvider

// @nuget: Microsoft.Extensions.TimeProvider.Testing

​

using System;

using System.Threading;

using System.Threading.Tasks;

using NUnit.Framework;

using Microsoft.Extensions.Time.Testing;

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, 11, false);

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

    }

    private static async Task ShouldTimeout()

    {

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

            await TestWithTimeout(11, 10, true);    

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

    }

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

    {

        var timeProvider = new FakeTimeProvider();

        var timeOutCallbackFired = false;

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

        {

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

            await Task.Delay(longRunningTaskInMs);

            timeProvider.Advance(TimeSpan.FromMilliseconds(longRunningTaskInMs));

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

            return "   returned in time";

        }).WithTimeout(waitTimeoutInMs, timeProvider: timeProvider, onTimeout: () =>

        {

            Console.WriteLine("   timeout!");

            timeOutCallbackFired = true;

            return "   fired";

        });

​

        Assert.That(actual, Is.EqualTo(expectedWaitCallbackFired ? "   fired" : "   returned in time"));

        Assert.That(timeOutCallbackFired, Is.EqualTo(expectedWaitCallbackFired));

    }

}

​

public static class TaskExtensions

{

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

    {

        using (var timeoutCancellation = new CancellationTokenSource())

        {

            try

            {

                return await task

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

                    .ConfigureAwait(false);

            }

            catch (TimeoutException timeoutException)

            {

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

            }

        }

    }

}