Merge pull request #557 from Cysharp/hadashiA/fix-cancel-immediate

Fix unintended returning to pool with cancelImmediately
2026-05-16 20:20:45 +00:00 · 2024-03-28 15:21:40 +09:00
parent 01c8fada1f 938ddd5356
commit b724a2aa84
7 changed files with 260 additions and 149 deletions
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/External/Addressables/AddressablesAsyncExtensions.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/External/Addressables/AddressablesAsyncExtensions.cs
@@ -110,6 +110,7 @@ namespace Cysharp.Threading.Tasks
            CancellationTokenRegistration cancellationTokenRegistration;
            IProgress<float> progress;
            bool autoReleaseWhenCanceled;
+            bool cancelImmediately;
            bool completed;

            UniTaskCompletionSourceCore<AsyncUnit> core;
@@ -134,8 +135,9 @@ namespace Cysharp.Threading.Tasks
                result.handle = handle;
                result.progress = progress;
                result.cancellationToken = cancellationToken;
-                result.completed = false;
+                result.cancelImmediately = cancelImmediately;
                result.autoReleaseWhenCanceled = autoReleaseWhenCanceled;
+                result.completed = false;
                
                if (cancelImmediately && cancellationToken.CanBeCanceled)
                {
@@ -169,33 +171,41 @@ namespace Cysharp.Threading.Tasks

                if (completed)
                {
-                    TryReturn();
+                    return;
+                }
+                
+                completed = true;
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    if (autoReleaseWhenCanceled && handle.IsValid())
+                    {
+                        Addressables.Release(handle);
+                    }
+                    core.TrySetCanceled(cancellationToken);
+                }
+                else if (handle.Status == AsyncOperationStatus.Failed)
+                {
+                    core.TrySetException(handle.OperationException);
                }
                else
                {
-                    completed = true;
-                    if (cancellationToken.IsCancellationRequested)
-                    {
-                        if (autoReleaseWhenCanceled && handle.IsValid())
-                        {
-                            Addressables.Release(handle);
-                        }
-                        core.TrySetCanceled(cancellationToken);
-                    }
-                    else if (handle.Status == AsyncOperationStatus.Failed)
-                    {
-                        core.TrySetException(handle.OperationException);
-                    }
-                    else
-                    {
-                        core.TrySetResult(AsyncUnit.Default);
-                    }
+                    core.TrySetResult(AsyncUnit.Default);
                }
            }

            public void GetResult(short token)
            {
-                core.GetResult(token);
+                try
+                {
+                    core.GetResult(token);
+                }
+                finally
+                {
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
+                }
            }

            public UniTaskStatus GetStatus(short token)
@@ -217,7 +227,6 @@ namespace Cysharp.Threading.Tasks
            {
                if (completed)
                {
-                    TryReturn();
                    return false;
                }

@@ -304,6 +313,7 @@ namespace Cysharp.Threading.Tasks
            CancellationTokenRegistration cancellationTokenRegistration;
            IProgress<float> progress;
            bool autoReleaseWhenCanceled;
+            bool cancelImmediately;
            bool completed;

            UniTaskCompletionSourceCore<T> core;
@@ -330,6 +340,7 @@ namespace Cysharp.Threading.Tasks
                result.completed = false;
                result.progress = progress;
                result.autoReleaseWhenCanceled = autoReleaseWhenCanceled;
+                result.cancelImmediately = cancelImmediately;
                
                if (cancelImmediately && cancellationToken.CanBeCanceled)
                {
@@ -363,33 +374,38 @@ namespace Cysharp.Threading.Tasks

                if (completed)
                {
-                    TryReturn();
+                    return;
+                }
+                completed = true;
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    if (autoReleaseWhenCanceled && handle.IsValid())
+                    {
+                        Addressables.Release(handle);
+                    }
+                    core.TrySetCanceled(cancellationToken);
+                }
+                else if (argHandle.Status == AsyncOperationStatus.Failed)
+                {
+                    core.TrySetException(argHandle.OperationException);
                }
                else
                {
-                    completed = true;
-                    if (cancellationToken.IsCancellationRequested)
-                    {
-                        if (autoReleaseWhenCanceled && handle.IsValid())
-                        {
-                            Addressables.Release(handle);
-                        }
-                        core.TrySetCanceled(cancellationToken);
-                    }
-                    else if (argHandle.Status == AsyncOperationStatus.Failed)
-                    {
-                        core.TrySetException(argHandle.OperationException);
-                    }
-                    else
-                    {
-                        core.TrySetResult(argHandle.Result);
-                    }
+                    core.TrySetResult(argHandle.Result);
                }
            }

            public T GetResult(short token)
            {
-                return core.GetResult(token);
+                try
+                {
+                    return core.GetResult(token);
+                }
+                finally
+                {
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                        TryReturn();
+                }
            }

            void IUniTaskSource.GetResult(short token)
