fix: MQTTTransport Overhaul

[Carter Tinney (cartertinney)]

Description of the problem

Paho 2.1 exposed defects in the SDK's MQTT integration that went beyond callback signature changes. The failures and their corresponding fixes are listed together below.

Paho 2.x compatibility

• The SDK selected deprecated callback API v1. Creating an MQTT client emitted a customer-visible deprecation warning, while moving directly to v2 would break callbacks that expected v1 arguments and raw MQTT 3.1.1 return codes.
  • How fixed: Require paho-mqtt>=2.1.0,<3.0.0, instantiate TCP and WebSocket clients with CallbackAPIVersion.VERSION2, adapt every callback signature, and map Paho's normalized MQTT 3.1.1 reasons into SDK exceptions.
• Paho and the SDK could both own reconnect. The SDK worked around Paho's automatic reconnect with a two-hour delay, leaving two competing recovery mechanisms and timing-dependent behavior.
  • How fixed: Disable Paho automatic reconnect with reconnect_on_failure=False and leave reconnect policy entirely with the SDK pipeline.

Connection lifecycle

• No component owned the complete result of connect(). The transport returned after starting Paho's loop, while callbacks and a separate pipeline watchdog independently decided CONNACK success, failure, and timeout. Callers could observe transport success before MQTT establishment, and boundary races could produce conflicting pipeline state.
  • How fixed: Make MQTTTransport.connect() block for CONNACK and resolve the attempt through a reusable synchronized lifecycle state machine. It returns only after accepted CONNACK, raises mapped rejected-CONNACK or disconnect errors, and enforces the 60-second CONNACK timeout. The pipeline completes ConnectOperation from that result and maps ConnectionTimeoutError to OperationTimeout.
• Main relied on the pipeline to serialize connection lifecycle calls. Direct or future non-pipeline callers could overlap connect(), disconnect(), or shutdown() and race Paho's socket, packet, and network-loop state.
  • How fixed: Serialize all three public transport lifecycle methods with a transport-owned lock while retaining the pipeline's existing operation serialization. Shutdown remains ordered behind an active lifecycle call, matching main's behavior.
• Explicit disconnect completion depended entirely on Paho's callback. Paho 2.1 can return MQTT_ERR_NO_CONN without invoking on_disconnect; a reauthorization that started while already disconnected could therefore leave its internal DisconnectOperation pending forever.
  • How fixed: MQTTTransport.disconnect() blocks until the network loop exits, and MQTTTransportStage completes an explicit disconnect from that method's return. Reauthorization remains an MQTT-stage-owned disconnect-then-connect sequence, so higher pipeline stages do not define what reauthorization means. A deferred pipeline invocation places the return fallback behind any connection-drop callback already submitted by Paho, preserving event-before-operation ordering and preventing duplicate completion.
• The transport callback conflated requested disconnects with connection drops. Expected and unexpected closure shared one SDK-facing handler even though explicit lifecycle methods now have authoritative synchronous results.
  • How fixed: Narrow the SDK-facing callback to on_mqtt_connection_dropped_handler. Explicit disconnects complete from disconnect(); only unsolicited closure reaches the drop handler. Drop-only tracking cleanup and background-error reporting always run for a transport-classified drop, even if a DisconnectOperation becomes pending before the pipeline processes the callback. Those old-connection effects finish before the pending operation completes, because its callbacks may synchronously reconnect and register new work. The transport synthesizes ConnectionDroppedError if Paho anomalously reports a successful reason for an unsolicited closure.
• Callback ordering could lose a real disconnect or affect the wrong operation. A socket close immediately after accepted CONNACK could be discarded after the connected callback cleared the pending operation; a delayed callback from an older connection could instead complete or fail a replacement operation. Rejected CONNACK callback sequences could also produce multiple terminal reports.
  • How fixed: Keep first-terminal-outcome arbitration inside the reusable transport lifecycle. A close before CONNACK is a connection failure, a close after accepted CONNACK but before connect() returns is a dropped connection, and later outcomes are ignored. Join the previous network thread before beginning a new attempt so old callbacks cannot complete a retry.
• Paho 2.1 hides one protocol-version refusal. With automatic reconnect disabled, Paho returns MQTT_ERR_PROTOCOL before invoking on_connect, then reports the same Unspecified error through on_disconnect that it uses for an ordinary pre-CONNACK network loss. Callback API v2 therefore does not expose which case occurred.
  • How fixed: Cover every CONNACK reason delivered through on_connect; explicitly classify the indistinguishable callback path as ConnectionFailedError rather than claiming a protocol-specific result.
