Support awaiting additional engine callbacks

Add EngineCallbackTiming enumeration with support for Application.onBeforeRender, Canvas.willRenderCanvases & Canvas.preWillRenderCanvases callbacks.
Extract PlayerLoopRunner core logic into new base class ContinuationRunner.
Add new EngineCallbackRunner class for running callbacks supported by EngineCallbackTiming enum.
Modify YieldPromise to support creating a promise for the new enum.

src/UniTask/Assets/Plugins/UniTask/Runtime/Internal/ContinuationRunner.cs

@@ -0,0 +1,172 @@
+
+using System;
+using UnityEngine;
+
+namespace Cysharp.Threading.Tasks.Internal
+{
+    internal abstract class ContinuationRunner
+    {
+        const int InitialSize = 16;
+
+        readonly object runningAndQueueLock = new object();
+        readonly object arrayLock = new object();
+        readonly Action<Exception> unhandledExceptionCallback;
+
+        int tail = 0;
+        bool running = false;
+        IPlayerLoopItem[] loopItems = new IPlayerLoopItem[InitialSize];
+        MinimumQueue<IPlayerLoopItem> waitQueue = new MinimumQueue<IPlayerLoopItem>(InitialSize);
+
+
+        public ContinuationRunner()
+        {
+            this.unhandledExceptionCallback = ex => Debug.LogException(ex);
+        }
+
+        public void AddAction(IPlayerLoopItem item)
+        {
+            lock (runningAndQueueLock)
+            {
+                if (running)
+                {
+                    waitQueue.Enqueue(item);
+                    return;
+                }
+            }
+
+            lock (arrayLock)
+            {
+                // Ensure Capacity
+                if (loopItems.Length == tail)
+                {
+                    Array.Resize(ref loopItems, checked(tail * 2));
+                }
+                loopItems[tail++] = item;
+            }
+        }
+
+        public int Clear()
+        {
+            lock (arrayLock)
+            {
+                var rest = 0;
+
+                for (var index = 0; index < loopItems.Length; index++)
+                {
+                    if (loopItems[index] != null)
+                    {
+                        rest++;
+                    }
+
+                    loopItems[index] = null;
+                }
+
+                tail = 0;
+                return rest;
+            }
+        }
+
+        [System.Diagnostics.DebuggerHidden]
+        protected void RunCore()
+        {
+            lock (runningAndQueueLock)
+            {
+                running = true;
+            }
+
+            lock (arrayLock)
+            {
+                var j = tail - 1;
+
+                for (int i = 0; i < loopItems.Length; i++)
+                {
+                    var action = loopItems[i];
+                    if (action != null)
+                    {
+                        try
+                        {
+                            if (!action.MoveNext())
+                            {
+                                loopItems[i] = null;
+                            }
+                            else
+                            {
+                                continue; // next i 
+                            }
+                        }
+                        catch (Exception ex)
+                        {
+                            loopItems[i] = null;
+                            try
+                            {
+                                unhandledExceptionCallback(ex);
+                            }
+                            catch { }
+                        }
+                    }
+
+                    // find null, loop from tail
+                    while (i < j)
+                    {
+                        var fromTail = loopItems[j];
+                        if (fromTail != null)
+                        {
+                            try
+                            {
+                                if (!fromTail.MoveNext())
+                                {
+                                    loopItems[j] = null;
+                                    j--;
+                                    continue; // next j
+                                }
+                                else
+                                {
+                                    // swap
+                                    loopItems[i] = fromTail;
+                                    loopItems[j] = null;
+                                    j--;
+                                    goto NEXT_LOOP; // next i
+                                }
+                            }
+                            catch (Exception ex)
+                            {
+                                loopItems[j] = null;
+                                j--;
+                                try
+                                {
+                                    unhandledExceptionCallback(ex);
+                                }
+                                catch { }
+                                continue; // next j
+                            }
+                        }
+                        else
+                        {
+                            j--;
+                        }
+                    }
+
+                    tail = i; // loop end
+                    break; // LOOP END
+
+                NEXT_LOOP:
+                    continue;
+                }
+
+
+                lock (runningAndQueueLock)
+                {
+                    running = false;
+                    while (waitQueue.Count != 0)
+                    {
+                        if (loopItems.Length == tail)
+                        {
+                            Array.Resize(ref loopItems, checked(tail * 2));
+                        }
+                        loopItems[tail++] = waitQueue.Dequeue();
+                    }
+                }
+            }
+        }
+    }
+}

