From f9630ef3473df8fc1c248120364078a60154e12c Mon Sep 17 00:00:00 2001 From: Christian Kratky Date: Fri, 30 Oct 2020 00:19:30 +0100 Subject: [PATCH] Improve internal session/connection management and connection takeover. --- Source/MQTTnet/Server/MqttClientConnection.cs | 86 +++---- .../Server/MqttClientSessionsManager.cs | 216 +++++++++++------- Source/MQTTnet/Server/MqttServer.cs | 4 +- 3 files changed, 175 insertions(+), 131 deletions(-) diff --git a/Source/MQTTnet/Server/MqttClientConnection.cs b/Source/MQTTnet/Server/MqttClientConnection.cs index fb7c4714..c346e285 100644 --- a/Source/MQTTnet/Server/MqttClientConnection.cs +++ b/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; - var task = _packageReceiverTask; + if (IsTakeOver) + { + // 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) - { - return task; - } - - return PlatformAbstractionLayer.CompletedTask; + await _packageReceiverTask; } 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(); - - _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); - } - } + _packetDispatcher.Cancel(); _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 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) - { - await StopAsync().ConfigureAwait(false); - } + StopInternal(); } } diff --git a/Source/MQTTnet/Server/MqttClientSessionsManager.cs b/Source/MQTTnet/Server/MqttClientSessionsManager.cs index 05c614e7..3e6d95f0 100644 --- a/Source/MQTTnet/Server/MqttClientSessionsManager.cs +++ b/Source/MQTTnet/Server/MqttClientSessionsManager.cs @@ -19,9 +19,9 @@ namespace MQTTnet.Server { readonly AsyncQueue _messageQueue = new AsyncQueue(); - readonly AsyncLock _createConnectionGate = new AsyncLock(); - readonly ConcurrentDictionary _connections = new ConcurrentDictionary(); - readonly ConcurrentDictionary _sessions = new ConcurrentDictionary(); + readonly object _createConnectionSyncRoot = new object(); + readonly Dictionary _connections = new Dictionary(); + readonly Dictionary _sessions = new Dictionary(); readonly IDictionary _serverSessionItems = new ConcurrentDictionary(); @@ -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 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 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(); - foreach (var connection in _connections.Values) + lock (_connections) { - var clientStatus = new MqttClientStatus(connection); - connection.FillStatus(clientStatus); + foreach (var connection in _connections.Values) + { + var clientStatus = new MqttClientStatus(connection); + connection.FillStatus(clientStatus); - var sessionStatus = new MqttSessionStatus(connection.Session, this); - connection.Session.FillStatus(sessionStatus); - clientStatus.Session = sessionStatus; + var sessionStatus = new MqttSessionStatus(connection.Session, this); + connection.Session.FillStatus(sessionStatus); + clientStatus.Session = sessionStatus; - result.Add(clientStatus); + result.Add(clientStatus); + } } return Task.FromResult((IList)result); @@ -99,12 +112,15 @@ namespace MQTTnet.Server { var result = new List(); - foreach (var session in _sessions.Values) + lock (_sessions) { - var sessionStatus = new MqttSessionStatus(session, this); - session.FillStatus(sessionStatus); + foreach (var session in _sessions.Values) + { + var sessionStatus = new MqttSessionStatus(session, this); + session.FillStatus(sessionStatus); - result.Add(sessionStatus); + result.Add(sessionStatus); + } } return Task.FromResult((IList)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)) - { - throw new InvalidOperationException($"Client session '{clientId}' is unknown."); - } - - return session.SubscribeAsync(topicFilters); + return GetSession(clientId).SubscribeAsync(topicFilters); } public Task UnsubscribeAsync(string clientId, IEnumerable 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)) - { - throw new InvalidOperationException($"Client session '{clientId}' is unknown."); - } - - return session.UnsubscribeAsync(topicFilters); + return GetSession(clientId).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( - applicationMessage, - sender?.ClientId, - false); - - if (isSubscribed) + foreach (var clientSession in _sessions.Values) { - 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 CreateClientConnectionAsync(MqttConnectPacket connectPacket, MqttConnectionValidatorContext connectionValidatorContext, IMqttChannelAdapter channelAdapter) + MqttClientConnection CreateClientConnection(MqttConnectPacket connectPacket, MqttConnectionValidatorContext connectionValidatorContext, IMqttChannelAdapter channelAdapter) { - MqttClientConnection connection; - MqttClientConnection existingConnection = null; - - using (await _createConnectionGate.WaitAsync(_cancellationToken).ConfigureAwait(false)) + lock (_createConnectionSyncRoot) { - var session = _sessions.AddOrUpdate( - connectPacket.ClientId, - key => + MqttClientSession session; + lock (_sessions) + { + if (!_sessions.TryGetValue(connectPacket.ClientId, out session)) { - _logger.Verbose("Created a new session for client '{0}'.", key); - return new MqttClientSession(key, connectionValidatorContext.SessionItems, _eventDispatcher, _options, _retainedMessagesManager, _rootLogger); - }, - (key, existingSession) => + _logger.Verbose("Created a new session for client '{0}'.", connectPacket.ClientId); + session = CreateSession(connectPacket.ClientId, connectionValidatorContext); + } + 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); - return existingSession; - }); + _sessions[connectPacket.ClientId] = session; + } - 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( - connectPacket.ClientId, - key => connection, - (key, tempExistingConnection) => - { - existingConnection = tempExistingConnection; - return connection; - }); + _connections[connectPacket.ClientId] = connection; + } + + if (existingConnection != null) + { + existingConnection.IsTakeOver = true; + 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 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); + } } } \ No newline at end of file diff --git a/Source/MQTTnet/Server/MqttServer.cs b/Source/MQTTnet/Server/MqttServer.cs index cf22f8d2..5eed8674 100644 --- a/Source/MQTTnet/Server/MqttServer.cs +++ b/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;