Fix several deadlocks

2026-10-19 02:53:06 +08:00 · 2018-04-24 20:54:18 +02:00
parent ffb1645f14
commit 9e90111b37
20 changed files with 428 additions and 288 deletions
@@ -3,11 +3,15 @@ using MQTTnet.Serializer;

 namespace MQTTnet.Server
 {
+    // TODO: Rename to "RegisteredClient"
+    // TODO: Add IsConnected
+    // TODO: Add interface
+
    public class ConnectedMqttClient
    {
        public string ClientId { get; set; }

-        public MqttProtocolVersion ProtocolVersion { get; set; }
+        public MqttProtocolVersion? ProtocolVersion { get; set; }

        public TimeSpan LastPacketReceived { get; set; }

@@ -93,7 +93,7 @@ namespace MQTTnet.Server
                    return;
                }

-                await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, new[] { packet }).ConfigureAwait(false);
+                await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, new[] { packet }, cancellationToken).ConfigureAwait(false);

                _logger.Verbose<MqttClientPendingMessagesQueue>("Enqueued packet sent (ClientId: {0}).", _clientSession.ClientId);
            }
@@ -16,11 +16,11 @@ namespace MQTTnet.Server
    public sealed class MqttClientSession : IDisposable
    {
        private readonly MqttPacketIdentifierProvider _packetIdentifierProvider = new MqttPacketIdentifierProvider();
-        private readonly IMqttServerOptions _options;
-        private readonly IMqttNetLogger _logger;
        private readonly MqttRetainedMessagesManager _retainedMessagesManager;
+        private readonly IMqttNetLogger _logger;
+        private readonly IMqttServerOptions _options;
+        private readonly MqttClientSessionsManager _sessionsManager;

-        private IMqttChannelAdapter _adapter;
        private CancellationTokenSource _cancellationTokenSource;
        private MqttApplicationMessage _willMessage;
        private bool _wasCleanDisconnect;
@@ -28,22 +28,22 @@ namespace MQTTnet.Server
        public MqttClientSession(
            string clientId,
            IMqttServerOptions options,
+            MqttClientSessionsManager sessionsManager,
            MqttRetainedMessagesManager retainedMessagesManager,
            IMqttNetLogger logger)
        {
            _options = options ?? throw new ArgumentNullException(nameof(options));
+            _sessionsManager = sessionsManager;
            _retainedMessagesManager = retainedMessagesManager ?? throw new ArgumentNullException(nameof(retainedMessagesManager));
            _logger = logger ?? throw new ArgumentNullException(nameof(logger));

            ClientId = clientId;

            KeepAliveMonitor = new MqttClientKeepAliveMonitor(clientId, StopDueToKeepAliveTimeoutAsync, _logger);
-            SubscriptionsManager = new MqttClientSubscriptionsManager(_options, clientId);
+            SubscriptionsManager = new MqttClientSubscriptionsManager(clientId, _options, sessionsManager.Server);
            PendingMessagesQueue = new MqttClientPendingMessagesQueue(_options, this, _logger);
        }

-        public Func<MqttClientSession, MqttApplicationMessage, Task> ApplicationMessageReceivedCallback { get; set; }
-
        public MqttClientSubscriptionsManager SubscriptionsManager { get; }

        public MqttClientPendingMessagesQueue PendingMessagesQueue { get; }
@@ -52,9 +52,9 @@ namespace MQTTnet.Server

        public string ClientId { get; }

-        public MqttProtocolVersion? ProtocolVersion => _adapter?.PacketSerializer.ProtocolVersion;
+        public MqttProtocolVersion? ProtocolVersion { get; private set; }

-        public bool IsConnected => _adapter != null;
+        public bool IsConnected { get; private set; }

        public async Task<bool> RunAsync(MqttConnectPacket connectPacket, IMqttChannelAdapter adapter)
        {
@@ -63,66 +63,73 @@ namespace MQTTnet.Server

            try
            {
-                var cancellationTokenSource = new CancellationTokenSource();
+                _cancellationTokenSource = new CancellationTokenSource();

                _wasCleanDisconnect = false;
                _willMessage = connectPacket.WillMessage;
-                _adapter = adapter;
-                _cancellationTokenSource = cancellationTokenSource;

-                PendingMessagesQueue.Start(adapter, cancellationTokenSource.Token);
-                KeepAliveMonitor.Start(connectPacket.KeepAlivePeriod, cancellationTokenSource.Token);
+                IsConnected = true;
+                ProtocolVersion = adapter.PacketSerializer.ProtocolVersion;

-                await ReceivePacketsAsync(adapter, cancellationTokenSource.Token).ConfigureAwait(false);
+                PendingMessagesQueue.Start(adapter, _cancellationTokenSource.Token);
+                KeepAliveMonitor.Start(connectPacket.KeepAlivePeriod, _cancellationTokenSource.Token);
+
+                await ReceivePacketsAsync(adapter, _cancellationTokenSource.Token).ConfigureAwait(false);
            }
            catch (OperationCanceledException)
            {
            }
            catch (MqttCommunicationException exception)
            {
-                _logger.Warning<MqttClientSession>(exception, "Client '{0}': Communication exception while processing client packets.", ClientId);
+                _logger.Warning<MqttClientSession>(exception,
+                    "Client '{0}': Communication exception while processing client packets.", ClientId);
            }
            catch (Exception exception)
            {
-                _logger.Error<MqttClientSession>(exception, "Client '{0}': Unhandled exception while processing client packets.", ClientId);
+                _logger.Error<MqttClientSession>(exception,
+                    "Client '{0}': Unhandled exception while processing client packets.", ClientId);
+            }
+            finally
+            {
+                ProtocolVersion = null;
+                IsConnected = false;
+
+                _cancellationTokenSource?.Dispose();
+                _cancellationTokenSource = null;
            }

            return _wasCleanDisconnect;
        }

-        public async Task StopAsync(bool wasCleanDisconnect = false)
+        public Task StopAsync(bool wasCleanDisconnect = false)
        {
            try
            {
                if (_cancellationTokenSource == null)
                {
-                    return;
+                    return Task.FromResult(0);
                }

                _wasCleanDisconnect = wasCleanDisconnect;

                _cancellationTokenSource?.Cancel(false);
-
                PendingMessagesQueue.WaitForCompletion();
                KeepAliveMonitor.WaitForCompletion();

-                _cancellationTokenSource?.Dispose();
-                _cancellationTokenSource = null;
-
-                _adapter = null;
-
-                _logger.Info<MqttClientSession>("Client '{0}': Session stopped.", ClientId);
-            }
-            finally
-            {
                var willMessage = _willMessage;
                _willMessage = null; // clear willmessage so it is send just once

-                if (willMessage != null && !wasCleanDisconnect)
+                if (_willMessage != null && !wasCleanDisconnect)
                {
-                    await ApplicationMessageReceivedCallback(this, willMessage).ConfigureAwait(false);
+                    _sessionsManager.StartDispatchApplicationMessage(this, willMessage);
                }
            }
+            finally
+            {
+                _logger.Info<MqttClientSession>("Client '{0}': Session stopped.", ClientId);
+            }
+
+            return Task.FromResult(0);
        }

        public async Task EnqueueApplicationMessageAsync(MqttApplicationMessage applicationMessage)
@@ -170,8 +177,6 @@ namespace MQTTnet.Server

        public void Dispose()
        {
-            ApplicationMessageReceivedCallback = null;
-
            SubscriptionsManager?.Dispose();
            PendingMessagesQueue?.Dispose();

@@ -219,7 +224,7 @@ namespace MQTTnet.Server

            if (packet is MqttPingReqPacket)
            {
-                return adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, new[] { new MqttPingRespPacket() });
+                return adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, new[] { new MqttPingRespPacket() }, cancellationToken);
            }

            if (packet is MqttPubRelPacket pubRelPacket)
