From ea9e61c2e186b317ae5da667a9ff722262f69549 Mon Sep 17 00:00:00 2001
From: neuecc <ils@neue.cc>
Date: Wed, 20 May 2020 11:04:59 +0900
Subject: [PATCH] Add CancellationToken.WaitUntilCanceled

---
 README.md                                     | 10 +---
 .../CancellationTokenTest.cs                  | 51 +++++++++++++++++
 .../Runtime/CancellationTokenExtensions.cs    | 55 +++++++++++++++++--
 src/UniTask/Assets/Scenes/SandboxMain.cs      |  6 ++
 4 files changed, 109 insertions(+), 13 deletions(-)
 create mode 100644 src/UniTask.NetCoreTests/CancellationTokenTest.cs

diff --git a/README.md b/README.md
index b27f6cc..edfc115 100644
--- a/README.md
+++ b/README.md
@@ -378,15 +378,7 @@ ECS, PlayerLoop
 TODO:
 
 ```csharp
-// Setup Entities Loop.
-var loop = PlayerLoop.GetDefaultPlayerLoop();
-foreach (var world in Unity.Entities.World.All)
-{
-    ScriptBehaviourUpdateOrder.UpdatePlayerLoop(world, loop);
-    loop = PlayerLoop.GetCurrentPlayerLoop();
-}
-
-// UniTask PlayerLoop Initialize.
+var loop = PlayerLoop.GetCurrentPlayerLoop();
 PlayerLoopHelper.Initialize(ref loop);
 ```
 
diff --git a/src/UniTask.NetCoreTests/CancellationTokenTest.cs b/src/UniTask.NetCoreTests/CancellationTokenTest.cs
new file mode 100644
index 0000000..920fd83
--- /dev/null
+++ b/src/UniTask.NetCoreTests/CancellationTokenTest.cs
@@ -0,0 +1,51 @@
+using Cysharp.Threading.Tasks;
+using FluentAssertions;
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading;
+using System.Threading.Channels;
+using Cysharp.Threading.Tasks.Linq;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace NetCoreTests
+{
+    public class CancellationTokenTest
+    {
+        [Fact]
+        public async Task WaitUntilCanceled()
+        {
+            var cts = new CancellationTokenSource();
+
+            cts.CancelAfter(TimeSpan.FromSeconds(1.5));
+
+            var now = DateTime.UtcNow;
+
+            await cts.Token.WaitUntilCanceled();
+
+            var elapsed = DateTime.UtcNow - now;
+
+            elapsed.Should().BeGreaterThan(TimeSpan.FromSeconds(1));
+        }
+
+        [Fact]
+        public void AlreadyCanceled()
+        {
+            var cts = new CancellationTokenSource();
+
+            cts.Cancel();
+
+            cts.Token.WaitUntilCanceled().GetAwaiter().IsCompleted.Should().BeTrue();
+        }
+
+        [Fact]
+        public void None()
+        {
+            CancellationToken.None.WaitUntilCanceled().GetAwaiter().IsCompleted.Should().BeTrue();
+        }
+    }
+
+
+}
diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/CancellationTokenExtensions.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/CancellationTokenExtensions.cs
index c70940c..e87721f 100644
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/CancellationTokenExtensions.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/CancellationTokenExtensions.cs
@@ -1,6 +1,7 @@
 #pragma warning disable CS1591 // Missing XML comment for publicly visible type or member
 
 using System;
+using System.Runtime.CompilerServices;
 using System.Threading;
 
 namespace Cysharp.Threading.Tasks
@@ -9,15 +10,15 @@ namespace Cysharp.Threading.Tasks
     {
         static readonly Action<object> cancellationTokenCallback = Callback;
 
-        public static (UniTask, CancellationTokenRegistration) ToUniTask(this CancellationToken cts)
+        public static (UniTask, CancellationTokenRegistration) ToUniTask(this CancellationToken cancellationToken)
         {
-            if (cts.IsCancellationRequested)
+            if (cancellationToken.IsCancellationRequested)
             {
-                return (UniTask.FromCanceled(cts), default(CancellationTokenRegistration));
+                return (UniTask.FromCanceled(cancellationToken), default(CancellationTokenRegistration));
             }
 
             var promise = new UniTaskCompletionSource();
-            return (promise.Task, cts.RegisterWithoutCaptureExecutionContext(cancellationTokenCallback, promise));
+            return (promise.Task, cancellationToken.RegisterWithoutCaptureExecutionContext(cancellationTokenCallback, promise));
         }
 
         static void Callback(object state)
@@ -26,6 +27,11 @@ namespace Cysharp.Threading.Tasks
             promise.TrySetResult();
         }
 
+        public static CancellationTokenAwaitable WaitUntilCanceled(this CancellationToken cancellationToken)
+        {
+            return new CancellationTokenAwaitable(cancellationToken);
+        }
+
         public static CancellationTokenRegistration RegisterWithoutCaptureExecutionContext(this CancellationToken cancellationToken, Action callback)
         {
             var restoreFlow = false;
@@ -70,5 +76,46 @@ namespace Cysharp.Threading.Tasks
             }
         }
     }
+
+    public struct CancellationTokenAwaitable
+    {
+        CancellationToken cancellationToken;
+
+        public CancellationTokenAwaitable(CancellationToken cancellationToken)
+        {
+            this.cancellationToken = cancellationToken;
+        }
+
+        public Awaiter GetAwaiter()
+        {
+            return new Awaiter(cancellationToken);
+        }
+
+        public struct Awaiter : ICriticalNotifyCompletion
+        {
+            CancellationToken cancellationToken;
+
+            public Awaiter(CancellationToken cancellationToken)
+            {
+                this.cancellationToken = cancellationToken;
+            }
+
+            public bool IsCompleted => !cancellationToken.CanBeCanceled || cancellationToken.IsCancellationRequested;
+
+            public void GetResult()
+            {
+            }
+
+            public void OnCompleted(Action continuation)
+            {
+                UnsafeOnCompleted(continuation);
+            }
+
+            public void UnsafeOnCompleted(Action continuation)
+            {
+                cancellationToken.RegisterWithoutCaptureExecutionContext(continuation);
+            }
+        }
+    }
 }
 
diff --git a/src/UniTask/Assets/Scenes/SandboxMain.cs b/src/UniTask/Assets/Scenes/SandboxMain.cs
index c6f48cc..bbf75bb 100644
--- a/src/UniTask/Assets/Scenes/SandboxMain.cs
+++ b/src/UniTask/Assets/Scenes/SandboxMain.cs
@@ -38,6 +38,11 @@ public enum MyEnum
 
 
 
+
+
+
+
+
 public static partial class UnityUIComponentExtensions
 {
 
@@ -123,6 +128,7 @@ public class SandboxMain : MonoBehaviour
 
 
 
+         
 
         UnityEngine.Debug.Log("OK");
         await scheduled; // .ConfigureAwait(PlayerLoopTiming.Update); // .WaitAsync(PlayerLoopTiming.Update);