Improve RPC client handling.

Build/MQTTnet.nuspec

@@ -17,6 +17,7 @@
 * [Client] Exposed Dup flag in application messages.
 * [Client] Improved keep alive message sending.
 * [RpcClient] Fixed an issue when using a custom application message reveived handler (#1093).
+* [RpcClient] Fixed a race condition when using timeout and cancellation token at the same time (BREAKING CHANGE!).
 * [Server] Improved keep alive message sending.
 
 Git commit: $gitCommit

Source/MQTTnet.Extensions.Rpc/IMqttRpcClient.cs

@@ -7,12 +7,8 @@ namespace MQTTnet.Extensions.Rpc
 {
     public interface IMqttRpcClient : IDisposable
     {
-        Task<byte[]> ExecuteAsync(TimeSpan timeout, string methodName, string payload, MqttQualityOfServiceLevel qualityOfServiceLevel);
-
-        Task<byte[]> ExecuteAsync(TimeSpan timeout, string methodName, string payload, MqttQualityOfServiceLevel qualityOfServiceLevel, CancellationToken cancellationToken);
-
         Task<byte[]> ExecuteAsync(TimeSpan timeout, string methodName, byte[] payload, MqttQualityOfServiceLevel qualityOfServiceLevel);
 
-        Task<byte[]> ExecuteAsync(TimeSpan timeout, string methodName, byte[] payload, MqttQualityOfServiceLevel qualityOfServiceLevel, CancellationToken cancellationToken);
+        Task<byte[]> ExecuteAsync(string methodName, byte[] payload, MqttQualityOfServiceLevel qualityOfServiceLevel, CancellationToken cancellationToken);
     }
 }

Source/MQTTnet.Extensions.Rpc/MqttRpcClient.cs

@@ -5,9 +5,9 @@ using MQTTnet.Extensions.Rpc.Options.TopicGeneration;
 using MQTTnet.Protocol;
 using System;
 using System.Collections.Concurrent;
-using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using MQTTnet.Client.Subscribing;
 using MQTTnet.Implementations;
 
 namespace MQTTnet.Extensions.Rpc
@@ -36,22 +36,27 @@ namespace MQTTnet.Extensions.Rpc
             _mqttClient.ApplicationMessageReceivedHandler = _applicationMessageReceivedHandler;
         }
 
-        public Task<byte[]> ExecuteAsync(TimeSpan timeout, string methodName, string payload, MqttQualityOfServiceLevel qualityOfServiceLevel)
+        public async Task<byte[]> ExecuteAsync(TimeSpan timeout, string methodName, byte[] payload, MqttQualityOfServiceLevel qualityOfServiceLevel)
         {
-            return ExecuteAsync(timeout, methodName, Encoding.UTF8.GetBytes(payload), qualityOfServiceLevel, CancellationToken.None);
+            using (var timeoutToken = new CancellationTokenSource(timeout))
+            {
+                try
+                {
+                    return await ExecuteAsync(methodName, payload, qualityOfServiceLevel, timeoutToken.Token).ConfigureAwait(false);
+                }
+                catch (OperationCanceledException exception)
+                {
+                    if (timeoutToken.IsCancellationRequested)
+                    {
+                        throw new MqttCommunicationTimedOutException(exception);
+                    }
+
+                    throw;
+                }
+            }
         }
 