• Failed-connect cleanup could replace the actual failure. An exception from disconnect() or loop_stop() masked socket, authorization, Paho return-code, rejected-CONNACK, disconnect, or timeout errors and could change retry behavior.
  • How fixed: Make normal failed-connect teardown best-effort and preserve the original terminal error. Keep strict cleanup only for partial network-thread startup, where failure signals that the Paho client must be replaced.
• Repeated Paho disconnect callbacks could report the same connection drop more than once. Duplicate or racing closure notifications could emit multiple DisconnectedEvent instances and treat a later notification as a new drop.
  • How fixed: Keep closure classification and duplicate suppression inside ConnectionLifecycle. Requested disconnect callbacks are suppressed at the transport boundary; unsolicited drops are reported once.
• Partial startup and orphaned callbacks could leak resources or leave an unusable client. A failed Thread.start() left Paho in a state previously repaired through its private _thread field; a callback after MQTTTransport collection could leave the Paho client and loop alive.
  • How fixed: Close and replace a client left unusable by partial thread startup using public lifecycle behavior, without mutating _thread. Weak-reference callback cleanup now disconnects orphaned clients and stops their network loops.

MQTT operations and persistent sessions

• Paho v2 operation results were interpreted incorrectly. PUBLISH handling relied on tuple-style results, and rejected SUBACK reason codes were ignored, causing failed subscriptions to complete successfully.
  • How fixed: Read MQTTMessageInfo.rc and .mid, inspect every v2 SUBACK reason code, and complete rejected subscriptions with ProtocolClientError.
• Completion callbacks could arrive before MID registration. Paho may invoke PUBLISH, SUBSCRIBE, or UNSUBSCRIBE callbacks before the initiating method returns its MID. The operation manager could lose those completions or, for SUBACK, lose the rejection error.
  • How fixed: Retain unmatched completions together with any error and claim them when the corresponding MID is registered.
• A stale completion could complete an unrelated operation after MID reuse. Late callbacks from locally cancelled work were stored as unknown completions; when Paho reused that 16-bit MID, a new operation could complete immediately from the stale result.
  • How fixed: Retain cancelled MIDs as tombstones, discard their late completions, and clear an unused tombstone when Paho legitimately reuses the MID for newly registered work.
• Operation callbacks ran while the tracking lock was held. Callback re-entry could deadlock or interfere with the surrounding completion sequence.
  • How fixed: Remove tracking state under the lock, then invoke callbacks after releasing it.
• Disconnected QoS 1/2 publishes were failed even though Paho still owned them. On MQTT_ERR_NO_CONN, Paho retained those packets for delivery after reconnect, but the SDK discarded their completion tracking and told callers the work had failed.
  • How fixed: Preserve QoS 1/2 publish tracking so retained packets can complete later; disconnected QoS 0 still fails immediately because Paho does not retain it.
• origin/main preserves every in-flight operation when connection retry is enabled, including work Paho cannot resume. QoS 0 PUBLISH, SUBSCRIBE, and UNSUBSCRIBE callbacks can therefore remain pending forever after a disconnect, while QoS 1/2 PUBLISH tracking must survive for Paho's retransmission state.
  • How fixed: Record publish QoS in the pending-operation type (PUBLISH_QOS_0, PUBLISH_QOS_1, or PUBLISH_QOS_2). Disconnect cleanup preserves only QoS 1/2 publishes; it removes and tombstones QoS 0, SUBSCRIBE, and UNSUBSCRIBE MIDs, then completes those SDK callbacks as cancelled after tracking and lifecycle locks are released. If retry is disabled, all outstanding SDK callbacks are completed as cancelled.

Incorrect documented contract

• The subscribe() docstring described QoS 0 PUBLISH behavior. It gave callers the wrong exception contract.
  • How fixed: Document that subscribe() raises NoConnectionError whenever the client is not connected.

Known boundaries

Completing an SDK-tracked operation as cancelled only ends local completion tracking. It cannot cancel MQTT work already accepted by Paho, so a retained QoS 1/2 publish may still be delivered after a later connection. Hard-disconnect semantics remain a follow-up because local cancellation cannot retract a publish already accepted by Paho.

