Allocation-free timeout on Task

Question

I'd like to add a timeout to ChannelReader.ReadyAsync. Here are two solutions I've found:

var cts = new CancellationTokenSource();
cts.CancelAfter(2000);
try {
  var data = chan.ReadAsync(cts.Token);
} catch (OperationCanceledException) {
  // timeout
}

var tasks = new Task[] { Task.Delay(2000), chan.ReadAsync(CancellationToken.None) };
var completedTask = await Task.WhenAny(tasks);
if (completedTask == tasks[0])
  // timeout
else
  var data = ((T)completedTask).Result;

However, both these solutions aren't allocation-free. The first allocates a CancellationTokenSource and the second one a Timer in the Task.Delay. Is there a way to make a similar code without any allocation?

EDIT 1: dotTrace output when using the first solution

Don't use the second one. If your timeout occurs, the `ReadAsync` task will still run while using its resources. Also before .NET 4.5 if `ReadAsync` failed with an exception, this unobserved exception would have terminated the process. As this behaviour can be enabled again on 4.5+, it's a really bad idea to use approach 2 in a library. — ckuri, Mar 17 '20 at 13:41
Why are you so worried about the allocation? Have you measured the performance of your code and determined that this allocation is a hot-spot? — Sean, Mar 17 '20 at 13:42
I don't have an answer, sorry. But I think your question is "incomplete", because `new Task[]` is also an allocation, also (in the other example) the use of an Exception is an allocation. It might also be worth investigating if the whole `async-await` mechanism is allocation-free (I wouldn't know; maybe it is, maybe it isn't). So if you're that worried about allocations, then you need to take **all** of them in consideration, and not just a few obvious ones... — Peter B, Mar 17 '20 at 13:43
@ckuri Thanks, I didn't realize that. I've edited my question to include dotTrace output with the first solution. Most of my allocations come from `NonBlockingBufferedTransport.SendBuffers` which contains the solution one code. — Greg, Mar 17 '20 at 15:55
Although timer creation [has been optimized](https://github.com/dotnet/coreclr/pull/14527), optimizing timer management would be much trickier, let alone doing it allocation-free (which at the very least requires some sort of pooling). It's something you'd probably have to do yourself. — Jeroen Mostert, Mar 17 '20 at 16:05

Greg · Accepted Answer · 2022-07-24T12:05:35.893

Thanks for your answers, they made me think again of exactly I was looking for: reuse CancellationTokenSource. Once a CancellationTokenSource is cancelled, you can't reuse it. But in my case, ChannelReader.ReadAsync would most of the time return before the timeout triggers so I've used the fact that CancelAfter doesn't recreate a timer the second time you call it to avoid cancelling the CancellationTokenSource after ChannelReader.ReadAsync returns.

var timeoutCancellation = new CancellationTokenSource();

while (true)
{
    if (timeoutCancellation.IsCancellationRequested)
    {
        timeoutCancellation.Dispose();
        timeoutCancellation = new CancellationTokenSource();
    }

    T data;
    try
    {
        timeoutCancellation.CancelAfter(2000);
        data = await _queue.Reader.ReadAsync(timeoutCancellation.Token);
        // make sure it doesn't get cancelled so it can be reused in the next iteration
        // Timeout.Infinite won't work because it would delete the underlying timer
        timeoutCancellation.CancelAfter(int.MaxValue);
    }
    catch (OperationCanceledException) // timeout reached
    {
        // handle timeout
        continue;
    }

    // process data
}

This is not allocation-free but it reduces drastically the number of allocated objects.

EDIT 1: In .NET 6 you can also use CancellationTokenSource.TryReset to reuse a CancellationTokenSource.

Theodor Zoulias · Answer 2 · 2021-08-04T10:10:56.443

The ChannelReader<T> class has a ReadAllAsync method that exposes the data of the reader as an IAsyncEnumerable<T>. Below is an overload of this method that accepts also a timeout parameter. This parameter has the effect that in case the reader fails to emit any items during the specified span of time, a TimeoutException is thrown.

For reducing the allocations it uses the same clever technique from Greg's answer, with a single CancellationTokenSource that is rescheduled for cancellation after each iteration. After some thought I removed the line CancelAfter(int.MaxValue) because it's probably more harmful than useful in the general case, but I might be wrong.

public static async IAsyncEnumerable<TSource> ReadAllAsync<TSource>(
    this ChannelReader<TSource> source, TimeSpan timeout,
    [EnumeratorCancellation] CancellationToken cancellationToken = default)
{
    while (true)
    {
        using var cts = CancellationTokenSource
            .CreateLinkedTokenSource(cancellationToken);
        cts.CancelAfter(timeout);
        while (true)
        {
            try
            {
                if (!await source.WaitToReadAsync(cts.Token).ConfigureAwait(false))
                    yield break;
            }
            catch (OperationCanceledException)
            {
                cancellationToken.ThrowIfCancellationRequested();
                throw new TimeoutException();
            }
            while (source.TryRead(out var item))
            {
                yield return item;
                cancellationToken.ThrowIfCancellationRequested();
            }

            cts.CancelAfter(timeout);
            // It is possible that the CTS timed-out during the yielding
            if (cts.IsCancellationRequested) break; // Start a new loop with a new CTS
        }
    }
}

As a side note, the System.Interactive.Async package includes a Timeout operator with the signature shown below, that could be used in combination with the built-in ReadAllAsync, ~~and provide the same functionality with the above implementation~~. That method is not optimized for low allocations though.

public static IAsyncEnumerable<TSource> Timeout<TSource>(
    this IAsyncEnumerable<TSource> source, TimeSpan timeout);

Note: In retrospect the ReadAllAsync().Timeout() idea is dangerous, because the ReadAllAsync is a consuming method. In other words enumerating it has the side-effect of removing items from the channel. The Timeout operator is unaware of what's happening internally in the source sequence, so a timeout occurring at an unfortunate moment could cause an item to be lost. This leaves the implementation on the top as the only robust solution to the problem (in the scope of this answer).

Allocation-free timeout on Task

2 Answers2