refactor: replacing custom pool implementation with ObjectPool [MTT-6263]

Assets/Scripts/Infrastructure/NetworkObjectPool.cs

@@ -1,53 +1,62 @@
 using System;
 using System.Collections.Generic;
-using System.Runtime.CompilerServices;
 using Unity.Netcode;
 using UnityEngine;
 using UnityEngine.Assertions;
+using UnityEngine.Pool;
 
 namespace Unity.BossRoom.Infrastructure
 {
     /// <summary>
-    /// Object Pool for networked objects, used for controlling how objects are spawned by Netcode. Netcode by default will allocate new memory when spawning new
-    /// objects. With this Networked Pool, we're using custom spawning to reuse objects.
+    /// Object Pool for networked objects, used for controlling how objects are spawned by Netcode. Netcode by default
+    /// will allocate new memory when spawning new objects. With this Networked Pool, we're using the ObjectPool to
+    /// reuse objects.
     /// Boss Room uses this for projectiles. In theory it should use this for imps too, but we wanted to show vanilla spawning vs pooled spawning.
-    /// Hooks to NetworkManager's prefab handler to intercept object spawning and do custom actions
+    /// Hooks to NetworkManager's prefab handler to intercept object spawning and do custom actions.
     /// </summary>
     public class NetworkObjectPool : NetworkBehaviour
     {
-        private static NetworkObjectPool _instance;
-
-        public static NetworkObjectPool Singleton { get { return _instance; } }
+        public static NetworkObjectPool Singleton { get; private set; }
 
         [SerializeField]
         List<PoolConfigObject> PooledPrefabsList;
 
-        HashSet<GameObject> prefabs = new HashSet<GameObject>();
-
-        Dictionary<GameObject, Queue<NetworkObject>> pooledObjects = new Dictionary<GameObject, Queue<NetworkObject>>();
+        HashSet<GameObject> m_Prefabs = new HashSet<GameObject>();
 
-        private bool m_HasInitialized = false;
+        Dictionary<GameObject, ObjectPool<NetworkObject>> m_PooledObjects = new Dictionary<GameObject, ObjectPool<NetworkObject>>();
 
         public void Awake()
         {
-            if (_instance != null && _instance != this)
+            if (Singleton != null && Singleton != this)
             {
-                Destroy(this.gameObject);
+                Destroy(gameObject);
             }
             else
             {
-                _instance = this;
+                Singleton = this;
             }
         }
 
         public override void OnNetworkSpawn()
         {
-            InitializePool();
+            // Registers all objects in PooledPrefabsList to the cache.
+            foreach (var configObject in PooledPrefabsList)
+            {
+                RegisterPrefabInternal(configObject.Prefab, configObject.PrewarmCount);
+            }
         }
 
         public override void OnNetworkDespawn()
         {
-            ClearPool();
+            // Unregisters all objects in PooledPrefabsList from the cache.
+            foreach (var prefab in m_Prefabs)
+            {
+                // Unregister Netcode Spawn handlers
+                NetworkManager.Singleton.PrefabHandler.RemoveHandler(prefab);
+                m_PooledObjects[prefab].Clear();
+            }
+            m_PooledObjects.Clear();
+            m_Prefabs.Clear();
         }
 
         public void OnValidate()
@@ -65,16 +74,12 @@ public void OnValidate()
         /// <summary>
         /// Gets an instance of the given prefab from the pool. The prefab must be registered to the pool.
         /// </summary>
-        /// <param name="prefab"></param>
-        /// <returns></returns>
-        public NetworkObject GetNetworkObject(GameObject prefab)
-        {
-            return GetNetworkObjectInternal(prefab, Vector3.zero, Quaternion.identity);
-        }
-
-        /// <summary>
-        /// Gets an instance of the given prefab from the pool. The prefab must be registered to the pool.
-        /// </summary>
+        /// <remarks>
+        /// To spawn a NetworkObject from one of the pools, this must be called on the server, then the instance
+        /// returned from it must be spawned on the server. This method will then also be called on the client by the
+        /// PooledPrefabInstanceHandler when the client receives a spawn message for a prefab that has been registered
+        /// here.
+        /// </remarks>
         /// <param name="prefab"></param>
         /// <param name="position">The position to spawn the object at.</param>
         /// <param name="rotation">The rotation to spawn the object with.</param>
@@ -89,107 +94,71 @@ public NetworkObject GetNetworkObject(GameObject prefab, Vector3 position, Quate
         /// </summary>
         public void ReturnNetworkObject(NetworkObject networkObject, GameObject prefab)
         {
-            var go = networkObject.gameObject;
-            go.SetActive(false);
-            pooledObjects[prefab].Enqueue(networkObject);
+            m_PooledObjects[prefab].Release(networkObject);
         }
 
         /// <summary>
-        /// Adds a prefab to the list of spawnable prefabs.
+        /// Builds up the cache for a prefab.
         /// </summary>
-        /// <param name="prefab">The prefab to add.</param>
-        /// <param name="prewarmCount"></param>
-        public void AddPrefab(GameObject prefab, int prewarmCount = 0)
+        void RegisterPrefabInternal(GameObject prefab, int prewarmCount)
         {
-            var networkObject = prefab.GetComponent<NetworkObject>();
+            NetworkObject CreateFunc()
+            {
+                return Instantiate(prefab).GetComponent<NetworkObject>();
+            }
 
-            Assert.IsNotNull(networkObject, $"{nameof(prefab)} must have {nameof(networkObject)} component.");
-            Assert.IsFalse(prefabs.Contains(prefab), $"Prefab {prefab.name} is already registered in the pool.");
+            void ActionOnGet(NetworkObject networkObject)
+            {
+                networkObject.gameObject.SetActive(true);
+            }
 
-            RegisterPrefabInternal(prefab, prewarmCount);
-        }
+            void ActionOnRelease(NetworkObject networkObject)
+            {
+                networkObject.gameObject.SetActive(false);
+            }
 
-        /// <summary>
-        /// Builds up the cache for a prefab.
-        /// </summary>
-        private void RegisterPrefabInternal(GameObject prefab, int prewarmCount)
-        {
-            prefabs.Add(prefab);
+            void ActionOnDestroy(NetworkObject networkObject)
+            {
+                Destroy(networkObject.gameObject);
+            }
 
-            var prefabQueue = new Queue<NetworkObject>();
-            pooledObjects[prefab] = prefabQueue;
-            for (int i = 0; i < prewarmCount; i++)
+            m_Prefabs.Add(prefab);
+
+            // Create the pool
+            m_PooledObjects[prefab] = new ObjectPool<NetworkObject>(CreateFunc, ActionOnGet, ActionOnRelease, ActionOnDestroy, defaultCapacity: prewarmCount);
+
+            // Populate the pool
+            var prewarmNetworkObjects = new List<NetworkObject>();
+            for (var i = 0; i < prewarmCount; i++)
+            {
+                prewarmNetworkObjects.Add(m_PooledObjects[prefab].Get());
+            }
+            foreach (var networkObject in prewarmNetworkObjects)
             {
-                var go = CreateInstance(prefab);
-                ReturnNetworkObject(go.GetComponent<NetworkObject>(), prefab);
+                m_PooledObjects[prefab].Release(networkObject);
             }
 
             // Register Netcode Spawn handlers
             NetworkManager.Singleton.PrefabHandler.AddHandler(prefab, new PooledPrefabInstanceHandler(prefab, this));
         }
 
-        [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private GameObject CreateInstance(GameObject prefab)
-        {
-            return Instantiate(prefab);
-        }
-
         /// <summary>
         /// This matches the signature of <see cref="NetworkSpawnManager.SpawnHandlerDelegate"/>
         /// </summary>
         /// <param name="prefab"></param>
         /// <param name="position"></param>
         /// <param name="rotation"></param>
         /// <returns></returns>
-        private NetworkObject GetNetworkObjectInternal(GameObject prefab, Vector3 position, Quaternion rotation)
+        NetworkObject GetNetworkObjectInternal(GameObject prefab, Vector3 position, Quaternion rotation)
         {
-            var queue = pooledObjects[prefab];
-
-            NetworkObject networkObject;
-            if (queue.Count > 0)
-            {
-                networkObject = queue.Dequeue();
-            }
-            else
-            {
-                networkObject = CreateInstance(prefab).GetComponent<NetworkObject>();
-            }
-
-            // Here we must reverse the logic in ReturnNetworkObject.
-            var go = networkObject.gameObject;
-            go.SetActive(true);
+            var networkObject = m_PooledObjects[prefab].Get();
 
-            go.transform.position = position;
-            go.transform.rotation = rotation;
+            var noTransform = networkObject.transform;
+            noTransform.position = position;
+            noTransform.rotation = rotation;
 
             return networkObject;
         }
-
-        /// <summary>
-        /// Registers all objects in <see cref="PooledPrefabsList"/> to the cache.
-        /// </summary>
-        public void InitializePool()
-        {
-            if (m_HasInitialized) return;
-            foreach (var configObject in PooledPrefabsList)
-            {
-                RegisterPrefabInternal(configObject.Prefab, configObject.PrewarmCount);
-            }
-            m_HasInitialized = true;
-        }
-
-        /// <summary>
-        /// Unregisters all objects in <see cref="PooledPrefabsList"/> from the cache.
-        /// </summary>
-        public void ClearPool()
-        {
-            foreach (var prefab in prefabs)
-            {
-                // Unregister Netcode Spawn handlers
-                NetworkManager.Singleton.PrefabHandler.RemoveHandler(prefab);
-            }
-            pooledObjects.Clear();
-        }
     }
 
     [Serializable]
@@ -212,8 +181,7 @@ public PooledPrefabInstanceHandler(GameObject prefab, NetworkObjectPool pool)
 
         NetworkObject INetworkPrefabInstanceHandler.Instantiate(ulong ownerClientId, Vector3 position, Quaternion rotation)
         {
-            var netObject = m_Pool.GetNetworkObject(m_Prefab, position, rotation);
-            return netObject;
+            return m_Pool.GetNetworkObject(m_Prefab, position, rotation);
         }
 
         void INetworkPrefabInstanceHandler.Destroy(NetworkObject networkObject)

CHANGELOG.md

@@ -11,6 +11,7 @@ Additional documentation and release notes are available at [Multiplayer Documen
 ### Changed
 * Replaced our polling for lobby updates with a subscription to the new Websocket based LobbyEvents (#805). This saves up a significant amount of bandwidth usage to and from the service, since updates are infrequent in this game. Now clients and hosts only use up bandwidth on the Lobby service when it is needed. With polling, we used to send a GET request per client once every 2s. The responses were between ~550 bytes and 900 bytes, so if we suppose an average of 725 bytes and 100 000 concurrent users (CCU), this amounted to around 725B * 30 calls per minute * 100 000 CCU = 2.175 GB per minute. Scaling this to a month would get us 93.96 TB per month. In our case, since the only changes to the lobbies happen when a user connects or disconnects, most of that data was not necessary and can be saved to reduce bandwidth usage. Since the cost of using the Lobby service depends on bandwidth usage, this would also save money on an actual game.
 * Simplified reconnection flow by offloading responsibility to ConnectionMethod (#804). Now the ClientReconnectingState uses the ConnectionMethod it is configured with to handle setting up reconnection (i.e. reconnecting to the Lobby before trying to reconnect to the Relay server if it is using Relay and Lobby). It can now also fail early and stop retrying if the lobby doesn't exist anymore.
+* Replaced our custom pool implementation using queues with ObjectPool (#824)
 
 ### Cleanup
 * Clarified a TODO comment inside ClientCharacter, detailing how anticipation should only be executed on owning client players (#786)