diff --git a/src/UniTask.NetCoreTests/Linq/CreateTest.cs b/src/UniTask.NetCoreTests/Linq/CreateTest.cs index 17ba22f..e1319eb 100644 --- a/src/UniTask.NetCoreTests/Linq/CreateTest.cs +++ b/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(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 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 index 61b7afd..fa34774 100644 --- a/src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/Create.cs +++ b/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 + sealed class AsyncWriter : IUniTaskSource, IAsyncWriter, IDisposable { readonly _Create enumerator; @@ -137,6 +138,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) {