UniRx.ToolKit.ObjectPool相当の機能を別の実装で代替する

何がしたいのか

UniRxの開発が終了してしまったため、今後はUniRxではなくR3を使うことが推奨されます。
しかしUniRxからR3への移行は機械的にコードを置き換えていくだけでは駄目で、UniRxしか存在しない機能を他の何かで代替する必要があります。

今回は「UniRx.ToolKit.ObjectPool」とだいたい同じ挙動をするオブジェクトプールをR3またはUniTaskで再現してみた、という内容です。

使用ライブラリ

• R3
• UniTask
• uPools

ObjectPoolの再現

さっそくですが、同期版のオブジェクトプールUniRx.ToolKit.ObjectPoolをuPoolsをベースに再現してみました。

R3実装

まずは「UniRxをそのままR3に置き換えて使いたい」という需要のために、UniRx版のインタフェースをだいたいそのまま踏襲してR3で実装したものが次です。

using System;
using System.Threading;
using UnityEngine;
using uPools;
using R3;

namespace UniRxLikePool
{
    /// <summary>
    /// UniRx.ToolKit.ObjectPoolをR3のみで再現したもの
    /// </summary>
    public abstract class R3ObjectPool<T> : ObjectPoolBase<T> where T : Component
    {
        private readonly CancellationTokenSource _cts = new();

        protected override void OnDestroy(T instance)
        {
            if (instance == null) return;
            var go = instance.gameObject;
            if (go == null) return;

            UnityEngine.Object.Destroy(go);
        }

        protected override void OnRent(T instance)
        {
            instance.gameObject.SetActive(true);
        }

        protected override void OnReturn(T instance)
        {
            instance.gameObject.SetActive(false);
        }

        /// <summary>
        /// プールの事前拡張（フレーム分散）
        /// </summary>
        /// <param name="preloadCount">プール内の目標オブジェクト数</param>
        /// <param name="threshold">1フレームあたりに生成する最大数</param>
        public Observable<Unit> PreloadAsync(
            int preloadCount,
            int threshold)
        {
            if (IsDisposed) return Observable.Empty<Unit>();
            if (Count >= preloadCount) return Observable.Empty<Unit>();
            if (threshold <= 0) threshold = 1;

            return Observable.Defer(() => Observable.EveryUpdate(_cts.Token)
                .SelectMany(_ => Observable.Repeat(Unit.Default, threshold))
                .TakeWhile(_ => Count < preloadCount)
                .ForEachAsync(_ =>
                {
                    var instance = CreateInstance();
                    Return(instance);
                })
                .ToObservable());
        }

        /// <summary>
        /// プールの縮小を行う
        /// </summary>
        /// <param name="instanceCountRatio">縮小する割合（0.0ですべて消す ~ 1.0で何も消さない）</param>
        /// <param name="minSize">プールに残す最小数</param>
        /// <param name="callOnRent">オブジェクト削除時にOnRentを呼び出すか</param>
        public void Shrink(float instanceCountRatio, int minSize, bool callOnRent = false)
        {
            if (IsDisposed) return;

            if (instanceCountRatio <= 0) instanceCountRatio = 0;
            if (instanceCountRatio >= 1.0f) instanceCountRatio = 1.0f;

            var size = (int)(Count * instanceCountRatio);
            size = Math.Max(minSize, size);

            while (Count > size)
            {
                if (!stack.TryPop(out var obj)) break;
                if (callOnRent)
                {
                    OnRent(obj);
                }

                OnDestroy(obj);
            }
        }

        /// <summary>
        /// 一定の時間間隔でプールの縮小を行う
        /// </summary>
        /// <param name="checkInterval">縮小の時間間隔</param>
        /// <param name="instanceCountRatio">縮小する割合（0.0ですべて消す ~ 1.0で何も消さない）</param>
        /// <param name="minSize">プールに残す最小数</param>
        /// <param name="callOnRent">オブジェクト削除時にOnRentを呼び出すか</param>
        public IDisposable StartShrinkTimer(TimeSpan checkInterval,
            float instanceCountRatio,
            int minSize,
            bool callOnRent = false)
        {
            return Observable.Interval(checkInterval)
                .TakeWhile(_ => !IsDisposed)
                .Subscribe(_ => Shrink(instanceCountRatio, minSize, callOnRent));
        }

