Proposal: IAsyncEnumerable<T>.WithCancellation extension method

[stephentoub]

Background

Per recent design revisions, IAsyncEnumerable<T>.GetAsyncEnumerator will now accept a CancellationToken. To get a cancellation token into the GetAsyncEnumerator call, code can use GetAsyncEnumerator explicitly, e.g.

IAsyncEnumerator<T> enumerator = enumerable.GetAsyncEnumerator(cancellationToken);
try
{
     while (await enumerator.MoveNextAsync())
     {
         T current = enumerator.Current;
         …
     }
}
finally
{
    await enumerator.DisposeAsync();
}

To use await foreach, as there’s no language-provided facility for passing a CancellationToken into the compiler-generated GetAsyncEnumerator call, to get a CancellationToken in we need to take advantage of the pattern-matching support await foreach supplies. By returning a struct that exposes the right pattern, we can smuggle a CancellationToken in, stored in that struct, such that when the compiler-generated code invokes that struct’s GetAsyncEnumerator(), it in turn calls GetAsyncEnumerator(_cancellationToken) on the underlying enumerable.

To make that simple, we should add a WithCancellation extension method on IAsyncEnumerable<T> that accepts the CancellationToken, allowing you to write:

await foreach (T item in enumerable.WithCancellation(cancellationToken))
{
    …
}

with that generating the equivalent of:

var enumerator = enumerable.WithCancellation(cancellationToken).GetAsyncEnumerator();
try
{
     while (await enumerator.MoveNextAsync())
     {
         T current = enumerator.Current;
         …
     }
}
finally
{
    await enumerator.DisposeAsync();
}

However, we also need to deal with the ability to ConfigureAwait(false) an enumerable. We’ve already added the ConfigureAwait(bool) extension method for IAsyncEnumerable<T>, and we need to be able to allow developers to use both it and WithCancellation. Since both of these would return structs that by design do not implement IAsyncEnumerable<T>, we need another mechanism to allow chaining them, e.g.

await foreach (T item in enumerable.WithCancellation(cancellationToken).ConfigureAwait(false))
{
    …
}

Proposal

As part of adding the async-related interfaces, we also added a ConfigureAwait extension method and a new ConfiguredAsyncEnumerable struct type. We replace those with a slightly updated version:

-namespace System.Threading.Tasks
-{
-    public static class TaskExtensions // previously existing class
-    {
-        public static ConfiguredAsyncEnumerable<T> ConfigureAwait<T>(this IAsyncEnumerable<T> source, bool continueOnCapturedContext);
-    }
-}
-
-namespace System.Runtime.CompilerServices
-{
-    [StructLayout(LayoutKind.Auto)]
-    public readonly struct ConfiguredAsyncEnumerable<T>
-    {
-        public Enumerator GetAsyncEnumerator(CancellationToken cancellationToken = default);
-        public readonly struct Enumerator
-        {
-            public ConfiguredValueTaskAwaitable<bool> MoveNextAsync();
-            public T Current { get; }
-            public ConfiguredValueTaskAwaitable DisposeAsync();
-        }
-    }
-}
+namespace System.Threading.Tasks
+{
+  public static class TaskExtensions // previously existing class
+    {
+        public static ConfiguredCancelableAsyncEnumerable<T> WithCancellation<T>(this IAsyncEnumerable<T> source, CancellationToken cancellationToken);
+        public static ConfiguredCancelableAsyncEnumerable<T> ConfigureAwait<T>(this IAsyncEnumerable<T> source, bool continueOnCapturedContext);
+    }
+}
+
+namespace System.Runtime.CompilerServices
+{
+   [StructLayout(LayoutKind.Auto)]
+   public readonly struct ConfiguredCancelableAsyncEnumerable<T>
+    {
+        public ConfiguredCancelableAsyncEnumerable<T> WithCancellation(CancellationToken cancellationToken);
+        public ConfiguredCancelableAsyncEnumerable<T> ConfigureAwait(bool continueOnCapturedContext);
+
+        public Enumerator GetAsyncEnumerator();
+        public readonly struct Enumerator
+        {
+            public ConfiguredValueTaskAwaitable<bool> MoveNextAsync();
+            public T Current { get; }
+            public ConfiguredValueTaskAwaitable DisposeAsync();
+        }
+    }
+}

This let’s you write:

await foreach (T item in source) { … }
await foreach (T item in source.ConfigureAwait(false)) { … }
await foreach (T item in source.WithCancellation(token)) { … }
await foreach (T item in source.ConfigureAwait(false).WithCancellation(token)) { … }
await foreach (T item in source.WithCancellation(token).ConfigureAwait(false)) { … }

It of course also lets you write something a bit non-sensical, like:

await foreach (T item in source.ConfigureAwait(false).ConfigureAwait(true).WithCancellation(token1).WithCancellation(token2)) { … }

and in such a case we would just use the last value specified, e.g. this would end up being equivalent to:

await foreach (T item in source.ConfigureAwait(true).WithCancellation(token2)) { … }

Note that the GetAsyncEnumerator() on the struct is defined to be parameterless (this assumes the compiler's pattern-matching support allows that... if not, it'll be defined to take a token as well). Since this struct is only ever meant to be used via the provided extension methods with the await foreach construct, there’s little use for also accepting a CancellationToken into GetAsyncEnumerator(), and allowing that would just be confusing, as that argument will be ignored by the implementation.

