diff --git a/dotnet/src/webdriver/BiDi/Broker.cs b/dotnet/src/webdriver/BiDi/Broker.cs index a121dde38732d..2cfa324838317 100644 --- a/dotnet/src/webdriver/BiDi/Broker.cs +++ b/dotnet/src/webdriver/BiDi/Broker.cs @@ -17,6 +17,7 @@ // under the License. // +using System.Buffers; using System.Collections.Concurrent; using System.Text.Json; using System.Text.Json.Serialization.Metadata; @@ -47,7 +48,7 @@ public Broker(ITransport transport, IBiDi bidi, Func sessionProv _eventDispatcher = new EventDispatcher(sessionProvider); _receiveMessagesCancellationTokenSource = new CancellationTokenSource(); - _receivingMessageTask = Task.Run(() => ReceiveMessagesAsync(_receiveMessagesCancellationTokenSource.Token)); + _receivingMessageTask = Task.Run(() => ReceiveMessagesLoopAsync(_receiveMessagesCancellationTokenSource.Token)); } public Task SubscribeAsync(string eventName, EventHandler eventHandler, SubscriptionOptions? options, JsonTypeInfo jsonTypeInfo, CancellationToken cancellationToken) @@ -71,7 +72,13 @@ public async Task ExecuteCommandAsync(TCommand comma var timeout = options?.Timeout ?? TimeSpan.FromSeconds(30); cts.CancelAfter(timeout); - var data = JsonSerializer.SerializeToUtf8Bytes(command, jsonCommandTypeInfo); + using var sendBuffer = new PooledBufferWriter(); + + using (var writer = new Utf8JsonWriter(sendBuffer)) + { + JsonSerializer.Serialize(writer, command, jsonCommandTypeInfo); + } + var commandInfo = new CommandInfo(tcs, jsonResultTypeInfo); _pendingCommands[command.Id] = commandInfo; @@ -83,7 +90,16 @@ public async Task ExecuteCommandAsync(TCommand comma try { - await _transport.SendAsync(data, cts.Token).ConfigureAwait(false); + if (_logger.IsEnabled(LogEventLevel.Trace)) + { +#if NET8_0_OR_GREATER + _logger.Trace($"BiDi SND --> {System.Text.Encoding.UTF8.GetString(sendBuffer.WrittenMemory.Span)}"); +#else + _logger.Trace($"BiDi SND --> {System.Text.Encoding.UTF8.GetString(sendBuffer.WrittenMemory.ToArray())}"); +#endif + } + + await _transport.SendAsync(sendBuffer.WrittenMemory, cts.Token).ConfigureAwait(false); } catch { @@ -116,7 +132,7 @@ public async ValueTask DisposeAsync() GC.SuppressFinalize(this); } - private void ProcessReceivedMessage(byte[] data) + private void ProcessReceivedMessage(ReadOnlySpan data) { const int TypeSuccess = 1; const int TypeEvent = 2; @@ -210,20 +226,20 @@ private void ProcessReceivedMessage(byte[] data) { if (_logger.IsEnabled(LogEventLevel.Warn)) { - _logger.Warn($"The remote end responded with 'success' message type, but no pending command with id {id} was found. Message content: {System.Text.Encoding.UTF8.GetString(data)}"); + _logger.Warn($"The remote end responded with 'success' message type, but no pending command with id {id} was found. Message content: {System.Text.Encoding.UTF8.GetString(data.ToArray())}"); } } break; case TypeEvent: - if (method is null) throw new BiDiException($"The remote end responded with 'event' message type, but missed required 'method' property. Message content: {System.Text.Encoding.UTF8.GetString(data)}"); + if (method is null) throw new BiDiException($"The remote end responded with 'event' message type, but missed required 'method' property. Message content: {System.Text.Encoding.UTF8.GetString(data.ToArray())}"); if (!_eventDispatcher.TryGetJsonTypeInfo(method, out var jsonTypeInfo)) { if (_logger.IsEnabled(LogEventLevel.Warn)) { - _logger.Warn($"Received BiDi event with method '{method}', but no event type mapping was found. Event will be ignored. Message content: {System.Text.Encoding.UTF8.GetString(data)}"); + _logger.Warn($"Received BiDi event with method '{method}', but no event type mapping was found. Event will be ignored. Message content: {System.Text.Encoding.UTF8.GetString(data.ToArray())}"); } break; @@ -238,7 +254,7 @@ private void ProcessReceivedMessage(byte[] data) break; case TypeError: - if (id is null) throw new BiDiException($"The remote end responded with 'error' message type, but missed required 'id' property. Message content: {System.Text.Encoding.UTF8.GetString(data)}"); + if (id is null) throw new BiDiException($"The remote end responded with 'error' message type, but missed required 'id' property. Message content: {System.Text.Encoding.UTF8.GetString(data.ToArray())}"); if (_pendingCommands.TryGetValue(id.Value, out var errorCommand)) { @@ -249,7 +265,7 @@ private void ProcessReceivedMessage(byte[] data) { if (_logger.IsEnabled(LogEventLevel.Warn)) { - _logger.Warn($"The remote end responded with 'error' message type, but no pending command with id {id} was found. Message content: {System.Text.Encoding.UTF8.GetString(data)}"); + _logger.Warn($"The remote end responded with 'error' message type, but no pending command with id {id} was found. Message content: {System.Text.Encoding.UTF8.GetString(data.ToArray())}"); } } @@ -258,24 +274,37 @@ private void ProcessReceivedMessage(byte[] data) default: if (_logger.IsEnabled(LogEventLevel.Warn)) { - _logger.Warn($"The remote end responded with unknown message type. Message content: {System.Text.Encoding.UTF8.GetString(data)}"); + _logger.Warn($"The remote end responded with unknown message type. Message content: {System.Text.Encoding.UTF8.GetString(data.ToArray())}"); } break; } } - private async Task ReceiveMessagesAsync(CancellationToken cancellationToken) + private async Task ReceiveMessagesLoopAsync(CancellationToken cancellationToken) { + using var receiveBufferWriter = new PooledBufferWriter(); + try { while (!cancellationToken.IsCancellationRequested) { - var data = await _transport.ReceiveAsync(cancellationToken).ConfigureAwait(false); + receiveBufferWriter.Reset(); + + await _transport.ReceiveAsync(receiveBufferWriter, cancellationToken).ConfigureAwait(false); + + if (_logger.IsEnabled(LogEventLevel.Trace)) + { +#if NET8_0_OR_GREATER + _logger.Trace($"BiDi RCV <-- {System.Text.Encoding.UTF8.GetString(receiveBufferWriter.WrittenMemory.Span)}"); +#else + _logger.Trace($"BiDi RCV <-- {System.Text.Encoding.UTF8.GetString(receiveBufferWriter.WrittenMemory.ToArray())}"); +#endif + } try { - ProcessReceivedMessage(data); + ProcessReceivedMessage(receiveBufferWriter.WrittenMemory.Span); } catch (Exception ex) { @@ -307,4 +336,69 @@ private async Task ReceiveMessagesAsync(CancellationToken cancellationToken) } private readonly record struct CommandInfo(TaskCompletionSource TaskCompletionSource, JsonTypeInfo JsonResultTypeInfo); + + private sealed class PooledBufferWriter : IBufferWriter, IDisposable + { + private const int DefaultBufferSize = 1024 * 8; + + private byte[]? _buffer = ArrayPool.Shared.Rent(DefaultBufferSize); + private int _written; + + public ReadOnlyMemory WrittenMemory => new(_buffer, 0, _written); + + public void Reset() + { + var buffer = _buffer ?? throw new ObjectDisposedException(nameof(PooledBufferWriter)); + + if (buffer.Length > DefaultBufferSize) + { + ArrayPool.Shared.Return(buffer); + _buffer = ArrayPool.Shared.Rent(DefaultBufferSize); + } + + _written = 0; + } + + public void Advance(int count) => _written += count; + + public Memory GetMemory(int sizeHint = 0) + { + EnsureCapacity(sizeHint); + return _buffer.AsMemory(_written); + } + + public Span GetSpan(int sizeHint = 0) + { + EnsureCapacity(sizeHint); + return _buffer.AsSpan(_written); + } + + private void EnsureCapacity(int sizeHint) + { + var buffer = _buffer ?? throw new ObjectDisposedException(nameof(PooledBufferWriter)); + + if (sizeHint <= 0) sizeHint = buffer.Length - _written; + if (sizeHint <= 0) sizeHint = buffer.Length; + + if (_written + sizeHint > buffer.Length) + { + var newSize = Math.Max(buffer.Length * 2, _written + sizeHint); + var newBuffer = ArrayPool.Shared.Rent(newSize); + buffer.AsSpan(0, _written).CopyTo(newBuffer); + ArrayPool.Shared.Return(buffer); + _buffer = newBuffer; + } + } + + public void Dispose() + { + var buffer = _buffer; + + if (buffer is not null) + { + _buffer = null; + ArrayPool.Shared.Return(buffer); + } + } + } } diff --git a/dotnet/src/webdriver/BiDi/ITransport.cs b/dotnet/src/webdriver/BiDi/ITransport.cs index b40f531ad249b..0116e5e7c8ee4 100644 --- a/dotnet/src/webdriver/BiDi/ITransport.cs +++ b/dotnet/src/webdriver/BiDi/ITransport.cs @@ -17,11 +17,13 @@ // under the License. // +using System.Buffers; + namespace OpenQA.Selenium.BiDi; interface ITransport : IAsyncDisposable { - Task ReceiveAsync(CancellationToken cancellationToken); + Task ReceiveAsync(IBufferWriter writer, CancellationToken cancellationToken); - ValueTask SendAsync(ReadOnlyMemory data, CancellationToken cancellationToken); + Task SendAsync(ReadOnlyMemory data, CancellationToken cancellationToken); } diff --git a/dotnet/src/webdriver/BiDi/WebSocketTransport.cs b/dotnet/src/webdriver/BiDi/WebSocketTransport.cs index 11466bc17bbb2..f1d3d40fc01c9 100644 --- a/dotnet/src/webdriver/BiDi/WebSocketTransport.cs +++ b/dotnet/src/webdriver/BiDi/WebSocketTransport.cs @@ -19,7 +19,6 @@ using System.Buffers; using System.Net.WebSockets; -using System.Text; using OpenQA.Selenium.Internal.Logging; namespace OpenQA.Selenium.BiDi; @@ -27,10 +26,8 @@ namespace OpenQA.Selenium.BiDi; sealed class WebSocketTransport(ClientWebSocket webSocket) : ITransport { private readonly static ILogger _logger = Internal.Logging.Log.GetLogger(); - private readonly ClientWebSocket _webSocket = webSocket; private readonly SemaphoreSlim _socketSendSemaphoreSlim = new(1, 1); - private readonly MemoryStream _sharedMemoryStream = new(); public static async Task ConnectAsync(Uri uri, Action? configure, CancellationToken cancellationToken) { @@ -53,61 +50,63 @@ public static async Task ConnectAsync(Uri uri, Action ReceiveAsync(CancellationToken cancellationToken) + public async Task ReceiveAsync(IBufferWriter writer, CancellationToken cancellationToken) { - var receiveBuffer = ArrayPool.Shared.Rent(1024 * 8); +#if NET8_0_OR_GREATER + ValueWebSocketReceiveResult result; - try + do { - _sharedMemoryStream.SetLength(0); + var memory = writer.GetMemory(); - ArraySegment segment = new(receiveBuffer); + result = await _webSocket.ReceiveAsync(memory, cancellationToken).ConfigureAwait(false); - WebSocketReceiveResult result; - - do + if (result.MessageType == WebSocketMessageType.Close) { - result = await _webSocket.ReceiveAsync(segment, cancellationToken).ConfigureAwait(false); + await _webSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None).ConfigureAwait(false); - if (result.MessageType == WebSocketMessageType.Close) - { - await _webSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None).ConfigureAwait(false); + throw new WebSocketException(WebSocketError.ConnectionClosedPrematurely, + $"The remote end closed the WebSocket connection. Status: {_webSocket.CloseStatus}, Description: {_webSocket.CloseStatusDescription}"); + } - throw new WebSocketException(WebSocketError.ConnectionClosedPrematurely, - $"The remote end closed the WebSocket connection. Status: {result.CloseStatus}, Description: {result.CloseStatusDescription}"); - } + writer.Advance(result.Count); + } + while (!result.EndOfMessage); +#else + WebSocketReceiveResult result; - _sharedMemoryStream.Write(receiveBuffer, 0, result.Count); + do + { + var memory = writer.GetMemory(); + + if (!System.Runtime.InteropServices.MemoryMarshal.TryGetArray((ReadOnlyMemory)memory, out var segment)) + { + throw new InvalidOperationException($"The {nameof(IBufferWriter)} must provide array-backed memory."); } - while (!result.EndOfMessage); - byte[] data = _sharedMemoryStream.ToArray(); + result = await _webSocket.ReceiveAsync(segment, cancellationToken).ConfigureAwait(false); - if (_logger.IsEnabled(LogEventLevel.Trace)) + if (result.MessageType == WebSocketMessageType.Close) { - _logger.Trace($"BiDi RCV <-- {Encoding.UTF8.GetString(data)}"); + await _webSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None).ConfigureAwait(false); + + throw new WebSocketException(WebSocketError.ConnectionClosedPrematurely, + $"The remote end closed the WebSocket connection. Status: {result.CloseStatus}, Description: {result.CloseStatusDescription}"); } - return data; - } - finally - { - ArrayPool.Shared.Return(receiveBuffer); + writer.Advance(result.Count); } + while (!result.EndOfMessage); +#endif } - public async ValueTask SendAsync(ReadOnlyMemory data, CancellationToken cancellationToken) + public async Task SendAsync(ReadOnlyMemory data, CancellationToken cancellationToken) { await _socketSendSemaphoreSlim.WaitAsync(cancellationToken).ConfigureAwait(false); try { #if NET8_0_OR_GREATER - if (_logger.IsEnabled(LogEventLevel.Trace)) - { - _logger.Trace($"BiDi SND --> {Encoding.UTF8.GetString(data.Span)}"); - } - await _webSocket.SendAsync(data, WebSocketMessageType.Text, true, cancellationToken).ConfigureAwait(false); #else if (!System.Runtime.InteropServices.MemoryMarshal.TryGetArray(data, out ArraySegment segment)) @@ -115,11 +114,6 @@ public async ValueTask SendAsync(ReadOnlyMemory data, CancellationToken ca segment = new ArraySegment(data.ToArray()); } - if (_logger.IsEnabled(LogEventLevel.Trace)) - { - _logger.Trace($"BiDi SND --> {Encoding.UTF8.GetString(segment.Array!, segment.Offset, segment.Count)}"); - } - await _webSocket.SendAsync(segment, WebSocketMessageType.Text, true, cancellationToken).ConfigureAwait(false); #endif } @@ -151,7 +145,6 @@ public async ValueTask DisposeAsync() } _webSocket.Dispose(); - _sharedMemoryStream.Dispose(); _socketSendSemaphoreSlim.Dispose(); _disposed = true;