fix: Error when a client is disconnected as a result of attempting to send to it [MTT-4607] (#2495)

* fix: Error when a client is disconnected as a result of attempting to send to it

* Changelog

Author: ShadauxCat

com.unity.netcode.gameobjects/CHANGELOG.md

@@ -8,6 +8,12 @@ Additional documentation and release notes are available at [Multiplayer Documen
 
 ## [Unreleased]
 
+### Fixed
+
+- Fixed some errors that could occur if a connection is lost and the loss is detected when attempting to write to the socket. (#2495)
+
+## [1.4.0]
+
 ### Added
 
 - Added a way to access the GlobalObjectIdHash via PrefabIdHash for use in the Connection Approval Callback. (#2437)

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

@@ -58,6 +58,8 @@ public SendQueueItem(NetworkDelivery delivery, int writerSize, Allocator writerA
         private Dictionary<Type, uint> m_MessageTypes = new Dictionary<Type, uint>();
         private Dictionary<ulong, NativeList<SendQueueItem>> m_SendQueues = new Dictionary<ulong, NativeList<SendQueueItem>>();
 
+        private HashSet<ulong> m_DisconnectedClients = new HashSet<ulong>();
+
         // This is m_PerClientMessageVersion[clientId][messageType] = version
         private Dictionary<ulong, Dictionary<Type, int>> m_PerClientMessageVersions = new Dictionary<ulong, Dictionary<Type, int>>();
         private Dictionary<uint, Type> m_MessagesByHash = new Dictionary<uint, Type>();
@@ -155,8 +157,9 @@ public unsafe void Dispose()
             // from the queue.
             foreach (var kvp in m_SendQueues)
             {
-                CleanupDisconnectedClient(kvp.Key);
+                ClientDisconnected(kvp.Key);
             }
+            CleanupDisconnectedClients();
 
             for (var queueIndex = 0; queueIndex < m_IncomingMessageQueue.Length; ++queueIndex)
             {
@@ -464,41 +467,37 @@ internal void ClientConnected(ulong clientId)
         }
 
         internal void ClientDisconnected(ulong clientId)
+        {
+            m_DisconnectedClients.Add(clientId);
+        }
+
+        private void CleanupDisconnectedClient(ulong clientId)
         {
             if (!m_SendQueues.ContainsKey(clientId))
             {
                 return;
             }
-            CleanupDisconnectedClient(clientId);
-            m_SendQueues.Remove(clientId);
-        }
 
-        private void CleanupDisconnectedClient(ulong clientId)
-        {
             var queue = m_SendQueues[clientId];
             for (var i = 0; i < queue.Length; ++i)
             {
                 queue.ElementAt(i).Writer.Dispose();
             }
 
             queue.Dispose();
+            m_SendQueues.Remove(clientId);
+
+            m_PerClientMessageVersions.Remove(clientId);
         }
 
         internal void CleanupDisconnectedClients()
         {
-            var removeList = new NativeList<ulong>(Allocator.Temp);
-            foreach (var clientId in m_PerClientMessageVersions.Keys)
+            foreach (var clientId in m_DisconnectedClients)
             {
-                if (!m_SendQueues.ContainsKey(clientId))
-                {
-                    removeList.Add(clientId);
-                }
+                CleanupDisconnectedClient(clientId);
             }
 
-            foreach (var clientId in removeList)
-            {
-                m_PerClientMessageVersions.Remove(clientId);
-            }
+            m_DisconnectedClients.Clear();
         }
 
         public static int CreateMessageAndGetVersion<T>() where T : INetworkMessage, new()
@@ -637,6 +636,10 @@ internal unsafe int SendPreSerializedMessage<TMessageType>(in FastBufferWriter t
 
             for (var i = 0; i < clientIds.Count; ++i)
             {
+                if (m_DisconnectedClients.Contains(clientIds[i]))
+                {
+                    continue;
+                }
                 var messageVersion = 0;
                 // Special case because this is the message that carries the version info - thus the version info isn't
                 // populated yet when we get this. The first part of this message always has to be the version data
@@ -788,6 +791,14 @@ internal unsafe void ProcessSendQueues()
                 for (var i = 0; i < sendQueueItem.Length; ++i)
                 {
                     ref var queueItem = ref sendQueueItem.ElementAt(i);
+                    // this is checked at every iteration because
+                    // 1) each writer needs to be disposed, so we have to do the full loop regardless, and
+                    // 2) the call to m_MessageSender.Send() may result in calling ClientDisconnected(), so the result of this check may change partway through iteration
+                    if (m_DisconnectedClients.Contains(clientId))
+                    {
+                        queueItem.Writer.Dispose();
+                        continue;
+                    }
                     if (queueItem.BatchHeader.BatchCount == 0)
                     {
                         queueItem.Writer.Dispose();

com.unity.netcode.gameobjects/Tests/Editor/Messaging/DisconnectOnSendTests.cs

@@ -0,0 +1,90 @@
+using System.Collections.Generic;
+using NUnit.Framework;
+
+namespace Unity.Netcode.EditorTests
+{
+    public class DisconnectOnSendTests
+    {
+        private struct TestMessage : INetworkMessage, INetworkSerializeByMemcpy
+        {
+            public void Serialize(FastBufferWriter writer, int targetVersion)
+            {
+            }
+
+            public bool Deserialize(FastBufferReader reader, ref NetworkContext context, int receivedMessageVersion)
+            {
+                return true;
+            }
+
+            public void Handle(ref NetworkContext context)
+            {
+            }
+
+            public int Version => 0;
+        }
+
+        private class DisconnectOnSendMessageSender : IMessageSender
+        {
+            public MessagingSystem MessagingSystem;
+
+            public void Send(ulong clientId, NetworkDelivery delivery, FastBufferWriter batchData)
+            {
+                MessagingSystem.ClientDisconnected(clientId);
+            }
+        }
+        private class TestMessageProvider : IMessageProvider
+        {
+            // Keep track of what we sent
+            private List<MessagingSystem.MessageWithHandler> m_MessageList = new List<MessagingSystem.MessageWithHandler>{
+                new MessagingSystem.MessageWithHandler
+                {
+                    MessageType = typeof(TestMessage),
+                    Handler = MessagingSystem.ReceiveMessage<TestMessage>,
+                    GetVersion = MessagingSystem.CreateMessageAndGetVersion<TestMessage>
+                }
+            };
+
+            public List<MessagingSystem.MessageWithHandler> GetMessages()
+            {
+                return m_MessageList;
+            }
+        }
+
+        private TestMessageProvider m_TestMessageProvider;
+        private DisconnectOnSendMessageSender m_MessageSender;
+        private MessagingSystem m_MessagingSystem;
+        private ulong[] m_Clients = { 0 };
+
+        [SetUp]
+        public void SetUp()
+        {
+            m_MessageSender = new DisconnectOnSendMessageSender();
+            m_TestMessageProvider = new TestMessageProvider();
+            m_MessagingSystem = new MessagingSystem(m_MessageSender, this, m_TestMessageProvider);
+            m_MessageSender.MessagingSystem = m_MessagingSystem;
+            m_MessagingSystem.ClientConnected(0);
+            m_MessagingSystem.SetVersion(0, XXHash.Hash32(typeof(TestMessage).FullName), 0);
+        }
+
+        [TearDown]
+        public void TearDown()
+        {
+            m_MessagingSystem.Dispose();
+        }
+
+        private TestMessage GetMessage()
+        {
+            return new TestMessage();
+        }
+
+        [Test]
+        public void WhenDisconnectIsCalledDuringSend_NoErrorsOccur()
+        {
+            var message = GetMessage();
+            m_MessagingSystem.SendMessage(ref message, NetworkDelivery.Reliable, m_Clients);
+
+            // This is where an exception would be thrown and logged.
+            m_MessagingSystem.ProcessSendQueues();
+        }
+    }
+}