refactor and fix

Fixing several issues with the disconnect flow.
Not removing the transport id until later. The server now only does this within OnClientDisconnectFromServer.
When disconnecting a client with a reason, the client is not completely disconnected until the end of the frame when the server has finished sending all pending messages.
Fixing some missed MessageDeliveryType replacements.
Removing the MockSkippingApproval as it is no longer needed with the fixes in place.

Author: NoelStephensUnity

com.unity.netcode.gameobjects/Runtime/Connection/NetworkConnectionManager.cs

@@ -594,9 +594,11 @@ internal void DisconnectEventHandler(ulong transportClientId)
             // Check to see if the client has already been removed from the table but
             // do not remove it just yet.
             var (clientId, isConnectedClient) = TransportIdToClientId(transportClientId);
-            if (!isConnectedClient)
+
+            // If the client is not registered and we are the server
+            if (!isConnectedClient && NetworkManager.IsServer)
             {
-                // If so, exit early
+                // Then exit early
                 return;
             }
 
@@ -631,15 +633,6 @@ internal void DisconnectEventHandler(ulong transportClientId)
             {
                 // We need to process the disconnection before notifying
                 OnClientDisconnectFromServer(clientId);
-
-                // Now notify the client has disconnected.
-                // (transport id cleanup is handled within)
-                InvokeOnClientDisconnectCallback(clientId);
-
-                if (LocalClient.IsHost)
-                {
-                    InvokeOnPeerDisconnectedCallback(clientId);
-                }
             }
             else
             {
@@ -806,12 +799,15 @@ private IEnumerator ApprovalTimeout(ulong clientId)
         /// </summary>
         internal void ApproveConnection(ref ConnectionRequestMessage connectionRequestMessage, ref NetworkContext context)
         {
+            if (ConnectionApprovalCallback == null)
+            {
+                return;
+            }
             // Note: Delegate creation allocates.
             // Note: ToArray() also allocates. :(
             var response = new NetworkManager.ConnectionApprovalResponse();
             ClientsToApprove[context.SenderId] = response;
-
-            ConnectionApprovalCallback(
+            ConnectionApprovalCallback?.Invoke(
                 new NetworkManager.ConnectionApprovalRequest
                 {
                     Payload = connectionRequestMessage.ConnectionData,
@@ -865,13 +861,6 @@ internal void ProcessPendingApprovals()
             }
         }
 
-        /// <summary>
-        /// Adding this because message hooks cannot happen fast enough under certain scenarios
-        /// where the message is sent and responded to before the hook is in place.
-        /// </summary>
-        internal bool MockSkippingApproval;
-
-
         /// <summary>
         /// Server Side: Handles the denial of a client who sent a connection request
         /// </summary>
@@ -886,12 +875,36 @@ private void HandleConnectionDisconnect(ulong ownerClientId, string reason = "")
                 {
                     Reason = reason
                 };
-                SendMessage(ref disconnectReason, NetworkDelivery.Reliable, ownerClientId);
+                SendMessage(ref disconnectReason, MessageDeliveryType<DisconnectReasonMessage>.DefaultDelivery, ownerClientId);
+                m_ClientsToDisconnect.Add(ownerClientId);
+                return;
             }
 
             DisconnectRemoteClient(ownerClientId);
         }
 
+        private List<ulong> m_ClientsToDisconnect = new List<ulong>();
+
+        internal void ProcessClientsToDisconnect()
+        {
+            if (m_ClientsToDisconnect.Count == 0)
+            {
+                return;
+            }
+            foreach (var clientId in m_ClientsToDisconnect)
+            {
+                try
+                {
+                    DisconnectRemoteClient(clientId);
+                }
+                catch (Exception ex)
+                {
+                    Debug.LogException(ex);
+                }
+            }
+            m_ClientsToDisconnect.Clear();
+        }
+
         /// <summary>
         /// Server Side: Handles the approval of a client
         /// </summary>
@@ -952,12 +965,6 @@ internal void HandleConnectionApproval(ulong ownerClientId, bool createPlayerObj
             // Server doesn't send itself the connection approved message
             if (ownerClientId != NetworkManager.ServerClientId)
             {
-                if (MockSkippingApproval)
-                {
-                    NetworkLog.LogInfo("Mocking server not responding with connection approved...");
-                    return;
-                }
-
                 SendConnectionApprovedMessage(ownerClientId);
 
                 // If scene management is disabled, then we are done and notify the local host-server the client is connected
@@ -1053,7 +1060,7 @@ private void SendConnectionApprovedMessage(ulong approvedClientId)
                 }
             }
 
-            SendMessage(ref message, NetworkDelivery.ReliableFragmentedSequenced, approvedClientId);
+            SendMessage(ref message, MessageDeliveryType<ConnectionApprovedMessage>.DefaultDelivery, approvedClientId);
 
             message.MessageVersions.Dispose();
             message.ConnectedClientIds.Dispose();
@@ -1124,13 +1131,9 @@ internal NetworkClient AddClient(ulong clientId)
                 return ConnectedClients[clientId];
             }
 
