fix: MQTTTransport Overhaul

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

@@ -235,9 +235,7 @@ class ConnectOperation(PipelineOperation):
     Even though this is an base operation, it will most likely be handled by a more specific stage (such as an IoTHub or MQTT stage).
     """
 
-    def __init__(self, callback):
-        self.watchdog_timer = None
-        super().__init__(callback)
+    pass
 
 
 class ReauthorizeConnectionOperation(PipelineOperation):

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

@@ -780,7 +780,7 @@ def __init__(self):
         self.timeout_intervals = {
             pipeline_ops_mqtt.MQTTSubscribeOperation: 10,
             pipeline_ops_mqtt.MQTTUnsubscribeOperation: 10,
-            # Only Sub and Unsub are here because MQTT auto retries pub
+            # Only Sub and Unsub are here because MQTT client will resend QoS 1 publishes after reconnect automatically
         }
 
     @pipeline_thread.runs_on_pipeline_thread
@@ -838,7 +838,7 @@ def __init__(self):
         self.retry_intervals = {
             pipeline_ops_mqtt.MQTTSubscribeOperation: 20,
             pipeline_ops_mqtt.MQTTUnsubscribeOperation: 20,
-            # Only Sub and Unsub are here because MQTT auto retries pub
+            # Only Sub and Unsub are here because MQTT client will resend QoS 1 publishes after reconnect automatically
         }
         self.ops_waiting_to_retry = []
 

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

@@ -16,11 +16,13 @@
 This module contains decorators that are used to marshal code into pipeline and
 callback threads and to assert that code is being called in the correct thread.
 
+The `invoke_on_pipeline_thread`, `invoke_on_pipeline_thread_nowait`, and
+`invoke_on_pipeline_thread_deferred` decorators cause decorated functions to run on
+the pipeline thread.
+
 The intention of these decorators is to ensure the following:
 
-1. All pipeline functions execute in a single thread, known as the "pipeline
-  thread".  The `invoke_on_pipeline_thread` and `invoke_on_pipeline_thread_nowait`
-  decorators cause the decorated function to run on the pipeline thread.
+1. All pipeline functions execute in a single thread, known as the "pipeline thread".
 
 2. If the pipeline thread is busy running a different function, the invoke
   decorators will wait until that function is complete before invoking another
@@ -84,11 +86,13 @@ def _get_named_executor(thread_name):
     return _executors[thread_name]
 
 
-def _invoke_on_executor_thread(func, thread_name, block=True):
+def _invoke_on_executor_thread(func, thread_name, block=True, always_queue=False):
     """
     Return wrapper to run the function on a given thread.  If block==False,
     the call returns immediately without waiting for the decorated function to complete.
     If block==True, the call waits for the decorated function to complete before returning.
+    If always_queue==True, the function is submitted even when called from the target thread.
+    The block argument still determines whether the caller waits for the submitted function.
     """
 
     # Mocks and other callable objects don't have a __name__ attribute.
@@ -100,7 +104,7 @@ def _invoke_on_executor_thread(func, thread_name, block=True):
 
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
-        if threading.current_thread().name is not thread_name:
+        if always_queue or threading.current_thread().name is not thread_name:
             logger.debug("Starting {} in {} thread".format(function_name, thread_name))
 
             def thread_proc():
@@ -153,6 +157,16 @@ def invoke_on_pipeline_thread_nowait(func):
     return _invoke_on_executor_thread(func=func, thread_name="pipeline", block=False)
 
 
+def invoke_on_pipeline_thread_deferred(func):
+    """
+    Queue the decorated function for later execution on the pipeline thread, even if it is already
+    on the pipeline thread. Do not wait for it to complete.
+    """
+    return _invoke_on_executor_thread(
+        func=func, thread_name="pipeline", block=False, always_queue=True
+    )
+
+
 def invoke_on_callback_thread_nowait(func):
     """
     Run the decorated function on the callback thread, but don't wait for it to complete
@@ -178,9 +192,7 @@ def _assert_executor_thread(func, thread_name):
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
 
-        assert (
-            threading.current_thread().name == thread_name
-        ), """
+        assert threading.current_thread().name == thread_name, """
             Function {function_name} is not running inside {thread_name} thread.
             It should be. You should use invoke_on_{thread_name}_thread(_nowait) to enter the
             {thread_name} thread before calling this function.  If you're hitting this from

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

@@ -14,6 +14,12 @@ class ConnectionFailedError(Exception):
     pass
 
 
