How to force an IAsyncEnumerable to respect a CancellationToken

Question

I have an async iterator method that produces an IAsyncEnumerable<int> (a stream of numbers), one number every 200 msec. The caller of this method consumes the stream, but wants to stop the enumeration after 1000 msec. So a CancellationTokenSource is used, and the token is passed as an argument to the WithCancellation extension method. But the token is not respected. The enumeration continues until all the numbers are consumed:

static async IAsyncEnumerable<int> GetSequence()
{
    for (int i = 1; i <= 10; i++)
    {
        await Task.Delay(200);
        yield return i;
    }
}

var cts = new CancellationTokenSource(1000);
await foreach (var i in GetSequence().WithCancellation(cts.Token))
{
    Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} > {i}");
}

Output:

12:55:17.506 > 1
12:55:17.739 > 2
12:55:17.941 > 3
12:55:18.155 > 4
12:55:18.367 > 5
12:55:18.570 > 6
12:55:18.772 > 7
12:55:18.973 > 8
12:55:19.174 > 9
12:55:19.376 > 10

The expected output is a TaskCanceledException to occur after number 5. It seems that I have misunderstood what the WithCancellation is actually doing. The method just passes the supplied token to the iterator method, if that method accepts one. Otherwise, like with the method GetSequence() in my example, the token is ignored. I suppose that the solution in my case is to interrogate manually the token inside the body of the enumeration:

var cts = new CancellationTokenSource(1000);
await foreach (var i in GetSequence())
{
    cts.Token.ThrowIfCancellationRequested();
    Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} > {i}");
}

This is simple and works well. But in any case I wonder if it would be possible to create an extension method that does what I expected the WithCancellation to do, to bake the token inside the ensuing enumeration. This is the signature of the needed method:

public static IAsyncEnumerable<T> WithEnforcedCancellation<T>(
    this IAsyncEnumerable<T> source, CancellationToken cancellationToken)
{
    // Is it possible?
}

Answer 1

IAsyncEnumerable explicitly provides for this mechanism with the EnumeratorCancellation attribute:

static async IAsyncEnumerable<int> GetSequence([EnumeratorCancellation] CancellationToken ct = default) {
    for (int i = 1; i <= 10; i++) {
        ct.ThrowIfCancellationRequested();
        await Task.Delay(200);    // or `Task.Delay(200, ct)` if this wasn't an example
        yield return i;
    }
}

In fact, the compiler is helpful enough to issue a warning if you give the method a CancellationToken parameter, but do not add the attribute.

Note that the token passed to .WithCancellation will override any local token passed to the method. The specs have the details on this.

Of course, this will still only work if the enumeration actually accepts a CancellationToken -- but the fact that cancellation only really works if done cooperatively is true of any async work. Yeldar's answer is good for "forcing" some measure of cancellation into an enumerable that doesn't support it, but the preferred solution should be to modify the enumeration to support cancellation by itself -- the compiler does everything to help you out.

Answer 2

You can just extract your logic into an extension method like this:

public static async IAsyncEnumerable<T> WithEnforcedCancellation<T>(
    this IAsyncEnumerable<T> source, CancellationToken cancellationToken)
{
    if (source == null)
        throw new ArgumentNullException(nameof(source));

    cancellationToken.ThrowIfCancellationRequested();

    await foreach (var item in source)
    {
        cancellationToken.ThrowIfCancellationRequested();
        yield return item;
    }
}

Answer 3

I think it's important to reiterate that you're not supposed to do this. It's always better to make the async method support cancellation tokens, then cancellation is immediate as you would expect. If that's impossible, I still recommend trying one of the other answers before trying this one.

With that said, if you can't add cancellation support to the async method, and you absolutely do need immediate termination of the foreach, then you can hack your way around it.

One trick is to use Task.WhenAny with two arguments:

1. the task you get from IAsyncEnumerator.MoveNextAsync()
2. another task that does support cancellation

Here's the short version

// Start the 'await foreach' without the new syntax
// because we need access to the ValueTask returned by MoveNextAsync()
var enumerator = source.GetAsyncEnumerator(cancellationToken);

// Combine MoveNextAsync() with another Task that can be awaited indefinitely,
// until it throws OperationCanceledException
var untilCanceled = UntilCanceled(cancellationToken);
while (await await Task.WhenAny(enumerator.MoveNextAsync().AsTask(), untilCanceled))
{
    yield return enumerator.Current;
}

Long version with ConfigureAwait(false) and DisposeAsync() for completeness, should work if you run it locally.

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;

public static class AsyncStreamHelper
{
    public static async IAsyncEnumerable<T> WithEnforcedCancellation<T>(this IAsyncEnumerable<T> source, [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        if (source == null)
            throw new ArgumentNullException(nameof(source));
        cancellationToken.ThrowIfCancellationRequested();

        // Start the 'await foreach' without the new syntax
        // because we need access to the ValueTask returned by MoveNextAsync()
        var enumerator = source.GetAsyncEnumerator(cancellationToken);
        Task<bool> moveNext = null;

        // Combine MoveNextAsync() with another Task that can be awaited indefinitely,
        // until it throws OperationCanceledException
        var untilCanceled = UntilCanceled(cancellationToken);
        try
        {
            while (
                await (
                    await Task.WhenAny(
                        (
                            moveNext = enumerator.MoveNextAsync().AsTask()
                        ),
                        untilCanceled
                    ).ConfigureAwait(false)
                )
            )
            {
                yield return enumerator.Current;
            }
        }
        finally
        {
            if (moveNext != null && !moveNext.IsCompleted)
            {
                // Disable warning CS4014 "Because this call is not awaited, execution of the current method continues before the call is completed"
#pragma warning disable 4014 // This is the behavior we want!

                moveNext.ContinueWith(async _ =>
                {
                    await enumerator.DisposeAsync();
                }, TaskScheduler.Default);
#pragma warning restore 4014
            }
            else if (enumerator != null)
            {
                await enumerator.DisposeAsync();
            }
        }
    }

    private static Task<bool> UntilCanceled(CancellationToken cancellationToken)
    {
        // This is just one possible implementation... feel free to swap out for something else
        return new Task<bool>(() => true, cancellationToken);
    }
}

public class Program
{
    public static async Task Main()
    {
        var cts = new CancellationTokenSource(500);
        var stopwatch = Stopwatch.StartNew();
        try
        {
            await foreach (var i in GetSequence().WithEnforcedCancellation(cts.Token))
            {
                Console.WriteLine($"{stopwatch.Elapsed:m':'ss'.'fff} > {i}");
            }
        }
        catch (OperationCanceledException)
        {
            Console.WriteLine($"{stopwatch.Elapsed:m':'ss'.'fff} > Canceled");
        }
    }

    static async IAsyncEnumerable<int> GetSequence()
    {
        for (int i = 1; i <= 10; i++)
        {
            await Task.Delay(200);
            yield return i;
        }
    }
}

Caveats

The enumerator returns a ValueTask for improved performance (uses fewer allocations than regular Tasks), but a ValueTask cannot be used with Task.WhenAny(), so AsTask() is used which degrades performance by introducing allocation overhead.

The enumerator can only be disposed if the most recent MoveNextAsync() is completed. It's more likely that the Task is still running when cancellation is requested. That's why I added another call to DisposeAsync in a continuation task.

In this scenario, the enumerator is not yet disposed when the WithEnforcedCancellation() method exits. It will be disposed some indeterminate amount of time after the enumeration is abandoned. If DisposeAsync() throws an exception , the exception will be lost. It cannot bubble up the call stack, because there is no call stack.