Merge pull request #547 from vbBerni/develop

Fix Managed Client should dispose inner mqtt client #527

Source/MQTTnet.Extensions.ManagedClient/ManagedMqttClient.cs

@@ -27,9 +27,11 @@ namespace MQTTnet.Extensions.ManagedClient
 
         private CancellationTokenSource _connectionCancellationToken;
         private CancellationTokenSource _publishingCancellationToken;
+        private Task _maintainConnectionTask;
 
         private ManagedMqttClientStorageManager _storageManager;
 
+        private bool _disposed = false;
         private bool _subscriptionsNotPushed;
 
         public ManagedMqttClient(IMqttClient mqttClient, IMqttNetChildLogger logger)
@@ -73,6 +75,8 @@ namespace MQTTnet.Extensions.ManagedClient
 
         public async Task StartAsync(IManagedMqttClientOptions options)
         {
+            ThrowIfDisposed();
+
             if (options == null) throw new ArgumentNullException(nameof(options));
             if (options.ClientOptions == null) throw new ArgumentException("The client options are not set.", nameof(options));
 
@@ -81,7 +85,7 @@ namespace MQTTnet.Extensions.ManagedClient
                 throw new NotSupportedException("The managed client does not support existing sessions.");
             }
 
-            if (_connectionCancellationToken != null) throw new InvalidOperationException("The managed client is already started.");
+            if (!_maintainConnectionTask?.IsCompleted ?? false) throw new InvalidOperationException("The managed client is already started.");
 
             Options = options;
 
@@ -99,24 +103,32 @@ namespace MQTTnet.Extensions.ManagedClient
             _connectionCancellationToken = new CancellationTokenSource();
 
 #pragma warning disable CS4014 // Because this call is not awaited, execution of the current method continues before the call is completed
-            Task.Run(() => MaintainConnectionAsync(_connectionCancellationToken.Token), _connectionCancellationToken.Token);
+            _maintainConnectionTask = Task.Run(() => MaintainConnectionAsync(_connectionCancellationToken.Token), _connectionCancellationToken.Token);
 #pragma warning restore CS4014 // Because this call is not awaited, execution of the current method continues before the call is completed
 
             _logger.Info("Started");
         }
 
-        public Task StopAsync()
+        public async Task StopAsync()
         {
+            ThrowIfDisposed();
+
             StopPublishing();
             StopMaintainingConnection();
 
             _messageQueue.Clear();
 
-            return Task.FromResult(0);
+            if (_maintainConnectionTask != null)
+            {
+                await Task.WhenAny(_maintainConnectionTask);
+                _maintainConnectionTask = null;
+            }
         }
 
         public async Task<MqttClientPublishResult> PublishAsync(MqttApplicationMessage applicationMessage, CancellationToken cancellationToken)
         {
+            ThrowIfDisposed();
+
             if (applicationMessage == null) throw new ArgumentNullException(nameof(applicationMessage));
 
             await PublishAsync(new ManagedMqttApplicationMessageBuilder().WithApplicationMessage(applicationMessage).Build()).ConfigureAwait(false);
@@ -125,6 +137,8 @@ namespace MQTTnet.Extensions.ManagedClient
 
         public async Task PublishAsync(ManagedMqttApplicationMessage applicationMessage)
         {
+            ThrowIfDisposed();
+
             if (applicationMessage == null) throw new ArgumentNullException(nameof(applicationMessage));
 
             ManagedMqttApplicationMessage removedMessage = null;
@@ -180,6 +194,8 @@ namespace MQTTnet.Extensions.ManagedClient
 
         public Task SubscribeAsync(IEnumerable<TopicFilter> topicFilters)
         {
+            ThrowIfDisposed();
+
             if (topicFilters == null) throw new ArgumentNullException(nameof(topicFilters));
 
             lock (_subscriptions)
@@ -196,6 +212,8 @@ namespace MQTTnet.Extensions.ManagedClient
 
         public Task UnsubscribeAsync(IEnumerable<string> topics)
         {
+            ThrowIfDisposed();
+
             if (topics == null) throw new ArgumentNullException(nameof(topics));
 
             lock (_subscriptions)
@@ -215,8 +233,31 @@ namespace MQTTnet.Extensions.ManagedClient
 
         public void Dispose()
         {
-            _connectionCancellationToken?.Dispose();
-            _publishingCancellationToken?.Dispose();
+            if (_disposed)
+            {
+                return;
+            }
+
+            _disposed = true;
+
+            StopPublishing();
+            StopMaintainingConnection();
+
+            if (_maintainConnectionTask != null)
+            {
+                Task.WaitAny(_maintainConnectionTask);
+                _maintainConnectionTask = null;
+            }
+
+            _mqttClient.Dispose();
+        }
+
+        private void ThrowIfDisposed()
+        {
+            if (_disposed)
+            {
+                throw new ObjectDisposedException(nameof(ManagedMqttClient));
+            }
         }
 
         private async Task MaintainConnectionAsync(CancellationToken cancellationToken)
@@ -237,16 +278,19 @@ namespace MQTTnet.Extensions.ManagedClient
             }
             finally
             {
-                try
+                if (!_disposed)
                 {
-                    await _mqttClient.DisconnectAsync().ConfigureAwait(false);
-                }
-                catch (Exception exception)
-                {
-                    _logger.Error(exception, "Error while disconnecting.");
-                }
+                    try
+                    {
+                        await _mqttClient.DisconnectAsync().ConfigureAwait(false);
+                    }
+                    catch (Exception exception)
+                    {
+                        _logger.Error(exception, "Error while disconnecting.");
+                    }
 
-                _logger.Info("Stopped");
+                    _logger.Info("Stopped");
+                }
             }
         }
 
@@ -377,7 +421,7 @@ namespace MQTTnet.Extensions.ManagedClient
 
         private async Task SynchronizeSubscriptionsAsync()
         {
-            _logger.Info(nameof(ManagedMqttClient), "Synchronizing subscriptions");
+            _logger.Info("Synchronizing subscriptions");
 
             List<TopicFilter> subscriptions;
             HashSet<string> unsubscriptions;
@@ -473,4 +517,4 @@ namespace MQTTnet.Extensions.ManagedClient
             _connectionCancellationToken = null;
         }
     }
