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

* 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.

Author: ShadauxCat

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;
             }
         }
 
@@ -976,7 +979,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;
 
@@ -996,13 +999,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
@@ -1079,6 +1099,7 @@ public void Shutdown()
 
             if (SpawnManager != null)
             {
+                SpawnManager.CleanupAllTriggers();
                 SpawnManager.DestroyNonSceneObjects();
                 SpawnManager.ServerResetShudownStateForSceneObjects();
 
@@ -1167,6 +1188,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)
@@ -1183,9 +1209,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

@@ -417,7 +417,54 @@ internal void ReadIndex(in SnapshotDataMessage message)
             }
         }
 
-        internal void ReadSpawns(in SnapshotDataMessage message)
+        internal void SpawnObject(SnapshotSpawnCommand spawnCommand, ulong srcClientId)
+        {
+            if (m_NetworkManager)
+            {
+                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, 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
+            {
+                MockDespawnObject(despawnCommand);
+            }
+        }
+
+
+        internal void ReadSpawns(in SnapshotDataMessage message, ulong srcClientId)
         {
             SnapshotSpawnCommand spawnCommand;
             SnapshotDespawnCommand despawnCommand;

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);

testproject/Assets/Tests/Runtime/NoMemoryLeakOnNetworkManagerShutdownTest.cs

@@ -0,0 +1,125 @@
+using System.Collections;
+using Unity.Netcode;
+using Unity.Netcode.RuntimeTests;
+using NUnit.Framework;
+using TestProject.RuntimeTests.Support;
+using UnityEngine;
+using UnityEngine.TestTools;
+
+namespace TestProject.RuntimeTests
+{
+    public class NoMemoryLeakOnNetworkManagerShutdownTest
+    {
+        private GameObject m_Prefab;
+
+        [SetUp]
+        public void Setup()
+        {
+            ShutdownDuringOnNetworkSpawnBehaviour.SpawnCount = 0;
+            ShutdownDuringOnNetworkSpawnBehaviour.ClientRpcsCalled = 0;
+            ShutdownDuringOnNetworkSpawnBehaviour.ServerRpcsCalled = 0;
+            ShutdownDuringOnNetworkSpawnBehaviour.ShutdownImmediately = false;
+        }
+
+        [UnityTearDown]
+        public IEnumerator Teardown()
+        {
+            MultiInstanceHelpers.Destroy();
+            // Shutdown and clean up both of our NetworkManager instances
+            if (m_Prefab)
+            {
+                Object.Destroy(m_Prefab);
+            }
+            yield break;
+        }
+
+        public IEnumerator RunTest()
+        {
+            // Must be 1 for this test.
+            const int numClients = 1;
+            Assert.True(MultiInstanceHelpers.Create(numClients, out NetworkManager server, out NetworkManager[] clients));
+            m_Prefab = new GameObject("Object");
+            m_Prefab.AddComponent<ShutdownDuringOnNetworkSpawnBehaviour>();
+            var networkObject = m_Prefab.AddComponent<NetworkObject>();
+
+            // Make it a prefab
+            MultiInstanceHelpers.MakeNetworkObjectTestPrefab(networkObject);
+
+            var validNetworkPrefab = new NetworkPrefab();
+            validNetworkPrefab.Prefab = m_Prefab;
+            server.NetworkConfig.NetworkPrefabs.Add(validNetworkPrefab);
+            foreach (var client in clients)
+            {
+                client.NetworkConfig.NetworkPrefabs.Add(validNetworkPrefab);
+            }
+
+            // Start the instances
+            if (!MultiInstanceHelpers.Start(false, server, clients))
+            {
+                Debug.LogError("Failed to start instances");
+                Assert.Fail("Failed to start instances");
+            }
+
+            // [Client-Side] Wait for a connection to the server
+            yield return MultiInstanceHelpers.Run(MultiInstanceHelpers.WaitForClientsConnected(clients, null, 512));
+
+            // [Host-Side] Check to make sure all clients are connected
+            yield return MultiInstanceHelpers.Run(
+                MultiInstanceHelpers.WaitForClientsConnectedToServer(server, clients.Length, null, 512));
+
+            var serverObject = Object.Instantiate(m_Prefab, Vector3.zero, Quaternion.identity);
+            NetworkObject serverNetworkObject = serverObject.GetComponent<NetworkObject>();
+            serverNetworkObject.NetworkManagerOwner = server;
+            serverNetworkObject.Spawn();
+
+            // Wait until all objects have spawned.
+            const int maxFrames = 240;
+            var doubleCheckTime = Time.realtimeSinceStartup + 5.0f;
+            while (ShutdownDuringOnNetworkSpawnBehaviour.SpawnCount < clients.Length + 1)
+            {
+                if (Time.frameCount > maxFrames)
+                {
+                    // This is here in the event a platform is running at a higher
+                    // frame rate than expected
+                    if (doubleCheckTime < Time.realtimeSinceStartup)
+                    {
+                        Assert.Fail("Did not successfully call all expected client RPCs");
+                        break;
+                    }
+                }
+                var nextFrameNumber = Time.frameCount + 1;
+                yield return new WaitUntil(() => Time.frameCount >= nextFrameNumber);
+            }
+
+            Assert.AreEqual(ShutdownDuringOnNetworkSpawnBehaviour.SpawnCount, clients.Length + 1);
+            // Extra frames to catch Native Container memory leak log message
+            var lastFrameNumber = Time.frameCount + 10;
+            yield return new WaitUntil(() => Time.frameCount >= lastFrameNumber);
+            Object.Destroy(serverObject);
+        }
+
+        [UnityTest]
+        public IEnumerator WhenNetworkManagerShutsDownWhileTriggeredMessagesArePending_MemoryDoesNotLeak()
+        {
+            yield return RunTest();
+            LogAssert.NoUnexpectedReceived();
+        }
+
+        [UnityTest]
+        public IEnumerator WhenNetworkManagerShutsDownWhileTriggeredMessagesArePending_MessagesAreStillProcessed()
+        {
+            yield return RunTest();
+            Assert.AreEqual(1, ShutdownDuringOnNetworkSpawnBehaviour.ClientRpcsCalled);
+            Assert.AreEqual(1, ShutdownDuringOnNetworkSpawnBehaviour.ServerRpcsCalled);
+        }
+
+        [UnityTest]
+        public IEnumerator WhenNetworkManagerShutsDownImmediatelyWhileTriggeredMessagesArePending_MessagesAreNotProcessed()
+        {
+            ShutdownDuringOnNetworkSpawnBehaviour.ShutdownImmediately = true;
+            yield return RunTest();
+            Assert.AreEqual(0, ShutdownDuringOnNetworkSpawnBehaviour.ClientRpcsCalled);
+            Assert.AreEqual(0, ShutdownDuringOnNetworkSpawnBehaviour.ServerRpcsCalled);
+        }
+    }
+}