Unityのコルーチンをawaitできるようにする(CancellationToken対応)

環境

Unity 2018.3
UniRx - Reactive Extensions for Unity / Ver 6.2.2

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

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

await CoroutineFunction()
がUnityのコルーチンの動作と異なるのが気になる人

Unityのコルーチンをawaitしたい！

2018/12/29 追記

この件、@toRisouPさんに教えていただきました！

あー、これ、IEnumeratorのAwaiterになっちゃうのか。ひとまずawait CoroutineTest().ToYieldInstruction(token)って書けばコルーチンを起動してawaitできるはずかな。 https://t.co/tLx94tYVIf

— とりすーぷ@2日目メ08a (@toRisouP) 2018年12月28日

実際に試してみましたところ
UniRxをusingすると、IEnumeratorの拡張メソッドとして追加されるToYieldInstructionを使うと
期待通りの動作をしましたので、この記事のモノを使うより、そっちをお使いください

以下の記事は記念に残しておきます

本来UniRx.Asyncの機能でIEnumeratorもawait出来るのだけども

自分は今、UniRx.Async.UniTaskの力を使い、コルーチンからasync/awaitに移行している途中
現在コルーチンとasync/awaitが混在している状態

そんななかコルーチンで書いた処理をawaitしようとすると
普通にUnityのコルーチンで実行したときと
UniTaskの機能を使いawaitで待った時とで挙動が違うことに気が付いた

簡単な例

public class TestCoroutineAwait : MonoBehaviour
{
    void Start()
    {
        // コルーチンで実行
        StartCoroutine(CoroutineTest());

        // async/awaitで実行
        Go(this.GetCancellationTokenOnDestroy()).Forget();
    }

    async UniTask Go(CancellationToken token)
    {
        try
        {
            // UniRxの機能のConfigureAwaitでCancellationTokenを渡す
            await CoroutineTest().ConfigureAwait(cancellationToken: token);
        }
        finally
        {
            Debug.Log("End");
        }
    }

    IEnumerator CoroutineTest()
    {
        // 10カウントする
        for (int i=0; i<10; ++i)
        {
            Debug.Log(i);
            yield return new WaitForSeconds(1);
        }
    }
}

どう挙動が違う？

上記コードは、コルーチンで実行するなら1秒に1つづつ10秒かけてカウントアップするが、
async/awaitの場合、瞬時に出力されて終わるようだ(1フレームに1つ出力されてる？)

また、WaitForEndOfFrameのタイミングもスルーするのか
以下のコードをasync/awaitで実行しようとすると
「このタイミングではスクリーンバッファはReadPixelできない」というエラーとなる

    IEnumerator CoroutineTest()
    {
        { // スクリーンショットをとる
            // 描画が終わったタイミングで処理が戻るWaitForEndOfFrameを使う
            yield return new WaitForEndOfFrame();

            var texture = new Texture2D(Screen.width, Screen.height, TextureFormat.ARGB32, false);
            RenderTexture.active = null;
            texture.ReadPixels(new Rect(0, 0, Screen.width, Screen.height), 0, 0);
            texture.Apply();
            GameObject.DestroyImmediate(texture);
        }
    }

どうするか

今後Coroutineは使わない方向になるだろうし
とりあえずの間に合わせに簡単に自前でコルーチンを待つAwaiterを書いてみた

もちろん、CancellationTokenへの対応は個人的に必須なのでそこは外さずいれてある

自前のコルーチンAwaiter

// (2018/12/31 17:02追記) CancellationToken.Registerの登録周りで問題があったので書き直した

// コルーチン用Awaiter
public struct CoroutineAwaiter : System.Runtime.CompilerServices.ICriticalNotifyCompletion
{
    CancellationToken token;
    IEnumerator e;

    public CoroutineAwaiter(IEnumerator e, CancellationToken token)
    {
        this.e = e;
        this.token = token;
    }

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

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

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

    public void UnsafeOnCompleted(Action continuation)
    {
        CancellationTokenRegistration regist = new CancellationTokenRegistration();

        Coroutine c = CoroutineExecutor.Begin(e, () =>
        {
            // CancellationToken.Registerで登録したのを解除
            regist.Dispose();
            continuation();
        });

        regist = token.Register(() =>
        {
            CoroutineExecutor.End(c);
            continuation();
        });
    }

    // 我はawaiterでありながらawaitableでもあるのだ
    public CoroutineAwaiter GetAwaiter() => this;
}

// コルーチンを実行するためのシングルトン
public class CoroutineExecutor : MonoBehaviour
{
    static CoroutineExecutor _instance;

    static CoroutineExecutor instance
    {
        get
        {
            if (_instance == null)
            {
                Init();
            }
            return _instance;
        }
    }

    static void Init()
    {
        var g = new GameObject();
        g.name = "CoroutineExecutor";
        GameObject.DontDestroyOnLoad(g);
        _instance = g.AddComponent<CoroutineExecutor>();
    }

    public static Coroutine Begin(IEnumerator e, System.Action continuation)
    {
        return instance.StartCoroutine(Exec(e, continuation));
    }

    public static void End(Coroutine c)
    {
        instance.StopCoroutine(c);
    }

    static IEnumerator Exec(IEnumerator e, System.Action continuation)
    {
        yield return e;
        continuation();
    }
}

自前のコルーチンAwaiterの使用方法

上記のコードをファイルに保存し以下のように記述する

await new CoroutineAwaiter(CoroutineTest(), token);

自前のコルーチンAwaiterの使用例

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

    async UniTask Go(CancellationToken token)
    {
        try
        {
            // CoroutineAwaiterにIEnumeratorとCancellationTokenを渡しAwaiterとして包み込む
            await new CoroutineAwaiter(CoroutineScreenShot(), token);

            await new CoroutineAwaiter(CoroutineTest(), token);
        }
        finally
        {
            Debug.Log("End");
        }
    }

    IEnumerator CoroutineScreenShot()
    {
        yield return new WaitForEndOfFrame();

        var texture = new Texture2D(Screen.width, Screen.height, TextureFormat.ARGB32, false);
        RenderTexture.active = null;
        texture.ReadPixels(new Rect(0, 0, Screen.width, Screen.height), 0, 0);
        texture.Apply();
        GameObject.DestroyImmediate(texture);
    }

    IEnumerator CoroutineTest()
    {
        // 10カウント
        for (int i=0; i<10; ++i)
        {
            Debug.Log(i);
            yield return new WaitForSeconds(1);
        }
    }
}

終わりに

これで正しいのかいまいち自信が無いのですが、とりあえず期待通りに動いてるので
コルーチンを駆逐するまでこれでごまかしていきたいという所信表明

また、自分がUniRx.Asyncの使い方をそもそも間違えてる可能性が大きいので
間違ってたらご指摘お願いします

ライセンス

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