Improve internal session/connection management and connection takeover.

Source/MQTTnet/Server/MqttClientConnection.cs

@@ -42,8 +42,6 @@ namespace MQTTnet.Server
         long _receivedApplicationMessagesCount;
         long _sentApplicationMessagesCount;
 
-        volatile bool _isTakeover;
-
         public MqttClientConnection(MqttConnectPacket connectPacket,
             IMqttChannelAdapter channelAdapter,
             MqttClientSession session,
@@ -74,6 +72,8 @@ namespace MQTTnet.Server
 
         public MqttConnectPacket ConnectPacket { get; }
 
+        public bool IsTakeOver { get; set; }
+
         public string ClientId => ConnectPacket.ClientId;
 
         public bool IsReadingPacket => _channelAdapter.IsReadingPacket;
@@ -82,19 +82,30 @@ namespace MQTTnet.Server
 
         public MqttClientSession Session { get; }
 
-        public Task StopAsync(bool isTakeover = false)
+        public async Task StopAsync()
         {
-            _isTakeover = isTakeover;
+            if (IsTakeOver)
-            var task = _packageReceiverTask;
+            {
+                // Is is very important to send the DISCONNECT packet here BEFORE cancelling the
+                // token because the entire connection is closed (disposed) as soon as the cancellation
+                // token is cancelled. To there is no chance that the DISCONNECT packet will ever arrive
+                // at the client!
+                try
+                {
+                    await _channelAdapter.SendPacketAsync(new MqttDisconnectPacket
+                    {
+                        ReasonCode = MqttDisconnectReasonCode.SessionTakenOver
+                    }, _serverOptions.DefaultCommunicationTimeout, CancellationToken.None).ConfigureAwait(false);
+                }
+                catch (Exception exception)
+                {
+                    _logger.Error(exception, "Client '{0}': Error while sending DISCONNECT packet after takeover.", ClientId);
+                }
+            }
 
             StopInternal();
 
-            if (task != null)
+            await _packageReceiverTask;
-            {
-                return task;
-            }
-
-            return PlatformAbstractionLayer.CompletedTask;
         }
 
         public void ResetStatistics()
@@ -129,11 +140,11 @@ namespace MQTTnet.Server
 
         public Task RunAsync()
         {
-            _packageReceiverTask = RunInternalAsync();
+            _packageReceiverTask = RunInternalAsync(_cancellationToken.Token);
             return _packageReceiverTask;
         }
 
-        async Task RunInternalAsync()
+        async Task RunInternalAsync(CancellationToken cancellationToken)
         {
             var disconnectType = MqttClientDisconnectType.NotClean;
             try
@@ -142,12 +153,10 @@ namespace MQTTnet.Server
 
                 Session.WillMessage = ConnectPacket.WillMessage;
 
-                var cancellationToken = _cancellationToken.Token;
+                await SendAsync(_channelAdapter.PacketFormatterAdapter.DataConverter.CreateConnAckPacket(_connectionValidatorContext), cancellationToken).ConfigureAwait(false);
 
                 Task.Run(() => SendPendingPacketsAsync(cancellationToken), cancellationToken).Forget(_logger);
 
-                await SendAsync(_channelAdapter.PacketFormatterAdapter.DataConverter.CreateConnAckPacket(_connectionValidatorContext), cancellationToken).ConfigureAwait(false);
-
                 Session.IsCleanSession = false;
 
                 while (!cancellationToken.IsCancellationRequested)
@@ -156,7 +165,7 @@ namespace MQTTnet.Server
                     if (packet == null)
                     {
                         // The client has closed the connection gracefully.
-                        break;
+                        return;
                     }
 
                     Interlocked.Increment(ref _sentPacketsCount);
@@ -209,7 +218,7 @@ namespace MQTTnet.Server
                         disconnectType = MqttClientDisconnectType.Clean;
 
                         StopInternal();
-                        break;
+                        return;
                     }
 
                     _packetDispatcher.Dispatch(packet);
@@ -233,39 +242,21 @@ namespace MQTTnet.Server
             }
             finally
             {
-                if (_isTakeover)
+                if (IsTakeOver)
                 {
                     disconnectType = MqttClientDisconnectType.Takeover;
                 }
-                
+
                 if (Session.WillMessage != null)
                 {
                     _sessionsManager.DispatchApplicationMessage(Session.WillMessage, this);
                     Session.WillMessage = null;
                 }
 
-                _packetDispatcher.Reset();
+                _packetDispatcher.Cancel();
-
-                _packageReceiverTask = null;
-
-                if (_isTakeover)
-                {
-                    try
-                    {
-                        // Don't use SendAsync here _cancellationToken is already cancelled.
-                        await _channelAdapter.SendPacketAsync(new MqttDisconnectPacket
-                        {
-                            ReasonCode = MqttDisconnectReasonCode.SessionTakenOver
-                        }, TimeSpan.Zero, CancellationToken.None).ConfigureAwait(false);
-                    }
-                    catch (Exception exception)
-                    {
-                        _logger.Error(exception, "Client '{0}': Error while sending DISCONNECT packet after takeover.", ClientId);
-                    }
-                }
 
                 _logger.Info("Client '{0}': Connection stopped.", ClientId);
-                
+
                 try
                 {
                     await _sessionsManager.CleanUpClient(ClientId, _channelAdapter, disconnectType);
@@ -279,7 +270,7 @@ namespace MQTTnet.Server
 
         void StopInternal()
         {
-            _cancellationToken.Cancel(false);
+            _cancellationToken.Cancel();
         }
 
         async Task EnqueueSubscribedRetainedMessagesAsync(ICollection<MqttTopicFilter> topicFilters)
@@ -454,7 +445,10 @@ namespace MQTTnet.Server
             }
             catch (Exception exception)
             {
-                if (exception is MqttCommunicationTimedOutException)
+                if (exception is OperationCanceledException)
+                {
+                }
+                else if (exception is MqttCommunicationTimedOutException)
                 {
                     _logger.Warning(exception, "Sending publish packet failed: Timeout (ClientId: {0}).", ClientId);
                 }
@@ -462,9 +456,6 @@ namespace MQTTnet.Server
                 {
                     _logger.Warning(exception, "Sending publish packet failed: Communication exception (ClientId: {0}).", ClientId);
                 }
-                else if (exception is OperationCanceledException)
-                {
-                }
                 else
                 {
                     _logger.Error(exception, "Sending publish packet failed (ClientId: {0}).", ClientId);
@@ -477,10 +468,7 @@ namespace MQTTnet.Server
                     Session.ApplicationMessagesQueue.Enqueue(queuedApplicationMessage);
                 }
 
-                if (!_cancellationToken.Token.IsCancellationRequested)
+                StopInternal();
-                {
-                    await StopAsync().ConfigureAwait(false);
-                }
             }
         }
 