-        public Task<byte[]> ExecuteAsync(TimeSpan timeout, string methodName, string payload, MqttQualityOfServiceLevel qualityOfServiceLevel, CancellationToken cancellationToken)
+        public async Task<byte[]> ExecuteAsync(string methodName, byte[] payload, MqttQualityOfServiceLevel qualityOfServiceLevel, CancellationToken cancellationToken)
-        {
-            return ExecuteAsync(timeout, methodName, Encoding.UTF8.GetBytes(payload), qualityOfServiceLevel, cancellationToken);
-        }
-
-        public Task<byte[]> ExecuteAsync(TimeSpan timeout, string methodName, byte[] payload, MqttQualityOfServiceLevel qualityOfServiceLevel)
-        {
-            return ExecuteAsync(timeout, methodName, payload, qualityOfServiceLevel, CancellationToken.None);
-        }
-
-        public async Task<byte[]> ExecuteAsync(TimeSpan timeout, string methodName, byte[] payload, MqttQualityOfServiceLevel qualityOfServiceLevel, CancellationToken cancellationToken)
         {
             if (methodName == null) throw new ArgumentNullException(nameof(methodName));
 
@@ -89,48 +94,33 @@ namespace MQTTnet.Extensions.Rpc
 
             try
             {
+#if NET452
                 var promise = new TaskCompletionSource<byte[]>();
+#else
+                var promise = new TaskCompletionSource<byte[]>(TaskCreationOptions.RunContinuationsAsynchronously);
+#endif
+                
                 if (!_waitingCalls.TryAdd(responseTopic, promise))
                 {
                     throw new InvalidOperationException();
                 }
 
-                await _mqttClient.SubscribeAsync(responseTopic, qualityOfServiceLevel).ConfigureAwait(false);
+                var subscribeOptions = new MqttClientSubscribeOptionsBuilder()
-                await _mqttClient.PublishAsync(requestMessage).ConfigureAwait(false);
+                    .WithTopicFilter(responseTopic, qualityOfServiceLevel)
+                    .Build();
 
-                using (var timeoutCts = new CancellationTokenSource(timeout))
+                await _mqttClient.SubscribeAsync(subscribeOptions, cancellationToken).ConfigureAwait(false);
-                using (var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token))
+                await _mqttClient.PublishAsync(requestMessage, cancellationToken).ConfigureAwait(false);
+
+                using (cancellationToken.Register(() => { promise.TrySetCanceled(); }))
                 {
-                    linkedCts.Token.Register(() =>
+                    return await promise.Task.ConfigureAwait(false);
-                    {
-                        if (!promise.Task.IsCompleted && !promise.Task.IsFaulted && !promise.Task.IsCanceled)
-                        {
-                            promise.TrySetCanceled();
-                        }
-                    });
-
-                    try
-                    {
-                        var result = await promise.Task.ConfigureAwait(false);
-                        timeoutCts.Cancel(false);
-                        return result;
-                    }
-                    catch (OperationCanceledException exception)
-                    {
-                        if (timeoutCts.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
-                        {
-                            throw new MqttCommunicationTimedOutException(exception);
-                        }
-                        else
-                        {
-                            throw;
-                        }
-                    }
                 }
             }
             finally
             {
                 _waitingCalls.TryRemove(responseTopic, out _);
+
                 await _mqttClient.UnsubscribeAsync(responseTopic).ConfigureAwait(false);
             }
         }
@@ -142,11 +132,15 @@ namespace MQTTnet.Extensions.Rpc
                 return PlatformAbstractionLayer.CompletedTask;
             }
 
+#if NET452
+            Task.Run(() => promise.TrySetResult(eventArgs.ApplicationMessage.Payload));
+#else
             promise.TrySetResult(eventArgs.ApplicationMessage.Payload);
+#endif
 
             // Set this message to handled to that other code can avoid execution etc.
             eventArgs.IsHandled = true;
-            
+
             return PlatformAbstractionLayer.CompletedTask;
         }
 
@@ -162,4 +156,4 @@ namespace MQTTnet.Extensions.Rpc
             _waitingCalls.Clear();
         }
     }
-}
+}

Source/MQTTnet.Extensions.Rpc/MqttRpcClientExtensions.cs

@@ -0,0 +1,19 @@
+using System;
+using System.Text;
+using System.Threading.Tasks;
+using MQTTnet.Protocol;
+
+namespace MQTTnet.Extensions.Rpc
+{
+    public static class MqttRpcClientExtensions
+    {
+        public static Task<byte[]> ExecuteAsync(this IMqttRpcClient client, TimeSpan timeout, string methodName, string payload, MqttQualityOfServiceLevel qualityOfServiceLevel)
+        {
+            if (client == null) throw new ArgumentNullException(nameof(client));
+
+            var buffer = Encoding.UTF8.GetBytes(payload ?? string.Empty);
+
+            return client.ExecuteAsync(timeout, methodName, buffer, qualityOfServiceLevel);
+        }
+    }
+}

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

@@ -9,6 +9,8 @@ using System.Diagnostics;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using MQTTnet.Extensions.Rpc;
+using MQTTnet.Extensions.Rpc.Options;
 using MQTTnet.LowLevelClient;
 
 namespace MQTTnet.Tests.Mockups
@@ -77,6 +79,11 @@ namespace MQTTnet.Tests.Mockups
             };
         }
 
