Unable to cancel a Task returned by TaskCompletionSource by injecting a timeout in C#

Question

I have the following method that waits for an event to happen to set the results by a TaskCompletionSource.

private async Task<string?> RequestDataAsync(string? remoteEntity, CancellationToken cancellationToken)
{
    var taskCompletionSource = new TaskCompletionSource<string?>();
    try
    {
        if (remoteEntity is not null)
        {
            var handler = GetHandler();
            handler.Event.OnResultReady += result =>
            {
                taskCompletionSource.SetResult(result);
            };
            await SendRequestToAcquireDataAsync(cancellationToken);
        }
        else
        {
            taskCompletionSource.SetResult(null);
        }
    }
    catch (Exception ex)
    {
        taskCompletionSource.SetException(ex);
    }
    return await taskCompletionSource.Task;
}

As the code states, we are subscribing to the OnResultReady event and after the subscription line we are sending the request to the remote server to acquire data. The server is usually prompt enough to provide the results but there were cases the server was not be responsive enough. So we want to terminate RequestDataAsync() after 10 seconds to prevent blocking the client app. So, I wrote the following wrapper code to call RequestDataAsync() method while employing a cancellation token that kicks in after 10 seconds:

    public async Task<string?> RequestDataWithTimeoutAsync(string? remoteEntity, CancellationToken cancellationToken, int requestLifetimeInMilliseconds = 10000)
    {
        var lifeTimeCancellationTokenSource = new CancellationTokenSource();
        using var linkedCancellationToken = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, lifeTimeCancellationTokenSource.Token);
        lifeTimeCancellationTokenSource.CancelAfter(requestLifetimeInMilliseconds);
        string? result = null;
        try
        {
            result = await RequestDataAsync(remoteEntity, linkedCancellationToken.Token);
        }
        catch (OperationCanceledException oex) when (oex.CancellationToken == linkedCancellationToken.Token)
        {
            if(cancellationToken.IsCancellationRequested)
            {
                throw new OperationCanceledException(oex.Message, oex, cancellationToken);
            }
            else if(lifeTimeCancellationTokenSource.IsCancellationRequested)
            {
                throw new TimeoutException();
            }
        }
        return result;
    }

SendRequestToAcquireDataAsync() posts a web request to a Web API whose backend code produces data asynchronously and make it available through the handler's event handler. Imagine it like a durable azure function.

The problem I've been facing is the wrapper method does not terminate the first method after 10 seconds and the first method hangs indefinitely. What is the solution to this problem?

Answer 1

If SendRequestToAcquireDataAsync is properly cancelling when its CancellationToken is cancelled, then you're probably just hanging around for OnResultReady to be fired, even after your CancellationToken has been cancelled.

So, you probably need to register an action with the CancellationToken to cancel the TaskCompletionSource when it's cancelled. Something like:

private async Task<string?> RequestDataAsync(string? remoteEntity, CancellationToken cancellationToken)
{
    var taskCompletionSource = new TaskCompletionSource<string?>();
    try
    {
        if (remoteEntity is not null)
        {
            var handler = GetHandler();
            handler.Event.OnResultReady += result =>
            {
                taskCompletionSource.TrySetResult(result);
            };
            await SendRequestToAcquireDataAsync(cancellationToken);
        }
        else
        {
            taskCompletionSource.SetResult(null);
        }
    }
    catch (Exception ex)
    {
        taskCompletionSource.SetException(ex);
    }

    using (cancellationToken.Register(() => taskCompletionSource.TrySetCanceled())
    {
        return await taskCompletionSource.Task;
    }
}

I've also changed a SetResult to a TrySetResult, to avoid an exception if the event gets fired after the CancellationToken is cancelled.

---

You should probably also be unsubscribing from that event once this is all over:

private async Task<string?> RequestDataAsync(string? remoteEntity, CancellationToken cancellationToken)
{
    var taskCompletionSource = new TaskCompletionSource<string?>();
    Handler? handler = null;

    void ResultReadyHandler(string result)
    {
        taskCompletionSource.TrySetResult(result);
    }

    try
    {
        if (remoteEntity is not null)
        {
            handler = GetHandler();
            handler.Event.OnResultReady += ResultReadyHandler;
            await SendRequestToAcquireDataAsync(cancellationToken);
        }
        else
        {
            taskCompletionSource.SetResult(null);
        }
    }
    catch (Exception ex)
    {
        taskCompletionSource.SetException(ex);
    }

    using (cancellationToken.Register(() => taskCompletionSource.TrySetCanceled())
    {
        try
        {
            return await taskCompletionSource.Task;
        }
        finally
        {
            if (handler != null)
            {
                handler.Event.OnResultReady -= ResultReadyHandler;
            }
        }
    }
}

(Obviously fix the type of Handler).

---

That said, you're making this more complex than it needs to be. In your catch, you're setting the exception on the TaskCompletionSource, and then awaiting that TaskCompletionSource.Task which just rethrows the exception. That's just a long-winded way of just letting the exception bubble up, which it does automatically. Same with returning null:

private async Task<string?> RequestDataAsync(string? remoteEntity, CancellationToken cancellationToken)
{
    if (remoteEntity is null)
    {
        return null;
    }
  
    var taskCompletionSource = new TaskCompletionSource<string?>();

    void ResultReadyHandler(string result)
    {
        taskCompletionSource.TrySetResult(result);
    }

    var handler = GetHandler();
    handler.Event.OnResultReady += ResultReadyHandler;

    try
    {
        await SendRequestToAcquireDataAsync(cancellationToken)
        using (cancellationToken.Register(() => taskCompletionSource.TrySetCanceled())
        {
            return await taskCompletionSource.Task;
        }
    }
    finally
    {
        handler.Event.OnResultReady -= ResultReadyHandler;
    }
}

---

That said, why does SendRequestToAcquireDataAsync have such an awkward API? Why does it return a Task that you have to await, and then fire an event which you have to subscribe to?