Merge pull request #517 from Cysharp/hadashiA/cancel-immediately-flag

Add a flag to cancel immediately instead of player loop

README.md

@@ -336,6 +336,18 @@ if (isCanceled)
 
 Note: Only suppress throws if you call directly into the most source method. Otherwise, the return value will be converted, but the entire pipeline will not suppress throws.
 
+Some features that use Unity's player loop, such as `UniTask.Yield` and `UniTask.Delay` etc, determines CancellationToken state on the player loop. 
+This means it does not cancel immediately upon `CancellationToken` fired. 
+
+If you want to change this behaviour, the cancellation to be immediate, set the `cancelImmediately` flag as an argument.
+
+```csharp
+await UniTask.Yield(cancellationToken, cancelImmediately: true);
+```
+
+Note: Setting `cancelImmediately` to true and detecting an immediate cancellation is more costly than the default behavior.
+This is because it uses `CancellationToken.Register`; it is heavier than checking CancellationToken on the player loop.
+
 Timeout handling
 ---
 Timeout is a variation of cancellation. You can set timeout by `CancellationTokenSouce.CancelAfterSlim(TimeSpan)` and pass CancellationToken to async methods.
@@ -363,7 +375,7 @@ If you want to use timeout with other source of cancellation, use `CancellationT
 
 ```csharp
 var cancelToken = new CancellationTokenSource();
-cancelButton.onClick.AddListener(()=>
+cancelButton.onClick.AddListener(() =>
 {
     cancelToken.Cancel(); // cancel from button click.
 });

src/UniTask/Assets/Plugins/UniTask/Runtime/External/Addressables/AddressablesAsyncExtensions.cs

@@ -25,7 +25,12 @@ public static UniTask WithCancellation(this AsyncOperationHandle handle, Cancell
             return ToUniTask(handle, cancellationToken: cancellationToken);
         }
 
-        public static UniTask ToUniTask(this AsyncOperationHandle handle, IProgress<float> progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken))
+        public static UniTask WithCancellation(this AsyncOperationHandle handle, CancellationToken cancellationToken, bool cancelImmediately)
+        {
+            return ToUniTask(handle, cancellationToken: cancellationToken, cancelImmediately: cancelImmediately);
+        }
+
+        public static UniTask ToUniTask(this AsyncOperationHandle handle, IProgress<float> progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false)
         {
             if (cancellationToken.IsCancellationRequested) return UniTask.FromCanceled(cancellationToken);
 
@@ -44,7 +49,7 @@ public static UniTask WithCancellation(this AsyncOperationHandle handle, Cancell
                 return UniTask.CompletedTask;
             }
 
-            return new UniTask(AsyncOperationHandleConfiguredSource.Create(handle, timing, progress, cancellationToken, out var token), token);
+            return new UniTask(AsyncOperationHandleConfiguredSource.Create(handle, timing, progress, cancellationToken, cancelImmediately, out var token), token);
         }
 
         public struct AsyncOperationHandleAwaiter : ICriticalNotifyCompletion
@@ -106,6 +111,7 @@ static AsyncOperationHandleConfiguredSource()
             readonly Action<AsyncOperationHandle> continuationAction;
             AsyncOperationHandle handle;
             CancellationToken cancellationToken;
+            CancellationTokenRegistration cancellationTokenRegistration;
             IProgress<float> progress;
             bool completed;
 
@@ -116,7 +122,7 @@ static AsyncOperationHandleConfiguredSource()
                 continuationAction = Continuation;
             }
 
