diff --git a/src/Fleck/Interfaces/ISocket.cs b/src/Fleck/Interfaces/ISocket.cs index b6481a0a..2a553b19 100644 --- a/src/Fleck/Interfaces/ISocket.cs +++ b/src/Fleck/Interfaces/ISocket.cs @@ -17,7 +17,7 @@ public interface ISocket bool NoDelay { get; set; } EndPoint LocalEndPoint { get; } - Task Accept(Action callback, Action error); + Task Accept(Action callback, Action error); Task Send(byte[] buffer, Action callback, Action error); Task Receive(byte[] buffer, Action callback, Action error, int offset = 0); Task Authenticate(X509Certificate2 certificate, SslProtocols enabledSslProtocols, Action callback, Action error); diff --git a/src/Fleck/SocketWrapper.cs b/src/Fleck/SocketWrapper.cs index eb8a2894..91d35bc9 100644 --- a/src/Fleck/SocketWrapper.cs +++ b/src/Fleck/SocketWrapper.cs @@ -130,14 +130,14 @@ public Task Receive(byte[] buffer, Action callback, Action } } - public Task Accept(Action callback, Action error) + public Task Accept(Action callback, Action error) { - Func end = r => _tokenSource.Token.IsCancellationRequested ? null : new SocketWrapper(_socket.EndAccept(r)); + Func end = r => _tokenSource.Token.IsCancellationRequested ? null : new SocketWrapper (_socket.EndAccept (r)); var task = _taskFactory.FromAsync(_socket.BeginAccept, end, null); - task.ContinueWith(t => callback(t.Result), TaskContinuationOptions.OnlyOnRanToCompletion) - .ContinueWith(t => error(t.Exception), TaskContinuationOptions.OnlyOnFaulted); + task.ContinueWith(t => error(t.Exception), TaskContinuationOptions.OnlyOnFaulted); - return task; + + return task.ContinueWith (t => callback (t.Result), TaskContinuationOptions.OnlyOnRanToCompletion); } public void Dispose() diff --git a/src/Fleck/WebSocketServer.cs b/src/Fleck/WebSocketServer.cs index c5ce5721..9594e266 100644 --- a/src/Fleck/WebSocketServer.cs +++ b/src/Fleck/WebSocketServer.cs @@ -4,6 +4,8 @@ using System.Security.Cryptography.X509Certificates; using System.Collections.Generic; using System.Security.Authentication; +using System.Threading; +using System.Threading.Tasks; using Fleck.Helpers; namespace Fleck @@ -93,26 +95,47 @@ public void Start(Action config) _config = config; } - private void ListenForClients() + private void TryRestart () { - ListenerSocket.Accept(OnClientConnect, e => { - FleckLog.Error("Listener socket is closed", e); - if(RestartAfterListenError){ - FleckLog.Info("Listener socket restarting"); - try - { - ListenerSocket.Dispose(); - var socket = new Socket(_locationIP.AddressFamily, SocketType.Stream, ProtocolType.IP); - ListenerSocket = new SocketWrapper(socket); - Start(_config); - FleckLog.Info("Listener socket restarted"); - } - catch (Exception ex) - { - FleckLog.Error("Listener could not be restarted", ex); - } - } - }); + FleckLog.Info ("Listener socket restarting"); + try { + ListenerSocket.Dispose (); + var socket = new Socket (_locationIP.AddressFamily, SocketType.Stream, ProtocolType.IP); + ListenerSocket = new SocketWrapper (socket); + Start (_config); + FleckLog.Info ("Listener socket restarted"); + } catch (Exception ex) { + FleckLog.Error ("Listener socket could not be restarted", ex); + } + } + + private void ListenForClients () + { + ManualResetEvent acceptDone = new ManualResetEvent (false); + bool running = true; + + Task.Run (() => { + + while (running) { + + acceptDone.Reset (); + + var task = ListenerSocket.Accept( + s => { + running = (s != null); + acceptDone.Set (); + OnClientConnect (s); }, + e => { FleckLog.Error ("Error while listening for new clients", e); + if (RestartAfterListenError) TryRestart (); + running = false; acceptDone.Set (); } + ); + + task.ContinueWith((t) => FleckLog.Warn ("Error during client connect", t.Exception), + TaskContinuationOptions.OnlyOnFaulted); + + acceptDone.WaitOne (); + } + }); } private void OnClientConnect(ISocket clientSocket) @@ -120,7 +143,6 @@ private void OnClientConnect(ISocket clientSocket) if (clientSocket == null) return; // socket closed FleckLog.Debug(String.Format("Client connected from {0}:{1}", clientSocket.RemoteIpAddress, clientSocket.RemotePort.ToString())); - ListenForClients(); WebSocketConnection connection = null;