@@ -234,7 +239,7 @@ namespace MQTTnet.Server
                    PacketIdentifier = pubRecPacket.PacketIdentifier
                };

-                return adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, new[] { responsePacket });
+                return adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, new[] { responsePacket }, cancellationToken);
            }

            if (packet is MqttPubAckPacket || packet is MqttPubCompPacket)
@@ -267,10 +272,19 @@ namespace MQTTnet.Server
            return StopAsync();
        }

+        private async Task EnqueueSubscribedRetainedMessagesAsync(ICollection<TopicFilter> topicFilters)
+        {
+            var retainedMessages = await _retainedMessagesManager.GetSubscribedMessagesAsync(topicFilters);
+            foreach (var applicationMessage in retainedMessages)
+            {
+                await EnqueueApplicationMessageAsync(applicationMessage).ConfigureAwait(false);
+            }
+        }
+
        private async Task HandleIncomingSubscribePacketAsync(IMqttChannelAdapter adapter, MqttSubscribePacket subscribePacket, CancellationToken cancellationToken)
        {
            var subscribeResult = await SubscriptionsManager.SubscribeAsync(subscribePacket).ConfigureAwait(false);
-            await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, new[] { subscribeResult.ResponsePacket }).ConfigureAwait(false);
+            await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, new[] { subscribeResult.ResponsePacket }, cancellationToken).ConfigureAwait(false);

            if (subscribeResult.CloseConnection)
            {
@@ -283,16 +297,7 @@ namespace MQTTnet.Server
        private async Task HandleIncomingUnsubscribePacketAsync(IMqttChannelAdapter adapter, MqttUnsubscribePacket unsubscribePacket, CancellationToken cancellationToken)
        {
            var unsubscribeResult = await SubscriptionsManager.UnsubscribeAsync(unsubscribePacket).ConfigureAwait(false);
-            await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, new[] { unsubscribeResult });
-        }
-
-        private async Task EnqueueSubscribedRetainedMessagesAsync(ICollection<TopicFilter> topicFilters)
-        {
-            var retainedMessages = await _retainedMessagesManager.GetSubscribedMessagesAsync(topicFilters);
-            foreach (var applicationMessage in retainedMessages)
-            {
-                await EnqueueApplicationMessageAsync(applicationMessage).ConfigureAwait(false);
-            }
+            await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, new[] { unsubscribeResult }, cancellationToken);
        }

        private Task HandleIncomingPublishPacketAsync(IMqttChannelAdapter adapter, MqttPublishPacket publishPacket, CancellationToken cancellationToken)
