[v2 - bug fix] Queue MQTT C2D messages if callback not set

iothub/device/src/IotHubBaseClient.cs

@@ -338,6 +338,10 @@ public async Task SetIncomingMessageCallbackAsync(
             finally
             {
                 _receiveMessageSemaphore.Release();
+                if (messageCallback != null)
+                {
+                    await InnerHandler.EnsurePendingMessagesAreDeliveredAsync(cancellationToken).ConfigureAwait(false);
+                }
 
                 if (Logging.IsEnabled)
                     Logging.Exit(this, messageCallback, nameof(SetIncomingMessageCallbackAsync));
@@ -780,18 +784,7 @@ internal async Task<MessageAcknowledgement> OnMessageReceivedAsync(IncomingMessa
 
             try
             {
-                Func<IncomingMessage, Task<MessageAcknowledgement>> callback = _receiveMessageCallback;
-
-                if (callback != null)
-                {
-                    return await callback.Invoke(message).ConfigureAwait(false);
-                }
-
-                // The SDK should only receive messages when the user sets a listener, so this should never happen.
-                if (Logging.IsEnabled)
-                    Logging.Error(this, $"Received a message when no listener was set. Abandoning message with message Id: {message.MessageId}.", nameof(OnMessageReceivedAsync));
-
-                return MessageAcknowledgement.Abandon;
+                return await _receiveMessageCallback.Invoke(message).ConfigureAwait(false);
             }
             catch (Exception ex)
             {

iothub/device/src/Pipeline/AmqpTransportHandler.cs

@@ -243,6 +243,13 @@ public override async Task EnableReceiveMessageAsync(CancellationToken cancellat
             }
         }
 
+        // This method is added to ensure that over MQTT devices can receive messages that were sent when it was disconnected.
+        // This behavior is available by default over AMQP, so no additional implementation is required here.
+        public override Task EnsurePendingMessagesAreDeliveredAsync(CancellationToken cancellationToken)
+        {
+            return Task.CompletedTask;
+        }
+
         public override async Task DisableReceiveMessageAsync(CancellationToken cancellationToken)
         {
             if (Logging.IsEnabled)

iothub/device/src/Pipeline/DefaultDelegatingHandler.cs

@@ -70,6 +70,14 @@ public virtual Task EnableReceiveMessageAsync(CancellationToken cancellationToke
             return NextHandler?.EnableReceiveMessageAsync(cancellationToken) ?? Task.CompletedTask;
         }
 
+        // This is to ensure that if device connects over MQTT with CleanSession flag set to false,
+        // then any message sent while the device was disconnected is delivered on the callback.
+        public virtual Task EnsurePendingMessagesAreDeliveredAsync(CancellationToken cancellationToken)
+        {
+            ThrowIfDisposed();
+            return NextHandler?.EnsurePendingMessagesAreDeliveredAsync(cancellationToken);
+        }
+
         public virtual Task DisableReceiveMessageAsync(CancellationToken cancellationToken)
         {
             ThrowIfDisposed();

iothub/device/src/Pipeline/ExceptionRemappingHandler.cs

@@ -45,6 +45,13 @@ public override Task EnableReceiveMessageAsync(CancellationToken cancellationTok
             return ExecuteWithExceptionRemappingAsync(() => base.EnableReceiveMessageAsync(cancellationToken));
         }
 
+        // This is to ensure that if device connects over MQTT with CleanSession flag set to false,
+        // then any message sent while the device was disconnected is delivered on the callback.
+        public override Task EnsurePendingMessagesAreDeliveredAsync(CancellationToken cancellationToken)
+        {
+            return ExecuteWithExceptionRemappingAsync(() => base.EnsurePendingMessagesAreDeliveredAsync(cancellationToken));
+        }
+
         public override Task DisableReceiveMessageAsync(CancellationToken cancellationToken)
         {
             return ExecuteWithExceptionRemappingAsync(() => base.DisableReceiveMessageAsync(cancellationToken));

iothub/device/src/Pipeline/IDelegatingHandler.cs

@@ -26,6 +26,10 @@ internal interface IDelegatingHandler : IContinuationProvider<IDelegatingHandler
 
         Task EnableReceiveMessageAsync(CancellationToken cancellationToken);
 
+        // This is to ensure that if device connects over MQTT with CleanSession flag set to false,
+        // then any message sent while the device was disconnected is delivered on the callback.
+        Task EnsurePendingMessagesAreDeliveredAsync(CancellationToken cancellationToken);
+
         Task DisableReceiveMessageAsync(CancellationToken cancellationToken);
 
         // Methods.

iothub/device/src/Pipeline/MqttTransportHandler.cs

@@ -33,6 +33,7 @@ internal sealed class MqttTransportHandler : TransportHandlerBase, IDisposable
         private const string ModuleToCloudMessagesTopicFormat = "devices/{0}/modules/{1}/messages/events/";
         private readonly string _deviceToCloudMessagesTopic;
         private readonly string _moduleToCloudMessagesTopic;
+        private readonly ConcurrentQueue<IncomingMessage> _messageQueue;
 
         // Topic names for receiving cloud-to-device messages.
 
@@ -103,6 +104,8 @@ internal sealed class MqttTransportHandler : TransportHandlerBase, IDisposable
         private readonly Timer _twinTimeoutTimer;
 
         private bool _isSubscribedToTwinResponses;
+        private bool _isDeviceReceiveMessageCallbackSet;
+        private int _messageQueueSize;
 
         private readonly string _deviceId;
         private readonly string _moduleId;
@@ -178,9 +181,11 @@ internal MqttTransportHandler(PipelineContext context, IDelegatingHandler nextHa
             var mqttFactory = new MqttFactory(new MqttLogger());
             _mqttClient = mqttFactory.CreateMqttClient();
             _mqttClientOptionsBuilder = new MqttClientOptionsBuilder();
+            _messageQueue = new ConcurrentQueue<IncomingMessage>();
 
             _hostName = context.IotHubConnectionCredentials.HostName;
             _connectionCredentials = context.IotHubConnectionCredentials;
+            _messageQueueSize = _mqttTransportSettings.MessageQueueSize;
 
             if (_mqttTransportSettings.Protocol == IotHubClientTransportProtocol.Tcp)
             {
@@ -579,6 +584,7 @@ public override async Task EnableReceiveMessageAsync(CancellationToken cancellat
                         await SubscribeAsync(_moduleEventMessageTopic, cancellationToken).ConfigureAwait(false);
                     }
                 }
+                _isDeviceReceiveMessageCallbackSet = true;
             }
             catch (Exception ex) when (ex is not IotHubClientException && ex is not OperationCanceledException)
             {
@@ -594,6 +600,26 @@ public override async Task EnableReceiveMessageAsync(CancellationToken cancellat
             }
         }
 
+        public override async Task EnsurePendingMessagesAreDeliveredAsync(CancellationToken cancellationToken)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+
+            // If the device connects with a CleanSession flag set to false, we will need to deliver the messages
+            // that were sent before the client had subscribed to the C2D message receive topic.
+            if (!_mqttTransportSettings.CleanSession)
+            {
+                IncomingMessage message = null;
+                // Received C2D messages are enqueued into _messageQueue.
+                while (!_messageQueue.IsEmpty)
+                {
+                    if(_messageQueue.TryDequeue(out message))
+                    {
+                        await HandleReceivedCloudToDeviceMessageAsync(message).ConfigureAwait(false);
+                    }
+                }
+            }
+        }
+
         public override async Task DisableReceiveMessageAsync(CancellationToken cancellationToken)
         {
             if (Logging.IsEnabled)
@@ -616,6 +642,7 @@ public override async Task DisableReceiveMessageAsync(CancellationToken cancella
                         await UnsubscribeAsync(_moduleEventMessageTopic, cancellationToken).ConfigureAwait(false);
                     }
                 }
+                _isDeviceReceiveMessageCallbackSet = false;
             }
             catch (Exception ex) when (ex is not IotHubClientException && ex is not OperationCanceledException)
             {
@@ -1054,14 +1081,31 @@ private Task HandleDisconnectionAsync(MqttClientDisconnectedEventArgs disconnect
             return Task.CompletedTask;
         }
 
+        private IncomingMessage ProcessC2DMessage(MqttApplicationMessageReceivedEventArgs receivedEventArgs)
+        {
+            byte[] payload = receivedEventArgs.ApplicationMessage.Payload;
+
+            var receivedCloudToDeviceMessage = new IncomingMessage(payload)
+            {
+                PayloadConvention = _payloadConvention,
+            };
+
+            PopulateMessagePropertiesFromMqttMessage(receivedCloudToDeviceMessage, receivedEventArgs.ApplicationMessage);
+
+            return receivedCloudToDeviceMessage;
+        }
+
         private async Task HandleReceivedMessageAsync(MqttApplicationMessageReceivedEventArgs receivedEventArgs)
         {
             receivedEventArgs.AutoAcknowledge = false;
             string topic = receivedEventArgs.ApplicationMessage.Topic;
 
             if (topic.StartsWith(_deviceBoundMessagesTopic, StringComparison.InvariantCulture))
             {
-                await HandleReceivedCloudToDeviceMessageAsync(receivedEventArgs).ConfigureAwait(false);
+                // If C2D message callback is not set, messages are added to a queue to be processed later.
+                // However, the messages are still being ack'ed right away.
+                await HandleReceivedCloudToDeviceMessageAsync(ProcessC2DMessage(receivedEventArgs)).ConfigureAwait(false);
+                await receivedEventArgs.AcknowledgeAsync(CancellationToken.None).ConfigureAwait(false);
             }
             else if (topic.StartsWith(TwinDesiredPropertiesPatchTopic, StringComparison.InvariantCulture))
             {
@@ -1088,18 +1132,9 @@ private async Task HandleReceivedMessageAsync(MqttApplicationMessageReceivedEven
             }
         }
 
-        private async Task HandleReceivedCloudToDeviceMessageAsync(MqttApplicationMessageReceivedEventArgs receivedEventArgs)
+        private async Task HandleReceivedCloudToDeviceMessageAsync(IncomingMessage receivedCloudToDeviceMessage)
         {
-            byte[] payload = receivedEventArgs.ApplicationMessage.Payload;
-
-            var receivedCloudToDeviceMessage = new IncomingMessage(payload)
-            {
-                PayloadConvention = _payloadConvention,
-            };
-
-            PopulateMessagePropertiesFromMqttMessage(receivedCloudToDeviceMessage, receivedEventArgs.ApplicationMessage);
-
-            if (_messageReceivedListener != null)
+            if (_messageReceivedListener != null && _isDeviceReceiveMessageCallbackSet)
             {
                 MessageAcknowledgement acknowledgementType = await _messageReceivedListener.Invoke(receivedCloudToDeviceMessage);
 
@@ -1108,18 +1143,18 @@ private async Task HandleReceivedCloudToDeviceMessageAsync(MqttApplicationMessag
                     if (Logging.IsEnabled)
                         Logging.Error(this, "Cannot 'reject' or 'abandon' a received message over MQTT. Message will be acknowledged as 'complete' instead.");
                 }
-
-                // Note that MQTT does not support Abandon or Reject, so always Complete by acknowledging the message like this.
-                try
-                {
-                    await receivedEventArgs.AcknowledgeAsync(CancellationToken.None).ConfigureAwait(false);
-                }
-                catch (Exception ex)
+            }
+            else if(!_isDeviceReceiveMessageCallbackSet)
+            {
+                _messageQueue.Enqueue(receivedCloudToDeviceMessage);
+                while(_messageQueue.Count > _messageQueueSize)
                 {
-                    // This likely happened because the connection was lost. The service will re-send this message so the user
-                    // can acknowledge it on the new connection.
+                    _messageQueue.TryDequeue(out _);
                     if (Logging.IsEnabled)
-                        Logging.Error(this, $"Failed to send the acknowledgement for a received cloud to device message {ex}"); ;
+                    {
+                        Logging.Info(this, $"Queue size of {_messageQueueSize} for C2D messages has been reached, removing oldest queued C2D message. " +
+                            $"To avoid losing further messages, set SetIncomingMessageCallbackAsync() to process the messages or increase message queue size in IotHubClientMqttSettings.");
+                    }
                 }
             }
             else if (Logging.IsEnabled)

iothub/device/src/Pipeline/RetryDelegatingHandler.cs

@@ -221,6 +221,43 @@ await _internalRetryHandler
             }
         }
 
+        // This is to ensure that if device connects over MQTT with CleanSession flag set to false,
+        // then any message sent while the device was disconnected is delivered on the callback.
+        public override async Task EnsurePendingMessagesAreDeliveredAsync(CancellationToken cancellationToken)
+        {
+            try
+            {
+                if (Logging.IsEnabled)
+                    Logging.Enter(this, cancellationToken, nameof(EnsurePendingMessagesAreDeliveredAsync));
+
+                await _internalRetryHandler
+                    .RunWithRetryAsync(
+                        async () =>
+                        {
+                            // Wait to acquire the _cloudToDeviceMessageSubscriptionSemaphore. This ensures that concurrently invoked API calls are invoked in a thread-safe manner.
+                            await _cloudToDeviceMessageSubscriptionSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);
+
+                            try
+                            {
+                                // Ensure that a callback for receiving messages has been previously set.
+                                Debug.Assert(_deviceReceiveMessageEnabled);
+                                await base.EnsurePendingMessagesAreDeliveredAsync(cancellationToken).ConfigureAwait(false);
+                            }
+                            finally
+                            {
+                                _cloudToDeviceMessageSubscriptionSemaphore?.Release();
+                            }
+                        },
+                        cancellationToken)
+                    .ConfigureAwait(false);
+            }
+            finally
+            {
+                if (Logging.IsEnabled)
+                    Logging.Exit(this, cancellationToken, nameof(EnsurePendingMessagesAreDeliveredAsync));
+            }
+        }
+
         public override async Task DisableReceiveMessageAsync(CancellationToken cancellationToken)
         {
             if (Logging.IsEnabled)

iothub/device/src/TransportSettings/IotHubClientMqttSettings.cs

@@ -95,6 +95,15 @@ public IotHubClientMqttSettings(IotHubClientTransportProtocol transportProtocol
         /// </remarks>
         public RemoteCertificateValidationCallback RemoteCertificateValidationCallback { get; set; }
 
+        /// <summary>
+        /// C2D message queue size when callback has not been set.
+        /// </summary>
+        /// <remarks>
+        /// If C2D message callback is not set, messages from service are received and stored in a queue until a callback is set. 
+        /// If the number of messages sent is greater than queue size, older messages are removed as latest messages are added.
+        /// </remarks>
+        public int MessageQueueSize { get; set; } = 10;
+
         /// <summary>
         /// Used by Edge runtime to specify an authentication chain for Edge-to-Edge connections.
         /// </summary>
@@ -116,6 +125,7 @@ internal override IotHubClientTransportSettings Clone()
                 WillMessage = WillMessage,
                 RemoteCertificateValidationCallback = RemoteCertificateValidationCallback,
                 AuthenticationChain = AuthenticationChain,
+                MessageQueueSize = MessageQueueSize,
             };
         }
     }