Skip to content

[dotnet] [bidi] ~Zero allocation per command/event #17214

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nvborisenko merged 7 commits into from
Mar 13, 2026
Merged

[dotnet] [bidi] ~Zero allocation per command/event #17214

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
a2eaf9d
Reuse allocated memory
nvborisenko Mar 12, 2026
31f8271
Pooled minimum buffer
nvborisenko Mar 12, 2026
0df02bf
Loop
nvborisenko Mar 12, 2026
1a688ad
ReadOnlySpan
nvborisenko Mar 12, 2026
87d9f1b
Pooled buffer for sending data
nvborisenko Mar 12, 2026
aef6426
Update Broker.cs
nvborisenko Mar 12, 2026
9fc48cf
Clean branches in receive
nvborisenko Mar 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 107 additions & 13 deletions dotnet/src/webdriver/BiDi/Broker.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
// under the License.
// </copyright>

using System.Buffers;
using System.Collections.Concurrent;
using System.Text.Json;
using System.Text.Json.Serialization.Metadata;
Expand Down Expand Up @@ -47,7 +48,7 @@ public Broker(ITransport transport, IBiDi bidi, Func<ISessionModule> sessionProv
_eventDispatcher = new EventDispatcher(sessionProvider);

_receiveMessagesCancellationTokenSource = new CancellationTokenSource();
_receivingMessageTask = Task.Run(() => ReceiveMessagesAsync(_receiveMessagesCancellationTokenSource.Token));
_receivingMessageTask = Task.Run(() => ReceiveMessagesLoopAsync(_receiveMessagesCancellationTokenSource.Token));
}

public Task<Subscription> SubscribeAsync<TEventArgs>(string eventName, EventHandler eventHandler, SubscriptionOptions? options, JsonTypeInfo<TEventArgs> jsonTypeInfo, CancellationToken cancellationToken)
Expand All @@ -71,7 +72,13 @@ public async Task<TResult> ExecuteCommandAsync<TCommand, TResult>(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;

Expand All @@ -83,7 +90,16 @@ public async Task<TResult> ExecuteCommandAsync<TCommand, TResult>(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
{
Expand Down Expand Up @@ -116,7 +132,7 @@ public async ValueTask DisposeAsync()
GC.SuppressFinalize(this);
}

private void ProcessReceivedMessage(byte[] data)
private void ProcessReceivedMessage(ReadOnlySpan<byte> data)
{
const int TypeSuccess = 1;
const int TypeEvent = 2;
Expand Down Expand Up @@ -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())}");
}
Comment thread
nvborisenko marked this conversation as resolved.
Resolved
}

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;
Expand All @@ -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))
{
Expand All @@ -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())}");
}
}

Expand All @@ -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)
{
Expand Down Expand Up @@ -307,4 +336,69 @@ private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
}

private readonly record struct CommandInfo(TaskCompletionSource<EmptyResult> TaskCompletionSource, JsonTypeInfo JsonResultTypeInfo);

private sealed class PooledBufferWriter : IBufferWriter<byte>, IDisposable
{
private const int DefaultBufferSize = 1024 * 8;

private byte[]? _buffer = ArrayPool<byte>.Shared.Rent(DefaultBufferSize);
private int _written;

public ReadOnlyMemory<byte> WrittenMemory => new(_buffer, 0, _written);

public void Reset()
{
var buffer = _buffer ?? throw new ObjectDisposedException(nameof(PooledBufferWriter));

if (buffer.Length > DefaultBufferSize)
{
ArrayPool<byte>.Shared.Return(buffer);
_buffer = ArrayPool<byte>.Shared.Rent(DefaultBufferSize);
}

_written = 0;
}

public void Advance(int count) => _written += count;

public Memory<byte> GetMemory(int sizeHint = 0)
{
EnsureCapacity(sizeHint);
return _buffer.AsMemory(_written);
Comment thread
nvborisenko marked this conversation as resolved.
Resolved
}

public Span<byte> 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<byte>.Shared.Rent(newSize);
buffer.AsSpan(0, _written).CopyTo(newBuffer);
ArrayPool<byte>.Shared.Return(buffer);
_buffer = newBuffer;
}
}