@@ -416,7 +432,6 @@ namespace Cysharp.Threading.Tasks
            {
                if (completed)
                {
-                    TryReturn();
                    return false;
                }

--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.Delay.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.Delay.cs
@@ -208,6 +208,7 @@ namespace Cysharp.Threading.Tasks

            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
            UniTaskCompletionSourceCore<object> core;

            YieldPromise()
@@ -227,6 +228,7 @@ namespace Cysharp.Threading.Tasks
                }

                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
                
                if (cancelImmediately && cancellationToken.CanBeCanceled)
                {
@@ -253,7 +255,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -290,6 +295,7 @@ namespace Cysharp.Threading.Tasks
                core.Reset();
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }
        }
@@ -309,6 +315,7 @@ namespace Cysharp.Threading.Tasks
            UniTaskCompletionSourceCore<AsyncUnit> core;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;

            NextFramePromise()
            {
@@ -328,6 +335,7 @@ namespace Cysharp.Threading.Tasks

                result.frameCount = PlayerLoopHelper.IsMainThread ? Time.frameCount : -1;
                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;

                if (cancelImmediately && cancellationToken.CanBeCanceled)
                {
@@ -354,7 +362,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -414,6 +425,7 @@ namespace Cysharp.Threading.Tasks
            UniTaskCompletionSourceCore<object> core;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;

            WaitForEndOfFramePromise()
            {
@@ -432,6 +444,7 @@ namespace Cysharp.Threading.Tasks
                }

                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;

                if (cancelImmediately && cancellationToken.CanBeCanceled)
                {
@@ -458,7 +471,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -533,6 +549,7 @@ namespace Cysharp.Threading.Tasks
            int delayFrameCount;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;

            int currentFrameCount;
            UniTaskCompletionSourceCore<AsyncUnit> core;
@@ -556,6 +573,7 @@ namespace Cysharp.Threading.Tasks
                result.delayFrameCount = delayFrameCount;
                result.cancellationToken = cancellationToken;
                result.initialFrame = PlayerLoopHelper.IsMainThread ? Time.frameCount : -1;
+                result.cancelImmediately = cancelImmediately;

                if (cancelImmediately && cancellationToken.CanBeCanceled)
                {
@@ -582,7 +600,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -653,6 +674,7 @@ namespace Cysharp.Threading.Tasks
                delayFrameCount = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }
        }
@@ -673,6 +695,7 @@ namespace Cysharp.Threading.Tasks
            float elapsed;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;

            UniTaskCompletionSourceCore<object> core;

@@ -696,6 +719,7 @@ namespace Cysharp.Threading.Tasks
                result.delayTimeSpan = (float)delayTimeSpan.TotalSeconds;
                result.cancellationToken = cancellationToken;
                result.initialFrame = PlayerLoopHelper.IsMainThread ? Time.frameCount : -1;
+                result.cancelImmediately = cancelImmediately;

                if (cancelImmediately && cancellationToken.CanBeCanceled)
                {
@@ -722,7 +746,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -775,6 +802,7 @@ namespace Cysharp.Threading.Tasks
                elapsed = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }
        }
@@ -795,6 +823,7 @@ namespace Cysharp.Threading.Tasks
            int initialFrame;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;

            UniTaskCompletionSourceCore<object> core;

@@ -818,6 +847,7 @@ namespace Cysharp.Threading.Tasks
                result.delayFrameTimeSpan = (float)delayFrameTimeSpan.TotalSeconds;
                result.initialFrame = PlayerLoopHelper.IsMainThread ? Time.frameCount : -1;
                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;

                if (cancelImmediately && cancellationToken.CanBeCanceled)
                {
@@ -844,7 +874,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -897,6 +930,7 @@ namespace Cysharp.Threading.Tasks
                elapsed = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }
        }
