fix: Fixed a few memory leak cases when shutting down NetworkManager during Incoming Message Queue processing (#1323) (#1466)

* fix: Fixed memory leak on shutdown

* fix: Fixed a few memory leak cases when shutting down NetworkManager during Incoming Message Queue processing

* standards check

* Fixed test not cleaning up quite correctly.

* -Reworked the way shutdown is handled so it can't be processed while handling incoming messages.
-Added tests for both ways to shut down to make sure they both work correctly
-Fixed an issue in the snapshot system that was causing the new tests to fail

* standards

* Fixed broken test, slight naming change to be more accurate.

Co-authored-by: Jaedyn Draper <284434+ShadauxCat@users.noreply.github.com>

Author: andrews-unity

com.unity.netcode.gameobjects/Runtime/Core/NetworkManager.cs

@@ -73,6 +73,9 @@ public NetworkPrefabHandler PrefabHandler
             }
         }
 
+        private bool m_ShuttingDown;
+        private bool m_StopProcessingMessages;
+
         private class NetworkManagerHooks : INetworkHooks
         {
             private NetworkManager m_NetworkManager;
@@ -116,7 +119,7 @@ public void OnAfterReceiveBatch(ulong senderId, int messageCount, int batchSizeI
 
             public bool OnVerifyCanSend(ulong destinationId, Type messageType, NetworkDelivery delivery)
             {
-                return true;
+                return !m_NetworkManager.m_StopProcessingMessages;
             }
 
             public bool OnVerifyCanReceive(ulong senderId, Type messageType)
@@ -134,7 +137,7 @@ public bool OnVerifyCanReceive(ulong senderId, Type messageType)
                     return false;
                 }
 
-                return true;
+                return !m_NetworkManager.m_StopProcessingMessages;
             }
         }
 
