diff --git a/dotnet/src/webdriver/BiDi/Broker.cs b/dotnet/src/webdriver/BiDi/Broker.cs index 8f01585ec04d9..162fc392fa920 100644 --- a/dotnet/src/webdriver/BiDi/Broker.cs +++ b/dotnet/src/webdriver/BiDi/Broker.cs @@ -113,7 +113,7 @@ public async ValueTask DisposeAsync() _receiveMessagesCancellationTokenSource.Dispose(); - _transport.Dispose(); + await _transport.DisposeAsync().ConfigureAwait(false); GC.SuppressFinalize(this); } @@ -269,6 +269,15 @@ private async Task ReceiveMessagesAsync(CancellationToken cancellationToken) _logger.Error($"Unhandled error occurred while receiving remote messages: {ex}"); } + // Fail all pending commands, as the connection is likely broken if we failed to receive messages. + foreach (var id in _pendingCommands.Keys) + { + if (_pendingCommands.TryRemove(id, out var pendingCommand)) + { + pendingCommand.TaskCompletionSource.TrySetException(ex); + } + } + throw; } } diff --git a/dotnet/src/webdriver/BiDi/ITransport.cs b/dotnet/src/webdriver/BiDi/ITransport.cs index bdf33406b3936..f202535253c7b 100644 --- a/dotnet/src/webdriver/BiDi/ITransport.cs +++ b/dotnet/src/webdriver/BiDi/ITransport.cs @@ -19,7 +19,7 @@ namespace OpenQA.Selenium.BiDi; -interface ITransport : IDisposable +interface ITransport : IAsyncDisposable { Task ReceiveAsync(CancellationToken cancellationToken); diff --git a/dotnet/src/webdriver/BiDi/WebSocketTransport.cs b/dotnet/src/webdriver/BiDi/WebSocketTransport.cs index 5f9ba5333e8e1..fcab07b9a7770 100644 --- a/dotnet/src/webdriver/BiDi/WebSocketTransport.cs +++ b/dotnet/src/webdriver/BiDi/WebSocketTransport.cs @@ -24,7 +24,7 @@ namespace OpenQA.Selenium.BiDi; -sealed class WebSocketTransport(ClientWebSocket webSocket) : ITransport, IDisposable +sealed class WebSocketTransport(ClientWebSocket webSocket) : ITransport { private readonly static ILogger _logger = Internal.Logging.Log.GetLogger(); @@ -67,6 +67,14 @@ public async Task ReceiveAsync(CancellationToken cancellationToken) { 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); } while (!result.EndOfMessage); @@ -107,26 +115,31 @@ public async Task SendAsync(byte[] data, CancellationToken cancellationToken) private bool _disposed; - public void Dispose() + public async ValueTask DisposeAsync() { - Dispose(true); - GC.SuppressFinalize(this); - } + if (_disposed) return; - private void Dispose(bool disposing) - { - if (_disposed) + if (_webSocket.State == WebSocketState.Open) { - return; + try + { + await _webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None).ConfigureAwait(false); + } + catch (Exception ex) + { + if (_logger.IsEnabled(LogEventLevel.Warn)) + { + _logger.Warn($"Error closing WebSocket gracefully: {ex.Message}"); + } + } } - if (disposing) - { - _webSocket.Dispose(); - _sharedMemoryStream.Dispose(); - _socketSendSemaphoreSlim.Dispose(); - } + _webSocket.Dispose(); + _sharedMemoryStream.Dispose(); + _socketSendSemaphoreSlim.Dispose(); _disposed = true; + + GC.SuppressFinalize(this); } }