Add WithCancellation for async enumerables (#21939)

4 files changed, 90 insertions, 51 deletions

diff --git a/src/System.Private.CoreLib/shared/System.Private.CoreLib.Shared.projitems b/src/System.Private.CoreLib/shared/System.Private.CoreLib.Shared.projitems
index bd67504204..5700d5569e 100644
--- a/src/System.Private.CoreLib/shared/System.Private.CoreLib.Shared.projitems
+++ b/src/System.Private.CoreLib/shared/System.Private.CoreLib.Shared.projitems
@@ -507,7 +507,7 @@
     <Compile Include="$(MSBuildThisFileDirectory)System\Runtime\CompilerServices\CompilationRelaxationsAttribute.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Runtime\CompilerServices\CompilerGeneratedAttribute.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Runtime\CompilerServices\CompilerGlobalScopeAttribute.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)System\Runtime\CompilerServices\ConfiguredAsyncEnumerable.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)System\Runtime\CompilerServices\ConfiguredCancelableAsyncEnumerable.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Runtime\CompilerServices\ConfiguredValueTaskAwaitable.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Runtime\CompilerServices\CustomConstantAttribute.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Runtime\CompilerServices\DateTimeConstantAttribute.cs" />
diff --git a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ConfiguredAsyncEnumerable.cs b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ConfiguredAsyncEnumerable.cs
deleted file mode 100644
index 25470656e5..0000000000
--- a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ConfiguredAsyncEnumerable.cs
+++ /dev/null
@@ -1,48 +0,0 @@
-// Licensed to the .NET Foundation under one or more agreements.
-// The .NET Foundation licenses this file to you under the MIT license.
-// See the LICENSE file in the project root for more information.
-
-using System.Collections.Generic;
-using System.Runtime.InteropServices;
-using System.Threading;
-using System.Threading.Tasks;
-using System.Threading.Tasks.Sources;
-
-namespace System.Runtime.CompilerServices
-{
-    [StructLayout(LayoutKind.Auto)]
-    public readonly struct ConfiguredAsyncEnumerable<T>
-    {
-        private readonly IAsyncEnumerable<T> _enumerable;
-        private readonly bool _continueOnCapturedContext;
-
-        internal ConfiguredAsyncEnumerable(IAsyncEnumerable<T> enumerable, bool continueOnCapturedContext)
-        {
-            _enumerable = enumerable;
-            _continueOnCapturedContext = continueOnCapturedContext;
-        }
-
-        public Enumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) =>
-            new Enumerator(_enumerable.GetAsyncEnumerator(cancellationToken), _continueOnCapturedContext);
-
-        public readonly struct Enumerator
-        {
-            private readonly IAsyncEnumerator<T> _enumerator;
-            private readonly bool _continueOnCapturedContext;
-
-            internal Enumerator(IAsyncEnumerator<T> enumerator, bool continueOnCapturedContext)
-            {
-                _enumerator = enumerator;
-                _continueOnCapturedContext = continueOnCapturedContext;
-            }
-
-            public ConfiguredValueTaskAwaitable<bool> MoveNextAsync() =>
-                _enumerator.MoveNextAsync().ConfigureAwait(_continueOnCapturedContext);
-
-            public T Current => _enumerator.Current;
-
-            public ConfiguredValueTaskAwaitable DisposeAsync() =>
-                _enumerator.DisposeAsync().ConfigureAwait(_continueOnCapturedContext);
-        }
-    }
-}
diff --git a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ConfiguredCancelableAsyncEnumerable.cs b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ConfiguredCancelableAsyncEnumerable.cs
new file mode 100644
index 0000000000..4f1bca7169
--- /dev/null
+++ b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ConfiguredCancelableAsyncEnumerable.cs
@@ -0,0 +1,78 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Collections.Generic;
+using System.Runtime.InteropServices;
+using System.Threading;
+
+namespace System.Runtime.CompilerServices
+{
+    /// <summary>Provides an awaitable async enumerable that enables cancelable iteration and configured awaits.</summary>
+    [StructLayout(LayoutKind.Auto)]
+    public readonly struct ConfiguredCancelableAsyncEnumerable<T>
+    {
+        private readonly IAsyncEnumerable<T> _enumerable;
+        private readonly CancellationToken _cancellationToken;
+        private readonly bool _continueOnCapturedContext;
+
+        internal ConfiguredCancelableAsyncEnumerable(IAsyncEnumerable<T> enumerable, bool continueOnCapturedContext, CancellationToken cancellationToken)
+        {
+            _enumerable = enumerable;
+            _continueOnCapturedContext = continueOnCapturedContext;
+            _cancellationToken = cancellationToken;
+        }
+
+        /// <summary>Configures how awaits on the tasks returned from an async iteration will be performed.</summary>
+        /// <param name="continueOnCapturedContext">Whether to capture and marshal back to the current context.</param>
+        /// <returns>The configured enumerable.</returns>
+        /// <remarks>This will replace any previous value set by <see cref="ConfigureAwait(bool)"/> for this iteration.</remarks>
+        public ConfiguredCancelableAsyncEnumerable<T> ConfigureAwait(bool continueOnCapturedContext) =>
+            new ConfiguredCancelableAsyncEnumerable<T>(_enumerable, continueOnCapturedContext, _cancellationToken);
+
+        /// <summary>Sets the <see cref="CancellationToken"/> to be passed to <see cref="IAsyncEnumerable{T}.GetAsyncEnumerator(CancellationToken)"/> when iterating.</summary>
+        /// <param name="cancellationToken">The <see cref="CancellationToken"/> to use.</param>
+        /// <returns>The configured enumerable.</returns>
+        /// <remarks>This will replace any previous <see cref="CancellationToken"/> set by <see cref="WithCancellation(CancellationToken)"/> for this iteration.</remarks>
+        public ConfiguredCancelableAsyncEnumerable<T> WithCancellation(CancellationToken cancellationToken) =>
+            new ConfiguredCancelableAsyncEnumerable<T>(_enumerable, _continueOnCapturedContext, cancellationToken);
+
+        public Enumerator GetAsyncEnumerator() =>
+            // as with other "configured" awaitable-related type in CompilerServices, we don't null check to defend against
+            // misuse like `default(ConfiguredCancelableAsyncEnumerable<T>).GetAsyncEnumerator()`, which will null ref by design.
+            new Enumerator(_enumerable.GetAsyncEnumerator(_cancellationToken), _continueOnCapturedContext);
+
+        /// <summary>Provides an awaitable async enumerator that enables cancelable iteration and configured awaits.</summary>
+        [StructLayout(LayoutKind.Auto)]
+        public readonly struct Enumerator
+        {
+            private readonly IAsyncEnumerator<T> _enumerator;
+            private readonly bool _continueOnCapturedContext;
+
+            internal Enumerator(IAsyncEnumerator<T> enumerator, bool continueOnCapturedContext)
+            {
+                _enumerator = enumerator;
+                _continueOnCapturedContext = continueOnCapturedContext;
+            }
+
+            /// <summary>Advances the enumerator asynchronously to the next element of the collection.</summary>
+            /// <returns>
+            /// A <see cref="ConfiguredValueTaskAwaitable{Boolean}"/> that will complete with a result of <c>true</c>
+            /// if the enumerator was successfully advanced to the next element, or <c>false</c> if the enumerator has
+            /// passed the end of the collection.
+            /// </returns>
+            public ConfiguredValueTaskAwaitable<bool> MoveNextAsync() =>
+                _enumerator.MoveNextAsync().ConfigureAwait(_continueOnCapturedContext);
+
+            /// <summary>Gets the element in the collection at the current position of the enumerator.</summary>
+            public T Current => _enumerator.Current;
+
+            /// <summary>
+            /// Performs application-defined tasks associated with freeing, releasing, or
+            /// resetting unmanaged resources asynchronously.
+            /// </summary>
+            public ConfiguredValueTaskAwaitable DisposeAsync() =>
+                _enumerator.DisposeAsync().ConfigureAwait(_continueOnCapturedContext);
+        }
+    }
+}
diff --git a/src/System.Private.CoreLib/shared/System/Threading/Tasks/TaskExtensions.cs b/src/System.Private.CoreLib/shared/System/Threading/Tasks/TaskExtensions.cs
index 6a2b82077d..9192b1433e 100644
--- a/src/System.Private.CoreLib/shared/System/Threading/Tasks/TaskExtensions.cs
+++ b/src/System.Private.CoreLib/shared/System/Threading/Tasks/TaskExtensions.cs
@@ -53,8 +53,17 @@ namespace System.Threading.Tasks
         /// <param name="source">The source enumerable being iterated.</param>
         /// <param name="continueOnCapturedContext">Whether to capture and marshal back to the current context.</param>
         /// <returns>The configured enumerable.</returns>