@@ -974,7 +977,7 @@ private void OnApplicationQuit()
         // Note that this gets also called manually by OnSceneUnloaded and OnApplicationQuit
         private void OnDestroy()
         {
-            Shutdown();
+            ShutdownInternal();
 
             UnityEngine.SceneManagement.SceneManager.sceneUnloaded -= OnSceneUnloaded;
 
@@ -994,13 +997,30 @@ private void DisconnectRemoteClient(ulong clientId)
         /// Globally shuts down the library.
         /// Disconnects clients if connected and stops server if running.
         /// </summary>
-        public void Shutdown()
+        /// <param name="discardMessageQueue">
+        /// If false, any messages that are currently in the incoming queue will be handled,
+        /// and any messages in the outgoing queue will be sent, before the shutdown is processed.
+        /// If true, NetworkManager will shut down immediately, and any unprocessed or unsent messages
+        /// will be discarded.
+        /// </param>
+        public void Shutdown(bool discardMessageQueue = false)
         {
             if (NetworkLog.CurrentLogLevel <= LogLevel.Developer)
             {
                 NetworkLog.LogInfo(nameof(Shutdown));
             }
 
+            m_ShuttingDown = true;
+            m_StopProcessingMessages = discardMessageQueue;
+        }
+
+        internal void ShutdownInternal()
+        {
+            if (NetworkLog.CurrentLogLevel <= LogLevel.Developer)
+            {
+                NetworkLog.LogInfo(nameof(ShutdownInternal));
+            }
+
             if (IsServer)
             {
                 // make sure all messages are flushed before transport disconnect clients
@@ -1077,6 +1097,7 @@ public void Shutdown()
 
             if (SpawnManager != null)
             {
+                SpawnManager.CleanupAllTriggers();
                 SpawnManager.DestroyNonSceneObjects();
                 SpawnManager.ServerResetShudownStateForSceneObjects();
 
@@ -1165,6 +1186,11 @@ private void OnNetworkPreUpdate()
                 return;
             }
 
+            if (m_ShuttingDown && m_StopProcessingMessages)
+            {
+                return;
+            }
+
             // Only update RTT here, server time is updated by time sync messages
             var reset = NetworkTimeSystem.Advance(Time.deltaTime);
             if (reset)
@@ -1181,9 +1207,18 @@ private void OnNetworkPreUpdate()
 
         private void OnNetworkPostLateUpdate()
         {
-            MessagingSystem.ProcessSendQueues();
-            NetworkMetrics.DispatchFrame();
+
+            if (!m_ShuttingDown || !m_StopProcessingMessages)
+            {
+                MessagingSystem.ProcessSendQueues();
+                NetworkMetrics.DispatchFrame();
+            }
             SpawnManager.CleanupStaleTriggers();
+
+            if (m_ShuttingDown)
+            {
+                ShutdownInternal();
+            }
         }
 
         /// <summary>

com.unity.netcode.gameobjects/Runtime/Core/SnapshotSystem.cs

@@ -465,30 +465,45 @@ internal void ReadIndex(in SnapshotDataMessage message)
             }
         }
 
-        internal void SpawnObject(SnapshotSpawnCommand spawnCommand)
+        internal void SpawnObject(SnapshotSpawnCommand spawnCommand, ulong srcClientId)
         {
             if (m_NetworkManager)
             {
-                var networkObject = m_NetworkManager.SpawnManager.CreateLocalNetworkObject(false,
-                    spawnCommand.GlobalObjectIdHash, spawnCommand.OwnerClientId, (spawnCommand.ParentNetworkId == spawnCommand.NetworkObjectId) ? spawnCommand.NetworkObjectId : spawnCommand.ParentNetworkId, spawnCommand.ObjectPosition,
-                    spawnCommand.ObjectRotation);
+                NetworkObject networkObject;
+                if (spawnCommand.ParentNetworkId == spawnCommand.NetworkObjectId)
+                {
+                    networkObject = m_NetworkManager.SpawnManager.CreateLocalNetworkObject(false,
+                            spawnCommand.GlobalObjectIdHash, spawnCommand.OwnerClientId, null, spawnCommand.ObjectPosition,
+                            spawnCommand.ObjectRotation);
+                }
+                else
+                {
+                    networkObject = m_NetworkManager.SpawnManager.CreateLocalNetworkObject(false,
+                            spawnCommand.GlobalObjectIdHash, spawnCommand.OwnerClientId, spawnCommand.ParentNetworkId, spawnCommand.ObjectPosition,
+                            spawnCommand.ObjectRotation);
+                }
+
                 m_NetworkManager.SpawnManager.SpawnNetworkObjectLocally(networkObject, spawnCommand.NetworkObjectId,
                     true, spawnCommand.IsPlayerObject, spawnCommand.OwnerClientId, false);
+                //todo: discuss with tools how to report shared bytes
+                m_NetworkManager.NetworkMetrics.TrackObjectSpawnReceived(srcClientId, networkObject, 8);
             }
             else
             {
                 MockSpawnObject(spawnCommand);
             }
         }
 
-        internal void DespawnObject(SnapshotDespawnCommand despawnCommand)
+        internal void DespawnObject(SnapshotDespawnCommand despawnCommand, ulong srcClientId)
         {
             if (m_NetworkManager)
             {
                 m_NetworkManager.SpawnManager.SpawnedObjects.TryGetValue(despawnCommand.NetworkObjectId,
                     out NetworkObject networkObject);
 
                 m_NetworkManager.SpawnManager.OnDespawnObject(networkObject, true);
+                //todo: discuss with tools how to report shared bytes
+                m_NetworkManager.NetworkMetrics.TrackObjectDestroyReceived(srcClientId, networkObject, 8);
             }
             else
             {
@@ -497,7 +512,7 @@ internal void DespawnObject(SnapshotDespawnCommand despawnCommand)
         }
 
 
-        internal void ReadSpawns(in SnapshotDataMessage message)
+        internal void ReadSpawns(in SnapshotDataMessage message, ulong srcClientId)
         {
             SnapshotSpawnCommand spawnCommand;
             SnapshotDespawnCommand despawnCommand;
@@ -516,7 +531,7 @@ internal void ReadSpawns(in SnapshotDataMessage message)
 
                 // Debug.Log($"[Spawn] {spawnCommand.NetworkObjectId} {spawnCommand.TickWritten}");
 
-                SpawnObject(spawnCommand);
+                SpawnObject(spawnCommand, srcClientId);
             }
             for (var i = 0; i < message.Despawns.Length; i++)
             {
@@ -532,7 +547,7 @@ internal void ReadSpawns(in SnapshotDataMessage message)
 
                 // Debug.Log($"[DeSpawn] {despawnCommand.NetworkObjectId} {despawnCommand.TickWritten}");
 
-                DespawnObject(despawnCommand);
+                DespawnObject(despawnCommand, srcClientId);
             }
         }
 
@@ -1077,7 +1092,7 @@ internal void HandleSnapshot(ulong clientId, in SnapshotDataMessage message)
             ReadBuffer(message);
             ReadIndex(message);
             ReadAcks(clientId, m_ClientData[clientId], message, GetConnectionRtt(clientId));
-            ReadSpawns(message);
+            ReadSpawns(message, clientId);
         }
 
         // todo --M1--

