Add support for arm64 vm extension

azure/defaults.go

@@ -288,15 +288,22 @@ func ManagedClusterID(subscriptionID, resourceGroup, managedClusterName string)
 // https://learn.microsoft.com/azure/virtual-machines/extensions/custom-script-windows for Windows.
 // This extension allows running arbitrary scripts on the VM.
 // Its role is to detect and report Kubernetes bootstrap failure or success.
-func GetBootstrappingVMExtension(osType string, cloud string, vmName string) *ExtensionSpec {
+func GetBootstrappingVMExtension(osType string, cloud string, vmName string, cpuArchitectureType string) *ExtensionSpec {
 	// currently, the bootstrap extension is only available in AzurePublicCloud.
 	if osType == LinuxOS && cloud == PublicCloudName {
 		// The command checks for the existence of the bootstrapSentinelFile on the machine, with retries and sleep between retries.
+		// We set the version to 1.1.1 for arm64 machines and 1.0 for x64. This is due to a known issue with newer versions of
+		// Go on Ubuntu 20.04. The issue is being tracked here: https://github.com/golang/go/issues/58550
+		// TODO: Remove this once the issue is fixed, or when Ubuntu 20.04 is no longer supported.
+		extensionVersion := "1.0"
+		if cpuArchitectureType == "arm64" {
+			extensionVersion = "1.1.1"
+		}
 		return &ExtensionSpec{
 			Name:      BootstrappingExtensionLinux,
 			VMName:    vmName,
 			Publisher: "Microsoft.Azure.ContainerUpstream",
-			Version:   "1.0",
+			Version:   extensionVersion,
 			ProtectedSettings: map[string]string{
 				"commandToExecute": LinuxBootstrapExtensionCommand,
 			},

azure/defaults_test.go

@@ -117,3 +117,70 @@ func TestMSCorrelationIDSendDecorator(t *testing.T) {
 		receivedReq.Header.Get(string(tele.CorrIDKeyVal)),
 	).To(Equal(string(corrID)))
 }
+
+func TestGetBootstrappingVMExtension(t *testing.T) {
+	testCases := []struct {
+		name            string
+		osType          string
+		cloud           string
+		vmName          string
+		cpuArchitecture string
+		expectedVersion string
+		expectNil       bool
+	}{
+		{
+			name:            "Linux OS, Public Cloud, x64 CPU Architecture",
+			osType:          LinuxOS,
+			cloud:           PublicCloudName,
+			vmName:          "test-vm",
+			cpuArchitecture: "x64",
+			expectedVersion: "1.0",
+		},
+		{
+			name:            "Linux OS, Public Cloud, arm64 CPU Architecture",
+			osType:          LinuxOS,
+			cloud:           PublicCloudName,
+			vmName:          "test-vm",
+			cpuArchitecture: "arm64",
+			expectedVersion: "1.1.1",
+		},
+		{
+			name:            "Windows OS, Public Cloud",
+			osType:          WindowsOS,
+			cloud:           PublicCloudName,
+			vmName:          "test-vm",
+			cpuArchitecture: "x64",
+			expectedVersion: "1.0",
+		},
+		{
+			name:            "Invalid OS Type",
+			osType:          "invalid",
+			cloud:           PublicCloudName,
+			vmName:          "test-vm",
+			cpuArchitecture: "x64",
+			expectedVersion: "1.0",
+			expectNil:       true,
+		},
+		{
+			name:            "Invalid Cloud",
+			osType:          LinuxOS,
+			cloud:           "invalid",
+			vmName:          "test-vm",
+			cpuArchitecture: "x64",
+			expectedVersion: "1.0",
+			expectNil:       true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			g := NewWithT(t)
+			actualExtension := GetBootstrappingVMExtension(tc.osType, tc.cloud, tc.vmName, tc.cpuArchitecture)
+			if tc.expectNil {
+				g.Expect(actualExtension).To(BeNil())
+			} else {
+				g.Expect(actualExtension.Version).To(Equal(tc.expectedVersion))
+			}
+		})
+	}
+}

azure/scope/machine.go

@@ -375,7 +375,8 @@ func (m *MachineScope) VMExtensionSpecs() []azure.ResourceSpecGetter {
 		})
 	}
 