@@ -916,6 +950,7 @@ namespace Cysharp.Threading.Tasks
            ValueStopwatch stopwatch;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;

            UniTaskCompletionSourceCore<AsyncUnit> core;

@@ -938,6 +973,7 @@ namespace Cysharp.Threading.Tasks
                result.stopwatch = ValueStopwatch.StartNew();
                result.delayTimeSpanTicks = delayTimeSpan.Ticks;
                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;

                if (cancelImmediately && cancellationToken.CanBeCanceled)
                {
@@ -964,7 +1000,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -1013,6 +1052,7 @@ namespace Cysharp.Threading.Tasks
                stopwatch = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }
        }
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.WaitUntil.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.WaitUntil.cs
@@ -49,6 +49,7 @@ namespace Cysharp.Threading.Tasks
            Func<bool> predicate;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;

            UniTaskCompletionSourceCore<object> core;

@@ -70,6 +71,7 @@ namespace Cysharp.Threading.Tasks

                result.predicate = predicate;
                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;

                if (cancelImmediately && cancellationToken.CanBeCanceled)
                {
@@ -96,7 +98,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -147,6 +152,7 @@ namespace Cysharp.Threading.Tasks
                predicate = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }
        }
@@ -165,6 +171,7 @@ namespace Cysharp.Threading.Tasks
            Func<bool> predicate;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;

            UniTaskCompletionSourceCore<object> core;

@@ -212,7 +219,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -263,6 +273,7 @@ namespace Cysharp.Threading.Tasks
                predicate = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }
        }
@@ -280,6 +291,7 @@ namespace Cysharp.Threading.Tasks

            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;

            UniTaskCompletionSourceCore<object> core;

@@ -287,7 +299,7 @@ namespace Cysharp.Threading.Tasks
            {
            }

-            public static IUniTaskSource Create(CancellationToken cancellationToken, PlayerLoopTiming timing, bool completeImmediately, out short token)
+            public static IUniTaskSource Create(CancellationToken cancellationToken, PlayerLoopTiming timing, bool cancelImmediately, out short token)
            {
                if (cancellationToken.IsCancellationRequested)
                {
@@ -300,8 +312,9 @@ namespace Cysharp.Threading.Tasks
                }

                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;

-                if (completeImmediately && cancellationToken.CanBeCanceled)
+                if (cancelImmediately && cancellationToken.CanBeCanceled)
                {
                    result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
                    {
@@ -326,7 +339,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -362,6 +378,7 @@ namespace Cysharp.Threading.Tasks
                core.Reset();
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }
        }
@@ -385,6 +402,7 @@ namespace Cysharp.Threading.Tasks
            IEqualityComparer<U> equalityComparer;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;

            UniTaskCompletionSourceCore<U> core;

@@ -410,6 +428,7 @@ namespace Cysharp.Threading.Tasks
                result.currentValue = monitorFunction(target);
                result.equalityComparer = equalityComparer ?? UnityEqualityComparer.GetDefault<U>();
                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
                
                if (cancelImmediately && cancellationToken.CanBeCanceled)
                {
@@ -436,7 +455,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -497,6 +519,7 @@ namespace Cysharp.Threading.Tasks
                equalityComparer = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }
        }
