From 38ba8517a36f06971f2464b560c750bc093b6738 Mon Sep 17 00:00:00 2001
From: Carter Tinney <catinney@microsoft.com>
Date: Wed, 2 Sep 2026 15:36:49 -0700
Subject: [PATCH 1/5] e2e: automate pre-teardown leak checks

Run opted-in leak tracking around pytest's call phase so test-local references are released before checking while the client is still alive. Remove the repeated lifecycle calls and leak-specific variable cleanup from individual tests.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 tests/e2e/iothub_e2e/aio/test_c2d.py          | 12 ++----
 .../iothub_e2e/aio/test_connect_disconnect.py | 18 --------
 .../aio/test_connect_disconnect_stress.py     |  9 ----
 tests/e2e/iothub_e2e/aio/test_methods.py      | 18 ++++----
 tests/e2e/iothub_e2e/aio/test_sas_renewal.py  |  4 --
 tests/e2e/iothub_e2e/aio/test_send_message.py | 43 -------------------
 .../aio/test_send_message_stress.py           | 12 ------
 tests/e2e/iothub_e2e/aio/test_twin.py         | 26 -----------
 tests/e2e/iothub_e2e/aio/test_twin_stress.py  | 18 --------
 tests/e2e/iothub_e2e/conftest.py              | 29 +++++++++++--
 tests/e2e/iothub_e2e/sync/test_sync_c2d.py    | 12 ++----
 .../sync/test_sync_connect_disconnect.py      | 15 -------
 .../e2e/iothub_e2e/sync/test_sync_methods.py  | 18 ++++----
 .../iothub_e2e/sync/test_sync_sas_renewal.py  |  3 --
 .../iothub_e2e/sync/test_sync_send_message.py | 43 -------------------
 tests/e2e/iothub_e2e/sync/test_sync_twin.py   | 17 --------
 16 files changed, 49 insertions(+), 248 deletions(-)

diff --git a/tests/e2e/iothub_e2e/aio/test_c2d.py b/tests/e2e/iothub_e2e/aio/test_c2d.py
index a67c95e3f..5163bff3b 100644
--- a/tests/e2e/iothub_e2e/aio/test_c2d.py
+++ b/tests/e2e/iothub_e2e/aio/test_c2d.py
@@ -21,18 +21,17 @@ class TestReceiveC2d(object):
     @pytest.mark.it("Can receive C2D")
     @pytest.mark.quicktest_suite
     async def test_receive_c2d(self, client, service_helper, leak_tracker):
-        leak_tracker.set_initial_object_list()
         event_loop = asyncio.get_running_loop()
 
         message = json.dumps(get_random_dict())
 
-        received_message = None
+        received_message_data = None
         received = asyncio.Event()
 
         async def handle_on_message_received(message):
-            nonlocal received_message, received
+            nonlocal received_message_data
             logger.info("received {}".format(message))
-            received_message = message
+            received_message_data = message.data.decode("utf-8")
             event_loop.call_soon_threadsafe(received.set)
 
         client.on_message_received = handle_on_message_received
@@ -42,7 +41,4 @@ async def handle_on_message_received(message):
         await asyncio.wait_for(received.wait(), timeout=const.E2E_TIMEOUT)
         assert received.is_set()
 
-        assert received_message.data.decode("utf-8") == message
-
-        received_message = None  # so this isn't tagged as a leak
-        leak_tracker.check_for_leaks()
+        assert received_message_data == message
diff --git a/tests/e2e/iothub_e2e/aio/test_connect_disconnect.py b/tests/e2e/iothub_e2e/aio/test_connect_disconnect.py
index a6326e412..9b7bf446d 100644
--- a/tests/e2e/iothub_e2e/aio/test_connect_disconnect.py
+++ b/tests/e2e/iothub_e2e/aio/test_connect_disconnect.py
@@ -21,8 +21,6 @@ class TestConnectDisconnect(object):
     async def test_connect_disconnect(self, brand_new_client, leak_tracker):
         client = brand_new_client
 
-        leak_tracker.set_initial_object_list()
-
         assert client
         logger.info("connecting")
         await client.connect()
@@ -34,8 +32,6 @@ async def test_connect_disconnect(self, brand_new_client, leak_tracker):
         await client.connect()
         assert client.connected
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it(
         "Can do a manual connect in the `on_connection_state_change` call that is notifying the user about a disconnect."
     )
