[dotnet] [bidi] Rent buffer for receiving bytes in WebSocket

[nvborisenko]

User description

Most likely the memory is already allocated.

💥 What does this PR do?

Instead of allocating new memory, try to use already allocated 8KB.

🔧 Implementation Notes

And return it back when disposing.

---

PR Type

Enhancement

---

Description

• Use ArrayPool to rent reusable buffer instead of allocating new memory

• Reduces memory allocations by reusing pre-allocated 8KB buffer

• Returns buffer to pool when WebSocketTransport is disposed

• Improves performance through reduced garbage collection pressure

---

Diagram Walkthrough

flowchart LR
  A["WebSocketTransport initialization"] -->|"Rent 8KB buffer"| B["ArrayPool.Shared"]
  B -->|"Reusable byte array"| C["_receiveBuffer"]
  C -->|"Wrap in ArraySegment"| D["ReceiveAsync operations"]
  D -->|"Return buffer on dispose"| B

Loading

File Walkthrough

Relevant files

Performance optimization

WebSocketTransport.cs Implement ArrayPool buffer renting for WebSocket operations dotnet/src/webdriver/BiDi/WebSocketTransport.cs Added System.Buffers namespace import for ArrayPool support Changed _receiveBuffer from ArraySegment with new allocation to rented byte[] from ArrayPool Created local ArraySegment wrapper in ReceiveAsync method for WebSocket API compatibility Added buffer return to ArrayPool in Dispose() method to release rented memory +7/-3

---

[qodo-code-review]

PR Compliance Guide 🔍

Below is a summary of compliance checks for this PR:

