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(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 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;