[dotnet] [bidi] Decouple transport and processing

[nvborisenko]

Always receiving from remote end, processing from queue.

💥 What does this PR do?

This pull request significantly refactors the Broker class in the WebDriver BiDi implementation to improve message handling, resource management, and error propagation. The changes introduce asynchronous message processing using channels, enhance buffer pooling, and ensure more robust disposal and error handling throughout the broker lifecycle. Additionally, related test code is updated for proper asynchronous disposal and a new test is added to verify idempotent disposal.

Message Handling and Processing Improvements:

• Replaced the single message receiving loop with two coordinated tasks: one for receiving messages (ReceiveMessagesAsync) and another for processing messages (ProcessMessagesAsync), communicating via a bounded channel to apply backpressure and improve throughput. This also introduces a buffer pool to reduce allocations and ensure proper buffer reuse. [1] [2] [3] [4]
• Added robust error propagation: if message receiving fails, a terminal exception is stored and propagated to pending commands and subsequent command executions. [1] [2]

Resource and Buffer Management:

• Implemented a buffer pool using a channel to efficiently reuse PooledBufferWriter instances, and ensured all buffers are disposed of during broker disposal. [1] [2] [3]
• Improved the Advance and EnsureCapacity methods of PooledBufferWriter to guard against invalid buffer operations and ensure correct buffer sizing. [1] [2]

Disposal and Test Enhancements:

• Refactored the disposal logic to await both receiving and processing tasks, dispose of all resources in the correct order, and drain the buffer pool.
• Updated test fixture teardown to dispose the driver asynchronously and added a test to verify that multiple calls to DisposeAsync are safe (idempotent). [1] [2]

🔄 Types of changes

• New feature (non-breaking change which adds functionality and tests!)

[qodo-code-review]

Review Summary by Qodo

Decouple BiDi transport and message processing with channel pattern

✨ Enhancement

Walkthroughs

Description

• Decouple message receiving and processing into separate async tasks
• Implement channel-based producer/consumer pattern for message handling
• Add buffer pooling to reduce memory allocations and improve performance
• Improve resource management with proper cleanup of pooled buffers

File Changes

1. dotnet/src/webdriver/BiDi/Broker.cs ✨ Enhancement +74/-21

Implement producer/consumer pattern with buffer pooling

• Added System.Threading.Channels import for channel-based architecture
• Introduced bounded Channel<PooledBufferWriter> for queuing received messages
• Implemented ConcurrentBag<PooledBufferWriter> buffer pool with RentBuffer and ReturnBuffer
 methods
• Split message handling into ReceiveMessagesAsync (producer) and ProcessMessagesAsync
 (consumer) tasks
• Refactored DisposeAsync to await both receiving and processing tasks and dispose pooled buffers
• Moved message processing logic from receive loop to separate processing task

dotnet/src/webdriver/BiDi/Broker.cs

---

[qodo-code-review]

Code Review by Qodo

🐞 Bugs (0) 📘 Rule violations (0) 📎 Requirement gaps (0) 🎨 UX Issues (0)

1. Premature command failure ☑ 🐞 Bug ≡ Correctness

Description

On a receive exception, ReceiveMessagesAsync faults and removes all pending commands immediately,
even though already-received messages may still be buffered in the channel and not yet processed.
This can incorrectly fail commands whose responses were successfully received but were still queued
for ProcessMessagesAsync.

Code

dotnet/src/webdriver/BiDi/Broker.cs[R344-345]

            throw;

Evidence

Broker decouples receiving and processing via a bounded channel (capacity 16) and a separate
processing task, so there can be up to 16 unprocessed response messages sitting in
_receivedMessages. If ReceiveMessagesAsync throws (e.g., transport disconnect) after it has already
enqueued responses, the catch block faults/removes all _pendingCommands before ProcessMessagesAsync
drains the queued responses, causing those responses to be ignored and their commands to surface an
exception instead of completing normally.

dotnet/src/webdriver/BiDi/Broker.cs[42-49]
dotnet/src/webdriver/BiDi/Broker.cs[299-346]
dotnet/src/webdriver/BiDi/Broker.cs[353-378]

Agent prompt

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

### Issue description
`ReceiveMessagesAsync` currently faults/removes all pending commands immediately when receiving fails, even though responses already queued in `_receivedMessages` may still complete those commands. This creates a race where valid buffered responses are ignored and callers see a transport exception instead.

### Issue Context
The receiver and processor run as separate tasks with a bounded channel between them. On receiver failure, you should allow the processor to drain already-received messages before failing whatever commands truly remain unresolved.

### Fix Focus Areas
- dotnet/src/webdriver/BiDi/Broker.cs[299-351]
- dotnet/src/webdriver/BiDi/Broker.cs[353-378]
- dotnet/src/webdriver/BiDi/Broker.cs[42-49]

### Suggested approach
- Do **not** fail `_pendingCommands` inside the receiver catch.
- Instead, capture the receive exception and complete the channel with it (`_receivedMessages.Writer.TryComplete(ex)`), then in `ProcessMessagesAsync` after the read loop ends:
 - await `reader.Completion` to observe whether completion was faulted
 - if faulted (or if you captured a receive exception), fail any **remaining** `_pendingCommands` at that point (after the channel has been drained).
This preserves correctness for responses that were already received while still failing commands when the connection is truly broken.

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

---

2. DisposeAsync disposes _eventDispatcher early ☑ 📘 Rule violation ☼ Reliability

Description

