async/awaitでメモリ確保の起きないUpdate Loopを書く

環境

Unity 2018.3
UniRx - Reactive Extensions for Unity / Ver 6.2.2

想定している読者ターゲット

async/awaitを理解している人
UniTaskを少し理解している人
Awaiterの仕組みをある程度理解している人

UniTask Yield(PlayerLoopTiming timing, CancellationToken cancellationToken)
がメモリ確保していることが気になる人

概要

async内で毎フレーム更新するループを記述したい！

async UniTask Loop(CancellationToken cancellationToken){
    while(true){
        // 1フレーム待機する
        // Cancelされたらこのループを終わらせる
        await OneFrame(cancellationToken);

        // 更新する
        transform.position += Vector.forward * Time.deltaTime;
    }
}

↑こんな感じで！

コルーチンで書くと↓こんな感じ

IEnumerator Loop(){
    while(true){
        // 1フレーム待機する
        yield return null;

        // 更新する
        transform.position += Vector.forward * Time.deltaTime;
    }
}

これをするには「1フレーム待機し、CancelされたらCancellation例外を送出するAwaiter」が必要となる

1フレーム待機させるには、AwaiterのUnsafeOnCompletedで渡されるcontinuationを
次のフレームで実行するような仕組みをつくればすぐ出来る

早速出来たコードを見てみよう

１フレームだけ待機するAwaiterのコード

// メモリ確保が起きないよう構造体で書こう
public struct OneFrameAwaitable
{
    // ここも構造体
    public struct OneFrameAwaiter : System.Runtime.CompilerServices.ICriticalNotifyCompletion
    {
        CancellationToken token;

        public OneFrameAwaiter(CancellationToken token)
        {
            this.token = token;
        }

        // falseを返すべきなんだ
        public bool IsCompleted => false;

        // GetResultでThrowIfCancellationRequestedを呼び出し、キャンセルされてたらCancellation例外送出
        public void GetResult() => token.ThrowIfCancellationRequested();

        // UnsafeOnCompletedを呼び出す
        public void OnCompleted(Action continuation) => UnsafeOnCompleted(continuation);

        // 後述のOneFrameManagerにcontinuationを登録する
        public void UnsafeOnCompleted(Action continuation) => OneFrameManager.Register(continuation);
    }

    OneFrameAwaiter awaiter;

    public OneFrameAwaitable(CancellationToken cancellationToken = default)
    {
        // このnewは構造体なのでアロケーションないよ
        awaiter = new OneFrameAwaiter(cancellationToken);
    }

    public OneFrameAwaiter GetAwaiter() => awaiter;
}

// OneFrameAwaiterのUnsafeOnCompletedで渡されたActionを次のフレームに実行するためのマネージャ
public class OneFrameManager : MonoBehaviour
{
    static List<Action> actions = new List<Action>();
    static List<Action> actions2 = new List<Action>();

    [RuntimeInitializeOnLoadMethod]
    static void Init()
    {
        var g = new GameObject();
        g.name = "OneFrameManager ";
        GameObject.DontDestroyOnLoad(g);
        g.AddComponent<OneFrameManager>();
    }

    public static void Register(Action a)
    {
        actions2.Add(a);
    }

    private void Start()
    {
        StartCoroutine(WaitForEndOfFrame());

        IEnumerator WaitForEndOfFrame()
        {
            // WaitForEndOfFrameを最初にnewしておいて使いまわす
            var a = new WaitForEndOfFrame(); 

            while (true)
            {
                yield return a;

                // WaitForEndOfFrameのタイミングでActionをSwapして次のUpdateで実行するようにする

                var temp1 = actions;
                actions = actions2;
                actions2 = temp1;
                actions2.Clear();
            }
        }
    }

    private void Update()
    {
        // 登録されたActionを全部実行する
        foreach (var a in actions)
        {
            a.Invoke();
        }

        actions.Clear();
    }
}

１フレーム待機するOneFrameAwaiterと、OneFrameAwaiterを返すOneFrameAwaitable
次のフレームに実行するためのマネージャのOneFrameManagerという構成でできている