com.unity.netcode.gameobjects/Runtime/Messaging/MessagingSystem.cs

@@ -100,7 +100,7 @@ public MessagingSystem(IMessageSender messageSender, object owner, IMessageProvi
             }
         }
 
-        public void Dispose()
+        public unsafe void Dispose()
         {
             if (m_Disposed)
             {
@@ -113,6 +113,14 @@ public void Dispose()
             {
                 CleanupDisconnectedClient(kvp.Key);
             }
+
+            for (var queueIndex = 0; queueIndex < m_IncomingMessageQueue.Length; ++queueIndex)
+            {
+                // Avoid copies...
+                ref var item = ref m_IncomingMessageQueue.GetUnsafeList()->ElementAt(queueIndex);
+                item.Reader.Dispose();
+            }
+
             m_IncomingMessageQueue.Dispose();
             m_Disposed = true;
         }
@@ -253,11 +261,15 @@ public void HandleMessage(in MessageHeader header, FastBufferReader reader, ulon
 
         internal unsafe void ProcessIncomingMessageQueue()
         {
-            for (var i = 0; i < m_IncomingMessageQueue.Length; ++i)
+            for (var index = 0; index < m_IncomingMessageQueue.Length; ++index)
             {
                 // Avoid copies...
-                ref var item = ref m_IncomingMessageQueue.GetUnsafeList()->ElementAt(i);
+                ref var item = ref m_IncomingMessageQueue.GetUnsafeList()->ElementAt(index);
                 HandleMessage(item.Header, item.Reader, item.SenderId, item.Timestamp);
+                if (m_Disposed)
+                {
+                    return;
+                }
             }
 
             m_IncomingMessageQueue.Clear();
@@ -461,16 +473,16 @@ internal unsafe void ProcessSendQueues()
                     try
                     {
                         m_MessageSender.Send(clientId, queueItem.NetworkDelivery, queueItem.Writer);
+
+                        for (var hookIdx = 0; hookIdx < m_Hooks.Count; ++hookIdx)
+                        {
+                            m_Hooks[hookIdx].OnAfterSendBatch(clientId, queueItem.BatchHeader.BatchSize, queueItem.Writer.Length, queueItem.NetworkDelivery);
+                        }
                     }
                     finally
                     {
                         queueItem.Writer.Dispose();
                     }
-
-                    for (var hookIdx = 0; hookIdx < m_Hooks.Count; ++hookIdx)
-                    {
-                        m_Hooks[hookIdx].OnAfterSendBatch(clientId, queueItem.BatchHeader.BatchSize, queueItem.Writer.Length, queueItem.NetworkDelivery);
-                    }
                 }
                 sendQueueItem.Clear();
             }

com.unity.netcode.gameobjects/Runtime/Spawning/NetworkSpawnManager.cs

@@ -154,6 +154,24 @@ internal unsafe void CleanupStaleTriggers()
                 m_Triggers.Remove(staleKeys[i]);
             }
         }
+        /// <summary>
+        /// Cleans up any trigger that's existed for more than a second.
+        /// These triggers were probably for situations where a request was received after a despawn rather than before a spawn.
+        /// </summary>
+        internal void CleanupAllTriggers()
+        {
+            foreach (var kvp in m_Triggers)
+            {
+                foreach (var data in kvp.Value.TriggerData)
+                {
+                    data.Reader.Dispose();
+                }
+
+                kvp.Value.TriggerData.Dispose();
+            }
+
+            m_Triggers.Clear();
+        }
 
         internal void RemoveOwnership(NetworkObject networkObject)
         {

com.unity.netcode.gameobjects/Tests/Editor/StartStopTests.cs

@@ -45,7 +45,7 @@ public void TestStartupServerState()
         public void TestFlagShutdown()
         {
             m_NetworkManager.StartServer();
-            m_NetworkManager.Shutdown();
+            m_NetworkManager.ShutdownInternal();
 
             Assert.False(m_NetworkManager.IsServer);
             Assert.False(m_NetworkManager.IsClient);