added cancellation if playstate changes while awaiting indefinite cha…

…nges (#211)
XRTK · Jun 8, 2019 · c358d59 · c358d59
1 parent ba3accb
commit c358d59
Showing 1 changed file with 34 additions and 3 deletions.
diff --git a/XRTK-Core/Packages/com.xrtk.core/Utilities/Async/Awaiters.cs b/XRTK-Core/Packages/com.xrtk.core/Utilities/Async/Awaiters.cs
@@ -74,9 +74,21 @@ void Exception()
                     }
 
                     cancellationTokenSource.Token.Register(Exception);
+#if UNITY_EDITOR
+                    var editorCancelled = false;
+                    UnityEditor.EditorApplication.playModeStateChanged += playModeStateChanged => editorCancelled = true;
+#endif
 
                     while (!cancellationTokenSource.IsCancellationRequested)
                     {
+
+#if UNITY_EDITOR
+                        if (editorCancelled)
+                        {
+                            tcs.TrySetCanceled(CancellationToken.None);
+                        }
+#endif
+
                         try
                         {
                             if (!predicate(element))
@@ -103,12 +115,31 @@ private static async Task WaitUntil_Indefinite<T>(T element, Func<T, bool> predi
         {
             var tcs = new TaskCompletionSource<object>();
 
+#if UNITY_EDITOR
+            var editorCancelled = false;
+            UnityEditor.EditorApplication.playModeStateChanged += playModeStateChanged => editorCancelled = true;
+#endif
+
             while (true)
             {
-                if (!predicate(element))
+#if UNITY_EDITOR
+                if (editorCancelled)
+                {
+                    tcs.TrySetCanceled(CancellationToken.None);
+                }
+#endif
+
+                try
+                {
+                    if (!predicate(element))
+                    {
+                        await Task.Delay(1);
+                        continue;
+                    }
+                }
+                catch (Exception e)
                 {
-                    await Task.Delay(1);
-                    continue;
+                    tcs.TrySetException(e);
                 }
 
                 tcs.TrySetResult(Task.CompletedTask);