@@ -303,7 +308,8 @@ namespace MQTTnet.Server
            {
                case MqttQualityOfServiceLevel.AtMostOnce:
                    {
-                        return ApplicationMessageReceivedCallback?.Invoke(this, applicationMessage);
+                        _sessionsManager.StartDispatchApplicationMessage(this, applicationMessage);
+                        return Task.FromResult(0);
                    }
                case MqttQualityOfServiceLevel.AtLeastOnce:
                    {
@@ -322,25 +328,25 @@ namespace MQTTnet.Server

        private async Task HandleIncomingPublishPacketWithQoS1(IMqttChannelAdapter adapter, MqttApplicationMessage applicationMessage, MqttPublishPacket publishPacket, CancellationToken cancellationToken)
        {
-            await ApplicationMessageReceivedCallback(this, applicationMessage).ConfigureAwait(false);
+            _sessionsManager.StartDispatchApplicationMessage(this, applicationMessage);

            var response = new MqttPubAckPacket { PacketIdentifier = publishPacket.PacketIdentifier };
-            await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, new[] { response }).ConfigureAwait(false);
+            await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, new[] { response }, cancellationToken).ConfigureAwait(false);
        }

        private async Task HandleIncomingPublishPacketWithQoS2(IMqttChannelAdapter adapter, MqttApplicationMessage applicationMessage, MqttPublishPacket publishPacket, CancellationToken cancellationToken)
        {
            // QoS 2 is implement as method "B" [4.3.3 QoS 2: Exactly once delivery]
-            await ApplicationMessageReceivedCallback(this, applicationMessage).ConfigureAwait(false);
+            _sessionsManager.StartDispatchApplicationMessage(this, applicationMessage);

            var response = new MqttPubRecPacket { PacketIdentifier = publishPacket.PacketIdentifier };
-            await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, new[] { response }).ConfigureAwait(false);
+            await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, new[] { response }, cancellationToken).ConfigureAwait(false);
        }

        private Task HandleIncomingPubRelPacketAsync(IMqttChannelAdapter adapter, MqttPubRelPacket pubRelPacket, CancellationToken cancellationToken)
        {
            var response = new MqttPubCompPacket { PacketIdentifier = pubRelPacket.PacketIdentifier };
-            return adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, new[] { response });
+            return adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, new[] { response }, cancellationToken);
        }
    }
 }