Security Compliance
⚪	Sensitive data remanence Description: Returning the rented receive buffer to ArrayPool without clearing it may expose previously received sensitive data to other pool consumers; consider using ArrayPool.Shared.Return(_receiveBuffer, clearArray: true) or manually clearing when the buffer may contain secrets. WebSocketTransport.cs [36-97] Referred Code private readonly byte[] _receiveBuffer = ArrayPool<byte>.Shared.Rent(1024 * 8); private readonly SemaphoreSlim _socketSendSemaphoreSlim = new(1, 1); private readonly MemoryStream _sharedMemoryStream = new(); public async Task ConnectAsync(CancellationToken cancellationToken) { await _webSocket.ConnectAsync(_uri, cancellationToken).ConfigureAwait(false); } public async Task<byte[]> ReceiveAsync(CancellationToken cancellationToken) { _sharedMemoryStream.SetLength(0); ArraySegment<byte> segment = new(_receiveBuffer); WebSocketReceiveResult result; do { result = await _webSocket.ReceiveAsync(segment, cancellationToken).ConfigureAwait(false); ... (clipped 41 lines)
Ticket Compliance
⚪	🎫 No ticket provided Create ticket/issue
Codebase Duplication Compliance
⚪	Codebase context is not defined Follow the guide to enable codebase context checks.
Custom Compliance
🟢	Generic: Meaningful Naming and Self-Documenting Code Objective: Ensure all identifiers clearly express their purpose and intent, making code self-documenting Status: Passed Learn more about managing compliance generic rules or creating your own custom rules
	Generic: Secure Error Handling Objective: To prevent the leakage of sensitive system information through error messages while providing sufficient detail for internal debugging. Status: Passed Learn more about managing compliance generic rules or creating your own custom rules
	Generic: Secure Logging Practices Objective: To ensure logs are useful for debugging and auditing without exposing sensitive information like PII, PHI, or cardholder data. Status: Passed Learn more about managing compliance generic rules or creating your own custom rules
⚪	Generic: Comprehensive Audit Trails Objective: To create a detailed and reliable record of critical system actions for security analysis and compliance. Status: No auditing: The newly added buffer rental/return and receive logic do not add or reference any audit logging for critical actions, and it's unclear whether such logging is handled elsewhere. Referred Code private readonly byte[] _receiveBuffer = ArrayPool<byte>.Shared.Rent(1024 * 8); private readonly SemaphoreSlim _socketSendSemaphoreSlim = new(1, 1); private readonly MemoryStream _sharedMemoryStream = new(); public async Task ConnectAsync(CancellationToken cancellationToken) { await _webSocket.ConnectAsync(_uri, cancellationToken).ConfigureAwait(false); } public async Task<byte[]> ReceiveAsync(CancellationToken cancellationToken) { _sharedMemoryStream.SetLength(0); ArraySegment<byte> segment = new(_receiveBuffer); WebSocketReceiveResult result; do { result = await _webSocket.ReceiveAsync(segment, cancellationToken).ConfigureAwait(false); ... (clipped 41 lines) Learn more about managing compliance generic rules or creating your own custom rules
	Generic: Robust Error Handling and Edge Case Management Objective: Ensure comprehensive error handling that provides meaningful context and graceful degradation Status: Missing handling: The new receive path writes directly from a rented segment without handling ReceiveAsync returning 0 bytes or WebSocket errors/close statuses, which may require explicit checks or logging. Referred Code ArraySegment<byte> segment = new(_receiveBuffer); WebSocketReceiveResult result; do { result = await _webSocket.ReceiveAsync(segment, cancellationToken).ConfigureAwait(false); _sharedMemoryStream.Write(segment.Array!, segment.Offset, result.Count); } while (!result.EndOfMessage); Learn more about managing compliance generic rules or creating your own custom rules
	Generic: Security-First Input Validation and Data Handling Objective: Ensure all data inputs are validated, sanitized, and handled securely to prevent vulnerabilities Status: Buffer return: Returning the rented buffer without clearing may retain sensitive message data in the pool; consider returning with clearArray:true or explicitly clearing before return if messages can contain sensitive data. Referred Code public void Dispose() { _webSocket.Dispose(); _sharedMemoryStream.Dispose(); _socketSendSemaphoreSlim.Dispose(); ArrayPool<byte>.Shared.Return(_receiveBuffer); } Learn more about managing compliance generic rules or creating your own custom rules
Update

Compliance status legend

🟢 - Fully Compliant
🟡 - Partial Compliant
🔴 - Not Compliant
⚪ - Requires Further Human Verification
🏷️ - Compliance label

[qodo-code-review]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	✅ Make the Dispose method idempotent Suggestion Impact: The commit added a _disposed flag, an early return if already disposed, and set the flag after disposing resources, matching the suggested change. code diff: + private bool _disposed; + public void Dispose() { + if (_disposed) + { + return; + } + _webSocket.Dispose(); _sharedMemoryStream.Dispose(); _socketSendSemaphoreSlim.Dispose(); ArrayPool<byte>.Shared.Return(_receiveBuffer); + _disposed = true; } Make the Dispose method idempotent by using a private flag. This prevents returning the same buffer to the ArrayPool multiple times, which would otherwise cause a critical error. dotnet/src/webdriver/BiDi/WebSocketTransport.cs [91-97] +private bool _disposed; + public void Dispose() { + if (_disposed) + { + return; + } + _webSocket.Dispose(); _sharedMemoryStream.Dispose(); _socketSendSemaphoreSlim.Dispose(); ArrayPool<byte>.Shared.Return(_receiveBuffer); + _disposed = true; } [To ensure code accuracy, apply this suggestion manually] Suggestion importance[1-10]: 9 __ Why: The suggestion correctly identifies a critical issue where calling Dispose multiple times would corrupt the ArrayPool by returning the same buffer more than once, potentially leading to memory corruption.	High
	Use explicit size for ArraySegment Create the ArraySegment with an explicit size matching the requested buffer size. This is necessary because ArrayPool.Rent() can return an array larger than requested. dotnet/src/webdriver/BiDi/WebSocketTransport.cs [50] -ArraySegment<byte> segment = new(_receiveBuffer); +const int BufferSize = 1024 * 8; +ArraySegment<byte> segment = new(_receiveBuffer, 0, BufferSize); Apply / Chat Suggestion importance[1-10]: 6 __ Why: The suggestion correctly points out that ArrayPool.Rent can return a larger array than requested, and the current code does not account for this, making it less robust.	Low
Update