-            var networkClient = LocalClient;
-
             // If this is not the local client then create a new one
-            if (clientId != NetworkManager.LocalClientId)
-            {
-                networkClient = new NetworkClient();
-            }
+            var networkClient = clientId == NetworkManager.LocalClientId ? LocalClient : new NetworkClient();
+
             networkClient.SetRole(clientId == NetworkManager.ServerClientId, isClient: true, NetworkManager);
             networkClient.ClientId = clientId;
             if (!ConnectedClients.ContainsKey(clientId))
@@ -1237,6 +1240,14 @@ internal void OnClientDisconnectFromServer(ulong clientId)
             // clean up as everything that needs to be destroyed will be during shutdown.
             if (NetworkManager.ShutdownInProgress && clientId == NetworkManager.ServerClientId)
             {
+                // Now notify the client has disconnected.
+                // (transport id cleanup is handled within)
+                InvokeOnClientDisconnectCallback(clientId);
+
+                if (LocalClient.IsHost)
+                {
+                    InvokeOnPeerDisconnectedCallback(clientId);
+                }
                 return;
             }
 
@@ -1429,17 +1440,18 @@ internal void OnClientDisconnectFromServer(ulong clientId)
             var (transportId, idExists) = ClientIdToTransportId(clientId);
             if (idExists)
             {
-                NetworkManager.NetworkConfig.NetworkTransport.DisconnectRemoteClient(transportId);
-
-                InvokeOnClientDisconnectCallback(clientId);
-
-                if (LocalClient.IsHost)
+                // Clean up the transport to client (and vice versa) mappings
+                var (transportIdDisconnected, wasRemoved) = TransportIdCleanUp(transportId);
+                if (wasRemoved)
                 {
-                    InvokeOnPeerDisconnectedCallback(clientId);
-                }
+                    NetworkManager.NetworkConfig.NetworkTransport.DisconnectRemoteClient(transportId);
+                    InvokeOnClientDisconnectCallback(clientId);
 
-                // Clean up the transport to client (and vice versa) mappings
-                TransportIdCleanUp(transportId);
+                    if (LocalClient.IsHost)
+                    {
+                        InvokeOnPeerDisconnectedCallback(clientId);
+                    }
+                }
             }
 
             // Assure the client id is no longer in the pending clients list
@@ -1487,16 +1499,6 @@ internal void DisconnectClient(ulong clientId, string reason = null)
                 return;
             }
 
-            if (!string.IsNullOrEmpty(reason))
-            {
-                var disconnectReason = new DisconnectReasonMessage
-                {
-                    Reason = reason
-                };
-                SendMessage(ref disconnectReason, NetworkDelivery.Reliable, clientId);
-            }
-
-            Transport.ClosingRemoteConnection();
             var transportId = ClientIdToTransportId(clientId);
             if (transportId.Item2)
             {

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

@@ -348,7 +348,6 @@ public void NetworkUpdate(NetworkUpdateStage updateStage)
                         AnticipationSystem.SetupForUpdate();
                         MessageManager.ProcessIncomingMessageQueue();
 
-                        MessageManager.CleanupDisconnectedClients();
                         AnticipationSystem.ProcessReanticipation();
 #if COM_UNITY_MODULES_PHYSICS || COM_UNITY_MODULES_PHYSICS2D
                         foreach (var networkObjectEntry in NetworkTransformFixedUpdate)
@@ -480,6 +479,18 @@ public void NetworkUpdate(NetworkUpdateStage updateStage)
                         // This is "ok" to invoke when not processing messages since it is just cleaning up messages that never got handled within their timeout period.
                         DeferredMessageManager.CleanupStaleTriggers();
 
+                        if (IsServer)
+                        {
+                            // Process any pending clients that need to be disconnected.
+                            // This is typically a disconnect with reason scenario where
+                            // we want the disconnect reason message to be sent prior to
+                            // completely shutting down the endpoint.
+                            ConnectionManager.ProcessClientsToDisconnect();
+                        }
+
+                        // Clean up disconnected clients last
+                        MessageManager.CleanupDisconnectedClients();
+
                         if (m_ShuttingDown)
                         {
                             // Host-server will disconnect any connected clients prior to finalizing its shutdown

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

@@ -492,6 +492,11 @@ private void CleanupDisconnectedClient(ulong clientId)
 
         internal void CleanupDisconnectedClients()
         {
+            if (m_DisconnectedClients.Count == 0)
+            {
+                return;
+            }
+
             foreach (var clientId in m_DisconnectedClients)
             {
                 CleanupDisconnectedClient(clientId);