TaskCompletionSource Pitfalls

Posted on May 14, 2023 in dotnet

Imagine you are implementing a client for a server that performs an operation on some data, and that operation is totally asynchronous, and the server has only one connection to it.

Contrary to an HTTP server, where a new connection is opened for each request (or, where requests are queued on one connection), requests must be pushed to the server with some sort of unique identifier, and the server will push responses back in whatever order it produces them.

Let us try to implement the client API:

public async Task<Response> SubmitRequestAsync(Request request)
{
    await server.SendRequestAsync(request);

    // now what?
}

Here, we have to return a Task<Response> that represents the client waiting for the response... but what shall we be waiting on? Due to the full asynchronous nature of the server, we probably want to store the request in some sort of lookup table, and then wait for the corresponding response:

public async Task<Response> SubmitRequestAsync(Request request)
{
    await server.SendRequestAsync(request);

    _requests[request.UniqueID] = request;

    // now what?
}

We need a background task that listens on the connection and receive responses:

public async Task HandleResponses()
{
    while (true)
    {
        var response = await server.ReceiveResponseAsync();

        // now what?
    }
}

Still, what should the SubmitRequestAsync return, that we could await? And how should the HandleResponse method complete the API call?

TaskCompletionSource

The TaskCompletionSource is a nice feature of .NET that allows a developer to create "something that creates a Task" and then explicitely decide when to complete, cancel or fail that Task. Using a completion source, we can finalize our API:

public async Task<Response> SubmitRequestAsync(Request request)
{
    await server.SendRequestAsync(request);

    var completion = new TaskCompletionSource<Response>();

    _requests[request.UniqueID] = request;
    _completions[request.UniqueID] = completion;

    return completion.Task;
}

Here, we create a completion source, and return its Task. This is a "logical" task: there is no executing code corresponding to it. It just represent "the execution of something". And that execution would be completed by the background task:

public async Task HandleResponses()
{
    while (true)
    {
        var response = await server.ReceiveResponseAsync();
        var completion = _completions[response.UniqueID];
        completion.TrySetResult(response);        
    }
}

The task retrieves the completion source corresponding to the unique request identifier, and completes it by assigning it a result. In a more elaborate scenario, it could also fail it by assigning it an exception, or even cancel it.

NOTE: this is a simplified version of the code. In real life, we would need to ensure thread safety, remove requests and completions from their lookup tables, etc.

Yes, but

We could end the article here, but... there's a but. What do you think happens when completion.TrySetResult(response) is invoked? Intuitively, we assume that:

• TrySetResult returns from the call and the while loop repeats;
• Whatever code was awaiting the completion's Task is unlocked and starts running.

In other words, invoking TrySetResults fork a new parallel code execution path that resumes whatever code was await-ing the completion's Task, while the main code execution path continues looping the while loop.

However—and here is the but—.NET tries to optimize asynchronous calls to avoid unnecessary and expensive management of tasks and threads. You have probably heard that when invoking the following method, the call to DoSomethingElse will run immediately and synchronously, and that the method will only return a Task when it encounters the first await statement:

public async Task DoSomething()
{
    DoSomethingElse();
    await DoYetAnotherThing();
}

It turns out that the very same thing happens with TrySetResult. .NET will immediately run the code that is await-ing on the completion's Task, and the call to TrySetResult will only complete after either that code is done running, or an await statement is encountered.

And then it all fails

Now look at the following code:

public async Task HandleResponses()
{
    while (true)
    {
        var response = await server.ReceiveResponseAsync();
        AquireSemaphore();
        var completion = _completions[response.UniqueID];
        completion.TrySetResult(response);        
        ReleaseSemaphore();
    }
}

public async Task UseTheAPI()
{
    var response = await SubmitRequestAsync(request);
    AquireSemaphore();
    // do something with the response
    ReleaseSemaphore();
}

What do you think will happen? Spoiler: it will deadlock. The call to TrySetResult will complete the SubmitRequestAsync call and flow and continue the UseTheAPI method. It will try to run AquireSemaphore in UseTheAPI before returning, and will fail to get the semaphore, because it has already been acquired. It will never return, and never release the semaphore, and everything hangs.

In other words: invoking TrySetResult is not a fire-and-forget thing. In most cases it will be OK, but in some situation it can lead to deadlocks and other oddities, where a lot happens between the moment TrySetResult is invoked and the moment it returns.

Uh, what shall we do?

In some situations, we really don't know what may be await-ing the completion's Task. What we want is a way to tell .NET: don't be clever. Submit whatever is await-ing to the thread pool and return immediately. Run it all on separate tasks. It may be a little more expensive, but at least invoking TrySetResult would become deterministic.

There is an option to do this, but it is not a TrySetResult option. Instead, it needs to be specified when creating the completion souce. The class constructor has an overload that supports a TaskCreationOptions that can be used to indicate how the tasks resulting from TrySetResult are supposed to run.

Especially, TaskCreationOptions.RunContinuationsAsynchronously specifies that "continuations added to the current task [are] to be executed asynchronously." We can create our continuation as such:

var completion = new TaskCompletionSource<Response>(
    TaskCreationOptions.RunContinuationsAsynchronously
);

When TrySetResult is invoked, it submits all the continuations on the completion's Task (i.e. all code await-ing that task) on the thread pool for execution, instead of running them, and returns immediately.

Nice, why is this not the default?

Because in most cases, this is an overkill. Whenever TaskCompletionSource instances are used in simple and controlled environment, everything should be OK and it's better to let .NET optimize everything.

Nevertheless, it is good to be aware of this pitfall.