        public new void Dispose()
        {
            base.Dispose();
            _cts.Cancel();
            _cts.Dispose();
        }
    }
}

使用例

using System;
using UnityEngine;
using Object = UnityEngine.Object;
using R3;

namespace Sandboxes
{
    public class Sandbox : MonoBehaviour
    {
        [SerializeField] private Transform _prefab;
        private IDisposable _poolShrinkDisposable;

        private void Start()
        {
            var pool = new SampleR3ObjectPool(_prefab);

            // プールの事前拡張
            // Subscribe() を忘れずに
            pool.PreloadAsync(100, 10).Subscribe();

            // 1秒おきに10%ずつプール内のオブジェクトを破棄する
            // 最低1個は残す
            _poolShrinkDisposable = pool.StartShrinkTimer(
                checkInterval: TimeSpan.FromSeconds(1),
                instanceCountRatio: 0.9f,
                minSize: 1);

            // 貸出
            var instance = pool.Rent();

            // 返却
            pool.Return(instance);
        }
        
        private void OnDestroy()
        {
            _poolShrinkDisposable?.Dispose();
        }
    }

    /// <summary>
    /// サンプル実装
    /// </summary>
    public class SampleR3ObjectPool : R3ObjectPool<Transform>
    {
        private readonly GameObject _root;
        private readonly Transform _prefab;

        public SampleR3ObjectPool(Transform prefab)
        {
            _root = new GameObject("R3_Root");
            _prefab = prefab;
        }

        protected override Transform CreateInstance()
        {
            var n = Object.Instantiate(_prefab, _root.transform);
            return n.transform;
        }

        protected override void OnDestroy(Transform instance)
        {
            // 削除時に実行される
            base.OnDestroy(instance);
        }

        protected override void OnRent(Transform instance)
        {
            // 貸出前に実行される

            base.OnRent(instance);
        }

        protected override void OnReturn(Transform instance)
        {
            // 返却前に実行される
            base.OnReturn(instance);
        }
    }
}

OnRentやOnReturn、OnDestroyと一部のイベントコールバックの名前が変わっていますが、UniRx版とだいたい同じ挙動をします。
（ただ非同期にObservableを使ってたりと現代風ではないので、可能なら次に紹介するUniTask版の方がよさそう）

UniTask版

こちらはUniRx.ToolKit.ObjectPoolの機能をだいたい引き継ぎつつ、async/awaitで扱えるようにUniTaskで実装したものです。
置き換える余力があるならR3版よりはこっちを使ったほうがよいかと。

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

namespace UniRxLikePool
{
    public abstract class UniTaskObjectPool<T> : ObjectPoolBase<T> where T : Component
    {
        private readonly CancellationTokenSource _cts = new();

        protected override void OnDestroy(T instance)
        {
            if (instance == null) return;
            var go = instance.gameObject;
            if (go == null) return;

            UnityEngine.Object.Destroy(go);
        }

        protected override void OnRent(T instance)
        {
            instance.gameObject.SetActive(true);
        }

        protected override void OnReturn(T instance)
        {
            instance.gameObject.SetActive(false);
        }

        /// <summary>
        /// プールの事前拡張（フレーム分散）
        /// </summary>
        /// <param name="preloadCount">新しく生成する個数</param>
        /// <param name="threshold">1フレームあたりに生成する最大数</param>
        /// <param name="cancellationToken">キャンセルに使う</param>
        public async UniTask PreloadAsync(
            int preloadCount,
            int threshold,
            CancellationToken cancellationToken = default)
        {
            if (IsDisposed) return;
            if (Count >= preloadCount) return;
            if (threshold <= 0) threshold = 1;

            var linkedCts =
                CancellationTokenSource.CreateLinkedTokenSource(_cts.Token, cancellationToken);


            for (var i = 0; Count < preloadCount; i++)
            {
                var instance = CreateInstance();
                Return(instance);
                
                if ((i + 1) % threshold == 0)
                {
                    await UniTask.Yield(cancellationToken: linkedCts.Token);
                }
            }
        }