-}
+}

Tests/MQTTnet.Core.Tests/ManagedMqttClientTests.cs

@@ -1,7 +1,12 @@
-using System.Threading.Tasks;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
+using MQTTnet.Client;
+using MQTTnet.Client.Options;
+using MQTTnet.Diagnostics;
 using MQTTnet.Extensions.ManagedClient;
 using MQTTnet.Server;
+using MQTTnet.Tests.Mockups;
 
 namespace MQTTnet.Tests
 {
@@ -40,5 +45,67 @@ namespace MQTTnet.Tests
                 await managedClient.StopAsync();
             }
         }
+
+        [TestMethod]
+        public async Task ManagedClients_Will_Message_Send()
+        {
+            using (var testEnvironment = new TestEnvironment())
+            {
+                var receivedMessagesCount = 0;
+
+                var factory = new MqttFactory();
+
+                await testEnvironment.StartServerAsync();
+
+                var willMessage = new MqttApplicationMessageBuilder().WithTopic("My/last/will").WithAtMostOnceQoS().Build();
+                var clientOptions = new MqttClientOptionsBuilder()
+                    .WithTcpServer("localhost", testEnvironment.ServerPort)
+                    .WithWillMessage(willMessage);
+                var dyingClient = testEnvironment.CreateClient();
+                var dyingManagedClient = new ManagedMqttClient(dyingClient, testEnvironment.ClientLogger.CreateChildLogger());
+                await dyingManagedClient.StartAsync(new ManagedMqttClientOptionsBuilder()
+                    .WithClientOptions(clientOptions)
+                    .Build());
+
+                var recievingClient = await testEnvironment.ConnectClientAsync();
+                await recievingClient.SubscribeAsync("My/last/will");
+
+                recievingClient.UseReceivedApplicationMessageHandler(context => Interlocked.Increment(ref receivedMessagesCount));
+
+                dyingManagedClient.Dispose();
+
+                await Task.Delay(1000);
+
+                Assert.AreEqual(1, receivedMessagesCount);
+            }
+        }
+
+        [TestMethod]
+        public async Task Start_Stop()
+        {
+            using (var testEnvironment = new TestEnvironment())
+            {
+                var factory = new MqttFactory();
+
+                var server = await testEnvironment.StartServerAsync();
+
+                var managedClient = new ManagedMqttClient(testEnvironment.CreateClient(), new MqttNetLogger().CreateChildLogger());
+                var clientOptions = new MqttClientOptionsBuilder()
+                    .WithTcpServer("localhost", testEnvironment.ServerPort);
+
+                TaskCompletionSource<bool> connected = new TaskCompletionSource<bool>();
+                managedClient.Connected += (s, e) => { connected.SetResult(true); };
+
+                await managedClient.StartAsync(new ManagedMqttClientOptionsBuilder()
+                    .WithClientOptions(clientOptions)
+                    .Build());
+
+                await connected.Task;
+
+                await managedClient.StopAsync();
+
+                Assert.AreEqual(0, (await server.GetClientStatusAsync()).Count);
+            }
+        }
     }
-}
+}

Tests/MQTTnet.Core.Tests/Mockups/TestEnvironment.cs

@@ -29,6 +29,10 @@ namespace MQTTnet.Tests.Mockups
 
         public int ServerPort { get; set; } = 1888;
 
+        public IMqttNetLogger ServerLogger => _serverLogger;
+
+        public IMqttNetLogger ClientLogger => _clientLogger;
+
         public TestEnvironment()
         {
             _serverLogger.LogMessagePublished += (s, e) =>
@@ -56,7 +60,9 @@ namespace MQTTnet.Tests.Mockups
 
         public IMqttClient CreateClient()
         {
-            return _mqttFactory.CreateMqttClient(_clientLogger);
+            var client = _mqttFactory.CreateMqttClient(_clientLogger);
+            _clients.Add(client);
+            return client;
         }
 
         public Task<IMqttServer> StartServerAsync()
@@ -84,10 +90,9 @@ namespace MQTTnet.Tests.Mockups
 
         public async Task<IMqttClient> ConnectClientAsync(MqttClientOptionsBuilder options)
         {
-            var client = _mqttFactory.CreateMqttClient(_clientLogger);
+            var client = CreateClient();
             await client.ConnectAsync(options.WithTcpServer("localhost", ServerPort).Build());
 
-            _clients.Add(client);
             return client;
         }
 
@@ -135,4 +140,4 @@ namespace MQTTnet.Tests.Mockups
             }
         }
     }
-}
+}

Tests/MQTTnet.Core.Tests/MqttFactoryTests.cs

@@ -64,9 +64,9 @@ namespace MQTTnet.Tests
             }
             finally
             {
-                MqttNetGlobalLogger.LogMessagePublished -= globalLog;
-
                 await managedClient.StopAsync();
+
+                MqttNetGlobalLogger.LogMessagePublished -= globalLog;
             }
 
             Assert.IsFalse(invalidLogIdOccured);