Unity
UniRx

Unity UniRxとasync/awaitでフレーム管理

やりたいこと

async/awaitをコルーチンの代わりに使いたい。
そのためにはコルーチンでできることをasync/awaitで書く必要がある。

それの動作確認メモ。UniRxを使ってます。

計測方法にミスがあったので修正

void Start()
{
    // MainThreadDispatcherの初期化タイミングが計測に影響しないように
    // 確実に初期化しておく
    MainThreadDispatcher.Initialize();

    Observable.NextFrame()
        .Subscribe(_ =>
        {
            // MainThreadDispatcherの初期化が終わってから実行する
            DoAsync();
        });
}

このような呼び出し方をしていたが、この場合はコルーチンの実行タイミングがずれそうな気もするので、次のような測定方法に変更。

void Start()
{
    // MainThreadDispatcherの初期化タイミングが計測に影響しないように
    // 確実に初期化しておく
    MainThreadDispatcher.Initialize();
    DoAsync();
}

void Update()
{
    ++frameCounter;
}


async void DoAsync()
{
    // 初期化待機
    await Task.Delay(1000);

    Debug.Log("[1] +" + frameCounter));

    await Observable.Return(Unit.Default, Scheduler.MainThread);

    Debug.Log("[2] +" + frameCounter));
}

また、 Time.frameCount も挙動が微妙に怪しいので、Update()内で自前でカウンタを増やすことにした。

やったこと

フレーム数の制御

例えば、コルーチンの場合はfor文とyield return nullを使って指定フレーム数後に処理を実行するということができる。

コルーチンで待機
IEnumerator WaitFrameCoroutine()
{
    var startFrame = frameCounter;

    Debug.Log("[1] +" + (frameCounter - startFrame));

    // 5フレーム後に実行
    for (var i = 0; i < 5; i++)
    {
        yield return null;
    }

    Debug.Log("[2] +" + (frameCounter - startFrame));
}
実行結果
[1] + 0
[2] + 5

これを async/await でやるにはどうしたらいいかを試してみる。

試作1: Observable.TimerFrameを使ってみる

とりあえず、シンプルに Observable.TimerFrame を使ってみる。

async void DoAsync()
{
    // 初期化待機
    await Task.Delay(1000);

    var startFrame = frameCounter;

    Debug.Log("[1] +" + (frameCounter - startFrame));

    await Observable.TimerFrame(5); //5フレーム待ちたいが…

    Debug.Log("[2] +" + (frameCounter - startFrame));
}
実行結果
[1] + 0
[2] + 6

1フレームずれてしまった。
これは恐らく、SynchronizationContext.Postを実行したときに、処理の継続が次のフレームに持ち越されてしまうのが原因であると考えられる。
つまり、「awaitに1フレーム」要しているということになる。

というわけで、「待機したいフレーム数-1」の値を設定すればObservable.TimerFrameは使えそうかな…?と思ったけど、ちょっと怖い部分はある。

Observable.TimerFrameの不安要素

Observable.TimerFrameのフレーム数カウントだが、コルーチン内で呼び出した場合は結果がズレる可能性がある。
というのも、Observable.TimerFrame自体が内部でコルーチンを使っているためである。

Update()内でObservable.TimerFrameを実行した場合、そのままそのフレーム中のコルーチンでフレーム数カウントが開始される。
だが、どこかのコルーチン中でObservable.TimerFrameを実行した場合、Observable.TimerFrameの起動は次のフレームに持ち越されてしまう。

そのため、結果が1フレームズレる可能性がある。なので使う場所をちゃんと見極めないと怖い。

試作2: BatchFrameを使ってみる

BatchFrameは、指定したフレーム数の間、OnNextメッセージをキャッシュしてくれるオペレータである。これを流用して無理矢理動かしてみた。

async void DoAsync()
{
    // 初期化待機
    await Task.Delay(1000);
    var startFrame = frameCounter;

    Debug.Log("[1] +" + (frameCounter - startFrame));

    await Observable.Return(Unit.Default).BatchFrame(5, FrameCountType.EndOfFrame);

    Debug.Log("[2] +" + (frameCounter - startFrame));
}
実行結果
[1] + 0
[2] + 0

なんと、何もawaitせずに通過してしまった。
それもそのはず。BatchFrameOnCompletedメッセージがきたら即処理を終了してしまうからである。
そのためこの手法は全然ダメ。

試作3: Observable.Rangeを使ってみる

