Listener fix

src/Fleck/Interfaces/ISocket.cs

@@ -17,7 +17,7 @@ public interface ISocket
         bool NoDelay { get; set; }
         EndPoint LocalEndPoint { get; }
 
-        Task<ISocket> Accept(Action<ISocket> callback, Action<Exception> error);
+        Task Accept(Action<ISocket> callback, Action<Exception> error);
         Task Send(byte[] buffer, Action callback, Action<Exception> error);
         Task<int> Receive(byte[] buffer, Action<int> callback, Action<Exception> error, int offset = 0);
         Task Authenticate(X509Certificate2 certificate, SslProtocols enabledSslProtocols, Action callback, Action<Exception> error);

src/Fleck/SocketWrapper.cs

@@ -130,14 +130,14 @@ public Task<int> Receive(byte[] buffer, Action<int> callback, Action<Exception>
             }
         }
 
-        public Task<ISocket> Accept(Action<ISocket> callback, Action<Exception> error)
+        public Task Accept(Action<ISocket> callback, Action<Exception> error)
         {
-            Func<IAsyncResult, ISocket> end = r => _tokenSource.Token.IsCancellationRequested ? null : new SocketWrapper(_socket.EndAccept(r));
+            Func<IAsyncResult, ISocket> 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()

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,34 +95,54 @@ public void Start(Action<IWebSocketConnection> 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)
         {
             if (clientSocket == null) return; // socket closed
 
             FleckLog.Debug(String.Format("Client connected from {0}:{1}", clientSocket.RemoteIpAddress, clientSocket.RemotePort.ToString()));
-            ListenForClients();
 
             WebSocketConnection connection = null;