UniTask2のAutoResetUniTaskCompletionSourceでは2回TrySetしてはいけない

この記事のまとめ

AutoResetUniTaskCompletionSourceをawaitして結果を取得したあとで、さらにTrySetResult, TrySetCanceled, TrySetExceptionなどを呼び出すと次回以降の動作がおかしくなる。

awaitして結果を取得した時点でそのAutoResetUniTaskCompletionSourceはすでに返却されているので触ってはいけない。

複数箇所でTrySetXXXを呼び出していると事故る可能性がある。

もう少し詳しい説明

AutoResetUniTaskCompletionSourceとは

AutoResetUniTaskCompletionSourceを使うとアロケーションを抑えてUniTaskCompletionSourceを使えます。

例えば以下のような形で使えます。

// 1秒後に1を返すメソッド
private async UniTask<int> GetIntAsync()
{
    AutoResetUniTaskCompletionSource<int> completionSource = AutoResetUniTaskCompletionSource<int>.Create();

    // 別スレッドで1秒待ってから結果として1をセット
    UniTask.Create(async () =>
    {
        await UniTask.Delay(TimeSpan.FromSeconds(1));
        completionSource.TrySetResult(1);
    });

    var result = await completionSource.Task;
    return result;
}

// GetIntAsyncの結果を出力する
public async UniTaskVoid Start()
{
    var int1 = await GetIntAsync();
    Debug.Log($"{int1}"); // 「1」が出力される
}

問題になる例

awaitした後でまたTrySetResultを呼び出すと、次にAutoResetUniTaskCompletionSourceを使ったときに意図通りに動かなくなります。

例えば以下のようなコードです。
AutoResetUniTaskCompletionSourceをawaitした後でさらにcompletionSource.TrySetResult(2);を呼び出してしまっています。

// 1秒後に1を返すメソッド
private async UniTask<int> GetIntAsync()
{
    AutoResetUniTaskCompletionSource<int> completionSource = AutoResetUniTaskCompletionSource<int>.Create();

    // 別スレッドで1秒待ってから結果として1をセット
    UniTask.Create(async () =>
    {
        await UniTask.Delay(TimeSpan.FromSeconds(1));
        completionSource.TrySetResult(1);
    });

    var result = await completionSource.Task;

    completionSource.TrySetResult(2); // ← この行を追加

    return result;
}

// GetIntAsyncの結果を出力する
public async UniTaskVoid Start()
{
    var int1 = await GetIntAsync();
    var int2 = await GetIntAsync();
    Debug.Log($"{int1}, {int2}"); // 「1, 2」が出力される
}

なぜ「1, 2」と出力されるのか

AutoResetUniTaskCompletionSource<int>.Create()は内部のオブジェクトプールから未使用のAutoResetUniTaskCompletionSourceが取り出されます。
そのおかげで毎回UniTaskCompletionSourceをnewするのを防ぐことができます。

そして取り出されたAutoResetUniTaskCompletionSourceはawaitして結果が返された時点でオブジェクトプールに返却されます。
上記のコードで言えばvar result = await completionSource.Task;でresultに結果が格納された時点になります。

しかし、返却された後でもCreateで借りてきたAutoResetUniTaskCompletionSource自体にはアクセスできてしまいます。
これに対して上記コードのcompletionSource.TrySetResult(2);のようにアクセスしてはいけません。
すでに状態がリセっトされてプールされているAutoResetUniTaskCompletionSourceに対してTrySetResult(2)が呼び出してしまっています。

それによって、次にAutoResetUniTaskCompletionSource<int>.Create()で借りてくるのはTrySetResult(2)が呼び出されてしまっているAutoResetUniTaskCompletionSourceになります。
なので2回目のGetIntAsyncは1秒待たずにすぐに2を返す動作となってしまっています。

もう少し実践的な例

前のコードはawaitした直後にTrySetResultを呼び出すという不自然すぎる例でした。
もうすこし気付きづらい例として次のコードでも似たような問題が発生します。