-            public static IUniTaskSource Create(AsyncOperationHandle handle, PlayerLoopTiming timing, IProgress<float> progress, CancellationToken cancellationToken, out short token)
+            public static IUniTaskSource Create(AsyncOperationHandle handle, PlayerLoopTiming timing, IProgress<float> progress, CancellationToken cancellationToken, bool cancelImmediately, out short token)
             {
                 if (cancellationToken.IsCancellationRequested)
                 {
@@ -132,6 +138,15 @@ public static IUniTaskSource Create(AsyncOperationHandle handle, PlayerLoopTimin
                 result.progress = progress;
                 result.cancellationToken = cancellationToken;
                 result.completed = false;
+
+                if (cancelImmediately && cancellationToken.CanBeCanceled)
+                {
+                    result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
+                    {
+                        var promise = (AsyncOperationHandleConfiguredSource)state;
+                        promise.core.TrySetCanceled(promise.cancellationToken);
+                    }, result);
+                }
 
                 TaskTracker.TrackActiveTask(result, 3);
 
@@ -219,6 +234,7 @@ bool TryReturn()
                 handle = default;
                 progress = default;
                 cancellationToken = default;
+                cancellationTokenRegistration.Dispose();
                 return pool.TryPush(this);
             }
         }
@@ -237,7 +253,12 @@ public static UniTask<T> WithCancellation<T>(this AsyncOperationHandle<T> handle
             return ToUniTask(handle, cancellationToken: cancellationToken);
         }
 
-        public static UniTask<T> ToUniTask<T>(this AsyncOperationHandle<T> handle, IProgress<float> progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken))
+        public static UniTask<T> WithCancellation<T>(this AsyncOperationHandle<T> handle, CancellationToken cancellationToken, bool cancelImmediately)
+        {
+            return ToUniTask(handle, cancellationToken: cancellationToken, cancelImmediately: cancelImmediately);
+        }
+
+        public static UniTask<T> ToUniTask<T>(this AsyncOperationHandle<T> handle, IProgress<float> progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false)
         {
             if (cancellationToken.IsCancellationRequested) return UniTask.FromCanceled<T>(cancellationToken);
 
@@ -255,7 +276,7 @@ public static UniTask<T> WithCancellation<T>(this AsyncOperationHandle<T> handle
                 return UniTask.FromResult(handle.Result);
             }
 
-            return new UniTask<T>(AsyncOperationHandleConfiguredSource<T>.Create(handle, timing, progress, cancellationToken, out var token), token);
+            return new UniTask<T>(AsyncOperationHandleConfiguredSource<T>.Create(handle, timing, progress, cancellationToken, cancelImmediately, out var token), token);
         }
 
         sealed class AsyncOperationHandleConfiguredSource<T> : IUniTaskSource<T>, IPlayerLoopItem, ITaskPoolNode<AsyncOperationHandleConfiguredSource<T>>
@@ -272,6 +293,7 @@ static AsyncOperationHandleConfiguredSource()
             readonly Action<AsyncOperationHandle<T>> continuationAction;
             AsyncOperationHandle<T> handle;
             CancellationToken cancellationToken;
+            CancellationTokenRegistration cancellationTokenRegistration;
             IProgress<float> progress;
             bool completed;
 
@@ -282,7 +304,7 @@ static AsyncOperationHandleConfiguredSource()
                 continuationAction = Continuation;
             }
 
-            public static IUniTaskSource<T> Create(AsyncOperationHandle<T> handle, PlayerLoopTiming timing, IProgress<float> progress, CancellationToken cancellationToken, out short token)
+            public static IUniTaskSource<T> Create(AsyncOperationHandle<T> handle, PlayerLoopTiming timing, IProgress<float> progress, CancellationToken cancellationToken, bool cancelImmediately, out short token)
             {
                 if (cancellationToken.IsCancellationRequested)
                 {
@@ -298,6 +320,15 @@ public static IUniTaskSource<T> Create(AsyncOperationHandle<T> handle, PlayerLoo
                 result.cancellationToken = cancellationToken;
                 result.completed = false;
                 result.progress = progress;
+
+                if (cancelImmediately && cancellationToken.CanBeCanceled)
+                {
+                    result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
+                    {
+                        var promise = (AsyncOperationHandleConfiguredSource<T>)state;
+                        promise.core.TrySetCanceled(promise.cancellationToken);
+                    }, result);
+                }
 
                 TaskTracker.TrackActiveTask(result, 3);
 
@@ -390,6 +421,7 @@ bool TryReturn()
                 handle = default;
                 progress = default;
                 cancellationToken = default;
+                cancellationTokenRegistration.Dispose();
                 return pool.TryPush(this);
             }
         }