-	bootstrapExtensionSpec := azure.GetBootstrappingVMExtension(m.AzureMachine.Spec.OSDisk.OSType, m.CloudEnvironment(), m.Name())
+	cpuArchitectureType, _ := m.cache.VMSKU.GetCapability(resourceskus.CPUArchitectureType)
+	bootstrapExtensionSpec := azure.GetBootstrappingVMExtension(m.AzureMachine.Spec.OSDisk.OSType, m.CloudEnvironment(), m.Name(), cpuArchitectureType)
 
 	if bootstrapExtensionSpec != nil {
 		extensionSpecs = append(extensionSpecs, &vmextensions.VMExtensionSpec{

azure/scope/machine_test.go

@@ -573,6 +573,9 @@ func TestMachineScope_VMExtensionSpecs(t *testing.T) {
 						},
 					},
 				},
+				cache: &MachineCache{
+					VMSKU: resourceskus.SKU{},
+				},
 			},
 			want: []azure.ResourceSpecGetter{
 				&vmextensions.VMExtensionSpec{
@@ -621,6 +624,9 @@ func TestMachineScope_VMExtensionSpecs(t *testing.T) {
 						},
 					},
 				},
+				cache: &MachineCache{
+					VMSKU: resourceskus.SKU{},
+				},
 			},
 			want: []azure.ResourceSpecGetter{},
 		},
@@ -655,6 +661,9 @@ func TestMachineScope_VMExtensionSpecs(t *testing.T) {
 						},
 					},
 				},
+				cache: &MachineCache{
+					VMSKU: resourceskus.SKU{},
+				},
 			},
 			want: []azure.ResourceSpecGetter{
 				&vmextensions.VMExtensionSpec{
@@ -703,6 +712,9 @@ func TestMachineScope_VMExtensionSpecs(t *testing.T) {
 						},
 					},
 				},
+				cache: &MachineCache{
+					VMSKU: resourceskus.SKU{},
+				},
 			},
 			want: []azure.ResourceSpecGetter{},
 		},
@@ -737,6 +749,9 @@ func TestMachineScope_VMExtensionSpecs(t *testing.T) {
 						},
 					},
 				},
+				cache: &MachineCache{
+					VMSKU: resourceskus.SKU{},
+				},
 			},
 			want: []azure.ResourceSpecGetter{},
 		},
@@ -771,6 +786,9 @@ func TestMachineScope_VMExtensionSpecs(t *testing.T) {
 						},
 					},
 				},
+				cache: &MachineCache{
+					VMSKU: resourceskus.SKU{},
+				},
 			},
 			want: []azure.ResourceSpecGetter{},
 		},
@@ -818,6 +836,9 @@ func TestMachineScope_VMExtensionSpecs(t *testing.T) {
 						},
 					},
 				},