Observable.Rangeは指定した回数、OnNextメッセージを発行するファクトリメソッドである。
このObservable.Rangeだが、実はスケジューラの設定が可能であり、MainThreadSchedulerを指定すると1フレームに1個メッセージを発行するという挙動になる。
この挙動をうまく使って指定フレーム数待機できるか確認してみた。

async void DoAsync()
{
    // 初期化待機
    await Task.Delay(1000);
    var startFrame = frameCounter;

    Debug.Log("[1] +" + (frameCounter - startFrame));

    await Observable.Range(0, 5, Scheduler.MainThread);

    Debug.Log("[2] +" + (frameCounter - startFrame));
}
実行結果
[1] + 0
[2] + 6

Observable.TimerFrameと同じ結論になってしまった。たぶん原因は同じ。

別のアプローチを試す

そもそも、 1フレームぴったり待機することができるのかを先に調べたほうがいい気がしたため検証してみた。

1フレーム待ってみる

UnitySynchronizationContextに処理がPost()されると1フレームずれる」という性質を悪用して、1フレーム強引に待機できないか試してみる。

Observable.Return
async void DoAsync()
{
    var startFrame = frameCounter;

    Debug.Log("[1] +" + (frameCounter - startFrame));

    await Observable.Return(Unit.Default, Scheduler.MainThread);

    Debug.Log("[2] +" + (frameCounter - startFrame));
}
実行結果
[1] + 0
[2] + 1

いけた。けどやりたいことに対して高コストすぎる気はする。

コルーチンの機能を使ってみる

コルーチンにはWaitWhile()という、指定条件を満たす間待機するコルーチンを作る機構がある。
これを使ってawaitしてしまう。

async void DoAsync()
{
    var startFrame = Time.frameCount;

    Debug.Log("[1] +" + (Time.frameCount - startFrame));

    await WaitFrame(5);

    Debug.Log("[2] +" + (Time.frameCount - startFrame));
}

/// <summary>
/// 指定フレーム数待機する
/// frame > 0 でないと正しくうごかない
/// </summary>
WaitWhile WaitFrame(int frame)
{
    var s = Time.frameCount;
    return new WaitWhile(() => (Time.frameCount - s) < frame);
}
実行結果
[1] + 0
[2] + 5

できた。

[フレーム数の制御]の結論

UniRxの機能を使ってゴリゴリやればフレーム数のカウントはできなくはなさそうだけど、結局はWaitWhile/WaitUntilawaitするのがラクそう。
ただし、このWaitWhile/WaitUntilawaitそのものもUniRxの機能ではある。

EndOfFrameをawaitで実現する

その場でコルーチンを起動して、EndOfFrameのタイミングで完了して、そしてSynchronizationContextを使わずにそのままINotifyCompletion.OnCompleted()を叩くようにできればイケる気はするが、手間がかかる。

じゃあUniRxでどうにかできないかな?とも思ったが、UniRxのAwaiterの実装そのままでは厳しそう。
というのも、UniRxではAwaiterとしてAsyncSubjectを使っている。
この実装を覗いてみると、

        class AwaitObserver : IObserver<T>
        {
            private readonly SynchronizationContext _context;
            private readonly Action _callback;

            public AwaitObserver(Action callback, bool originalContext)
            {
                if (originalContext)
                    _context = SynchronizationContext.Current;

                _callback = callback;
            }

            public void OnCompleted()
            {
                InvokeOnOriginalContext();
            }

            public void OnError(Exception error)
            {
                InvokeOnOriginalContext();
            }

            public void OnNext(T value)
            {
            }

            private void InvokeOnOriginalContext()
            {
                if (_context != null)
                {
                    //
                    // No need for OperationStarted and OperationCompleted calls here;
                    // this code is invoked through await support and will have a way
                    // to observe its start/complete behavior, either through returned
                    // Task objects or the async method builder's interaction with the
                    // SynchronizationContext object.
                    //
                    _context.Post(c => ((Action)c)(), _callback);
                }
                else
                {
                    _callback();
                }
            }
        }

AynscSubject.cs

となっている。AwaitObserverのコンストラクタの第二引数、bool originalContextfalseを指定できればなんとかなりそうである。

だが、awaitで待機を必要する場合、かならずtrueでここに到達するため、結局SynchronizationContextPost()されてしまい、処理のタイミングがズレてしまう。

/// <summary>
/// Specifies a callback action that will be invoked when the subject completes.
/// </summary>
/// <param name="continuation">Callback action that will be invoked when the subject completes.</param>
/// <exception cref="ArgumentNullException"><paramref name="continuation"/> is null.</exception>
public void OnCompleted(Action continuation)
{
    if (continuation == null)
        throw new ArgumentNullException("continuation");

    OnCompleted(continuation, true);
}

