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

Fix problem with finally in UniTaskAsyncEnumerable.Create not being executed
Cysharp · Aug 31, 2023 · 4c3d693 · 4c3d693
2 parents d210e3d + b448680
commit 4c3d693
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 1 deletion.
diff --git a/src/UniTask.NetCoreTests/Linq/CreateTest.cs b/src/UniTask.NetCoreTests/Linq/CreateTest.cs
@@ -159,6 +159,30 @@ public async Task ASyncManually()
             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++)

diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/Create.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/Create.cs
@@ -52,6 +52,7 @@ public _Create(Func<IAsyncWriter<T>, CancellationToken, UniTask> create, Cancell
             public UniTask DisposeAsync()
             {
                 TaskTracker.RemoveTracking(this);
+                writer.Dispose();
                 return default;
             }
 
@@ -127,7 +128,7 @@ public void SetResult(T value)
             }
         }
 
-        sealed class AsyncWriter : IUniTaskSource, IAsyncWriter<T>
+        sealed class AsyncWriter : IUniTaskSource, IAsyncWriter<T>, IDisposable
         {
             readonly _Create enumerator;
 
@@ -137,6 +138,15 @@ public AsyncWriter(_Create enumerator)
             {
                 this.enumerator = enumerator;
             }
+
+            public void Dispose()
+            {
+                var status = core.GetStatus(core.Version);
+                if (status == UniTaskStatus.Pending)
+                {
+                    core.TrySetCanceled();
+                }
+            }            
 
             public void GetResult(short token)
             {