Merge pull request #484 from Cysharp/hadashiA/fix-async-enumerable

Fix problem with finally in UniTaskAsyncEnumerable.Create not being executed

src/UniTask.NetCoreTests/Linq/CreateTest.cs

@@ -159,6 +159,30 @@ namespace NetCoreTests.Linq
             list.Should().Equal(100, 200, 300, 400);
         }
 
+        [Fact]
+        public async Task AwaitForeachBreak()
+        {
+            var finallyCalled = false;
+            var enumerable = UniTaskAsyncEnumerable.Create<int>(async (writer, _) =>
+            {
+                try
+                {
+                    await writer.YieldAsync(1);
+                }
+                finally
+                {
+                    finallyCalled = true;
+                }
+            });
+
+            await foreach (var x in enumerable)
+            {
+                x.Should().Be(1);
+                break;
+            }
+            finallyCalled.Should().BeTrue();
+        }
+
         async IAsyncEnumerable<int> Range(int from, int count)
         {
             for (int i = 0; i < count; i++)

src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/Create.cs

@@ -52,6 +52,7 @@ namespace Cysharp.Threading.Tasks.Linq
             public UniTask DisposeAsync()
             {
                 TaskTracker.RemoveTracking(this);
+                writer.Dispose();
                 return default;
             }
 
@@ -127,7 +128,7 @@ namespace Cysharp.Threading.Tasks.Linq
             }
         }
 
-        sealed class AsyncWriter : IUniTaskSource, IAsyncWriter<T>
+        sealed class AsyncWriter : IUniTaskSource, IAsyncWriter<T>, IDisposable
         {
             readonly _Create enumerator;
 
@@ -138,6 +139,15 @@ namespace Cysharp.Threading.Tasks.Linq
                 this.enumerator = enumerator;
             }
             
+            public void Dispose()
+            {
+                var status = core.GetStatus(core.Version);
+                if (status == UniTaskStatus.Pending)
+                {
+                    core.TrySetCanceled();
+                }
+            }            
+
             public void GetResult(short token)
             {
                 core.GetResult(token);