Source/MQTTnet/Server/MqttClientSessionsManager.cs

@@ -19,9 +19,9 @@ namespace MQTTnet.Server
     {
         readonly AsyncQueue<MqttEnqueuedApplicationMessage> _messageQueue = new AsyncQueue<MqttEnqueuedApplicationMessage>();
 
-        readonly AsyncLock _createConnectionGate = new AsyncLock();
+        readonly object _createConnectionSyncRoot = new object();
-        readonly ConcurrentDictionary<string, MqttClientConnection> _connections = new ConcurrentDictionary<string, MqttClientConnection>();
+        readonly Dictionary<string, MqttClientConnection> _connections = new Dictionary<string, MqttClientConnection>();
-        readonly ConcurrentDictionary<string, MqttClientSession> _sessions = new ConcurrentDictionary<string, MqttClientSession>();
+        readonly Dictionary<string, MqttClientSession> _sessions = new Dictionary<string, MqttClientSession>();
 
         readonly IDictionary<object, object> _serverSessionItems = new ConcurrentDictionary<object, object>();
 
@@ -56,9 +56,16 @@ namespace MQTTnet.Server
             Task.Run(() => TryProcessQueuedApplicationMessagesAsync(_cancellationToken), _cancellationToken).Forget(_logger);
         }
 
