From 340736795cee336ad0065e5f0910592509185fc3 Mon Sep 17 00:00:00 2001
From: neuecc <ils@neue.cc>
Date: Mon, 24 Oct 2022 21:51:23 +0900
Subject: [PATCH] UnityEvent.AsAsyncEnumerable handle cancel correctly #365

---
 .../Runtime/UnityAsyncExtensions.uGUI.cs      | 40 +++++++++++++++++--
 1 file changed, 36 insertions(+), 4 deletions(-)

diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.uGUI.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.uGUI.cs
index 11ebd63..9031b50 100644
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.uGUI.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.uGUI.cs
@@ -688,13 +688,27 @@ namespace Cysharp.Threading.Tasks
             static void OnCanceled1(object state)
             {
                 var self = (UnityEventHandlerAsyncEnumerator)state;
-                self.DisposeAsync().Forget();
+                try
+                {
+                    self.completionSource.TrySetCanceled(self.cancellationToken1);
+                }
+                finally
+                {
+                    self.DisposeAsync().Forget();
+                }
             }
 
             static void OnCanceled2(object state)
             {
                 var self = (UnityEventHandlerAsyncEnumerator)state;
-                self.DisposeAsync().Forget();
+                try
+                {
+                    self.completionSource.TrySetCanceled(self.cancellationToken2);
+                }
+                finally
+                {
+                    self.DisposeAsync().Forget();
+                }
             }
 
             public UniTask DisposeAsync()
@@ -706,6 +720,8 @@ namespace Cysharp.Threading.Tasks
                     registration1.Dispose();
                     registration2.Dispose();
                     unityEvent.RemoveListener(unityAction);
+
+                    completionSource.TrySetCanceled();
                 }
 
                 return default;
@@ -793,13 +809,27 @@ namespace Cysharp.Threading.Tasks
             static void OnCanceled1(object state)
             {
                 var self = (UnityEventHandlerAsyncEnumerator)state;
-                self.DisposeAsync().Forget();
+                try
+                {
+                    self.completionSource.TrySetCanceled(self.cancellationToken1);
+                }
+                finally
+                {
+                    self.DisposeAsync().Forget();
+                }
             }
 
             static void OnCanceled2(object state)
             {
                 var self = (UnityEventHandlerAsyncEnumerator)state;
-                self.DisposeAsync().Forget();
+                try
+                {
+                    self.completionSource.TrySetCanceled(self.cancellationToken2);
+                }
+                finally
+                {
+                    self.DisposeAsync().Forget();
+                }
             }
 
             public UniTask DisposeAsync()
@@ -815,6 +845,8 @@ namespace Cysharp.Threading.Tasks
                         disp.Dispose();
                     }
                     unityEvent.RemoveListener(unityAction);
+
+                    completionSource.TrySetCanceled();
                 }
 
                 return default;