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

[nvborisenko]

This is big improvement how we use memory.

💥 What does this PR do?

This pull request introduces a significant refactor to the BiDi transport and broker layers, focusing on performance improvements and memory management. The main changes include replacing byte array and memory stream usage with pooled buffer writers, updating interfaces for more efficient message handling, and enhancing logging for trace-level diagnostics.

These changes collectively improve the efficiency, reliability, and maintainability of the BiDi transport and broker components.

🔄 Types of changes

• Cleanup (formatting, renaming)
• New feature (non-breaking change which adds functionality and tests!)

[qodo-code-review]

Review Summary by Qodo

Implement pooled buffers for zero-allocation BiDi messaging

✨ Enhancement

Walkthroughs

Description

• Implement pooled buffer writers for zero-allocation message handling
• Replace byte array allocations with reusable buffer pools
• Update transport interface to use IBufferWriter pattern
• Enhance trace-level logging for BiDi commands and events

File Changes

1. dotnet/src/webdriver/BiDi/Broker.cs ✨ Enhancement +107/-13

Pooled buffer implementation and memory optimization

• Introduced PooledBufferWriter class implementing IBufferWriter<byte> for memory reuse via
 ArrayPool
• Replaced JsonSerializer.SerializeToUtf8Bytes() with Utf8JsonWriter using pooled buffers for
 command serialization
• Updated ReceiveMessagesAsync to ReceiveMessagesLoopAsync with pooled buffer reuse in receive
 loop
• Changed ProcessReceivedMessage parameter from byte[] to ReadOnlySpan<byte> for zero-copy
 operations
• Added conditional trace-level logging for sent/received messages with NET8_0_OR_GREATER
 optimization

dotnet/src/webdriver/BiDi/Broker.cs

---

2. dotnet/src/webdriver/BiDi/ITransport.cs ✨ Enhancement +4/-2

Transport interface refactoring for buffer writer pattern

• Changed ReceiveAsync signature to accept IBufferWriter<byte> instead of returning byte[]
• Changed SendAsync return type from ValueTask to Task for consistency
• Added System.Buffers namespace import

dotnet/src/webdriver/BiDi/ITransport.cs

---

3. dotnet/src/webdriver/BiDi/WebSocketTransport.cs ✨ Enhancement +15/-44

WebSocket transport buffer writer integration

• Removed _sharedMemoryStream field and replaced with direct buffer writer usage
• Refactored ReceiveAsync to write directly to provided IBufferWriter<byte> instead of
 accumulating in memory stream
• Removed duplicate trace-level logging from SendAsync (now handled in Broker.cs)
• Simplified receive loop by using writer.Advance() instead of manual buffer management
• Changed SendAsync return type from ValueTask to Task

dotnet/src/webdriver/BiDi/WebSocketTransport.cs

---

[qodo-code-review]

Code Review by Qodo

🐞 Bugs (1) 📘 Rule violations (2) 📎 Requirement gaps (0)

1. Fallback receive drops bytes ☑ 🐞 Bug ✓ Correctness

Description

WebSocketTransport.ReceiveAsync receives into a temporary array created by memory.ToArray() when
MemoryMarshal.TryGetArray fails, then calls writer.Advance(result.Count) without copying the
received bytes into the writer buffer, so the writer reports bytes written that were never written
to it.

Code

dotnet/src/webdriver/BiDi/WebSocketTransport.cs[R59-76]

+            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());
          }
-            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);

Evidence

In the fallback branch, the WebSocket writes into an ArraySegment backed by a newly allocated array,
but the only interaction with the provided IBufferWriter is Advance(count); no copy into the
writer’s Memory/Span occurs. Broker then reads from the writer’s WrittenMemory, so if the fallback
branch is taken it will parse stale/garbage bytes rather than the bytes received from the socket.

dotnet/src/webdriver/BiDi/WebSocketTransport.cs[59-77]
dotnet/src/webdriver/BiDi/ITransport.cs[24-29]
dotnet/src/webdriver/BiDi/Broker.cs[292-308]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
`WebSocketTransport.ReceiveAsync` has a fallback for non-array-backed `IBufferWriter<byte>` buffers. In that fallback, it receives into a temporary `byte[]` created by `memory.ToArray()`, but it never copies the received bytes into the writer’s buffer before calling `writer.Advance(result.Count)`. This can corrupt incoming messages if the fallback path is taken.
### Issue Context
`Broker.ReceiveMessagesLoopAsync` reads from `receiveBufferWriter.WrittenMemory` immediately after `ReceiveAsync` completes, so `ReceiveAsync` must guarantee that the bytes it counts as written are actually written into the provided writer.
### Fix Focus Areas
- dotnet/src/webdriver/BiDi/WebSocketTransport.cs[59-77]
### Suggested implementation direction
- Keep `var memory = writer.GetMemory(sizeHint)`.
- If `TryGetArray` fails:
- Allocate/rent a temp array segment of length `memory.Length`.
- Receive into that temp segment.
- Copy `temp.AsSpan(0, result.Count)` into `memory.Span` (or into `writer.GetSpan(result.Count)`).
- Then call `writer.Advance(result.Count)`.
- (Optional) Consider using `ArrayPool<byte>.Shared` instead of `memory.ToArray()` to avoid a GC allocation in the fallback path.

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