src/UniTask/Assets/Plugins/UniTask/Runtime/Internal/ContinuationRunner.cs.meta

@@ -0,0 +1,11 @@
+fileFormatVersion: 2
+guid: 4f21ae8b72659b348af5ebb66b84b07b
+MonoImporter:
+  externalObjects: {}
+  serializedVersion: 2
+  defaultReferences: []
+  executionOrder: 0
+  icon: {instanceID: 0}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 

src/UniTask/Assets/Plugins/UniTask/Runtime/Internal/EngineCallbackRunner.cs

@@ -0,0 +1,47 @@
+
+namespace Cysharp.Threading.Tasks.Internal
+{
+    internal sealed class EngineCallbackRunner : ContinuationRunner
+    {
+        readonly EngineCallbackTiming timing;
+
+
+        public EngineCallbackRunner(EngineCallbackTiming timing) : base()
+        {
+            this.timing = timing;
+        }
+
+        // delegate entrypoint.
+        public void Run()
+        {
+            // for debugging, create named stacktrace.
+#if DEBUG
+            switch (timing)
+            {
+                case EngineCallbackTiming.OnBeforeRender:
+                    OnBeforeRender();
+                    break;
+                case EngineCallbackTiming.WillRenderCanvases:
+                    WillRenderCanvases();
+                    break;
+#if UNITY_2021_3_OR_NEWER
+                case EngineCallbackTiming.PreWillRenderCanvases:
+                    PreWillRenderCanvases();
+                    break;
+#endif
+
+                default:
+                    break;
+            }
+#else
+            RunCore();
+#endif
+        }
+
+        void OnBeforeRender() => RunCore();
+        void WillRenderCanvases() => RunCore();
+#if UNITY_2021_3_OR_NEWER
+        void PreWillRenderCanvases() => RunCore();
+#endif
+    }
+}

src/UniTask/Assets/Plugins/UniTask/Runtime/Internal/EngineCallbackRunner.cs.meta

@@ -0,0 +1,11 @@
+fileFormatVersion: 2
+guid: 10e6588bde350b3418e72c3cd963f5f3
+MonoImporter:
+  externalObjects: {}
+  serializedVersion: 2
+  defaultReferences: []
+  executionOrder: 0
+  icon: {instanceID: 0}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 

src/UniTask/Assets/Plugins/UniTask/Runtime/Internal/PlayerLoopRunner.cs

@@ -1,74 +1,16 @@
 
