fix: time out and cancel propagation in managed client (#1987)

* fix: time out and cancel propagation in managed client

* Update release notes

* Refactor code

---------

Co-authored-by: christian <6939810+chkr1011@users.noreply.github.com>

.github/workflows/ReleaseNotes.md

@@ -6,4 +6,5 @@
 * [Server] The server will no longer send _NoMatchingSubscribers_ when the actual subscription was non success (#1965, BREAKING CHANGE!).
 * [Server] Fixed broken support for _null_ in _AddServer_ method in ASP.NET integration (#1981).
 * [ManagedClient] Added a new event (SubscriptionsChangedAsync) which is fired when a subscription or unsubscription was made (#1894, thanks to @pdufrene).
+* [ManagedClient] Fixed race condition when server shuts down while subscribing (#1987, thanks to @marve).
 * [TopicTemplate] Added new extension which provides a template engine for topics (#1932, thanks to @simonthum).

Source/MQTTnet.Extensions.ManagedClient/ManagedMqttClient.cs

@@ -19,6 +19,8 @@ namespace MQTTnet.Extensions.ManagedClient
 {
     public sealed class ManagedMqttClient : Disposable, IManagedMqttClient
     {
+        readonly MqttNetSourceLogger _logger;
+
         readonly AsyncEvent<InterceptingPublishMessageEventArgs> _interceptingPublishMessageEvent = new AsyncEvent<InterceptingPublishMessageEventArgs>();
         readonly AsyncEvent<ApplicationMessageProcessedEventArgs> _applicationMessageProcessedEvent = new AsyncEvent<ApplicationMessageProcessedEventArgs>();
         readonly AsyncEvent<ApplicationMessageSkippedEventArgs> _applicationMessageSkippedEvent = new AsyncEvent<ApplicationMessageSkippedEventArgs>();
@@ -27,9 +29,7 @@ namespace MQTTnet.Extensions.ManagedClient
         readonly AsyncEvent<ManagedProcessFailedEventArgs> _synchronizingSubscriptionsFailedEvent = new AsyncEvent<ManagedProcessFailedEventArgs>();
         readonly AsyncEvent<SubscriptionsChangedEventArgs> _subscriptionsChangedEvent = new AsyncEvent<SubscriptionsChangedEventArgs>();
 
-        readonly MqttNetSourceLogger _logger;
         readonly BlockingQueue<ManagedMqttApplicationMessage> _messageQueue = new BlockingQueue<ManagedMqttApplicationMessage>();
-
         readonly AsyncLock _messageQueueLock = new AsyncLock();
 
         /// <summary>
@@ -90,14 +90,12 @@ namespace MQTTnet.Extensions.ManagedClient
             remove => _applicationMessageProcessedEvent.RemoveHandler(value);
         }
 
-
         public event Func<InterceptingPublishMessageEventArgs, Task> InterceptPublishMessageAsync
         {
             add => _interceptingPublishMessageEvent.AddHandler(value);
             remove => _interceptingPublishMessageEvent.RemoveHandler(value);
         }
 
-
         public event Func<MqttApplicationMessageReceivedEventArgs, Task> ApplicationMessageReceivedAsync
         {
             add => InternalClient.ApplicationMessageReceivedAsync += value;
@@ -369,6 +367,13 @@ namespace MQTTnet.Extensions.ManagedClient
             return remainingTime < TimeSpan.Zero ? TimeSpan.Zero : remainingTime;
         }
 
+        CancellationTokenSource NewTimeoutToken(CancellationToken linkedToken)
+        {
+            var newTimeoutToken = CancellationTokenSource.CreateLinkedTokenSource(linkedToken);
+            newTimeoutToken.CancelAfter(Options.ClientOptions.Timeout);
+            return newTimeoutToken;
+        }
+
         async Task HandleSubscriptionExceptionAsync(Exception exception, List<MqttTopicFilter> addedSubscriptions, List<string> removedSubscriptions)
         {
             _logger.Warning(exception, "Synchronizing subscriptions failed.");
@@ -411,7 +416,7 @@ namespace MQTTnet.Extensions.ManagedClient
                     {
                         if (_isCleanDisconnect)
                         {
-                            using (var disconnectTimeout = new CancellationTokenSource(Options.ClientOptions.Timeout))
+                            using (var disconnectTimeout = NewTimeoutToken(CancellationToken.None))
                             {
                                 await InternalClient.DisconnectAsync(new MqttClientDisconnectOptions(), disconnectTimeout.Token).ConfigureAwait(false);
                             }
@@ -461,7 +466,7 @@ namespace MQTTnet.Extensions.ManagedClient
 
                     cancellationToken.ThrowIfCancellationRequested();
 
-                    await TryPublishQueuedMessageAsync(message).ConfigureAwait(false);
+                    await TryPublishQueuedMessageAsync(message, cancellationToken).ConfigureAwait(false);
                 }
             }
             catch (OperationCanceledException)
@@ -477,7 +482,7 @@ namespace MQTTnet.Extensions.ManagedClient
             }
         }
 
-        async Task PublishReconnectSubscriptionsAsync()
+        async Task PublishReconnectSubscriptionsAsync(CancellationToken cancellationToken)
         {
             _logger.Info("Publishing subscriptions at reconnect");
 
@@ -496,13 +501,13 @@ namespace MQTTnet.Extensions.ManagedClient
 
                         if (topicFilters.Count == Options.MaxTopicFiltersInSubscribeUnsubscribePackets)
                         {
-                            subscribeUnsubscribeResult = await SendSubscribeUnsubscribe(topicFilters, null).ConfigureAwait(false);
+                            subscribeUnsubscribeResult = await SendSubscribeUnsubscribe(topicFilters, null, cancellationToken).ConfigureAwait(false);
                             topicFilters.Clear();
                             await HandleSubscriptionsResultAsync(subscribeUnsubscribeResult).ConfigureAwait(false);
                         }
                     }
 
-                    subscribeUnsubscribeResult = await SendSubscribeUnsubscribe(topicFilters, null).ConfigureAwait(false);
+                    subscribeUnsubscribeResult = await SendSubscribeUnsubscribe(topicFilters, null, cancellationToken).ConfigureAwait(false);
                     await HandleSubscriptionsResultAsync(subscribeUnsubscribeResult).ConfigureAwait(false);
                 }
             }
@@ -555,13 +560,13 @@ namespace MQTTnet.Extensions.ManagedClient
 
                     if (addedTopicFilters.Count == Options.MaxTopicFiltersInSubscribeUnsubscribePackets)
                     {
-                        subscribeUnsubscribeResult = await SendSubscribeUnsubscribe(addedTopicFilters, null).ConfigureAwait(false);
+                        subscribeUnsubscribeResult = await SendSubscribeUnsubscribe(addedTopicFilters, null, cancellationToken).ConfigureAwait(false);
                         addedTopicFilters.Clear();
                         await HandleSubscriptionsResultAsync(subscribeUnsubscribeResult).ConfigureAwait(false);
                     }
                 }
 
-                subscribeUnsubscribeResult = await SendSubscribeUnsubscribe(addedTopicFilters, null).ConfigureAwait(false);
+                subscribeUnsubscribeResult = await SendSubscribeUnsubscribe(addedTopicFilters, null, cancellationToken).ConfigureAwait(false);
                 await HandleSubscriptionsResultAsync(subscribeUnsubscribeResult).ConfigureAwait(false);
 
                 var removedTopicFilters = new List<string>();
@@ -571,13 +576,13 @@ namespace MQTTnet.Extensions.ManagedClient
 
                     if (removedTopicFilters.Count == Options.MaxTopicFiltersInSubscribeUnsubscribePackets)
                     {
-                        subscribeUnsubscribeResult = await SendSubscribeUnsubscribe(null, removedTopicFilters).ConfigureAwait(false);
+                        subscribeUnsubscribeResult = await SendSubscribeUnsubscribe(null, removedTopicFilters, cancellationToken).ConfigureAwait(false);
                         removedTopicFilters.Clear();
                         await HandleSubscriptionsResultAsync(subscribeUnsubscribeResult).ConfigureAwait(false);
                     }
                 }
 
-                subscribeUnsubscribeResult = await SendSubscribeUnsubscribe(null, removedTopicFilters).ConfigureAwait(false);
+                subscribeUnsubscribeResult = await SendSubscribeUnsubscribe(null, removedTopicFilters, cancellationToken).ConfigureAwait(false);
                 await HandleSubscriptionsResultAsync(subscribeUnsubscribeResult).ConfigureAwait(false);
             }
         }
