【Unity】Awaitableの二度漬けには超ご用心

Unity 2023からAwaitableが導入され、フレームをまたぐ処理をasync/awaitスタイルで書きやすくなりました。

using System;
using System.Threading;
using UnityEngine;

public class Test : MonoBehaviour
{
    private void Start() => TestAsync(destroyCancellationToken);

    private async Awaitable TestAsync(CancellationToken ct)
    {
        try
        {
            Debug.Log($"frame {Time.frameCount}");

            await Awaitable.NextFrameAsync(ct);

            Debug.Log($"frame {Time.frameCount}");

            await Awaitable.NextFrameAsync(ct);

            Debug.Log($"frame {Time.frameCount}");
        }
        catch (Exception ex)
        {
            Debug.LogException(ex);
        }
    }
}

二度漬け禁止

Awaitableの実態は参照型として実装されたUnityEngine.Awaitableクラスです。参照型なので通常であれば作成ごとにヒープメモリ割り当てが発生するところですが、内部的にオブジェクトプーリングを利用してインスタンスを再利用することでメモリ割り当てを抑えています。

var awaitable1 = Awaitable.NextFrameAsync(ct); // 新しいAwaitableインスタンスが作成される 

await awaitable1; // awaitされたawaitable1は、完了時にプールされる

var awaitable2 = Awaitable.NextFrameAsync(ct); // プールから取り出されたawaitable1がawaitable2として再利用される

await awaitable2; // ふたたびプールに戻る

これを実現するために、Awaitableでは複数回のawaitが禁止されています。例えば次のコードは正しく動きません。

var awaitable1 = Awaitable.NextFrameAsync(ct);

await awaitable1; // awaitされたawaitable1は、完了時にプールされる
await awaitable1; // "InvalidOperationException: Awaitable is in detached state"

既にプールに戻った状態のAwaitableをawaitしようとしたので例外が吐かれました。

注意が必要なパターン

では、ちょっと意地悪して、プールに戻ったAwaitableを意図的にプールから出したうえで二度漬けしてみましょう。

var awaitable1 = Awaitable.NextFrameAsync(ct);

await awaitable1; // awaitされたawaitable1は、完了時にプールされる

var awaitable2 = Awaitable.WaitForSecondsAsync(10f, ct); // プールから取り出されたawaitable1がawaitable2として再利用される

await awaitable1; // 既に再利用されたawaitable2を待機してしまい……

結果として、二回目のawaitで例外は出ず、無関係のはずのWaitForSecondsAsyncのほうを待機してしまいます。

このくらい簡単なケースだと「そんなコード書かないだろ」と思うところですが、複雑なケースではうっかり二度漬けすることはあり得ます。例えば作成したAwaitableをすぐにawaitしないケースは要注意です。

さらに厄介なのが、実際にうっかり二度漬けしてしまった際の壊れ方です。二度漬けによって無関係のAwaitableをawaitしてしまった場合、本来そのAwaitableをawaitするはずだった側のコードでは、運が良ければ例外が出ますが、場合によってはさらに別の無関係のAwaitableをawaitしてしまうかもしれません。このように連鎖的な不整合の発生により、原因となった二度漬けコードを発見することが困難になる危険性があります。

ValueTask, UniTaskだと

ValueTask (IValueTaskSource使用時) や UniTaskでも同様のプーリング戦略が行われており、Awaitableと同じく二度漬け禁止ですが、うっかり二度漬けしても例外を吐けるように安全策が講じられています。

ValueTaskやUniTaskはAwaitableとは異なり構造体として定義されているため、それ自体のヒープアロケーションがありません。ただし、構造体の内部に参照型であるIValueTaskSourceやIUniTaskSourceといったオブジェクトの参照を保持しており、実際にプーリングされるのはこれらの"ソース型"のほうです。

https://github.com/Cysharp/UniTask/blob/08184af73794796192bf450a154ed0e1310a3e75/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.cs#L34-L38

さらに、ValueTask・UniTaskなどのタスク型とソース型はお互いの寿命を確認するためにtokenという数値を保持しています。このtokenはタスクの作成時にタスク型とソース型で同じ値がセットされるのですが、タスクが完了したタイミングで、ソース型側のtokenがインクリメントされるようになっています。結果としてタスク型とソース型は異なるtokenを持つことになり、タスク型はタスクが既に死んだことを知ることができ、適切に例外を吐くことができます。これによってソース型のオブジェクトを安全にプールに戻すことができるというわけです。

このパターンは以前に書いた記事で詳しい実装を紹介しているので、こちらもどうぞ👉
https://zenn.dev/ruccho/articles/e9876dadeca78a

余談

Awaitableを二度漬けした際の挙動は未定義と明記されています。

Important: The pooling of Awaitable instances means it’s never safe to await more than once on an Awaitable instance. Doing so can result in undefined behavior such as an exception or a deadlock.
https://docs.unity3d.com/6000.0/Documentation/Manual/async-awaitable-introduction.html

よってこれは意図的な設計であるわけですが、プーリングによるヒープメモリ確保回避と二度漬けセーフティを両立するための設計がすでにValueTaskやUniTaskで実現されているのに、Awaitableはなぜそうしなかったんでしょうね……？確かに構造体によるラップやtokenのチェックに多少のコストはかかるはずなので、Awaitableの設計にパフォーマンス上の利点がないとはいえないですが、ちょっと安全性を捨てすぎてない？と思ったりします。

おわり

Awaitableを使うときは二度漬けに超注意しましょう！