Add support for edge zones in Azure provider

Author: Jont828

cluster-autoscaler/cloudprovider/azure/azure_scale_set.go

@@ -254,6 +254,16 @@ func (scaleSet *ScaleSet) SetScaleSetSize(size int64) error {
 		Sku:      vmssInfo.Sku,
 		Location: vmssInfo.Location,
 	}
+
+	if vmssInfo.ExtendedLocation != nil {
+		op.ExtendedLocation = &compute.ExtendedLocation{
+			Name: vmssInfo.ExtendedLocation.Name,
+			Type: vmssInfo.ExtendedLocation.Type,
+		}
+
+		klog.V(3).Infof("Passing ExtendedLocation information if it is not nil, with Edge Zone name:(%s)", *op.ExtendedLocation.Name)
+	}
+
 	ctx, cancel := getContextWithTimeout(vmssContextTimeout)
 	defer cancel()
 	klog.V(3).Infof("Waiting for virtualMachineScaleSetsClient.CreateOrUpdateAsync(%s)", scaleSet.Name)

cluster-autoscaler/cloudprovider/azure/azure_scale_set_test.go

@@ -32,6 +32,11 @@ import (
 	"sigs.k8s.io/cloud-provider-azure/pkg/azureclients/vmssvmclient/mockvmssvmclient"
 )
 