@@ -519,6 +542,7 @@ namespace Cysharp.Threading.Tasks
            IEqualityComparer<U> equalityComparer;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;

            UniTaskCompletionSourceCore<U> core;

@@ -543,6 +567,7 @@ namespace Cysharp.Threading.Tasks
                result.currentValue = monitorFunction(target);
                result.equalityComparer = equalityComparer ?? UnityEqualityComparer.GetDefault<U>();
                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
                
                if (cancelImmediately && cancellationToken.CanBeCanceled)
                {
@@ -569,7 +594,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -630,6 +658,7 @@ namespace Cysharp.Threading.Tasks
                equalityComparer = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }
        }
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.AssetBundleRequestAllAssets.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.AssetBundleRequestAllAssets.cs
@@ -101,6 +101,7 @@ namespace Cysharp.Threading.Tasks
            IProgress<float> progress;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
            bool completed;

            UniTaskCompletionSourceCore<UnityEngine.Object[]> core;
@@ -127,6 +128,7 @@ namespace Cysharp.Threading.Tasks
                result.asyncOperation = asyncOperation;
                result.progress = progress;
                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
                result.completed = false;

                asyncOperation.completed += result.continuationAction;
@@ -156,7 +158,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -216,6 +221,7 @@ namespace Cysharp.Threading.Tasks
                progress = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }
            
@@ -223,19 +229,17 @@ namespace Cysharp.Threading.Tasks
            {
                if (completed)
                {
-                    TryReturn();
+                    return;
+                }
+                
+                completed = true;
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    core.TrySetCanceled(cancellationToken);
                }
                else
                {
-                    completed = true;
-                    if (cancellationToken.IsCancellationRequested)
-                    {
-                        core.TrySetCanceled(cancellationToken);
-                    }
-                    else
-                    {
-                        core.TrySetResult(asyncOperation.allAssets);
-                    }
+                    core.TrySetResult(asyncOperation.allAssets);
                }
            }
        }
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.AsyncGPUReadback.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.AsyncGPUReadback.cs
@@ -45,7 +45,7 @@ namespace Cysharp.Threading.Tasks
            AsyncGPUReadbackRequest asyncOperation;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
-
+            bool cancelImmediately;
            UniTaskCompletionSourceCore<AsyncGPUReadbackRequest> core;

            AsyncGPUReadbackRequestAwaiterConfiguredSource()
@@ -66,6 +66,7 @@ namespace Cysharp.Threading.Tasks

                result.asyncOperation = asyncOperation;
                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
                
                if (cancelImmediately && cancellationToken.CanBeCanceled)
                {
@@ -92,7 +93,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -146,6 +150,7 @@ namespace Cysharp.Threading.Tasks
                asyncOperation = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }
        }
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.cs
@@ -97,6 +97,7 @@ namespace Cysharp.Threading.Tasks
            IProgress<float> progress;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
            bool completed;

            UniTaskCompletionSourceCore<AsyncUnit> core;
@@ -123,6 +124,7 @@ namespace Cysharp.Threading.Tasks
                result.asyncOperation = asyncOperation;
                result.progress = progress;
                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
                result.completed = false;
                
                asyncOperation.completed += result.continuationAction;
@@ -152,7 +154,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -209,6 +214,7 @@ namespace Cysharp.Threading.Tasks
                progress = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }

@@ -216,19 +222,16 @@ namespace Cysharp.Threading.Tasks
            {
                if (completed)
                {
-                    TryReturn();
+                    return;
+                }
+                completed = true;
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    core.TrySetCanceled(cancellationToken);
                }
                else
                {
-                    completed = true;
-                    if (cancellationToken.IsCancellationRequested)
-                    {
-                        core.TrySetCanceled(cancellationToken);
-                    }
-                    else
-                    {
-                        core.TrySetResult(AsyncUnit.Default);
-                    }
+                    core.TrySetResult(AsyncUnit.Default);
                }
            }
        }