+class ConnectionTimeoutError(ConnectionFailedError):
+    """Connection was not established before the timeout expired."""
+
+    pass
+
+
 class ConnectionDroppedError(Exception):
     """
     Previously established connection was dropped

pyproject.toml

@@ -28,7 +28,7 @@ classifiers = [
 dependencies = [
     "deprecation>=2.1.0,<3.0.0",
     "janus>=2.0.0,<3.0.0",
-    "paho-mqtt>=2.0.0,<3.0.0",
+    "paho-mqtt>=2.1.0,<3.0.0",
     "PySocks",
     "requests>=2.32.3,<3.0.0",
     "requests-unixsocket>=0.4.1",

tests/e2e/iothub_e2e/aio/test_infrastructure.py

@@ -9,7 +9,7 @@
 class TestServiceHelper(object):
     @pytest.mark.it("returns None when wait_for_event_arrival times out")
     async def test_validate_wait_for_eventhub_arrival_timeout(
-        self, client, random_message, service_helper
+        self, client, random_message, service_helper, leak_tracker
     ):
         # Because we have to support py27, we can't use `threading.Condition.wait_for`.
         # make sure our stand-in functionality behaves the same way when dealing with

tests/e2e/iothub_e2e/aio/test_send_message.py

@@ -205,8 +205,9 @@ async def test_connects_after_automatic_disconnect_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.
-    async def test_fails_if_disconnect_before_sending(self, client, random_message, dropper):
+    async def test_fails_if_disconnect_before_sending(
+        self, client, random_message, dropper, service_helper, leak_tracker
+    ):
 
         assert client.connected
 
@@ -220,11 +221,18 @@ async def test_fails_if_disconnect_before_sending(self, client, random_message,
         with pytest.raises(OperationCancelled):
             await asyncio.wait_for(send_task, timeout=const.E2E_TIMEOUT)
 
+        # -----------------------------------------------------------------------------------------
+        # The SDK operation is cancelled, but Paho still owns the accepted QoS publish. Reconnect
+        # and let the MQTT exchange finish so the normal leak check sees no active session state.
+        dropper.restore_all()
+        await client.connect()
+        event = await service_helper.wait_for_eventhub_arrival(random_message.message_id)
+        assert json.dumps(event.message_body) == random_message.data
+
     @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
+        self, client, random_message, dropper, service_helper, leak_tracker
     ):
 
         assert client.connected
@@ -234,3 +242,11 @@ async def test_fails_if_drop_before_sending_retry_disabled(
             await client.send_message(random_message)
 
         assert not client.connected
+
+        # -----------------------------------------------------------------------------------------
+        # The SDK operation is cancelled, but Paho still owns the accepted QoS publish. Reconnect
+        # and let the MQTT exchange finish so the normal leak check sees no active session state.
+        dropper.restore_all()
+        await client.connect()
+        event = await service_helper.wait_for_eventhub_arrival(random_message.message_id)
+        assert json.dumps(event.message_body) == random_message.data

tests/e2e/iothub_e2e/aio/test_twin.py

@@ -107,9 +107,8 @@ 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
+        self, client, random_reported_props, dropper, service_helper, leak_tracker
     ):
 
         assert client.connected
@@ -138,9 +137,8 @@ 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
+        self, client, random_reported_props, dropper, service_helper, leak_tracker
     ):
 
         assert client.connected

tests/e2e/iothub_e2e/sync/test_sync_infrastructure.py

@@ -8,7 +8,9 @@
 @pytest.mark.describe("ServiceHelper object")
 class TestServiceHelper(object):
     @pytest.mark.it("returns None when wait_for_event_arrival times out")
-    def test_sync_wait_for_event_arrival(self, client, random_message, service_helper):
+    def test_sync_wait_for_event_arrival(
+        self, client, random_message, service_helper, leak_tracker
+    ):
 
         event = service_helper.wait_for_eventhub_arrival(uuid.uuid4(), timeout=2)
         assert event is None

tests/e2e/iothub_e2e/sync/test_sync_send_message.py

@@ -193,9 +193,14 @@ 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
+        self,
+        client,
+        random_message,
+        dropper,
+        run_in_daemon_thread,
+        service_helper,
+        leak_tracker,
     ):
 
         assert client.connected
@@ -208,11 +213,18 @@ def test_sync_fails_if_disconnect_before_sending_with_retry_disabled(
         with pytest.raises(OperationCancelled):
             send_task.result(timeout=const.E2E_TIMEOUT)
 
+        # -----------------------------------------------------------------------------------------
+        # The SDK operation is cancelled, but Paho still owns the accepted QoS publish. Reconnect
+        # and let the MQTT exchange finish so the normal leak check sees no active session state.
+        dropper.restore_all()
+        client.connect()
+        event = service_helper.wait_for_eventhub_arrival(random_message.message_id)
+        assert json.dumps(event.message_body) == random_message.data
+
     @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
+        self, client, random_message, dropper, service_helper, leak_tracker
     ):
 
         assert client.connected
@@ -222,3 +234,11 @@ def test_sync_fails_if_drop_before_sending_with_retry_disabled(
             client.send_message(random_message)
 
         assert not client.connected
+
+        # -----------------------------------------------------------------------------------------
+        # The SDK operation is cancelled, but Paho still owns the accepted QoS publish. Reconnect
+        # and let the MQTT exchange finish so the normal leak check sees no active session state.
+        dropper.restore_all()
+        client.connect()
+        event = service_helper.wait_for_eventhub_arrival(random_message.message_id)
+        assert json.dumps(event.message_body) == random_message.data

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

tests/unit/common/pipeline/conftest.py

@@ -8,6 +8,7 @@
     arbitrary_event,
     arbitrary_op,
     fake_pipeline_thread,
+    fake_pipeline_thread_queue,
     fake_non_pipeline_thread,
     pipeline_connected_mock,
     nucleus,

tests/unit/common/pipeline/fixtures.py

@@ -3,12 +3,14 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
+import concurrent.futures
 import pytest
 import threading
 from azure.iot.device.common.pipeline import (
     pipeline_events_base,
     pipeline_ops_base,
     pipeline_nucleus,
+    pipeline_thread,
 )
 
 
@@ -34,6 +36,34 @@ def arbitrary_op(mocker):
     return op
 
 
+class FakePipelineThreadQueue(object):
+    def __init__(self):
+        self._queued_calls = []
+
+    def submit(self, call):
+        future = concurrent.futures.Future()
+        self._queued_calls.append((call, future))
+        return future
+
+    def run_next(self):
+        call, future = self._queued_calls.pop(0)
+        if future.set_running_or_notify_cancel():
+            try:
+                result = call()
+            except BaseException as e:
+                future.set_exception(e)
+            else:
+                future.set_result(result)
+        return future
+
+    def run_all(self):
+        while self._queued_calls:
+            self.run_next()
+
+    def __len__(self):
+        return len(self._queued_calls)
+
+
 @pytest.fixture
 def pipeline_connected_mock(mocker):
     """This mock can have it's return value altered by any test to indicate whether or not the