Broker.DisposeAsync() disposes _eventDispatcher before awaiting _processingTask, so
ProcessMessagesAsync() can still call _eventDispatcher.EnqueueEvent() after the dispatcher has
completed its channel, silently dropping in-flight events. This is an unsafe shutdown ordering in
concurrent code that can lead to lost/partial event delivery during disposal.

Code

dotnet/src/webdriver/BiDi/Broker.cs[R126-137]

        try
        {
-            await _receivingMessageTask.ConfigureAwait(false);
+            await _receivingTask.ConfigureAwait(false);
        }
        catch (OperationCanceledException) when (_receiveMessagesCancellationTokenSource.IsCancellationRequested)
        {
            // Expected when cancellation is requested, ignore.
        }

+        await _processingTask.ConfigureAwait(false);
+

Evidence

The checklist requires correct synchronization/ordering for concurrent state transitions. In the new
shutdown sequence, the dispatcher is disposed (completing its channel) before the message-processing
task is awaited, while message processing can still enqueue events, causing dropped events due to
TryWrite on a completed channel.

dotnet/src/webdriver/BiDi/Broker.cs[121-137]
dotnet/src/webdriver/BiDi/Broker.cs[353-376]
dotnet/src/webdriver/BiDi/Broker.cs[250-269]
dotnet/src/webdriver/BiDi/EventDispatcher.cs[74-77]
dotnet/src/webdriver/BiDi/EventDispatcher.cs[118-127]
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
`Broker.DisposeAsync()` disposes `_eventDispatcher` before `_processingTask` has completed. Since `ProcessMessagesAsync()` can still call `_eventDispatcher.EnqueueEvent()` while draining `_receivedMessages`, events can be silently dropped because `EventDispatcher.DisposeAsync()` completes `_pendingEvents.Writer`.

## Issue Context
The new producer/consumer split means message processing can outlive receiving and can continue enqueuing events during shutdown. Disposal ordering should ensure the processing pipeline has fully drained (or is explicitly prevented from enqueuing) before tearing down the dispatcher.

## Fix Focus Areas
- dotnet/src/webdriver/BiDi/Broker.cs[121-137]
- dotnet/src/webdriver/BiDi/Broker.cs[353-376]
- dotnet/src/webdriver/BiDi/EventDispatcher.cs[74-77]
- dotnet/src/webdriver/BiDi/EventDispatcher.cs[118-127]

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

---

3. Dispose skips cleanup on fault ☑ 🐞 Bug ☼ Reliability

Description

DisposeAsync will abort before disposing the transport and pooled buffers if either background task
faults with a non-cancellation exception. This can leak the underlying transport and pooled arrays
when ReceiveMessagesAsync throws (it rethrows on non-cancellation failures).

Code

dotnet/src/webdriver/BiDi/Broker.cs[R126-146]

        try
        {
-            await _receivingMessageTask.ConfigureAwait(false);
+            await _receivingTask.ConfigureAwait(false);
        }
        catch (OperationCanceledException) when (_receiveMessagesCancellationTokenSource.IsCancellationRequested)
        {
            // Expected when cancellation is requested, ignore.
        }

+        await _processingTask.ConfigureAwait(false);
+
        _receiveMessagesCancellationTokenSource.Dispose();

        await _transport.DisposeAsync().ConfigureAwait(false);

+        while (_bufferPool.TryTake(out var buffer))
+        {
+            buffer.Dispose();
+        }
+

Evidence

DisposeAsync only swallows OperationCanceledException; any other exception from awaiting
_receivingTask or _processingTask will prevent execution of the subsequent cleanup code
(_transport.DisposeAsync() and draining/disposing _bufferPool). ReceiveMessagesAsync explicitly
rethrows on non-cancellation receive failures, making this path reachable under real
disconnect/error conditions.

dotnet/src/webdriver/BiDi/Broker.cs[121-148]
dotnet/src/webdriver/BiDi/Broker.cs[329-351]

Agent prompt

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

### Issue description
`Broker.DisposeAsync()` can exit early (throw) when `_receivingTask` or `_processingTask` faults, skipping disposal of `_transport` and pooled buffers. This can leak sockets and ArrayPool-rented memory.

### Issue Context
`ReceiveMessagesAsync` rethrows on non-cancellation exceptions, so the receiving task can fault during normal runtime failures (e.g., remote closes websocket).

### Fix Focus Areas
- dotnet/src/webdriver/BiDi/Broker.cs[121-148]
- dotnet/src/webdriver/BiDi/Broker.cs[299-351]

### Suggested approach
- Restructure `DisposeAsync` using `try/finally` so that transport and buffer cleanup always runs.
- Optionally capture exceptions from `_receivingTask`/`_processingTask` and rethrow **after** cleanup (or aggregate them) so disposal remains deterministic but still reports failures.
- Ensure `_processingTask` is awaited/observed even if `_receivingTask` faults, to avoid leaving work running in the background.

ⓘ 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

Refactors the .NET BiDi Broker message flow to decouple transport receiving from message processing, aiming to improve throughput and reduce allocations via buffering and pooling.

Changes:

• Introduced a bounded Channel<PooledBufferWriter> to queue received messages and process them on a dedicated task.
• Added a ConcurrentBag<PooledBufferWriter> pool with rent/return helpers to reuse buffers.
• Updated disposal to await both receiving and processing tasks and to dispose pooled buffers.

[Copilot]

Pull request overview

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

[nvborisenko]

I am here, let's continue.

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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