fix: add retries to wireserver, cache first time response to avoid spam

Author: awesomenix

parts/linux/cloud-init/artifacts/cse_config.sh

@@ -343,7 +343,7 @@ EOF
   if [ "${GPU_NODE}" = "true" ]; then
     # Check VM tag directly to determine if GPU drivers should be skipped
     export -f should_skip_nvidia_drivers
-    should_skip=$(retrycmd_silent 10 1 10 bash -cx should_skip_nvidia_drivers)
+    should_skip=$(should_skip_nvidia_drivers)
     if [ "$?" -eq 0 ] && [ "${should_skip}" = "true" ]; then
       echo "Generating non-GPU containerd config for GPU node due to VM tags"
       echo "${CONTAINERD_CONFIG_NO_GPU_CONTENT}" | base64 -d > /etc/containerd/config.toml || exit $ERR_FILE_WATCH_TIMEOUT
@@ -356,8 +356,8 @@ EOF
     echo "${CONTAINERD_CONFIG_CONTENT}" | base64 -d > /etc/containerd/config.toml || exit $ERR_FILE_WATCH_TIMEOUT
   fi
 
-  export -f e2e_mock_azure_china_cloud
-  E2EMockAzureChinaCloud=$(retrycmd_silent 10 1 10 bash -cx e2e_mock_azure_china_cloud)
+  export -f should_e2e_mock_azure_china_cloud
+  E2EMockAzureChinaCloud=$(should_e2e_mock_azure_china_cloud)
   if [ -n "${BOOTSTRAP_PROFILE_CONTAINER_REGISTRY_SERVER}" ]; then
     logs_to_events "AKS.CSE.ensureContainerd.configureContainerdRegistryHost" configureContainerdRegistryHost
   elif [ "${TARGET_CLOUD}" = "AzureChinaCloud" ] || [ "${E2EMockAzureChinaCloud}" = "true" ]; then
