diff --git a/Source/MQTTnet/Client/MqttClient.cs b/Source/MQTTnet/Client/MqttClient.cs
index ccedb5ae..93112598 100644
--- a/Source/MQTTnet/Client/MqttClient.cs
+++ b/Source/MQTTnet/Client/MqttClient.cs
@@ -222,8 +222,11 @@ namespace MQTTnet.Client
 
             try
             {
-                await WaitForTaskAsync(_packetReceiverTask, sender).ConfigureAwait(false);
-                await WaitForTaskAsync(_keepAliveMessageSenderTask, sender).ConfigureAwait(false);
+                if (!(exception is MqttCommunicationClosedGracefullyException))
+                {
+                    await WaitForTaskAsync(_packetReceiverTask, sender).ConfigureAwait(false);
+                    await WaitForTaskAsync(_keepAliveMessageSenderTask, sender).ConfigureAwait(false);
+                }
 
                 if (_adapter != null)
                 {
diff --git a/Source/MQTTnet/Implementations/MqttTcpChannel.cs b/Source/MQTTnet/Implementations/MqttTcpChannel.cs
index c48a06f6..07451393 100644
--- a/Source/MQTTnet/Implementations/MqttTcpChannel.cs
+++ b/Source/MQTTnet/Implementations/MqttTcpChannel.cs
@@ -87,8 +87,14 @@ namespace MQTTnet.Implementations
 
         public void Dispose()
         {
-            TryDispose(_stream, () => _stream = null);
-            TryDispose(_socket, () => _socket = null);
+            Cleanup(ref _stream, (s) => s.Dispose());
+            Cleanup(ref _socket, (s) => {
+                if (s.Connected)
+                {
+                    s.Shutdown(SocketShutdown.Both);
+                }
+                s.Dispose();
+            });
         }
 
         private bool InternalUserCertificateValidationCallback(object sender, X509Certificate x509Certificate, X509Chain chain, SslPolicyErrors sslPolicyErrors)
@@ -157,11 +163,16 @@ namespace MQTTnet.Implementations
             }
         }
 
-        private static void TryDispose(IDisposable disposable, Action afterDispose)
+        private static void Cleanup<T>(ref T item, Action<T> handler) where T : class
         {
+            var temp = item;
+            item = null;
             try
             {
-                disposable?.Dispose();
+                if (temp != null)
+                {
+                    handler(temp);
+                }
             }
             catch (ObjectDisposedException)
             {
@@ -169,10 +180,6 @@ namespace MQTTnet.Implementations
             catch (NullReferenceException)
             {
             }
-            finally
-            {
-                afterDispose();
-            }
         }
     }
 }
diff --git a/Source/MQTTnet/Server/MqttClientSession.cs b/Source/MQTTnet/Server/MqttClientSession.cs
index 429fdb45..05d552cb 100644
--- a/Source/MQTTnet/Server/MqttClientSession.cs
+++ b/Source/MQTTnet/Server/MqttClientSession.cs
@@ -29,6 +29,7 @@ namespace MQTTnet.Server
         private MqttApplicationMessage _willMessage;
         private bool _wasCleanDisconnect;
         private IMqttChannelAdapter _adapter;
+        private Task<bool> _run;
 
         public MqttClientSession(
             string clientId,
@@ -65,7 +66,13 @@ namespace MQTTnet.Server
             status.LastNonKeepAlivePacketReceived = _keepAliveMonitor.LastNonKeepAlivePacketReceived;
         }
 
-        public async Task<bool> RunAsync(MqttConnectPacket connectPacket, IMqttChannelAdapter adapter)
+        public Task<bool> RunAsync(MqttConnectPacket connectPacket, IMqttChannelAdapter adapter)
+        {
+            _run = RunInternalAsync(connectPacket, adapter);
+            return _run;
+        }
+
+        private async Task<bool> RunInternalAsync(MqttConnectPacket connectPacket, IMqttChannelAdapter adapter)
         {
             if (connectPacket == null) throw new ArgumentNullException(nameof(connectPacket));
             if (adapter == null) throw new ArgumentNullException(nameof(adapter));
@@ -122,6 +129,16 @@ namespace MQTTnet.Server
                 {
                     _adapter.ReadingPacketStarted -= OnAdapterReadingPacketStarted;
                     _adapter.ReadingPacketCompleted -= OnAdapterReadingPacketCompleted;
+
+                    try
+                    {
+                        await _adapter.DisconnectAsync(_options.DefaultCommunicationTimeout, CancellationToken.None).ConfigureAwait(false);
+                        _adapter.Dispose();
+                    }
+                    catch (Exception exception)
+                    {
+                        _logger.Error(exception, exception.Message);
+                    }
                 }
 
                 _adapter = null;
@@ -153,6 +170,8 @@ namespace MQTTnet.Server
                 }
 
                 _willMessage = null;
+
+                _run?.GetAwaiter().GetResult();
             }
             finally
             {
diff --git a/Source/MQTTnet/Server/MqttClientSessionsManager.cs b/Source/MQTTnet/Server/MqttClientSessionsManager.cs
index 34080a00..6b62e44b 100644
--- a/Source/MQTTnet/Server/MqttClientSessionsManager.cs
+++ b/Source/MQTTnet/Server/MqttClientSessionsManager.cs
@@ -249,16 +249,6 @@ namespace MQTTnet.Server
             }
             finally
             {
-                try
-                {
-                    await clientAdapter.DisconnectAsync(_options.DefaultCommunicationTimeout, CancellationToken.None).ConfigureAwait(false);
-                    clientAdapter.Dispose();
-                }
-                catch (Exception exception)
-                {
-                    _logger.Error(exception, exception.Message);
-                }
-
                 if (!_options.EnablePersistentSessions)
                 {
                     DeleteSession(clientId);
diff --git a/Source/MQTTnet/Server/MqttServer.cs b/Source/MQTTnet/Server/MqttServer.cs
index dd732f1b..6d532917 100644
--- a/Source/MQTTnet/Server/MqttServer.cs
+++ b/Source/MQTTnet/Server/MqttServer.cs
@@ -106,14 +106,14 @@ namespace MQTTnet.Server
 
                 _cancellationTokenSource.Cancel(false);
                 
+                _clientSessionsManager.Stop();
+
                 foreach (var adapter in _adapters)
                 {
                     adapter.ClientAccepted -= OnClientAccepted;
                     await adapter.StopAsync().ConfigureAwait(false);
                 }
 
-                _clientSessionsManager.Stop();
-
                 _logger.Info("Stopped.");
                 Stopped?.Invoke(this, EventArgs.Empty);
             }
diff --git a/Tests/MQTTnet.Core.Tests/MqttServerTests.cs b/Tests/MQTTnet.Core.Tests/MqttServerTests.cs
index 14f1076a..24626ab8 100644
--- a/Tests/MQTTnet.Core.Tests/MqttServerTests.cs
+++ b/Tests/MQTTnet.Core.Tests/MqttServerTests.cs
@@ -253,9 +253,11 @@ namespace MQTTnet.Core.Tests
 
             await c1.ConnectAsync(clientOptions);
 
+            await Task.Delay(500);
+
             await s.StopAsync();
 
-            await Task.Delay(500);
+            await Task.Delay(1000);
 
             Assert.IsTrue(disconnectCalled);
         }