-using System;
-using UnityEngine;
-
 namespace Cysharp.Threading.Tasks.Internal
 {
-    internal sealed class PlayerLoopRunner
+    internal sealed class PlayerLoopRunner : ContinuationRunner
     {
-        const int InitialSize = 16;
-
         readonly PlayerLoopTiming timing;
-        readonly object runningAndQueueLock = new object();
-        readonly object arrayLock = new object();
-        readonly Action<Exception> unhandledExceptionCallback;
-
-        int tail = 0;
-        bool running = false;
-        IPlayerLoopItem[] loopItems = new IPlayerLoopItem[InitialSize];
-        MinimumQueue<IPlayerLoopItem> waitQueue = new MinimumQueue<IPlayerLoopItem>(InitialSize);
 
 
-
+        public PlayerLoopRunner(PlayerLoopTiming timing) : base()
-        public PlayerLoopRunner(PlayerLoopTiming timing)
         {
-            this.unhandledExceptionCallback = ex => Debug.LogException(ex);
             this.timing = timing;
         }
 
-        public void AddAction(IPlayerLoopItem item)
-        {
-            lock (runningAndQueueLock)
-            {
-                if (running)
-                {
-                    waitQueue.Enqueue(item);
-                    return;
-                }
-            }
-
-            lock (arrayLock)
-            {
-                // Ensure Capacity
-                if (loopItems.Length == tail)
-                {
-                    Array.Resize(ref loopItems, checked(tail * 2));
-                }
-                loopItems[tail++] = item;
-            }
-        }
-
-        public int Clear()
-        {
-            lock (arrayLock)
-            {
-                var rest = 0;
-
-                for (var index = 0; index < loopItems.Length; index++)
-                {
-                    if (loopItems[index] != null)
-                    {
-                        rest++;
-                    }
-
-                    loopItems[index] = null;
-                }
-
-                tail = 0;
-                return rest;
-            }
-        }
-
         // delegate entrypoint.
         public void Run()
         {
@@ -152,109 +94,6 @@ namespace Cysharp.Threading.Tasks.Internal
         void TimeUpdate() => RunCore();
         void LastTimeUpdate() => RunCore();
 #endif
-
-        [System.Diagnostics.DebuggerHidden]
-        void RunCore()
-        {
-            lock (runningAndQueueLock)
-            {
-                running = true;
-            }
-
-            lock (arrayLock)
-            {
-                var j = tail - 1;
-
-                for (int i = 0; i < loopItems.Length; i++)
-                {
-                    var action = loopItems[i];
-                    if (action != null)
-                    {
-                        try
-                        {
-                            if (!action.MoveNext())
-                            {
-                                loopItems[i] = null;
-                            }
-                            else
-                            {
-                                continue; // next i 
-                            }
-                        }
-                        catch (Exception ex)
-                        {
-                            loopItems[i] = null;
-                            try
-                            {
-                                unhandledExceptionCallback(ex);
-                            }
-                            catch { }
-                        }
-                    }
-
-                    // find null, loop from tail
-                    while (i < j)
-                    {
-                        var fromTail = loopItems[j];
-                        if (fromTail != null)
-                        {
-                            try
-                            {
-                                if (!fromTail.MoveNext())
-                                {
-                                    loopItems[j] = null;
-                                    j--;
-                                    continue; // next j
-                                }
-                                else
-                                {
-                                    // swap
-                                    loopItems[i] = fromTail;
-                                    loopItems[j] = null;
-                                    j--;
-                                    goto NEXT_LOOP; // next i
-                                }
-                            }
-                            catch (Exception ex)
-                            {
-                                loopItems[j] = null;
-                                j--;
-                                try
-                                {
-                                    unhandledExceptionCallback(ex);
-                                }
-                                catch { }
-                                continue; // next j
-                            }
-                        }
-                        else
-                        {
-                            j--;
-                        }
-                    }
-
-                    tail = i; // loop end
-                    break; // LOOP END
-
-                    NEXT_LOOP:
-                    continue;
-                }
-
-
-                lock (runningAndQueueLock)
-                {
-                    running = false;
-                    while (waitQueue.Count != 0)
-                    {
-                        if (loopItems.Length == tail)
-                        {
-                            Array.Resize(ref loopItems, checked(tail * 2));
-                        }
-                        loopItems[tail++] = waitQueue.Dequeue();
-                    }
-                }
-            }
-        }
     }
 }
 

src/UniTask/Assets/Plugins/UniTask/Runtime/PlayerLoopHelper.cs

@@ -98,6 +98,16 @@ namespace Cysharp.Threading.Tasks
 #endif
     }
 
+    public enum EngineCallbackTiming
+    {
+        OnBeforeRender = 0,
+
+        WillRenderCanvases = 1,
+#if UNITY_2021_3_OR_NEWER
+        PreWillRenderCanvases = 2,
+#endif
+    }
+
     [Flags]
     public enum InjectPlayerLoopTimings
     {
@@ -171,6 +181,25 @@ namespace Cysharp.Threading.Tasks
 #endif
     }
 
+    [Flags]
+    public enum InjectEngineCallbackTimings
+    {
+        All =
+            OnBeforeRender |
+            WillRenderCanvases
+#if UNITY_2021_3_OR_NEWER
+            | PreWillRenderCanvases
+#endif
+        ,
+
+        OnBeforeRender = 1,
+
+        WillRenderCanvases = 2,
+#if UNITY_2021_3_OR_NEWER
+        PreWillRenderCanvases = 4,
+#endif
+    }
+
     public interface IPlayerLoopItem
     {
         bool MoveNext();
@@ -192,6 +221,7 @@ namespace Cysharp.Threading.Tasks
         static SynchronizationContext unitySynchronizationContext;
         static ContinuationQueue[] yielders;
         static PlayerLoopRunner[] runners;
+        static EngineCallbackRunner[] callbackRunners;
         internal static bool IsEditorApplicationQuitting { get; private set; }
         static PlayerLoopSystem[] InsertRunner(PlayerLoopSystem loopSystem,
             bool injectOnFirst,
@@ -395,7 +425,20 @@ namespace Cysharp.Threading.Tasks
             }
         }
 
