From d3428958d28fc1bcc0ac04d7b76c9e0eb69ace37 Mon Sep 17 00:00:00 2001
From: Cory Leach <cory.leach@gmail.com>
Date: Thu, 21 Nov 2019 10:31:43 -0500
Subject: [PATCH] Fix for WaitForBackground

Unit tests were sometimes failing because the WaitForBackground awaiter was sometimes resuming on the main thread.
---
 Runtime/Awaiters.cs            | 46 ++++++++++++++++++++++++++++++++--
 Runtime/UnityTaskUtil.cs       |  4 +--
 Tests/Runtime/AwaitersTests.cs | 24 ++++++++++++------
 3 files changed, 62 insertions(+), 12 deletions(-)

diff --git a/Runtime/Awaiters.cs b/Runtime/Awaiters.cs
index 48e9035..9c601e2 100644
--- a/Runtime/Awaiters.cs
+++ b/Runtime/Awaiters.cs
@@ -28,9 +28,51 @@ public static class Awaiters
 
     public class WaitForBackground
     {
-        public ConfiguredTaskAwaitable.ConfiguredTaskAwaiter GetAwaiter()
+        private class BackgroundThreadJoinAwaiter : IAwaitable
         {
-            return Task.Run(() => {}).ConfigureAwait(false).GetAwaiter();
+            private Action _continuation = null;
+            private bool isCompleted = false;
+            public bool IsCompleted => isCompleted;
+
+            public void Complete()
+            {
+                isCompleted = true;
+                _continuation?.Invoke();
+            }
+
+            public void GetResult()
+            {
+                do
+                {
+                } while (!isCompleted);
+            }
+
+            void INotifyCompletion.OnCompleted(Action continuation)
+            {
+                if (isCompleted)
+                {
+                    throw new InvalidOperationException("Continuation is invalid. This awaiter is already completed.");
+                }
+                _continuation = continuation;
+            }
+        }
+        
+        public IAwaitable GetAwaiter()
+        {
+            //Doing Task.Run(()=>{}).ConfigureAwait(false) will apparently sometimes still resume on the main thread
+            //Updated to the below pattern to ensure we actually will be running in the background when we resume
+            var awaiter = new BackgroundThreadJoinAwaiter();
+            Task.Run(async () =>
+            {
+                //Doing complete immediately without a yield appears to cause the awaiter to never resume
+                //I'm not entirely sure as to why.
+                //I suspected maybe Complete() was getting called before the the method doing the awaiting could add its continuation
+                //However when I added a check and exception for this I did not see it get thrown.
+                //Adding the await Task.Yield however appeared to get Unit tests to consistently pass.
+                await Task.Yield(); 
+                awaiter.Complete(); 
+            });
+            return awaiter;
         }
     }
 
diff --git a/Runtime/UnityTaskUtil.cs b/Runtime/UnityTaskUtil.cs
index 532f7e0..cbe69f7 100644
--- a/Runtime/UnityTaskUtil.cs
+++ b/Runtime/UnityTaskUtil.cs
@@ -19,8 +19,8 @@ private static void Install()
             UnityTaskFactory = new TaskFactory<UnityEngine.Object>(UnityTaskScheduler);
         }
 
-        public static bool CurrentThreadIsUnityThread => SynchronizationContext.Current == UnitySynchronizationContext;
-
+        public static bool CurrentThreadIsUnityThread => UnityThreadId == Thread.CurrentThread.ManagedThreadId;
+        
         public static TaskScheduler UnityTaskScheduler { get; private set; }
 
         public static TaskFactory<UnityEngine.Object> UnityTaskFactory { get; private set; }
diff --git a/Tests/Runtime/AwaitersTests.cs b/Tests/Runtime/AwaitersTests.cs
index b5df174..099d59b 100644
--- a/Tests/Runtime/AwaitersTests.cs
+++ b/Tests/Runtime/AwaitersTests.cs
@@ -1,4 +1,5 @@
 using System.Collections;
+using System.Threading;
 using System.Threading.Tasks;
 using NUnit.Framework;
 using UnityEngine;
@@ -8,25 +9,32 @@ namespace Gameframe.Async.Tests
 {
     public class AwaitersTests
     {
-        [UnityTest]
+        [UnityTest, Timeout(1000)]
         public IEnumerator TestBackgroundAndMainThreadAwaiters()
         {
+            yield return null;
             var task = DoTest();
             yield return task.AsIEnumerator();
         }
         
         private async Task DoTest()
         {
+            Assert.IsTrue(UnityTaskUtil.UnitySynchronizationContext != null);
+            
             //Start on unity thread
             Assert.IsTrue(UnityTaskUtil.CurrentThreadIsUnityThread);
+
+            //Test Migration Multiple times just to be sure
+            for (int i = 0; i < 10; i++)
+            {
+                //Migrate to background thread
+                await Awaiters.BackgroundThread;            
+                Assert.IsFalse(UnityTaskUtil.CurrentThreadIsUnityThread,$"Expected to be on background thread. CurrentThread:{Thread.CurrentThread.ManagedThreadId} UnityThreadId:{UnityTaskUtil.UnityThreadId}");
             
-            //Migrate to background thread
-            await Awaiters.BackgroundThread;            
-            Assert.IsFalse(UnityTaskUtil.CurrentThreadIsUnityThread);
-            
-            //Migrate back to main thread
-            await Awaiters.MainUnityThread;
-            Assert.IsTrue(UnityTaskUtil.CurrentThreadIsUnityThread);
+                //Migrate back to main thread
+                await Awaiters.MainUnityThread;
+                Assert.IsTrue(UnityTaskUtil.CurrentThreadIsUnityThread,$"Expected to be on main thread. CurrentThread:{Thread.CurrentThread.ManagedThreadId} UnityThreadId:{UnityTaskUtil.UnityThreadId}");
+            }
             
             //Await the main thread when already on the main thread should do nothing
             await Awaiters.MainUnityThread;