[dotnet] [bidi] Rent buffer for receiving bytes in WebSocket (#16650)

Author: nvborisenko

dotnet/src/webdriver/BUILD.bazel

@@ -61,6 +61,7 @@ csharp_library(
     ],
     deps = [
         nuget_package("Microsoft.Bcl.AsyncInterfaces"),
+        nuget_package("System.Buffers"),
         nuget_package("System.Threading.Tasks.Extensions"),
         nuget_package("System.Memory"),
         nuget_package("System.Runtime.CompilerServices.Unsafe"),
@@ -94,6 +95,7 @@ csharp_library(
     deps = [
         nuget_package("NETStandard.Library"),
         nuget_package("Microsoft.Bcl.AsyncInterfaces"),
+        nuget_package("System.Buffers"),
         nuget_package("System.Threading.Tasks.Extensions"),
         nuget_package("System.Memory"),
         nuget_package("System.Runtime.CompilerServices.Unsafe"),
@@ -150,6 +152,7 @@ csharp_library(
     ],
     deps = [
         nuget_package("Microsoft.Bcl.AsyncInterfaces"),
+        nuget_package("System.Buffers"),
         nuget_package("System.Threading.Tasks.Extensions"),
         nuget_package("System.Memory"),
         nuget_package("System.Runtime.CompilerServices.Unsafe"),
@@ -181,6 +184,7 @@ csharp_library(
     deps = [
         nuget_package("NETStandard.Library"),
         nuget_package("Microsoft.Bcl.AsyncInterfaces"),
+        nuget_package("System.Buffers"),
         nuget_package("System.Threading.Tasks.Extensions"),
         nuget_package("System.Memory"),
         nuget_package("System.Runtime.CompilerServices.Unsafe"),

dotnet/src/webdriver/BiDi/Module.cs

@@ -23,9 +23,9 @@ namespace OpenQA.Selenium.BiDi;
 
 public abstract class Module
 {
-    protected BiDi BiDi { get; private set; }
+    protected BiDi BiDi { get; private set; } = null!;
 
-    protected Broker Broker { get; private set; }
+    protected Broker Broker { get; private set; } = null!;
 
     protected abstract void Initialize(JsonSerializerOptions options);
 

dotnet/src/webdriver/BiDi/WebSocketTransport.cs

@@ -19,6 +19,7 @@
 
 using OpenQA.Selenium.Internal.Logging;
 using System;
+using System.Buffers;
 using System.IO;
 using System.Net.WebSockets;
 using System.Text;
@@ -32,7 +33,7 @@ class WebSocketTransport(Uri _uri) : ITransport, IDisposable
     private readonly static ILogger _logger = Internal.Logging.Log.GetLogger<WebSocketTransport>();
 
     private readonly ClientWebSocket _webSocket = new();
-    private readonly ArraySegment<byte> _receiveBuffer = new(new byte[1024 * 8]);
+    private readonly byte[] _receiveBuffer = ArrayPool<byte>.Shared.Rent(1024 * 8);
 
     private readonly SemaphoreSlim _socketSendSemaphoreSlim = new(1, 1);
     private readonly MemoryStream _sharedMemoryStream = new();
@@ -46,13 +47,15 @@ public async Task<byte[]> ReceiveAsync(CancellationToken cancellationToken)
     {
         _sharedMemoryStream.SetLength(0);
 
+        ArraySegment<byte> segment = new(_receiveBuffer);
+
         WebSocketReceiveResult result;
 
         do
         {
-            result = await _webSocket.ReceiveAsync(_receiveBuffer, cancellationToken).ConfigureAwait(false);
+            result = await _webSocket.ReceiveAsync(segment, cancellationToken).ConfigureAwait(false);
 
-            _sharedMemoryStream.Write(_receiveBuffer.Array!, _receiveBuffer.Offset, result.Count);
+            _sharedMemoryStream.Write(segment.Array!, segment.Offset, result.Count);
         }
         while (!result.EndOfMessage);
 
@@ -85,10 +88,32 @@ public async Task SendAsync(byte[] data, CancellationToken cancellationToken)
         }
     }
 
+    private bool _disposed;
+
     public void Dispose()
     {
+        if (_disposed)
+        {
+            return;
+        }
+
         _webSocket.Dispose();
         _sharedMemoryStream.Dispose();
         _socketSendSemaphoreSlim.Dispose();
+        ReleaseBuffer();
+        _disposed = true;
+    }
+
+    ~WebSocketTransport()
+    {
+        ReleaseBuffer();
+    }
+
+    private void ReleaseBuffer()
+    {
+        if (_receiveBuffer is not null)
+        {
+            ArrayPool<byte>.Shared.Return(_receiveBuffer);
+        }
     }
 }