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 creatingindefinitelyRunningRange
variable with the default cancellation token, the optional one - Since our code execution is deferred, the
foreach
is calling theGetAsyncEnumerator
with the provided token from theWithCancellation
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:
- If the token on the enumeration is the default one, use the provided token
- 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
- 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.