fix: release cancelled MQTT operations

Remove transport callbacks when operations complete through timeout or retry paths, and discard Paho QoS state when the SDK cancels in-flight work. This prevents cancelled publish and subscribe objects from surviving dropped-connection tests or being completed again by late acknowledgements.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: cartertinney

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

@@ -376,7 +376,7 @@ def shutdown(self):
         self._mqtt_client.on_disconnect = None
         # Now disconnect and do some additional cleanup.
         self._force_transport_disconnect_and_cleanup()
-        self._op_manager.cancel_all_operations()
+        self.cancel_all_operations()
 
     def connect(self, password=None):
         """
@@ -480,7 +480,7 @@ def disconnect(self, clear_inflight=False):
                 )
                 # Still clear inflight operations since we're effectively disconnected
                 if clear_inflight:
-                    self._op_manager.cancel_all_operations()
+                    self.cancel_all_operations()
             else:
                 # This could result in ConnectionDroppedError or ProtocolClientError
                 err = _create_error_from_rc_code(rc)
@@ -491,7 +491,21 @@ def disconnect(self, clear_inflight=False):
             # cause a force disconnect via the on_disconnect handler, thus it is safe to clear
             # ops here and now.
             if clear_inflight:
-                self._op_manager.cancel_all_operations()
+                self.cancel_all_operations()
+
+    def cancel_operation(self, callback):
+        """Stop tracking the operation associated with a transport callback."""
+        self._op_manager.cancel_operation(callback)
+
+    def cancel_all_operations(self):
+        """Cancel SDK operations and discard corresponding Paho QoS state."""
+        self._op_manager.cancel_all_operations()
+
+        # Paho retains QoS 1/2 messages for a future reconnect. Once the SDK has cancelled the
+        # corresponding operations, those messages must not be retried or kept alive.
+        with self._mqtt_client._out_message_mutex:
+            self._mqtt_client._out_messages.clear()
+            self._mqtt_client._inflight_messages = 0
 
     def subscribe(self, topic, qos=1, callback=None):
         """
@@ -509,7 +523,7 @@ def subscribe(self, topic, qos=1, callback=None):
         """
         logger.info("subscribing to {} with qos {}".format(topic, qos))
         try:
-            (rc, mid) = self._mqtt_client.subscribe(topic, qos=qos)
+            rc, mid = self._mqtt_client.subscribe(topic, qos=qos)
         except ValueError:
             raise
         except Exception as e:
@@ -534,7 +548,7 @@ def unsubscribe(self, topic, callback=None):
         """
         logger.info("unsubscribing from {}".format(topic))
         try:
-            (rc, mid) = self._mqtt_client.unsubscribe(topic)
+            rc, mid = self._mqtt_client.unsubscribe(topic)
         except ValueError:
             raise
         except Exception as e:
@@ -568,7 +582,7 @@ def publish(self, topic, payload, qos=1, callback=None):
         """
         logger.info("publishing on {}".format(topic))
         try:
-            (rc, mid) = self._mqtt_client.publish(topic=topic, payload=payload, qos=qos)
+            rc, mid = self._mqtt_client.publish(topic=topic, payload=payload, qos=qos)
         except ValueError:
             raise
         except TypeError:
@@ -660,9 +674,9 @@ def complete_operation(self, mid):
             else:
                 # Otherwise, store the mid as an unknown response
                 logger.debug("Response received for unknown MID: {}".format(mid))
-                self._unknown_operation_completions[
-                    mid
-                ] = mid  # TODO: set something more useful here
+                self._unknown_operation_completions[mid] = (
+                    mid  # TODO: set something more useful here
+                )
 
         # Now that the lock has been released, if the callback should be triggered,
         # go ahead and trigger it now.
@@ -680,6 +694,17 @@ def complete_operation(self, mid):
                 # fully expected.  QOS=1 means we might get 2 PUBACKs
                 logger.debug("No callback set for MID: {}".format(mid))
 
+    def cancel_operation(self, callback):
+        """Remove pending operations associated with a callback without invoking it."""
+        with self._lock:
+            matching_mids = [
+                mid
+                for mid, pending_callback in self._pending_operation_callbacks.items()
+                if pending_callback is callback
+            ]
+            for mid in matching_mids:
+                del self._pending_operation_callbacks[mid]
+
     def cancel_all_operations(self):
         """Complete all pending operations with cancellation, removing MID tracking"""
         logger.debug("Cancelling all pending operations")

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

@@ -124,6 +124,35 @@ def _cancel_connection_watchdog(self, op):
         except AttributeError:
             pass
 
+    def _create_transport_operation_callback(self, op, acknowledgement_name):
+        @pipeline_thread.invoke_on_pipeline_thread_nowait
+        def on_complete(cancelled=False):
+            if op.completed or op.completing:
+                logger.debug(
+                    "{}({}): Ignoring {} for an already-completed operation".format(
+                        self.name, op.name, acknowledgement_name
+                    )
+                )
+            elif cancelled:
+                op.complete(
+                    error=pipeline_exceptions.OperationCancelled(
+                        "Operation cancelled before {} received".format(acknowledgement_name)
+                    )
+                )
+            else:
+                logger.debug(
+                    "{}({}): {} received. completing op.".format(
+                        self.name, op.name, acknowledgement_name
+                    )
+                )
+                op.complete()
+
+        def remove_transport_callback(op, error):
+            self.transport.cancel_operation(on_complete)
+
+        op.add_callback(remove_transport_callback)
+        return on_complete
+
     @pipeline_thread.runs_on_pipeline_thread
     def _run_op(self, op):
         if isinstance(op, pipeline_ops_base.InitializePipelineOperation):
@@ -259,20 +288,7 @@ def on_disconnect_complete(op, error):
 
         elif isinstance(op, pipeline_ops_mqtt.MQTTPublishOperation):
             logger.debug("{}({}): publishing on {}".format(self.name, op.name, op.topic))
-
-            @pipeline_thread.invoke_on_pipeline_thread_nowait
-            def on_complete(cancelled=False):
-                if cancelled:
-                    op.complete(
-                        error=pipeline_exceptions.OperationCancelled(
-                            "Operation cancelled before PUBACK received"
-                        )
-                    )
-                else:
-                    logger.debug(
-                        "{}({}): PUBACK received. completing op.".format(self.name, op.name)
-                    )
-                    op.complete()
+            on_complete = self._create_transport_operation_callback(op, "PUBACK")
 
             try:
                 self.transport.publish(topic=op.topic, payload=op.payload, callback=on_complete)
@@ -281,20 +297,7 @@ def on_complete(cancelled=False):
 
         elif isinstance(op, pipeline_ops_mqtt.MQTTSubscribeOperation):
             logger.debug("{}({}): subscribing to {}".format(self.name, op.name, op.topic))
-
-            @pipeline_thread.invoke_on_pipeline_thread_nowait
-            def on_complete(cancelled=False):
-                if cancelled:
-                    op.complete(
-                        error=pipeline_exceptions.OperationCancelled(
-                            "Operation cancelled before SUBACK received"
-                        )
-                    )
-                else:
-                    logger.debug(
-                        "{}({}): SUBACK received. completing op.".format(self.name, op.name)
-                    )
-                    op.complete()
+            on_complete = self._create_transport_operation_callback(op, "SUBACK")
 
             try:
                 self.transport.subscribe(topic=op.topic, callback=on_complete)
@@ -303,20 +306,7 @@ def on_complete(cancelled=False):
 
         elif isinstance(op, pipeline_ops_mqtt.MQTTUnsubscribeOperation):
             logger.debug("{}({}): unsubscribing from {}".format(self.name, op.name, op.topic))
-
-            @pipeline_thread.invoke_on_pipeline_thread_nowait
-            def on_complete(cancelled=False):
-                if cancelled:
-                    op.complete(
-                        error=pipeline_exceptions.OperationCancelled(
-                            "Operation cancelled before UNSUBACK received"
-                        )
-                    )
-                else:
-                    logger.debug(
-                        "{}({}): UNSUBACK received.  completing op.".format(self.name, op.name)
-                    )
-                    op.complete()
+            on_complete = self._create_transport_operation_callback(op, "UNSUBACK")
 
             try:
                 self.transport.unsubscribe(topic=op.topic, callback=on_complete)
@@ -451,11 +441,7 @@ def _on_mqtt_disconnected(self, cause=None):
                         self.name
                     )
                 )
-                # TODO: Remove private access to the op manager (this layer shouldn't know about it)
-                # This is a stopgap. I didn't want to invest too much infrastructure into a cancel flow
-                # given that future development of individual operation cancels might affect the
-                # approach to cancelling inflight ops waiting in the transport.
-                self.transport._op_manager.cancel_all_operations()
+                self.transport.cancel_all_operations()
 
             # Regardless of cause, it is now a ConnectionDroppedError. Log it and swallow it.
             # Higher layers will see that we're disconnected and may reconnect as necessary.

tests/unit/common/pipeline/test_pipeline_stages_mqtt.py

@@ -8,6 +8,7 @@
 import sys
 import threading
 from azure.iot.device.common import transport_exceptions, handle_exceptions
+from azure.iot.device.common.mqtt_transport import OperationManager
 from azure.iot.device.common.pipeline import (
     pipeline_ops_base,
     pipeline_ops_mqtt,
@@ -616,6 +617,16 @@ def test_complete_with_cancel(self, mocker, stage, op):
         assert op.completed
         assert isinstance(op.error, pipeline_exceptions.OperationCancelled)
 
+    @pytest.mark.it("Removes its transport callback when completed by another stage")
+    def test_removes_transport_callback_on_external_completion(self, stage, op):
+        stage.run_op(op)
+        transport_callback = stage.transport.publish.call_args[1]["callback"]
+
+        op.complete()
+        transport_callback()
+
+        stage.transport.cancel_operation.assert_called_once_with(transport_callback)
+
     @pytest.mark.it(
         "Completes the operation using the exception that was raised, if an exception was raised from the MQTTTransport"
     )
@@ -676,6 +687,23 @@ def test_complete_with_cancel(self, mocker, stage, op):
         assert op.completed
         assert isinstance(op.error, pipeline_exceptions.OperationCancelled)
 
+    @pytest.mark.it("Removes its transport callback when completed by another stage")
+    def test_removes_transport_callback_on_external_completion(self, stage, op):
+        manager = OperationManager()
+        stage.transport.subscribe.side_effect = lambda topic, callback: manager.establish_operation(
+            mid=1, callback=callback
+        )
+        stage.transport.cancel_operation.side_effect = manager.cancel_operation
+        stage.run_op(op)
+        transport_callback = stage.transport.subscribe.call_args[1]["callback"]
+        assert manager._pending_operation_callbacks == {1: transport_callback}
+
+        op.complete()
+        transport_callback()
+
+        stage.transport.cancel_operation.assert_called_once_with(transport_callback)
+        assert manager._pending_operation_callbacks == {}
+
     @pytest.mark.it(
         "Completes the operation using the exception that was raised, if an exception was raised from the MQTTTransport"
     )
@@ -736,6 +764,16 @@ def test_complete_with_cancel(self, mocker, stage, op):
         assert op.completed
         assert isinstance(op.error, pipeline_exceptions.OperationCancelled)
 
+    @pytest.mark.it("Removes its transport callback when completed by another stage")
+    def test_removes_transport_callback_on_external_completion(self, stage, op):
+        stage.run_op(op)
+        transport_callback = stage.transport.unsubscribe.call_args[1]["callback"]
+
+        op.complete()
+        transport_callback()
+
+        stage.transport.cancel_operation.assert_called_once_with(transport_callback)
+
     @pytest.mark.it(
         "Completes the operation using the exception that was raised, if an exception was raised from the MQTTTransport"
     )
@@ -1168,25 +1206,22 @@ def cause(self, request, arbitrary_exception):
     @pytest.mark.it(
         "Cancels all in-flight operations in the transport, if connection retry has been disabled"
     )
-    def test_inflight_no_retry(self, mocker, stage, cause):
-        stage.transport._op_manager = mocker.MagicMock()
-        mock_cancel = stage.transport._op_manager.cancel_all_operations
+    def test_inflight_no_retry(self, stage, cause):
+        mock_cancel = stage.transport.cancel_all_operations
         stage.nucleus.pipeline_configuration.connection_retry = False
         assert stage._pending_connection_op is None
         assert mock_cancel.call_count == 0
 
         # Trigger disconnect
         stage.transport.on_mqtt_disconnected_handler(cause)
 
-        assert mock_cancel.call_count == 1
-        assert mock_cancel.call_args == mocker.call()
+        mock_cancel.assert_called_once_with()
 
     @pytest.mark.it(
         "Does not cancel any in-flight operations in the transport if connection retry has been enabled"
     )
-    def test_inflight_unexpected_with_retry(self, mocker, stage, cause):
-        stage.transport._op_manager = mocker.MagicMock()
-        mock_cancel = stage.transport._op_manager.cancel_all_operations
+    def test_inflight_unexpected_with_retry(self, stage, cause):
+        mock_cancel = stage.transport.cancel_all_operations
         stage.nucleus.pipeline_configuration.connection_retry = True
         assert stage._pending_connection_op is None
         assert mock_cancel.call_count == 0

tests/unit/common/test_mqtt_transport.py

@@ -1008,6 +1008,22 @@ def test_does_not_set_thread_to_none_on_exception_in_non_paho_thread(
         assert mock_mqtt_client._thread is not None
 
 
+@pytest.mark.describe("MQTTTransport - .cancel_all_operations()")
+class TestCancelAllOperations(object):
+    @pytest.mark.it("Cancels SDK operations and clears Paho outgoing message state")
+    def test_clears_sdk_and_paho_operations(self, mocker, mock_mqtt_client, transport):
+        cancel_sdk_operations = mocker.patch.object(transport._op_manager, "cancel_all_operations")
+        mock_mqtt_client._out_message_mutex = threading.Lock()
+        mock_mqtt_client._out_messages = {1: mqtt.MQTTMessage(mid=1)}
+        mock_mqtt_client._inflight_messages = 1
+
+        transport.cancel_all_operations()
+
+        cancel_sdk_operations.assert_called_once_with()
+        assert mock_mqtt_client._out_messages == {}
+        assert mock_mqtt_client._inflight_messages == 0
+
+
 @pytest.mark.describe("MQTTTransport - OCCURRENCE: Disconnect Completed")
 class TestEventDisconnectCompleted(object):
     @pytest.fixture(
@@ -2486,6 +2502,24 @@ def stop_tracking_mocks(*args):
         assert mocker.call.cb() not in calls_during_lock
 
 
+@pytest.mark.describe("OperationManager - .cancel_operation()")
+class TestOperationManagerCancelOperation(object):
+    @pytest.mark.it("Removes only operations associated with the provided callback")
+    def test_remove_matching_pending_operations(self, mocker):
+        manager = OperationManager()
+        callback_to_cancel = mocker.MagicMock()
+        callback_to_keep = mocker.MagicMock()
+        manager.establish_operation(mid=1, callback=callback_to_cancel)
+        manager.establish_operation(mid=2, callback=callback_to_keep)
+        manager.establish_operation(mid=3, callback=callback_to_cancel)
+
+        manager.cancel_operation(callback_to_cancel)
+
+        assert manager._pending_operation_callbacks == {2: callback_to_keep}
+        callback_to_cancel.assert_not_called()
+        callback_to_keep.assert_not_called()
+
+
 @pytest.mark.describe("OperationManager - .cancel_all_operations()")
 class TestOperationManagerCancelAllOperations(object):
     @pytest.mark.it("Removes all MID tracking for all pending operations")