How to cancel a TaskCompletionSource using a timeout

Question

I have the function that I call asynchronously using the await keyword:

public Task<StatePropertyEx> RequestStateForEntity(EntityKey entity, string propName)
{
    var tcs = new TaskCompletionSource<StateInfo>();
    try
    {
        var propInstance = BuildCacheKey(entity, propName);
        StateCacheItem cacheItem;
        if (_stateCache.TryGetValue(propInstance, out cacheItem))
        {
            tcs.SetResult( new StateInfo (cacheItem.State.Name, cacheItem.State.Value) );
            return tcs.Task;
        }

        //state not found in local cache so save the tcs for later and request the state
        var cacheKey = BuildCacheKey(entity, propName);
       _stateRequestItemList.TryAdd(cacheKey, new StateRequestItem(entity, propName, tcs));

        _evtClient.SubmitStateRequest(entity, propName);

        return tcs.Task;
    }
    catch (Exception ex)
    {
        tcs.SetException(ex);
        return tcs.Task;
    }
}

The function has a look to see if it has the information it needs and if it does it returns it. If it doesn’t have the details it sends a request out which should eventually come in as an event. At that point my code (not shown) finds the stored TaskCompletionSource item, sets the result and returns it. This all works fine but I have now been asked to consider a situation where a reply may never be returned when I request state via the “_evtClient.SubmitStateRequest(entity, propName);” line. I need to implement some sort of timeout mechanism so I can cancel the TCS task so the function caller can fail gracefully. I’ve been looking on SO and the internet and can’t find anything that looks right. I’m now not sure if I need to restructure the above code in a different way. Can anyone advise or point me to a similar scenario?

The code that calls the above function can call it in a single hit like this:

var stateProperty = await RequestStateForEntity(key, stateName);

or in a batch, like this:

await
    Task.WhenAll(
        stateDefinitions.Select(stateDefinition => stateDefinition.Name)
            .Select(
                stateName =>
                    Task.Factory.StartNew(
                        async () => results.Add(await RequestStateForEntity(key, stateName)))
                        .Unwrap())
            .ToArray());

Answer 1

First off, what you really want to enable is cancellation. The fact that the cancellation comes from a timeout is just a footnote.

.NET has some great built-in support for cancellation, and the Task-based Asynchronous Pattern prescribes how to use it.

Essentially, what you want to do is take a CancellationToken:

Task<StatePropertyEx> RequestStateForEntity(EntityKey entity, string propName,
    CancellationToken cancellationToken);

Next, you want to respond when that token is signaled. Ideally, you would want to just pass the CancellationToken down to the _evtClient so that the request is truly cancelled:

_evtClient.SubmitStateRequest(entity, propName, cancellationToken);

This is the normal way of implementing cancellation, and it works great if SubmitStateRequest already understands cancellation. Often the event arguments have a flag indicating cancellation (e.g., AsyncCompletedEventArgs.Cancelled). If at all possible, use this approach (i.e., change _evtClient as necessary to support cancellation).

But sometimes this just isn't possible. In this case, you can choose to pretend to support cancellation. What you're actually doing is just ignoring the request if it completes after it was cancelled. This is not the most ideal situation but sometimes you have no choice.

Personally, I don't really like this kind of approach since it makes the API "lie": the method signature claims to support cancellation but it really is just faking it. So first, I recommend documenting this. Put in a code comment apology explaining that _evtClient doesn't support cancellation, and the "cancellation" is actually just pretend cancellation.

Then, you'll need to hook into the CancellationToken yourself, after the state request item is in the list but before the actual request is sent:

var item = new StateRequestItem(entity, propName, tcs);
_stateRequestItemList.TryAdd(cacheKey, item);
item.CancellationRegistration = cancellationToken.Register(() =>
{
  StateRequestItem cancelledItem;
  if (!_stateRequestItemList.TryRemove(cacheKey, out cancelledItem))
    return;
  cancelledItem.TaskCompletionSource.TrySetCanceled();
});
_evtClient.SubmitStateRequest(entity, propName);

Finally, you'll need to update your event handler completion code (not shown) to ignore the situation where the state request item has already been removed, and to dispose the CancellationRegistration if the state request item is found.

---

Once your method supports cancellation, then it's easy to cancel via a timer:

var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
CancellationToken token = cts.Token;

or from any other kind of situation. Say, if the user cancels whatever (s)he's doing. Or if another part of the system decides it doesn't need that data anymore. Once your code supports cancellation, it can handle cancellation for any reason.