From c2824027d4d7188b449658245d2eb6c566c00427 Mon Sep 17 00:00:00 2001
From: neuecc <ils@neue.cc>
Date: Thu, 16 Jul 2020 05:36:21 +0900
Subject: [PATCH] more strict handle cancel timing on
 asyncOperation.WithCancellation()

---
 .../UniTask/Runtime/UnityAsyncExtensions.cs   | 40 ++++++++++++++++---
 .../UniTask/Runtime/UnityAsyncExtensions.tt   |  6 +++
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.cs
index 6799a9c..015f0e5 100644
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.cs
@@ -134,6 +134,12 @@ namespace Cysharp.Threading.Tasks
                 else
                 {
                     completed = true;
+                    if (cancellationToken.IsCancellationRequested)
+                    {
+                        core.TrySetCanceled(cancellationToken);
+                        return;
+                    }
+
                     core.TrySetResult(AsyncUnit.Default);
                 }
             }
@@ -160,7 +166,7 @@ namespace Cysharp.Threading.Tasks
             }
 
             public bool MoveNext()
-            {   
+            {
                 if (completed)
                 {
                     TryReturn();
@@ -420,6 +426,12 @@ namespace Cysharp.Threading.Tasks
                 else
                 {
                     completed = true;
+                    if (cancellationToken.IsCancellationRequested)
+                    {
+                        core.TrySetCanceled(cancellationToken);
+                        return;
+                    }
+
                     core.TrySetResult(asyncOperation.asset);
                 }
             }
@@ -450,7 +462,7 @@ namespace Cysharp.Threading.Tasks
             }
 
             public bool MoveNext()
-            {   
+            {
                 if (completed)
                 {
                     TryReturn();
@@ -714,6 +726,12 @@ namespace Cysharp.Threading.Tasks
                 else
                 {
                     completed = true;
+                    if (cancellationToken.IsCancellationRequested)
+                    {
+                        core.TrySetCanceled(cancellationToken);
+                        return;
+                    }
+
                     core.TrySetResult(asyncOperation.asset);
                 }
             }
@@ -744,7 +762,7 @@ namespace Cysharp.Threading.Tasks
             }
 
             public bool MoveNext()
-            {   
+            {
                 if (completed)
                 {
                     TryReturn();
@@ -1008,6 +1026,12 @@ namespace Cysharp.Threading.Tasks
                 else
                 {
                     completed = true;
+                    if (cancellationToken.IsCancellationRequested)
+                    {
+                        core.TrySetCanceled(cancellationToken);
+                        return;
+                    }
+
                     core.TrySetResult(asyncOperation.assetBundle);
                 }
             }
@@ -1038,7 +1062,7 @@ namespace Cysharp.Threading.Tasks
             }
 
             public bool MoveNext()
-            {   
+            {
                 if (completed)
                 {
                     TryReturn();
@@ -1325,6 +1349,12 @@ namespace Cysharp.Threading.Tasks
                 else
                 {
                     completed = true;
+                    if (cancellationToken.IsCancellationRequested)
+                    {
+                        core.TrySetCanceled(cancellationToken);
+                        return;
+                    }
+
                     var result = asyncOperation.webRequest;
                     if (result.IsError())
                     {
@@ -1363,7 +1393,7 @@ namespace Cysharp.Threading.Tasks
             }
 
             public bool MoveNext()
-            {   
+            {
                 if (completed)
                 {
                     TryReturn();
diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.tt b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.tt
index aa12b5a..a32a1f6 100644
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.tt
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.tt
@@ -205,6 +205,12 @@ namespace Cysharp.Threading.Tasks
                 else
                 {
                     completed = true;
+                    if (cancellationToken.IsCancellationRequested)
+                    {
+                        core.TrySetCanceled(cancellationToken);
+                        return;
+                    }
+
 <# if(IsUnityWebRequest(t)) { #>
                     var result = asyncOperation.webRequest;
                     if (result.IsError())