Revert "Revert "revert: defer MQTT cancellation cleanup""

This reverts commit 364a6ca.

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.cancel_all_operations()
+        self._op_manager.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.cancel_all_operations()
+                    self._op_manager.cancel_all_operations()
             else:
                 # This could result in ConnectionDroppedError or ProtocolClientError
                 err = _create_error_from_rc_code(rc)
@@ -491,21 +491,7 @@ 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.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
+                self._op_manager.cancel_all_operations()
 
     def subscribe(self, topic, qos=1, callback=None):
         """
@@ -523,7 +509,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:
@@ -548,7 +534,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:
@@ -582,7 +568,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:
@@ -674,9 +660,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.
@@ -694,17 +680,6 @@ 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,35 +124,6 @@ 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):
@@ -288,7 +259,20 @@ def on_disconnect_complete(op, error):
 
         elif isinstance(op, pipeline_ops_mqtt.MQTTPublishOperation):
             logger.debug("{}({}): publishing on {}".format(self.name, op.name, op.topic))
-            on_complete = self._create_transport_operation_callback(op, "PUBACK")
+
+            @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()
 
             try:
                 self.transport.publish(topic=op.topic, payload=op.payload, callback=on_complete)
@@ -297,7 +281,20 @@ def on_disconnect_complete(op, error):
 
         elif isinstance(op, pipeline_ops_mqtt.MQTTSubscribeOperation):
             logger.debug("{}({}): subscribing to {}".format(self.name, op.name, op.topic))
-            on_complete = self._create_transport_operation_callback(op, "SUBACK")
+
+            @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()
 
             try:
                 self.transport.subscribe(topic=op.topic, callback=on_complete)
@@ -306,7 +303,20 @@ def on_disconnect_complete(op, error):
 
         elif isinstance(op, pipeline_ops_mqtt.MQTTUnsubscribeOperation):
             logger.debug("{}({}): unsubscribing from {}".format(self.name, op.name, op.topic))
-            on_complete = self._create_transport_operation_callback(op, "UNSUBACK")
+
+            @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()
 
             try:
                 self.transport.unsubscribe(topic=op.topic, callback=on_complete)
@@ -441,7 +451,11 @@ def _on_mqtt_disconnected(self, cause=None):
                         self.name
                     )
                 )
-                self.transport.cancel_all_operations()
+                # 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()
 
             # 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/e2e/iothub_e2e/aio/test_send_message.py

@@ -205,9 +205,8 @@ async def test_connects_after_automatic_disconnect_retry_disabled(
 
     @pytest.mark.it("Fails if connection disconnects before sending")
     @pytest.mark.uses_iptables
-    async def test_fails_if_disconnect_before_sending(
-        self, client, random_message, dropper, leak_tracker
-    ):
+    # TODO: Re-enable leak tracking after the MQTT cancellation refactor.
+    async def test_fails_if_disconnect_before_sending(self, client, random_message, dropper):
 
         assert client.connected
 
@@ -223,8 +222,9 @@ async def test_fails_if_disconnect_before_sending(
 
     @pytest.mark.it("Fails if connection drops before sending")
     @pytest.mark.uses_iptables
+    # TODO: Re-enable leak tracking after the MQTT cancellation refactor.
     async def test_fails_if_drop_before_sending_retry_disabled(
-        self, client, random_message, dropper, leak_tracker
+        self, client, random_message, dropper
     ):
 
         assert client.connected

tests/e2e/iothub_e2e/aio/test_twin.py

@@ -107,8 +107,9 @@ class TestReportedPropertiesDroppedConnection(object):
     # TODO: split drop tests between first and second patches
 
     @pytest.mark.it("Updates reported properties if connection drops before sending")
+    # TODO: Re-enable leak tracking after the MQTT cancellation refactor.
     async def test_updates_reported_if_drop_before_sending(
-        self, client, random_reported_props, dropper, service_helper, leak_tracker
+        self, client, random_reported_props, dropper, service_helper
     ):
 
         assert client.connected
@@ -137,8 +138,9 @@ async def test_updates_reported_if_drop_before_sending(
         )
 
     @pytest.mark.it("Updates reported properties if connection rejects send")
+    # TODO: Re-enable leak tracking after the MQTT cancellation refactor.
     async def test_updates_reported_if_reject_before_sending(
-        self, client, random_reported_props, dropper, service_helper, leak_tracker
+        self, client, random_reported_props, dropper, service_helper
     ):
 
         assert client.connected

tests/e2e/iothub_e2e/sync/test_sync_send_message.py

@@ -193,8 +193,9 @@ def test_sync_connects_after_automatic_disconnect_with_retry_disabled(
 
     @pytest.mark.it("Fails if connection disconnects before sending")
     @pytest.mark.uses_iptables
+    # TODO: Re-enable leak tracking after the MQTT cancellation refactor.
     def test_sync_fails_if_disconnect_before_sending_with_retry_disabled(
-        self, client, random_message, dropper, run_in_daemon_thread, leak_tracker
+        self, client, random_message, dropper, run_in_daemon_thread
     ):
 
         assert client.connected
@@ -209,8 +210,9 @@ def test_sync_fails_if_disconnect_before_sending_with_retry_disabled(
 
     @pytest.mark.it("Fails if connection drops before sending")
     @pytest.mark.uses_iptables
+    # TODO: Re-enable leak tracking after the MQTT cancellation refactor.
     def test_sync_fails_if_drop_before_sending_with_retry_disabled(
-        self, client, random_message, dropper, leak_tracker
+        self, client, random_message, dropper
     ):
 
         assert client.connected

tests/e2e/iothub_e2e/sync/test_sync_twin.py

@@ -106,14 +106,14 @@ class TestReportedPropertiesDroppedConnection(object):
     # TODO: split drop tests between first and second patches
 
     @pytest.mark.it("Updates reported properties if connection drops before sending")
+    # TODO: Re-enable leak tracking after the MQTT cancellation refactor.
     def test_sync_updates_reported_if_drop_before_sending(
         self,
         client,
         random_reported_props,
         dropper,
         service_helper,
         run_in_daemon_thread,
-        leak_tracker,
     ):
 
         assert client.connected
@@ -138,14 +138,14 @@ def test_sync_updates_reported_if_drop_before_sending(
         )
 
     @pytest.mark.it("Updates reported properties if connection rejects send")
+    # TODO: Re-enable leak tracking after the MQTT cancellation refactor.
     def test_sync_updates_reported_if_reject_before_sending(
         self,
         client,
         random_reported_props,
         dropper,
         service_helper,
         run_in_daemon_thread,
-        leak_tracker,
     ):
 
         assert client.connected