-        public async Task StopAsync()
+        public async Task CloseAllConnectionsAsync()
         {
-            foreach (var connection in _connections.Values)
+            List<MqttClientConnection> connections;
+            lock (_connections)
+            {
+                connections = _connections.Values.ToList();
+                _connections.Clear();
+            }
+
+            foreach (var connection in connections)
             {
                 await connection.StopAsync().ConfigureAwait(false);
             }
@@ -66,7 +73,10 @@ namespace MQTTnet.Server
 
         public List<MqttClientConnection> GetConnections()
         {
-            return _connections.Values.ToList();
+            lock (_connections)
+            {
+                return _connections.Values.ToList();
+            }
         }
 
         public Task HandleClientConnectionAsync(IMqttChannelAdapter clientAdapter)
@@ -80,16 +90,19 @@ namespace MQTTnet.Server
         {
             var result = new List<IMqttClientStatus>();
 
-            foreach (var connection in _connections.Values)
+            lock (_connections)
             {
-                var clientStatus = new MqttClientStatus(connection);
+                foreach (var connection in _connections.Values)
-                connection.FillStatus(clientStatus);
+                {
+                    var clientStatus = new MqttClientStatus(connection);
+                    connection.FillStatus(clientStatus);
 
-                var sessionStatus = new MqttSessionStatus(connection.Session, this);
+                    var sessionStatus = new MqttSessionStatus(connection.Session, this);
-                connection.Session.FillStatus(sessionStatus);
+                    connection.Session.FillStatus(sessionStatus);
-                clientStatus.Session = sessionStatus;
+                    clientStatus.Session = sessionStatus;
 
-                result.Add(clientStatus);
+                    result.Add(clientStatus);
+                }
             }
 
             return Task.FromResult((IList<IMqttClientStatus>)result);
@@ -99,12 +112,15 @@ namespace MQTTnet.Server
         {
             var result = new List<IMqttSessionStatus>();
 
-            foreach (var session in _sessions.Values)
+            lock (_sessions)
             {
-                var sessionStatus = new MqttSessionStatus(session, this);
+                foreach (var session in _sessions.Values)
-                session.FillStatus(sessionStatus);
+                {
+                    var sessionStatus = new MqttSessionStatus(session, this);
+                    session.FillStatus(sessionStatus);
 
-                result.Add(sessionStatus);
+                    result.Add(sessionStatus);
+                }
             }
 
             return Task.FromResult((IList<IMqttSessionStatus>)result);
@@ -122,12 +138,7 @@ namespace MQTTnet.Server
             if (clientId == null) throw new ArgumentNullException(nameof(clientId));
             if (topicFilters == null) throw new ArgumentNullException(nameof(topicFilters));
 
-            if (!_sessions.TryGetValue(clientId, out var session))
+            return GetSession(clientId).SubscribeAsync(topicFilters);
-            {
-                throw new InvalidOperationException($"Client session '{clientId}' is unknown.");
-            }
-
-            return session.SubscribeAsync(topicFilters);
         }
 
         public Task UnsubscribeAsync(string clientId, IEnumerable<string> topicFilters)
@@ -135,25 +146,27 @@ namespace MQTTnet.Server
             if (clientId == null) throw new ArgumentNullException(nameof(clientId));
             if (topicFilters == null) throw new ArgumentNullException(nameof(topicFilters));
 
-            if (!_sessions.TryGetValue(clientId, out var session))
+            return GetSession(clientId).UnsubscribeAsync(topicFilters);
-            {
-                throw new InvalidOperationException($"Client session '{clientId}' is unknown.");
-            }
-
-            return session.UnsubscribeAsync(topicFilters);
         }
 
         public async Task DeleteSessionAsync(string clientId)
         {
-            if (_connections.TryGetValue(clientId, out var connection))
+            MqttClientConnection connection;
+            lock (_connections)
+            {
+                _connections.TryGetValue(clientId, out connection);
+            }
+            
+            lock (_sessions)
+            {
+                _sessions.Remove(clientId);
+            }
+
+            if (connection != null)
             {
                 await connection.StopAsync().ConfigureAwait(false);
             }
-
+            
-            if (_sessions.TryRemove(clientId, out _))
-            {
-            }
-
             _logger.Verbose("Session for client '{0}' deleted.", clientId);
         }
 
@@ -228,19 +241,22 @@ namespace MQTTnet.Server
 
                 var deliveryCount = 0;
 
-                foreach (var clientSession in _sessions.Values)
+                lock (_sessions)
                 {
-                    var isSubscribed = clientSession.EnqueueApplicationMessage(
+                    foreach (var clientSession in _sessions.Values)
-                        applicationMessage,
-                        sender?.ClientId,
-                        false);
-
-                    if (isSubscribed)
                     {
-                        deliveryCount++;
+                        var isSubscribed = clientSession.EnqueueApplicationMessage(
+                            applicationMessage,
+                            sender?.ClientId,
+                            false);
+
+                        if (isSubscribed)
+                        {
+                            deliveryCount++;
+                        }
                     }
                 }
-
+                
                 if (deliveryCount == 0)
                 {
                     var undeliveredMessageInterceptor = _options.UndeliveredMessageInterceptor;
@@ -295,7 +311,7 @@ namespace MQTTnet.Server
                     return;
                 }
 
-                var connection = await CreateClientConnectionAsync(connectPacket, connectionValidatorContext, channelAdapter).ConfigureAwait(false);
+                var connection = CreateClientConnection(connectPacket, connectionValidatorContext, channelAdapter);
                 await _eventDispatcher.SafeNotifyClientConnectedAsync(connectPacket.ClientId).ConfigureAwait(false);
                 await connection.RunAsync().ConfigureAwait(false);
             }