+        public async Task<IMqttRpcClient> ConnectRpcClientAsync(IMqttRpcClientOptions options)
+        {
+            return new MqttRpcClient(await ConnectClientAsync(), options);
+        }
+
         public IMqttClient CreateClient()
         {
             lock (_clients)

Tests/MQTTnet.Core.Tests/Mockups/TestServerWrapper.cs

@@ -109,15 +109,6 @@ namespace MQTTnet.Tests.Mockups
             return Implementation.StartAsync(options);
         }
 
-        public void ConnectionValidator(MqttConnectionValidatorContext ctx)
-        {
-            if (!ctx.ClientId.StartsWith(TestContext.TestName))
-            {
-                TestEnvironment.TrackException(new InvalidOperationException($"Invalid client ID used ({ctx.ClientId}). It must start with UnitTest name."));
-                ctx.ReasonCode = Protocol.MqttConnectReasonCode.ClientIdentifierNotValid;
-            }
-        }
-
         public Task StopAsync()
         {
             return Implementation.StopAsync();
@@ -137,5 +128,14 @@ namespace MQTTnet.Tests.Mockups
         {
             Implementation.Dispose();
         }
+        
+        void ConnectionValidator(MqttConnectionValidatorContext ctx)
+        {
+            if (!ctx.ClientId.StartsWith(TestContext.TestName))
+            {
+                TestEnvironment.TrackException(new InvalidOperationException($"Invalid client ID used ({ctx.ClientId}). It must start with UnitTest name."));
+                ctx.ReasonCode = Protocol.MqttConnectReasonCode.ClientIdentifierNotValid;
+            }
+        }
     }
 }

Tests/MQTTnet.Core.Tests/RPC_Tests.cs

@@ -81,18 +81,19 @@ namespace MQTTnet.Tests
 
                 var requestSender = await testEnvironment.ConnectClientAsync();
 
-                var rpcClient = new MqttRpcClient(requestSender, new MqttRpcClientOptionsBuilder().WithTopicGenerationStrategy(new TestTopicStrategy()) .Build());
+                var rpcClient = await testEnvironment.ConnectRpcClientAsync(new MqttRpcClientOptionsBuilder().WithTopicGenerationStrategy(new TestTopicStrategy()) .Build());
+                
                 await rpcClient.ExecuteAsync(TimeSpan.FromSeconds(2), "ping", "", MqttQualityOfServiceLevel.AtMostOnce);
             }
         }
 
-        private async Task Execute_Success(MqttQualityOfServiceLevel qosLevel, MqttProtocolVersion protocolVersion)
+        async Task Execute_Success(MqttQualityOfServiceLevel qosLevel, MqttProtocolVersion protocolVersion)
         {
             using (var testEnvironment = new TestEnvironment(TestContext))
             {
                 await testEnvironment.StartServerAsync();
                 var responseSender = await testEnvironment.ConnectClientAsync(new MqttClientOptionsBuilder().WithProtocolVersion(protocolVersion));
-                await responseSender.SubscribeAsync("MQTTnet.RPC/+/ping");
+                await responseSender.SubscribeAsync("MQTTnet.RPC/+/ping", qosLevel);
 
                 responseSender.ApplicationMessageReceivedHandler = new MqttApplicationMessageReceivedHandlerDelegate(async e =>
                 {
@@ -101,14 +102,16 @@ namespace MQTTnet.Tests
 
                 var requestSender = await testEnvironment.ConnectClientAsync();
 
-                var rpcClient = new MqttRpcClient(requestSender, new MqttRpcClientOptionsBuilder().Build());
+                using (var rpcClient = new MqttRpcClient(requestSender, new MqttRpcClientOptionsBuilder().Build()))
-                var response = await rpcClient.ExecuteAsync(TimeSpan.FromSeconds(5), "ping", "", qosLevel);
+                {
-
+                    var response = await rpcClient.ExecuteAsync(TimeSpan.FromSeconds(5), "ping", "", qosLevel);
-                Assert.AreEqual("pong", Encoding.UTF8.GetString(response));
+                    
+                    Assert.AreEqual("pong", Encoding.UTF8.GetString(response));
+                }
             }
         }
 
-        private class TestTopicStrategy : IMqttRpcClientTopicGenerationStrategy
+        class TestTopicStrategy : IMqttRpcClientTopicGenerationStrategy
         {
             public MqttRpcTopicPair CreateRpcTopics(TopicGenerationContext context)
             {