Retry transient token/connection failures in az iot ops upgrade

[huguesBouvier]

Description

az iot ops upgrade can abort with exit 1 before doing any work when ARM token acquisition (AzureCliCredential → az account get-access-token) fails with a transient Connection reset by peer, or when a subsequent ARM call hits a connection-reset / DNS / read-timeout flake. These are environmental issues, not product or user errors, and today they surface as a hard command failure.

This mirrors, on the product side, a test-side retry we added for the same transient failure (internal bug 38787405 / recurring infra pattern).

Change

• New bounded helper retry_on_transient_error in azext_edge/edge/util/az_client.py (default 3 attempts, linear 15s × attempt backoff).
  • Retries only connection/token-acquisition transients: ServiceRequestError, ServiceResponseError, builtin ConnectionError/TimeoutError, and connection/DNS message markers (e.g. connection reset, reset by peer, temporary failure in name resolution).
  • HTTP 429/5xx are intentionally excluded — the azure-core transport already retries those via its RetryPolicy (verified by the existing test_ops_upgrade_retry_assertion), so retrying here would double-retry.
  • Non-transient errors (e.g. az login needed, validation errors) fail fast.
• upgrade2.py wraps its idempotent operations with the helper:
  • UpgradeManager construction + analyze_cluster (ARM reads — this is the "before doing any work" token-acquisition point).
  • Each extension create/update/delete, the instance update, default registry-endpoint creation, and secretsync migration (idempotent ARM PUT/DELETE).
• Idempotency: the wrapped calls are ARM reads or PUT/DELETE, so re-running converges the same state.

Testing

• New unit tests for is_transient_error (transient markers, non-transient auth/validation, and that 429/5xx are not retried here) and retry_on_transient_error (succeeds after retries, fails fast on non-transient, exhausts attempts).
• pytest azext_edge/tests/utility/test_az_client_unit.py azext_edge/tests/edge/orchestration/test_upgrade2_unit.py → 212 passed.
• flake8 clean; pylint introduces no new findings (pre-existing R0917s only).

History

• Updated HISTORY.rst.

[digimaun]

The 429/5xx exclusion is documented but not enforced. is_transient_error has no HttpResponseError check, so a status error falls through to the substring markers. A 504 message lowercases to contain "gateway timeout", which matches the bare "timeout" marker and comes back transient. The azure core transport already retries 5xx, which test_ops_upgrade_retry_assertion proves for the 503 case, so a persistent 504 gets retried again here. That is the double retry this was meant to avoid, plus about 45s of added sleep. The same breadth retries any genuine failure whose message mentions a timeout or connection error, which is at odds with the "genuine errors still fail fast" note in HISTORY. We should return False for a status bearing HttpResponseError before the marker fallback, and keep the fallback to the credential and connection errors that have no status. The test at test_az_client_unit.py:63 uses message="service error", so it never exercises the real 504 wording. A case asserting False on a "gateway timeout" message would pin it.

[digimaun]

The retry wraps each begin_* plus wait_for_terminal_state as a whole, so an exhausted transport error abandons the poller and resubmits the whole operation. On top of the transport's own retries, one failing request can reach around a dozen HTTP attempts, and the fixed 15s then 30s sleeps add up. Worst case across the extension ops plus instance, registry and secretsync is several minutes of pure added sleep on an interactive command. The case this PR describes, a token failure before any work during the UpgradeManager build and analyze_cluster, is the clear win and worth keeping. For the ARM writes the transport already retries connection and 5xx, so the extra layer is less clearly earning the latency. We should scope the helper to the reads before any work, or at least use a smaller backoff and fewer attempts for the writes.

[digimaun]

migrate_to_v2 runs a PATCH, then zero or more SecretSync PATCHes, then a delete of the legacy SPC, working from state cached at construction. Wrapping the whole thing means a transient in the delete restarts the sequence from the top rather than resuming. The re-run replays the PATCHes fine since they recompute the same content, but the delete is the concern. If the first delete succeeded server side and only the response was lost, the retry re-issues it and the vendored client maps a 404 to an error, so a converged migration could still surface as a hard failure. I could not confirm the exact RP behavior for a duplicate delete statically. We should tolerate an already deleted SPC on retry, or resume rather than restart, so a lost delete response does not fail a completed migration.

[digimaun]

This project does not use HISTORY.rst. Release notes are captured on the GitHub releases, so we should drop this change. The PR checklist still asking for a HISTORY.rst entry is misleading, but that is a separate cleanup.

[digimaun]

Following on the latency point in the other comment, and thinking about the approach more broadly, it may be worth splitting this rather than tuning the helper. The transient failures are really two kinds, and azure core already handles one.