@@ -315,7 +331,10 @@ namespace MQTTnet.Server
                 // in case it is a takeover _connections already contains the new connection
                 if (disconnectType != MqttClientDisconnectType.Takeover)
                 {
-                    _connections.TryRemove(clientId, out _);
+                    lock (_connections)
+                    {
+                        _connections.Remove(clientId);
+                    }
 
                     if (!_options.EnablePersistentSessions)
                     {
@@ -360,52 +379,52 @@ namespace MQTTnet.Server
             return context;
         }
 
-        async Task<MqttClientConnection> CreateClientConnectionAsync(MqttConnectPacket connectPacket, MqttConnectionValidatorContext connectionValidatorContext, IMqttChannelAdapter channelAdapter)
+        MqttClientConnection CreateClientConnection(MqttConnectPacket connectPacket, MqttConnectionValidatorContext connectionValidatorContext, IMqttChannelAdapter channelAdapter)
         {
-            MqttClientConnection connection;
+            lock (_createConnectionSyncRoot)
-            MqttClientConnection existingConnection = null;
-
-            using (await _createConnectionGate.WaitAsync(_cancellationToken).ConfigureAwait(false))
             {
-                var session = _sessions.AddOrUpdate(
+                MqttClientSession session;
-                    connectPacket.ClientId,
+                lock (_sessions)
-                    key =>
+                {
+                    if (!_sessions.TryGetValue(connectPacket.ClientId, out session))
                     {
-                        _logger.Verbose("Created a new session for client '{0}'.", key);
+                        _logger.Verbose("Created a new session for client '{0}'.", connectPacket.ClientId);
-                        return new MqttClientSession(key, connectionValidatorContext.SessionItems, _eventDispatcher, _options, _retainedMessagesManager, _rootLogger);
+                        session = CreateSession(connectPacket.ClientId, connectionValidatorContext);
-                    },
+                    }
-                    (key, existingSession) =>
+                    else
                     {
                         if (connectPacket.CleanSession)
                         {
                             _logger.Verbose("Deleting existing session of client '{0}'.", connectPacket.ClientId);
-                            return new MqttClientSession(key, connectionValidatorContext.SessionItems, _eventDispatcher, _options, _retainedMessagesManager, _rootLogger);
+                            session = CreateSession(connectPacket.ClientId, connectionValidatorContext);
                         }
+                        else
+                        {
+                            _logger.Verbose("Reusing existing session of client '{0}'.", connectPacket.ClientId);
+                        }
+                    }
 
-                        _logger.Verbose("Reusing existing session of client '{0}'.", connectPacket.ClientId);
+                    _sessions[connectPacket.ClientId] = session;
-                        return existingSession;
+                }
-                    });
 
-                connection = new MqttClientConnection(connectPacket, channelAdapter, session, connectionValidatorContext, _options, this, _retainedMessagesManager, _rootLogger);
+                MqttClientConnection existingConnection;
+                MqttClientConnection connection;
+                lock (_connections)
+                {
+                    _connections.TryGetValue(connectPacket.ClientId, out existingConnection);
+                    connection = CreateConnection(connectPacket, channelAdapter, session, connectionValidatorContext);
 
-                _connections.AddOrUpdate(
+                    _connections[connectPacket.ClientId] = connection;
-                    connectPacket.ClientId,
+                }
-                    key => connection,
+
-                    (key, tempExistingConnection) =>
+                if (existingConnection != null)
-                    {
+                {
-                        existingConnection = tempExistingConnection;
+                    existingConnection.IsTakeOver = true;
-                        return connection;
+                    existingConnection.StopAsync().GetAwaiter().GetResult();
-                    });
+                }
+
+                return connection;
             }
