WithCancellation

src/UniTask.NetCore/IAsyncEnumerable.cs

@@ -1,4 +1,5 @@
-using System.Threading;
+using System.Runtime.InteropServices;
+using System.Threading;
 
 namespace Cysharp.Threading.Tasks
 {
@@ -17,4 +18,55 @@ namespace Cysharp.Threading.Tasks
     {
         UniTask DisposeAsync();
     }
+
+    public static class UniTaskAsyncEnumerableExtensions
+    {
+        public static UniTaskCancelableAsyncEnumerable<T> WithCancellation<T>(this IUniTaskAsyncEnumerable<T> source, CancellationToken cancellationToken)
+        {
+            return new UniTaskCancelableAsyncEnumerable<T>(source, cancellationToken);
+        }
+    }
+
+    [StructLayout(LayoutKind.Auto)]
+    public readonly struct UniTaskCancelableAsyncEnumerable<T>
+    {
+        private readonly IUniTaskAsyncEnumerable<T> enumerable;
+        private readonly CancellationToken cancellationToken;
+
+        internal UniTaskCancelableAsyncEnumerable(IUniTaskAsyncEnumerable<T> enumerable, CancellationToken cancellationToken)
+        {
+            this.enumerable = enumerable;
+            this.cancellationToken = cancellationToken;
+        }
+
+        public Enumerator GetAsyncEnumerator()
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+            return new Enumerator(enumerable.GetAsyncEnumerator(cancellationToken));
+        }
+
+        [StructLayout(LayoutKind.Auto)]
+        public readonly struct Enumerator
+        {
+            private readonly IUniTaskAsyncEnumerator<T> enumerator;
+
+            internal Enumerator(IUniTaskAsyncEnumerator<T> enumerator)
+            {
+                this.enumerator = enumerator;
+            }
+
+            public T Current => enumerator.Current;
+
+            public UniTask<bool> MoveNextAsync()
+            {
+                return enumerator.MoveNextAsync();
+            }
+
+
+            public UniTask DisposeAsync()
+            {
+                return enumerator.DisposeAsync();
+            }
+        }
+    }
 }

src/UniTask.NetCore/Linq/Range.cs

@@ -55,7 +55,7 @@ namespace Cysharp.Threading.Tasks.Linq
 
             public UniTask<bool> MoveNextAsync()
             {
-                if (cancellationToken.IsCancellationRequested) return CompletedTasks.False;
+                cancellationToken.ThrowIfCancellationRequested();
 
                 current++;
 

src/UniTask.NetCore/Linq/Repeat.cs

@@ -50,7 +50,7 @@ namespace Cysharp.Threading.Tasks.Linq
 
             public UniTask<bool> MoveNextAsync()
             {
-                if (cancellationToken.IsCancellationRequested) return CompletedTasks.False;
+                cancellationToken.ThrowIfCancellationRequested();
 
                 if (remaining-- != 0)
                 {

src/UniTask.NetCoreSandbox/Program.cs

@@ -32,17 +32,24 @@ namespace NetCoreSandbox
 
         static async Task Main(string[] args)
         {
+            var cts = new CancellationTokenSource();
 
 
+            await foreach (var item in UniTaskAsyncEnumerable.Range(1, 3).WithCancellation(cts.Token))
+            {
+                Console.WriteLine(item);
+                cts.Cancel();
+            }
 
 
-            await UniTaskAsyncEnumerable.Range(1, 3).ForEachAsync(x =>
+            /*
+            .ForEachAsync(x =>
            {
                if (x == 2) throw new Exception();
 
                Console.WriteLine(x);
            });
-
+           */