cc: @jcouv, @MadsTorgersen, @terrajobst, @bartonjs

[jcouv]

1. The type returned by ConfigureAwait(...) is ConfiguredCancelableAsyncEnumerable<T>, so I didn't understand how you chain a call source.ConfigureAwait(false).WithCancellation(token).
2. I didn't understand the purpose or necessity of using a dedicated task-like type (ConfiguredValueTaskAwaitable<T>).
3. ConfiguredCancelableAsyncEnumerable<T>.GetAsyncEnumerator() seems to be missing the CancellationToken parameter.

[stephentoub]

so I didn't understand how you chain a call

There were supposed to be instance methods on the struct as well mirroring the extension methods. I've fixed that above.

seems to be missing the CancellationToken parameter

See the "Note that the..." paragraph at the end.

I didn't understand the purpose or necessity of using a dedicated task-like type

It's for ConfigureAwait, so that MoveNextAsync returns the awaitable that's aware of the configured context settings.

[jcouv]

Thanks.

seems to be missing the CancellationToken parameter

I was assuming that pattern-based foreach would match a ... GetAsyncEnumerator(CancellationToken) method, but not a ... GetAsyncEnumerator() method.
Supporting the latter only helps for folks implementing enumerator types by hand which provide their own token. Do you expect such types besides the one proposed here?

We can raise the question in LDM.

[svick]

I assume that .ConfigureAwait(false) was named like that (and not e.g. .ContinueOnCapturedContext(false)) so that it could be extended later with another parameter. I think this proposal is pretty much doing that, but it's not using the seemingly established pattern. Is that the right way to go?

I can see three alternatives:

1. Go all in on ConfigureAwait by having ConfigureAwait(bool continueOnCapturedContext, CancellationToken cancellationToken) and also some combination of overloads and optional parameters for convenience.
2. Stay consistent with Task by having ConfigureAwait(bool), but go with the more explicit WithCancellation(CancellationToken) for the new option (the above proposal).
3. Break consistency with Task and make both options explicit by having ContinueOnCapturedContext(bool) and WithCancellation(CancellationToken).

All three alternatives sound mostly reasonable to me, so I'm not sure which one of them would be best.

[stephentoub]

I think this proposal is pretty much doing that, but it's not using the seemingly established pattern

The token doesn't actually impact the behavior of the awaits, though.

[svick]

@stephentoub I think that from a certain point of view, it does, especially if it's decided that compiler-generated async enumerable will check the CancellationToken on every iteration: continueOnCapturedContext configures where the code after await will execute and CancellationToken is about configuring if the code after await will execute.

You could claim (and I think you are) that continueOnCapturedContext is about configuring the await foreach (or the awaits it expands into), while CancellationToken is about configuring the async enumerable, but I think that's a fairly subtle distinction to make and it might be too much about the mechanisms used to implement await foreach, instead of being about user intentions.

[stephentoub]

I think that's a fairly subtle distinction to make and it might be too much about the mechanisms used to implement await foreach, instead of being about user intentions.

If the CancellationToken were actually configuring the await, it would, for example, cancel the await operation even if the awaited thing hasn't completed yet. That's always the behavior we've talked about when considering whether to have a ConfigureAwait overload that accepted a CancellationToken, and that's very much not what happens here. I want to avoid that connotation.

[svick]

@stephentoub Yeah, that does make sense. Thanks for the explanation.

[bartdesmet]

Looks good. I ran into the exact problem when building various aggregation functions for LINQ to IAsyncEnumerable.

However, I wonder if it would be better for these to live in some AsyncEnumerableExtensions class in System.Collections.Generic next to the interface definition, to avoid having to include the seemingly unrelated System.Threading.Tasks namespace. In fact, for WithCancellation, the System.Threading.Tasks namespace makes even less sense, given that CancellationToken is defined in System.Threading.

I see this has been brought up before when introducing ConfigureAwait in https://github.com/dotnet/corefx/issues/32684 but I'm not sure if merely avoiding a potential clash with an AsyncEnumerable type for LINQ has been the reason to put it in System.Threading.Tasks. Maybe we can reserve that name and use AsyncEnumerableExtensions or something to avoid the clash (short of having a means to define extension methods without requiring to come up with a containing class).

It may not matter all that much, given that VS offers to include the namespace when trying to use an extension method (and System.Threading.Tasks is in the template for new .cs files), but it may be good to give it an extra consideration before shipping. Of course, going the way of sticking extension methods next to the original type would also imply some System.AsyncDisposableExtensions type for IAsyncDisposable's ConfigureAwait, so one could argue it gets unwieldy.

[jcouv]

As mentioned by @yyjdelete here, the DisposeAsync in ConfiguredCancelableAsyncEnumerable.Enumerator will not be called.
After the recent LDM decision, we only bind Dispose instance methods in foreach for ref structs, pattern-based. Since ref structs are disallowed in async methods, await foreach never binds to DisposeAsync instance methods, pattern-based. (Correction: foreach and await foreach never bind pattern-based disposal at the moment, only using binds pattern-based disposal, and that's restricted to ref structs)

await foreach only binds via IAsyncDIsposable interface.

[stephentoub]

@jcouv, not sure how I missed that in the ldm discussion, but I don't think that can be the answer. How does ConfigureAwait work in that model?