-        public static void Initialize(ref PlayerLoopSystem playerLoop, InjectPlayerLoopTimings injectTimings = InjectPlayerLoopTimings.All)
+        static bool GetInjectCallback(InjectEngineCallbackTimings injectTimings, InjectEngineCallbackTimings targetTimings,
+            int index, EngineCallbackTiming engineCallbackTiming, out EngineCallbackRunner runner)
+        {
+            runner = null;
+            if ((injectTimings & targetTimings) == targetTimings)
+            {
+                runner = (callbackRunners[index] = new EngineCallbackRunner(engineCallbackTiming));
+                return true;
+            }
+            return runner != null;
+        }
+
+        public static void Initialize(ref PlayerLoopSystem playerLoop, InjectPlayerLoopTimings injectTimings = InjectPlayerLoopTimings.All,
+            InjectEngineCallbackTimings injectCallbackTimings = InjectEngineCallbackTimings.All)
         {
 #if UNITY_2020_2_OR_NEWER
             yielders = new ContinuationQueue[16];
@@ -405,6 +448,12 @@ namespace Cysharp.Threading.Tasks
             runners = new PlayerLoopRunner[14];
 #endif
 
+#if UNITY_2021_3_OR_NEWER
+            callbackRunners = new EngineCallbackRunner[3];
+#else
+            callbackRunners = new EngineCallbackRunner[2];
+#endif
+
             var copyList = playerLoop.subSystemList.ToArray();
 
             // Initialization
@@ -487,6 +536,26 @@ namespace Cysharp.Threading.Tasks
 
             playerLoop.subSystemList = copyList;
             PlayerLoop.SetPlayerLoop(playerLoop);
+
+            if (GetInjectCallback(injectCallbackTimings, InjectEngineCallbackTimings.OnBeforeRender,
+                0, EngineCallbackTiming.OnBeforeRender, out var onBeforeRenderRunner))
+            {
+                Application.onBeforeRender += onBeforeRenderRunner.Run;
+            }
+
+            if (GetInjectCallback(injectCallbackTimings, InjectEngineCallbackTimings.WillRenderCanvases,
+                1, EngineCallbackTiming.WillRenderCanvases, out var willRenderCanvasesRunner))
+            {
+                Canvas.willRenderCanvases += willRenderCanvasesRunner.Run;
+            }
+
+#if UNITY_2021_3_OR_NEWER
+            if (GetInjectCallback(injectCallbackTimings, InjectEngineCallbackTimings.PreWillRenderCanvases,
+                2, EngineCallbackTiming.PreWillRenderCanvases, out var preWillRenderCanvasesRunner))
+            {
+                Canvas.preWillRenderCanvases += preWillRenderCanvasesRunner.Run;
+            }
+#endif
         }
 
         public static void AddAction(PlayerLoopTiming timing, IPlayerLoopItem action)
@@ -499,11 +568,26 @@ namespace Cysharp.Threading.Tasks
             runner.AddAction(action);
         }
 
+        public static void AddAction(EngineCallbackTiming timing, IPlayerLoopItem action)
+        {
+            var runner = callbackRunners[(int)timing];
+            if (runner == null)
+            {
+                ThrowInvalidCallbackTiming(timing);
+            }
+            runner.AddAction(action);
+        }
+
         static void ThrowInvalidLoopTiming(PlayerLoopTiming playerLoopTiming)
         {
             throw new InvalidOperationException("Target playerLoopTiming is not injected. Please check PlayerLoopHelper.Initialize. PlayerLoopTiming:" + playerLoopTiming);
         }
 
