byte[]

.github/copilot-instructions.md

@@ -0,0 +1,4 @@
+# Copilot Instructions
+
+## Project Guidelines
+- All new byte[xx] calls should use using var data = ArrayPool<byte>.Shared.Lease(xx); instead of directly allocating byte arrays

src/WireMock.Net.Minimal/WebSockets/WebSocketMessageBuilder.cs

@@ -58,6 +58,12 @@ internal class WebSocketMessageBuilder : IWebSocketMessageBuilder
         return this;
     }
 
+    public IWebSocketMessageBuilder Close()
+    {
+        ShouldClose = true;
+        return this;
+    }
+
     internal enum MessageType
     {
         Text,

src/WireMock.Net.Shared/Extensions/ArrayPoolExtensions.cs

@@ -0,0 +1,20 @@
+// Copyright © WireMock.Net
+
+namespace System.Buffers;
+
+internal sealed class Lease<T>(ArrayPool<T> pool, int length) : IDisposable
+{
+    public T[] Rented { get; } = pool.Rent(length);
+
+    public static implicit operator T[](Lease<T> lease) => lease.Rented;
+
+    public void Dispose()
+    {
+        pool.Return(Rented, true);
+    }
+}
+
+internal static class ArrayPoolExtensions
+{
+    public static Lease<T> Lease<T>(this ArrayPool<T> source, int length) => new(source, length);
+}

src/WireMock.Net.Shared/WebSockets/IWebSocketMessageBuilder.cs

@@ -49,4 +49,10 @@ public interface IWebSocketMessageBuilder
     /// </summary>
     [PublicAPI]
     IWebSocketMessageBuilder AndClose();
+
+    /// <summary>
+    /// Close the WebSocket connection.
+    /// </summary>
+    [PublicAPI]
+    IWebSocketMessageBuilder Close();
 }

test/WireMock.Net.Tests/WebSockets/ClientWebSocketExtensions.cs

@@ -1,5 +1,6 @@
 // Copyright © WireMock.Net
 
+using System.Buffers;
 using System.Net.WebSockets;
 using System.Text;
 
@@ -12,9 +13,9 @@ internal static class ClientWebSocketExtensions
         return client.SendAsync(new ArraySegment<byte>(Encoding.UTF8.GetBytes(text)), WebSocketMessageType.Text, endOfMessage, cancellationToken);
     }
 
-    internal static async Task<string> ReceiveAsTextAsync(this ClientWebSocket client, CancellationToken cancellationToken = default)
+    internal static async Task<string> ReceiveAsTextAsync(this ClientWebSocket client, int bufferSize = 1024, CancellationToken cancellationToken = default)
     {
-        var receiveBuffer = new byte[1024];
+        using var receiveBuffer = ArrayPool<byte>.Shared.Lease(1024);
         var result = await client.ReceiveAsync(new ArraySegment<byte>(receiveBuffer), cancellationToken);
 
         if (result.MessageType != WebSocketMessageType.Text)
@@ -32,7 +33,7 @@ internal static class ClientWebSocketExtensions
 
     internal static async Task<byte[]> ReceiveAsBytesAsync(this ClientWebSocket client, int bufferSize = 1024, CancellationToken cancellationToken = default)
     {
-        var receiveBuffer = new byte[bufferSize];
+        using var receiveBuffer = ArrayPool<byte>.Shared.Lease(1024);
         var result = await client.ReceiveAsync(new ArraySegment<byte>(receiveBuffer), cancellationToken);
 
         if (result.MessageType != WebSocketMessageType.Binary)