fix: MTT-2991: Flush UnityTransport send queues on disconnect (#1847)

* fix: MTT-2991: Flush UnityTransport send queues on disconnect

* Add PR number to CHANGELOG entry

* Update CHANGELOG entry

* Fix existing disconnect tests

* Don't send more than necessary in overflow test

Author: simon-lemay-unity

com.unity.netcode.gameobjects/CHANGELOG.md

@@ -22,6 +22,7 @@ Additional documentation and release notes are available at [Multiplayer Documen
 - Removed `com.unity.modules.animation`, `com.unity.modules.physics` and `com.unity.modules.physics2d` dependencies from the package (#1812)
 
 ### Fixed
+- Fixed issue where `UnityTransport` send queues were not flushed when calling `DisconnectLocalClient` or `DisconnectRemoteClient`. (#1847)
 - Fixed NetworkBehaviour dependency verification check for an existing NetworkObject not searching from root parent transform relative GameObject. (#1841)
 - Fixed issue where entries were not being removed from the NetworkSpawnManager.OwnershipToObjectsTable. (#1838)
 - Fixed ClientRpcs would always send to all connected clients by default as opposed to only sending to the NetworkObject's Observers list by default. (#1836)

com.unity.netcode.gameobjects/Runtime/Transports/UTP/UnityTransport.cs

@@ -876,10 +876,23 @@ private void ClearSendQueuesForClientId(ulong clientId)
             }
         }
 
+        private void FlushSendQueuesForClientId(ulong clientId)
+        {
+            foreach (var kvp in m_SendQueue)
+            {
+                if (kvp.Key.ClientId == clientId)
+                {
+                    SendBatchedMessages(kvp.Key, kvp.Value);
+                }
+            }
+        }
+
         public override void DisconnectLocalClient()
         {
             if (m_State == State.Connected)
             {
+                FlushSendQueuesForClientId(m_ServerClientId);
+
                 if (m_Driver.Disconnect(ParseClientId(m_ServerClientId)) == 0)
                 {
                     m_State = State.Disconnected;
@@ -904,15 +917,16 @@ public override void DisconnectRemoteClient(ulong clientId)
 
             if (m_State == State.Listening)
             {
-                var connection = ParseClientId(clientId);
+                FlushSendQueuesForClientId(clientId);
+
+                m_ReliableReceiveQueues.Remove(clientId);
+                ClearSendQueuesForClientId(clientId);
 
+                var connection = ParseClientId(clientId);
                 if (m_Driver.GetConnectionState(connection) != NetworkConnection.State.Disconnected)
                 {
                     m_Driver.Disconnect(connection);
                 }
-
-                m_ReliableReceiveQueues.Remove(clientId);
-                ClearSendQueuesForClientId(clientId);
             }
         }
 

com.unity.netcode.gameobjects/Tests/Runtime/Transports/UnityTransportConnectionTests.cs

@@ -303,7 +303,16 @@ public IEnumerator ServerDisconnectWithDataInQueue()
 
             m_Server.DisconnectRemoteClient(m_ServerEvents[0].ClientID);
 
-            yield return WaitForNetworkEvent(NetworkEvent.Disconnect, m_ClientsEvents[0]);
+            yield return WaitForNetworkEvent(NetworkEvent.Data, m_ClientsEvents[0]);
+
+            if (m_ClientsEvents[0].Count >= 3)
+            {
+                Assert.AreEqual(NetworkEvent.Disconnect, m_ClientsEvents[0][2].Type);
+            }
+            else
+            {
+                yield return WaitForNetworkEvent(NetworkEvent.Disconnect, m_ClientsEvents[0]);
+            }
         }
 
         // Check client disconnection with data in send queue.
@@ -323,7 +332,16 @@ public IEnumerator ClientDisconnectWithDataInQueue()
 
             m_Clients[0].DisconnectLocalClient();
 
-            yield return WaitForNetworkEvent(NetworkEvent.Disconnect, m_ServerEvents);
+            yield return WaitForNetworkEvent(NetworkEvent.Data, m_ServerEvents);
+
+            if (m_ServerEvents.Count >= 3)
+            {
+                Assert.AreEqual(NetworkEvent.Disconnect, m_ServerEvents[2].Type);
+            }
+            else
+            {
+                yield return WaitForNetworkEvent(NetworkEvent.Disconnect, m_ServerEvents);
+            }
         }
 
         // Check that a server can disconnect a client after another client has disconnected.

com.unity.netcode.gameobjects/Tests/Runtime/Transports/UnityTransportTests.cs

@@ -283,7 +283,7 @@ public IEnumerator DisconnectOnReliableSendQueueOverflow()
 
             m_Server.Shutdown();
 
-            var numSends = (UnityTransport.InitialMaxSendQueueSize / 1024) + 1;
+            var numSends = (UnityTransport.InitialMaxSendQueueSize / 1024);
 
             for (int i = 0; i < numSends; i++)
             {
@@ -415,5 +415,47 @@ public IEnumerator SendQueuesFlushedOnShutdown([ValueSource("k_DeliveryParameter
 
             yield return null;
         }
+
+        [UnityTest]
+        public IEnumerator SendQueuesFlushedOnLocalClientDisconnect([ValueSource("k_DeliveryParameters")] NetworkDelivery delivery)
+        {
+            InitializeTransport(out m_Server, out m_ServerEvents);
+            InitializeTransport(out m_Client1, out m_Client1Events);
+
+            m_Server.StartServer();
+            m_Client1.StartClient();
+
+            yield return WaitForNetworkEvent(NetworkEvent.Connect, m_Client1Events);
+
+            var data = new ArraySegment<byte>(new byte[] { 42 });
+            m_Client1.Send(m_Client1.ServerClientId, data, delivery);
+
+            m_Client1.DisconnectLocalClient();
+
+            yield return WaitForNetworkEvent(NetworkEvent.Data, m_ServerEvents);
+
+            yield return null;
+        }
+
+        [UnityTest]
+        public IEnumerator SendQueuesFlushedOnRemoteClientDisconnect([ValueSource("k_DeliveryParameters")] NetworkDelivery delivery)
+        {
+            InitializeTransport(out m_Server, out m_ServerEvents);
+            InitializeTransport(out m_Client1, out m_Client1Events);
+
+            m_Server.StartServer();
+            m_Client1.StartClient();
+
+            yield return WaitForNetworkEvent(NetworkEvent.Connect, m_Client1Events);
+
+            var data = new ArraySegment<byte>(new byte[] { 42 });
+            m_Server.Send(m_Client1.ServerClientId, data, delivery);
+
+            m_Server.DisconnectRemoteClient(m_ServerEvents[0].ClientID);
+
+            yield return WaitForNetworkEvent(NetworkEvent.Data, m_Client1Events);
+
+            yield return null;
+        }
     }
 }