@@ -88,6 +118,15 @@ def fake_pipeline_thread():
     this_thread.name = old_name
 
 
+@pytest.fixture
+def fake_pipeline_thread_queue(mocker, fake_pipeline_thread):
+    """Capture work queued for deferred execution on the pipeline thread."""
+    thread_queue = FakePipelineThreadQueue()
+    executor = mocker.patch.object(pipeline_thread, "_get_named_executor").return_value
+    executor.submit.side_effect = thread_queue.submit
+    return thread_queue
+
+
 @pytest.fixture
 def fake_non_pipeline_thread():
     """

tests/unit/common/pipeline/test_pipeline_ops_base.py

@@ -61,18 +61,10 @@ def init_kwargs(self, mocker):
         return kwargs
 
 
-class ConnectOperationInstantiationTests(ConnectOperationTestConfig):
-    @pytest.mark.it("Initializes 'watchdog_timer' attribute to 'None'")
-    def test_retry_timer(self, cls_type, init_kwargs):
-        op = cls_type(**init_kwargs)
-        assert op.watchdog_timer is None
-
-
 pipeline_ops_test.add_operation_tests(
     test_module=this_module,
     op_class_under_test=pipeline_ops_base.ConnectOperation,
     op_test_config_class=ConnectOperationTestConfig,
-    extended_op_instantiation_test_class=ConnectOperationInstantiationTests,
 )
 
 

tests/unit/common/pipeline/test_pipeline_stages_base.py

@@ -2538,6 +2538,10 @@ def test_op_completes_success(self, stage, op):
             pytest.param(
                 [pipeline_ops_base.DisconnectOperation(callback=None)], id="Single op waiting"
             ),
+            pytest.param(
+                [pipeline_ops_base.ShutdownPipelineOperation(callback=None)],
+                id="Shutdown waiting",
+            ),
             pytest.param(
                 [
                     pipeline_ops_base.ReauthorizeConnectionOperation(callback=None),