Cancellation Tokens with IAsyncEnumerable

This article will remind us what cancellation tokens are and how to use them in C#. After that, we are going to talk about how we can use them with the IAsyncEnumerable interface and the cases when we should be careful while using them.

So, let’s start.

What are Cancellation Tokens in C#?

Let’s briefly introduce cancellation tokens in C#. It is good to handle them when dealing with asynchronous operations if something goes wrong. Whether it is a long-running operation or we want to stop its execution on purpose.

To have such a feature in our code, we can write a class:

public class CancellationToken
{
    public bool IsCancelled { get; set; }

    public void Cancel()
    {
        IsCancelled = true;
    }

    public void ThrowIfCancelled()
    {
        if (IsCancelled)
        {
            throw new OperationCanceledException();
        }
    }
}

Let’s say we have an indefinitely running method. We can easily cancel that method using our class, and our code won’t hang:

async Task IndefinitelyRunningTask(CancellationToken cancellationToken)
{
    while (true)
    {
        await Task.Delay(5000);
        cancellationToken.ThrowIfCancelled();
    }
}

As we can see, if we want to cancel it, we could call the Cancel method on our cancellation token.

Our implementation is not perfect, e.g., it is not thread-safe. Luckily we don’t need to continue our implementation as CLR already has a CancellationToken, perfect and ready for us.

Yet, it absences the Cancel method. That’s because cancellation is a concern of another class, CancellationTokenSource. This decoupling gives us a level of safety since a method with access to cancellation tokens cannot request its cancellation.

We already talked about using them in real-life scenarios within the HttpClient class. You can read more about it here: Canceling HTTP Requests in ASP.NET Core with CancellationToken.

Requesting the Cancellation Within the IAsyncEnumerable Interface

C# enables us to iterate through the collections asynchronously. You can read more about how and why here: IAsyncEnumerable with yield in C#.

Let’s write an indefinite method GetIndefinitelyRunningRangeAsync that is returning IAsyncEnumerable interface:

private static async IAsyncEnumerable<int> GetIndefinitelyRunningRangeAsync()
{
    while (true)
    {
        await Task.Delay(5000);
        yield return index++;
    }
}

And let’s write another method IndefinitelyRunningMethod where we will iterate through that enumeration using await foreach syntax:

public static async Task IndefinitelyRunningMethod()
{
    var indefinitelyRunningRange = GetIndefinitelyRunningRangeAsync();

    await foreach (int index in indefinitelyRunningRange)
    {
        // Do something with the index 
    }
}

When we compile our code, the compiler will generate the GetAsyncEnumerator call for us, and our method will look kind of like this after the compilation:

private static async Task CompilerGeneratingExample()
{
    var indefinitelyRunningRange = GetIndefinitelyRunningRangeAsync();

    IAsyncEnumerator<int> enumerator = indefinitelyRunningRange.GetAsyncEnumerator();

    try
    {
        while (await enumerator.MoveNextAsync()) { };
    }
    finally
    {
        if (enumerator != null)
        {
            await enumerator.DisposeAsync();
        }
    }
}

As we can see, we are not controlling the GetAsyncEnumerator method and its parameters. But we are still not in trouble since we can extend our GetIndefinitelyRunningRangeAsync method. It could accept an optional cancellation token, the same way as the GetAsyncEnumerator in the IAsyncEnumerable interface does:

private static async IAsyncEnumerable<int> GetIndefinitelyRunningRangeAsync(
    System.Threading.CancellationToken cancellationToken = default)
{
    int index = 0;
    while (true)
    {
        await Task.Delay(5000, cancellationToken);
        yield return index++;
    }
}

Great, we can as well extend our IndefinitelyRunningMethod to be able to cancel the enumeration:

public static async Task IndefinitelyRunningMethodCancelled()
{
    var cancellationTokenSource = new CancellationTokenSource();
    cancellationTokenSource.CancelAfter(7000);

    var indefinitelyRunningRange = GetIndefinitelyRunningRangeAsync(cancellationTokenSource.Token);

    await foreach (int index in indefinitelyRunningRange)
    {
        // Do something with the index 
    }
}

Since this method is no longer indefinitely running, we renamed it to IndefinitelyRunningMethodCancelled as well.

And if we run the project, we can see that TaskCanceledException is thrown:

That was easy. But what if we are not in control of how enumeration is created? What if we are dealing with some third-party library? What if some other developer encapsulated GetIndefinitelyRunningRangeAsync call in a wrapper method GetIndefinitelyRunningRangeWrapperAsync:

private static IAsyncEnumerable<int> GetIndefinitelyRunningRangeWrapperAsync()
{
    return GetIndefinitelyRunningRangeAsync();
}

Let’s add another method IndefinitelyRunningWrappedMethodCancelled where we utilize this method:

public static async Task IndefinitelyRunningWrappedMethodCancelled()
{
    var cancellationTokenSource = new CancellationTokenSource();
    cancellationTokenSource.CancelAfter(7000);

    var indefinitelyRunningRange = GetIndefinitelyRunningRangeWrapperAsync();

    await foreach (int index in indefinitelyRunningRange)
    {
        // Do something with the index 
    }
}

This method is almost the same as IndefinitleyRunningMethodCancelled, we are only setting up our indefinitelyRunningRange variable differently.

It is apparent that we can’t pass the token inside. Lucky for us, Microsoft provided us WithCancellation method.

WithCancellation Method Demystification

Microsoft’s developers added the WithCancellation extension method to support passing the cancellation token when the compiler generates the GetAsyncEnumerator call. Or even as in our case, when we are not in control of the enumeration at all.

As stated in the documentation, it is enough that we extend our IndefinitelyRunningWrappedMethodCancelled method as:

public static async Task IndefinitelyRunningWrappedMethodCancelled()
{
    var cancellationTokenSource = new CancellationTokenSource();
    cancellationTokenSource.CancelAfter(7000);

    var indefinitelyRunningRange = GetIndefinitelyRunningRangeWrapperAsync();
    await foreach (int index in indefinitelyRunningRange.WithCancellation(cancellationTokenSource.Token))
    {
        // Do something with the index
    }
}

But unfortunately, it is not enough. If we run the previous code, we are iterating through the enumeration indefinitely. Nothing validates a cancellation token state, and the exception is not thrown, although we managed to pass a cancellation token. At least, we think we are.

Let’s try to find out what has happened.

Decompilation of the WithCancellation Method Call

If we go to the implementation of the WithCancellation extension method, we can see that it is just a wrapper method that creates a new instance of the ConfiguredCancelableAsyncEnumerable class where the cancellation token is passed as an argument. It is used after in the GetAsyncEnumerator method where it is passed again in the GetAsyncEnumerator of the extended enumerable. So everything seems correct, yet our code is not working as expected.

Let’s dig deeper and decompile our code:

[DebuggerHidden]
IAsyncEnumerator<int> IAsyncEnumerable<int>.GetAsyncEnumerator(
    CancellationToken cancellationToken)
{
    Program.\u003CGetIndefinitelyRunningRangeAsync\u003Ed__1 runningRangeAsyncD1;
    if (this.\u003C\u003E1__state == -2 && this.\u003C\u003El__initialThreadId == Environment.CurrentManagedThreadId)
    {
        this.\u003C\u003E1__state = -3;
        this.\u003C\u003Et__builder = AsyncIteratorMethodBuilder.Create();
        this.\u003C\u003Ew__disposeMode = false;
        runningRangeAsyncD1 = this;
    }
    else
        runningRangeAsyncD1 = new Program.\u003CGetIndefinitelyRunningRangeAsync\u003Ed__1(-3);
        runningRangeAsyncD1.cancellationToken = this.\u003C\u003E3__cancellationToken;  
        return (IAsyncEnumerator<int>) runningRangeAsyncD1;
}

