Private/Azure/Invoke-AzCommandWithRetry.ps1
|
function Invoke-AzCommandWithRetry { <# .SYNOPSIS Execute Azure command with automatic retry logic for throttling. .DESCRIPTION Wraps Azure cmdlet execution with exponential backoff retry mechanism. Handles HTTP 429 (Too Many Requests) and transient network errors automatically. Tracks statistics for monitoring (TotalCalls, ThrottledCalls, FailedCalls). Respects Retry-After and x-ms-ratelimit-remaining headers. .PARAMETER Command ScriptBlock containing the Azure command to execute. .PARAMETER MaxRetries Maximum number of retry attempts. Default: 5. .PARAMETER InitialDelaySeconds Initial delay between retries in seconds. Doubles with each retry (exponential backoff). Default: 2. .PARAMETER OperationName Name of the operation for statistics tracking (e.g., "PolicyDefinition", "PolicySet"). .EXAMPLE $policy = Invoke-AzCommandWithRetry -Command { Get-AzPolicyDefinition -Id $id } -OperationName "PolicyDefinition" .OUTPUTS Object - Result from the executed command #> [CmdletBinding()] param( [Parameter(Mandatory = $true)] [scriptblock]$Command, [int]$MaxRetries = 5, [int]$InitialDelaySeconds = 2, [string]$OperationName = "Azure API" ) # Initialize statistics if not already done if (-not $script:ApiCallStats) { Initialize-ApiStatistics } # Increment total calls counter $script:ApiCallStats.TotalCalls++ $attempt = 0 $delay = $InitialDelaySeconds $maxDelaySeconds = 60 # Cap delay at 60 seconds while ($attempt -lt $MaxRetries) { try { $attempt++ Write-Debug "[$OperationName] Attempt $attempt/$MaxRetries" $result = & $Command # Success after retry if ($attempt -gt 1) { Write-Debug "✅ [$OperationName] succeeded after $($attempt - 1) retries" } return $result } catch { $exception = $_.Exception $shouldRetry = $false $retryDelaySeconds = $delay # Extract HTTP status code if available $statusCode = $null if ($exception.Response) { $statusCode = $exception.Response.StatusCode.value__ } # Check for throttling (429) or rate limit errors $isThrottling = $statusCode -eq 429 -or $exception.Message -match 'TooManyRequests|Rate limit|throttl' if ($isThrottling) { $shouldRetry = $true $script:ApiCallStats.ThrottledCalls++ # Check for Retry-After header (Azure provides this) if ($exception.Response.Headers -and $exception.Response.Headers["Retry-After"]) { $retryAfter = $exception.Response.Headers["Retry-After"] # Retry-After can be in seconds (int) or HTTP-date (string) if ($retryAfter -match '^\d+$') { $retryDelaySeconds = [int]$retryAfter Write-Debug "Using Retry-After header: ${retryDelaySeconds}s" } } # Check for x-ms-ratelimit-remaining-* headers if ($exception.Response.Headers -and $exception.Response.Headers["x-ms-ratelimit-remaining-subscription-reads"]) { $remaining = $exception.Response.Headers["x-ms-ratelimit-remaining-subscription-reads"] if ($remaining -match '^\d+$' -and [int]$remaining -eq 0) { # Rate limit exhausted, wait longer $retryDelaySeconds = [Math]::Max($retryDelaySeconds, 5) Write-Debug "Rate limit exhausted (x-ms-ratelimit-remaining: 0), waiting ${retryDelaySeconds}s" } } # Cap delay at maximum $retryDelaySeconds = [Math]::Min($retryDelaySeconds, $maxDelaySeconds) if ($attempt -ge $MaxRetries) { $script:ApiCallStats.FailedCalls++ throw "Maximum retries exceeded for $OperationName after throttling. Error: $_" } Write-Warning "⚠️ [$OperationName] Throttling detected (HTTP $statusCode)" Write-Warning " Waiting ${retryDelaySeconds}s before retry $attempt/$MaxRetries..." Start-Sleep -Seconds $retryDelaySeconds # Exponential backoff for next retry (if Retry-After not provided) if (-not $exception.Response.Headers["Retry-After"]) { $delay = [Math]::Min($delay * 2, $maxDelaySeconds) } } elseif ($exception.Message -match 'timeout|network|connection') { # Transient network errors $shouldRetry = $true Write-Warning "⚠️ [$OperationName] Network error (attempt $attempt/$MaxRetries), waiting ${retryDelaySeconds}s..." } # Retry or throw if ($shouldRetry -and $attempt -lt $MaxRetries) { if (-not $isThrottling) { Start-Sleep -Seconds $retryDelaySeconds $delay = [Math]::Min($delay * 2, $maxDelaySeconds) # Exponential backoff } } else { # Final failure $script:ApiCallStats.FailedCalls++ throw } } } $script:ApiCallStats.FailedCalls++ throw "Command failed after $MaxRetries attempts: $OperationName" } |