ここ

(INotifyCompletion.OnCompleted(Action continuation)は、await時に対象がまだ動作中に場合に呼び出されるコールバック登録用のメソッドである)

自前でAwaiterを作ってみる

じゃあ、SynchronizationContextを使わないAwaiterを用意すればいいのでは?ということで試してみる。

SynchronizationContextを使わずにコールバックを呼び出す実装
using System;
using System.Runtime.CompilerServices;
using UniRx;

namespace UniRxExtensions
{
    public static class UniRxExtension
    {
        public static ConfigureAwaitObservableAwaiter<T> NoSyncAwait<T>(this IObservable<T> observable)
        {
            return new ConfigureAwaitObservableAwaiter<T>(observable);
        }
    }

    public class ConfigureAwaitObservableAwaiter<T> : INotifyCompletion
    {
        private AsyncSubject<T> _originalAwaiter;

        public bool IsCompleted => _originalAwaiter.IsCompleted;

        public ConfigureAwaitObservableAwaiter(IObservable<T> observable)
        {
            _originalAwaiter = observable.GetAwaiter();
        }

        public ConfigureAwaitObservableAwaiter<T> GetAwaiter()
        {
            return this;
        }

        public void OnCompleted(Action continuation)
        {
            if (continuation == null)
                throw new ArgumentNullException("continuation");

            _originalAwaiter.Subscribe(_ => continuation());
        }

        public T GetResult()
        {
            return _originalAwaiter.GetResult();
        }
    }
}

ほぼほぼAsyncSubjectの機能に乗っかりつつ、OnCompletedの部分だけ書き換えた。
これでObservableの実行コンテキスト上でそのまま処理が進むはず。

使い方
//NoSyncAwait()を末尾につけてawaitすると、同期コンテキストを常に無視する
await Observable.Return(Unit.Default, Scheduler.MainThreadEndOfFrame).NoSyncAwait();

計測

どのタイミングで処理が終わったのかを明示的にするため、各タイミングでログを出すように変更。
その上で先程のNoSyncAwait()を挟んで実験。

EndOfFrameを待てるか
using System.Collections;
using UniRx;
using UniRxExtensions;
using UnityEngine;

namespace AwaitTest
{
    public class AwaitChecker : MonoBehaviour
    {
        public bool IsEnabled = false;

        int _frameCount = 0;

        void Start()
        {
            // MainThreadDispatcherの初期化タイミングが計測に影響しないように
            MainThreadDispatcher.Initialize();

            StartCoroutine(Coroutine());
        }

        /// <summary>
        /// 計測用の処理
        /// </summary>
        async void DoAsync()
        {
            Debug.Log("[-----start------] +" + _frameCount);

            await Observable.Return(Unit.Default, Scheduler.MainThreadEndOfFrame).NoSyncAwait();

            Debug.Log("[-----end------] +" + _frameCount);
        }

        void Update()
        {
            if (IsEnabled)
            {
                Debug.Log("Update + " + _frameCount);
            }

            if (_frameCount == 3)
            {
                // MainThreadDispatcher.Initialize()の初期化直後は計測がブレるので、
                // 数フレーム待ってから実行
                DoAsync();
            }
        }

        void LateUpdate()
        {
            if (IsEnabled)
            {
                Debug.Log("LateUpdate + " + _frameCount);
            }
        }

        IEnumerator Coroutine()
        {
            while (true)
            {
                yield return new WaitForEndOfFrame();
                if (IsEnabled)
                {
                    Debug.Log("EndOfFrame + " + _frameCount);
                }
                ++_frameCount;
            }
        }
    }
}

実行結果(該当部分のみ)
Update + 3
[-----start------] +3
LateUpdate + 3
EndOfFrame + 3
Update + 4
LateUpdate + 4
[-----end------] +4
EndOfFrame + 4

----end----- というログがLateUpdateEndOfFrameの間に出力されていることから、ちゃんと待てているっぽい。
フレーム数がずれてるが、これはWaitForEndOfFrame使ったときはもとからこうだった気もする。

[EndOfFrameをawaitで実現する]の結論

UniRxに手を加えて自前でAwaiterを作ればできなくはないけど、これだったらもうコルーチン使ったほうが楽じゃない?

まとめ

厳密なフレーム数の制御や、フレーム内でのタイミング制御をasync/awaitで実行するのはムズカシイ。
なので、それらの処理は素直にコルーチンで書いたほうが良さそう。

参考資料

An other world awaits you