fix: cancel non-resumable MQTT operations

Author: cartertinney

azure-iot-device/azure/iot/device/common/mqtt_transport.py

@@ -726,7 +726,9 @@ def disconnect(self, clear_inflight=False):
         """
         Disconnect from the MQTT Server and wait for the network loop to stop.
 
-        Optionally, clear any inflight operation tracking if clear_inflight is True.
+        If clear_inflight is True, complete all tracked operations as cancelled. Otherwise,
+        preserve resumable QoS 1 and QoS 2 publishes and complete non-resumable operations as
+        cancelled.
 
         :raises: ProtocolClientError if there is some client error.
         :raises: ConnectionDroppedError in unexpected cases.
@@ -759,7 +761,7 @@ def disconnect(self, clear_inflight=False):
                 if clear_inflight:
                     self._op_manager.complete_all_tracked_operations_as_cancelled()
                 else:
-                    self._op_manager.stop_tracking_non_publish_operations()
+                    self._op_manager.complete_non_resumable_operations_as_cancelled()
                 if self._connection_lifecycle:
                     self._connection_lifecycle.finish_disconnect()
             else:
@@ -774,7 +776,7 @@ def disconnect(self, clear_inflight=False):
             if clear_inflight:
                 self._op_manager.complete_all_tracked_operations_as_cancelled()
             else:
-                self._op_manager.stop_tracking_non_publish_operations()
+                self._op_manager.complete_non_resumable_operations_as_cancelled()
             if self._connection_lifecycle:
                 self._connection_lifecycle.finish_disconnect()
 
@@ -893,12 +895,20 @@ def publish(self, topic, payload, qos=1, callback=None):
                 "Paho retained QoS {} PUBLISH with MID {} for the next connection".format(qos, mid)
             )
         self._op_manager.register_operation(
-            mid=mid, callback=callback, operation_type=OperationType.PUBLISH
+            mid=mid,
+            callback=callback,
+            operation_type={
+                0: OperationType.PUBLISH_QOS_0,
+                1: OperationType.PUBLISH_QOS_1,
+                2: OperationType.PUBLISH_QOS_2,
+            }[qos],
         )
 
 
 class OperationType(Enum):
-    PUBLISH = "PUBLISH"
+    PUBLISH_QOS_0 = "PUBLISH_QOS_0"
+    PUBLISH_QOS_1 = "PUBLISH_QOS_1"
+    PUBLISH_QOS_2 = "PUBLISH_QOS_2"
     SUBSCRIBE = "SUBSCRIBE"
     UNSUBSCRIBE = "UNSUBSCRIBE"
 
@@ -1073,22 +1083,26 @@ def complete_operation(self, mid, error=None):
                 # Completion callbacks are optional.
                 logger.debug("No callback set for Paho MID {}".format(mid))
 
-    def stop_tracking_non_publish_operations(self):
-        """Stop tracking SUBSCRIBE and UNSUBSCRIBE operations without invoking callbacks.
+    def complete_non_resumable_operations_as_cancelled(self):
+        """Complete operations Paho cannot resume as cancelled.
 
-        Paho does not retain these operations for a later connection. PUBLISH operations remain
-        tracked because Paho owns their MQTT 3.1.1 QoS retransmission state.
+        Paho retains MQTT 3.1.1 QoS 1 and QoS 2 PUBLISH operations for a later connection.
+        SUBSCRIBE, UNSUBSCRIBE, and QoS 0 PUBLISH operations are not retained by Paho, so remove
+        their tracking, tombstone their MIDs, and invoke their callbacks with ``cancelled=True``.
         """
+        logger.debug("Completing non-resumable tracked operations as cancelled")
         with self._lock:
-            matching_mids = [
-                mid
+            pending_ops = [
+                (mid, pending_operation)
                 for mid, pending_operation in self._pending_operations.items()
                 if pending_operation.operation_type
-                in (OperationType.SUBSCRIBE, OperationType.UNSUBSCRIBE)
+                not in (OperationType.PUBLISH_QOS_1, OperationType.PUBLISH_QOS_2)
             ]
-            for mid in matching_mids:
+            for mid, _ in pending_ops:
                 del self._pending_operations[mid]
-            self._cancelled_operation_mids.update(matching_mids)
+            self._cancelled_operation_mids.update(mid for mid, _ in pending_ops)
+
+        self._defer_or_invoke_cancellation_callbacks(pending_ops)
 
     def complete_all_tracked_operations_as_cancelled(self):
         """Complete all tracked SDK operations as cancelled and clear unknown completions.
@@ -1107,13 +1121,16 @@ def complete_all_tracked_operations_as_cancelled(self):
             self._pending_operations.clear()
             self._unknown_operation_completions.clear()
 
+        self._defer_or_invoke_cancellation_callbacks(pending_ops)
+
+    def _defer_or_invoke_cancellation_callbacks(self, pending_operations):
         deferred_operations = getattr(
             self._deferred_cancellation_callbacks, "pending_operations", None
         )
         if deferred_operations is not None:
-            deferred_operations.extend(pending_ops)
+            deferred_operations.extend(pending_operations)
         else:
-            self._invoke_cancellation_callbacks(pending_ops)
+            self._invoke_cancellation_callbacks(pending_operations)
 
     def _invoke_cancellation_callbacks(self, pending_operations):
         """Invoke callbacks for operations whose tracking was cancelled."""
@@ -1140,6 +1157,7 @@ def _invoke_cancellation_callbacks(self, pending_operations):
 
 # TODO: Clarify hard-disconnect semantics because cancelling an SDK publish operation does not
 # prevent Paho from delivering a retained QoS 1 or QoS 2 message after a later connection.
+# Re-evaluate the inclusion of "hard" disconnect.
 
 # NOTE: Connection lifecycle calls are deliberately serialized here and by ConnectionStateStage.
 # CONNECTION_TIMEOUT bounds the wait for CONNACK, allowing queued lifecycle operations such as

azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_mqtt.py

@@ -330,11 +330,11 @@ def _reconcile_mqtt_operation_tracking_after_connection_drop(self):
             self.transport._op_manager.complete_all_tracked_operations_as_cancelled()
         else:
             logger.debug(
-                "{}: Connection Retry enabled - preserving PUBLISH tracking and stopping SUBSCRIBE and UNSUBSCRIBE tracking".format(
+                "{}: Connection Retry enabled - preserving resumable PUBLISH tracking and completing non-resumable MQTT operations as cancelled".format(
                     self.name
                 )
             )
-            self.transport._op_manager.stop_tracking_non_publish_operations()
+            self.transport._op_manager.complete_non_resumable_operations_as_cancelled()
 
     @pipeline_thread.runs_on_pipeline_thread
     def _complete_pending_connection_op_after_disconnect(self, cause=None):

tests/unit/common/pipeline/test_pipeline_stages_mqtt.py

@@ -645,7 +645,10 @@ def test_does_not_apply_connection_drop_handling(self, stage, op, fake_pipeline_
         assert (
             stage.transport._op_manager.complete_all_tracked_operations_as_cancelled.call_count == 0
         )
-        assert stage.transport._op_manager.stop_tracking_non_publish_operations.call_count == 0
+        assert (
+            stage.transport._op_manager.complete_non_resumable_operations_as_cancelled.call_count
+            == 0
+        )
         assert stage.report_background_exception.call_count == 0
 
     @pytest.mark.it(
@@ -966,13 +969,19 @@ def test_connection_drop_handling(
                 stage.transport._op_manager.complete_all_tracked_operations_as_cancelled.call_count
                 == 0
             )
-            assert stage.transport._op_manager.stop_tracking_non_publish_operations.call_count == 1
+            assert (
+                stage.transport._op_manager.complete_non_resumable_operations_as_cancelled.call_count
+                == 1
+            )
         else:
             assert (
                 stage.transport._op_manager.complete_all_tracked_operations_as_cancelled.call_count
                 == 1
             )
-            assert stage.transport._op_manager.stop_tracking_non_publish_operations.call_count == 0
+            assert (
+                stage.transport._op_manager.complete_non_resumable_operations_as_cancelled.call_count
+                == 0
+            )
 
         assert stage.report_background_exception.call_count == 1
         background_exception = stage.report_background_exception.call_args.args[0]
@@ -992,7 +1001,7 @@ def test_connection_drop_handling(
             ),
             pytest.param(
                 True,
-                "stop_tracking_non_publish_operations",
+                "complete_non_resumable_operations_as_cancelled",
                 id="Connection retry enabled",
             ),
         ],
@@ -1084,22 +1093,24 @@ def test_completes_tracked_operations_without_retry(self, mocker, stage, arbitra
         assert mock_cancel.call_args == mocker.call()
 
     @pytest.mark.it(
-        "Preserves publishes and stops tracking other MQTT operations if connection retry is enabled"
+        "Preserves resumable publishes and cancels other MQTT operations if connection retry is enabled"
     )
-    def test_preserves_publish_tracking_with_retry(self, mocker, stage, arbitrary_exception):
+    def test_cancels_non_resumable_operations_with_retry(self, mocker, stage, arbitrary_exception):
         stage.transport._op_manager = mocker.MagicMock()
         mock_cancel = stage.transport._op_manager.complete_all_tracked_operations_as_cancelled
-        mock_stop_non_publish = stage.transport._op_manager.stop_tracking_non_publish_operations
+        mock_cancel_non_resumable = (
+            stage.transport._op_manager.complete_non_resumable_operations_as_cancelled
+        )
         stage.nucleus.pipeline_configuration.connection_retry = True
         assert stage._pending_connection_op is None
         assert mock_cancel.call_count == 0
-        assert mock_stop_non_publish.call_count == 0
+        assert mock_cancel_non_resumable.call_count == 0
 
         # Trigger disconnect
         stage.transport.on_mqtt_connection_dropped_handler(arbitrary_exception)
 
         assert mock_cancel.call_count == 0
-        assert mock_stop_non_publish.call_args == mocker.call()
+        assert mock_cancel_non_resumable.call_args == mocker.call()
 
     @pytest.mark.it("Raises a ConnectionDroppedError as a background exception")
     def test_background_exception_raised(self, stage, arbitrary_exception):