From d572490a677b2cb941a03cd02926db15c987ec92 Mon Sep 17 00:00:00 2001
From: Christian <6939810+chkr1011@users.noreply.github.com>
Date: Sun, 4 Dec 2022 13:39:50 +0100
Subject: [PATCH] Use the net452 supporting async task completion source
 (#1607)

---
 .../MQTTnet.Extensions.Rpc/MqttRpcClient.cs   | 81 +++++++++----------
 1 file changed, 39 insertions(+), 42 deletions(-)

diff --git a/Source/MQTTnet.Extensions.Rpc/MqttRpcClient.cs b/Source/MQTTnet.Extensions.Rpc/MqttRpcClient.cs
index 205e1591..e3ac567a 100644
--- a/Source/MQTTnet.Extensions.Rpc/MqttRpcClient.cs
+++ b/Source/MQTTnet.Extensions.Rpc/MqttRpcClient.cs
@@ -2,23 +2,23 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-using MQTTnet.Client;
-using MQTTnet.Exceptions;
-using MQTTnet.Protocol;
 using System;
 using System.Collections.Concurrent;
 using System.Threading;
 using System.Threading.Tasks;
+using MQTTnet.Client;
+using MQTTnet.Exceptions;
 using MQTTnet.Internal;
+using MQTTnet.Protocol;
 
 namespace MQTTnet.Extensions.Rpc
 {
     public sealed class MqttRpcClient : IMqttRpcClient
     {
-        readonly ConcurrentDictionary<string, TaskCompletionSource<byte[]>> _waitingCalls = new ConcurrentDictionary<string, TaskCompletionSource<byte[]>>();
         readonly IMqttClient _mqttClient;
         readonly MqttRpcClientOptions _options;
-        
+        readonly ConcurrentDictionary<string, AsyncTaskCompletionSource<byte[]>> _waitingCalls = new ConcurrentDictionary<string, AsyncTaskCompletionSource<byte[]>>();
+
         public MqttRpcClient(IMqttClient mqttClient, MqttRpcClientOptions options)
         {
             _mqttClient = mqttClient ?? throw new ArgumentNullException(nameof(mqttClient));
@@ -27,6 +27,18 @@ namespace MQTTnet.Extensions.Rpc
             _mqttClient.ApplicationMessageReceivedAsync += HandleApplicationMessageReceivedAsync;
         }
 
+        public void Dispose()
+        {
+            _mqttClient.ApplicationMessageReceivedAsync -= HandleApplicationMessageReceivedAsync;
+
+            foreach (var tcs in _waitingCalls)
+            {
+                tcs.Value.TrySetCanceled();
+            }
+
+            _waitingCalls.Clear();
+        }
+
         public async Task<byte[]> ExecuteAsync(TimeSpan timeout, string methodName, byte[] payload, MqttQualityOfServiceLevel qualityOfServiceLevel)
         {
             using (var timeoutToken = new CancellationTokenSource(timeout))
@@ -49,15 +61,19 @@ namespace MQTTnet.Extensions.Rpc
 
         public async Task<byte[]> ExecuteAsync(string methodName, byte[] payload, MqttQualityOfServiceLevel qualityOfServiceLevel, CancellationToken cancellationToken = default)
         {
-            if (methodName == null) throw new ArgumentNullException(nameof(methodName));
-            
-            var topicNames = _options.TopicGenerationStrategy.CreateRpcTopics(new TopicGenerationContext
+            if (methodName == null)
             {
-                MethodName = methodName,
-                QualityOfServiceLevel = qualityOfServiceLevel,
-                MqttClient = _mqttClient,
-                Options = _options
-            });
+                throw new ArgumentNullException(nameof(methodName));
+            }
+
+            var topicNames = _options.TopicGenerationStrategy.CreateRpcTopics(
+                new TopicGenerationContext
+                {
+                    MethodName = methodName,
+                    QualityOfServiceLevel = qualityOfServiceLevel,
+                    MqttClient = _mqttClient,
+                    Options = _options
+                });
 
             var requestTopic = topicNames.RequestTopic;
             var responseTopic = topicNames.ResponseTopic;
@@ -72,8 +88,7 @@ namespace MQTTnet.Extensions.Rpc
                 throw new MqttProtocolViolationException("RPC response topic is empty.");
             }
 
-            var requestMessage = new MqttApplicationMessageBuilder()
-                .WithTopic(requestTopic)
+            var requestMessage = new MqttApplicationMessageBuilder().WithTopic(requestTopic)
                 .WithPayload(payload)
                 .WithQualityOfServiceLevel(qualityOfServiceLevel)
                 .WithResponseTopic(responseTopic)
@@ -81,25 +96,23 @@ namespace MQTTnet.Extensions.Rpc
 
             try
             {
-#if NET452
-                var awaitable = new TaskCompletionSource<byte[]>();
-#else
-                var awaitable = new TaskCompletionSource<byte[]>(TaskCreationOptions.RunContinuationsAsynchronously);
-#endif
+                var awaitable = new AsyncTaskCompletionSource<byte[]>();
 
                 if (!_waitingCalls.TryAdd(responseTopic, awaitable))
                 {
                     throw new InvalidOperationException();
                 }
 
-                var subscribeOptions = new MqttClientSubscribeOptionsBuilder()
-                    .WithTopicFilter(responseTopic, qualityOfServiceLevel)
-                    .Build();
+                var subscribeOptions = new MqttClientSubscribeOptionsBuilder().WithTopicFilter(responseTopic, qualityOfServiceLevel).Build();
 
                 await _mqttClient.SubscribeAsync(subscribeOptions, cancellationToken).ConfigureAwait(false);
                 await _mqttClient.PublishAsync(requestMessage, cancellationToken).ConfigureAwait(false);
 
-                using (cancellationToken.Register(() => { awaitable.TrySetCanceled(); }))
+                using (cancellationToken.Register(
+                           () =>
+                           {
+                               awaitable.TrySetCanceled();
+                           }))
                 {
                     return await awaitable.Task.ConfigureAwait(false);
                 }
@@ -107,8 +120,8 @@ namespace MQTTnet.Extensions.Rpc
             finally
             {
                 _waitingCalls.TryRemove(responseTopic, out _);
-                
-                await _mqttClient.UnsubscribeAsync(responseTopic).ConfigureAwait(false);
+
+                await _mqttClient.UnsubscribeAsync(responseTopic, CancellationToken.None).ConfigureAwait(false);
             }
         }
 
@@ -119,28 +132,12 @@ namespace MQTTnet.Extensions.Rpc
                 return CompletedTask.Instance;
             }
 
-#if NET452
-            Task.Run(() => awaitable.TrySetResult(eventArgs.ApplicationMessage.Payload));
-#else
             awaitable.TrySetResult(eventArgs.ApplicationMessage.Payload);
-#endif
 
             // Set this message to handled to that other code can avoid execution etc.
             eventArgs.IsHandled = true;
 
             return CompletedTask.Instance;
         }
-
-        public void Dispose()
-        {
-            _mqttClient.ApplicationMessageReceivedAsync -= HandleApplicationMessageReceivedAsync;
-
-            foreach (var tcs in _waitingCalls)
-            {
-                tcs.Value.TrySetCanceled();
-            }
-
-            _waitingCalls.Clear();
-        }
     }
 }
\ No newline at end of file