diff --git a/Assets/Mirror/Core/Batching/Batcher.cs b/Assets/Mirror/Core/Batching/Batcher.cs index 5e30f08c8..41f467dbf 100644 --- a/Assets/Mirror/Core/Batching/Batcher.cs +++ b/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(); + } } } diff --git a/Assets/Mirror/Core/NetworkClient.cs b/Assets/Mirror/Core/NetworkClient.cs index 69b3f6e2a..c8f01205c 100644 --- a/Assets/Mirror/Core/NetworkClient.cs +++ b/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. diff --git a/Assets/Mirror/Core/NetworkConnection.cs b/Assets/Mirror/Core/NetworkConnection.cs index 6317eca0a..851beaea0 100644 --- a/Assets/Mirror/Core/NetworkConnection.cs +++ b/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})"; } } diff --git a/Assets/Mirror/Core/NetworkServer.cs b/Assets/Mirror/Core/NetworkServer.cs index 638799e85..23933b093 100644 --- a/Assets/Mirror/Core/NetworkServer.cs +++ b/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}"); diff --git a/Assets/Mirror/Tests/Editor/Batching/BatcherTests.cs b/Assets/Mirror/Tests/Editor/Batching/BatcherTests.cs index 4925c251e..f609bdc98 100644 --- a/Assets/Mirror/Tests/Editor/Batching/BatcherTests.cs +++ b/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(new byte[]{0x01}), TimeStamp); + batcher.AddMessage(new ArraySegment(new byte[]{0x02}), TimeStamp); + batcher.AddMessage(new ArraySegment(new byte[]{0x03}), TimeStamp); + Assert.That(NetworkWriterPool.Count, Is.LessThan(previousCount)); + + // clear + batcher.Clear(); + Assert.That(NetworkWriterPool.Count, Is.EqualTo(previousCount)); + } } }