テストコード

public class TestTween : MonoBehaviour
{
    void Start()
    {
        Sequence(this.GetCancellationTokenOnDestroy()).Forget();
    }

    private void Update()
    {
        Debug.Log("Update");
    }

    async UniTask Sequence(CancellationToken cancellationToken)
    {
        try
        {
            while (true)
            {
                await new OneFrameAwaitable(cancellationToken: cancellationToken);
                Debug.Log("async");
            }
        }
        finally
        {
            Debug.Log("End");
        }
    }
}

テストコードを実行すると、asyncとUpdateが交互にログ出力されていることがわかる
これで、次のフレームまで待機するAwaiterが完成した

UniTaskに無いの？

ある

UniTask.Yieldがそれだ
PlayerLoopTimingというものが渡せるため、上で示したものより柔軟性があってFixedUpdateで実行したりなんて指定も出来る

じゃあそれを使えばいいじゃないか、おしまい

...
と、したいのだが、一つだけ問題がある

UniTask.Yieldの引数にCancellationTokenを渡さないバージョンはGC.aloocが起きないのだが、
CancellationTokenを渡すバージョンはメモリ確保が起きてしまうことだ

64B x await UniTask.Yield(cancellationToken: cancellationToken)の数だけメモリ確保が起きる

これだとUpdate Loopに使うのは少し抵抗が出てしまう

これを回避するには await UniTask.Yield(); と引数を渡さずに呼び出し
その直後にcancellationToken.ThrowIfCancellationRequested();を呼び出せばいいが
少し面倒だ

UniTask.Yield

なぜメモリ確保が起きるのか、実際のコードを見てみよう

public static YieldAwaitable Yield(PlayerLoopTiming timing = PlayerLoopTiming.Update)
{
    // optimized for single continuation
    return new YieldAwaitable(timing);
}

public static UniTask Yield(PlayerLoopTiming timing, CancellationToken cancellationToken)
{
    return new UniTask(new YieldPromise(timing, cancellationToken));
}

YieldAwaitableはstructだが、YieldPromiseはclassなため、Yieldを呼び出すたびにメモリ確保が発生するようだ

この記事を書いてるうちに気づいたが

UniTaskの仕組みを拝借すれば、OneFrameManager相当のものを無くし
FixedUpdateやUpdateなどに処理を振り分ける機能を追加した上でメモリ確保を無くしたものができそうだ

メモリ確保を無くし、UniTaskの機能の上で動く1フレーム待機するAwaiter

public struct YieldAwaitable
{
    public struct YieldAwaiter : System.Runtime.CompilerServices.ICriticalNotifyCompletion
    {
        CancellationToken token;
        PlayerLoopTiming timing;

        public YieldAwaiter(PlayerLoopTiming timing, CancellationToken token)
        {
            this.token = token;
            this.timing = timing;
        }

        // falseを返すべきなんだ
        public bool IsCompleted => false;

        // GetResultでThrowIfCancellationRequestedを呼び出し、キャンセルされてたら例外
        public void GetResult() => token.ThrowIfCancellationRequested();

        // UnsafeOnCompletedを呼び出す
        public void OnCompleted(Action continuation) => UnsafeOnCompleted(continuation);

        // UniRx.AsyncのPlayerLoopHelperに登録する
        public void UnsafeOnCompleted(Action continuation) => PlayerLoopHelper.AddContinuation(timing, continuation);
    }

    YieldAwaiter awaiter;

    public YieldAwaitable(PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default)
    {
        // このnewは構造体なのでアロケーションないよ
        awaiter = new YieldAwaiter(timing, cancellationToken);
    }

    public YieldAwaiter GetAwaiter() => awaiter;
}

出来た

終わりに

自分の期待通りに動いてはいるのだけど、これが正しい実装と言えるのかは自信がまったくないです

間違ってるところとかありましたらご指摘お願いいたします

ライセンス

この記事の私の書いた部分のソースコードのライセンスはMIT LicenseかApache 2.0 Licenseのお好きな方でどうぞ