From 6b1e951b640de3355bc4bb989d69256f146bb502 Mon Sep 17 00:00:00 2001
From: Tommaso Checchi <tommaso.checchi1@gmail.com>
Date: Wed, 9 Jul 2025 21:29:12 -0700
Subject: [PATCH] RunOnThreadPool: fixed an issue where a cancellation could
 prevent returning to the main thread

---
 .../Plugins/UniTask/Runtime/UniTask.Run.cs    | 34 +++++++++----------
 1 file changed, 16 insertions(+), 18 deletions(-)

diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.Run.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.Run.cs
index ac3e7958..7fbbff0a 100644
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.Run.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.Run.cs
@@ -66,12 +66,11 @@ public static async UniTask RunOnThreadPool(Action action, bool configureAwait =
 
             await UniTask.SwitchToThreadPool();
 
-            cancellationToken.ThrowIfCancellationRequested();
-
             if (configureAwait)
             {
                 try
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
                     action();
                 }
                 finally
@@ -81,6 +80,7 @@ public static async UniTask RunOnThreadPool(Action action, bool configureAwait =
             }
             else
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 action();
             }
 
@@ -94,12 +94,11 @@ public static async UniTask RunOnThreadPool(Action<object> action, object state,
 
             await UniTask.SwitchToThreadPool();
 
-            cancellationToken.ThrowIfCancellationRequested();
-
             if (configureAwait)
             {
                 try
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
                     action(state);
                 }
                 finally
@@ -109,6 +108,7 @@ public static async UniTask RunOnThreadPool(Action<object> action, object state,
             }
             else
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 action(state);
             }
 
@@ -122,12 +122,11 @@ public static async UniTask RunOnThreadPool(Func<UniTask> action, bool configure
 
             await UniTask.SwitchToThreadPool();
 
-            cancellationToken.ThrowIfCancellationRequested();
-
             if (configureAwait)
             {
                 try
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
                     await action();
                 }
                 finally
@@ -137,6 +136,7 @@ public static async UniTask RunOnThreadPool(Func<UniTask> action, bool configure
             }
             else
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 await action();
             }
 
@@ -150,12 +150,11 @@ public static async UniTask RunOnThreadPool(Func<object, UniTask> action, object
 
             await UniTask.SwitchToThreadPool();
 
-            cancellationToken.ThrowIfCancellationRequested();
-
             if (configureAwait)
             {
                 try
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
                     await action(state);
                 }
                 finally
@@ -165,6 +164,7 @@ public static async UniTask RunOnThreadPool(Func<object, UniTask> action, object
             }
             else
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 await action(state);
             }
 
@@ -178,12 +178,11 @@ public static async UniTask<T> RunOnThreadPool<T>(Func<T> func, bool configureAw
 
             await UniTask.SwitchToThreadPool();
 
-            cancellationToken.ThrowIfCancellationRequested();
-
             if (configureAwait)
             {
                 try
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
                     return func();
                 }
                 finally
@@ -194,6 +193,7 @@ public static async UniTask<T> RunOnThreadPool<T>(Func<T> func, bool configureAw
             }
             else
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 return func();
             }
         }
@@ -205,23 +205,22 @@ public static async UniTask<T> RunOnThreadPool<T>(Func<UniTask<T>> func, bool co
 
             await UniTask.SwitchToThreadPool();
 
-            cancellationToken.ThrowIfCancellationRequested();
-
             if (configureAwait)
             {
                 try
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
                     return await func();
                 }
                 finally
                 {
-                    cancellationToken.ThrowIfCancellationRequested();
                     await UniTask.Yield();
                     cancellationToken.ThrowIfCancellationRequested();
                 }
             }
             else
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 var result = await func();
                 cancellationToken.ThrowIfCancellationRequested();
                 return result;
@@ -235,12 +234,11 @@ public static async UniTask<T> RunOnThreadPool<T>(Func<object, T> func, object s
 
             await UniTask.SwitchToThreadPool();
 
-            cancellationToken.ThrowIfCancellationRequested();
-
             if (configureAwait)
             {
                 try
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
                     return func(state);
                 }
                 finally
@@ -251,6 +249,7 @@ public static async UniTask<T> RunOnThreadPool<T>(Func<object, T> func, object s
             }
             else
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 return func(state);
             }
         }
@@ -262,23 +261,22 @@ public static async UniTask<T> RunOnThreadPool<T>(Func<object, UniTask<T>> func,
 
             await UniTask.SwitchToThreadPool();
 
-            cancellationToken.ThrowIfCancellationRequested();
-
             if (configureAwait)
             {
                 try
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
                     return await func(state);
                 }
                 finally
                 {
-                    cancellationToken.ThrowIfCancellationRequested();
                     await UniTask.Yield();
                     cancellationToken.ThrowIfCancellationRequested();
                 }
             }
             else
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 var result = await func(state);
                 cancellationToken.ThrowIfCancellationRequested();
                 return result;