private async UniTask<int> GetIntAsync()
{
    AutoResetUniTaskCompletionSource<int> completionSource = AutoResetUniTaskCompletionSource<int>.Create();

    // 別スレッドで1秒待ってから結果として1をセット
    UniTask.Create(async () =>
    {
        await UniTask.Delay(TimeSpan.FromSeconds(1));
        completionSource.TrySetResult(1);
    });

    // 別スレッドで2秒待ってから結果として2をセット
    UniTask.Create(async () =>
    {
        await UniTask.Delay(TimeSpan.FromSeconds(2));
        completionSource.TrySetResult(2);
    });

    var result = await completionSource.Task;
    return result;
}

// GetIntAsyncの結果を出力する
public async UniTaskVoid Start()
{
    var int1 = await GetIntAsync();
    await UniTask.Delay(TimeSpan.FromSeconds(3));
    var int2 = await GetIntAsync();
    Debug.Log($"{int1}, {int2}"); // 「1, 2」が出力される
}

今回のGetIntAsyncは2つスレッドが走ってどちらか早い方の結果を返す形になっています。
最初のGetIntAsyncの呼び出しでは1秒待って1がセットされます。

そしてStartメソッドの中で3秒待ってから2回目のGetIntAsyncを呼び出しています。
この3秒待ちの間にもう1つのスレッドの方でオブジェクトプール内のAutoResetUniTaskCompletionSourceに2がセットされてしまいます。
結果として2回目のGetIntAsyncではすぐに2が返ってきてしまいます。

今回はTrySetResultが2回呼ばれる例でしたが、TrySetCanceledやTrySetExceptionも同様です。
TrySetResult、TrySetCanceled、TrySetExceptionのどれかが1度だけ呼ばれるように実装しなければなりません。

回避策

スレッド内で中断処理を実装して1度しかTrySetResultが呼ばれないようにすれば意図通りに動きます。

private async UniTask<int> GetIntAsync()
{
    AutoResetUniTaskCompletionSource<int> completionSource = AutoResetUniTaskCompletionSource<int>.Create();

    var cancellationTokenSource = new CancellationTokenSource();

    // 別スレッドで1秒待ってから結果として1をセット
    UniTask.Create(async () =>
    {
        await UniTask.Delay(TimeSpan.FromSeconds(1), cancellationToken:cancellationTokenSource.Token);
        lock (cancellationTokenSource)
        {
            if (cancellationTokenSource.IsCancellationRequested)
            {
                return;
            }
            cancellationTokenSource.Cancel();
        }
        completionSource.TrySetResult(1);
    });

    // 別スレッドで2秒待ってから結果として2をセット
    UniTask.Create(async () =>
    {
        await UniTask.Delay(TimeSpan.FromSeconds(2), cancellationToken:cancellationTokenSource.Token);
        lock (cancellationTokenSource)
        {
            if (cancellationTokenSource.IsCancellationRequested)
            {
                return;
            }
            cancellationTokenSource.Cancel();
        }
        completionSource.TrySetResult(2);
    });

    var result = await completionSource.Task;
    return result;
}

// GetIntAsyncの結果を出力する
public async UniTaskVoid Start()
{
    var int1 = await GetIntAsync();
    await UniTask.Delay(TimeSpan.FromSeconds(3));
    var int2 = await GetIntAsync();
    Debug.Log($"{int1}, {int2}"); // 「1, 1」が出力される
}

メモ

issueとしても投げたがオブジェクトプール内でTrySetXXXが呼ばれたら例外投げてくれるとわかりやすい気がする。

と思ったけどSystem.Buffers.ArrayPoolなどもReturnした後でアクセスすればおかしくなるのでそういうものか？
AutoResetUniTaskCompletionSourceは返却されるタイミングが非明示的でわかりづらいのと、非同期やマルチスレッドで使われることが多いのが問題を大きくしている。

色々と面倒なのでAutoResetUniTaskCompletionSourceではなく常にUniTaskCompletionSourceを使うようにした方が潜在的なバグを防げそう。