For the ARM HTTP transients, connection reset, read timeout and 5xx on retryable methods, RetryPolicy already retries, which is where that double retry comes from. oci_client.get_oci_retry_policy already sets one up in this package, so the get*_client factories could carry a RetryPolicy and the write path wrapping could go away.

The token acquisition case this PR is really about is the part RetryPolicy cannot cover. AzureCliCredential.get_token failing while shelling out to az account get-access-token surfaces as ClientAuthenticationError, and RetryPolicy re-raises that without retrying (_retry.py send). Since get_token runs inside the pipeline, no transport retry helps.

So the more elegant shape is to lean on RetryPolicy for the HTTP transients and drop the per call site wrapping, then handle the one real gap centrally by wrapping the shared AZURE_CLI_CREDENTIAL with a small retrying get_token. That fixes the token flake once for every command instead of only upgrade, and removes the lambda boilerplate.

[zzcodework]

Azure SRE Agent - automated review

Reviewing the full diff (head 96b7b7fe, base dev).

Well-scoped change: the helper is documented, deliberately excludes HTTP status errors so it does not double-retry on top of azure-core's RetryPolicy, the loop-variable capture in apply_upgrades is correctly bound (lambda ext=ext, op_type=op_type: ...), HISTORY.rst is updated, and there is real unit coverage for both the transient and fail-fast paths. One correctness gap stands out.

Blocking

1. azext_edge/edge/util/az_client.py::is_transient_error — the documented contract ("HTTP status errors (429/5xx) are excluded on purpose") is not actually enforced. The isinstance checks only cover ServiceRequestError / ServiceResponseError / ConnectionError / TimeoutError; an HttpResponseError then falls through to the substring scan, and the marker list contains "timeout" and "timed out". So a real service response like 504 Gateway Timeout, or any HttpResponseError whose message embeds the word "timeout" (very common in ARM error bodies — "The operation has timed out", "Gateway Timeout"), is classified as transient and retried here on top of the transport retry policy. Same trap for a 400 whose body mentions a connection error from a downstream service. Fix is one line before the fallback:

   from azure.core.exceptions import HttpResponseError
   if isinstance(error, HttpResponseError):
       return False

   Note this is not caught by the new tests: test_is_transient_error_http_status_not_retried_here only ever constructs HttpResponseError(message="service error"), which contains none of the markers, so it passes for the wrong reason. Please add a case with message="Gateway Timeout" / "The operation has timed out" — it should fail today.

Suggestions

2. azext_edge/edge/providers/orchestration/upgrade2.py — retry_on_transient_error is applied to self._apply_single_operation(...) for ExtensionOperation.CREATE/UPDATE/DELETE and to _apply_instance_update / _create_default_registry_endpoint / migrate_to_v2. The helper's docstring asserts "func must be idempotent; all call sites wrap idempotent ARM reads or PUT/DELETE operations". That holds for ARM PUT/DELETE, but a ServiceResponseError means the request may well have been delivered and only the response was lost — so the retry can re-issue a mutation whose first attempt is still in flight, and for a long-running extension op the second PUT can collide with the first (409 / conflicting provisioning state). Worth stating explicitly in the PR description which of these are safe to re-issue mid-flight, or restricting the mutation call sites to ServiceRequestError only (request never sent) while keeping the broader set for the read-only analyze_cluster / UpgradeManager(...) construction.

3. _TRANSIENT_ERROR_MARKERS contains both "timeout" and "timed out" — the former subsumes the latter, and "connection error" is broad enough to match a message that merely mentions a downstream connection error in an otherwise permanent failure. Trimming to the specific socket-level strings (and relying on the isinstance checks for the rest) reduces the false-positive surface, which matters because every false positive costs the user 45s of sleep before the real error surfaces.

4. retry_on_transient_error sleeps up to 15 + 30 = 45s per wrapped call while the Rich progress display in apply_upgrades is live. With several extensions failing transiently this is minutes of an apparently stalled progress bar; the logger.warning may also render awkwardly inside the live display. Consider surfacing the retry through the progress task description, or at minimum call out the worst-case added wall-clock time in the changelog entry.

Nits

5. azext_edge/edge/providers/orchestration/upgrade2.py — the new from ...util.az_client import retry_on_transient_error is appended after the from .targets import InitTargets block; parent-relative imports elsewhere in this file are grouped above the sibling .-relative ones. Move it up to keep the import block consistent.

6. azext_edge/tests/utility/test_az_client_unit.py::test_is_transient_error_service_request takes an unused mocker fixture.

Summary — the retry helper and its call sites are structurally sound and well tested, but is_transient_error does not enforce the HTTP-status exclusion it documents (finding 1), and the existing test passes for the wrong reason. Worth also clarifying the mid-flight-mutation retry semantics before this ships.

This is an automated review and may be incomplete; please verify findings before acting on them.