@@ -8,38 +8,34 @@ using MQTTnet.Diagnostics;
 using MQTTnet.Exceptions;
 using MQTTnet.Packets;
 using MQTTnet.Protocol;
-using MQTTnet.Serializer;

 namespace MQTTnet.Server
 {
    public sealed class MqttClientSessionsManager : IDisposable
    {
        private readonly Dictionary<string, MqttClientSession> _sessions = new Dictionary<string, MqttClientSession>();
-        private readonly SemaphoreSlim _semaphore = new SemaphoreSlim(1, 1);
+        private readonly SemaphoreSlim _sessionsLock = new SemaphoreSlim(1, 1);

-        private readonly IMqttServerOptions _options;
        private readonly MqttRetainedMessagesManager _retainedMessagesManager;
+        private readonly IMqttServerOptions _options;
        private readonly IMqttNetLogger _logger;

-        public MqttClientSessionsManager(IMqttServerOptions options, MqttRetainedMessagesManager retainedMessagesManager, IMqttNetLogger logger)
+        public MqttClientSessionsManager(IMqttServerOptions options, MqttServer server, MqttRetainedMessagesManager retainedMessagesManager, IMqttNetLogger logger)
        {
            _logger = logger ?? throw new ArgumentNullException(nameof(logger));
            _options = options ?? throw new ArgumentNullException(nameof(options));
+            Server = server ?? throw new ArgumentNullException(nameof(server));
            _retainedMessagesManager = retainedMessagesManager ?? throw new ArgumentNullException(nameof(retainedMessagesManager));
        }

-        public Action<ConnectedMqttClient> ClientConnectedCallback { get; set; }
-        public Action<ConnectedMqttClient, bool> ClientDisconnectedCallback { get; set; }
-        public Action<string, TopicFilter> ClientSubscribedTopicCallback { get; set; }
-        public Action<string, string> ClientUnsubscribedTopicCallback { get; set; }
-        public Action<string, MqttApplicationMessage> ApplicationMessageReceivedCallback { get; set; }
+        public MqttServer Server { get; }

        public async Task RunSessionAsync(IMqttChannelAdapter clientAdapter, CancellationToken cancellationToken)
        {
            var clientId = string.Empty;
            var wasCleanDisconnect = false;
            MqttClientSession clientSession = null;
-            
+
            try
            {
                if (!(await clientAdapter.ReceivePacketAsync(_options.DefaultCommunicationTimeout, cancellationToken)
@@ -57,30 +53,30 @@ namespace MQTTnet.Server
                var connectReturnCode = ValidateConnection(connectPacket);
                if (connectReturnCode != MqttConnectReturnCode.ConnectionAccepted)
                {
-                    await clientAdapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, new[]
+                    await clientAdapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, new[]
                    {
                        new MqttConnAckPacket
                        {
                            ConnectReturnCode = connectReturnCode
                        }
-                    }).ConfigureAwait(false);
+                    }, cancellationToken).ConfigureAwait(false);

                    return;
                }

-                var result = await GetOrCreateClientSessionAsync(connectPacket).ConfigureAwait(false);
+                var result = await PrepareClientSessionAsync(connectPacket).ConfigureAwait(false);
                clientSession = result.Session;

-                await clientAdapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, new[]
+                await clientAdapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, new[]
                {
                    new MqttConnAckPacket
                    {
                        ConnectReturnCode = connectReturnCode,
                        IsSessionPresent = result.IsExistingSession
                    }
-                }).ConfigureAwait(false);
+                }, cancellationToken).ConfigureAwait(false);

