feat(linux): reduce prefetch optimization runtime within VHD builds and add retry logic to handle AIB rate limits

.pipelines/templates/.builder-release-template.yaml

@@ -295,6 +295,7 @@ steps:
       SIG_IMAGE_NAME: $(SIG_IMAGE_NAME)
       SKU_NAME: $(SKU_NAME)
       STORAGE_ACCOUNT_BLOB_URL: $(CLASSIC_BLOB)
+      STORAGE_ACCOUNT_BLOB_URL_STAGING: $(CLASSIC_BLOB_STAGING)
       VHD_NAME: $(VHD_NAME)
       IMAGE_BUILDER_IDENTITY_ID: $(AZURE_MSI_RESOURCE_STRING)
       BUILD_RUN_NUMBER: $(Build.BuildNumber)

vhdbuilder/prefetch/scripts/optimize.sh

@@ -11,6 +11,7 @@ set -uxo pipefail
 [ -z "${SIG_IMAGE_NAME:-}" ] && echo "SIG_IMAGE_NAME is not set" && exit 1
 [ -z "${SKU_NAME:-}" ] && echo "SKU_NAME is not set" && exit 1
 [ -z "${STORAGE_ACCOUNT_BLOB_URL:-}" ] && echo "STORAGE_ACCOUNT_BLOB_URL is not set" && exit 1
+[ -z "${STORAGE_ACCOUNT_BLOB_URL_STAGING:-}" ] && echo "STORAGE_ACCOUNT_BLOB_URL_STAGING is not set" && exit 1
 [ -z "${VHD_NAME:-}" ] && echo "VHD_NAME is not set" && exit 1
 [ -z "${IMAGE_BUILDER_IDENTITY_ID:-}" ] && echo "IMAGE_BUILDER_IDENTITY_ID is not set" && exit 1
 [ -z "${BUILD_RUN_NUMBER:-}" ] && echo "BUILD_RUN_NUMBER is not set" && exit 1
@@ -28,50 +29,69 @@ CAPTURED_SIG_VERSION_ID="/subscriptions/${SUBSCRIPTION_ID}/resourceGroups/${SIG_
 IMAGE_BUILDER_RG_NAME="image-builder-${CAPTURED_SIG_VERSION}-${BUILD_RUN_NUMBER}"
 IMAGE_BUILDER_TEMPLATE_NAME="template-${CAPTURED_SIG_VERSION}-${BUILD_RUN_NUMBER}"
 IMAGE_BUILDER_TEMPLATE_NAME="${IMAGE_BUILDER_TEMPLATE_NAME:0:64}"
+IMAGE_BUILDER_TEMPLATE_ID="/subscriptions/${SUBSCRIPTION_ID}/resourceGroups/${IMAGE_BUILDER_RG_NAME}/providers/Microsoft.VirtualMachineImages/imageTemplates/${IMAGE_BUILDER_TEMPLATE_NAME}"
 VHD_URI="${STORAGE_ACCOUNT_BLOB_URL}/${VHD_NAME}"
 