@@ -592,7 +597,7 @@ namespace MQTTnet.Extensions.ManagedClient
             MqttClientConnectResult connectResult = null;
             try
             {
-                using (var connectTimeout = new CancellationTokenSource(Options.ClientOptions.Timeout))
+                using (var connectTimeout = NewTimeoutToken(cancellationToken))
                 {
                     connectResult = await InternalClient.ConnectAsync(Options.ClientOptions, connectTimeout.Token).ConfigureAwait(false);
                 }
@@ -611,7 +616,7 @@ namespace MQTTnet.Extensions.ManagedClient
             }
         }
 
-        async Task<SendSubscribeUnsubscribeResult> SendSubscribeUnsubscribe(List<MqttTopicFilter> addedSubscriptions, List<string> removedSubscriptions)
+        async Task<SendSubscribeUnsubscribeResult> SendSubscribeUnsubscribe(List<MqttTopicFilter> addedSubscriptions, List<string> removedSubscriptions, CancellationToken cancellationToken)
         {
             var subscribeResults = new List<MqttClientSubscribeResult>();
             var unsubscribeResults = new List<MqttClientUnsubscribeResult>();
@@ -626,8 +631,11 @@ namespace MQTTnet.Extensions.ManagedClient
                         unsubscribeOptionsBuilder.WithTopicFilter(removedSubscription);
                     }
 
