Improved Wait Strategy for Publish Transactions

Source/HiveMQtt/Client/HiveMQClient.cs

@@ -207,37 +207,25 @@
         else if (message.QoS == QualityOfService.AtLeastOnceDelivery)
         {
             // QoS 1: Acknowledged Delivery
-            var taskCompletionSource = new TaskCompletionSource<PubAckPacket>();
-            void TaskHandler(object? sender, OnPublishQoS1CompleteEventArgs args) => taskCompletionSource.SetResult(args.PubAckPacket);
-            EventHandler<OnPublishQoS1CompleteEventArgs> eventHandler = TaskHandler;
-            publishPacket.OnPublishQoS1Complete += eventHandler;
-
             Logger.Trace($"Queuing packet for send: {publishPacket.GetType().Name} id={publishPacket.PacketIdentifier}");
             this.OutgoingPublishQueue.Enqueue(publishPacket);
 
-            var pubAckPacket = await taskCompletionSource.Task.WaitAsync(TimeSpan.FromSeconds(60)).ConfigureAwait(false);
-
-            publishPacket.OnPublishQoS1Complete -= eventHandler;
+            // Wait on the QoS 1 handshake
+            var pubAckPacket = await publishPacket.OnPublishQoS1CompleteTCS.Task.WaitAsync(TimeSpan.FromSeconds(60)).ConfigureAwait(false);
             return new PublishResult(publishPacket.Message, pubAckPacket);
         }
         else if (message.QoS == QualityOfService.ExactlyOnceDelivery)
         {
             // QoS 2: Assured Delivery
             PublishResult? publishResult = null;
-            var taskCompletionSource = new TaskCompletionSource<List<ControlPacket>>();
-            void TaskHandler(object? sender, OnPublishQoS2CompleteEventArgs args) => taskCompletionSource.SetResult(args.PacketList);
-            EventHandler<OnPublishQoS2CompleteEventArgs> eventHandler = TaskHandler;
-            publishPacket.OnPublishQoS2Complete += eventHandler;
-
             Logger.Trace($"Queuing packet for send: {publishPacket.GetType().Name} id={publishPacket.PacketIdentifier}");
             this.OutgoingPublishQueue.Enqueue(publishPacket);
 
             List<ControlPacket> packetList;
             try
             {
                 // Wait on the QoS 2 handshake
-                // FIXME: Timeout value
-                packetList = await taskCompletionSource.Task.WaitAsync(TimeSpan.FromSeconds(60)).ConfigureAwait(false);
+                packetList = await publishPacket.OnPublishQoS2CompleteTCS.Task.WaitAsync(TimeSpan.FromSeconds(60)).ConfigureAwait(false);
             }
             catch (TimeoutException)
             {
@@ -254,7 +242,6 @@
                 {
                     QoS2ReasonCode = null,
                 };
-                publishPacket.OnPublishQoS2Complete -= eventHandler;
                 return publishResult;
             }
 
@@ -271,8 +258,6 @@
                 throw new HiveMQttClientException("PublishAsync: QoS 2 complete but no PubRec packet received.");
             }
 
-            // Remove our wait handler
-            publishPacket.OnPublishQoS2Complete -= eventHandler;
             return publishResult;
         }
 