+# the image builder template distributes the prefetch-optimized image as a VHD without specifying a
+# target uri, which causes image builder to publish the VHD to a storage account within its own
+# staging resource group. we then copy that staging VHD into the target storage account ourselves,
+# which avoids the slow second copy image builder performs when distributing directly to an external
+# target storage account. the name of the run output is used to look up the resulting blob
+# (artifactUri) after the template run completes.
+DISTRIBUTE_RUN_OUTPUT_NAME="VHD"
+
+# the destination vhd container may have an immutability policy, which prevents azcopy from writing the VHD
+# into it directly (azcopy creates the page blob and then writes to it, which the policy rejects as a
+# modification of an existing blob). To work around this, we azcopy the VHD into the staging container
+# (STORAGE_ACCOUNT_BLOB_URL_STAGING, which has no immutability policy), then perform a single server-side
+# blob copy from the staging container into the immutable vhd container. The server-side copy creates the
+# destination blob in one atomic operation, which the immutability policy allows.
+DESTINATION_STORAGE_ACCOUNT_NAME="${STORAGE_ACCOUNT_BLOB_URL#*://}"
+DESTINATION_STORAGE_ACCOUNT_NAME="${DESTINATION_STORAGE_ACCOUNT_NAME%%.*}"
+DESTINATION_CONTAINER_NAME="${STORAGE_ACCOUNT_BLOB_URL##*/}"
+STAGING_VHD_URI="${STORAGE_ACCOUNT_BLOB_URL_STAGING}/${VHD_NAME}"
+
 main() {
     # for idempotency, check to see if the VHD we're trying to create already exists.
     # if it's already in the expected state, this will cause the script to exit early.
     # otherwise, we delete any existing VHD in an unexpected state and retry the whole
     # optimization + conversion flow
     check_for_existing_vhd
 
-    # attempt to perform prefetch optimization and VHD conversion
+    # attempt to perform prefetch optimization, distributing the optimized image as a VHD into the
+    # image builder staging storage account, then copy that VHD into the target storage account
     ensure_image_builder_rg || exit $?
     run_image_builder_template || exit $?
+    copy_optimized_vhd || exit $?
 }
 
 check_for_existing_vhd() {
     vhd_info="$(az storage blob show --blob-url "${VHD_URI}" --auth-mode login)"
-    if [ -n "${vhd_info}" ]; then
-        echo "VHD already exists at: ${VHD_URI}"
-        image_builder_source="$(jq -r '.metadata.VMImageBuilderSource' <<< "${vhd_info}")"
-        if [ -n "${image_builder_source}" ] && [ "${image_builder_source}" != "null" ]; then
-            echo "VHD ${VHD_URI} has already been produced by a previous image builder template run"
-            copy_status="$(jq -r '.properties.copy.status' <<< "${vhd_info}")"
-            if [ "${copy_status,,}" = "success" ]; then
-                echo "VHD ${VHD_URI} has been successfully copied from image builder storage, nothing to do"
-                exit 0
-            fi
-            if [ "${copy_status,,}" = "pending" ]; then
-                echo "echo VHD ${VHD_URI} is currently being copied from image builder storage, will wait for copy completion"
-                if wait_for_vhd_copy; then
-                    exit 0
-                fi
-                echo "pending copy operation was not successful, will delete existing blob and attempt to retry optimization and VHD creation"
-                delete_vhd || exit $?
-            else
-                echo "VHD ${VHD_URI} has a bad copy state: ${copy_status}, will delete it and recreate"
-                delete_vhd || exit $?
-            fi
-        else
-            echo "VHD ${VHD_URI} exists but was not produced by an image builder template run, will delete before proceeding"
-            delete_vhd || exit $?
+    if [ -z "${vhd_info}" ]; then
+        echo "no existing VHD was found at: ${VHD_URI}, will proceed with optimization and VHD creation"
+        return 0
+    fi
+    echo "VHD already exists at: ${VHD_URI}"
+
+    # the VHD is produced by a server-side blob copy from the staging container, so its copy status tells
+    # us whether a previous run already completed. A successful copy means there is nothing to do; a
+    # pending copy just needs to be waited on. Any other state means the blob is in a bad state and must
+    # be recreated.
+    copy_status="$(jq -r '.properties.copy.status' <<< "${vhd_info}")"
+    if [ "${copy_status,,}" = "success" ]; then
+        echo "VHD ${VHD_URI} was already produced by a previous prefetch optimization run, nothing to do"
+        exit 0
+    fi
+    if [ "${copy_status,,}" = "pending" ]; then
+        echo "VHD ${VHD_URI} is currently being copied by a previous prefetch optimization run, will wait for completion"
+        if wait_for_vhd_copy; then
+            exit 0
         fi
+        echo "pending copy over ${VHD_URI} did not complete successfully, will delete existing blob and retry"
     else
-        echo "no existing VHD was found at: ${VHD_URI}, will proceed with optimization and VHD creation"
+        echo "VHD ${VHD_URI} exists in an unexpected copy state: '${copy_status}', will delete before proceeding"
     fi
+    delete_vhd || exit $?
 }
 
 ensure_image_builder_rg() {
@@ -92,7 +112,6 @@ run_image_builder_template() {
             -e "s#<SOURCE_TYPE>#${SOURCE_TYPE}#g" \
             -e "s#<SOURCE_ID_KEY>#${SOURCE_ID_KEY}#g" \
             -e "s#<SOURCE_ID>#${SOURCE_ID}#g" \
-            -e "s#<VHD_URI>#${VHD_URI}#g" \
             "${IMAGE_BUILDER_TEMPLATE_PATH}" > input.json || return $?
 
         if [ ! -f "input.json" ]; then
@@ -126,7 +145,7 @@ run_image_builder_template() {
         return 1
     fi
 
-    echo "template ${IMAGE_BUILDER_TEMPLATE_NAME} has ran to completion, VHD has been published to: ${VHD_URI}"
+    echo "template ${IMAGE_BUILDER_TEMPLATE_NAME} has ran to completion, optimized VHD has been published to image builder staging storage"
 }
 
 need_new_template() {
@@ -275,18 +294,58 @@ create_temp_storage() {
     TEMP_VHD_URI="${temp_vhd_uri}"
 }
 
+# copy_optimized_vhd copies the prefetch-optimized VHD that the image builder template published to its
+# staging storage account into the target storage account. Because the template distributes a VHD
+# without specifying a target uri, image builder publishes the blob to a storage account within its own
+# staging resource group and exposes its location via the run output's artifactUri. The destination vhd
+# container may have an immutability policy that blocks azcopy from writing to it directly, so we:
+# 1. azcopy the optimized VHD into the staging container (STORAGE_ACCOUNT_BLOB_URL_STAGING)
+# 2. perform a single server-side blob copy from the staging container into the immutable vhd container
+copy_optimized_vhd() {
+    artifact_uri="$(az resource show --ids "${IMAGE_BUILDER_TEMPLATE_ID}/runOutputs/${DISTRIBUTE_RUN_OUTPUT_NAME}" --api-version "${IMAGE_BUILDER_API_VERSION}" --query "properties.artifactUri" -o tsv 2>/dev/null)"
+    if [ -z "${artifact_uri}" ] || [ "${artifact_uri,,}" = "null" ] || [ "${artifact_uri,,}" = "none" ]; then
+        echo "unable to determine artifactUri for run output ${DISTRIBUTE_RUN_OUTPUT_NAME}, cannot continue"
+        return 1
+    fi
+    echo "setting azcopy environment variables with pool identity: ${IMAGE_BUILDER_IDENTITY_ID}"
+    export AZCOPY_AUTO_LOGIN_TYPE="AZCLI"
+    export AZCOPY_CONCURRENCY_VALUE="AUTO"
+    echo "copying optimized VHD from image builder staging storage to ${STAGING_VHD_URI}"
+    azcopy copy "${artifact_uri}" "${STAGING_VHD_URI}" --recursive=true
+    azcopy_exit_code=$?
+
+    if [ "${azcopy_exit_code}" -ne 0 ]; then
+        echo "failed to copy optimized VHD to staging location ${STAGING_VHD_URI}"
+        return "${azcopy_exit_code}"
+    fi
+
+    # server-side copy the VHD from the staging container into the immutable vhd container. This creates
+    # the destination blob in a single operation, which is permitted by the container's immutability policy.
+    echo "starting server-side copy of ${STAGING_VHD_URI} to ${VHD_URI}"
+    az storage blob copy start \
+        --account-name "${DESTINATION_STORAGE_ACCOUNT_NAME}" \
+        --auth-mode login \
+        --destination-container "${DESTINATION_CONTAINER_NAME}" \
+        --destination-blob "${VHD_NAME}" \
+        --source-uri "${STAGING_VHD_URI}" || return $?
+
+    wait_for_vhd_copy || return $?
+
+    echo "optimized VHD has been published to: ${VHD_URI}"
+}
+
 wait_for_vhd_copy() {
     copy_status="$(az storage blob show --blob-url "${VHD_URI}" --auth-mode login | jq -r '.properties.copy.status')"
     while [ "${copy_status,,}" = "pending" ]; do
-        echo "VHD ${VHD_URI} is still undergoing a pending copy operation"
+        echo "server-side copy to ${VHD_URI} is still pending, will wait 30s before checking again"
         sleep 30s
         copy_status="$(az storage blob show --blob-url "${VHD_URI}" --auth-mode login | jq -r '.properties.copy.status')"
     done
     if [ "${copy_status,,}" != "success" ]; then
-        echo "VHD copy over ${VHD_URI} finished with unexpected status: ${copy_status}"
+        echo "server-side copy to ${VHD_URI} finished with unexpected status: ${copy_status}"
         return 1
     fi
-    echo "pending copy operation over ${VHD_URI} has completed successfully"
+    echo "server-side copy to ${VHD_URI} completed successfully"
 }
 
 delete_vhd() {

vhdbuilder/prefetch/templates/optimize.json

@@ -34,7 +34,6 @@
         "distribute": [
             {
                 "type": "VHD",
-                "uri": "<VHD_URI>",
                 "runOutputName": "VHD"
             }
         ]