-                    var unsubscribeResult = await InternalClient.UnsubscribeAsync(unsubscribeOptionsBuilder.Build()).ConfigureAwait(false);
+                    using (var unsubscribeTimeout = NewTimeoutToken(cancellationToken))
-                    unsubscribeResults.Add(unsubscribeResult);
+                    {
+                        var unsubscribeResult = await InternalClient.UnsubscribeAsync(unsubscribeOptionsBuilder.Build(), unsubscribeTimeout.Token).ConfigureAwait(false);
+                        unsubscribeResults.Add(unsubscribeResult);
+                    }
 
                     //clear because these worked, maybe the subscribe below will fail, only report those
                     removedSubscriptions.Clear();
@@ -642,8 +650,11 @@ namespace MQTTnet.Extensions.ManagedClient
                         subscribeOptionsBuilder.WithTopicFilter(addedSubscription);
                     }
 
-                    var subscribeResult = await InternalClient.SubscribeAsync(subscribeOptionsBuilder.Build()).ConfigureAwait(false);
+                    using (var subscribeTimeout = NewTimeoutToken(cancellationToken))
-                    subscribeResults.Add(subscribeResult);
+                    {
+                        var subscribeResult = await InternalClient.SubscribeAsync(subscribeOptionsBuilder.Build(), subscribeTimeout.Token).ConfigureAwait(false);
+                        subscribeResults.Add(subscribeResult);
+                    }
                 }
             }
             catch (Exception exception)
@@ -705,7 +716,7 @@ namespace MQTTnet.Extensions.ManagedClient
                 }
                 else if (connectionState == ReconnectionResult.Reconnected)
                 {
-                    await PublishReconnectSubscriptionsAsync().ConfigureAwait(false);
+                    await PublishReconnectSubscriptionsAsync(cancellationToken).ConfigureAwait(false);
                     StartPublishing();
                 }
                 else if (connectionState == ReconnectionResult.Recovered)