        /// <summary>
        /// プールの縮小を行う
        /// </summary>
        /// <param name="instanceCountRatio">縮小する割合（0.0ですべて消す ~ 1.0で何も消さない）</param>
        /// <param name="minSize">プールに残す最小数</param>
        /// <param name="callOnRent">オブジェクト削除時にOnRentを呼び出すか</param>
        public void Shrink(float instanceCountRatio, int minSize, bool callOnRent = false)
        {
            if (IsDisposed) return;

            if (instanceCountRatio <= 0) instanceCountRatio = 0;
            if (instanceCountRatio >= 1.0f) instanceCountRatio = 1.0f;

            var size = (int)(Count * instanceCountRatio);
            size = Math.Max(minSize, size);

            while (Count > size)
            {
                if (!stack.TryPop(out var obj)) break;
                if (callOnRent)
                {
                    OnRent(obj);
                }

                OnDestroy(obj);
            }
        }

        /// <summary>
        /// 一定の時間間隔でプールの縮小を行う
        /// </summary>
        /// <param name="checkInterval">縮小の時間間隔</param>
        /// <param name="instanceCountRatio">縮小する割合（0.0ですべて消す ~ 1.0で何も消さない）</param>
        /// <param name="minSize">プールに残す最小数</param>
        /// <param name="callOnRent">オブジェクト削除時にOnRentを呼び出すか</param>
        /// <param name="cancellationToken">キャンセルする場合に使用</param>
        public async UniTask StartShrinkTimerAsync(TimeSpan checkInterval,
            float instanceCountRatio,
            int minSize,
            bool callOnRent = false,
            CancellationToken cancellationToken = default)
        {
            if (IsDisposed) return;

            var linkedCts =
                CancellationTokenSource.CreateLinkedTokenSource(_cts.Token, cancellationToken);
            while (!linkedCts.IsCancellationRequested)
            {
                await UniTask.Delay(checkInterval, cancellationToken: linkedCts.Token);
                if (IsDisposed) return;
                Shrink(instanceCountRatio, minSize, callOnRent);
            }
        }

        public new void Dispose()
        {
            base.Dispose();
            _cts.Cancel();
            _cts.Dispose();
        }
    }
}

使用例

using System;
using Cysharp.Threading.Tasks;
using UnityEngine;
using Object = UnityEngine.Object;
using Samples.UniRxLikePool;

namespace Sandboxes
{
    public class Sandbox : MonoBehaviour
    {
        [SerializeField] private Transform _prefab;

        private async UniTaskVoid Start()
        {
            var pool = new SampleUniTaskObjectPool(_prefab);

            // プールの事前拡張
            // awaitで拡張完了を待つ（Forgetしてもいい）
            await pool.PreloadAsync(100, 10);

            // 1秒おきに10%ずつプール内のオブジェクトを破棄する
            // 最低1個は残す
            pool.StartShrinkTimerAsync(
                    checkInterval: TimeSpan.FromSeconds(1),
                    instanceCountRatio: 0.9f,
                    minSize: 1,
                    callOnRent: false,
                    destroyCancellationToken)
                .Forget();

            // 貸出
            var instance = pool.Rent();

            // 返却
            pool.Return(instance);
        }
    }

    /// <summary>
    /// サンプル実装
    /// </summary>
    public class SampleUniTaskObjectPool : UniTaskObjectPool<Transform>
    {
        private readonly GameObject _root;
        private readonly Transform _prefab;

        public SampleUniTaskObjectPool(Transform prefab)
        {
            _root = new GameObject("R3_Root");
            _prefab = prefab;
        }

        protected override Transform CreateInstance()
        {
            var n = Object.Instantiate(_prefab, _root.transform);
            return n.transform;
        }

        protected override void OnDestroy(Transform instance)
        {
            // 削除時に実行される
            base.OnDestroy(instance);
        }

        protected override void OnRent(Transform instance)
        {
            // 貸出前に実行される

            base.OnRent(instance);
        }

        protected override void OnReturn(Transform instance)
        {
            // 返却前に実行される
            base.OnReturn(instance);
        }
    }
}

AsyncObjectPoolの再現

非同期版のUniRx.ToolKit.AsyncObjectPoolも再現してみました。
こちらは最初からasync/awaitで使うべきだと感じたのでUniTask版のみです。

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