public void Dispose()
{
var buffer = _buffer;

if (buffer is not null)
{
_buffer = null;
ArrayPool<byte>.Shared.Return(buffer);
}
}
}
}
6 changes: 4 additions & 2 deletions dotnet/src/webdriver/BiDi/ITransport.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,13 @@
// under the License.
// </copyright>

using System.Buffers;

namespace OpenQA.Selenium.BiDi;

interface ITransport : IAsyncDisposable
{
Task<byte[]> ReceiveAsync(CancellationToken cancellationToken);
Task ReceiveAsync(IBufferWriter<byte> writer, CancellationToken cancellationToken);

ValueTask SendAsync(ReadOnlyMemory<byte> data, CancellationToken cancellationToken);
Task SendAsync(ReadOnlyMemory<byte> data, CancellationToken cancellationToken);
}
59 changes: 15 additions & 44 deletions dotnet/src/webdriver/BiDi/WebSocketTransport.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,18 +19,15 @@

using System.Buffers;
using System.Net.WebSockets;
using System.Text;
using OpenQA.Selenium.Internal.Logging;

namespace OpenQA.Selenium.BiDi;

sealed class WebSocketTransport(ClientWebSocket webSocket) : ITransport
{
private readonly static ILogger _logger = Internal.Logging.Log.GetLogger<WebSocketTransport>();

private readonly ClientWebSocket _webSocket = webSocket;
private readonly SemaphoreSlim _socketSendSemaphoreSlim = new(1, 1);
private readonly MemoryStream _sharedMemoryStream = new();

public static async Task<ITransport> ConnectAsync(Uri uri, Action<ClientWebSocketOptions>? configure, CancellationToken cancellationToken)
{
Expand All @@ -53,73 +50,48 @@ public static async Task<ITransport> ConnectAsync(Uri uri, Action<ClientWebSocke
return webSocketTransport;
}

public async Task<byte[]> ReceiveAsync(CancellationToken cancellationToken)
public async Task ReceiveAsync(IBufferWriter<byte> writer, CancellationToken cancellationToken)
{
var receiveBuffer = ArrayPool<byte>.Shared.Rent(1024 * 8);
WebSocketReceiveResult result;

try
do
{
_sharedMemoryStream.SetLength(0);

ArraySegment<byte> segment = new(receiveBuffer);
var memory = writer.GetMemory();

WebSocketReceiveResult result;

do
if (!System.Runtime.InteropServices.MemoryMarshal.TryGetArray((ReadOnlyMemory<byte>)memory, out var segment))
{
result = await _webSocket.ReceiveAsync(segment, cancellationToken).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: {result.CloseStatus}, Description: {result.CloseStatusDescription}");
}

_sharedMemoryStream.Write(receiveBuffer, 0, result.Count);
segment = new ArraySegment<byte>(memory.ToArray());
}
Comment thread
nvborisenko marked this conversation as resolved.
Resolved
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<byte>.Shared.Return(receiveBuffer);
writer.Advance(result.Count);
Comment thread
qodo-code-review[bot] marked this conversation as resolved.
Resolved
}
while (!result.EndOfMessage);
}

public async ValueTask SendAsync(ReadOnlyMemory<byte> data, CancellationToken cancellationToken)
public async Task SendAsync(ReadOnlyMemory<byte> 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<byte> segment))
{
segment = new ArraySegment<byte>(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
}
Expand Down Expand Up @@ -151,7 +123,6 @@ public async ValueTask DisposeAsync()
}

_webSocket.Dispose();
_sharedMemoryStream.Dispose();
_socketSendSemaphoreSlim.Dispose();

_disposed = true;
Expand Down
Loading