diff --git a/src/Discord.Net/Net/WebSockets/BuiltInEngine.cs b/src/Discord.Net/Net/WebSockets/BuiltInEngine.cs index 7dbea06f5..ff2cba729 100644 --- a/src/Discord.Net/Net/WebSockets/BuiltInEngine.cs +++ b/src/Discord.Net/Net/WebSockets/BuiltInEngine.cs @@ -21,6 +21,7 @@ namespace Discord.Net.WebSockets private readonly DiscordConfig _config; private readonly ConcurrentQueue _sendQueue; private WebSocketClient _webSocket; + private Task _tempTask; public event EventHandler BinaryMessage = delegate { }; public event EventHandler TextMessage = delegate { }; @@ -35,18 +36,15 @@ namespace Discord.Net.WebSockets _sendQueue = new ConcurrentQueue(); } - public Task Connect(string host, CancellationToken cancelToken) + public async Task Connect(string host, CancellationToken cancelToken) { - return Task.Run(async () => - { - _webSocket = new WebSocketClient(); - _webSocket.Options.Proxy = null; - _webSocket.Options.SetRequestHeader("User-Agent", _config.UserAgent); - _webSocket.Options.KeepAliveInterval = TimeSpan.Zero; - await _webSocket.ConnectAsync(new Uri(host), cancelToken)//.ConfigureAwait(false); - .ContinueWith(t => ReceiveAsync(cancelToken)).ConfigureAwait(false); - //TODO: ContinueWith is a temporary hack, may be a bug related to https://github.com/dotnet/corefx/issues/4429 - }); + _webSocket = new WebSocketClient(); + _webSocket.Options.Proxy = null; + _webSocket.Options.SetRequestHeader("User-Agent", _config.UserAgent); + _webSocket.Options.KeepAliveInterval = TimeSpan.Zero; + _tempTask = await _webSocket.ConnectAsync(new Uri(host), cancelToken)//.ConfigureAwait(false); + .ContinueWith(t => ReceiveAsync(cancelToken)).ConfigureAwait(false); + //TODO: ContinueWith is a temporary hack, may be a bug related to https://github.com/dotnet/corefx/issues/4429 } public Task Disconnect() @@ -61,7 +59,7 @@ namespace Discord.Net.WebSockets } public IEnumerable GetTasks(CancellationToken cancelToken) - => new Task[] { /*ReceiveAsync(cancelToken),*/ SendAsync(cancelToken) }; + => new Task[] { /*ReceiveAsync(cancelToken),*/ _tempTask, SendAsync(cancelToken) }; private Task ReceiveAsync(CancellationToken cancelToken) { diff --git a/src/Discord.Net/Net/WebSockets/GatewaySocket.cs b/src/Discord.Net/Net/WebSockets/GatewaySocket.cs index 35776a2db..9e0a5c9c5 100644 --- a/src/Discord.Net/Net/WebSockets/GatewaySocket.cs +++ b/src/Discord.Net/Net/WebSockets/GatewaySocket.cs @@ -63,7 +63,7 @@ namespace Discord.Net.WebSockets } catch (OperationCanceledException) { } } - public Task Disconnect() => _taskManager.Stop(true); + public Task Disconnect() => _taskManager.Stop(); protected override async Task Run() { @@ -75,7 +75,7 @@ namespace Discord.Net.WebSockets protected override Task Cleanup() { var ex = _taskManager.Exception; - if (ex != null && (ex as WebSocketException)?.Code != 1012) + if (ex == null || (ex as WebSocketException)?.Code != 1012) //if (ex == null || (ex as WebSocketException)?.Code != 1012) _sessionId = null; //Reset session unless close code 1012 return base.Cleanup(); } diff --git a/src/Discord.Net/Net/WebSockets/WS4NetEngine.cs b/src/Discord.Net/Net/WebSockets/WS4NetEngine.cs index 3edfeba26..c419af0b4 100644 --- a/src/Discord.Net/Net/WebSockets/WS4NetEngine.cs +++ b/src/Discord.Net/Net/WebSockets/WS4NetEngine.cs @@ -16,7 +16,7 @@ namespace Discord.Net.WebSockets private readonly ConcurrentQueue _sendQueue; private readonly TaskManager _taskManager; private WebSocketClient _webSocket; - private ManualResetEventSlim _waitUntilConnect; + private ManualResetEventSlim _waitUntilConnect, _waitUntilDisconnect; public event EventHandler BinaryMessage = delegate { }; public event EventHandler TextMessage = delegate { }; @@ -31,6 +31,7 @@ namespace Discord.Net.WebSockets _taskManager = taskManager; _sendQueue = new ConcurrentQueue(); _waitUntilConnect = new ManualResetEventSlim(); + _waitUntilDisconnect = new ManualResetEventSlim(true); } public Task Connect(string host, CancellationToken cancelToken) @@ -62,13 +63,15 @@ namespace Discord.Net.WebSockets _webSocket = null; if (socket != null) { - //We dont want a slow disconnect to mess up the next connection, so lets just unregister events + socket.Close(); + socket.Opened -= OnWebSocketOpened; socket.DataReceived -= OnWebSocketBinary; socket.MessageReceived -= OnWebSocketText; + + _waitUntilDisconnect.Wait(); //We need the next two events to raise this one socket.Error -= OnWebSocketError; socket.Closed -= OnWebSocketClosed; - socket.Opened -= OnWebSocketOpened; - socket.Close(); + socket.Dispose(); } return TaskHelper.CompletedTask; @@ -78,6 +81,7 @@ namespace Discord.Net.WebSockets { _taskManager.SignalError(e.Exception); _waitUntilConnect.Set(); + _waitUntilDisconnect.Set(); } private void OnWebSocketClosed(object sender, EventArgs e) { @@ -85,12 +89,16 @@ namespace Discord.Net.WebSockets if (e is ClosedEventArgs) ex = new WebSocketException((e as ClosedEventArgs).Code, (e as ClosedEventArgs).Reason); else - ex = new Exception($"Connection lost"); + ex = new Exception("Connection lost"); _taskManager.SignalError(ex); _waitUntilConnect.Set(); + _waitUntilDisconnect.Set(); } private void OnWebSocketOpened(object sender, EventArgs e) - => _waitUntilConnect.Set(); + { + _waitUntilDisconnect.Reset(); + _waitUntilConnect.Set(); + } private void OnWebSocketText(object sender, MessageReceivedEventArgs e) => OnTextMessage(e.Message); private void OnWebSocketBinary(object sender, DataReceivedEventArgs e)