@@ -497,13 +482,13 @@
     /// Close the socket and set the connect state to disconnected.
     /// </summary>
     /// <param name="clean">Indicates whether the disconnect was intended or not.</param>
     private async Task<bool> HandleDisconnectionAsync(bool clean = true)
     {
         Logger.Debug($"HandleDisconnection: Handling disconnection. clean={clean}.");
 
         // Cancel all background tasks and close the socket
         this.ConnectState = ConnectState.Disconnected;
         this.cancellationTokenSource.Cancel();
         this.CloseSocket();
 
         if (clean)

Source/HiveMQtt/MQTT5/Packets/PublishPacket.cs

@@ -39,6 +39,20 @@
     {
         this.PacketIdentifier = (ushort)packetIdentifier;
         this.Message = message;
+
+        // Setup the QoS 1 TaskCompletionSource so users can simply call
+        //
+        //   await PublishPacket.OnPublishQoS1CompleteTCS
+        //
+        // to wait for the QoS 1 publish to complete.
+        this.OnPublishQoS1Complete += (sender, args) => this.OnPublishQoS1CompleteTCS.SetResult(args.PubAckPacket);
+
+        // Setup the QoS 2 TaskCompletionSource so users can simply call
+        //
+        //   await PublishPacket.OnPublishQoS2CompleteTCS
+        //
+        // to wait for the QoS 2 publish to complete.
+        this.OnPublishQoS2Complete += (sender, args) => this.OnPublishQoS2CompleteTCS.SetResult(args.PacketList);
     }
 
     /// <summary>
@@ -66,23 +80,61 @@
 
     internal virtual void OnPublishQoS1CompleteEventLauncher(PubAckPacket packet)
     {
-        var eventArgs = new OnPublishQoS1CompleteEventArgs(packet);
-        Logger.Trace("OnPublishQoS1CompleteEventLauncher");
-        this.OnPublishQoS1Complete?.Invoke(this, eventArgs);
+        if (this.OnPublishQoS1Complete != null && this.OnPublishQoS1Complete.GetInvocationList().Length > 0)
+        {
+            var eventArgs = new OnPublishQoS1CompleteEventArgs(packet);
+            Logger.Trace("OnPublishQoS1CompleteEventLauncher");
+            _ = Task.Run(() => this.OnPublishQoS1Complete?.Invoke(this, eventArgs)).ContinueWith(
+                t =>
+                {
+                    if (t.IsFaulted)
+                    {
+                        Logger.Error("OnPublishQoS1CompleteEventLauncher exception: " + t.Exception.Message);
+                    }
+                },
+                TaskScheduler.Default);
+        }
     }
 
+    /// <summary>
+    /// Gets the awaitable TaskCompletionSource for the QoS 1 publish transaction.
+    /// <para>
+    /// Valid for outgoing Publish messages QoS 1.  A TaskCompletionSource that is set when the QoS 1 publish transaction is complete.
+    /// </para>
+    /// </summary>
+    public TaskCompletionSource<PubAckPacket> OnPublishQoS1CompleteTCS { get; } = new();
+
     /// <summary>
     /// Valid for outgoing Publish messages QoS 2.  An event that is fired after the the QoS 2 PubComp is received.
     /// </summary>
     public event EventHandler<OnPublishQoS2CompleteEventArgs> OnPublishQoS2Complete = new((client, e) => { });
 
     internal virtual void OnPublishQoS2CompleteEventLauncher(List<ControlPacket> packetList)
     {
-        var eventArgs = new OnPublishQoS2CompleteEventArgs(packetList);
-        Logger.Trace("OnPublishQoS2CompleteEventLauncher");
-        this.OnPublishQoS2Complete?.Invoke(this, eventArgs);
+        if (this.OnPublishQoS2Complete != null && this.OnPublishQoS2Complete.GetInvocationList().Length > 0)
+        {
+            var eventArgs = new OnPublishQoS2CompleteEventArgs(packetList);
+            Logger.Trace("OnPublishQoS2CompleteEventLauncher");
+            _ = Task.Run(() => this.OnPublishQoS2Complete?.Invoke(this, eventArgs)).ContinueWith(
+                t =>
+                {
+                    if (t.IsFaulted)
+                    {
+                        Logger.Error("OnPublishQoS2CompleteEventLauncher exception: " + t.Exception.Message);
+                    }
+                },
+                TaskScheduler.Default);
+        }
     }
 
+    /// <summary>
+    /// Gets the awaitable TaskCompletionSource for the QoS 2 publish transaction.
+    /// <para>
+    /// Valid for outgoing Publish messages QoS 2.  A TaskCompletionSource that is set when the QoS 2 publish transaction is complete.
+    /// </para>
+    /// </summary>
+    public TaskCompletionSource<List<ControlPacket>> OnPublishQoS2CompleteTCS { get; } = new();
+
     /// <summary>
     /// Decode the received MQTT Publish packet.
     /// </summary>