@@ -53,7 +49,6 @@ async def test_connect_in_the_middle_of_disconnect(
         client = brand_new_client
         assert client
 
-        leak_tracker.set_initial_object_list()
         event_loop = asyncio.get_running_loop()
 
         reconnected_event = asyncio.Event()
@@ -98,9 +93,6 @@ async def handle_on_connection_state_change():
         event = await service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert event
 
-        random_message = None  # so this isn't flagged as a leak
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it(
         "Can do a manual disconnect in the `on_connection_state_change` call that is notifying the user about a connect."
     )
@@ -127,7 +119,6 @@ async def test_disconnect_in_the_middle_of_connect(
         assert client
         disconnect_on_next_connect_event = False
 
-        leak_tracker.set_initial_object_list()
         event_loop = asyncio.get_running_loop()
 
         disconnected_event = asyncio.Event()
@@ -178,9 +169,6 @@ async def handle_on_connection_state_change():
         event = await service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert event
 
-        random_message = None  # So this doesn't get flagged as a leak.
-        leak_tracker.check_for_leaks()
-
     # TODO: Add connect/disconnect stress, multiple times with connect inside disconnect and disconnect inside connect.
 
 
@@ -194,7 +182,6 @@ async def test_disconnect_on_drop_outgoing(self, client, dropper, leak_tracker):
         This test verifies that the client will disconnect (eventually) if the network starts
         dropping packets
         """
-        leak_tracker.set_initial_object_list()
 
         await client.connect()
         assert client.connected
@@ -211,15 +198,12 @@ async def test_disconnect_on_drop_outgoing(self, client, dropper, leak_tracker):
             lambda: client.connected, timeout=const.E2E_TIMEOUT
         )
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("disconnects when network rejects all outgoing packets")
     async def test_disconnect_on_reject_outgoing(self, client, dropper, leak_tracker):
         """
         This test verifies that the client will disconnect (eventually) if the network starts
         rejecting packets
         """
-        leak_tracker.set_initial_object_list()
 
         await client.connect()
         assert client.connected
@@ -235,5 +219,3 @@ async def test_disconnect_on_reject_outgoing(self, client, dropper, leak_tracker
         await wait_helpers.async_wait_for_condition(
             lambda: client.connected, timeout=const.E2E_TIMEOUT
         )
-
-        leak_tracker.check_for_leaks()
diff --git a/tests/e2e/iothub_e2e/aio/test_connect_disconnect_stress.py b/tests/e2e/iothub_e2e/aio/test_connect_disconnect_stress.py
index ccd756e4d..46460b535 100644
--- a/tests/e2e/iothub_e2e/aio/test_connect_disconnect_stress.py
+++ b/tests/e2e/iothub_e2e/aio/test_connect_disconnect_stress.py
@@ -20,21 +20,17 @@ class TestConnectDisconnectStress(object):
     async def test_non_overlapped_connect_disconnect_stress(
         self, client, iteration_count, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         for _ in range(iteration_count):
             await client.connect()
             await client.disconnect()
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.parametrize("iteration_count", [20, 250])
     @pytest.mark.it("Can do many overlapped connects and disconnects")
     @pytest.mark.timeout(600)
     async def test_overlapped_connect_disconnect_stress(
         self, client, iteration_count, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         futures = []
         for _ in range(iteration_count):
@@ -46,15 +42,12 @@ async def test_overlapped_connect_disconnect_stress(
         finally:
             await task_cleanup.cleanup_tasks(futures)
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.parametrize("iteration_count", [20, 500])
     @pytest.mark.it("Can do many overlapped random connects and disconnects")
     @pytest.mark.timeout(600)
     async def test_overlapped_random_connect_disconnect_stress(
         self, client, iteration_count, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         futures = []
         for _ in range(iteration_count):
@@ -67,5 +60,3 @@ async def test_overlapped_random_connect_disconnect_stress(
             await asyncio.gather(*futures)
         finally:
             await task_cleanup.cleanup_tasks(futures)
-
-        leak_tracker.check_for_leaks()
diff --git a/tests/e2e/iothub_e2e/aio/test_methods.py b/tests/e2e/iothub_e2e/aio/test_methods.py
index 12fcbbd37..4fc355ea3 100644
--- a/tests/e2e/iothub_e2e/aio/test_methods.py
+++ b/tests/e2e/iothub_e2e/aio/test_methods.py
@@ -36,9 +36,9 @@ async def test_handle_method_call(
         service_helper,
         leak_tracker,
     ):
-        leak_tracker.set_initial_object_list()
 
-        actual_request = None
+        actual_request_name = None
+        actual_request_payload = None
 
         if include_request_payload:
             request_payload = get_random_dict()
@@ -51,9 +51,10 @@ async def test_handle_method_call(
             response_payload = None
 
         async def handle_on_method_request_received(request):
-            nonlocal actual_request
+            nonlocal actual_request_name, actual_request_payload
             logger.info("Method request for {} received".format(request.name))
-            actual_request = request
+            actual_request_name = request.name
+            actual_request_payload = request.payload
             logger.info("Sending response")
             await client.send_method_response(
                 MethodResponse.create_from_method_request(
@@ -67,15 +68,12 @@ async def handle_on_method_request_received(request):
         method_response = await service_helper.invoke_method(method_name, request_payload)
 
         # verify that the method request arrived correctly
-        assert actual_request.name == method_name
+        assert actual_request_name == method_name
         if request_payload:
-            assert actual_request.payload == request_payload
+            assert actual_request_payload == request_payload
         else:
-            assert not actual_request.payload
+            assert not actual_request_payload
 
         # and make sure the response came back successfully
         assert method_response.status == method_response_status
         assert method_response.payload == response_payload
-
-        actual_request = None  # so this isn't tagged as a leak
-        leak_tracker.check_for_leaks()
diff --git a/tests/e2e/iothub_e2e/aio/test_sas_renewal.py b/tests/e2e/iothub_e2e/aio/test_sas_renewal.py
index 2e3429e29..390fade16 100644
--- a/tests/e2e/iothub_e2e/aio/test_sas_renewal.py
+++ b/tests/e2e/iothub_e2e/aio/test_sas_renewal.py
@@ -24,7 +24,6 @@ class TestSasRenewal(object):
     @pytest.mark.parametrize(*parametrize.connection_retry_disabled_and_enabled)
     @pytest.mark.parametrize(*parametrize.auto_connect_disabled_and_enabled)
     async def test_sas_renews(self, client, service_helper, random_message, leak_tracker):
-        leak_tracker.set_initial_object_list()
         event_loop = asyncio.get_running_loop()
 
         connected_event = asyncio.Event()
@@ -80,6 +79,3 @@ async def handle_on_connection_state_change():
         # TODO incoming_event_queue.get should check thread future
         event = await service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert json.dumps(event.message_body) == random_message.data
-
-        random_message = None  # so this isn't flagged as a leak
-        leak_tracker.check_for_leaks()
diff --git a/tests/e2e/iothub_e2e/aio/test_send_message.py b/tests/e2e/iothub_e2e/aio/test_send_message.py
index b390a4c47..bc1545594 100644
--- a/tests/e2e/iothub_e2e/aio/test_send_message.py
+++ b/tests/e2e/iothub_e2e/aio/test_send_message.py
@@ -20,22 +20,16 @@ class TestSendMessage(object):
     @pytest.mark.quicktest_suite
     async def test_send_simple_message(self, client, random_message, service_helper, leak_tracker):
 
-        leak_tracker.set_initial_object_list()
-
         await client.send_message(random_message)
 
         event = await service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert event.system_properties["message-id"] == random_message.message_id
         assert json.dumps(event.message_body) == random_message.data
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Connects the transport if necessary")
     @pytest.mark.quicktest_suite
     async def test_connect_if_necessary(self, client, random_message, service_helper, leak_tracker):
 
-        leak_tracker.set_initial_object_list()
-
         await client.disconnect()
         assert not client.connected
 
@@ -45,11 +39,8 @@ async def test_connect_if_necessary(self, client, random_message, service_helper
         event = await service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert json.dumps(event.message_body) == random_message.data
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Raises correct exception for un-serializable payload")
     async def test_bad_payload_raises(self, client, leak_tracker):
-        leak_tracker.set_initial_object_list()
 
         # There's no way to serialize a function.
         def thing_that_cant_serialize():
@@ -59,12 +50,8 @@ def thing_that_cant_serialize():
             await client.send_message(thing_that_cant_serialize)
         assert isinstance(e_info.value.__cause__, TypeError)
 
-        # TODO: investigate leak
-        # leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Can send a JSON-formatted string that isn't wrapped in a Message object")
     async def test_sends_json_string(self, client, service_helper, leak_tracker):
-        leak_tracker.set_initial_object_list()
 
         message = json.dumps(dev_utils.get_random_dict())
 
@@ -76,11 +63,8 @@ async def test_sends_json_string(self, client, service_helper, leak_tracker):
         )
         assert json.dumps(event.message_body) == message
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Can send a random string that isn't wrapped in a Message object")
     async def test_sends_random_string(self, client, service_helper, leak_tracker):
-        leak_tracker.set_initial_object_list()
 
         message = dev_utils.get_random_string(16)
 
@@ -91,8 +75,6 @@ async def test_sends_random_string(self, client, service_helper, leak_tracker):
         )
         assert event.message_body == message
 
-        leak_tracker.check_for_leaks()
-
 
 @pytest.mark.dropped_connection
 @pytest.mark.describe("Client send_message method with dropped connections")
@@ -103,7 +85,6 @@ class TestSendMessageDroppedConnection(object):
     async def test_sends_if_drop_before_sending(
         self, client, random_message, dropper, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         assert client.connected
 
@@ -132,14 +113,11 @@ async def test_sends_if_drop_before_sending(
 
         logger.info("Success")
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Sends if connection rejects send")
     @pytest.mark.uses_iptables
     async def test_sends_if_reject_before_sending(
         self, client, random_message, dropper, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         assert client.connected
 
@@ -168,8 +146,6 @@ async def test_sends_if_reject_before_sending(
 
         logger.info("Success")
 
-        leak_tracker.check_for_leaks()
-
 
 @pytest.mark.describe("Client send_message with reconnect disabled")
 @pytest.mark.keep_alive(5)
@@ -186,20 +162,16 @@ async def reconnect_after_test(self, dropper, client):
     async def test_send_message_retry_disabled(
         self, client, random_message, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         await client.send_message(random_message)
 
         event = await service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert json.dumps(event.message_body) == random_message.data
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Automatically connects if transport manually disconnected before sending")
     async def test_connect_if_necessary_retry_disabled(
         self, client, random_message, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         await client.disconnect()
         assert not client.connected
@@ -210,14 +182,11 @@ async def test_connect_if_necessary_retry_disabled(
         event = await service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert json.dumps(event.message_body) == random_message.data
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Automatically connects if transport automatically disconnected before sending")
     @pytest.mark.uses_iptables
     async def test_connects_after_automatic_disconnect_retry_disabled(
         self, client, random_message, dropper, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         assert client.connected
 
@@ -234,14 +203,11 @@ async def test_connects_after_automatic_disconnect_retry_disabled(
         event = await service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert json.dumps(event.message_body) == random_message.data
 
-        leak_tracker.check_for_leaks()
-
     @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
     ):
-        leak_tracker.set_initial_object_list()
 
         assert client.connected
 
@@ -255,16 +221,11 @@ async def test_fails_if_disconnect_before_sending(
         with pytest.raises(OperationCancelled):
             await asyncio.wait_for(send_task, timeout=const.E2E_TIMEOUT)
 
-        random_message = None  # so this doesn't get tagged as a leak
-        # TODO: investigate leak
-        # leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Fails if connection drops before sending")
     @pytest.mark.uses_iptables
     async def test_fails_if_drop_before_sending_retry_disabled(
         self, client, random_message, dropper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         assert client.connected
 
@@ -273,7 +234,3 @@ async def test_fails_if_drop_before_sending_retry_disabled(
             await client.send_message(random_message)
 
         assert not client.connected
-
-        random_message = None  # so this doesn't get tagged as a leak
-        # TODO: investigate leak
-        # leak_tracker.check_for_leaks()
diff --git a/tests/e2e/iothub_e2e/aio/test_send_message_stress.py b/tests/e2e/iothub_e2e/aio/test_send_message_stress.py
index b8e58d4f3..a41cc4b89 100644
--- a/tests/e2e/iothub_e2e/aio/test_send_message_stress.py
+++ b/tests/e2e/iothub_e2e/aio/test_send_message_stress.py
@@ -239,8 +239,6 @@ async def test_stress_send_continuous_telemetry(
         limits of the code
         """
 
-        leak_tracker.set_initial_object_list()
-
         await self.send_and_verify_continuous_telemetry(
             client=client,
             service_helper=service_helper,
@@ -248,8 +246,6 @@ async def test_stress_send_continuous_telemetry(
             test_length_in_seconds=test_length_in_seconds,
         )
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("send {} messages all at once".format(ALL_AT_ONCE_MESSAGE_COUNT))
     @pytest.mark.timeout(ALL_AT_ONCE_TOTAL_ELAPSED_TIME_FAILURE_TRIGGER)
     async def test_stress_send_message_all_at_once(
@@ -265,16 +261,12 @@ async def test_stress_send_message_all_at_once(
         handle large volumes of outstanding messages.
         """
 
-        leak_tracker.set_initial_object_list()
-
         await self.send_and_verify_many_telemetry_messages(
             client=client,
             service_helper=service_helper,
             message_count=message_count,
         )
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it(
         "regular message delivery with flaky network {} messages per second for {} seconds".format(
             SEND_TELEMETRY_FLAKY_NETWORK_MESSAGES_PER_SECOND,
@@ -302,8 +294,6 @@ async def test_stress_send_message_with_flaky_network(
         that they always arrive.
         """
 
-        leak_tracker.set_initial_object_list()
-
         await asyncio.gather(
             self.do_periodic_network_disconnects(
                 client=client,
@@ -319,5 +309,3 @@ async def test_stress_send_message_with_flaky_network(
                 test_length_in_seconds=test_length_in_seconds,
             ),
         )
-
-        leak_tracker.check_for_leaks()
diff --git a/tests/e2e/iothub_e2e/aio/test_twin.py b/tests/e2e/iothub_e2e/aio/test_twin.py
index b0d09e593..a368a2fe5 100644
--- a/tests/e2e/iothub_e2e/aio/test_twin.py
+++ b/tests/e2e/iothub_e2e/aio/test_twin.py
@@ -25,7 +25,6 @@ class TestReportedProperties(object):
     async def test_sends_simple_reported_patch(
         self, client, random_reported_props, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         # patch properties
         await client.patch_twin_reported_properties(random_reported_props)
@@ -41,12 +40,8 @@ async def test_sends_simple_reported_patch(
         twin = await client.get_twin()
         assert twin[const.REPORTED][const.TEST_CONTENT] == random_reported_props[const.TEST_CONTENT]
 
-        # TODO: investigate leak
-        # leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Raises correct exception for un-serializable patch")
     async def test_bad_reported_patch_raises(self, client, leak_tracker):
-        leak_tracker.set_initial_object_list()
 
         # There's no way to serialize a function.
         def thing_that_cant_serialize():
@@ -56,15 +51,11 @@ def thing_that_cant_serialize():
             await client.patch_twin_reported_properties(thing_that_cant_serialize)
         assert isinstance(e_info.value.__cause__, TypeError)
 
-        # TODO: investigate leak
-        # leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Can clear a reported property")
     @pytest.mark.quicktest_suite
     async def test_clear_property(
         self, client, random_reported_props, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         # patch properties and verify that the service received the patch
         await client.patch_twin_reported_properties(random_reported_props)
@@ -86,14 +77,11 @@ async def test_clear_property(
         twin = await client.get_twin()
         assert const.TEST_CONTENT not in twin[const.REPORTED]
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Connects the transport if necessary")
     @pytest.mark.quicktest_suite
     async def test_patch_reported_connect_if_necessary(
         self, client, random_reported_props, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         await client.disconnect()
 
@@ -110,8 +98,6 @@ async def test_patch_reported_connect_if_necessary(
         twin = await client.get_twin()
         assert twin[const.REPORTED][const.TEST_CONTENT] == random_reported_props[const.TEST_CONTENT]
 
-        leak_tracker.check_for_leaks()
-
 
 @pytest.mark.dropped_connection
 @pytest.mark.describe("Client Reported Properties with dropped connection")
@@ -124,7 +110,6 @@ class TestReportedPropertiesDroppedConnection(object):
     async def test_updates_reported_if_drop_before_sending(
         self, client, random_reported_props, dropper, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         assert client.connected
         dropper.drop_outgoing()
@@ -151,14 +136,10 @@ async def test_updates_reported_if_drop_before_sending(
             == random_reported_props[const.TEST_CONTENT]
         )
 
-        # TODO: investigate leak
-        # leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Updates reported properties if connection rejects send")
     async def test_updates_reported_if_reject_before_sending(
         self, client, random_reported_props, dropper, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         assert client.connected
         dropper.reject_outgoing()
@@ -185,16 +166,12 @@ async def test_updates_reported_if_reject_before_sending(
             == random_reported_props[const.TEST_CONTENT]
         )
 
-        # TODO: investigate leak
-        # leak_tracker.check_for_leaks()
-
 
 @pytest.mark.describe("Client Desired Properties")
 class TestDesiredProperties(object):
     @pytest.mark.it("Receives a patch for a simple desired property")
     @pytest.mark.quicktest_suite
     async def test_receives_simple_desired_patch(self, client, service_helper, leak_tracker):
-        leak_tracker.set_initial_object_list()
         event_loop = asyncio.get_running_loop()
 
         received_patch = None
@@ -221,8 +198,5 @@ async def handle_on_patch_received(patch):
         twin = await client.get_twin()
         assert twin[const.DESIRED][const.TEST_CONTENT] == random_dict
 
-        # TODO: investigate leak
-        # leak_tracker.check_for_leaks()
-
 
 # TODO: etag tests, version tests
diff --git a/tests/e2e/iothub_e2e/aio/test_twin_stress.py b/tests/e2e/iothub_e2e/aio/test_twin_stress.py
index fcd15f584..870edf169 100644
--- a/tests/e2e/iothub_e2e/aio/test_twin_stress.py
+++ b/tests/e2e/iothub_e2e/aio/test_twin_stress.py
@@ -51,7 +51,6 @@ async def test_stress_serial_reported_property_updates(
         Send reported property updates, one at a time, and verify that each one
         has been received at the service. Do not overlap these calls.
         """
-        leak_tracker.set_initial_object_list()
 
         await call_with_retry(client, client.patch_twin_reported_properties, reset_reported_props)
 
@@ -83,8 +82,6 @@ async def test_stress_serial_reported_property_updates(
                         "Wrong patch received. Expecting {}, got {}".format(patch, received_patch)
                     )
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.parametrize(
         "iteration_count, batch_size",
         [
@@ -101,7 +98,6 @@ async def test_stress_parallel_reported_property_updates(
         with `batch_size` overlapped calls in a batch. Verify that the updates arrive
         at the service.
         """
-        leak_tracker.set_initial_object_list()
 
         await call_with_retry(client, client.patch_twin_reported_properties, reset_reported_props)
 
@@ -155,8 +151,6 @@ async def test_stress_parallel_reported_property_updates(
                                     )
                                 )
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.parametrize(
         "iteration_count", [pytest.param(10, id="10 updates"), pytest.param(50, id="50 updates")]
     )
@@ -168,7 +162,6 @@ async def test_stress_serial_desired_property_updates(
         Update desired properties, one at a time, and verify that the desired property arrives
         at the client before the next update.
         """
-        leak_tracker.set_initial_object_list()
         event_loop = asyncio.get_running_loop()
 
         patches = asyncio.Queue()
@@ -193,8 +186,6 @@ async def handle_on_patch_received(patch):
             received_patch = await asyncio.wait_for(patches.get(), timeout=const.E2E_TIMEOUT)
             assert received_patch[const.TEST_CONTENT] == property_value
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.parametrize(
         "iteration_count, batch_size",
         [
@@ -212,7 +203,6 @@ async def test_stress_parallel_desired_property_updates(
         Update desired properties in batches. Each batch updates `batch_size` properties,
         with each property being updated in it's own `PATCH`.
         """
-        leak_tracker.set_initial_object_list()
         event_loop = asyncio.get_running_loop()
 
         patches = asyncio.Queue()
@@ -267,8 +257,6 @@ async def handle_on_patch_received(patch):
                                 )
                             )
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.parametrize(
         "iteration_count", [pytest.param(10, id="10 updates"), pytest.param(50, id="50 updates")]
     )
@@ -281,7 +269,6 @@ async def test_stress_serial_get_twin_calls(
         calls `get_twin()` `iteration_count` times. Once a reported property shows up in the
         twin, that property is updated to be verified in future `get_twin` calls.
         """
-        leak_tracker.set_initial_object_list()
 
         last_property_value = None
         current_property_value = None
@@ -315,8 +302,6 @@ async def test_stress_serial_get_twin_calls(
 
         assert last_property_value, "No patches with updated properties were received"
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.parametrize(
         "iteration_count, batch_size",
         [
@@ -334,7 +319,6 @@ async def test_stress_parallel_get_twin_calls(
         calls `get_twin()` `iteration_count` times. Once a reported property shows up in the
         twin, that property is updated to be verified in future `get_twin` calls.
         """
-        leak_tracker.set_initial_object_list()
 
         last_property_value = None
         current_property_value = get_random_property_value()
@@ -401,5 +385,3 @@ async def test_stress_parallel_get_twin_calls(
             if got_a_match:
                 last_property_value = current_property_value
                 current_property_value = None
-
-        leak_tracker.check_for_leaks()
diff --git a/tests/e2e/iothub_e2e/conftest.py b/tests/e2e/iothub_e2e/conftest.py
index baa4ba47f..c52ed10e1 100644
--- a/tests/e2e/iothub_e2e/conftest.py
+++ b/tests/e2e/iothub_e2e/conftest.py
@@ -121,6 +121,7 @@ def leak_tracker_filter(leaks):
 
 @pytest.fixture(scope="function")
 def leak_tracker():
+    """Opt the requesting test into a leak check before fixture teardown."""
     tracker = leak_tracker_module.LeakTracker()
     tracker.track_module("azure.iot.device")
     tracker.track_module("paho")
@@ -225,11 +226,11 @@ def pytest_runtest_setup(item):
     # 1. The `outer_leak_tracker` object attached to tests is called after `disconnect` or
     #    `shutdown` is called. This means it can only detect objects that survive `shutdown`.
     #
-    # 2. The `leak_tracker` fixture is used within tests and needs to be manually invoked.
-    #    This means it gets called before `shutdown`, so it can detect leaks that might otherwise
-    #    get cleaned up.
+    # 2. The `leak_tracker` fixture opts tests into a check immediately after the test body.
+    #    This means it runs before `shutdown`, so it can detect leaks that might otherwise get
+    #    cleaned up.
     #
-    # Of these 2, the `leak_tracker` fixture is more useful, but it does require manual steps.
+    # Of these 2, the `leak_tracker` fixture is more useful.
     #
     item.outer_leak_tracker = leak_tracker_module.LeakTracker()
     item.outer_leak_tracker.track_module("azure.iot.device")
@@ -238,6 +239,26 @@ def pytest_runtest_setup(item):
     item.outer_leak_tracker.set_initial_object_list()
 
 
+@pytest.hookimpl(wrapper=True)
+def pytest_runtest_call(item):
+    """
+    Run the inner leak check around tests that request the `leak_tracker` fixture.
+
+    The baseline is captured after fixture setup. The check runs after the test frame has been
+    released, but before fixture teardown shuts down the client.
+    """
+    tracker = item.funcargs.get("leak_tracker")
+    if tracker is not None:
+        tracker.set_initial_object_list()
+
+    result = yield
+
+    if tracker is not None:
+        tracker.check_for_leaks()
+
+    return result
+
+
 @pytest.hookimpl(hookwrapper=True)
 def pytest_exception_interact(node, call, report):
     e = call.excinfo.value
diff --git a/tests/e2e/iothub_e2e/sync/test_sync_c2d.py b/tests/e2e/iothub_e2e/sync/test_sync_c2d.py
index a5ac53ade..edc2deffe 100644
--- a/tests/e2e/iothub_e2e/sync/test_sync_c2d.py
+++ b/tests/e2e/iothub_e2e/sync/test_sync_c2d.py
@@ -20,17 +20,16 @@ class TestReceiveC2d(object):
     @pytest.mark.it("Can receive C2D")
     @pytest.mark.quicktest_suite
     def test_sync_receive_c2d(self, client, service_helper, leak_tracker):
-        leak_tracker.set_initial_object_list()
 
         message = json.dumps(get_random_dict())
 
-        received_message = None
+        received_message_data = None
         received = threading.Event()
 
         def handle_on_message_received(message):
-            nonlocal received_message, received
+            nonlocal received_message_data
             logger.info("received {}".format(message))
-            received_message = message
+            received_message_data = message.data.decode("utf-8")
             received.set()
 
         client.on_message_received = handle_on_message_received
@@ -39,7 +38,4 @@ def handle_on_message_received(message):
 
         assert received.wait(timeout=const.E2E_TIMEOUT)
 
-        assert received_message.data.decode("utf-8") == message
-
-        received_message = None  # so this isn't tagged as a leak
-        leak_tracker.check_for_leaks()
+        assert received_message_data == message
diff --git a/tests/e2e/iothub_e2e/sync/test_sync_connect_disconnect.py b/tests/e2e/iothub_e2e/sync/test_sync_connect_disconnect.py
index c6222161a..702322249 100644
--- a/tests/e2e/iothub_e2e/sync/test_sync_connect_disconnect.py
+++ b/tests/e2e/iothub_e2e/sync/test_sync_connect_disconnect.py
@@ -20,7 +20,6 @@ class TestConnectDisconnect(object):
     @pytest.mark.parametrize(*parametrize.auto_connect_disabled_and_enabled)
     @pytest.mark.quicktest_suite
     def test_sync_connect_disconnect(self, brand_new_client, leak_tracker):
-        leak_tracker.set_initial_object_list()
 
         client = brand_new_client
 
@@ -33,8 +32,6 @@ def test_sync_connect_disconnect(self, brand_new_client, leak_tracker):
         client.connect()
         assert client.connected
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it(
         "Can do a manual connect in the `on_connection_state_change` call that is notifying the user about a disconnect."
     )
@@ -49,7 +46,6 @@ def test_sync_connect_in_the_middle_of_disconnect(
         Explanation: People will call `connect` inside `on_connection_state_change` handlers.
         We have to make sure that we can handle this without getting stuck in a bad state.
         """
-        leak_tracker.set_initial_object_list()
 
         client = brand_new_client
         assert client
@@ -98,8 +94,6 @@ def handle_on_connection_state_change():
         event = service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert event
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it(
         "Can do a manual disconnect in the `on_connection_state_change` call that is notifying the user about a connect."
     )
@@ -117,7 +111,6 @@ def test_sync_disconnect_in_the_middle_of_connect(
         less likely to be a user scenario, but it lets us test with unusual-but-specific timing
         on the call to `disconnect`.
         """
-        leak_tracker.set_initial_object_list()
 
         client = brand_new_client
         assert client
@@ -172,8 +165,6 @@ def handle_on_connection_state_change():
         event = service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert event
 
-        leak_tracker.check_for_leaks()
-
 
 @pytest.mark.dropped_connection
 @pytest.mark.describe("Client object with dropped connection")
@@ -185,7 +176,6 @@ def test_sync_disconnect_on_drop_outgoing(self, client, dropper, leak_tracker):
         This test verifies that the client will disconnect (eventually) if the network starts
         dropping packets
         """
-        leak_tracker.set_initial_object_list()
 
         client.connect()
         assert client.connected
@@ -198,15 +188,12 @@ def test_sync_disconnect_on_drop_outgoing(self, client, dropper, leak_tracker):
         dropper.restore_all()
         wait_helpers.wait_for_condition(lambda: client.connected, timeout=const.E2E_TIMEOUT)
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("disconnects when network rejects all outgoing packets")
     def test_sync_disconnect_on_reject_outgoing(self, client, dropper, leak_tracker):
         """
         This test verifies that the client will disconnect (eventually) if the network starts
         rejecting packets
         """
-        leak_tracker.set_initial_object_list()
 
         client.connect()
         assert client.connected
@@ -218,5 +205,3 @@ def test_sync_disconnect_on_reject_outgoing(self, client, dropper, leak_tracker)
         # have a pending ConnectOperation floating around and this would get tagged as a leak.
         dropper.restore_all()
         wait_helpers.wait_for_condition(lambda: client.connected, timeout=const.E2E_TIMEOUT)
-
-        leak_tracker.check_for_leaks()
diff --git a/tests/e2e/iothub_e2e/sync/test_sync_methods.py b/tests/e2e/iothub_e2e/sync/test_sync_methods.py
index a82a0e47c..332d8a188 100644
--- a/tests/e2e/iothub_e2e/sync/test_sync_methods.py
+++ b/tests/e2e/iothub_e2e/sync/test_sync_methods.py
@@ -35,9 +35,9 @@ def test_sync_handle_method_call(
         service_helper,
         leak_tracker,
     ):
-        leak_tracker.set_initial_object_list()
 
-        actual_request = None
+        actual_request_name = None
+        actual_request_payload = None
 
         if include_request_payload:
             request_payload = get_random_dict()
@@ -50,9 +50,10 @@ def test_sync_handle_method_call(
             response_payload = None
 
         def handle_on_method_request_received(request):
-            nonlocal actual_request
+            nonlocal actual_request_name, actual_request_payload
             logger.info("Method request for {} received".format(request.name))
-            actual_request = request
+            actual_request_name = request.name
+            actual_request_payload = request.payload
             logger.info("Sending response")
             client.send_method_response(
                 MethodResponse.create_from_method_request(
@@ -66,15 +67,12 @@ def handle_on_method_request_received(request):
         method_response = service_helper.invoke_method(method_name, request_payload)
 
         # verify that the method request arrived correctly
-        assert actual_request.name == method_name
+        assert actual_request_name == method_name
         if request_payload:
-            assert actual_request.payload == request_payload
+            assert actual_request_payload == request_payload
         else:
-            assert not actual_request.payload
+            assert not actual_request_payload
 
         # and make sure the response came back successfully
         assert method_response.status == method_response_status
         assert method_response.payload == response_payload
-
-        actual_request = None  # so this isn't tagged as a leak
-        leak_tracker.check_for_leaks()
diff --git a/tests/e2e/iothub_e2e/sync/test_sync_sas_renewal.py b/tests/e2e/iothub_e2e/sync/test_sync_sas_renewal.py
index ab97eabdf..774b15984 100644
--- a/tests/e2e/iothub_e2e/sync/test_sync_sas_renewal.py
+++ b/tests/e2e/iothub_e2e/sync/test_sync_sas_renewal.py
@@ -24,7 +24,6 @@ class TestSasRenewal(object):
     @pytest.mark.parametrize(*parametrize.connection_retry_disabled_and_enabled)
     @pytest.mark.parametrize(*parametrize.auto_connect_disabled_and_enabled)
     def test_sync_sas_renews(self, client, service_helper, random_message, leak_tracker):
-        leak_tracker.set_initial_object_list()
 
         connected_event = threading.Event()
         disconnected_event = threading.Event()
@@ -75,5 +74,3 @@ def handle_on_connection_state_change():
         # and verify that the message arrived at the service
         event = service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert json.dumps(event.message_body) == random_message.data
-
-        leak_tracker.check_for_leaks()
diff --git a/tests/e2e/iothub_e2e/sync/test_sync_send_message.py b/tests/e2e/iothub_e2e/sync/test_sync_send_message.py
index 7c8ebfc02..7522765fa 100644
--- a/tests/e2e/iothub_e2e/sync/test_sync_send_message.py
+++ b/tests/e2e/iothub_e2e/sync/test_sync_send_message.py
@@ -18,19 +18,15 @@ class TestSendMessage(object):
     @pytest.mark.it("Can send a simple message")
     @pytest.mark.quicktest_suite
     def test_sync_send_message_simple(self, client, random_message, service_helper, leak_tracker):
-        leak_tracker.set_initial_object_list()
 
         client.send_message(random_message)
 
         event = service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert json.dumps(event.message_body) == random_message.data
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Connects the transport if necessary")
     @pytest.mark.quicktest_suite
     def test_sync_connect_if_necessary(self, client, random_message, service_helper, leak_tracker):
-        leak_tracker.set_initial_object_list()
 
         client.disconnect()
         assert not client.connected
@@ -41,11 +37,8 @@ def test_sync_connect_if_necessary(self, client, random_message, service_helper,
         event = service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert json.dumps(event.message_body) == random_message.data
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Raises correct exception for un-serializable payload")
     def test_sync_bad_payload_raises(self, client, leak_tracker):
-        leak_tracker.set_initial_object_list()
 
         # There's no way to serialize a function.
         def thing_that_cant_serialize():
@@ -55,12 +48,8 @@ def thing_that_cant_serialize():
             client.send_message(thing_that_cant_serialize)
         assert isinstance(e_info.value.__cause__, TypeError)
 
-        # TODO; investigate this leak
-        # leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Can send a JSON-formatted string that isn't wrapped in a Message object")
     def test_sync_sends_json_string(self, client, service_helper, leak_tracker):
-        leak_tracker.set_initial_object_list()
 
         message = json.dumps(dev_utils.get_random_dict())
 
@@ -72,11 +61,8 @@ def test_sync_sends_json_string(self, client, service_helper, leak_tracker):
         )
         assert json.dumps(event.message_body) == message
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Can send a random string that isn't wrapped in a Message object")
     def test_sync_sends_random_string(self, client, service_helper, leak_tracker):
-        leak_tracker.set_initial_object_list()
 
         message = dev_utils.get_random_string(16)
 
@@ -87,8 +73,6 @@ def test_sync_sends_random_string(self, client, service_helper, leak_tracker):
         )
         assert event.message_body == message
 
-        leak_tracker.check_for_leaks()
-
 
 @pytest.mark.dropped_connection
 @pytest.mark.describe("Client send_message method with dropped connections")
@@ -105,7 +89,6 @@ def test_sync_sends_if_drop_before_sending(
         run_in_daemon_thread,
         leak_tracker,
     ):
-        leak_tracker.set_initial_object_list()
 
         assert client.connected
 
@@ -124,9 +107,6 @@ def test_sync_sends_if_drop_before_sending(
         event = service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert json.dumps(event.message_body) == random_message.data
 
-        random_message = None  # so this doesn't get tagged as a leak
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Sends if connection rejects send")
     @pytest.mark.uses_iptables
     def test_sync_sends_if_reject_before_sending(
@@ -138,7 +118,6 @@ def test_sync_sends_if_reject_before_sending(
         run_in_daemon_thread,
         leak_tracker,
     ):
-        leak_tracker.set_initial_object_list()
 
         assert client.connected
 
@@ -157,9 +136,6 @@ def test_sync_sends_if_reject_before_sending(
         event = service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert json.dumps(event.message_body) == random_message.data
 
-        random_message = None  # so this doesn't get tagged as a leak
-        leak_tracker.check_for_leaks()
-
 
 @pytest.mark.describe("Client send_message with reconnect disabled")
 @pytest.mark.keep_alive(5)
@@ -176,20 +152,16 @@ def reconnect_after_test(self, dropper, client):
     def test_sync_send_message_simple_with_retry_disabled(
         self, client, random_message, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         client.send_message(random_message)
 
         event = service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert json.dumps(event.message_body) == random_message.data
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Automatically connects if transport manually disconnected before sending")
     def test_sync_connect_if_necessary_with_retry_disabled(
         self, client, random_message, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         client.disconnect()
         assert not client.connected
@@ -200,14 +172,11 @@ def test_sync_connect_if_necessary_with_retry_disabled(
         event = service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert json.dumps(event.message_body) == random_message.data
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Automatically connects if transport automatically disconnected before sending")
     @pytest.mark.uses_iptables
     def test_sync_connects_after_automatic_disconnect_with_retry_disabled(
         self, client, random_message, dropper, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         assert client.connected
 
@@ -222,14 +191,11 @@ def test_sync_connects_after_automatic_disconnect_with_retry_disabled(
         event = service_helper.wait_for_eventhub_arrival(random_message.message_id)
         assert json.dumps(event.message_body) == random_message.data
 
-        leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Fails if connection disconnects before sending")
     @pytest.mark.uses_iptables
     def test_sync_fails_if_disconnect_before_sending_with_retry_disabled(
         self, client, random_message, dropper, run_in_daemon_thread, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         assert client.connected
 
@@ -241,16 +207,11 @@ def test_sync_fails_if_disconnect_before_sending_with_retry_disabled(
         with pytest.raises(OperationCancelled):
             send_task.result(timeout=const.E2E_TIMEOUT)
 
-        random_message = None  # So this doesn't get tagged as a leak
-        # TODO: investigate this leak
-        # leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Fails if connection drops before sending")
     @pytest.mark.uses_iptables
     def test_sync_fails_if_drop_before_sending_with_retry_disabled(
         self, client, random_message, dropper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         assert client.connected
 
@@ -259,7 +220,3 @@ def test_sync_fails_if_drop_before_sending_with_retry_disabled(
             client.send_message(random_message)
 
         assert not client.connected
-
-        random_message = None  # So this doesn't get tagged as a leak
-        # TODO: investigate this leak
-        # leak_tracker.check_for_leaks()
diff --git a/tests/e2e/iothub_e2e/sync/test_sync_twin.py b/tests/e2e/iothub_e2e/sync/test_sync_twin.py
index 3d9d4031c..af4b7174e 100644
--- a/tests/e2e/iothub_e2e/sync/test_sync_twin.py
+++ b/tests/e2e/iothub_e2e/sync/test_sync_twin.py
@@ -26,7 +26,6 @@ class TestReportedProperties(object):
     def test_sync_sends_simple_reported_patch(
         self, client, random_reported_props, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         # patch properties
         client.patch_twin_reported_properties(random_reported_props)
@@ -44,7 +43,6 @@ def test_sync_sends_simple_reported_patch(
 
     @pytest.mark.it("Raises correct exception for un-serializable patch")
     def test_sync_bad_reported_patch_raises(self, client, leak_tracker):
-        leak_tracker.set_initial_object_list()
 
         # There's no way to serialize a function.
         def thing_that_cant_serialize():
@@ -57,7 +55,6 @@ def thing_that_cant_serialize():
     @pytest.mark.it("Can clear a reported property")
     @pytest.mark.quicktest_suite
     def test_sync_clear_property(self, client, random_reported_props, service_helper, leak_tracker):
-        leak_tracker.set_initial_object_list()
 
         # patch properties and verify that the service received the patch
         client.patch_twin_reported_properties(random_reported_props)
@@ -84,7 +81,6 @@ def test_sync_clear_property(self, client, random_reported_props, service_helper
     def test_sync_patch_reported_connect_if_necessary(
         self, client, random_reported_props, service_helper, leak_tracker
     ):
-        leak_tracker.set_initial_object_list()
 
         client.disconnect()
 
@@ -101,8 +97,6 @@ def test_sync_patch_reported_connect_if_necessary(
         twin = client.get_twin()
         assert twin[const.REPORTED][const.TEST_CONTENT] == random_reported_props[const.TEST_CONTENT]
 
-        leak_tracker.check_for_leaks()
-
 
 @pytest.mark.dropped_connection
 @pytest.mark.describe("Client Reported Properties with dropped connection")
@@ -121,7 +115,6 @@ def test_sync_updates_reported_if_drop_before_sending(
         run_in_daemon_thread,
         leak_tracker,
     ):
-        leak_tracker.set_initial_object_list()
 
         assert client.connected
         dropper.drop_outgoing()
@@ -144,9 +137,6 @@ def test_sync_updates_reported_if_drop_before_sending(
             == random_reported_props[const.TEST_CONTENT]
         )
 
-        # TODO: investigate leak
-        # leak_tracker.check_for_leaks()
-
     @pytest.mark.it("Updates reported properties if connection rejects send")
     def test_sync_updates_reported_if_reject_before_sending(
         self,
@@ -157,7 +147,6 @@ def test_sync_updates_reported_if_reject_before_sending(
         run_in_daemon_thread,
         leak_tracker,
     ):
-        leak_tracker.set_initial_object_list()
 
         assert client.connected
         dropper.reject_outgoing()
@@ -180,9 +169,6 @@ def test_sync_updates_reported_if_reject_before_sending(
             == random_reported_props[const.TEST_CONTENT]
         )
 
-        # TODO: investigate leak
-        # leak_tracker.check_for_leaks()
-
 
 @pytest.mark.describe("Client Desired Properties")
 class TestDesiredProperties(object):
@@ -190,7 +176,6 @@ class TestDesiredProperties(object):
     @pytest.mark.quicktest_suite
     def test_sync_receives_simple_desired_patch(self, client, service_helper, leak_tracker):
         received_patches = queue.Queue()
-        leak_tracker.set_initial_object_list()
 
         def handle_on_patch_received(patch):
             nonlocal received_patches
@@ -222,7 +207,5 @@ def handle_on_patch_received(patch):
                 assert twin[const.DESIRED][const.TEST_CONTENT] == random_dict
                 break
 
-        leak_tracker.check_for_leaks()
-
 
 # TODO: etag tests, version tests

From 7f4960cd48f9adb2104f7117ef7521c922314964 Mon Sep 17 00:00:00 2001
From: Carter Tinney <catinney@microsoft.com>
Date: Wed, 2 Sep 2026 16:53:16 -0700
Subject: [PATCH 2/5] 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>
---
 .../azure/iot/device/common/mqtt_transport.py | 43 +++++++---
 .../common/pipeline/pipeline_stages_mqtt.py   | 80 ++++++++-----------
 .../pipeline/test_pipeline_stages_mqtt.py     | 51 ++++++++++--
 tests/unit/common/test_mqtt_transport.py      | 34 ++++++++
 4 files changed, 144 insertions(+), 64 deletions(-)

diff --git a/azure-iot-device/azure/iot/device/common/mqtt_transport.py b/azure-iot-device/azure/iot/device/common/mqtt_transport.py
index be1f3abc9..1c3ae7e99 100644
--- a/azure-iot-device/azure/iot/device/common/mqtt_transport.py
+++ b/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")
diff --git a/azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_mqtt.py b/azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_mqtt.py
index 236ca244b..a543576ee 100644
--- a/azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_mqtt.py
+++ b/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.
diff --git a/tests/unit/common/pipeline/test_pipeline_stages_mqtt.py b/tests/unit/common/pipeline/test_pipeline_stages_mqtt.py
index ae65958fb..d8a4b3d4c 100644
--- a/tests/unit/common/pipeline/test_pipeline_stages_mqtt.py
+++ b/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,9 +1206,8 @@ 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
@@ -1178,15 +1215,13 @@ def test_inflight_no_retry(self, mocker, stage, cause):
         # 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
diff --git a/tests/unit/common/test_mqtt_transport.py b/tests/unit/common/test_mqtt_transport.py
index 41d58d904..d499296db 100644
--- a/tests/unit/common/test_mqtt_transport.py
+++ b/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")

From 2bb3263aa07d1de474353f585ca04727c59b3341 Mon Sep 17 00:00:00 2001
From: Carter Tinney <catinney@microsoft.com>
Date: Wed, 2 Sep 2026 17:29:41 -0700
Subject: [PATCH 3/5] revert: defer MQTT cancellation cleanup

Remove the MQTT cancellation changes while that module is being refactored. Keep the automatic leak-check infrastructure, but leave the eight tests that expose the known cancellation leaks opted out until the refactor lands.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../azure/iot/device/common/mqtt_transport.py | 43 +++-------
 .../common/pipeline/pipeline_stages_mqtt.py   | 80 +++++++++++--------
 tests/e2e/iothub_e2e/aio/test_send_message.py |  8 +-
 tests/e2e/iothub_e2e/aio/test_twin.py         |  6 +-
 .../iothub_e2e/sync/test_sync_send_message.py |  6 +-
 tests/e2e/iothub_e2e/sync/test_sync_twin.py   |  4 +-
 .../pipeline/test_pipeline_stages_mqtt.py     | 51 ++----------
 tests/unit/common/test_mqtt_transport.py      | 34 --------
 8 files changed, 78 insertions(+), 154 deletions(-)

diff --git a/azure-iot-device/azure/iot/device/common/mqtt_transport.py b/azure-iot-device/azure/iot/device/common/mqtt_transport.py
index 1c3ae7e99..be1f3abc9 100644
--- a/azure-iot-device/azure/iot/device/common/mqtt_transport.py
+++ b/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")
diff --git a/azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_mqtt.py b/azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_mqtt.py
index a543576ee..236ca244b 100644
--- a/azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_mqtt.py
+++ b/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.
diff --git a/tests/e2e/iothub_e2e/aio/test_send_message.py b/tests/e2e/iothub_e2e/aio/test_send_message.py
index bc1545594..7450f8754 100644
--- a/tests/e2e/iothub_e2e/aio/test_send_message.py
+++ b/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
diff --git a/tests/e2e/iothub_e2e/aio/test_twin.py b/tests/e2e/iothub_e2e/aio/test_twin.py
index a368a2fe5..bef8a03a2 100644
--- a/tests/e2e/iothub_e2e/aio/test_twin.py
+++ b/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
diff --git a/tests/e2e/iothub_e2e/sync/test_sync_send_message.py b/tests/e2e/iothub_e2e/sync/test_sync_send_message.py
index 7522765fa..804dc523a 100644
--- a/tests/e2e/iothub_e2e/sync/test_sync_send_message.py
+++ b/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
diff --git a/tests/e2e/iothub_e2e/sync/test_sync_twin.py b/tests/e2e/iothub_e2e/sync/test_sync_twin.py
index af4b7174e..c7815d0ad 100644
--- a/tests/e2e/iothub_e2e/sync/test_sync_twin.py
+++ b/tests/e2e/iothub_e2e/sync/test_sync_twin.py
@@ -106,6 +106,7 @@ 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,
@@ -113,7 +114,6 @@ def test_sync_updates_reported_if_drop_before_sending(
         dropper,
         service_helper,
         run_in_daemon_thread,
-        leak_tracker,
     ):
 
         assert client.connected
@@ -138,6 +138,7 @@ 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,
@@ -145,7 +146,6 @@ def test_sync_updates_reported_if_reject_before_sending(
         dropper,
         service_helper,
         run_in_daemon_thread,
-        leak_tracker,
     ):
 
         assert client.connected
diff --git a/tests/unit/common/pipeline/test_pipeline_stages_mqtt.py b/tests/unit/common/pipeline/test_pipeline_stages_mqtt.py
index d8a4b3d4c..ae65958fb 100644
--- a/tests/unit/common/pipeline/test_pipeline_stages_mqtt.py
+++ b/tests/unit/common/pipeline/test_pipeline_stages_mqtt.py
@@ -8,7 +8,6 @@
 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,
@@ -617,16 +616,6 @@ 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"
     )
@@ -687,23 +676,6 @@ 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"
     )
@@ -764,16 +736,6 @@ 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"
     )
@@ -1206,8 +1168,9 @@ 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, stage, cause):
-        mock_cancel = stage.transport.cancel_all_operations
+    def test_inflight_no_retry(self, mocker, stage, cause):
+        stage.transport._op_manager = mocker.MagicMock()
+        mock_cancel = stage.transport._op_manager.cancel_all_operations
         stage.nucleus.pipeline_configuration.connection_retry = False
         assert stage._pending_connection_op is None
         assert mock_cancel.call_count == 0
@@ -1215,13 +1178,15 @@ def test_inflight_no_retry(self, stage, cause):
         # Trigger disconnect
         stage.transport.on_mqtt_disconnected_handler(cause)
 
-        mock_cancel.assert_called_once_with()
+        assert mock_cancel.call_count == 1
+        assert mock_cancel.call_args == mocker.call()
 
     @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, stage, cause):
-        mock_cancel = stage.transport.cancel_all_operations
+    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
         stage.nucleus.pipeline_configuration.connection_retry = True
         assert stage._pending_connection_op is None
         assert mock_cancel.call_count == 0
diff --git a/tests/unit/common/test_mqtt_transport.py b/tests/unit/common/test_mqtt_transport.py
index d499296db..41d58d904 100644
--- a/tests/unit/common/test_mqtt_transport.py
+++ b/tests/unit/common/test_mqtt_transport.py
@@ -1008,22 +1008,6 @@ 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(
@@ -2502,24 +2486,6 @@ 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")

From 364a6cab60a886ee60a057b3215224b1ba46354a Mon Sep 17 00:00:00 2001
From: Carter Tinney <catinney@microsoft.com>
Date: Thu, 3 Sep 2026 07:29:23 -0700
Subject: [PATCH 4/5] Revert "revert: defer MQTT cancellation cleanup"

This reverts commit 2bb3263aa07d1de474353f585ca04727c59b3341.
---
 .../azure/iot/device/common/mqtt_transport.py | 43 +++++++---
 .../common/pipeline/pipeline_stages_mqtt.py   | 80 ++++++++-----------
 tests/e2e/iothub_e2e/aio/test_send_message.py |  8 +-
 tests/e2e/iothub_e2e/aio/test_twin.py         |  6 +-
 .../iothub_e2e/sync/test_sync_send_message.py |  6 +-
 tests/e2e/iothub_e2e/sync/test_sync_twin.py   |  4 +-
 .../pipeline/test_pipeline_stages_mqtt.py     | 51 ++++++++++--
 tests/unit/common/test_mqtt_transport.py      | 34 ++++++++
 8 files changed, 154 insertions(+), 78 deletions(-)

diff --git a/azure-iot-device/azure/iot/device/common/mqtt_transport.py b/azure-iot-device/azure/iot/device/common/mqtt_transport.py
index be1f3abc9..1c3ae7e99 100644
--- a/azure-iot-device/azure/iot/device/common/mqtt_transport.py
+++ b/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")
diff --git a/azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_mqtt.py b/azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_mqtt.py
index 236ca244b..a543576ee 100644
--- a/azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_mqtt.py
+++ b/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.
diff --git a/tests/e2e/iothub_e2e/aio/test_send_message.py b/tests/e2e/iothub_e2e/aio/test_send_message.py
index 7450f8754..bc1545594 100644
--- a/tests/e2e/iothub_e2e/aio/test_send_message.py
+++ b/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, leak_tracker
+    ):
 
         assert client.connected
 
@@ -222,9 +223,8 @@ async def test_fails_if_disconnect_before_sending(self, client, random_message,
 
     @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, leak_tracker
     ):
 
         assert client.connected
diff --git a/tests/e2e/iothub_e2e/aio/test_twin.py b/tests/e2e/iothub_e2e/aio/test_twin.py
index bef8a03a2..a368a2fe5 100644
--- a/tests/e2e/iothub_e2e/aio/test_twin.py
+++ b/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
diff --git a/tests/e2e/iothub_e2e/sync/test_sync_send_message.py b/tests/e2e/iothub_e2e/sync/test_sync_send_message.py
index 804dc523a..7522765fa 100644
--- a/tests/e2e/iothub_e2e/sync/test_sync_send_message.py
+++ b/tests/e2e/iothub_e2e/sync/test_sync_send_message.py
@@ -193,9 +193,8 @@ 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, leak_tracker
     ):
 
         assert client.connected
@@ -210,9 +209,8 @@ 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
+        self, client, random_message, dropper, leak_tracker
     ):
 
         assert client.connected
diff --git a/tests/e2e/iothub_e2e/sync/test_sync_twin.py b/tests/e2e/iothub_e2e/sync/test_sync_twin.py
index c7815d0ad..af4b7174e 100644
--- a/tests/e2e/iothub_e2e/sync/test_sync_twin.py
+++ b/tests/e2e/iothub_e2e/sync/test_sync_twin.py
@@ -106,7 +106,6 @@ 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,
@@ -114,6 +113,7 @@ def test_sync_updates_reported_if_drop_before_sending(
         dropper,
         service_helper,
         run_in_daemon_thread,
+        leak_tracker,
     ):
 
         assert client.connected
@@ -138,7 +138,6 @@ 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,
@@ -146,6 +145,7 @@ def test_sync_updates_reported_if_reject_before_sending(
         dropper,
         service_helper,
         run_in_daemon_thread,
+        leak_tracker,
     ):
 
         assert client.connected
diff --git a/tests/unit/common/pipeline/test_pipeline_stages_mqtt.py b/tests/unit/common/pipeline/test_pipeline_stages_mqtt.py
index ae65958fb..d8a4b3d4c 100644
--- a/tests/unit/common/pipeline/test_pipeline_stages_mqtt.py
+++ b/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,9 +1206,8 @@ 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
@@ -1178,15 +1215,13 @@ def test_inflight_no_retry(self, mocker, stage, cause):
         # 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
diff --git a/tests/unit/common/test_mqtt_transport.py b/tests/unit/common/test_mqtt_transport.py
index 41d58d904..d499296db 100644
--- a/tests/unit/common/test_mqtt_transport.py
+++ b/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")

From 6f3b1cc4e57838bbaded1815fe6c937087abf1af Mon Sep 17 00:00:00 2001
From: Carter Tinney <catinney@microsoft.com>
Date: Thu, 3 Sep 2026 07:32:35 -0700
Subject: [PATCH 5/5] Revert "Revert "revert: defer MQTT cancellation cleanup""

This reverts commit 364a6cab60a886ee60a057b3215224b1ba46354a.
---
 .../azure/iot/device/common/mqtt_transport.py | 43 +++-------
 .../common/pipeline/pipeline_stages_mqtt.py   | 80 +++++++++++--------
 tests/e2e/iothub_e2e/aio/test_send_message.py |  8 +-
 tests/e2e/iothub_e2e/aio/test_twin.py         |  6 +-
 .../iothub_e2e/sync/test_sync_send_message.py |  6 +-
 tests/e2e/iothub_e2e/sync/test_sync_twin.py   |  4 +-
 .../pipeline/test_pipeline_stages_mqtt.py     | 51 ++----------
 tests/unit/common/test_mqtt_transport.py      | 34 --------
 8 files changed, 78 insertions(+), 154 deletions(-)

diff --git a/azure-iot-device/azure/iot/device/common/mqtt_transport.py b/azure-iot-device/azure/iot/device/common/mqtt_transport.py
index 1c3ae7e99..be1f3abc9 100644
--- a/azure-iot-device/azure/iot/device/common/mqtt_transport.py
+++ b/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")
diff --git a/azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_mqtt.py b/azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_mqtt.py
index a543576ee..236ca244b 100644
--- a/azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_mqtt.py
+++ b/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.
diff --git a/tests/e2e/iothub_e2e/aio/test_send_message.py b/tests/e2e/iothub_e2e/aio/test_send_message.py
index bc1545594..7450f8754 100644
--- a/tests/e2e/iothub_e2e/aio/test_send_message.py
+++ b/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
diff --git a/tests/e2e/iothub_e2e/aio/test_twin.py b/tests/e2e/iothub_e2e/aio/test_twin.py
index a368a2fe5..bef8a03a2 100644
--- a/tests/e2e/iothub_e2e/aio/test_twin.py
+++ b/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
diff --git a/tests/e2e/iothub_e2e/sync/test_sync_send_message.py b/tests/e2e/iothub_e2e/sync/test_sync_send_message.py
index 7522765fa..804dc523a 100644
--- a/tests/e2e/iothub_e2e/sync/test_sync_send_message.py
+++ b/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
diff --git a/tests/e2e/iothub_e2e/sync/test_sync_twin.py b/tests/e2e/iothub_e2e/sync/test_sync_twin.py
index af4b7174e..c7815d0ad 100644
--- a/tests/e2e/iothub_e2e/sync/test_sync_twin.py
+++ b/tests/e2e/iothub_e2e/sync/test_sync_twin.py
@@ -106,6 +106,7 @@ 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,
@@ -113,7 +114,6 @@ def test_sync_updates_reported_if_drop_before_sending(
         dropper,
         service_helper,
         run_in_daemon_thread,
-        leak_tracker,
     ):
 
         assert client.connected
@@ -138,6 +138,7 @@ 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,
@@ -145,7 +146,6 @@ def test_sync_updates_reported_if_reject_before_sending(
         dropper,
         service_helper,
         run_in_daemon_thread,
-        leak_tracker,
     ):
 
         assert client.connected
diff --git a/tests/unit/common/pipeline/test_pipeline_stages_mqtt.py b/tests/unit/common/pipeline/test_pipeline_stages_mqtt.py
index d8a4b3d4c..ae65958fb 100644
--- a/tests/unit/common/pipeline/test_pipeline_stages_mqtt.py
+++ b/tests/unit/common/pipeline/test_pipeline_stages_mqtt.py
@@ -8,7 +8,6 @@
 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,
@@ -617,16 +616,6 @@ 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"
     )
@@ -687,23 +676,6 @@ 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"
     )
@@ -764,16 +736,6 @@ 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"
     )
@@ -1206,8 +1168,9 @@ 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, stage, cause):
-        mock_cancel = stage.transport.cancel_all_operations
+    def test_inflight_no_retry(self, mocker, stage, cause):
+        stage.transport._op_manager = mocker.MagicMock()
+        mock_cancel = stage.transport._op_manager.cancel_all_operations
         stage.nucleus.pipeline_configuration.connection_retry = False
         assert stage._pending_connection_op is None
         assert mock_cancel.call_count == 0
@@ -1215,13 +1178,15 @@ def test_inflight_no_retry(self, stage, cause):
         # Trigger disconnect
         stage.transport.on_mqtt_disconnected_handler(cause)
 
-        mock_cancel.assert_called_once_with()
+        assert mock_cancel.call_count == 1
+        assert mock_cancel.call_args == mocker.call()
 
     @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, stage, cause):
-        mock_cancel = stage.transport.cancel_all_operations
+    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
         stage.nucleus.pipeline_configuration.connection_retry = True
         assert stage._pending_connection_op is None
         assert mock_cancel.call_count == 0
diff --git a/tests/unit/common/test_mqtt_transport.py b/tests/unit/common/test_mqtt_transport.py
index d499296db..41d58d904 100644
--- a/tests/unit/common/test_mqtt_transport.py
+++ b/tests/unit/common/test_mqtt_transport.py
@@ -1008,22 +1008,6 @@ 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(
@@ -2502,24 +2486,6 @@ 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")