-        public static ConfiguredAsyncEnumerable<T> ConfigureAwait<T>(
+        public static ConfiguredCancelableAsyncEnumerable<T> ConfigureAwait<T>(
             this IAsyncEnumerable<T> source, bool continueOnCapturedContext) =>
-            new ConfiguredAsyncEnumerable<T>(source, continueOnCapturedContext);
+            new ConfiguredCancelableAsyncEnumerable<T>(source, continueOnCapturedContext, cancellationToken: default);
+
+        /// <summary>Sets the <see cref="CancellationToken"/> to be passed to <see cref="IAsyncEnumerable{T}.GetAsyncEnumerator(CancellationToken)"/> when iterating.</summary>
+        /// <typeparam name="T">The type of the objects being iterated.</typeparam>
+        /// <param name="source">The source enumerable being iterated.</param>
+        /// <param name="cancellationToken">The <see cref="CancellationToken"/> to use.</param>
+        /// <returns>The configured enumerable.</returns>
+        public static ConfiguredCancelableAsyncEnumerable<T> WithCancellation<T>(
+            this IAsyncEnumerable<T> source, CancellationToken cancellationToken) =>
+            new ConfiguredCancelableAsyncEnumerable<T>(source, continueOnCapturedContext: true, cancellationToken);
     }
 }