Add WithCancellation for async enumerables (#21939)

1 files changed, 11 insertions, 2 deletions

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);
     }
 }