【UniTask】Addressableのロードキャンセル時に起こりがちなメモリリークをautoReleaseWhenCanceledで自動解放して防ぐ

UniTaskでAddressableのロードキャンセル時に起こりがちなメモリリークをautoReleaseWhenCanceledで自動解放して防ぐ方法についてまとめました。

キャンセル時にはアンロードが必要
挙動を確認する
実行結果
関連

キャンセル時にはアンロードが必要

いま、Addressable アセットシステムでアセットをロードする際に、途中で処理をキャンセルすることを考えます。

しかし Addressables のロードメソッドの引数にはそもそも CancellationToken が用意されていないので、自身で以下のようなメソッドを用意する必要があります。

using System.Threading;
using Cysharp.Threading.Tasks;
using UnityEngine;
using UnityEngine.AddressableAssets;

public class Example : MonoBehaviour
{
    private async void Start()
    {
        var cancellationTokenSource = new CancellationTokenSource();
        var task = LoadAssetAsync<GameObject>("ExamplePrefab", cancellationTokenSource.Token);
        cancellationTokenSource.Cancel();
        var prefab = await task;
    }

    // キャンセル付きのロードメソッド
    private async UniTask<T> LoadAssetAsync<T>(string key, CancellationToken cancellationToken = default)
    {
        var handle = Addressables.LoadAssetAsync<T>(key);
        handle.Completed += op =>
        {
            // ロード完了時にキャンセルされていたらリリースする
            if (cancellationToken.IsCancellationRequested)
                Addressables.Release(op);
        };
        return await handle;
    }
}

上記では、リソースのロードが完了した際にキャンセルがリクエストされていたらAddressableのリソース解放処理を呼ぶことでキャンセルを制御しています。
これをやらないと、メソッドの外側から処理をキャンセルした時にリソースが解放されず、メモリリークに繋がります。

このように少々面倒なAddressableのキャンセル処理ですが、UniTaskを使うと以下のように簡潔に書くことができます。

using System.Threading;
using Cysharp.Threading.Tasks;
using UnityEngine;
using UnityEngine.AddressableAssets;

public class Example : MonoBehaviour
{
    private async void Start()
    {
        var cancellationTokenSource = new CancellationTokenSource();
        var task = LoadAssetAsync<GameObject>("ExamplePrefab", cancellationTokenSource.Token);
        cancellationTokenSource.Cancel();
        var prefab = await task;
    }

    // キャンセル付きのロードメソッド
    private async UniTask<T> LoadAssetAsync<T>(string key, CancellationToken cancellationToken = default)
    {
        return await Addressables.LoadAssetAsync<T>(key)
            .ToUniTask(
                cancellationToken: cancellationToken, // ToUniTaskの引数にCancellationTokenを渡せる
                autoReleaseWhenCanceled: true // これをtrueにするとキャンセル時にAddressable解放処理を行う
            );
    }
}

このように、ToUniTask の autoReleaseWhenCanceled を true にすることでキャンセル時にリソース解放処理をやってくれるので、安全に取り扱うことができます。
以下はこの機能が実装された時の作者の方のX投稿です。

UniTask v2.5.2出しました！Addressableのawait拡張に、キャンセル時に同時にReleaseするautoReleaseWhenCancelledフラグが追加されました。またv2.5.1で入ったValueTaskとの変換はUnityのバグを回避するため UNITY_2022_3_OR_NEWER の場合のみ有効になりました。https://t.co/OCl6N3AEvz
— neuecc (@neuecc) 2024年1月25日

挙動を確認する

さてそれではこの autoReleaseWhenCanceled の挙動を実際に確認してみます。

Addressable では正確に解放処理が行われると、そのリソースの参照カウントがゼロになります。
これを確認するために以下の DiagnosticCallback で参照カウントをログ出力して確認します。

light11.hatenadiary.com

これを使った確認用のコードが以下になります。
説明はコメントに書いてあります。

using System;
using System.Threading;
using Cysharp.Threading.Tasks;
using UnityEngine;
using UnityEngine.AddressableAssets;
using UnityEngine.ResourceManagement;

public class Example : MonoBehaviour
{
    private void Awake()
    {
        Addressables.ResourceManager.RegisterDiagnosticCallback(DiagnosticCallback);
    }

    private void Start()
    {
        var cancellationTokenSource = new CancellationTokenSource();

        // ロード開始
        Addressables.LoadAssetAsync<GameObject>("ExamplePrefab")
            .ToUniTask(cancellationToken: cancellationTokenSource.Token, autoReleaseWhenCanceled: true);

        // ロードが終わる前にすぐキャンセル
        cancellationTokenSource.Cancel();
    }

    private void OnDestroy()
    {
        Addressables.ResourceManager.UnregisterDiagnosticCallback(DiagnosticCallback);
    }

    private static void DiagnosticCallback(ResourceManager.DiagnosticEventContext context)
    {
        // ExamplePrefab関連のログのみ表示
        var isTarget = context.Location != null && context.Location.PrimaryKey == "ExamplePrefab";
        if (!isTarget)
            return;

        // ReferenceCountの変化のみを表示
        if (context.Type != ResourceManager.DiagnosticEventType.AsyncOperationReferenceCount)
            return;

        Debug.Log(
            "Reference count changed" +
            $"{Environment.NewLine}Location: {context.Location}" +
            $"{Environment.NewLine}Reference count: {context.EventValue}");
    }
}