Merge pull request #7 from tjgruber/6-sporadic-upload-chunk-failure-t…

…o-azure-storage-blob

6 sporadic upload chunk failure to azure storage blob

Private/Invoke-AzureStorageBlobUpload.ps1

@@ -7,21 +7,21 @@ function Invoke-AzureStorageBlobUpload {
         Upload and commit .intunewin file into Azure Storage blob container.
 
         This is a modified function that was originally developed by Dave Falkus and is available here:
-        https://github.com/microsoftgraph/powershell-intune-samples/blob/master/LOB_Application/Win32_Application_Add.ps1        
+        https://github.com/microsoftgraph/powershell-intune-samples/blob/master/LOB_Application/Win32_Application_Add.ps1
 
     .NOTES
         Author:      Nickolaj Andersen
         Contact:     @NickolajA
         Created:     2020-01-04
-        Updated:     2023-09-04
+        Updated:     2024-06-04
 
         Version history:
         1.0.0 - (2020-01-04) Function created
         1.0.1 - (2020-09-20) Fixed an issue where the System.IO.BinaryReader wouldn't open a file path containing whitespaces
         1.0.2 - (2021-03-15) Fixed an issue where SAS Uri renewal wasn't working correctly
-        1.0.3 - (2022-09-03) Added access token refresh functionality when a token is about to expire, to prevent uploads from failing due to an expire access token
-        1.0.4 - (2023-09-04) Updated with Test-AccessToken function
-    #>    
+        1.0.3 - (2022-09-03) Added access token refresh functionality when a token is about to expire, to prevent uploads from failing due to an expired access token
+        1.0.5 - (2024-06-04) Added retry logic for chunk uploads and finalization steps to enhance reliability (thanks to @tjgruber)
+    #>
     param(
         [parameter(Mandatory = $true)]
         [ValidateNotNullOrEmpty()]
@@ -36,6 +36,8 @@ function Invoke-AzureStorageBlobUpload {
         [string]$Resource
     )
     $ChunkSizeInBytes = 1024l * 1024l * 6l;
+    $RetryCount = 5
+    $RetryDelay = 10
 
     # Start the timer for SAS URI renewal
     $SASRenewalTimer = [System.Diagnostics.Stopwatch]::StartNew()
@@ -46,7 +48,7 @@ function Invoke-AzureStorageBlobUpload {
     $BinaryReader = New-Object -TypeName System.IO.BinaryReader([System.IO.File]::Open($FilePath, [System.IO.FileMode]::Open, [System.IO.FileAccess]::Read, [System.IO.FileShare]::ReadWrite))
     $Position = $BinaryReader.BaseStream.Seek(0, [System.IO.SeekOrigin]::Begin)
 
-    # Upload each chunk and dheck whether a SAS URI renewal is required after each chunk is uploaded and renew if needed
+    # Upload each chunk and check whether a SAS URI renewal is required after each chunk is uploaded and renew if needed
     $ChunkIDs = @()
     for ($Chunk = 0; $Chunk -lt $ChunkCount; $Chunk++) {
         Write-Verbose -Message "SAS Uri renewal timer has elapsed for: $($SASRenewalTimer.Elapsed.Minutes) minute $($SASRenewalTimer.Elapsed.Seconds) seconds"
@@ -75,12 +77,38 @@ function Invoke-AzureStorageBlobUpload {
 
         Write-Progress -Activity "Uploading file to Azure Storage blob" -Status "Uploading chunk $($CurrentChunk) of $($ChunkCount)" -PercentComplete ($CurrentChunk / $ChunkCount * 100)
         Write-Verbose -Message "Uploading file to Azure Storage blob, processing chunk '$($CurrentChunk)' of '$($ChunkCount)'"
-        $UploadResponse = Invoke-AzureStorageBlobUploadChunk -StorageUri $StorageUri -ChunkID $ChunkID -Bytes $Bytes
-
+
+        $UploadSuccess = $false
+        for ($i = 0; $i -lt $RetryCount; $i++) {
+            try {
+                $UploadResponse = Invoke-AzureStorageBlobUploadChunk -StorageUri $StorageUri -ChunkID $ChunkID -Bytes $Bytes
+                $UploadSuccess = $true
+                break
+            } catch {
+                Write-Warning "Failed to upload chunk. Attempt $($i + 1) of $RetryCount. Error: $_"
+                Start-Sleep -Seconds $RetryDelay
+                Write-Warning "Retrying upload of chunk '$($CurrentChunk)' of '$($ChunkCount)'"
+            }
+        }
+
+        if (-not $UploadSuccess) {
+            Write-Error "Failed to upload chunk after $RetryCount attempts. Aborting upload."
+            return
+        }
+
         if (($CurrentChunk -lt $ChunkCount) -and ($SASRenewalTimer.ElapsedMilliseconds -ge 450000)) {
             Write-Verbose -Message "SAS Uri renewal is required, attempting to renew"
-            $RenewedSASUri = Invoke-AzureStorageBlobUploadRenew -Resource $Resource
-            $SASRenewalTimer.Restart()
+            try {
+                $RenewedSASUri = Invoke-AzureStorageBlobUploadRenew -Resource $Resource
+                if ($null -ne $RenewedSASUri) {
+                    $StorageUri = $RenewedSASUri
+                    $SASRenewalTimer.Restart()
+                } else {
+                    Write-Warning "SAS Uri renewal failed, continuing with existing Uri"
+                }
+            } catch {
+                Write-Warning "SAS Uri renewal attempt failed with error: $_. Continuing with existing Uri."
+            }
         }
     }
 
@@ -91,9 +119,24 @@ function Invoke-AzureStorageBlobUpload {
     Write-Progress -Completed -Activity "Uploading File to Azure Storage blob"
 
     # Finalize the upload of the content file to Azure Storage blob
-    Invoke-AzureStorageBlobUploadFinalize -StorageUri $StorageUri -ChunkID $ChunkIDs
+    $FinalizeSuccess = $false
+    for ($i = 0; $i -lt $RetryCount; $i++) {
+        try {
+            Invoke-AzureStorageBlobUploadFinalize -StorageUri $StorageUri -ChunkID $ChunkIDs
+            $FinalizeSuccess = $true
+            break
+        } catch {
+            Write-Warning "Failed to finalize Azure Storage blob upload. Attempt $($i + 1) of $RetryCount. Error: $_"
+            Start-Sleep -Seconds $RetryDelay
+        }
+    }
+
+    if (-not $FinalizeSuccess) {
+        Write-Error "Failed to finalize upload after $RetryCount attempts. Aborting upload."
+        return
+    }
 
     # Close and dispose binary reader object
     $BinaryReader.Close()
     $BinaryReader.Dispose()
-}
+}

Private/Invoke-AzureStorageBlobUploadChunk.ps1

@@ -19,7 +19,8 @@ function Invoke-AzureStorageBlobUploadChunk {
         1.0.0 - (2020-01-04) Function created
         1.0.1 - (2021-04-02) Added UseBasicParsing to support conditions where IE first run experience have not been completed
         1.0.2 - (2024-01-10) Fixed issue described in #128 - thanks to @jaspain for finding the solution
-    #>    
+        1.0.3 - (2024-06-03) Added exception throwing on failure to support retry logic in the upload process (thanks to @tjgruber)
+    #>
     param(
         [parameter(Mandatory = $true)]
         [ValidateNotNullOrEmpty()]
@@ -33,18 +34,19 @@ function Invoke-AzureStorageBlobUploadChunk {
         [ValidateNotNullOrEmpty()]
         [System.Object]$Bytes
     )
-	$Uri = "$($StorageUri)&comp=block&blockid=$($ChunkID)"
-	$ISOEncoding = [System.Text.Encoding]::GetEncoding("iso-8859-1")
-	$EncodedBytes = $ISOEncoding.GetString($Bytes)
+    $Uri = "$($StorageUri)&comp=block&blockid=$($ChunkID)"
+    $ISOEncoding = [System.Text.Encoding]::GetEncoding("iso-8859-1")
+    $EncodedBytes = $ISOEncoding.GetString($Bytes)
     $Headers = @{
-		"content-type" = "text/plain; charset=iso-8859-1"
-		"x-ms-blob-type" = "BlockBlob"
-	}
+        "content-type" = "text/plain; charset=iso-8859-1"
+        "x-ms-blob-type" = "BlockBlob"
+    }
 
-	try	{
-		$WebResponse = Invoke-WebRequest $Uri -Method "Put" -Headers $Headers -Body $EncodedBytes -UseBasicParsing -ErrorAction Stop
-	}
-	catch {
+    try {
+        $WebResponse = Invoke-WebRequest $Uri -Method "Put" -Headers $Headers -Body $EncodedBytes -UseBasicParsing -ErrorAction Stop
+        return $WebResponse
+    } catch {
         Write-Warning -Message "Failed to upload chunk to Azure Storage blob. Error message: $($_.Exception.Message)"
-	} 
-}
+        throw $_
+    }
+}

Private/Invoke-AzureStorageBlobUploadFinalize.ps1

@@ -18,6 +18,7 @@ function Invoke-AzureStorageBlobUploadFinalize {
         Version history:
         1.0.0 - (2020-01-04) Function created
         1.0.1 - (2024-05-29) Added content-type header to the REST request to ensure correct handling of the request body (thanks to @tjgruber)
+        1.0.2 - (2024-06-03) Added exception throwing on failure to support retry logic in the finalization process (thanks to @tjgruber)
     #>
     param(
         [parameter(Mandatory = $true)]
@@ -45,8 +46,9 @@ function Invoke-AzureStorageBlobUploadFinalize {
 
     try {
         $WebResponse = Invoke-RestMethod -Uri $Uri -Method "Put" -Body $XML -Headers $Headers -ErrorAction Stop
-    }
-    catch {
+        return $WebResponse
+    } catch {
         Write-Warning -Message "Failed to finalize Azure Storage blob upload. Error message: $($_.Exception.Message)"
+        throw $_
     }
 }

Private/Invoke-AzureStorageBlobUploadRenew.ps1

@@ -13,17 +13,34 @@ function Invoke-AzureStorageBlobUploadRenew {
         Author:      Nickolaj Andersen
         Contact:     @NickolajA
         Created:     2020-01-04
-        Updated:     2021-03-15
+        Updated:     2024-06-03
 
         Version history:
         1.0.0 - (2020-01-04) Function created
         1.0.1 - (2021-03-15) Fixed an issue where SAS Uri renewal wasn't working correctly
-    #>    
+        1.0.2 - (2024-06-03) Added loop to check the status of the SAS URI renewal
+    #>
     param(
         [parameter(Mandatory = $true)]
         [ValidateNotNullOrEmpty()]
         [string]$Resource
     )
     $RenewSASURIRequest = Invoke-IntuneGraphRequest -APIVersion "Beta" -Resource "$($Resource)/renewUpload" -Method "POST" -Body "{}"
-    $FilesProcessingRequest = Wait-IntuneWin32AppFileProcessing -Stage "AzureStorageUriRenewal" -Resource $Resource
-}
+
+    # Loop to wait for the renewal process to complete and check the status
+    $attempts = 0
+    $maxAttempts = 3
+    $waitTime = 5 # seconds
+
+    while ($attempts -lt $maxAttempts) {
+        $FilesProcessingRequest = Invoke-IntuneGraphRequest -APIVersion "Beta" -Resource "$($Resource)" -Method "GET"
+        if ($FilesProcessingRequest.uploadState -eq "azureStorageUriRenewalSuccess") {
+            return $FilesProcessingRequest.azureStorageUri
+        } elseif ($FilesProcessingRequest.uploadState -eq "azureStorageUriRenewalFailed") {
+            throw "SAS Uri renewal failed"
+        }
+        $attempts++
+        Start-Sleep -Seconds $waitTime
+    }
+    throw "SAS Uri renewal did not complete in the expected time"
+}