diff --git a/src/Discord.Net.DebugTools/UnstableWebSocketClient.cs b/src/Discord.Net.DebugTools/UnstableWebSocketClient.cs index e12e103c5..a0f28ba0a 100644 --- a/src/Discord.Net.DebugTools/UnstableWebSocketClient.cs +++ b/src/Discord.Net.DebugTools/UnstableWebSocketClient.cs @@ -28,7 +28,7 @@ namespace Discord.Net.Providers.UnstableWebSocket private Task _task; private CancellationTokenSource _cancelTokenSource; private CancellationToken _cancelToken, _parentToken; - private bool _isDisposed; + private bool _isDisposed, _isDisconnecting; public UnstableWebSocketClient() { @@ -101,22 +101,44 @@ namespace Discord.Net.Providers.UnstableWebSocket { try { _cancelTokenSource.Cancel(false); } catch { } - if (!isDisposing) + _isDisconnecting = true; + try + { await (_task ?? Task.Delay(0)).ConfigureAwait(false); + _task = null; + } + finally { _isDisconnecting = false; } - if (_client != null && _client.State == WebSocketState.Open) + if (_client != null) { - var token = new CancellationToken(); if (!isDisposing) { - try { await _client.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "", token); } + try { await _client.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "", new CancellationToken()); } catch { } } try { _client.Dispose(); } catch { } + _client = null; } } + private async Task OnClosed(Exception ex) + { + if (_isDisconnecting) + return; //Ignore, this disconnect was requested. + + System.Diagnostics.Debug.WriteLine("OnClosed - " + ex.Message); + await _lock.WaitAsync().ConfigureAwait(false); + try + { + await DisconnectInternalAsync(false); + } + finally + { + _lock.Release(); + } + await Closed(ex); + } public void SetHeader(string key, string value) { @@ -173,10 +195,7 @@ namespace Discord.Net.Providers.UnstableWebSocket int resultCount; if (socketResult.MessageType == WebSocketMessageType.Close) - { - var _ = Closed(new WebSocketClosedException((int)socketResult.CloseStatus, socketResult.CloseStatusDescription)); - return; - } + throw new WebSocketClosedException((int)socketResult.CloseStatus, socketResult.CloseStatusDescription); if (!socketResult.EndOfMessage) { @@ -219,13 +238,13 @@ namespace Discord.Net.Providers.UnstableWebSocket } catch (Win32Exception ex) when (ex.HResult == HR_TIMEOUT) { - var _ = Closed(new Exception("Connection timed out.", ex)); + var _ = OnClosed(new Exception("Connection timed out.", ex)); } catch (OperationCanceledException) { } catch (Exception ex) { //This cannot be awaited otherwise we'll deadlock when DiscordApiClient waits for this task to complete. - var _ = Closed(ex); + var _ = OnClosed(ex); } } diff --git a/src/Discord.Net.WebSocket/Net/DefaultWebSocketClient.cs b/src/Discord.Net.WebSocket/Net/DefaultWebSocketClient.cs index b2ebffabe..aa8bb6986 100644 --- a/src/Discord.Net.WebSocket/Net/DefaultWebSocketClient.cs +++ b/src/Discord.Net.WebSocket/Net/DefaultWebSocketClient.cs @@ -26,7 +26,7 @@ namespace Discord.Net.WebSockets private Task _task; private CancellationTokenSource _cancelTokenSource; private CancellationToken _cancelToken, _parentToken; - private bool _isDisposed; + private bool _isDisposed, _isDisconnecting; public DefaultWebSocketClient() { @@ -98,22 +98,43 @@ namespace Discord.Net.WebSockets { try { _cancelTokenSource.Cancel(false); } catch { } - if (!isDisposing) + _isDisconnecting = true; + try + { await (_task ?? Task.Delay(0)).ConfigureAwait(false); + _task = null; + } + finally { _isDisconnecting = false; } - if (_client != null && _client.State == WebSocketState.Open) + if (_client != null) { - var token = new CancellationToken(); if (!isDisposing) { - try { await _client.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "", token); } + try { await _client.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "", new CancellationToken()); } catch { } } try { _client.Dispose(); } catch { } + _client = null; } } + private async Task OnClosed(Exception ex) + { + if (_isDisconnecting) + return; //Ignore, this disconnect was requested. + + await _lock.WaitAsync().ConfigureAwait(false); + try + { + await DisconnectInternalAsync(false); + } + finally + { + _lock.Release(); + } + await Closed(ex); + } public void SetHeader(string key, string value) { @@ -167,10 +188,7 @@ namespace Discord.Net.WebSockets int resultCount; if (socketResult.MessageType == WebSocketMessageType.Close) - { - var _ = Closed(new WebSocketClosedException((int)socketResult.CloseStatus, socketResult.CloseStatusDescription)); - return; - } + throw new WebSocketClosedException((int)socketResult.CloseStatus, socketResult.CloseStatusDescription); if (!socketResult.EndOfMessage) { @@ -217,13 +235,13 @@ namespace Discord.Net.WebSockets } catch (Win32Exception ex) when (ex.HResult == HR_TIMEOUT) { - var _ = Closed(new Exception("Connection timed out.", ex)); + var _ = OnClosed(new Exception("Connection timed out.", ex)); } catch (OperationCanceledException) { } catch (Exception ex) { //This cannot be awaited otherwise we'll deadlock when DiscordApiClient waits for this task to complete. - var _ = Closed(ex); + var _ = OnClosed(ex); } } }