@@ -320,6 +323,7 @@ namespace Cysharp.Threading.Tasks
            IProgress<float> progress;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
            bool completed;

            UniTaskCompletionSourceCore<UnityEngine.Object> core;
@@ -346,6 +350,7 @@ namespace Cysharp.Threading.Tasks
                result.asyncOperation = asyncOperation;
                result.progress = progress;
                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
                result.completed = false;
                
                asyncOperation.completed += result.continuationAction;
@@ -375,7 +380,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -436,6 +444,7 @@ namespace Cysharp.Threading.Tasks
                progress = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }

@@ -443,19 +452,16 @@ namespace Cysharp.Threading.Tasks
            {
                if (completed)
                {
-                    TryReturn();
+                    return;
+                }
+                completed = true;
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    core.TrySetCanceled(cancellationToken);
                }
                else
                {
-                    completed = true;
-                    if (cancellationToken.IsCancellationRequested)
-                    {
-                        core.TrySetCanceled(cancellationToken);
-                    }
-                    else
-                    {
-                        core.TrySetResult(asyncOperation.asset);
-                    }
+                    core.TrySetResult(asyncOperation.asset);
                }
            }
        }
@@ -548,6 +554,7 @@ namespace Cysharp.Threading.Tasks
            IProgress<float> progress;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
            bool completed;

            UniTaskCompletionSourceCore<UnityEngine.Object> core;
@@ -574,6 +581,7 @@ namespace Cysharp.Threading.Tasks
                result.asyncOperation = asyncOperation;
                result.progress = progress;
                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
                result.completed = false;
                
                asyncOperation.completed += result.continuationAction;
@@ -603,7 +611,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -664,6 +675,7 @@ namespace Cysharp.Threading.Tasks
                progress = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }

@@ -671,19 +683,16 @@ namespace Cysharp.Threading.Tasks
            {
                if (completed)
                {
-                    TryReturn();
+                    return;
+                }
+                completed = true;
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    core.TrySetCanceled(cancellationToken);
                }
                else
                {
-                    completed = true;
-                    if (cancellationToken.IsCancellationRequested)
-                    {
-                        core.TrySetCanceled(cancellationToken);
-                    }
-                    else
-                    {
-                        core.TrySetResult(asyncOperation.asset);
-                    }
+                    core.TrySetResult(asyncOperation.asset);
                }
            }
        }
@@ -777,6 +786,7 @@ namespace Cysharp.Threading.Tasks
            IProgress<float> progress;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
            bool completed;

            UniTaskCompletionSourceCore<AssetBundle> core;
@@ -803,6 +813,7 @@ namespace Cysharp.Threading.Tasks
                result.asyncOperation = asyncOperation;
                result.progress = progress;
                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
                result.completed = false;
                
                asyncOperation.completed += result.continuationAction;
@@ -832,7 +843,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -893,6 +907,7 @@ namespace Cysharp.Threading.Tasks
                progress = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }

@@ -900,19 +915,16 @@ namespace Cysharp.Threading.Tasks
            {
                if (completed)
                {
-                    TryReturn();
+                    return;
+                }
+                completed = true;
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    core.TrySetCanceled(cancellationToken);
                }
                else
                {
-                    completed = true;
-                    if (cancellationToken.IsCancellationRequested)
-                    {
-                        core.TrySetCanceled(cancellationToken);
-                    }
-                    else
-                    {
-                        core.TrySetResult(asyncOperation.assetBundle);
-                    }
+                    core.TrySetResult(asyncOperation.assetBundle);
                }
            }
        }