src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/UnityExtensions/EveryUpdate.cs

@@ -4,38 +4,50 @@ namespace Cysharp.Threading.Tasks.Linq
 {
     public static partial class UniTaskAsyncEnumerable
     {
-        public static IUniTaskAsyncEnumerable<AsyncUnit> EveryUpdate(PlayerLoopTiming updateTiming = PlayerLoopTiming.Update)
+        public static IUniTaskAsyncEnumerable<AsyncUnit> EveryUpdate(PlayerLoopTiming updateTiming = PlayerLoopTiming.Update, bool cancelImmediately = false)
         {
-            return new EveryUpdate(updateTiming);
+            return new EveryUpdate(updateTiming, cancelImmediately);
         }
     }
 
     internal class EveryUpdate : IUniTaskAsyncEnumerable<AsyncUnit>
     {
         readonly PlayerLoopTiming updateTiming;
+        readonly bool cancelImmediately;
 
-        public EveryUpdate(PlayerLoopTiming updateTiming)
+        public EveryUpdate(PlayerLoopTiming updateTiming, bool cancelImmediately)
         {
             this.updateTiming = updateTiming;
+            this.cancelImmediately = cancelImmediately;
         }
 
         public IUniTaskAsyncEnumerator<AsyncUnit> GetAsyncEnumerator(CancellationToken cancellationToken = default)
         {
-            return new _EveryUpdate(updateTiming, cancellationToken);
+            return new _EveryUpdate(updateTiming, cancellationToken, cancelImmediately);
         }
 
         class _EveryUpdate : MoveNextSource, IUniTaskAsyncEnumerator<AsyncUnit>, IPlayerLoopItem
         {
             readonly PlayerLoopTiming updateTiming;
-            CancellationToken cancellationToken;
+            readonly CancellationToken cancellationToken;
+            readonly CancellationTokenRegistration cancellationTokenRegistration;
 
             bool disposed;
 
-            public _EveryUpdate(PlayerLoopTiming updateTiming, CancellationToken cancellationToken)
+            public _EveryUpdate(PlayerLoopTiming updateTiming, CancellationToken cancellationToken, bool cancelImmediately)
             {
                 this.updateTiming = updateTiming;
                 this.cancellationToken = cancellationToken;
 
+                if (cancelImmediately && cancellationToken.CanBeCanceled)
+                {
+                    cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
+                    {
+                        var source = (_EveryUpdate)state;
+                        source.completionSource.TrySetCanceled(source.cancellationToken);
+                    }, this);
+                }
+
                 TaskTracker.TrackActiveTask(this, 2);
                 PlayerLoopHelper.AddAction(updateTiming, this);
             }
@@ -44,17 +56,22 @@ public _EveryUpdate(PlayerLoopTiming updateTiming, CancellationToken cancellatio
 
             public UniTask<bool> MoveNextAsync()
             {
-                // return false instead of throw
-                if (disposed || cancellationToken.IsCancellationRequested) return CompletedTasks.False;
-
+                if (disposed) return CompletedTasks.False;
+
                 completionSource.Reset();
+
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    completionSource.TrySetCanceled(cancellationToken);
+                }
                 return new UniTask<bool>(this, completionSource.Version);
             }
 
             public UniTask DisposeAsync()
             {
                 if (!disposed)
                 {
+                    cancellationTokenRegistration.Dispose();
                     disposed = true;
                     TaskTracker.RemoveTracking(this);
                 }
@@ -63,7 +80,13 @@ public UniTask DisposeAsync()
 
             public bool MoveNext()
             {
-                if (disposed || cancellationToken.IsCancellationRequested)
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    completionSource.TrySetCanceled(cancellationToken);
+                    return false;
+                }
+
+                if (disposed)
                 {
                     completionSource.TrySetResult(false);
                     return false;