CONNECTION_TIMEOUT bounds the wait for CONNACK, but Paho socket setup and loop_stop() are blocking calls with no safe cancellation API. Lifecycle calls, including shutdown, are therefore serialized rather than racing Paho teardown.

Validation

• Full unit suite: 5,445 passed, 6 skipped
• MQTT transport + MQTT stage suites: 482 passed
• Black passed
• Ruff passed
• git diff --check passed
• Python E2E build 163421 passed all matrix jobs after the reusable lifecycle fix; CI will validate the latest PR head

[Copilot]

🟡 Changes recommended

Critical MID-reuse and lifecycle races can complete unrelated operations or suppress valid disconnects.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Migrates the MQTT transport to Paho callback API v2 while hardening lifecycle, operation tracking, and cleanup behavior.

Changes:

• Adopts Paho 2.1 and v2 callbacks.
• Improves connection, SUBACK, publish, and MID handling.
• Expands regression and leak-test coverage.

File summaries

File	Description
uv.lock	Updates locked dependency metadata.
tests/unit/iothub/test_sync_clients.py	Tests callback deprecation warnings.
tests/unit/common/test_mqtt_transport.py	Covers transport behavior and races.
tests/unit/common/pipeline/test_pipeline_stages_mqtt.py	Tests pipeline lifecycle handling.
tests/e2e/iothub_e2e/sync/test_sync_twin.py	Re-enables sync twin leak checks.
tests/e2e/iothub_e2e/sync/test_sync_send_message.py	Re-enables sync messaging leak checks.
tests/e2e/iothub_e2e/sync/test_sync_infrastructure.py	Supports sync infrastructure leak checks.
tests/e2e/iothub_e2e/aio/test_twin.py	Re-enables async twin leak checks.
tests/e2e/iothub_e2e/aio/test_send_message.py	Re-enables async messaging leak checks.
tests/e2e/iothub_e2e/aio/test_infrastructure.py	Supports async infrastructure leak checks.
pyproject.toml	Requires Paho MQTT 2.1 or newer.
azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_mqtt.py	Updates connection and disconnect coordination.
azure-iot-device/azure/iot/device/common/pipeline/pipeline_stages_base.py	Documents publish resend behavior.
azure-iot-device/azure/iot/device/common/mqtt_transport.py	Implements v2 callbacks and operation lifecycle handling.

Review details

Suppressed comments (3)

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

• Add the missing space after the comma.

            # `wait_for_publish()`,but that is only supported for PUBLISH.

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

• This rationale omits QoS 2 even though the transport now deliberately preserves and resends both QoS 1 and QoS 2 publishes after reconnect. Mentioning only QoS 1 makes the timeout behavior misleading.

            # Only Sub and Unsub are here because MQTT client will resend QoS 1 publishes after reconnect automatically

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

• This rationale omits QoS 2 even though the transport now deliberately preserves and resends both QoS 1 and QoS 2 publishes after reconnect. Mentioning only QoS 1 makes the retry behavior misleading.

            # Only Sub and Unsub are here because MQTT client will resend QoS 1 publishes after reconnect automatically

• Files reviewed: 13/14 changed files
• Comments generated: 3
• Review effort level: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

Four moderate connection-lifecycle, cleanup, documentation, and disconnect-race issues remain unresolved.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (4)

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

• When disconnect() raises, this path stops the loop but exits before clear_inflight=True completes tracked operations as cancelled. A hard disconnect can therefore leave publish/subscribe pipeline operations pending indefinitely despite the documented cancellation semantics. Ensure hard-disconnect cleanup runs from a finally path even when Paho teardown fails, while preserving the disconnect exception.

        try:
            paho_error_code = self._mqtt_client.disconnect()
        except Exception as e:
            raise exceptions.ProtocolClientError("Unexpected Paho failure during disconnect") from e
        finally:
            try:
                # Always stop and join the network thread, even if disconnect() fails.
                self._mqtt_client.loop_stop()
            finally:

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

• This subscribe() contract describes the QoS 0 publish() special case. subscribe() raises NoConnectionError whenever Paho returns MQTT_ERR_NO_CONN, regardless of requested subscription QoS.

        :raises: NoConnectionError if a QoS 0 message is published while the client is not connected.

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

• A non-MQTT_ERR_NO_CONN return exits here before either cancellation or non-publish retirement runs, even though loop_stop() has already ended response processing. This leaves tracked operations unable to complete (including all operations for a hard disconnect). Perform operation-manager cleanup regardless of the disconnect return code before propagating the mapped error.

            else:
                # This could result in ConnectionDroppedError or ProtocolClientError
                err = _create_error_from_paho_error_code(paho_error_code)
                raise err

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

