fix: Don't coalesce small unreliable messages into packets larger than the MTU (#2784)

* Modify test to use mixed small and large packets

* Add new test for BatchedSendQueue

* Limit unreliable coalescing of small messages to MTU

Author: simon-lemay-unity

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

@@ -198,43 +198,69 @@ public bool PushMessage(ArraySegment<byte> message)
         /// could lead to a corrupted queue.
         /// </remarks>
         /// <param name="writer">The <see cref="DataStreamWriter"/> to write to.</param>
+        /// <param name="softMaxBytes">
+        /// Maximum number of bytes to copy (0 means writer capacity). This is a soft limit only.
+        /// If a message is larger than that but fits in the writer, it will be written. In effect,
+        /// this parameter is the maximum size that small messages can be coalesced together.
+        /// </param>
         /// <returns>How many bytes were written to the writer.</returns>
-        public int FillWriterWithMessages(ref DataStreamWriter writer)
+        public int FillWriterWithMessages(ref DataStreamWriter writer, int softMaxBytes = 0)
         {
             if (!IsCreated || Length == 0)
             {
                 return 0;
             }
 
+            softMaxBytes = softMaxBytes == 0 ? writer.Capacity : Math.Min(softMaxBytes, writer.Capacity);
+
             unsafe
             {
                 var reader = new DataStreamReader(m_Data.AsArray());
-
-                var writerAvailable = writer.Capacity;
                 var readerOffset = HeadIndex;
 
-                while (readerOffset < TailIndex)
+                reader.SeekSet(readerOffset);
+                var messageLength = reader.ReadInt();
+                var bytesToWrite = messageLength + sizeof(int);
+
+                // Our behavior here depends on the size of the first message in the queue. If it's
+                // larger than the soft limit, then add only that message to the writer (we want
+                // large payloads to be fragmented on their own). Otherwise coalesce all small
+                // messages until we hit the soft limit (which presumably means they won't be
+                // fragmented, which is the desired behavior for smaller messages).
+
+                if (bytesToWrite > softMaxBytes && bytesToWrite <= writer.Capacity)
                 {
-                    reader.SeekSet(readerOffset);
-                    var messageLength = reader.ReadInt();
+                    writer.WriteInt(messageLength);
+                    WriteBytes(ref writer, (byte*)m_Data.GetUnsafePtr() + reader.GetBytesRead(), messageLength);
 
-                    if (writerAvailable < sizeof(int) + messageLength)
+                    return bytesToWrite;
+                }
+                else
+                {
+                    var bytesWritten = 0;
+
+                    while (readerOffset < TailIndex)
                     {
-                        break;
+                        reader.SeekSet(readerOffset);
+                        messageLength = reader.ReadInt();
+                        bytesToWrite = messageLength + sizeof(int);
+
+                        if (bytesWritten + bytesToWrite <= softMaxBytes)
+                        {
+                            writer.WriteInt(messageLength);
+                            WriteBytes(ref writer, (byte*)m_Data.GetUnsafePtr() + reader.GetBytesRead(), messageLength);
+
+                            readerOffset += bytesToWrite;
+                            bytesWritten += bytesToWrite;
+                        }
+                        else
+                        {
+                            break;
+                        }
                     }
-                    else
-                    {
-                        writer.WriteInt(messageLength);
-
-                        var messageOffset = reader.GetBytesRead();
-                        WriteBytes(ref writer, (byte*)m_Data.GetUnsafePtr() + messageOffset, messageLength);
 
-                        writerAvailable -= sizeof(int) + messageLength;
-                        readerOffset += sizeof(int) + messageLength;
-                    }
+                    return bytesWritten;
                 }
-
-                return writer.Capacity - writerAvailable;
             }
         }
 

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

@@ -753,7 +753,7 @@ public void Execute()
                     // in the stream (the send queue does that automatically) we are sure they'll be
                     // reassembled properly at the other end. This allows us to lift the limit of ~44KB
                     // on reliable payloads (because of the reliable window size).
-                    var written = pipeline == ReliablePipeline ? Queue.FillWriterWithBytes(ref writer, MTU) : Queue.FillWriterWithMessages(ref writer);
+                    var written = pipeline == ReliablePipeline ? Queue.FillWriterWithBytes(ref writer, MTU) : Queue.FillWriterWithMessages(ref writer, MTU);
 
                     result = Driver.EndSend(writer);
                     if (result == written)