@@ -735,7 +746,7 @@ namespace MQTTnet.Extensions.ManagedClient
             }
         }
 
-        async Task TryPublishQueuedMessageAsync(ManagedMqttApplicationMessage message)
+        async Task TryPublishQueuedMessageAsync(ManagedMqttApplicationMessage message, CancellationToken cancellationToken)
         {
             Exception transmitException = null;
             bool acceptPublish = true;
@@ -750,7 +761,10 @@ namespace MQTTnet.Extensions.ManagedClient
 
                 if (acceptPublish)
                 {
-                    await InternalClient.PublishAsync(message.ApplicationMessage).ConfigureAwait(false);
+                    using (var publishTimeout = NewTimeoutToken(cancellationToken))
+                    {
+                        await InternalClient.PublishAsync(message.ApplicationMessage, publishTimeout.Token).ConfigureAwait(false);
+                    }
                 }
 
                 using (await _messageQueueLock.EnterAsync().ConfigureAwait(false)) //lock to avoid conflict with this.PublishAsync

Source/MQTTnet.Tests/Clients/ManagedMqttClient/ManagedMqttClient_Tests.cs

@@ -7,12 +7,14 @@ using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 using MQTTnet.Client;
 using MQTTnet.Diagnostics;
 using MQTTnet.Extensions.ManagedClient;
 using MQTTnet.Formatter;
 using MQTTnet.Internal;
+using MQTTnet.Packets;
 using MQTTnet.Protocol;
 using MQTTnet.Server;
 using MQTTnet.Tests.Mockups;
@@ -574,16 +576,152 @@ namespace MQTTnet.Tests.Clients.ManagedMqttClient
             }
         }
 