• A successful CONNACK followed immediately by socket closure queues both callbacks with the same ConnectOperation snapshot. Connected processing clears _pending_connection_op, so this identity check then drops the valid disconnect, leaving upper stages marked connected and MQTT operations unresolved. Resolve lifecycle ordering in the transport with a per-attempt state/generation (and only report an accepted connection after that state is committed) rather than using pending-operation identity as the connection generation.

        if connection_op_snapshot is not self._pending_connection_op:
            logger.info(
                "{}: Ignoring disconnected callback for a connection operation that is no longer pending".format(
                    self.name
                )
            )
            return

• Files reviewed: 12/14 changed files
• Comments generated: 2
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

The unresolved failed-CONNACK mapping and duplicate-disconnect race must be fixed before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

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

• Add the missing space after the comma.

            # `wait_for_publish()`,but that is only supported for PUBLISH.

• Files reviewed: 15/17 changed files
• Comments generated: 2
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

A critical disconnected-reauthorization stall and moderate reused-MID completion race remain unresolved.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

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

• A reused MID can still lose the new operation's completion. If a cancelled tombstone remains, Paho reuses that MID, and the new callback fires before the initiating API returns (an ordering this manager explicitly supports), this branch consumes the tombstone and discards the new completion. Registration then stores the operation as pending with no completion left, so it hangs until an outer timeout. Coordinate callback processing with the in-progress Paho call/registration so a tombstone is cleared before callbacks for a legitimately reused MID can be classified.

            if mid in self._cancelled_operation_mids:
                logger.debug("Discarding completion for cancelled Paho MID {}".format(mid))
                self._cancelled_operation_mids.remove(mid)

• Files reviewed: 15/17 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

A critical disconnect/drop race can skip cleanup and background-error reporting.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

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

• Add the missing space after the comma in this comment.

            # `wait_for_publish()`,but that is only supported for PUBLISH.

• Files reviewed: 19/21 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

Drop cleanup can incorrectly cancel or remove operations created by a newly started connection.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

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

• Add a space after the comma in this comment.

            # `wait_for_publish()`,but that is only supported for PUBLISH.

• Files reviewed: 19/21 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🔵 Needs a closer look

The broad lifecycle and concurrency changes require human approval, and the MQTT_ERR_NOMEM handling issue remains unresolved.

Review details

Suppressed comments (2)

Previously missed (2) — in code that hasn't changed since the last review.

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

• MQTT_ERR_NOMEM is the only Paho 2.1 error code omitted from this mapping, so an out-of-memory result now produces the misleading message Unknown Paho error code=1. Add the recognized code to preserve a meaningful ProtocolClientError (and include it in the parameterized error-code cases).
  azure-iot-device/azure/iot/device/common/mqtt_transport.py:851
• Add the missing space after the comma in this comment.

• Files reviewed: 18/20 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

Cancellation callbacks can deadlock under the lifecycle lock, and an early completion after MID reuse can be lost.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

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

• A cancelled MID can also be legitimately reused by a new operation whose Paho completion arrives before the initiating API returns—the early-completion race handled by register_operation(). In that ordering this branch consumes the old tombstone and discards the new operation's completion; registration then finds neither a tombstone nor an unknown completion and leaves the new callback pending forever. Cancellation/reuse tracking needs enough generation or operation-state information to distinguish a late old completion from an early completion for the reuse, and this combined ordering needs a regression test.

            if mid in self._cancelled_operation_mids:
                logger.debug("Discarding completion for cancelled Paho MID {}".format(mid))
                self._cancelled_operation_mids.remove(mid)

• Files reviewed: 18/20 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

The critical disconnect race and QoS 0 tracking defect must be resolved before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

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

• QoS 0 publishes are also retained here because tracking records only OperationType.PUBLISH. Paho does not retransmit QoS 0 after a disconnect, so if the connection drops before its on_publish callback, this method leaves the callback pending forever. Preserve only resumable QoS 1/2 publishes; QoS 0 tracking must be retired/completed on disconnect, which requires recording the publish QoS/resumability in PendingOperation.

• Files reviewed: 18/20 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

A critical teardown race can leave SUBSCRIBE or UNSUBSCRIBE operations pending forever.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 18/20 changed files
• Comments generated: 1
• Review effort level: Balanced