namespace UniRxLikePool
{
    public abstract class UniTaskAsyncObjectPool<T> : AsyncObjectPoolBase<T> where T : Component
    {
        private readonly CancellationTokenSource _cts = new();

        protected override void OnDestroy(T instance)
        {
            if (instance == null) return;
            var go = instance.gameObject;
            if (go == null) return;

            UnityEngine.Object.Destroy(go);
        }

        protected override void OnRent(T instance)
        {
            instance.gameObject.SetActive(true);
        }

        protected override void OnReturn(T instance)
        {
            instance.gameObject.SetActive(false);
        }

        /// <summary>
        /// プールの事前拡張（フレーム分散）
        /// </summary>
        /// <param name="preloadCount">プール内の目標オブジェクト数</param>
        /// <param name="threshold">1フレームあたりに生成する最大数</param>
        /// <param name="ct">CancellationToken</param>
        public async UniTask PreloadAsync(
            int preloadCount,
            int threshold,
            CancellationToken ct = default)
        {
            if (IsDisposed) return;
            if (Count >= preloadCount) return;
            if (threshold <= 0) threshold = 1;

            var token = CancellationTokenSource.CreateLinkedTokenSource(_cts.Token, ct).Token;

            var loaders = new UniTask[threshold];
            while (Count < preloadCount && !token.IsCancellationRequested)
            {
                var requireCount = preloadCount - Count;
                if (requireCount <= 0) break;

                var createCount = Math.Min(requireCount, threshold);

                for (var i = 0; i < threshold; i++)
                {
                    if (i < createCount)
                    {
                        loaders[i] = CreateInstanceAsync(token).ContinueWith(Return);
                    }
                    else
                    {
                        loaders[i] = UniTask.CompletedTask;
                    }
                }

                await UniTask.WhenAll(loaders);
            }
        }


        /// <summary>
        /// プールの縮小を行う
        /// </summary>
        /// <param name="instanceCountRatio">縮小する割合（0.0ですべて消す ~ 1.0で何も消さない）</param>
        /// <param name="minSize">プールに残す最小数</param>
        /// <param name="callOnRent">オブジェクト削除時にOnRentを呼び出すか</param>
        public void Shrink(float instanceCountRatio, int minSize, bool callOnRent = false)
        {
            if (IsDisposed) return;

            if (instanceCountRatio <= 0) instanceCountRatio = 0;
            if (instanceCountRatio >= 1.0f) instanceCountRatio = 1.0f;

            var size = (int)(Count * instanceCountRatio);
            size = Math.Max(minSize, size);

            while (Count > size)
            {
                if (!stack.TryPop(out var obj)) break;
                if (callOnRent)
                {
                    OnRent(obj);
                }

                OnDestroy(obj);
            }
        }

        /// <summary>
        /// 一定の時間間隔でプールの縮小を行う
        /// </summary>
        /// <param name="checkInterval">縮小の時間間隔</param>
        /// <param name="instanceCountRatio">縮小する割合（0.0ですべて消す ~ 1.0で何も消さない）</param>
        /// <param name="minSize">プールに残す最小数</param>
        /// <param name="callOnRent">オブジェクト削除時にOnRentを呼び出すか</param>
        /// <param name="cancellationToken">キャンセルする場合に使用</param>
        public async UniTask StartShrinkTimerAsync(TimeSpan checkInterval,
            float instanceCountRatio,
            int minSize,
            bool callOnRent = false,
            CancellationToken cancellationToken = default)
        {
            if (IsDisposed) return;

            var linkedCts =
                CancellationTokenSource.CreateLinkedTokenSource(_cts.Token, cancellationToken);
            while (!linkedCts.IsCancellationRequested)
            {
                await UniTask.Delay(checkInterval, cancellationToken: linkedCts.Token);
                if (IsDisposed) return;
                Shrink(instanceCountRatio, minSize, callOnRent);
            }
        }

        public new void Dispose()
        {
            base.Dispose();
            _cts.Cancel();
            _cts.Dispose();
        }
    }
}

まとめ

UniRxが提供するオブジェクトプール自体はそんな複雑なことはしてないので、自前で実装するか他のオブジェクトプールに乗り換えるのも全然アリだとは思います。
今回は「UniRxからR3に移行する時に、オブジェクトプールを最小工数で置き換えたい」という需要に応えるための実装紹介でした。