+        [TestMethod]
+        public async Task Subscribe_Does_Not_Hang_On_Server_Stop()
+        {
+            var timeout = TimeSpan.FromSeconds(2);
+            var testTimeout = timeout * 2;
+            const string topic = "test_topic_2";
+            using (var testEnvironment = CreateTestEnvironment())
+            using (var managedClient = await CreateManagedClientAsync(testEnvironment, timeout: timeout))
+            {
+                testEnvironment.IgnoreClientLogErrors = true;
+                bool reject = true;
+                var receivedOnServer = new SemaphoreSlim(0, 1);
+                var failedOnClient = new SemaphoreSlim(0, 1);
+                testEnvironment.Server.InterceptingInboundPacketAsync += e =>
+                {
+                    if (e.Packet is MqttSubscribePacket)
+                    {
+                        if (reject)
+                        {
+                            e.ProcessPacket = false;
+                        }
+                        receivedOnServer.Release();
+                    }
+                    return Task.CompletedTask;
+                };
+                managedClient.SynchronizingSubscriptionsFailedAsync += e =>
+                {
+                    failedOnClient.Release();
+                    return Task.CompletedTask;
+                };
+
+                await managedClient.SubscribeAsync(topic);
+                Assert.IsTrue(await receivedOnServer.WaitAsync(testTimeout));
+                Assert.IsTrue(await failedOnClient.WaitAsync(testTimeout));
+
+                reject = false;
+                await managedClient.SubscribeAsync(topic);
+                Assert.IsTrue(await receivedOnServer.WaitAsync(testTimeout));
+            }
+        }
+
+        [TestMethod]
+        public async Task Unsubscribe_Does_Not_Hang_On_Server_Stop()
+        {
+            var timeout = TimeSpan.FromSeconds(2);
+            var testTimeout = timeout * 2;
+            const string topic = "test_topic_2";
+            using (var testEnvironment = CreateTestEnvironment())
+            using (var managedClient = await CreateManagedClientAsync(testEnvironment, timeout: timeout))
+            {
+                testEnvironment.IgnoreClientLogErrors = true;
+                bool reject = true;
+                var receivedOnServer = new SemaphoreSlim(0, 1);
+                var failedOnClient = new SemaphoreSlim(0, 1);
+                testEnvironment.Server.InterceptingInboundPacketAsync += e =>
+                {
+                    if (e.Packet is MqttUnsubscribePacket)
+                    {
+                        if (reject)
+                        {
+                            e.ProcessPacket = false;
+                        }
+                        receivedOnServer.Release();
+                    }
+                    else if (e.Packet is MqttSubscribePacket)
+                    {
+                        receivedOnServer.Release();
+                    }
+                    return Task.CompletedTask;
+                };
+                managedClient.SynchronizingSubscriptionsFailedAsync += e =>
+                {
+                    failedOnClient.Release();
+                    return Task.CompletedTask;
+                };
+
+                await managedClient.SubscribeAsync(topic);
+                Assert.IsTrue(await receivedOnServer.WaitAsync(testTimeout));
+
+                await managedClient.UnsubscribeAsync(topic);
+                Assert.IsTrue(await receivedOnServer.WaitAsync(testTimeout));
+                Assert.IsTrue(await failedOnClient.WaitAsync(testTimeout));
+
+                reject = false;
+                await managedClient.UnsubscribeAsync(topic);
+                Assert.IsTrue(await receivedOnServer.WaitAsync(testTimeout));
+            }
+        }
+
+        [TestMethod]
+        public async Task Publish_Does_Not_Hang_On_Server_Error()
+        {
+            var timeout = TimeSpan.FromSeconds(2);
+            var testTimeout = timeout * 2;
+
+            const string topic = "test_topic_42";
+
+            using (var testEnvironment = CreateTestEnvironment())
+            using (var managedClient = await CreateManagedClientAsync(testEnvironment, timeout: timeout))
+            {
+                testEnvironment.IgnoreClientLogErrors = true;
+                bool reject = true;
+                var receivedOnServer = new TaskCompletionSource();
+                managedClient.ApplicationMessageProcessedAsync += e => Task.FromResult(reject &= e.Exception is null);
+                testEnvironment.Server.InterceptingInboundPacketAsync += e => 
+                {
+                    if (e.Packet is MqttPublishPacket)
+                    {
+                        if (reject)
+                        {
+                            e.ProcessPacket = false;
+                        }
+                        else
+                        {
+                            receivedOnServer.TrySetResult();
+                        }
+                    }
+                    return Task.CompletedTask;
+                };
+
+
+                await managedClient.EnqueueAsync(new MqttApplicationMessage { Topic = topic, Payload = new byte[] { 1 }, Retain = true, QualityOfServiceLevel = MqttQualityOfServiceLevel.AtLeastOnce });
+
+                var timeoutTask = Task.Delay(testTimeout);
+
+                var firstDone = await Task.WhenAny(receivedOnServer.Task, timeoutTask);
+                Assert.AreEqual(receivedOnServer.Task, firstDone, "Client is hung on publish!");
+            }
+        }
+
         async Task<MQTTnet.Extensions.ManagedClient.ManagedMqttClient> CreateManagedClientAsync(
             TestEnvironment testEnvironment,
             IMqttClient underlyingClient = null,
             TimeSpan? connectionCheckInterval = null,
-            string host = "localhost")
+            string host = "localhost",
+            TimeSpan? timeout = null)
         {
             await testEnvironment.StartServer();
 
             var clientOptions = new MqttClientOptionsBuilder().WithTcpServer(host, testEnvironment.ServerPort);
 
+            if (timeout != null)
+            {
+                clientOptions.WithTimeout(timeout.Value);
+            }
+
             var managedOptions = new ManagedMqttClientOptionsBuilder().WithClientOptions(clientOptions).Build();
 
             // Use a short connection check interval so that subscription operations are performed quickly