Fix unintended returning to pool with cancelImmediately (WaitUntil,etc)

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,8 @@ namespace Cysharp.Threading.Tasks
                 }
                 finally
                 {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                        TryReturn();
                 }
             }
 
@@ -147,6 +150,7 @@ namespace Cysharp.Threading.Tasks
                 predicate = default;
                 cancellationToken = default;
                 cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                 return pool.TryPush(this);
             }
         }
@@ -165,6 +169,7 @@ namespace Cysharp.Threading.Tasks
             Func<bool> predicate;
             CancellationToken cancellationToken;
             CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
 
             UniTaskCompletionSourceCore<object> core;
 
@@ -212,7 +217,8 @@ namespace Cysharp.Threading.Tasks
                 }
                 finally
                 {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                        TryReturn();
                 }
             }
 
@@ -263,6 +269,7 @@ namespace Cysharp.Threading.Tasks
                 predicate = default;
                 cancellationToken = default;
                 cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                 return pool.TryPush(this);
             }
         }
@@ -280,6 +287,7 @@ namespace Cysharp.Threading.Tasks
 
             CancellationToken cancellationToken;
             CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
 
             UniTaskCompletionSourceCore<object> core;
 
@@ -287,7 +295,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 +308,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 +335,8 @@ namespace Cysharp.Threading.Tasks
                 }
                 finally
                 {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                        TryReturn();
                 }
             }
 
@@ -362,6 +372,7 @@ namespace Cysharp.Threading.Tasks
                 core.Reset();
                 cancellationToken = default;
                 cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                 return pool.TryPush(this);
             }
         }
@@ -385,6 +396,7 @@ namespace Cysharp.Threading.Tasks
             IEqualityComparer<U> equalityComparer;
             CancellationToken cancellationToken;
             CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
 
             UniTaskCompletionSourceCore<U> core;
 
@@ -410,6 +422,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 +449,8 @@ namespace Cysharp.Threading.Tasks
                 }
                 finally
                 {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                        TryReturn();
                 }
             }
 
@@ -497,6 +511,7 @@ namespace Cysharp.Threading.Tasks
                 equalityComparer = default;
                 cancellationToken = default;
                 cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                 return pool.TryPush(this);
             }
         }
@@ -519,6 +534,7 @@ namespace Cysharp.Threading.Tasks
             IEqualityComparer<U> equalityComparer;
             CancellationToken cancellationToken;
             CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
 
             UniTaskCompletionSourceCore<U> core;
 
@@ -543,6 +559,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 +586,8 @@ namespace Cysharp.Threading.Tasks
                 }
                 finally
                 {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                        TryReturn();
                 }
             }
 
@@ -630,6 +648,7 @@ namespace Cysharp.Threading.Tasks
                 equalityComparer = default;
                 cancellationToken = default;
                 cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                 return pool.TryPush(this);
             }
         }

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,8 @@ namespace Cysharp.Threading.Tasks
                 }
                 finally
                 {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                        TryReturn();
                 }
             }
 
@@ -216,6 +219,7 @@ namespace Cysharp.Threading.Tasks
                 progress = default;
                 cancellationToken = default;
                 cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                 return pool.TryPush(this);
             }
             

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,8 @@ namespace Cysharp.Threading.Tasks
                 }
                 finally
                 {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                        TryReturn();
                 }
             }
 
@@ -146,6 +148,7 @@ namespace Cysharp.Threading.Tasks
                 asyncOperation = default;
                 cancellationToken = default;
                 cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                 return pool.TryPush(this);
             }
         }