@@ -1238,7 +1238,7 @@ public override void Initialize(NetworkManager networkManager = null)
             // We also increase the maximum resend timeout since the default one in UTP is very
             // aggressive (optimized for latency and low bandwidth). With NGO, it's too low and
             // we sometimes notice a lot of useless resends, especially if using Relay. (We can
-            // only do this with UTP 2.0 because 1.X doesn't support that parameter.)    
+            // only do this with UTP 2.0 because 1.X doesn't support that parameter.)
             m_NetworkSettings.WithReliableStageParameters(
                 windowSize: 64
 #if UTP_TRANSPORT_2_0_ABOVE

com.unity.netcode.gameobjects/Tests/Editor/Transports/BatchedSendQueueTests.cs

@@ -237,6 +237,35 @@ public void BatchedSendQueue_FillWriterWithMessages_PartialPushedMessages()
             AssertIsTestMessage(data);
         }
 
+        [Test]
+        public void BatchedSendQueue_FillWriterWithMessages_StopOnSoftMaxBytes()
+        {
+            var smallMessage = new ArraySegment<byte>(new byte[10]);
+            var largeMessage = new ArraySegment<byte>(new byte[3000]);
+
+            var smallMessageSize = smallMessage.Count + BatchedSendQueue.PerMessageOverhead;
+            var largeMessageSize = largeMessage.Count + BatchedSendQueue.PerMessageOverhead;
+
+            using var q = new BatchedSendQueue(k_TestQueueCapacity);
+            using var data = new NativeArray<byte>(largeMessageSize, Allocator.Temp);
+
+            q.PushMessage(smallMessage);
+            q.PushMessage(largeMessage);
+            q.PushMessage(smallMessage);
+
+            var writer = new DataStreamWriter(data);
+            Assert.AreEqual(smallMessageSize, q.FillWriterWithMessages(ref writer, 1000));
+            q.Consume(smallMessageSize);
+
+            writer = new DataStreamWriter(data);
+            Assert.AreEqual(largeMessageSize, q.FillWriterWithMessages(ref writer, 1000));
+            q.Consume(largeMessageSize);
+
+            writer = new DataStreamWriter(data);
+            Assert.AreEqual(smallMessageSize, q.FillWriterWithMessages(ref writer, 1000));
+            q.Consume(smallMessageSize);
+        }
+
         [Test]
         public void BatchedSendQueue_FillWriterWithBytes_NoopIfNoData()
         {

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

@@ -192,19 +192,31 @@ public IEnumerator MultipleSendsSingleFrame(
 
             yield return WaitForNetworkEvent(NetworkEvent.Connect, m_Client1Events);
 
-            var data1 = new ArraySegment<byte>(new byte[] { 11 });
-            m_Client1.Send(m_Client1.ServerClientId, data1, delivery);
+            var data1 = new byte[10];
+            data1[0] = 11;
+            m_Client1.Send(m_Client1.ServerClientId, new ArraySegment<byte>(data1), delivery);
 
-            var data2 = new ArraySegment<byte>(new byte[] { 22 });
-            m_Client1.Send(m_Client1.ServerClientId, data2, delivery);
+            var data2 = new byte[3000];
+            data2[0] = 22;
+            m_Client1.Send(m_Client1.ServerClientId, new ArraySegment<byte>(data2), delivery);
+
+            var data3 = new byte[10];
+            data3[0] = 33;
+            m_Client1.Send(m_Client1.ServerClientId, new ArraySegment<byte>(data3), delivery);
 
             yield return WaitForNetworkEvent(NetworkEvent.Data, m_ServerEvents);
 
-            Assert.AreEqual(3, m_ServerEvents.Count);
-            Assert.AreEqual(NetworkEvent.Data, m_ServerEvents[2].Type);
+            Assert.AreEqual(4, m_ServerEvents.Count);
+            Assert.AreEqual(NetworkEvent.Data, m_ServerEvents[3].Type);
 
             Assert.AreEqual(11, m_ServerEvents[1].Data.First());
+            Assert.AreEqual(10, m_ServerEvents[1].Data.Count);
+
             Assert.AreEqual(22, m_ServerEvents[2].Data.First());
+            Assert.AreEqual(3000, m_ServerEvents[2].Data.Count);
+
+            Assert.AreEqual(33, m_ServerEvents[3].Data.First());
+            Assert.AreEqual(10, m_ServerEvents[3].Data.Count);
 
             yield return null;
         }