fix: unified transport and bootstrap memory leaks (#4219)

fix: release unified transport and bootstrap memory between sessions

Dispose each connection's send queue when the connection is removed or the
transport shuts down, release the safety handle created per received packet,
reuse the UnifiedNetcodeTransport component across restarts, and reset the
UnifiedBootstrap statics on load so a disabled domain reload does not keep
the previous NetworkManager alive.

Author: NoelStephensUnity

com.unity.netcode.gameobjects/Runtime/Components/Helpers/UnifiedBootstrap.cs

@@ -31,6 +31,20 @@ internal class UnifiedBootstrap : ClientServerBootstrap
         // creates its own bootstrap, so the worlds created for other NetworkManagers are registered again.
         private static readonly List<NetcodeWorld> k_CreatedWorlds = new List<NetcodeWorld>();
 
+#if UNITY_EDITOR
+        [RuntimeInitializeOnLoadMethod(RuntimeInitializeLoadType.SubsystemRegistration)]
+        private static void ResetStaticsOnLoad()
+        {
+            Instance = null;
+            OnInitialized = null;
+            Port = 7979;
+            CurrentNetworkManagerForInitialization = null;
+            LastCreatedWorld = null;
+            s_WorldCounter = 0;
+            k_CreatedWorlds.Clear();
+        }
+#endif
+
         private static void RegisterCreatedWorlds()
         {
             for (int i = k_CreatedWorlds.Count - 1; i >= 0; i--)

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

@@ -1266,7 +1266,12 @@ internal void Initialize(bool server)
             NetworkConfig.InitializePrefabs();
             if (NetworkConfig.Prefabs.HasGhostPrefabs)
             {
-                NetworkConfig.NetworkTransport = gameObject.AddComponent<UnifiedNetcodeTransport>();
+                // Reused so that restarting a session does not add another transport.
+                if (!TryGetComponent(out UnifiedNetcodeTransport unifiedTransport))
+                {
+                    unifiedTransport = gameObject.AddComponent<UnifiedNetcodeTransport>();
+                }
+                NetworkConfig.NetworkTransport = unifiedTransport;
             }
 #endif
 #if UNIFIED_NETCODE
@@ -1400,6 +1405,7 @@ internal void InitializeNetcodeWorld()
             TypeManager.Initialize();
             UnifiedBootstrap.CurrentNetworkManagerForInitialization = this;
             new UnifiedBootstrap().Initialize("Default World");
+            UnifiedBootstrap.CurrentNetworkManagerForInitialization = null;
 
             if (NetcodeWorld != null && NetcodeWorld.IsCreated)
             {

com.unity.netcode.gameobjects/Runtime/Transports/Unified/UnifiedNetcodeTransport.cs

@@ -215,6 +215,16 @@ private class ConnectionInfo
 
         private Dictionary<int, ConnectionInfo> m_Connections;
 
+        private bool RemoveConnection(int connectionId)
+        {
+            if (!m_Connections.Remove(connectionId, out ConnectionInfo connectionInfo))
+            {
+                return false;
+            }
+            connectionInfo.SendQueue.Dispose();
+            return true;
+        }
+
         internal void DispatchMessage(int connectionId, in FixedBytes1280 buffer)
         {
             // The connection is removed on disconnect, which can be ahead of its last messages.
@@ -223,7 +233,7 @@ internal void DispatchMessage(int connectionId, in FixedBytes1280 buffer)
                 return;
             }
 
-            using var arr = FixedBytes1280.ToNativeArray(buffer);
+            var arr = FixedBytes1280.ToNativeArray(buffer);
             var reader = new DataStreamReader(arr);
             if (connectionInfo.ReceiveQueue == null)
             {
@@ -233,6 +243,11 @@ internal void DispatchMessage(int connectionId, in FixedBytes1280 buffer)
             {
                 connectionInfo.ReceiveQueue.PushReader(reader);
             }
+#if ENABLE_UNITY_COLLECTIONS_CHECKS
+            // Disposing an Allocator.None array does not release its safety handle. The queue has copied the data.
+            var safety = NativeArrayUnsafeUtility.GetAtomicSafetyHandle(arr);
+            CollectionHelper.DisposeSafetyHandle(ref safety);
+#endif
 
             var message = connectionInfo.ReceiveQueue.PopMessage();
             while (message.Count != 0)
@@ -278,6 +293,7 @@ public override NetworkEvent PollEvent(out ulong clientId, out ArraySegment<byte
 
         private void OnClientConnectedToServer(Connection connection, NetcodeConnectionEvent connectionEvent)
         {
+            RemoveConnection(connectionEvent.Id.Value);
             m_Connections[connectionEvent.Id.Value] = new ConnectionInfo
             {
                 ReceiveQueue = null,
@@ -292,6 +308,7 @@ private void OnClientConnectedToServer(Connection connection, NetcodeConnectionE
 
         private void OnServerNewClientConnection(Connection connection, NetcodeConnectionEvent connectionEvent)
         {
+            RemoveConnection(connectionEvent.Id.Value);
             m_Connections[connectionEvent.Id.Value] = new ConnectionInfo
             {
                 ReceiveQueue = null,
@@ -375,13 +392,14 @@ private void OnClientDisconnectFromServer(Connection connection, NetcodeConnecti
                 GetDisconnectMessageFromNetworkStreamDisconnectReason(connectionEvent.DisconnectReason)
             );
             // Removed before notifying, so the shutdown this triggers does not notify again from DisconnectLocalClient.
-            m_Connections.Remove(connectionEvent.Id.Value);
+            RemoveConnection(connectionEvent.Id.Value);
             InvokeOnTransportEvent(NetworkEvent.Disconnect, (ulong)connectionEvent.Id.Value, default, m_RealTimeProvider.RealTimeSinceStartup);
         }
 
         private void OnServerClientDisconnected(Connection connection, NetcodeConnectionEvent connectionEvent)
         {
             InvokeOnTransportEvent(NetworkEvent.Disconnect, (ulong)connectionEvent.Id.Value, default, m_RealTimeProvider.RealTimeSinceStartup);
+            RemoveConnection(connectionEvent.Id.Value);
         }
 
         private void OnClientConnectionEvent(Connection connection, NetcodeConnectionEvent connectionEvent)
@@ -435,14 +453,18 @@ public override bool StartServer()
 
         public override void DisconnectRemoteClient(ulong clientId)
         {
-            m_NetworkManager.NetcodeWorld.DisconnectAClient(m_Connections[(int)clientId].Connection);
-            m_Connections.Remove((int)clientId);
+            // The connection is already removed when the client disconnected first.
+            if (m_Connections.TryGetValue((int)clientId, out ConnectionInfo connectionInfo))
+            {
+                m_NetworkManager.NetcodeWorld.DisconnectAClient(connectionInfo.Connection);
+                RemoveConnection((int)clientId);
+            }
         }
 
         public override void DisconnectLocalClient()
         {
             // Remove the connection 1st (the world might not be available)
-            var wasConnected = m_Connections.Remove((int)ServerClientId);
+            var wasConnected = RemoveConnection((int)ServerClientId);
 
             // TODO-FIX-REVIEW-ME:
             // This was causing errors to occur upon shutdown during an integration test.
@@ -495,6 +517,15 @@ public override void Shutdown()
                 netcodeWorld.OnConnectionEvent -= OnClientConnectionEvent;
                 netcodeWorld.OnConnectionEvent -= OnServerConnectionEvent;
             }
+
+            if (m_Connections != null)
+            {
+                foreach (var connectionInfo in m_Connections.Values)
+                {
+                    connectionInfo.SendQueue.Dispose();
+                }
+                m_Connections.Clear();
+            }
         }
     }
 }