-                ClientConnectedCallback?.Invoke(new ConnectedMqttClient
+                Server.OnClientConnected(new ConnectedMqttClient
                {
                    ClientId = clientId,
                    ProtocolVersion = clientAdapter.PacketSerializer.ProtocolVersion
@@ -99,15 +95,15 @@ namespace MQTTnet.Server
            {
                try
                {
-                    await clientAdapter.DisconnectAsync(_options.DefaultCommunicationTimeout).ConfigureAwait(false);
+                    await clientAdapter.DisconnectAsync(_options.DefaultCommunicationTimeout, CancellationToken.None).ConfigureAwait(false);
                    clientAdapter.Dispose();
                }
                catch (Exception exception)
                {
                    _logger.Error<MqttClientSessionsManager>(exception, exception.Message);
                }
-                
-                ClientDisconnectedCallback?.Invoke(new ConnectedMqttClient
+
+                Server.OnClientDisconnected(new ConnectedMqttClient
                {
                    ClientId = clientId,
                    ProtocolVersion = clientAdapter.PacketSerializer.ProtocolVersion,
@@ -119,7 +115,7 @@ namespace MQTTnet.Server

        public async Task StopAsync()
        {
-            await _semaphore.WaitAsync().ConfigureAwait(false);
+            await _sessionsLock.WaitAsync().ConfigureAwait(false);
            try
            {
                foreach (var session in _sessions)
@@ -131,19 +127,19 @@ namespace MQTTnet.Server
            }
            finally
            {
-                _semaphore.Release();
+                _sessionsLock.Release();
            }
        }

        public async Task<IList<ConnectedMqttClient>> GetConnectedClientsAsync()
        {
-            await _semaphore.WaitAsync().ConfigureAwait(false);
+            await _sessionsLock.WaitAsync().ConfigureAwait(false);
            try
            {
                return _sessions.Where(s => s.Value.IsConnected).Select(s => new ConnectedMqttClient
                {
                    ClientId = s.Value.ClientId,
-                    ProtocolVersion = s.Value.ProtocolVersion ?? MqttProtocolVersion.V311,
+                    ProtocolVersion = s.Value.ProtocolVersion,
                    LastPacketReceived = s.Value.KeepAliveMonitor.LastPacketReceived,
                    LastNonKeepAlivePacketReceived = s.Value.KeepAliveMonitor.LastNonKeepAlivePacketReceived,
                    PendingApplicationMessages = s.Value.PendingMessagesQueue.Count
@@ -151,49 +147,13 @@ namespace MQTTnet.Server
            }
            finally
            {
-                _semaphore.Release();
+                _sessionsLock.Release();
            }
        }

-        public async Task DispatchApplicationMessageAsync(MqttClientSession senderClientSession, MqttApplicationMessage applicationMessage)
+        public void StartDispatchApplicationMessage(MqttClientSession senderClientSession, MqttApplicationMessage applicationMessage)
        {
-            try
-            {
-                var interceptorContext = InterceptApplicationMessage(senderClientSession, applicationMessage);
-                if (interceptorContext.CloseConnection)
-                {
-                    await senderClientSession.StopAsync().ConfigureAwait(false);
-                }
-
-                if (interceptorContext.ApplicationMessage == null || !interceptorContext.AcceptPublish)
-                {
-                    return;
-                }
-
-                if (applicationMessage.Retain)
-                {
-                    await _retainedMessagesManager.HandleMessageAsync(senderClientSession?.ClientId, applicationMessage).ConfigureAwait(false);
-                }
-
-                ApplicationMessageReceivedCallback?.Invoke(senderClientSession?.ClientId, applicationMessage);
-            }
-            catch (Exception exception)
-            {
-                _logger.Error<MqttClientSessionsManager>(exception, "Error while processing application message");
-            }
-
-            await _semaphore.WaitAsync().ConfigureAwait(false);
-            try
-            {
-                foreach (var clientSession in _sessions.Values)
-                {
-                    await clientSession.EnqueueApplicationMessageAsync(applicationMessage).ConfigureAwait(false);
-                }
-            }
-            finally
-            {
-                _semaphore.Release();
-            }
+            Task.Run(() => DispatchApplicationMessageAsync(senderClientSession, applicationMessage));
        }

        public async Task SubscribeAsync(string clientId, IList<TopicFilter> topicFilters)
@@ -201,19 +161,19 @@ namespace MQTTnet.Server
            if (clientId == null) throw new ArgumentNullException(nameof(clientId));
            if (topicFilters == null) throw new ArgumentNullException(nameof(topicFilters));

-            await _semaphore.WaitAsync().ConfigureAwait(false);
+            await _sessionsLock.WaitAsync().ConfigureAwait(false);
            try
            {
                if (!_sessions.TryGetValue(clientId, out var session))
                {
-                    throw new InvalidOperationException($"Client session {clientId} is unknown.");
+                    throw new InvalidOperationException($"Client session '{clientId}' is unknown.");
                }

-                await session.SubscribeAsync(topicFilters);
+                await session.SubscribeAsync(topicFilters).ConfigureAwait(false);
            }
            finally
            {
-                _semaphore.Release();
+                _sessionsLock.Release();
            }
        }

@@ -222,36 +182,25 @@ namespace MQTTnet.Server
            if (clientId == null) throw new ArgumentNullException(nameof(clientId));
            if (topicFilters == null) throw new ArgumentNullException(nameof(topicFilters));

-            await _semaphore.WaitAsync().ConfigureAwait(false);
+            await _sessionsLock.WaitAsync().ConfigureAwait(false);
            try
            {
                if (!_sessions.TryGetValue(clientId, out var session))
                {
-                    throw new InvalidOperationException($"Client session {clientId} is unknown.");
+                    throw new InvalidOperationException($"Client session '{clientId}' is unknown.");
                }

-                await session.UnsubscribeAsync(topicFilters);
+                await session.UnsubscribeAsync(topicFilters).ConfigureAwait(false);
            }
            finally
            {
-                _semaphore.Release();
+                _sessionsLock.Release();
            }
        }

-        private MqttApplicationMessageInterceptorContext InterceptApplicationMessage(MqttClientSession senderClientSession, MqttApplicationMessage applicationMessage)
+        public void Dispose()
        {
-            var interceptorContext = new MqttApplicationMessageInterceptorContext(
-                senderClientSession?.ClientId,
-                applicationMessage);
-
-            var interceptor = _options.ApplicationMessageInterceptor;
-            if (interceptor == null)
-            {
-                return interceptorContext;
-            }
-            
-            interceptor(interceptorContext);
-            return interceptorContext;
+            _sessionsLock?.Dispose();
        }

        private MqttConnectReturnCode ValidateConnection(MqttConnectPacket connectPacket)
@@ -271,9 +220,9 @@ namespace MQTTnet.Server
            return context.ReturnCode;
        }

-        private async Task<GetOrCreateClientSessionResult> GetOrCreateClientSessionAsync(MqttConnectPacket connectPacket)
+        private async Task<GetOrCreateClientSessionResult> PrepareClientSessionAsync(MqttConnectPacket connectPacket)
        {
-            await _semaphore.WaitAsync().ConfigureAwait(false);
+            await _sessionsLock.WaitAsync().ConfigureAwait(false);
            try
            {
                var isSessionPresent = _sessions.TryGetValue(connectPacket.ClientId, out var clientSession);
@@ -300,14 +249,7 @@ namespace MQTTnet.Server
                {
                    isExistingSession = false;

-                    clientSession = new MqttClientSession(connectPacket.ClientId, _options, _retainedMessagesManager, _logger)
-                    {
-                        ApplicationMessageReceivedCallback = DispatchApplicationMessageAsync
-                    };
-
-                    clientSession.SubscriptionsManager.TopicSubscribedCallback = ClientSubscribedTopicCallback;
-                    clientSession.SubscriptionsManager.TopicUnsubscribedCallback = ClientUnsubscribedTopicCallback;
-
+                    clientSession = new MqttClientSession(connectPacket.ClientId, _options, this, _retainedMessagesManager, _logger);
                    _sessions[connectPacket.ClientId] = clientSession;

                    _logger.Verbose<MqttClientSessionsManager>("Created a new session for client '{0}'.", connectPacket.ClientId);
@@ -317,19 +259,65 @@ namespace MQTTnet.Server
            }
            finally
            {
-                _semaphore.Release();
+                _sessionsLock.Release();
            }
        }

-        public void Dispose()
+        private async Task DispatchApplicationMessageAsync(MqttClientSession senderClientSession, MqttApplicationMessage applicationMessage)
        {
-            ClientConnectedCallback = null;
-            ClientDisconnectedCallback = null;
-            ClientSubscribedTopicCallback = null;
-            ClientUnsubscribedTopicCallback = null;
-            ApplicationMessageReceivedCallback = null;
+            try
+            {
+                var interceptorContext = InterceptApplicationMessage(senderClientSession, applicationMessage);
+                if (interceptorContext.CloseConnection)
+                {
+                    await senderClientSession.StopAsync().ConfigureAwait(false);
+                }

-            _semaphore?.Dispose();
+                if (interceptorContext.ApplicationMessage == null || !interceptorContext.AcceptPublish)
+                {
+                    return;
+                }
+
+                if (applicationMessage.Retain)
+                {
+                    await _retainedMessagesManager.HandleMessageAsync(senderClientSession?.ClientId, applicationMessage).ConfigureAwait(false);
+                }
+
+                Server.OnApplicationMessageReceived(senderClientSession?.ClientId, applicationMessage);
+            }
+            catch (Exception exception)
+            {
+                _logger.Error<MqttClientSessionsManager>(exception, "Error while processing application message");
+            }
+
+            await _sessionsLock.WaitAsync().ConfigureAwait(false);
+            try
+            {
+                foreach (var clientSession in _sessions.Values)
+                {
+                    await clientSession.EnqueueApplicationMessageAsync(applicationMessage).ConfigureAwait(false);
+                }
+            }
+            finally
+            {
+                _sessionsLock.Release();
+            }
+        }
+
+        private MqttApplicationMessageInterceptorContext InterceptApplicationMessage(MqttClientSession senderClientSession, MqttApplicationMessage applicationMessage)
+        {
+            var interceptorContext = new MqttApplicationMessageInterceptorContext(
+                senderClientSession?.ClientId,
+                applicationMessage);
+
+            var interceptor = _options.ApplicationMessageInterceptor;
+            if (interceptor == null)
+            {
+                return interceptorContext;
+            }
+
+            interceptor(interceptorContext);
+            return interceptorContext;
        }
    }
 }
@@ -13,17 +13,16 @@ namespace MQTTnet.Server
        private readonly Dictionary<string, MqttQualityOfServiceLevel> _subscriptions = new Dictionary<string, MqttQualityOfServiceLevel>();
        private readonly SemaphoreSlim _semaphore = new SemaphoreSlim(1, 1);
        private readonly IMqttServerOptions _options;
+        private readonly MqttServer _server;
        private readonly string _clientId;

-        public MqttClientSubscriptionsManager(IMqttServerOptions options, string clientId)
+        public MqttClientSubscriptionsManager(string clientId, IMqttServerOptions options, MqttServer server)
        {
-            _options = options ?? throw new ArgumentNullException(nameof(options));
            _clientId = clientId ?? throw new ArgumentNullException(nameof(clientId));
+            _options = options ?? throw new ArgumentNullException(nameof(options));
+            _server = server;
        }

-        public Action<string, TopicFilter> TopicSubscribedCallback { get; set; }
-        public Action<string, string> TopicUnsubscribedCallback { get; set; }
-
        public async Task<MqttClientSubscribeResult> SubscribeAsync(MqttSubscribePacket subscribePacket)
        {
            if (subscribePacket == null) throw new ArgumentNullException(nameof(subscribePacket));
@@ -61,7 +60,7 @@ namespace MQTTnet.Server
                    if (interceptorContext.AcceptSubscription)
                    {
                        _subscriptions[topicFilter.Topic] = topicFilter.QualityOfServiceLevel;
-                        TopicSubscribedCallback?.Invoke(_clientId, topicFilter);
+                        _server.OnClientSubscribedTopic(_clientId, topicFilter);
                    }
                }
            }
@@ -83,7 +82,7 @@ namespace MQTTnet.Server
                foreach (var topicFilter in unsubscribePacket.TopicFilters)
                {
                    _subscriptions.Remove(topicFilter);
-                    TopicUnsubscribedCallback?.Invoke(_clientId, topicFilter);
+                    _server.OnClientUnsubscribedTopic(_clientId, topicFilter);
                }
            }
            finally
@@ -57,7 +57,7 @@ namespace MQTTnet.Server
            return _clientSessionsManager.UnsubscribeAsync(clientId, topicFilters);
        }

-        public async Task PublishAsync(IEnumerable<MqttApplicationMessage> applicationMessages)
+        public Task PublishAsync(IEnumerable<MqttApplicationMessage> applicationMessages)
        {
            if (applicationMessages == null) throw new ArgumentNullException(nameof(applicationMessages));

@@ -65,8 +65,10 @@ namespace MQTTnet.Server

            foreach (var applicationMessage in applicationMessages)
            {
-                await _clientSessionsManager.DispatchApplicationMessageAsync(null, applicationMessage);
+                _clientSessionsManager.StartDispatchApplicationMessage(null, applicationMessage);
            }
+
+            return Task.FromResult(0);
        }

        public async Task StartAsync(IMqttServerOptions options)
@@ -80,14 +82,7 @@ namespace MQTTnet.Server
            _retainedMessagesManager = new MqttRetainedMessagesManager(Options, _logger);
            await _retainedMessagesManager.LoadMessagesAsync();

-            _clientSessionsManager = new MqttClientSessionsManager(Options, _retainedMessagesManager, _logger)
-            {
-                ClientConnectedCallback = OnClientConnected,
-                ClientDisconnectedCallback = OnClientDisconnected,
-                ClientSubscribedTopicCallback = OnClientSubscribedTopic,
-                ClientUnsubscribedTopicCallback = OnClientUnsubscribedTopic,
-                ApplicationMessageReceivedCallback = OnApplicationMessageReceived
-            };
+            _clientSessionsManager = new MqttClientSessionsManager(Options, this, _retainedMessagesManager, _logger);

            foreach (var adapter in _adapters)
            {
@@ -132,29 +127,29 @@ namespace MQTTnet.Server
            }
        }

-        private void OnClientConnected(ConnectedMqttClient client)
+        internal void OnClientConnected(ConnectedMqttClient client)
        {
            _logger.Info<MqttServer>("Client '{0}': Connected.", client.ClientId);
            ClientConnected?.Invoke(this, new MqttClientConnectedEventArgs(client));
        }

-        private void OnClientDisconnected(ConnectedMqttClient client, bool wasCleanDisconnect)
+        internal void OnClientDisconnected(ConnectedMqttClient client, bool wasCleanDisconnect)
        {
            _logger.Info<MqttServer>("Client '{0}': Disconnected (clean={1}).", client.ClientId, wasCleanDisconnect);
            ClientDisconnected?.Invoke(this, new MqttClientDisconnectedEventArgs(client, wasCleanDisconnect));
        }

-        private void OnClientSubscribedTopic(string clientId, TopicFilter topicFilter)
+        internal void OnClientSubscribedTopic(string clientId, TopicFilter topicFilter)
        {
            ClientSubscribedTopic?.Invoke(this, new MqttClientSubscribedTopicEventArgs(clientId, topicFilter));
        }

-        private void OnClientUnsubscribedTopic(string clientId, string topicFilter)
+        internal void OnClientUnsubscribedTopic(string clientId, string topicFilter)
        {
            ClientUnsubscribedTopic?.Invoke(this, new MqttClientUnsubscribedTopicEventArgs(clientId, topicFilter));
        }

-        private void OnApplicationMessageReceived(string clientId, MqttApplicationMessage applicationMessage)
+        internal void OnApplicationMessageReceived(string clientId, MqttApplicationMessage applicationMessage)
        {
            ApplicationMessageReceived?.Invoke(this, new MqttApplicationMessageReceivedEventArgs(clientId, applicationMessage));
        }