-
-            // Disconnect the client outside of the lock so that new clients can still connect while a single
-            // one is being disconnected.
-            if (existingConnection != null)
-            {
-                await existingConnection.StopAsync(true).ConfigureAwait(false);
-            }
-
-            return connection;
         }
 
         async Task<MqttApplicationMessageInterceptorContext> InterceptApplicationMessageAsync(MqttClientConnection senderConnection, MqttApplicationMessage applicationMessage)
@@ -447,5 +466,42 @@ namespace MQTTnet.Server
                 _logger.Error(exception, "Error while disconnecting client channel.");
             }
         }
+
+        MqttClientSession GetSession(string clientId)
+        {
+            lock (_sessions)
+            {
+                if (!_sessions.TryGetValue(clientId, out var session))
+                {
+                    throw new InvalidOperationException($"Client session '{clientId}' is unknown.");
+                }
+
+                return session;
+            }
+        }
+
+        MqttClientConnection CreateConnection(MqttConnectPacket connectPacket, IMqttChannelAdapter channelAdapter, MqttClientSession session, MqttConnectionValidatorContext connectionValidatorContext)
+        {
+            return new MqttClientConnection(
+                connectPacket,
+                channelAdapter,
+                session,
+                connectionValidatorContext,
+                _options,
+                this,
+                _retainedMessagesManager,
+                _rootLogger);
+        }
+
+        MqttClientSession CreateSession(string clientId, MqttConnectionValidatorContext connectionValidatorContext)
+        {
+            return new MqttClientSession(
+                clientId,
+                connectionValidatorContext.SessionItems,
+                _eventDispatcher,
+                _options,
+                _retainedMessagesManager,
+                _rootLogger);
+        }
     }
 }

Source/MQTTnet/Server/MqttServer.cs

@@ -186,10 +186,10 @@ namespace MQTTnet.Server
                     return;
                 }
 
-                await _clientSessionsManager.StopAsync().ConfigureAwait(false);
-
                 _cancellationTokenSource.Cancel(false);
 
+                await _clientSessionsManager.CloseAllConnectionsAsync().ConfigureAwait(false);
+                
                 foreach (var adapter in _adapters)
                 {
                     adapter.ClientHandler = null;