Improve session message injection methods

Author: Anton Smolkov

Source/MQTTnet.Server/Exceptions/MqttPendingMessagesOverflowException.cs

@@ -0,0 +1,19 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+namespace MQTTnet.Server.Exceptions
+{
+    public class MqttPendingMessagesOverflowException : Exception
+    {
+        public string SessionId { get; }
+        public MqttPendingMessagesOverflowStrategy OverflowStrategy { get; }
+
+        public MqttPendingMessagesOverflowException(string sessionId, MqttPendingMessagesOverflowStrategy overflowStrategy)
+            : base($"Send buffer max pending messages overflow occurred for session '{sessionId}'. Strategy: {overflowStrategy}.")
+        {
+            SessionId = sessionId;
+            OverflowStrategy = overflowStrategy;
+        }
+    }
+}

Source/MQTTnet.Server/Internal/MqttSession.cs

@@ -6,6 +6,7 @@
 using MQTTnet.Internal;
 using MQTTnet.Packets;
 using MQTTnet.Protocol;
+using MQTTnet.Server.Exceptions;
 
 namespace MQTTnet.Server.Internal;
 
@@ -109,10 +110,11 @@ public void EnqueueControlPacket(MqttPacketBusItem packetBusItem)
 
     public EnqueueDataPacketResult EnqueueDataPacket(MqttPacketBusItem packetBusItem)
     {
-        if (_packetBus.ItemsCount(MqttPacketBusPartition.Data) >= _serverOptions.MaxPendingMessagesPerClient)
+        if (PendingDataPacketsCount >= _serverOptions.MaxPendingMessagesPerClient)
         {
             if (_serverOptions.PendingMessagesOverflowStrategy == MqttPendingMessagesOverflowStrategy.DropNewMessage)
             {
+                packetBusItem.Fail(new MqttPendingMessagesOverflowException(Id, _serverOptions.PendingMessagesOverflowStrategy));
                 return EnqueueDataPacketResult.Dropped;
             }
 
@@ -121,10 +123,14 @@ public EnqueueDataPacketResult EnqueueDataPacket(MqttPacketBusItem packetBusItem
                 // Only drop from the data partition. Dropping from control partition might break the connection
                 // because the client does not receive PINGREQ packets etc. any longer.
                 var firstItem = _packetBus.DropFirstItem(MqttPacketBusPartition.Data);
-                if (firstItem != null && _eventContainer.QueuedApplicationMessageOverwrittenEvent.HasHandlers)
+                if (firstItem != null)
                 {
-                    var eventArgs = new QueueMessageOverwrittenEventArgs(Id, firstItem.Packet);
-                    _eventContainer.QueuedApplicationMessageOverwrittenEvent.InvokeAsync(eventArgs).ConfigureAwait(false);
+                    firstItem.Fail(new MqttPendingMessagesOverflowException(Id, _serverOptions.PendingMessagesOverflowStrategy));
+                    if (_eventContainer.QueuedApplicationMessageOverwrittenEvent.HasHandlers)
+                    {
+                        var eventArgs = new QueueMessageOverwrittenEventArgs(Id, firstItem.Packet);
+                        _eventContainer.QueuedApplicationMessageOverwrittenEvent.InvokeAsync(eventArgs).ConfigureAwait(false);
+                    }
                 }
             }
         }

Source/MQTTnet.Server/Status/MqttSessionStatus.cs

@@ -4,6 +4,7 @@
 
 using System.Collections;
 using MQTTnet.Internal;
+using MQTTnet.Packets;
 using MQTTnet.Server.Internal;
 using MQTTnet.Server.Internal.Formatter;
 
@@ -40,22 +41,57 @@ public Task DeleteAsync()
         return _session.DeleteAsync();
     }
 