We can see that the GetAsyncEnumerator method is not using cancellationToken argument at all.

This is what is happening under the hood:

• When we are calling GetIndefinitelyRunningRangeAsync, we are creating indefinitelyRunningRange variable with the default cancellation token, the optional one
• Since our code execution is deferred, the foreach is calling the GetAsyncEnumerator with the provided token from the WithCancellation method
• GetAsyncEnumerator is ignoring that token, and it is using the default one

Were all that documentation reading and digging through the GitHub futile? Why did people from Microsoft give us an extension that doesn’t work?

Well, it wasn’t, and they didn’t.

[EnumeratorCancellation] Attribute Usage

Long story short, we are missing the EnumeratorCancellation attribute in the implementation of the GetIndefinitelyRunningRangeAsync method:

private static async IAsyncEnumerable<int> GetIndefinitelyRunningRangeAsync(
    [EnumeratorCancellation] System.Threading.CancellationToken cancellationToken = default)
{
    int index = 0;
    while (true)
    {
        await Task.Delay(5000, cancellationToken);
        yield return index++;
    }
}

Since we now know that we have two tokens, default one on the enumeration and one that the WithCancellation method is providing, compiler somehow needs to know which one to use.

That’s where the EnumeratorCancellation attribute shines. It tells the compiler to generate the code in the GetAsyncEnumerator with three use cases:

1. If the token on the enumeration is the default one, use the provided token
2. If the token on the enumeration is the same as the provided one or provided token is the default one, use the token on the enumeration
3. In any other case, combine the two tokens with the CreateLinkedTokenSource method

We will show what is happening with the decompiled code as well:

have[DebuggerHidden]
IAsyncEnumerator<int> IAsyncEnumerable<int>.GetAsyncEnumerator(
    System.Threading.CancellationToken cancellationToken)
{
    Program.\u003CGetIndefinitelyRunningRangeAsync\u003Ed__1 runningRangeAsyncD1;
    if (this.\u003C\u003E1__state == -2 && this.\u003C\u003El__initialThreadId == Environment.CurrentManagedThreadId)
    {
        this.\u003C\u003E1__state = -3;
        this.\u003C\u003Et__builder = AsyncIteratorMethodBuilder.Create();
        this.\u003C\u003Ew__disposeMode = false;
        runningRangeAsyncD1 = this;
    }
    else
        runningRangeAsyncD1 = new Program.\u003CGetIndefinitelyRunningRangeAsync\u003Ed__1(-3);
    if (this.\u003C\u003E3__cancellationToken.Equals(new System.Threading.CancellationToken()))
        runningRangeAsyncD1.cancellationToken = cancellationToken;
    else if (cancellationToken.Equals(this.\u003C\u003E3__cancellationToken) || cancellationToken.Equals(new System.Threading.CancellationToken()))
    {
        runningRangeAsyncD1.cancellationToken = this.\u003C\u003E3__cancellationToken;
    }
    else
    {
        this.\u003C\u003Ex__combinedTokens = CancellationTokenSource.CreateLinkedTokenSource(this.\u003C\u003E3__cancellationToken, cancellationToken);
        runningRangeAsyncD1.cancellationToken = this.\u003C\u003Ex__combinedTokens.Token;
    }
    return (IAsyncEnumerator<int>) runningRangeAsyncD1;
}

In the highlighted code we can see three use cases we mention above.

The cancellation token is really not ignored anymore, and the iteration in our IndefinitelyRunningWrappedMethodCancelled method is canceled.

Conclusion

To use cancellation tokens with the IAsyncEnumerable when we are not in control of the enumerable, we can use WithCancellation extension method.

We should be careful when our methods are accepting cancellation tokens. For them to work as expected, we should add the EnumeratorCancellation attribute.

Lucky for us, the compiler will give us a CS8425 warning to remind us.

Until the next one.

All the best.