Merge 019f8aaf30 into 4c3d6938ed

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;
 
@@ -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)
             {

src/UniTask/Assets/Plugins/UniTask/Runtime/TriggerEvent.cs

@@ -20,8 +20,6 @@ namespace Cysharp.Threading.Tasks
     {
         ITriggerHandler<T> head; // head.prev is last
         ITriggerHandler<T> iteratingHead;
-
-        bool preserveRemoveSelf;
         ITriggerHandler<T> iteratingNode;
 
         void LogError(Exception ex)
@@ -44,6 +42,7 @@ namespace Cysharp.Threading.Tasks
             while (h != null)
             {
                 iteratingNode = h;
+                var next = h.Next;
 
                 try
                 {
@@ -55,18 +54,7 @@ namespace Cysharp.Threading.Tasks
                     Remove(h);
                 }
 
-                if (preserveRemoveSelf)
-                {
-                    preserveRemoveSelf = false;
-                    iteratingNode = null;
-                    var next = h.Next;
-                    Remove(h);
-                    h = next;
-                }
-                else
-                {
-                    h = h.Next;
-                }
+                h = next;
             }
 
             iteratingNode = null;
@@ -96,8 +84,7 @@ namespace Cysharp.Threading.Tasks
                 {
                     LogError(ex);
                 }
-
-                preserveRemoveSelf = false;
+                
                 iteratingNode = null;
                 var next = h.Next;
                 Remove(h);
@@ -131,8 +118,7 @@ namespace Cysharp.Threading.Tasks
                 {
                     LogError(ex);
                 }
-
-                preserveRemoveSelf = false;
+                
                 iteratingNode = null;
                 var next = h.Next;
                 Remove(h);
@@ -166,8 +152,7 @@ namespace Cysharp.Threading.Tasks
                 {
                     LogError(ex);
                 }
-
-                preserveRemoveSelf = false;
+                
                 iteratingNode = null;
                 var next = h.Next;
                 Remove(h);
@@ -240,72 +225,64 @@ namespace Cysharp.Threading.Tasks
         public void Remove(ITriggerHandler<T> handler)
         {
             if (handler == null) throw new ArgumentNullException(nameof(handler));
+            
+            var prev = handler.Prev;
+            var next = handler.Next;
 
-            if (iteratingNode != null && iteratingNode == handler)
+            if (next != null)
             {
-                // if remove self, reserve remove self after invoke completed.
-                preserveRemoveSelf = true;
+                next.Prev = prev;
+            }
+
+            if (handler == head)
+            {
+                head = next;
+            }
+            else if (handler == iteratingHead)
+            {
+                iteratingHead = next;
             }
             else
             {
-                var prev = handler.Prev;
-                var next = handler.Next;
-
-                if (next != null)
+                // when handler is head, prev indicate last so don't use it.
+                if (prev != null)
                 {
-                    next.Prev = prev;
+                    prev.Next = next;
                 }
-
-                if (handler == head)
-                {
-                    head = next;
-                }
-                else if (handler == iteratingHead)
-                {
-                    iteratingHead = next;
-                }
-                else
-                {
-                    // when handler is head, prev indicate last so don't use it.
-                    if (prev != null)
-                    {
-                        prev.Next = next;
-                    }
-                }
-
-                if (head != null)
-                {
-                    if (head.Prev == handler)
-                    {
-                        if (prev != head)
-                        {
-                            head.Prev = prev;
-                        }
-                        else
-                        {
-                            head.Prev = null;
-                        }
-                    }
-                }
-
-                if (iteratingHead != null)
-                {
-                    if (iteratingHead.Prev == handler)
-                    {
-                        if (prev != iteratingHead.Prev)
-                        {
-                            iteratingHead.Prev = prev;
-                        }
-                        else
-                        {
-                            iteratingHead.Prev = null;
-                        }
-                    }
-                }
-
-                handler.Prev = null;
-                handler.Next = null;
             }
+
+            if (head != null)
+            {
+                if (head.Prev == handler)
+                {
+                    if (prev != head)
+                    {
+                        head.Prev = prev;
+                    }
+                    else
+                    {
+                        head.Prev = null;
+                    }
+                }
+            }
+
+            if (iteratingHead != null)
+            {
+                if (iteratingHead.Prev == handler)
+                {
+                    if (prev != iteratingHead.Prev)
+                    {
+                        iteratingHead.Prev = prev;
+                    }
+                    else
+                    {
+                        iteratingHead.Prev = null;
+                    }
+                }
+            }
+
+            handler.Prev = null;
+            handler.Next = null;
         }
     }
 }

src/UniTask/Assets/Tests/AsyncReactivePropertyTest.cs

@@ -0,0 +1,63 @@
+using Cysharp.Threading.Tasks;
+using FluentAssertions;
+using System;
+using System.Collections;
+using System.Threading;
+using UnityEngine.TestTools;
+
+namespace Cysharp.Threading.TasksTests
+{
+    public class AsyncReactivePropertyTest
+    {
+        private int _callCounter;
+        
+        [UnityTest]
+        public IEnumerator WaitCancelWait() => UniTask.ToCoroutine(async () =>
+        {
+            // Test case for https://github.com/Cysharp/UniTask/issues/444
+            
+            var property = new AsyncReactiveProperty<int>(0);
+            
+            var cts1 = new CancellationTokenSource();
+            var cts2 = new CancellationTokenSource();
+            WaitForProperty(property, cts1.Token);
+            WaitForProperty(property, cts2.Token);
+
+            _callCounter = 0;
+            property.Value = 1;
+            _callCounter.Should().Be(2);
+
+            cts2.Cancel();
+            cts2.Dispose();
+            cts1.Cancel();
+            cts1.Dispose();
+
+            var cts3 = new CancellationTokenSource();
+            WaitForProperty(property, cts3.Token);
+
+            _callCounter = 0;
+            property.Value = 2;
+            _callCounter.Should().Be(1);
+            
+            cts3.Cancel();
+            cts3.Dispose();
+            await UniTask.CompletedTask;
+        });
+
+        private async void WaitForProperty(AsyncReactiveProperty<int> property, CancellationToken token)
+        {
+            while (true)
+            {
+                try
+                {
+                    await property.WaitAsync(token);
+                    _callCounter++;
+                }
+                catch (OperationCanceledException)
+                {
+                    break;
+                }
+            }
+        }
+    }
+}

src/UniTask/Assets/Tests/AsyncReactivePropertyTest.cs.meta

@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: 27665955eefb4448969b8cc4dd204600
+timeCreated: 1676129650