+				cache: &MachineCache{
+					VMSKU: resourceskus.SKU{},
+				},
 			},
 			want: []azure.ResourceSpecGetter{
 				&vmextensions.VMExtensionSpec{

azure/scope/machinepool.go

@@ -32,6 +32,7 @@ import (
 	infrav1 "sigs.k8s.io/cluster-api-provider-azure/api/v1beta1"
 	"sigs.k8s.io/cluster-api-provider-azure/azure"
 	machinepool "sigs.k8s.io/cluster-api-provider-azure/azure/scope/strategies/machinepool_deployments"
+	"sigs.k8s.io/cluster-api-provider-azure/azure/services/resourceskus"
 	"sigs.k8s.io/cluster-api-provider-azure/azure/services/roleassignments"
 	"sigs.k8s.io/cluster-api-provider-azure/azure/services/scalesets"
 	"sigs.k8s.io/cluster-api-provider-azure/azure/services/virtualmachineimages"
@@ -61,6 +62,7 @@ type (
 		MachinePool      *expv1.MachinePool
 		AzureMachinePool *infrav1exp.AzureMachinePool
 		ClusterScope     azure.ClusterScoper
+		Caches           *MachinePoolCache
 	}
 
 	// MachinePoolScope defines a scope defined around a machine pool and its cluster.
@@ -72,6 +74,12 @@ type (
 		patchHelper                *patch.Helper
 		capiMachinePoolPatchHelper *patch.Helper
 		vmssState                  *azure.VMSS
+		cache                      *MachinePoolCache
+	}
+
+	// MachinePoolCache stores common machine pool information so we don't have to hit the API multiple times within the same reconcile loop.
+	MachinePoolCache struct {
+		VMSKU resourceskus.SKU
 	}
 
 	// NodeStatus represents the status of a Kubernetes node.
@@ -116,6 +124,29 @@ func NewMachinePoolScope(params MachinePoolScopeParams) (*MachinePoolScope, erro
 	}, nil
 }
 
+// InitMachinePoolCache sets cached information about the machine pool to be used in the scope.
+func (m *MachinePoolScope) InitMachinePoolCache(ctx context.Context) error {
+	ctx, _, done := tele.StartSpanWithLogger(ctx, "scope.MachinePoolScope.InitMachinePoolCache")
+	defer done()
+
+	if m.cache == nil {
+		var err error
+		m.cache = &MachinePoolCache{}
+
+		skuCache, err := resourceskus.GetCache(m, m.Location())
+		if err != nil {
+			return err
+		}
+
+		m.cache.VMSKU, err = skuCache.Get(ctx, m.AzureMachinePool.Spec.Template.VMSize, resourceskus.VirtualMachines)
+		if err != nil {
+			return errors.Wrapf(err, "failed to get VM SKU %s in compute api", m.AzureMachinePool.Spec.Template.VMSize)
+		}
+	}
+
+	return nil
+}
+
 // ScaleSetSpec returns the scale set spec.
 func (m *MachinePoolScope) ScaleSetSpec() azure.ScaleSetSpec {
 	return azure.ScaleSetSpec{
@@ -699,7 +730,8 @@ func (m *MachinePoolScope) VMSSExtensionSpecs() []azure.ResourceSpecGetter {
 		})
 	}
 
-	bootstrapExtensionSpec := azure.GetBootstrappingVMExtension(m.AzureMachinePool.Spec.Template.OSDisk.OSType, m.CloudEnvironment(), m.Name())
+	cpuArchitectureType, _ := m.cache.VMSKU.GetCapability(resourceskus.CPUArchitectureType)
+	bootstrapExtensionSpec := azure.GetBootstrappingVMExtension(m.AzureMachinePool.Spec.Template.OSDisk.OSType, m.CloudEnvironment(), m.Name(), cpuArchitectureType)
 
 	if bootstrapExtensionSpec != nil {
 		extensionSpecs = append(extensionSpecs, &scalesets.VMSSExtensionSpec{

azure/scope/machinepool_test.go

@@ -34,6 +34,7 @@ import (
 	infrav1 "sigs.k8s.io/cluster-api-provider-azure/api/v1beta1"
 	"sigs.k8s.io/cluster-api-provider-azure/azure"
 	"sigs.k8s.io/cluster-api-provider-azure/azure/mock_azure"
+	"sigs.k8s.io/cluster-api-provider-azure/azure/services/resourceskus"
 	"sigs.k8s.io/cluster-api-provider-azure/azure/services/roleassignments"
 	"sigs.k8s.io/cluster-api-provider-azure/azure/services/scalesets"
 	infrav1exp "sigs.k8s.io/cluster-api-provider-azure/exp/api/v1beta1"
@@ -886,6 +887,9 @@ func TestMachinePoolScope_VMSSExtensionSpecs(t *testing.T) {
 						},
 					},
 				},
+				cache: &MachinePoolCache{
+					VMSKU: resourceskus.SKU{},
+				},
 			},
 			want: []azure.ResourceSpecGetter{
 				&scalesets.VMSSExtensionSpec{
@@ -932,6 +936,9 @@ func TestMachinePoolScope_VMSSExtensionSpecs(t *testing.T) {
 						},
 					},
 				},
+				cache: &MachinePoolCache{
+					VMSKU: resourceskus.SKU{},
+				},
 			},
 			want: []azure.ResourceSpecGetter{},
 		},
@@ -966,6 +973,9 @@ func TestMachinePoolScope_VMSSExtensionSpecs(t *testing.T) {
 						},
 					},
 				},