@@ -428,23 +428,10 @@ ensureDHCPv6() {
 }
 
 getPrimaryNicIP() {
-    local sleepTime=1
-    local maxRetries=10
-    local i=0
     local ip=""
-
-    while [ "$i" -lt "$maxRetries" ]; do
-        ip=$(curl -sSL -H "Metadata: true" "http://169.254.169.254/metadata/instance/network/interface?api-version=2021-02-01")
-        if [ "$?" -eq 0 ]; then
-            ip=$(echo "$ip" | jq -r '.[0].ipv4.ipAddress[0].privateIpAddress')
-            if [ -n "$ip" ]; then
-                break
-            fi
-        fi
-        sleep $sleepTime
-        i=$((i+1))
-    done
-    echo "$ip"
+    export -f get_primary_nic_ip
+    ip=$(get_primary_nic_ip)
+    echo "${ip}"
 }
 
 generateSelfSignedKubeletServingCertificate() {
@@ -470,7 +457,7 @@ configureKubeletServing() {
 
     # check if kubelet serving certificate rotation is disabled by customer-specified nodepool tags
     export -f should_disable_kubelet_serving_certificate_rotation
-    DISABLE_KUBELET_SERVING_CERTIFICATE_ROTATION=$(retrycmd_silent 10 1 10 bash -cx should_disable_kubelet_serving_certificate_rotation)
+    DISABLE_KUBELET_SERVING_CERTIFICATE_ROTATION=$(should_disable_kubelet_serving_certificate_rotation)
     if [ "$?" -ne 0 ]; then
         echo "failed to determine if kubelet serving certificate rotation should be disabled by nodepool tags"
         exit $ERR_LOOKUP_DISABLE_KUBELET_SERVING_CERTIFICATE_ROTATION_TAG
@@ -1131,18 +1118,18 @@ configCredentialProvider() {
 }
 
 setKubeletNodeIPFlag() {
-    imdsOutput=$(curl -s -H Metadata:true --noproxy "*" --max-time 5 "http://169.254.169.254/metadata/instance/network/interface?api-version=2021-02-01" 2> /dev/null)
-    if [ "$?" -eq 0 ]; then
-        nodeIPAddrs=()
-        ipv4Addr=$(echo $imdsOutput | jq -r '.[0].ipv4.ipAddress[0].privateIpAddress // ""')
-        [ -n "$ipv4Addr" ] && nodeIPAddrs+=("$ipv4Addr")
-        ipv6Addr=$(echo $imdsOutput | jq -r '.[0].ipv6.ipAddress[0].privateIpAddress // ""')
-        [ -n "$ipv6Addr" ] && nodeIPAddrs+=("$ipv6Addr")
-        nodeIPArg=$(IFS=, ; echo "${nodeIPAddrs[*]}") # join, comma-separated
-        if [ -n "$nodeIPArg" ]; then
-            echo "Adding --node-ip=$nodeIPArg to kubelet flags"
-            KUBELET_FLAGS="$KUBELET_FLAGS --node-ip=$nodeIPArg"
-        fi
+    local imdsOutput
+    export -f get_imds_network_metadata
+    imdsOutput=$(get_imds_network_metadata)
+    nodeIPAddrs=()
+    ipv4Addr=$(echo $imdsOutput | jq -r '.[0].ipv4.ipAddress[0].privateIpAddress // ""')
+    [ -n "$ipv4Addr" ] && nodeIPAddrs+=("$ipv4Addr")
+    ipv6Addr=$(echo $imdsOutput | jq -r '.[0].ipv6.ipAddress[0].privateIpAddress // ""')
+    [ -n "$ipv6Addr" ] && nodeIPAddrs+=("$ipv6Addr")
+    nodeIPArg=$(IFS=, ; echo "${nodeIPAddrs[*]}") # join, comma-separated
+    if [ -n "$nodeIPArg" ]; then
+        echo "Adding --node-ip=$nodeIPArg to kubelet flags"
+        KUBELET_FLAGS="$KUBELET_FLAGS --node-ip=$nodeIPArg"
     fi
 }
 

parts/linux/cloud-init/artifacts/cse_helpers.sh

@@ -121,6 +121,8 @@ ERR_ORAS_IMDS_TIMEOUT=210 # Error timeout waiting for IMDS response
 ERR_ORAS_PULL_NETWORK_TIMEOUT=211 # Error pulling oras tokens for login
 ERR_ORAS_PULL_UNAUTHORIZED=212 # Error pulling artifact with oras from registry with authorization issue
 
+ERR_IMDS_FETCH_FAILED=231 # Error fetching or caching IMDS instance metadata
+
 # Error checking nodepools tags for whether we need to disable kubelet serving certificate rotation
 ERR_LOOKUP_DISABLE_KUBELET_SERVING_CERTIFICATE_ROTATION_TAG=213
 
@@ -651,70 +653,128 @@ logs_to_events() {
     fi
 }
 
-should_skip_nvidia_drivers() {
+# Cache file for IMDS instance metadata to avoid redundant network calls.
+# The IMDS endpoint provides VM metadata that doesn't change during provisioning,
+# so we can safely cache it for the duration of CSE execution.
+IMDS_INSTANCE_METADATA_CACHE_FILE="/opt/azure/containers/imds_instance_metadata_cache.json"
+
+# Fetches IMDS instance metadata and caches it to a file.
+# Exits with ERR_IMDS_FETCH_FAILED on failure.
+# The cache is stored in IMDS_INSTANCE_METADATA_CACHE_FILE.
+fetch_and_cache_imds_instance_metadata() {
     set -x
-    body=$(curl -fsSL -H "Metadata: true" --noproxy "*" "http://169.254.169.254/metadata/instance?api-version=2021-02-01")
-    ret=$?
-    if [ "$ret" -ne 0 ]; then
-      return $ret
+    # If cache file already exists and is valid, skip fetching
+    if [ -f "$IMDS_INSTANCE_METADATA_CACHE_FILE" ]; then
+        # Validate the cache file contains valid JSON
+        if jq -e . "$IMDS_INSTANCE_METADATA_CACHE_FILE" > /dev/null 2>&1; then
+            echo "IMDS instance metadata cache already exists and is valid"
+            return 0
+        else
+            echo "IMDS instance metadata cache exists but is invalid, refetching"
+            rm -f "$IMDS_INSTANCE_METADATA_CACHE_FILE"
+        fi
     fi
-    should_skip=$(echo "$body" | jq -e '.compute.tagsList | map(select(.name | test("SkipGpuDriverInstall"; "i")))[0].value // "false" | test("true"; "i")')
+
+    local body
+    body=$(curl -fsSL -H "Metadata: true" --noproxy "*" --retry 20 --retry-delay 2 --retry-connrefused --connect-timeout 5 --max-time 5 "http://169.254.169.254/metadata/instance?api-version=2021-02-01")
+    if [ "$?" -ne 0 ]; then
+        echo "Failed to fetch IMDS instance metadata"
+        exit $ERR_IMDS_FETCH_FAILED
+    fi
+
+    # Validate response is valid JSON before caching
+    if ! echo "$body" | jq -e . > /dev/null 2>&1; then
+        echo "IMDS response is not valid JSON"
+        exit $ERR_IMDS_FETCH_FAILED
+    fi
+
+    echo "$body" > "$IMDS_INSTANCE_METADATA_CACHE_FILE"
+    echo "IMDS instance metadata cached successfully"
+}
+
+# Retrieves IMDS network interface metadata from the cached instance metadata.
+# The /metadata/instance endpoint contains both compute and network data,
+# so we extract .network.interface from the instance metadata cache.
+# Outputs:
+#   The JSON network interface array
+# Returns:
+#   0 on success, non-zero on failure
+get_imds_network_metadata() {
+    local network_metadata=""
+    network_metadata=$(jq -r '.network.interface' "$IMDS_INSTANCE_METADATA_CACHE_FILE")
+    echo "$network_metadata"
+}
+
+# Retrieves the primary NIC's private IPv4 address from IMDS instance metadata.
+# Outputs:
+#   The private IPv4 address or empty string if not found.
+# Returns:
+#   0 on success, non-zero on failure
+get_primary_nic_ip() {
+    local primary_ip=""
+    primary_ip=$(jq -r '.network.interface[0].ipv4.ipAddress[0].privateIpAddress // ""' "$IMDS_INSTANCE_METADATA_CACHE_FILE")
+    echo "$primary_ip"
+}
+
+# Retrieves a VM tag value from the cached IMDS instance metadata.
+# Arguments:
+#   $1 - tag_name: The name of the tag to look up (case-sensitive match)
+#   $2 - case_insensitive_name: Optional. If "true", performs case-insensitive name matching
+#   $3 - case_insensitive_value: Optional. If "true", performs case-insensitive value matching for "true"
+# Outputs:
+#   The tag value (lowercased) or empty string if not found.
+#   For case_insensitive_value="true", outputs "true" or "false" based on whether value matches "true" case-insensitively.
+# Returns:
+#   0 on success, 1 if cache file doesn't exist or tag not found
+get_imds_vm_tag_value() {
+    local tag_name="$1"
+    local tag_value
+    tag_value=$(jq -r --arg name "$tag_name" '.compute.tagsList | map(select(.name | test($name; "i")))[0].value // "false" | test("true"; "i")' "$IMDS_INSTANCE_METADATA_CACHE_FILE")
+    # Output lowercase value for consistency
+    echo "${tag_value,,}"
+}
+
+should_skip_nvidia_drivers() {
+    set -x
+    # Case-insensitive match for both tag name and value
+    local should_skip
+    should_skip=$(get_imds_vm_tag_value "SkipGpuDriverInstall")
     echo "$should_skip"
 }
 
 should_disable_kubelet_serving_certificate_rotation() {
     set -x
-    body=$(curl -fsSL -H "Metadata: true" --noproxy "*" "http://169.254.169.254/metadata/instance?api-version=2021-02-01")
-    ret=$?
-    if [ "$ret" -ne 0 ]; then
-      return $ret
-    fi
-    should_disable=$(echo "$body" | jq -r '.compute.tagsList[] | select(.name == "aks-disable-kubelet-serving-certificate-rotation") | .value')
-    echo "${should_disable,,}"
+    local should_disable
+    should_disable=$(get_imds_vm_tag_value "aks-disable-kubelet-serving-certificate-rotation")
+    echo "$should_disable"
 }
 
 should_skip_binary_cleanup() {
     set -x
-    body=$(curl -fsSL -H "Metadata: true" --noproxy "*" "http://169.254.169.254/metadata/instance?api-version=2021-02-01")
-    ret=$?
-    if [ "$ret" -ne 0 ]; then
-      return $ret
-    fi
-    should_skip=$(echo "$body" | jq -r '.compute.tagsList[] | select(.name == "SkipBinaryCleanup") | .value')
-    echo "${should_skip,,}"
+    local should_skip
+    should_skip=$(get_imds_vm_tag_value "SkipBinaryCleanup")
+    echo "$should_skip"
 }
 
 should_enforce_kube_pmc_install() {
     set -x
-    body=$(curl -fsSL -H "Metadata: true" --noproxy "*" "http://169.254.169.254/metadata/instance?api-version=2021-02-01")
-    ret=$?
-    if [ "$ret" -ne 0 ]; then
-      return $ret
-    fi
-    should_enforce=$(echo "$body" | jq -r '.compute.tagsList[] | select(.name == "ShouldEnforceKubePMCInstall") | .value')
-    echo "${should_enforce,,}"
+    local should_enforce
+    should_enforce=$(get_imds_vm_tag_value "ShouldEnforceKubePMCInstall")
+    echo "$should_enforce"
 }
 
-e2e_mock_azure_china_cloud() {
+should_e2e_mock_azure_china_cloud() {
     set -x
-    body=$(curl -fsSL -H "Metadata: true" --noproxy "*" "http://169.254.169.254/metadata/instance?api-version=2021-02-01")
-    ret=$?
-    if [ "$ret" -ne 0 ]; then
-      return $ret
-    fi
-    should_enforce=$(echo "$body" | jq -r '.compute.tagsList[] | select(.name == "E2EMockAzureChinaCloud") | .value')
-    echo "${should_enforce,,}"
+    local should_enforce
+    should_enforce=$(get_imds_vm_tag_value "E2EMockAzureChinaCloud")
+    echo "$should_enforce"
 }
 
-enableManagedGPUExperience() {
+should_enable_managed_gpu_experience() {
     set -x
-    body=$(curl -fsSL -H "Metadata: true" --noproxy "*" "http://169.254.169.254/metadata/instance?api-version=2021-02-01")
-    ret=$?
-    if [ "$ret" -ne 0 ]; then
-      return $ret
-    fi
-    should_enforce=$(echo "$body" | jq -r '.compute.tagsList[] | select(.name == "EnableManagedGPUExperience") | .value')
-    echo "${should_enforce,,}"
+    local should_enforce
+    should_enforce=$(get_imds_vm_tag_value "EnableManagedGPUExperience")
+    echo "$should_enforce"
 }
 
 isMarinerOrAzureLinux() {

parts/linux/cloud-init/artifacts/cse_main.sh

@@ -8,6 +8,9 @@ if [ -f /opt/azure/containers/provision.complete ]; then
     exit 0
 fi
 
+# Cleanup cache file to force fetch fresh instance metadata from IMDS
+rm -f /opt/azure/containers/imds_instance_metadata_cache.json
+
 for i in $(seq 1 120); do
     if [ -s "${CSE_HELPERS_FILEPATH}" ]; then
         grep -Fq '#HELPERSEOF' "${CSE_HELPERS_FILEPATH}" && break
@@ -130,7 +133,7 @@ function basePrep {
     export -f getInstallModeAndCleanupContainerImages
     export -f should_skip_binary_cleanup
 
-    SKIP_BINARY_CLEANUP=$(retrycmd_silent 10 1 10 bash -cx should_skip_binary_cleanup)
+    SKIP_BINARY_CLEANUP=$(should_skip_binary_cleanup)
     # this needs better fix to separate logs and return value;
     FULL_INSTALL_REQUIRED=$(getInstallModeAndCleanupContainerImages "$SKIP_BINARY_CLEANUP" "$IS_VHD" | tail -1)
     if [ "$?" -ne 0 ]; then
@@ -163,7 +166,7 @@ function basePrep {
     # Added as a temporary workaround to test installing packages from PMC prior to 1.34.0 GA.
     # TODO: Remove tag and usages once 1.34.0 is GA.
     export -f should_enforce_kube_pmc_install
-    SHOULD_ENFORCE_KUBE_PMC_INSTALL=$(retrycmd_silent 10 1 10 bash -cx should_enforce_kube_pmc_install)
+    SHOULD_ENFORCE_KUBE_PMC_INSTALL=$(should_enforce_kube_pmc_install)
     logs_to_events "AKS.CSE.configureKubeletAndKubectl" configureKubeletAndKubectl
 
     createKubeManifestDir
@@ -310,6 +313,7 @@ EOF
 # After this stage the node should be fully integrated into the cluster.
 # IMPORTANT: This stage should only run when actually joining a node to the cluster. This step should not be run when creating a VHD image
 function nodePrep {
+    logs_to_events "AKS.CSE.fetch_and_cache_imds_instance_metadata" fetch_and_cache_imds_instance_metadata
     # IMPORTANT NOTE: We do this here since this function can mutate kubelet flags and node labels,
     # which is used by configureK8s and other functions. Thus, we need to make sure flag and label content is correct beforehand.
     logs_to_events "AKS.CSE.configureKubeletServing" configureKubeletServing
@@ -337,7 +341,7 @@ function nodePrep {
 
     # Determine if GPU driver installation should be skipped
     export -f should_skip_nvidia_drivers
-    skip_nvidia_driver_install=$(retrycmd_silent 10 1 10 bash -cx should_skip_nvidia_drivers)
+    skip_nvidia_driver_install=$(should_skip_nvidia_drivers)
 
     if [ "$?" -ne 0 ]; then
         echo "Failed to determine if nvidia driver install should be skipped"
@@ -395,8 +399,8 @@ function nodePrep {
         echo $(date),$(hostname), "End configuring GPU drivers"
     fi
 
-    export -f enableManagedGPUExperience
-    ENABLE_MANAGED_GPU_EXPERIENCE=$(retrycmd_silent 10 1 10 bash -cx enableManagedGPUExperience)
+    export -f should_enable_managed_gpu_experience
+    ENABLE_MANAGED_GPU_EXPERIENCE=$(should_enable_managed_gpu_experience)
     if [ "$?" -ne 0 ] && [ "${GPU_NODE}" = "true" ] && [ "${skip_nvidia_driver_install}" != "true" ]; then
         echo "failed to determine if managed GPU experience should be enabled by nodepool tags"
         exit $ERR_LOOKUP_ENABLE_MANAGED_GPU_EXPERIENCE_TAG

parts/linux/cloud-init/artifacts/cse_start.sh

@@ -49,6 +49,9 @@ JSON_STRING=$( jq -n \
 mkdir -p /var/log/azure/aks
 echo $JSON_STRING | tee /var/log/azure/aks/provision.json
 
+# Cleanup cache file
+rm -f /opt/azure/containers/imds_instance_metadata_cache.json || true
+
 # Create stage marker for two-stage workflow
 if [ "${PRE_PROVISION_ONLY}" = "true" ]; then
     # Stage 1: Create marker indicating Stage 2 is needed
@@ -118,4 +121,4 @@ else
     upload_logs &
 fi
 
-exit "$EXIT_CODE"
+exit "$EXIT_CODE"