---

2. Trace logs full payload 📘 Rule violation ⛨ Security

Description

New/modified logging emits full BiDi JSON payloads (send/receive and warning paths), which can leak
sensitive data into logs and create large/unstable log output. This conflicts with the requirement
to keep logs resilient and avoid potentially sensitive/unstable message content.

Code

dotnet/src/webdriver/BiDi/Broker.cs[R93-102]

+            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);

Evidence

PR Compliance ID 13 requires resilient logging that avoids sensitive/unstable content; the changed
code logs entire JSON messages via Encoding.UTF8.GetString(...) for send/receive and includes
message content in warnings/exceptions.

dotnet/src/webdriver/BiDi/Broker.cs[93-102]
dotnet/src/webdriver/BiDi/Broker.cs[296-303]
dotnet/src/webdriver/BiDi/Broker.cs[229-229]
Best Practice: Learned patterns

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
The Broker emits full BiDi JSON payloads to logs (trace and warn/exception message-content paths). This can leak sensitive data and produce unstable/very large logs.
## Issue Context
Compliance requires resilient logging that avoids sensitive/unstable message content and favors stable identifiers.
## Fix Focus Areas
- dotnet/src/webdriver/BiDi/Broker.cs[93-102]
- dotnet/src/webdriver/BiDi/Broker.cs[296-303]
- dotnet/src/webdriver/BiDi/Broker.cs[226-279]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

---

3. No tests for refactor 📘 Rule violation ⛯ Reliability

Description

The transport/message handling was refactored (new pooled buffer writer, changed ITransport
receive/send contracts), but no corresponding tests are included in the provided changes. This
increases the risk of regressions in message framing and command/event processing.

Code

dotnet/src/webdriver/BiDi/ITransport.cs[R24-28]

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);

Evidence

PR Compliance ID 3 requires adding/updating tests when behavior changes; the diff shows significant
behavioral/interface changes to transport and Broker message handling without any test changes in
the provided PR diff.

AGENTS.md
dotnet/src/webdriver/BiDi/ITransport.cs[24-28]
dotnet/src/webdriver/BiDi/WebSocketTransport.cs[53-79]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
Transport and Broker message handling were refactored (buffer-writer based receive, pooled writer, updated async contracts), but tests were not updated/added in the provided changes.
## Issue Context
These are core message-path changes that can regress framing, partial reads, and payload handling. Prefer small/unit tests and avoid mocks that distort contracts.
## Fix Focus Areas
- dotnet/src/webdriver/BiDi/ITransport.cs[24-28]
- dotnet/src/webdriver/BiDi/WebSocketTransport.cs[53-79]
- dotnet/src/webdriver/BiDi/Broker.cs[72-103]
- dotnet/src/webdriver/BiDi/Broker.cs[284-308]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

---

ⓘ The new review experience is currently in Beta. Learn more

[Copilot]

Pull request overview

This PR refactors the .NET BiDi transport layer to reduce per-message allocations by switching from byte[]-returning receive APIs to writing into a pooled IBufferWriter<byte> and by pooling JSON serialization buffers for outbound messages.

Changes:

• Updated ITransport to receive into an IBufferWriter<byte> and adjusted send to use Task.
• Reworked WebSocketTransport.ReceiveAsync to write directly into a provided buffer instead of allocating/aggregating via a MemoryStream.
• Added a pooled buffer writer in Broker and switched outbound command serialization to Utf8JsonWriter over pooled buffers (plus adjusted trace logging).

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
dotnet/src/webdriver/BiDi/WebSocketTransport.cs	Refactors receive/send signatures and implements buffered receive into an IBufferWriter<byte>.
dotnet/src/webdriver/BiDi/ITransport.cs	Updates the transport interface to support pooled/buffered receive and Task-based send.
dotnet/src/webdriver/BiDi/Broker.cs	Pools outbound/inbound buffers and switches command serialization to write into a pooled buffer.