C# extension method to do some action if a target operation takes too long time to finish

Question

I have a potentially long running operation and I want to trigger some action if takes too long time. Also I want to reuse this logic. Operation to check:

public static async Task Do(CancellationToken ct, TimeSpan actionDuration)
{
    await Task.Delay(actionDuration, ct);
    Debug.WriteLine("action finished");
}

Example of usage:

[Fact]
public async Task AsyncTaskMethodWithoutException_LongerThanCheck_CheckTriggeredActionSucceed()
{
    Func<Task> action = () => Do(CancellationToken.None, TimeSpan.FromMilliseconds(300));
    
    var checkResult = await action.Check(TimeSpan.FromMilliseconds(100),
        () => { Debug.WriteLine("too long"); return Task.CompletedTask; },
        CancellationToken.None);

    checkResult.IsOnTooLongFired.Should().BeTrue();
    checkResult.ActionTask.Status.Should().Be(TaskStatus.RanToCompletion);
    checkResult.CheckTask.Status.Should().Be(TaskStatus.Canceled);
}

Here is my code:

public class TooLongCheckResult
{
    public static TooLongCheckResult ActionFailedInstantly = new TooLongCheckResult(Task.CompletedTask, Task.CompletedTask);

    public TooLongCheckResult(Task actionTask, Task checkTask)
    {
        ActionTask = actionTask;
        CheckTask = checkTask;
    }

    public Task ActionTask { get; }
    public Task CheckTask { get; }
    public bool IsOnTooLongFired { get; private set; }

    public void OnTooLongFired() => IsOnTooLongFired = true;
}

public sealed class TooLongCheckResult<T>
{
    public static TooLongCheckResult<T> ActionFailedInstantly = new TooLongCheckResult<T>(Task.CompletedTask, Task.CompletedTask);

    public TooLongCheckResult(Task actionTask, Task checkTask)
    {
        ActionTask = actionTask;
        CheckTask = checkTask;
    }

    public Task ActionTask { get; }
    public Task CheckTask { get; }
    public bool IsOnTooLongFired { get; private set; }
    public T Result { get; private set; }

    public void SetResult(T result) => Result = result;
    public void OnTooLongFired() => IsOnTooLongFired = true;
}

public static class TooLongTaskExt
{
    public static async Task<TooLongCheckResult> Check(this Func<Task> action, TimeSpan tooLongThreshold, Func<Task> onTooLong, CancellationToken ct)
    {
        var childCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        childCts.CancelAfter(tooLongThreshold);
        Task actionTask;
        try
        {
            actionTask = action.Invoke();
        }
        catch (Exception ex)
        {
            //most likely it's because method in action isn't declared with async keyword
            Debug.WriteLine("instant error: " + ex.GetType().Name);
            return TooLongCheckResult.ActionFailedInstantly;
        }

        var timerTask = Task.Delay(Timeout.InfiniteTimeSpan, childCts.Token);
        var result = new TooLongCheckResult(actionTask, timerTask);
        var firstFinished = await Task.WhenAny(timerTask, actionTask);
        if (firstFinished == timerTask)
        {
            await onTooLong.Invoke();
            result.OnTooLongFired();
        }
        else
        {
            childCts.Cancel();
        }

        try
        {
            await Task.WhenAll(timerTask, actionTask);
        }
        catch (TaskCanceledException)
        {
            Debug.WriteLine("cancellation is triggered");
        }
        catch (Exception ex)
        {
            Debug.WriteLine("awaited error: " + ex.GetType().Name);
            throw;
        }

        Debug.WriteLine("actionTask: " + actionTask.Status);
        Debug.WriteLine("timerTask: " + timerTask.Status);

        return result;
    }

    public static async Task<TooLongCheckResult<T>> CheckWithResult<T>(this Func<Task<T>> action, TimeSpan tooLongThreshold, Func<Task> onTooLong, CancellationToken ct)
    {
        var childCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        childCts.CancelAfter(tooLongThreshold);
        Task<T> actionTask;
        try
        {
            actionTask = action.Invoke();
        }
        catch (Exception ex)
        {
            //most likely it's because method in action isn't declared with async keyword
            Debug.WriteLine("instant error: " + ex.GetType().Name);
            return TooLongCheckResult<T>.ActionFailedInstantly;
        }

        var timerTask = Task.Delay(Timeout.InfiniteTimeSpan, childCts.Token);
        var result = new TooLongCheckResult<T>(actionTask, timerTask);
        var firstFinished = await Task.WhenAny(timerTask, actionTask);
        if (firstFinished == timerTask)
        {
            await onTooLong.Invoke();
            result.OnTooLongFired();
        }
        else
        {
            childCts.Cancel();
        }

        try
        {
            await Task.WhenAll(timerTask, actionTask);
        }
        catch (TaskCanceledException)
        {
            Debug.WriteLine("cancellation is triggered");
        }
        catch (Exception ex)
        {
            Debug.WriteLine("awaited error: " + ex.GetType().Name);
            throw;
        }

        Debug.WriteLine("actionTask: " + actionTask.Status);
        Debug.WriteLine("timerTask: " + timerTask.Status);

        if (actionTask.Status == TaskStatus.RanToCompletion)
        {
            result.SetResult(actionTask.Result);
        }

        return result;
    }
}

Answer 1

Just a suggestion to consider as improvement.

Abstract example:

public static class TaskExtensions
{
    public static Task InvokeAfterAsync(this Task task, double thresholdMilliseconds, Func<Task> callbackAsync, CancellationToken token = CancellationToken.None)
        => task.InvokeAfterAsync(TimeSpan.FromMilliseconds(thresholdMilliseconds), callbackAsync, token);

    public static async Task InvokeAfterAsync(this Task task, TimeSpan threshold, Func<Task> callbackAsync, CancellationToken token = CancellationToken.None)
    {
        using (CancellationTokenSource cts = CancellationTokenSource.CreateLinkedTokenSource(token))
        {
            Task fastTask = await Task.WhenAny(task, Task.Delay(threshold, cts.Token));
            if (fastTask != task)
            {
                await Task.WhenAll(callbackAsync(), task);
            }
            else
            {
                cts.Cancel();
            }
        }
    }
}

Possible usage

await DoMyJobAsync().InvokeAfterAsync(100, myCallbackAsync);

• CTS implements IDisposable, consider to Dispose() it or decorate with using statement.
• Don't use shorten names if possible. Make code more readable.
• Use less var statements. Use it only with = new if you want to store exactly the same type (not interface or parent type), otherwise avoid it. (but it's up to you)
• Let exceptions to be handled outside. E.g. if caller want to Cancel() its own CTS, it must care about to catch the exception.
• Use method overloads and default values to simplify the method usage.