feat: use a shared memory pool for reading

This commit is contained in:
Anna 2022-07-31 19:47:23 -04:00
parent 1c5a3870c2
commit 1dbc60ba7c
1 changed files with 24 additions and 13 deletions

View File

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