+				cache: &MachinePoolCache{
+					VMSKU: resourceskus.SKU{},
+				},
 			},
 			want: []azure.ResourceSpecGetter{
 				&scalesets.VMSSExtensionSpec{
@@ -1013,6 +1023,9 @@ func TestMachinePoolScope_VMSSExtensionSpecs(t *testing.T) {
 						},
 					},
 				},
+				cache: &MachinePoolCache{
+					VMSKU: resourceskus.SKU{},
+				},
 			},
 			want: []azure.ResourceSpecGetter{},
 		},
@@ -1046,6 +1059,9 @@ func TestMachinePoolScope_VMSSExtensionSpecs(t *testing.T) {
 						},
 					},
 				},
+				cache: &MachinePoolCache{
+					VMSKU: resourceskus.SKU{},
+				},
 			},
 			want: []azure.ResourceSpecGetter{},
 		},
@@ -1079,6 +1095,9 @@ func TestMachinePoolScope_VMSSExtensionSpecs(t *testing.T) {
 						},
 					},
 				},
+				cache: &MachinePoolCache{
+					VMSKU: resourceskus.SKU{},
+				},
 			},
 			want: []azure.ResourceSpecGetter{},
 		},
@@ -1128,6 +1147,9 @@ func TestMachinePoolScope_VMSSExtensionSpecs(t *testing.T) {
 						},
 					},
 				},
+				cache: &MachinePoolCache{
+					VMSKU: resourceskus.SKU{},
+				},
 			},
 			want: []azure.ResourceSpecGetter{
 				&scalesets.VMSSExtensionSpec{

azure/services/resourceskus/sku.go

@@ -72,6 +72,8 @@ const (
 	TrustedLaunchDisabled = "TrustedLaunchDisabled"
 	// ConfidentialComputingType identifies the capability for confidentical computing.
 	ConfidentialComputingType = "ConfidentialComputingType"
+	// CPUArchitectureType identifies the capability for cpu architecture.
+	CPUArchitectureType = "CpuArchitectureType"
 )
 
 // HasCapability return true for a capability which can be either

exp/controllers/azuremachinepool_controller.go

@@ -35,6 +35,7 @@ import (
 	"sigs.k8s.io/cluster-api-provider-azure/util/tele"
 	clusterv1 "sigs.k8s.io/cluster-api/api/v1beta1"
 	kubeadmv1 "sigs.k8s.io/cluster-api/bootstrap/kubeadm/api/v1beta1"
+	capierrors "sigs.k8s.io/cluster-api/errors"
 	expv1 "sigs.k8s.io/cluster-api/exp/api/v1beta1"
 	"sigs.k8s.io/cluster-api/util"
 	"sigs.k8s.io/cluster-api/util/annotations"
@@ -293,6 +294,22 @@ func (ampr *AzureMachinePoolReconciler) reconcileNormal(ctx context.Context, mac
 		return reconcile.Result{}, nil
 	}
 
+	var reconcileError azure.ReconcileError
+
+	// Initialize the cache to be used by the AzureMachine services.
+	err := machinePoolScope.InitMachinePoolCache(ctx)
+	if err != nil {
+		if errors.As(err, &reconcileError) && reconcileError.IsTerminal() {
+			ampr.Recorder.Eventf(machinePoolScope.AzureMachinePool, corev1.EventTypeWarning, "SKUNotFound", errors.Wrap(err, "failed to initialize machine cache").Error())
+			log.Error(err, "Failed to initialize machine cache")
+			machinePoolScope.SetFailureReason(capierrors.InvalidConfigurationMachineError)
+			machinePoolScope.SetFailureMessage(err)
+			machinePoolScope.SetNotReady()
+			return reconcile.Result{}, nil
+		}
+		return reconcile.Result{}, errors.Wrap(err, "failed to init machine scope cache")
+	}
+
 	ams, err := ampr.createAzureMachinePoolService(machinePoolScope)
 	if err != nil {
 		return reconcile.Result{}, errors.Wrap(err, "failed creating a newAzureMachinePoolService")