+        static void ThrowInvalidCallbackTiming(EngineCallbackTiming engineCallbackTiming)
+        {
+            throw new InvalidOperationException("Target engineCallbackTiming is not injected. Please check PlayerLoopHelper.Initialize. EngineCallbackTiming:" + engineCallbackTiming);
+        }
+
         public static void AddContinuation(PlayerLoopTiming timing, Action continuation)
         {
             var q = yielders[(int)timing];
@@ -514,6 +598,16 @@ namespace Cysharp.Threading.Tasks
             q.Enqueue(continuation);
         }
 
+        public static void AddContinuation(EngineCallbackTiming timing, Action continuation)
+        {
+            var q = yielders[(int)timing];
+            if (q == null)
+            {
+                ThrowInvalidCallbackTiming(timing);
+            }
+            q.Enqueue(continuation);
+        }
+
         // Diagnostics helper
 
 #if UNITY_2019_3_OR_NEWER

src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.Delay.cs

@@ -195,6 +195,23 @@ namespace Cysharp.Threading.Tasks
             }
         }
 
+        public static UniTask WaitForOnBeforeRender(CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false)
+        {
+            return new UniTask(YieldPromise.Create(EngineCallbackTiming.OnBeforeRender, cancellationToken, cancelImmediately, out var token), token);
+        }
+
+        public static UniTask WaitForWillRenderCanvases(CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false)
+        {
+            return new UniTask(YieldPromise.Create(EngineCallbackTiming.WillRenderCanvases, cancellationToken, cancelImmediately, out var token), token);
+        }
+
+#if UNITY_2021_3_OR_NEWER
+        public static UniTask WaitForPreWillRenderCanvases(CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false)
+        {
+            return new UniTask(YieldPromise.Create(EngineCallbackTiming.PreWillRenderCanvases, cancellationToken, cancelImmediately, out var token), token);
+        }
+#endif
+
         sealed class YieldPromise : IUniTaskSource, IPlayerLoopItem, ITaskPoolNode<YieldPromise>
         {
             static TaskPool<YieldPromise> pool;
@@ -214,20 +231,15 @@ namespace Cysharp.Threading.Tasks
             {
             }
 
-            public static IUniTaskSource Create(PlayerLoopTiming timing, CancellationToken cancellationToken, bool cancelImmediately, out short token)
+            private static YieldPromise Create(CancellationToken cancellationToken, bool cancelImmediately)
             {
-                if (cancellationToken.IsCancellationRequested)
-                {
-                    return AutoResetUniTaskCompletionSource.CreateFromCanceled(cancellationToken, out token);
-                }
-
                 if (!pool.TryPop(out var result))
                 {
                     result = new YieldPromise();
                 }
 
                 result.cancellationToken = cancellationToken;
-                
+
                 if (cancelImmediately && cancellationToken.CanBeCanceled)
                 {
                     result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
@@ -237,6 +249,35 @@ namespace Cysharp.Threading.Tasks
                     }, result);
                 }
 
+                return result;
+            }
+
+            public static IUniTaskSource Create(EngineCallbackTiming timing, CancellationToken cancellationToken, bool cancelImmediately, out short token)
+            {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    return AutoResetUniTaskCompletionSource.CreateFromCanceled(cancellationToken, out token);
+                }
+
+                var result = Create(cancellationToken, cancelImmediately);
+
+                TaskTracker.TrackActiveTask(result, 3);
+
+                PlayerLoopHelper.AddAction(timing, result);
+
+                token = result.core.Version;
+                return result;
+            }
+
+            public static IUniTaskSource Create(PlayerLoopTiming timing, CancellationToken cancellationToken, bool cancelImmediately, out short token)
+            {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    return AutoResetUniTaskCompletionSource.CreateFromCanceled(cancellationToken, out token);
+                }
+
+                var result = Create(cancellationToken, cancelImmediately);
+
                 TaskTracker.TrackActiveTask(result, 3);
 
                 PlayerLoopHelper.AddAction(timing, result);

src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.Threading.cs

@@ -52,6 +52,14 @@ namespace Cysharp.Threading.Tasks
             PlayerLoopHelper.AddContinuation(timing, action);
         }
 
+        /// <summary>
+        /// Queue the action to an Engine Callback.
+        /// </summary>
+        public static void Post(Action action, EngineCallbackTiming timing)
+        {
+            PlayerLoopHelper.AddContinuation(timing, action);
+        }
+
 #endif
 
         public static SwitchToThreadPoolAwaitable SwitchToThreadPool()