+const (
+	testASG      = "test-asg"
+	testLocation = "eastus"
+)
+
 func newTestScaleSet(manager *AzureManager, name string) *ScaleSet {
 	return &ScaleSet{
 		azureRef: azureRef{
@@ -44,6 +49,18 @@ func newTestScaleSet(manager *AzureManager, name string) *ScaleSet {
 	}
 }
 
+func newTestScaleSetMinSizeZero(manager *AzureManager, name string) *ScaleSet {
+	return &ScaleSet{
+		azureRef: azureRef{
+			Name: name,
+		},
+		manager:           manager,
+		minSize:           0,
+		maxSize:           5,
+		enableForceDelete: manager.config.EnableForceDelete,
+	}
+}
+
 func newTestVMSSList(cap int64, name, loc string, orchmode compute.OrchestrationMode) []compute.VirtualMachineScaleSet {
 	return []compute.VirtualMachineScaleSet{
 		{
@@ -61,6 +78,22 @@ func newTestVMSSList(cap int64, name, loc string, orchmode compute.Orchestration
 	}
 }
 
+func newTestVMSSListForEdgeZones(capacity int64, name string) *compute.VirtualMachineScaleSet {
+	return &compute.VirtualMachineScaleSet{
+		Name: to.StringPtr(name),
+		Sku: &compute.Sku{
+			Capacity: to.Int64Ptr(capacity),
+			Name:     to.StringPtr("Standard_D4_v2"),
+		},
+		VirtualMachineScaleSetProperties: &compute.VirtualMachineScaleSetProperties{},
+		Location:                         to.StringPtr(testLocation),
+		ExtendedLocation: &compute.ExtendedLocation{
+			Name: to.StringPtr("losangeles"),
+			Type: compute.ExtendedLocationTypes("EdgeZone"),
+		},
+	}
+}
+
 func newTestVMSSVMList(count int) []compute.VirtualMachineScaleSetVM {
 	var vmssVMList []compute.VirtualMachineScaleSetVM
 	for i := 0; i < count; i++ {
@@ -122,6 +155,15 @@ func TestMinSize(t *testing.T) {
 	assert.Equal(t, provider.NodeGroups()[0].MinSize(), 1)
 }
 
+func TestMinSizeZero(t *testing.T) {
+	provider := newTestProvider(t)
+	registered := provider.azureManager.RegisterNodeGroup(
+		newTestScaleSetMinSizeZero(provider.azureManager, testASG))
+	assert.True(t, registered)
+	assert.Equal(t, len(provider.NodeGroups()), 1)
+	assert.Equal(t, provider.NodeGroups()[0].MinSize(), 0)
+}
+
 func TestTargetSize(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
@@ -177,24 +219,33 @@ func TestIncreaseSize(t *testing.T) {
 	for _, orchMode := range orchestrationModes {
 
 		provider := newTestProvider(t)
-		expectedScaleSets := newTestVMSSList(3, "test-asg", "eastus", orchMode)
+		expectedScaleSets := newTestVMSSList(3, testASG, testLocation, orchMode)
+
+		// Include Edge Zone scenario here, testing scale from 3 to 5 and scale from zero cases.
+		expectedEdgeZoneScaleSets := newTestVMSSListForEdgeZones(3, "edgezone-vmss")
+		expectedEdgeZoneMinZeroScaleSets := newTestVMSSListForEdgeZones(0, "edgezone-minzero-vmss")
+		expectedScaleSets = append(expectedScaleSets, *expectedEdgeZoneScaleSets, *expectedEdgeZoneMinZeroScaleSets)
 
 		mockVMSSClient := mockvmssclient.NewMockInterface(ctrl)
 		mockVMSSClient.EXPECT().List(gomock.Any(), provider.azureManager.config.ResourceGroup).Return(expectedScaleSets, nil).AnyTimes()
-		mockVMSSClient.EXPECT().CreateOrUpdateAsync(gomock.Any(), provider.azureManager.config.ResourceGroup, "test-asg", gomock.Any()).Return(nil, nil)
+		mockVMSSClient.EXPECT().CreateOrUpdateAsync(gomock.Any(), provider.azureManager.config.ResourceGroup, testASG, gomock.Any()).Return(nil, nil)
+		// This should be Anytimes() because the parent function of this call - updateVMSSCapacity() is a goroutine
+		// and this test doesn't wait on goroutine, hence, it is difficult to write exact expected number (which is 3 here)
+		// before we return from this this.
+		// This is a future TODO: sync.WaitGroup should be used in actual code and make code easily testable
 		mockVMSSClient.EXPECT().WaitForCreateOrUpdateResult(gomock.Any(), gomock.Any(), provider.azureManager.config.ResourceGroup).Return(&http.Response{StatusCode: http.StatusOK}, nil).AnyTimes()
 		provider.azureManager.azClient.virtualMachineScaleSetsClient = mockVMSSClient
 
 		if orchMode == compute.Uniform {
 
 			mockVMSSVMClient := mockvmssvmclient.NewMockInterface(ctrl)
-			mockVMSSVMClient.EXPECT().List(gomock.Any(), provider.azureManager.config.ResourceGroup, "test-asg", gomock.Any()).Return(expectedVMSSVMs, nil).AnyTimes()
+			mockVMSSVMClient.EXPECT().List(gomock.Any(), provider.azureManager.config.ResourceGroup, testASG, gomock.Any()).Return(expectedVMSSVMs, nil).AnyTimes()
 			provider.azureManager.azClient.virtualMachineScaleSetVMsClient = mockVMSSVMClient
 		} else {
 
 			provider.azureManager.config.EnableVmssFlex = true
 			mockVMClient := mockvmclient.NewMockInterface(ctrl)
-			mockVMClient.EXPECT().ListVmssFlexVMsWithoutInstanceView(gomock.Any(), "test-asg").Return(expectedVMs, nil).AnyTimes()
+			mockVMClient.EXPECT().ListVmssFlexVMsWithoutInstanceView(gomock.Any(), testASG).Return(expectedVMs, nil).AnyTimes()
 			provider.azureManager.azClient.virtualMachinesClient = mockVMClient
 		}
 		err := provider.azureManager.forceRefresh()
@@ -206,23 +257,63 @@ func TestIncreaseSize(t *testing.T) {
 		assert.Equal(t, expectedErr, err)
 
 		registered := provider.azureManager.RegisterNodeGroup(
-			newTestScaleSet(provider.azureManager, "test-asg"))
+			newTestScaleSet(provider.azureManager, testASG))
 		assert.True(t, registered)
 		assert.Equal(t, len(provider.NodeGroups()), 1)
 
-		// current target size is 2.
+		// Current target size is 3.
 		targetSize, err := provider.NodeGroups()[0].TargetSize()
 		assert.NoError(t, err)
 		assert.Equal(t, 3, targetSize)
 
-		// increase 3 nodes.
+		// Increase 2 nodes.
 		err = provider.NodeGroups()[0].IncreaseSize(2)
 		assert.NoError(t, err)
 
-		// new target size should be 5.
+		// New target size should be 5.
 		targetSize, err = provider.NodeGroups()[0].TargetSize()
 		assert.NoError(t, err)
 		assert.Equal(t, 5, targetSize)
+
+		// Testing Edge Zone scenario. Scale from 3 to 5.
+		registeredForEdgeZone := provider.azureManager.RegisterNodeGroup(
+			newTestScaleSet(provider.azureManager, "edgezone-vmss"))
+		assert.True(t, registeredForEdgeZone)
+		assert.Equal(t, len(provider.NodeGroups()), 2)
+
+		targetSizeForEdgeZone, err := provider.NodeGroups()[1].TargetSize()
+		assert.NoError(t, err)
+		assert.Equal(t, 3, targetSizeForEdgeZone)
+
+		mockVMSSClient.EXPECT().CreateOrUpdateAsync(gomock.Any(), provider.azureManager.config.ResourceGroup,
+			"edgezone-vmss", gomock.Any()).Return(nil, nil)
+		err = provider.NodeGroups()[1].IncreaseSize(2)
+		assert.NoError(t, err)
+
+		targetSizeForEdgeZone, err = provider.NodeGroups()[1].TargetSize()
+		assert.NoError(t, err)
+		assert.Equal(t, 5, targetSizeForEdgeZone)
+
+		// Testing Edge Zone scenario scaleFromZero case. Scale from 0 to 2.
+		registeredForEdgeZoneMinZero := provider.azureManager.RegisterNodeGroup(
+			newTestScaleSetMinSizeZero(provider.azureManager, "edgezone-minzero-vmss"))
+		assert.True(t, registeredForEdgeZoneMinZero)
+		assert.Equal(t, len(provider.NodeGroups()), 3)
+
+		// Current target size is 0.
+		targetSizeForEdgeZoneMinZero, err := provider.NodeGroups()[2].TargetSize()
+		assert.NoError(t, err)
+		assert.Equal(t, 0, targetSizeForEdgeZoneMinZero)
+
+		mockVMSSClient.EXPECT().CreateOrUpdateAsync(gomock.Any(), provider.azureManager.config.ResourceGroup,
+			"edgezone-minzero-vmss", gomock.Any()).Return(nil, nil)
+		err = provider.NodeGroups()[2].IncreaseSize(2)
+		assert.NoError(t, err)
+
+		// New target size should be 2.
+		targetSizeForEdgeZoneMinZero, err = provider.NodeGroups()[2].TargetSize()
+		assert.NoError(t, err)
+		assert.Equal(t, 2, targetSizeForEdgeZoneMinZero)
 	}
 }