@@ -1021,6 +1033,7 @@ namespace Cysharp.Threading.Tasks
            IProgress<float> progress;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
            bool completed;

            UniTaskCompletionSourceCore<UnityWebRequest> core;
@@ -1047,6 +1060,7 @@ namespace Cysharp.Threading.Tasks
                result.asyncOperation = asyncOperation;
                result.progress = progress;
                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
                result.completed = false;
                
                asyncOperation.completed += result.continuationAction;
@@ -1077,7 +1091,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -1146,6 +1163,7 @@ namespace Cysharp.Threading.Tasks
                progress = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }

@@ -1153,23 +1171,20 @@ namespace Cysharp.Threading.Tasks
            {
                if (completed)
                {
-                    TryReturn();
+                    return;
+                }
+                completed = true;
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    core.TrySetCanceled(cancellationToken);
+                }
+                else if (asyncOperation.webRequest.IsError())
+                {
+                    core.TrySetException(new UnityWebRequestException(asyncOperation.webRequest));
                }
                else
                {
-                    completed = true;
-                    if (cancellationToken.IsCancellationRequested)
-                    {
-                        core.TrySetCanceled(cancellationToken);
-                    }
-                    else if (asyncOperation.webRequest.IsError())
-                    {
-                        core.TrySetException(new UnityWebRequestException(asyncOperation.webRequest));
-                    }
-                    else
-                    {
-                        core.TrySetResult(asyncOperation.webRequest);
-                    }
+                    core.TrySetResult(asyncOperation.webRequest);
                }
            }
        }
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.tt
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.tt
@@ -165,6 +165,7 @@ namespace Cysharp.Threading.Tasks
            IProgress<float> progress;
            CancellationToken cancellationToken;
            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
            bool completed;

            UniTaskCompletionSourceCore<<#= IsVoid(t) ? "AsyncUnit" : t.returnType #>> core;
@@ -191,6 +192,7 @@ namespace Cysharp.Threading.Tasks
                result.asyncOperation = asyncOperation;
                result.progress = progress;
                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
                result.completed = false;
                
                asyncOperation.completed += result.continuationAction;
@@ -227,7 +229,10 @@ namespace Cysharp.Threading.Tasks
                }
                finally
                {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
                }
            }

@@ -304,6 +309,7 @@ namespace Cysharp.Threading.Tasks
                progress = default;
                cancellationToken = default;
                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                return pool.TryPush(this);
            }

@@ -311,31 +317,28 @@ namespace Cysharp.Threading.Tasks
            {
                if (completed)
                {
-                    TryReturn();
+                    return;
+                }
+                completed = true;
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    core.TrySetCanceled(cancellationToken);
+                }
+<# if(IsUnityWebRequest(t)) { #>
+                else if (asyncOperation.webRequest.IsError())
+                {
+                    core.TrySetException(new UnityWebRequestException(asyncOperation.webRequest));
                }
                else
                {
-                    completed = true;
-                    if (cancellationToken.IsCancellationRequested)
-                    {
-                        core.TrySetCanceled(cancellationToken);
-                    }
-<# if(IsUnityWebRequest(t)) { #>
-                    else if (asyncOperation.webRequest.IsError())
-                    {
-                        core.TrySetException(new UnityWebRequestException(asyncOperation.webRequest));
-                    }
-                    else
-                    {
-                        core.TrySetResult(asyncOperation.webRequest);
-                    }
-<# } else { #>
-                    else
-                    {
-                        core.TrySetResult(<#= IsVoid(t) ? "AsyncUnit.Default" : $"asyncOperation.{t.returnField}" #>);
-                    }
-<# } #>
+                    core.TrySetResult(asyncOperation.webRequest);
                }
+<# } else { #>
+                else
+                {
+                    core.TrySetResult(<#= IsVoid(t) ? "AsyncUnit.Default" : $"asyncOperation.{t.returnField}" #>);
+                }
+<# } #>
            }
        }