-    public Task DeliverApplicationMessageAsync(MqttApplicationMessage applicationMessage)
+    /// <summary>
+    /// Delivers an application message immediately to the session.
+    /// </summary>
+    /// <param name="applicationMessage">The application message to deliver.</param>
+    /// <returns>
+    /// A task that represents the asynchronous operation. The result contains the delivered MQTT publish packet.
+    /// </returns>
+    public async Task<MqttPublishPacket> DeliverApplicationMessageAsync(MqttApplicationMessage applicationMessage)
     {
         ArgumentNullException.ThrowIfNull(applicationMessage);
 
         var packetBusItem = new MqttPacketBusItem(MqttPublishPacketFactory.Create(applicationMessage));
         _session.EnqueueDataPacket(packetBusItem);
 
-        return packetBusItem.WaitAsync();
+        var mqttPacket = await packetBusItem.WaitAsync();
+
+        return (MqttPublishPacket)mqttPacket;
     }
 
-    public Task EnqueueApplicationMessageAsync(MqttApplicationMessage applicationMessage)
+    /// <summary>
+    /// Attempts to enqueue an application message to the session's send buffer.
+    /// </summary>
+    /// <param name="applicationMessage">The application message to enqueue.</param>
+    /// <param name="publishPacket">The corresponding MQTT publish packed, if the operation was successful.</param>
+    /// <returns><c>true</c> if the message was successfully enqueued; otherwise, <c>false</c>.</returns>
+    /// <remarks>
+    /// When <see cref="MqttServerOptions.PendingMessagesOverflowStrategy"/> is set to <see cref="MqttPendingMessagesOverflowStrategy.DropOldestQueuedMessage"/>,
+    /// this method always returns <c>true</c>.
+    /// However, an existing message in the queue may be <b>dropped later</b> to make room for the newly enqueued message.
+    /// Such dropped messages can be tracked by subscribing to <see cref="MqttServer.QueuedApplicationMessageOverwrittenAsync"/> event.
+    /// </remarks>
+    public bool TryEnqueueApplicationMessage(MqttApplicationMessage applicationMessage, out MqttPublishPacket publishPacket)
     {
         ArgumentNullException.ThrowIfNull(applicationMessage);
 
-        _session.EnqueueDataPacket(new MqttPacketBusItem(MqttPublishPacketFactory.Create(applicationMessage)));
+        publishPacket = MqttPublishPacketFactory.Create(applicationMessage);
+        var enqueueDataPacketResult = _session.EnqueueDataPacket(new MqttPacketBusItem(publishPacket));
 
+        if (enqueueDataPacketResult == EnqueueDataPacketResult.Enqueued)
+        {
+            return true;
+        }
+
+        publishPacket = null;
+        return false;
+    }
+
+    [Obsolete("This method is obsolete. Use TryEnqueueApplicationMessage instead.")]
+    public Task EnqueueApplicationMessageAsync(MqttApplicationMessage applicationMessage)
+    {
+        TryEnqueueApplicationMessage(applicationMessage, out _);
         return CompletedTask.Instance;
     }
 }

Source/MQTTnet/Internal/MqttPacketBus.cs

@@ -141,14 +141,6 @@ public List<MqttPacket> ExportPackets(MqttPacketBusPartition partition)
             }
         }
 
-        public int ItemsCount(MqttPacketBusPartition partition)
-        {
-            lock (_syncRoot)
-            {
-                return _partitions[(int)partition].Count;
-            }
-        }
-
         public int PartitionItemsCount(MqttPacketBusPartition partition)
         {
             lock (_syncRoot)

Source/MQTTnet/Internal/MqttPacketBusItem.cs

@@ -10,8 +10,8 @@ namespace MQTTnet.Internal
 {
     public sealed class MqttPacketBusItem
     {
-        readonly AsyncTaskCompletionSource<bool> _promise = new AsyncTaskCompletionSource<bool>();
-        
+        readonly AsyncTaskCompletionSource<MqttPacket> _promise = new AsyncTaskCompletionSource<MqttPacket>();
+
         public MqttPacketBusItem(MqttPacket packet)
         {
             Packet = packet ?? throw new ArgumentNullException(nameof(packet));
@@ -28,7 +28,7 @@ public void Cancel()
 
         public void Complete()
         {
-            _promise.TrySetResult(true);
+            _promise.TrySetResult(Packet);
             Completed?.Invoke(this, EventArgs.Empty);
         }
 
@@ -37,7 +37,7 @@ public void Fail(Exception exception)
             _promise.TrySetException(exception);
         }
 
-        public Task WaitAsync()
+        public Task<MqttPacket> WaitAsync()
         {
             return _promise.Task;
         }