From 8b7f832c0f27a2ba51986f70bd0a74ba33f7cbd3 Mon Sep 17 00:00:00 2001
From: neuecc <ils@neue.cc>
Date: Mon, 25 May 2020 09:58:06 +0900
Subject: [PATCH] JobHandle.WaitAsync accepts CancellationToken

---
 .../Runtime/UnityAsyncExtensions.Jobs.cs      |  7 +++++--
 src/UniTask/Assets/Scenes/SandboxMain.cs      | 19 ++++++++++++-------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.Jobs.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.Jobs.cs
index 370e446..a5dafb8 100644
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.Jobs.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.Jobs.cs
@@ -9,10 +9,11 @@ namespace Cysharp.Threading.Tasks
 {
     public static partial class UnityAsyncExtensions
     {
-        public static async UniTask WaitAsync(this JobHandle jobHandle, PlayerLoopTiming waitTiming)
+        public static async UniTask WaitAsync(this JobHandle jobHandle, PlayerLoopTiming waitTiming, CancellationToken cancellationToken = default)
         {
             await UniTask.Yield(waitTiming);
             jobHandle.Complete();
+            cancellationToken.ThrowIfCancellationRequested(); // call cancel after Complete.
         }
 
         public static UniTask.Awaiter GetAwaiter(this JobHandle jobHandle)
@@ -28,7 +29,9 @@ namespace Cysharp.Threading.Tasks
 
             return new UniTask(handler, token).GetAwaiter();
         }
-     
+
+        // can not pass CancellationToken because can't handle JobHandle's Complete and NativeArray.Dispose.
+
         public static UniTask ToUniTask(this JobHandle jobHandle, PlayerLoopTiming waitTiming)
         {
             var handler = JobHandlePromise.Create(jobHandle, out var token);
diff --git a/src/UniTask/Assets/Scenes/SandboxMain.cs b/src/UniTask/Assets/Scenes/SandboxMain.cs
index e9dcf15..72625fe 100644
--- a/src/UniTask/Assets/Scenes/SandboxMain.cs
+++ b/src/UniTask/Assets/Scenes/SandboxMain.cs
@@ -193,15 +193,20 @@ public class SandboxMain : MonoBehaviour
     async UniTask RunJobAsync()
     {
         var job = new MyJob() { loopCount = 999, inOut = new NativeArray<int>(1, Allocator.TempJob) };
-        JobHandle.ScheduleBatchedJobs();
+        try
+        {
+            JobHandle.ScheduleBatchedJobs();
 
-        var scheduled = job.Schedule();
+            var scheduled = job.Schedule();
 
-        UnityEngine.Debug.Log("OK");
-        await scheduled; // .ConfigureAwait(PlayerLoopTiming.Update); // .WaitAsync(PlayerLoopTiming.Update);
-        UnityEngine.Debug.Log("OK2");
-
-        job.inOut.Dispose();
+            UnityEngine.Debug.Log("OK");
+            await scheduled; // .ConfigureAwait(PlayerLoopTiming.Update); // .WaitAsync(PlayerLoopTiming.Update);
+            UnityEngine.Debug.Log("OK2");
+        }
+        finally
+        {
+            job.inOut.Dispose();
+        }
     }