test/Assets/PerfectWorld/Scripts/PoolingManager/PoolManager.cs

using System;
using System.Collections;
using System.Collections.Generic;
using System.Threading.Tasks;
using UnityEngine;

namespace BrewMonster.Scripts
{
    public sealed class PoolManager : MonoSingleton<PoolManager>
    {
        private readonly Dictionary<string, ObjectPool> _pools = new();
        private readonly Dictionary<GameObject, ObjectPool> _instanceToPool = new();

        private AddressableManager _addressableManager;
        private Transform _poolContainer;

        protected override void Initialize()
        {
            base.Initialize();

            _addressableManager = AddressableManager.Instance;
            _addressableManager.OnDispose += ReleaseAllPools;

            _poolContainer = new GameObject("Addressables Object Pools").transform;
            _poolContainer.SetParent(transform, false);
        }

        /// <summary>
        /// Spawns an Addressables prefab from its pool. The returned task completes after the prefab is loaded if needed.
        /// </summary>
        public async Task<GameObject> SpawnAsync(
            string addressableKey,
            Vector3 position,
            Quaternion rotation,
            float memoryReleaseTTL,
            float autoDespawnTime = 0f,
            Transform parent = null)
        {
            if (string.IsNullOrEmpty(addressableKey))
            {
                BMLogger.LogError("PoolManager: Cannot spawn with a null or empty Addressables key.");
                return null;
            }

            ObjectPool pool = GetOrCreatePool(addressableKey, memoryReleaseTTL);
            pool.UpdateMemoryReleaseTTL(memoryReleaseTTL);

            GameObject instance = await pool.SpawnAsync(position, rotation, parent);
            if (instance == null)
            {
                return null;
            }

            _instanceToPool[instance] = pool;

            if (autoDespawnTime > 0f)
            {
                StartCoroutine(AutoDespawnAfter(pool, instance, pool.GetSpawnVersion(instance), autoDespawnTime));
            }

            return instance;
        }

        /// <summary>
        /// Coroutine-friendly spawn API for callers that do not use async/await.
        /// </summary>
        public void Spawn(
            string addressableKey,
            Vector3 position,
            Quaternion rotation,
            float memoryReleaseTTL,
            float autoDespawnTime,
            Action<GameObject> onComplete,
            Transform parent = null)
        {
            StartCoroutine(SpawnRoutine(addressableKey, position, rotation, memoryReleaseTTL, autoDespawnTime, onComplete, parent));
        }

        public bool Despawn(string addressableKey, GameObject instance)
        {
            if (string.IsNullOrEmpty(addressableKey) || instance == null)
            {
                return false;
            }

            if (!_pools.TryGetValue(addressableKey, out ObjectPool pool))
            {
                return false;
            }

            bool despawned = pool.Despawn(instance);
            if (despawned)
            {
                _instanceToPool[instance] = pool;
            }

            return despawned;
        }

        public bool Despawn(GameObject instance)
        {
            if (instance == null || !_instanceToPool.TryGetValue(instance, out ObjectPool pool))
            {
                return false;
            }

            return pool.Despawn(instance);
        }

        public bool TryGetPoolCounts(string addressableKey, out int activeCount, out int idleCount)
        {
            activeCount = 0;
            idleCount = 0;

            if (!_pools.TryGetValue(addressableKey, out ObjectPool pool))
            {
                return false;
            }

            activeCount = pool.ActiveCount;
            idleCount = pool.IdleCount;
            return true;
        }

        public void ReleasePool(string addressableKey)
        {
            if (!_pools.TryGetValue(addressableKey, out ObjectPool pool))
            {
                return;
            }

            UnregisterPoolInstances(pool);
            pool.ReleaseNow();
            _pools.Remove(addressableKey);
        }

        public void ReleaseAllPools()
        {
            List<ObjectPool> pools = new(_pools.Values);
            for (int i = 0; i < pools.Count; i++)
            {
                pools[i].ReleaseNow();
            }

            _pools.Clear();
            _instanceToPool.Clear();
        }

        internal void RemovePool(ObjectPool pool)
        {
            if (pool == null)
            {
                return;
            }

            UnregisterPoolInstances(pool);
            _pools.Remove(pool.AddressableKey);
        }

        private IEnumerator SpawnRoutine(
            string addressableKey,
            Vector3 position,
            Quaternion rotation,
            float memoryReleaseTTL,
            float autoDespawnTime,
            Action<GameObject> onComplete,
            Transform parent)
        {
            Task<GameObject> spawnTask = SpawnAsync(addressableKey, position, rotation, memoryReleaseTTL, autoDespawnTime, parent);
            while (!spawnTask.IsCompleted)
            {
                yield return null;
            }

            if (spawnTask.Exception != null)
            {
                BMLogger.LogError($"PoolManager: Spawn failed for '{addressableKey}': {spawnTask.Exception}");
                onComplete?.Invoke(null);
                yield break;
            }

            onComplete?.Invoke(spawnTask.Result);
        }

        private IEnumerator AutoDespawnAfter(ObjectPool pool, GameObject instance, int spawnVersion, float autoDespawnTime)
        {
            yield return new WaitForSeconds(autoDespawnTime);

            if (pool != null && pool.IsActiveInstance(instance, spawnVersion))
            {
                Despawn(instance);
            }
        }

        private ObjectPool GetOrCreatePool(string addressableKey, float memoryReleaseTTL)
        {
            if (_pools.TryGetValue(addressableKey, out ObjectPool pool))
            {
                return pool;
            }

            Transform poolRoot = new GameObject(GetPoolRootName(addressableKey)).transform;
            poolRoot.SetParent(_poolContainer, false);

            pool = new ObjectPool(addressableKey, this, _addressableManager, poolRoot, memoryReleaseTTL);
            _pools[addressableKey] = pool;
            return pool;
        }

        private void UnregisterPoolInstances(ObjectPool pool)
        {
            foreach (GameObject instance in pool.GetKnownInstances())
            {
                if (instance != null)
                {
                    _instanceToPool.Remove(instance);
                }
            }
        }

        private static string GetPoolRootName(string addressableKey)
        {
            string name = addressableKey.Replace('\\', '/');
            int lastSlash = name.LastIndexOf('/');
            if (lastSlash >= 0 && lastSlash < name.Length - 1)
            {
                name = name.Substring(lastSlash + 1);
            }

            return $"Pool - {name}";
        }

        protected override void OnDestroy()
        {
            if (_addressableManager != null)
            {
                _addressableManager.OnDispose -= ReleaseAllPools;
            }

            ReleaseAllPools();
            base.OnDestroy();
        }
    }
}