From 5c903e6fbf2aea3518c1af0e794ab94606ea979f Mon Sep 17 00:00:00 2001
From: Christian Kratky <christian.kratky@googlemail.com>
Date: Sat, 20 Apr 2019 21:12:30 +0200
Subject: [PATCH] Improve connection handling.

---
 .../Implementations/MqttTcpServerListener.cs   | 18 ++++++++++--------
 .../Server/MqttClientSessionsManager.cs        | 12 ++----------
 2 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/Source/MQTTnet/Implementations/MqttTcpServerListener.cs b/Source/MQTTnet/Implementations/MqttTcpServerListener.cs
index 8d771904..b4595364 100644
--- a/Source/MQTTnet/Implementations/MqttTcpServerListener.cs
+++ b/Source/MQTTnet/Implementations/MqttTcpServerListener.cs
@@ -82,7 +82,9 @@ namespace MQTTnet.Implementations
 #else
                     var clientSocket = await _socket.AcceptAsync().ConfigureAwait(false);
 #endif
-                    await HandleClientConnection(clientSocket).ConfigureAwait(false);
+#pragma warning disable 4014
+                    Task.Run(() => TryHandleClientConnectionAsync(clientSocket), cancellationToken);
+#pragma warning restore 4014
                 }
                 catch (Exception exception)
                 {
@@ -92,12 +94,17 @@ namespace MQTTnet.Implementations
             }
         }
 
-        private async Task HandleClientConnection(Socket clientSocket)
+        private async Task TryHandleClientConnectionAsync(Socket clientSocket)
         {
             Stream stream = null;
 
             try
             {
+                _logger.Verbose("Client '{0}' accepted by TCP listener '{1}, {2}'.",
+                    clientSocket.RemoteEndPoint,
+                    _socket.LocalEndPoint,
+                    _addressFamily == AddressFamily.InterNetwork ? "ipv4" : "ipv6");
+
                 clientSocket.NoDelay = _options.NoDelay;
 
                 stream = new NetworkStream(clientSocket, true);
@@ -109,11 +116,6 @@ namespace MQTTnet.Implementations
                     stream = sslStream;
                 }
 
-                _logger.Verbose("Client '{0}' accepted by TCP listener '{1}, {2}'.",
-                    clientSocket.RemoteEndPoint,
-                    _socket.LocalEndPoint,
-                    _addressFamily == AddressFamily.InterNetwork ? "ipv4" : "ipv6");
-
                 var clientAdapter = new MqttChannelAdapter(new MqttTcpChannel(stream), new MqttPacketFormatterAdapter(), _logger);
                 ClientAcceptedHandler?.Invoke(new MqttServerAdapterClientAcceptedEventArgs(clientAdapter));
             }
@@ -138,7 +140,7 @@ namespace MQTTnet.Implementations
                 }
                 catch (Exception disposeException)
                 {
-                    throw new AggregateException(exception, disposeException);
+                    _logger.Error(disposeException, "Error while cleanup of broken connection.");
                 }
             }
         }
diff --git a/Source/MQTTnet/Server/MqttClientSessionsManager.cs b/Source/MQTTnet/Server/MqttClientSessionsManager.cs
index b6fb1431..a6c8c5d7 100644
--- a/Source/MQTTnet/Server/MqttClientSessionsManager.cs
+++ b/Source/MQTTnet/Server/MqttClientSessionsManager.cs
@@ -274,17 +274,9 @@ namespace MQTTnet.Server
             {
                 _connections.TryRemove(clientId, out _);
 
-                ////connection?.ReferenceCounter.Decrement();
-                ////if (connection?.ReferenceCounter.HasReferences == true)
-                ////{
-                ////    disconnectType = MqttClientDisconnectType.Takeover;
-                ////}
-                ////else
+                if (!_options.EnablePersistentSessions)
                 {
-                    if (!_options.EnablePersistentSessions)
-                    {
-                        await DeleteSessionAsync(clientId).ConfigureAwait(false);
-                    }
+                    await DeleteSessionAsync(clientId).ConfigureAwait(false);
                 }
 
                 await TryCleanupChannelAsync(channelAdapter).ConfigureAwait(false);