fix(NetworkConnection): batches are now properly returned to NetworkWriterPool before destroying the connection (#3815)

Assets/Mirror/Core/Batching/Batcher.cs

@@ -146,5 +146,22 @@ public bool GetBatch(NetworkWriter writer)
             // nothing was written
             return false;
         }
+
+        // return all batches to the pool for cleanup
+        public void Clear()
+        {
+            // return batch in progress
+            if (batch != null)
+            {
+                NetworkWriterPool.Return(batch);
+                batch = null;
+            }
+
+            // return all queued batches
+            foreach (NetworkWriterPooled queued in batches)
+                NetworkWriterPool.Return(queued);
+
+            batches.Clear();
+        }
     }
 }

Assets/Mirror/Core/NetworkClient.cs

@@ -446,6 +446,7 @@ internal static void OnTransportDisconnected()
             // now that everything was handled, clear the connection.
             // previously this was done in Disconnect() already, but we still
             // need it for the above OnDisconnectedEvent.
+            connection?.Cleanup();
             connection = null;
 
             // transport handlers are only added when connecting.

Assets/Mirror/Core/NetworkConnection.cs

@@ -202,6 +202,18 @@ internal virtual void Update()
         // then later the transport events will do the clean up.
         public abstract void Disconnect();
 
+        // cleanup is called before the connection is removed.
+        // return any batches' pooled writers before the connection disappears.
+        // otherwise if a connection disappears before flushing, writers would
+        // never be returned to the pool.
+        public virtual void Cleanup()
+        {
+            foreach (Batcher batcher in batches.Values)
+            {
+                batcher.Clear();
+            }
+        }
+
         public override string ToString() => $"connection({connectionId})";
     }
 }

Assets/Mirror/Core/NetworkServer.cs

@@ -824,6 +824,7 @@ internal static void OnTransportDisconnected(int connectionId)
             // Debug.Log($"Server disconnect client:{connectionId}");
             if (connections.TryGetValue(connectionId, out NetworkConnectionToClient conn))
             {
+                conn.Cleanup();
                 RemoveConnection(connectionId);
                 // Debug.Log($"Server lost client:{connectionId}");
 

Assets/Mirror/Tests/Editor/Batching/BatcherTests.cs

@@ -284,5 +284,21 @@ public void MessageSerializationMismatch()
             reader = new NetworkReader(message);
             Assert.That(reader.ReadByte(), Is.EqualTo(3));
         }
+
+        [Test]
+        public void ClearReturnsToPool()
+        {
+            int previousCount = NetworkWriterPool.Count;
+
+            // add a few messages
+            batcher.AddMessage(new ArraySegment<byte>(new byte[]{0x01}), TimeStamp);
+            batcher.AddMessage(new ArraySegment<byte>(new byte[]{0x02}), TimeStamp);
+            batcher.AddMessage(new ArraySegment<byte>(new byte[]{0x03}), TimeStamp);
+            Assert.That(NetworkWriterPool.Count, Is.LessThan(previousCount));
+
+            // clear
+            batcher.Clear();
+            Assert.That(NetworkWriterPool.Count, Is.EqualTo(previousCount));
+        }
     }
 }