diff --git a/client/ExtraChat/Ext.cs b/client/ExtraChat/Ext.cs index f973c87..7744de9 100644 --- a/client/ExtraChat/Ext.cs +++ b/client/ExtraChat/Ext.cs @@ -1,4 +1,5 @@ -using System.Net.WebSockets; +using System.Buffers; +using System.Net.WebSockets; using ExtraChat.Protocol; using MessagePack; @@ -16,20 +17,30 @@ public static class Ext { } public static async Task ReceiveMessage(this ClientWebSocket client) { - var bytes = new List(2048); - var buffer = new ArraySegment(new byte[2048]); + const int maxSize = 1_048_576; + // rent a 1 MiB byte array from the shared runtime pool + var bytes = ArrayPool.Shared.Rent(maxSize); - WebSocketReceiveResult result; - do { - result = await client.ReceiveAsync(buffer, CancellationToken.None); - bytes.AddRange(buffer[..result.Count]); + try { + WebSocketReceiveResult result; + var i = 0; + do { + result = await client.ReceiveAsync(bytes[i..maxSize], CancellationToken.None); + i += result.Count; - // 1 MiB - if (bytes.Count > 1_048_576) { - throw new Exception(); - } - } while (!result.EndOfMessage); + // break if we've filled up the buffer, even if message isn't complete + if (i == maxSize) { + break; + } - return MessagePackSerializer.Deserialize(bytes.ToArray()); + if (i > maxSize) { + throw new Exception("read too many bytes for one message"); + } + } while (!result.EndOfMessage); + + return MessagePackSerializer.Deserialize(bytes); + } finally { + ArrayPool.Shared.Return(bytes); + } } }