From 1dbc60ba7c1191c17de859aafd42b70db541e45d Mon Sep 17 00:00:00 2001
From: Anna <git@anna.lgbt>
Date: Sun, 31 Jul 2022 19:47:23 -0400
Subject: [PATCH] feat: use a shared memory pool for reading

---
 client/ExtraChat/Ext.cs | 37 ++++++++++++++++++++++++-------------
 1 file changed, 24 insertions(+), 13 deletions(-)

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<ResponseContainer> ReceiveMessage(this ClientWebSocket client) {
-        var bytes = new List<byte>(2048);
-        var buffer = new ArraySegment<byte>(new byte[2048]);
+        const int maxSize = 1_048_576;
+        // rent a 1 MiB byte array from the shared runtime pool
+        var bytes = ArrayPool<byte>.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<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);
+        }
     }
 }