Repository: kubernetes-sigs/cloud-provider-azure
Pull request #2416 (2022-09-28 07:03:27)

[release-1.1] chore: remove node from VMSS VM cache when it is deleted
(cherry picked from commit 8b53dd895e555fa0fbbe1ab8854853e086871b02)

<!--  Thanks for sending a pull request!  Here are some tips for you:

1. If this is your first time, please read our contributor guidelines: https://git.k8s.io/community/contributors/guide/first-contribution.md#your-first-contribution and developer guide https://git.k8s.io/community/contributors/devel/development.md#development-guide
2. Please label this pull request according to what type of issue you are addressing, especially if this is a release targeted pull request. For reference on required PR/issue labels, read here:
https://git.k8s.io/community/contributors/devel/sig-release/release.md#issuepr-kind-label
3. Ensure you have added or ran the appropriate tests for your PR: https://git.k8s.io/community/contributors/devel/sig-testing/testing.md
4. If you want *faster* PR reviews, read how: https://git.k8s.io/community/contributors/guide/pull-requests.md#best-practices-for-faster-reviews
5. If the PR is unfinished, see how to mark it: https://git.k8s.io/community/contributors/guide/pull-requests.md#marking-unfinished-pull-requests
-->

#### What type of PR is this?

<!--
Add one of the following kinds:
/kind bug
/kind cleanup
/kind documentation
/kind feature
/kind design

Optionally add one or more of the following kinds if applicable:
/kind api-change
/kind deprecation
/kind failing-test
/kind flake
/kind regression
-->

/kind cleanup

#### What this PR does / why we need it:

Cherry pick of #2402: Remove the node from the VMSS VM cache if it is deleted. This will help reduce the chance of conflict between aks and cloud provider azure.

#### Which issue(s) this PR fixes:
<!--
*Automatically closes linked issue when PR is merged.
Usage: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
_If PR is about `failing-tests or flakes`, please post the related issues/tests in a comment and do not use `Fixes`_*
-->
Fixes #

#### Special notes for your reviewer:

#### Does this PR introduce a user-facing change?
<!--
If no, just write "NONE" in the release-note block below.
If yes, a release note is required:
Enter your extended release note in the block below. If the PR requires additional action from users switching to the new release, include the string "action required".

For more information on release notes see: https://git.k8s.io/community/contributors/guide/release-notes.md
-->
```release-note
NONE
```

#### Additional documentation e.g., KEPs (Kubernetes Enhancement Proposals), usage docs, etc.:

<!--
This section can be blank if this pull request does not require a release note.

When adding links which point to resources within git repositories, like
KEPs or supporting documentation, please reference a specific commit and avoid
linking directly to the master branch. This ensures that links reference a
specific point in time, rather than a document that may change over time.

See here for guidance on getting permanent links to files: https://help.github.com/en/articles/getting-permanent-links-to-files

Please use the following format for linking documentation:
- [KEP]: <link>
- [Usage]: <link>
- [Other doc]: <link>
-->
```docs

```



Patch:
diff --git a/pkg/provider/azure.go b/pkg/provider/azure.go
index 8951f4633..ef7e25063 100644
--- a/pkg/provider/azure.go
+++ b/pkg/provider/azure.go
@@ -1000,6 +1000,9 @@ func (az *Cloud) SetInformers(informerFactory informers.SharedInformerFactory) {
 				}
 			}
 			az.updateNodeCaches(node, nil)
+
+			klog.V(4).Infof("Removing node %s from VMSet cache.", node.Name)
+			_ = az.VMSet.DeleteCacheForNode(node.Name)
 		},
 	})
 	az.nodeInformerSynced = nodeInformer.HasSynced
diff --git a/pkg/provider/azure_controller_standard.go b/pkg/provider/azure_controller_standard.go
index c85b96079..8a403c550 100644
--- a/pkg/provider/azure_controller_standard.go
+++ b/pkg/provider/azure_controller_standard.go
@@ -113,7 +113,7 @@ func (as *availabilitySet) AttachDisk(nodeName types.NodeName, diskMap map[strin
 
 	// Invalidate the cache right after updating
 	defer func() {
-		_ = as.cloud.vmCache.Delete(vmName)
+		_ = as.DeleteCacheForNode(vmName)
 	}()
 
 	future, rerr := as.VirtualMachinesClient.UpdateAsync(ctx, nodeResourceGroup, vmName, newVM, "attach_disk")
@@ -134,6 +134,11 @@ func (as *availabilitySet) AttachDisk(nodeName types.NodeName, diskMap map[strin
 	return future, nil
 }
 
+func (as *availabilitySet) DeleteCacheForNode(nodeName string) error {
+	_ = as.cloud.vmCache.Delete(nodeName)
+	return nil
+}
+
 // WaitForUpdateResult waits for the response of the update request
 func (as *availabilitySet) WaitForUpdateResult(ctx context.Context, future *azure.Future, resourceGroupName, source string) error {
 	if rerr := as.VirtualMachinesClient.WaitForUpdateResult(ctx, future, resourceGroupName, source); rerr != nil {
@@ -203,7 +208,7 @@ func (as *availabilitySet) DetachDisk(nodeName types.NodeName, diskMap map[strin
 
 	// Invalidate the cache right after updating
 	defer func() {
-		_ = as.cloud.vmCache.Delete(vmName)
+		_ = as.DeleteCacheForNode(vmName)
 	}()
 
 	future, rerr := as.VirtualMachinesClient.UpdateAsync(ctx, nodeResourceGroup, vmName, newVM, "detach_disk")
@@ -237,7 +242,7 @@ func (as *availabilitySet) UpdateVM(nodeName types.NodeName) error {
 
 	// Invalidate the cache right after updating
 	defer func() {
-		_ = as.cloud.vmCache.Delete(vmName)
+		_ = as.DeleteCacheForNode(vmName)
 	}()
 
 	rerr := as.VirtualMachinesClient.Update(ctx, nodeResourceGroup, vmName, compute.VirtualMachineUpdate{}, "update_vm")
diff --git a/pkg/provider/azure_controller_vmss.go b/pkg/provider/azure_controller_vmss.go
index 74085c6f9..4cea14a04 100644
--- a/pkg/provider/azure_controller_vmss.go
+++ b/pkg/provider/azure_controller_vmss.go
@@ -116,7 +116,7 @@ func (ss *ScaleSet) AttachDisk(nodeName types.NodeName, diskMap map[string]*Atta
 
 	// Invalidate the cache right after updating
 	defer func() {
-		_ = ss.deleteCacheForNode(vmName)
+		_ = ss.DeleteCacheForNode(vmName)
 	}()
 
 	klog.V(2).Infof("azureDisk - update(%s): vm(%s) - attach disk list(%s)", nodeResourceGroup, nodeName, diskMap)
@@ -207,7 +207,7 @@ func (ss *ScaleSet) DetachDisk(nodeName types.NodeName, diskMap map[string]strin
 
 	// Invalidate the cache right after updating
 	defer func() {
-		_ = ss.deleteCacheForNode(vmName)
+		_ = ss.DeleteCacheForNode(vmName)
 	}()
 
 	klog.V(2).Infof("azureDisk - update(%s): vm(%s) - detach disk list(%s)", nodeResourceGroup, nodeName, diskMap)
@@ -247,7 +247,7 @@ func (ss *ScaleSet) UpdateVM(nodeName types.NodeName) error {
 
 	// Invalidate the cache right after updating
 	defer func() {
-		_ = ss.deleteCacheForNode(vmName)
+		_ = ss.DeleteCacheForNode(vmName)
 	}()
 
 	klog.V(2).Infof("azureDisk - update(%s): vm(%s)", nodeResourceGroup, nodeName)
diff --git a/pkg/provider/azure_vmsets.go b/pkg/provider/azure_vmsets.go
index 5b98dc3ab..a70588767 100644
--- a/pkg/provider/azure_vmsets.go
+++ b/pkg/provider/azure_vmsets.go
@@ -100,4 +100,7 @@ type VMSet interface {
 
 	// GetAgentPoolVMSetNames returns all vmSet names according to the nodes
 	GetAgentPoolVMSetNames(nodes []*v1.Node) (*[]string, error)
+
+	// DeleteCacheForNode removes the node entry from cache.
+	DeleteCacheForNode(nodeName string) error
 }
diff --git a/pkg/provider/azure_vmss.go b/pkg/provider/azure_vmss.go
index 07969e41e..fb65054e8 100644
--- a/pkg/provider/azure_vmss.go
+++ b/pkg/provider/azure_vmss.go
@@ -461,7 +461,7 @@ func (ss *ScaleSet) GetZoneByNodeName(name string) (cloudprovider.Zone, error) {
 	} else {
 		err = fmt.Errorf("failed to get zone info")
 		klog.Errorf("GetZoneByNodeName: got unexpected error %v", err)
-		_ = ss.deleteCacheForNode(name)
+		_ = ss.DeleteCacheForNode(name)
 		return cloudprovider.Zone{}, err
 	}
 
@@ -1307,7 +1307,7 @@ func (ss *ScaleSet) EnsureHostsInPool(service *v1.Service, nodes []*v1.Node, bac
 
 		// Invalidate the cache since the VMSS VM would be updated.
 		defer func() {
-			_ = ss.deleteCacheForNode(localNodeName)
+			_ = ss.DeleteCacheForNode(localNodeName)
 		}()
 	}
 
@@ -1592,7 +1592,7 @@ func (ss *ScaleSet) EnsureBackendPoolDeleted(service *v1.Service, backendPoolID,
 
 		// Invalidate the cache since the VMSS VM would be updated.
 		defer func() {
-			_ = ss.deleteCacheForNode(nodeName)
+			_ = ss.DeleteCacheForNode(nodeName)
 		}()
 	}
 
@@ -1781,7 +1781,7 @@ func (ss *ScaleSet) GetAgentPoolVMSetNames(nodes []*v1.Node) (*[]string, error)
 			if err != nil {
 				return nil, fmt.Errorf("GetAgentPoolVMSetNames: failed to get availabilitySetNodesCache")
 			}
-			vms := cached.(availabilitySetNodeEntry).vms
+			vms := cached.(*availabilitySetNodeEntry).vms
 			names, err = as.getAgentPoolAvailabilitySets(vms, []*v1.Node{node})
 			if err != nil {
 				return nil, fmt.Errorf("GetAgentPoolVMSetNames: failed to execute getAgentPoolAvailabilitySets: %w", err)
diff --git a/pkg/provider/azure_vmss_cache.go b/pkg/provider/azure_vmss_cache.go
index 8cb141d61..50c94cf32 100644
--- a/pkg/provider/azure_vmss_cache.go
+++ b/pkg/provider/azure_vmss_cache.go
@@ -266,29 +266,59 @@ func (ss *ScaleSet) newVMSSVirtualMachinesCache(resourceGroupName, vmssName, cac
 	return azcache.NewTimedcache(vmssVirtualMachinesCacheTTL, getter)
 }
 
-func (ss *ScaleSet) deleteCacheForNode(nodeName string) error {
+func (ss *ScaleSet) DeleteCacheForNode(nodeName string) error {
+	managedByAS, err := ss.isNodeManagedByAvailabilitySet(nodeName, azcache.CacheReadTypeUnsafe)
+	if err != nil {
+		klog.Warningf("DeleteCacheForNode(%s): failed to check if the node is managed by AvailabilitySet", nodeName)
+		return nil
+	}
+	if managedByAS {
+		cached, err := ss.availabilitySetNodesCache.Get(consts.AvailabilitySetNodesKey, azcache.CacheReadTypeDefault)
+		if err != nil {
+			klog.Warningf("DeleteCacheForNode(%s): failed to get availabilitySetNodesCache: %s", err.Error())
+			return nil
+		}
+
+		ss.lockMap.LockEntry(consts.AvailabilitySetNodesKey)
+		defer ss.lockMap.UnlockEntry(consts.AvailabilitySetNodesKey)
+
+		entry := cached.(*availabilitySetNodeEntry)
+		entry.nodeNames.Delete(nodeName)
+		entry.vmNames.Delete(nodeName)
+		for i := len(entry.vms) - 1; i >= 0; i-- {
+			if strings.EqualFold(to.String(entry.vms[i].Name), nodeName) {
+				entry.vms = append(entry.vms[:i], entry.vms[i+1:]...)
+				break
+			}
+		}
+
+		_ = ss.vmCache.Delete(nodeName)
+
+		return nil
+	}
+
 	node, err := ss.getNodeIdentityByNodeName(nodeName, azcache.CacheReadTypeUnsafe)
 	if err != nil {
-		klog.Errorf("deleteCacheForNode(%s) failed with error: %v", nodeName, err)
+		klog.Errorf("DeleteCacheForNode(%s) failed with error: %v", nodeName, err)
 		return err
 	}
 
 	cacheKey, timedcache, err := ss.getVMSSVMCache(node.resourceGroup, node.vmssName)
 	if err != nil {
-		klog.Errorf("deleteCacheForNode(%s) failed with error: %v", nodeName, err)
+		klog.Errorf("DeleteCacheForNode(%s) failed with error: %v", nodeName, err)
 		return err
 	}
 
 	vmcache, err := timedcache.Get(cacheKey, azcache.CacheReadTypeUnsafe)
 	if err != nil {
-		klog.Errorf("deleteCacheForNode(%s) failed with error: %v", nodeName, err)
+		klog.Errorf("DeleteCacheForNode(%s) failed with error: %v", nodeName, err)
 		return err
 	}
 	virtualMachines := vmcache.(*sync.Map)
 	virtualMachines.Delete(nodeName)
 
 	if err := ss.gcVMSSVMCache(); err != nil {
-		klog.Errorf("deleteCacheForNode(%s) failed to gc stale vmss caches: %v", nodeName, err)
+		klog.Errorf("DeleteCacheForNode(%s) failed to gc stale vmss caches: %v", nodeName, err)
 	}
 
 	return nil
@@ -322,7 +352,7 @@ func (ss *ScaleSet) newAvailabilitySetNodesCache() (*azcache.TimedCache, error)
 			return nil, err
 		}
 
-		localCache := availabilitySetNodeEntry{
+		localCache := &availabilitySetNodeEntry{
 			vmNames:   vmNames,
 			nodeNames: nodeNames,
 			vms:       vmList,
@@ -349,7 +379,7 @@ func (ss *ScaleSet) isNodeManagedByAvailabilitySet(nodeName string, crt azcache.
 		return false, err
 	}
 
-	cachedNodes := cached.(availabilitySetNodeEntry).nodeNames
+	cachedNodes := cached.(*availabilitySetNodeEntry).nodeNames
 	// if the node is not in the cache, assume the node has joined after the last cache refresh and attempt to refresh the cache.
 	if !cachedNodes.Has(nodeName) {
 		klog.V(2).Infof("Node %s has joined the cluster since the last VM cache refresh, refreshing the cache", nodeName)
@@ -359,6 +389,6 @@ func (ss *ScaleSet) isNodeManagedByAvailabilitySet(nodeName string, crt azcache.
 		}
 	}
 
-	cachedVMs := cached.(availabilitySetNodeEntry).vmNames
+	cachedVMs := cached.(*availabilitySetNodeEntry).vmNames
 	return cachedVMs.Has(nodeName), nil
 }


Test patch:
diff --git a/pkg/provider/azure_controller_vmss_test.go b/pkg/provider/azure_controller_vmss_test.go
index 170b65916..0d76ae543 100644
--- a/pkg/provider/azure_controller_vmss_test.go
+++ b/pkg/provider/azure_controller_vmss_test.go
@@ -29,6 +29,7 @@ import (
 	"k8s.io/apimachinery/pkg/types"
 	cloudprovider "k8s.io/cloud-provider"
 
+	"sigs.k8s.io/cloud-provider-azure/pkg/azureclients/vmclient/mockvmclient"
 	"sigs.k8s.io/cloud-provider-azure/pkg/azureclients/vmssclient/mockvmssclient"
 	"sigs.k8s.io/cloud-provider-azure/pkg/azureclients/vmssvmclient/mockvmssvmclient"
 	azcache "sigs.k8s.io/cloud-provider-azure/pkg/cache"
@@ -111,6 +112,8 @@ func TestAttachDiskWithVMSS(t *testing.T) {
 		mockVMSSClient.EXPECT().List(gomock.Any(), testCloud.ResourceGroup).Return([]compute.VirtualMachineScaleSet{expectedVMSS}, nil).AnyTimes()
 		mockVMSSClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, scaleSetName).Return(expectedVMSS, nil).MaxTimes(1)
 		mockVMSSClient.EXPECT().CreateOrUpdate(gomock.Any(), testCloud.ResourceGroup, scaleSetName, gomock.Any()).Return(nil).MaxTimes(1)
+		mockVMClient := testCloud.VirtualMachinesClient.(*mockvmclient.MockInterface)
+		mockVMClient.EXPECT().List(gomock.Any(), ss.ResourceGroup).Return([]compute.VirtualMachine{}, nil).AnyTimes()
 
 		expectedVMSSVMs, _, _ := buildTestVirtualMachineEnv(testCloud, scaleSetName, "", 0, test.vmssVMList, "succeeded", false)
 		for _, vmssvm := range expectedVMSSVMs {
@@ -266,6 +269,9 @@ func TestDetachDiskWithVMSS(t *testing.T) {
 			mockVMSSVMClient.EXPECT().UpdateAsync(gomock.Any(), testCloud.ResourceGroup, scaleSetName, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes()
 		}
 
+		mockVMClient := testCloud.VirtualMachinesClient.(*mockvmclient.MockInterface)
+		mockVMClient.EXPECT().List(gomock.Any(), ss.ResourceGroup).Return([]compute.VirtualMachine{}, nil).AnyTimes()
+
 		diskMap := map[string]string{}
 		for _, diskName := range test.disks {
 			diskURI := fmt.Sprintf("/subscriptions/%s/resourceGroups/%s/providers/Microsoft.Compute/disks/%s",
@@ -375,6 +381,9 @@ func TestUpdateVMWithVMSS(t *testing.T) {
 			mockVMSSVMClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, scaleSetName, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 		}
 
+		mockVMClient := testCloud.VirtualMachinesClient.(*mockvmclient.MockInterface)
+		mockVMClient.EXPECT().List(gomock.Any(), ss.ResourceGroup).Return([]compute.VirtualMachine{}, nil).AnyTimes()
+
 		err = ss.UpdateVM(test.vmssvmName)
 		assert.Equal(t, test.expectedErr, err != nil, "TestCase[%d]: %s, err: %v", i, test.desc, err)
 		if test.expectedErr {
diff --git a/pkg/provider/azure_mock_vmsets_test.go b/pkg/provider/azure_mock_vmsets_test.go
index e5bfb3584..89a04f3ee 100644
--- a/pkg/provider/azure_mock_vmsets_test.go
+++ b/pkg/provider/azure_mock_vmsets_test.go
@@ -53,124 +53,94 @@ func (m *MockVMSet) EXPECT() *MockVMSetMockRecorder {
 	return m.recorder
 }
 
-// GetInstanceIDByNodeName mocks base method.
-func (m *MockVMSet) GetInstanceIDByNodeName(name string) (string, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetInstanceIDByNodeName", name)
-	ret0, _ := ret[0].(string)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// GetInstanceIDByNodeName indicates an expected call of GetInstanceIDByNodeName.
-func (mr *MockVMSetMockRecorder) GetInstanceIDByNodeName(name interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetInstanceIDByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetInstanceIDByNodeName), name)
-}
-
-// GetInstanceTypeByNodeName mocks base method.
-func (m *MockVMSet) GetInstanceTypeByNodeName(name string) (string, error) {
+// AttachDisk mocks base method.
+func (m *MockVMSet) AttachDisk(nodeName types.NodeName, diskMap map[string]*AttachDiskOptions) (*azure.Future, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetInstanceTypeByNodeName", name)
-	ret0, _ := ret[0].(string)
+	ret := m.ctrl.Call(m, "AttachDisk", nodeName, diskMap)
+	ret0, _ := ret[0].(*azure.Future)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-// GetInstanceTypeByNodeName indicates an expected call of GetInstanceTypeByNodeName.
-func (mr *MockVMSetMockRecorder) GetInstanceTypeByNodeName(name interface{}) *gomock.Call {
+// AttachDisk indicates an expected call of AttachDisk.
+func (mr *MockVMSetMockRecorder) AttachDisk(nodeName, diskMap interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetInstanceTypeByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetInstanceTypeByNodeName), name)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AttachDisk", reflect.TypeOf((*MockVMSet)(nil).AttachDisk), nodeName, diskMap)
 }
 
-// GetIPByNodeName mocks base method.
-func (m *MockVMSet) GetIPByNodeName(name string) (string, string, error) {
+// DeleteCacheForNode mocks base method.
+func (m *MockVMSet) DeleteCacheForNode(nodeName string) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetIPByNodeName", name)
-	ret0, _ := ret[0].(string)
-	ret1, _ := ret[1].(string)
-	ret2, _ := ret[2].(error)
-	return ret0, ret1, ret2
-}
-
-// GetIPByNodeName indicates an expected call of GetIPByNodeName.
-func (mr *MockVMSetMockRecorder) GetIPByNodeName(name interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetIPByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetIPByNodeName), name)
-}
-
-// GetPrimaryInterface mocks base method.
-func (m *MockVMSet) GetPrimaryInterface(nodeName string) (network.Interface, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetPrimaryInterface", nodeName)
-	ret0, _ := ret[0].(network.Interface)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
+	ret := m.ctrl.Call(m, "DeleteCacheForNode", nodeName)
+	ret0, _ := ret[0].(error)
+	return ret0
 }
 
-// GetPrimaryInterface indicates an expected call of GetPrimaryInterface.
-func (mr *MockVMSetMockRecorder) GetPrimaryInterface(nodeName interface{}) *gomock.Call {
+// DeleteCacheForNode indicates an expected call of DeleteCacheForNode.
+func (mr *MockVMSetMockRecorder) DeleteCacheForNode(nodeName interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPrimaryInterface", reflect.TypeOf((*MockVMSet)(nil).GetPrimaryInterface), nodeName)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteCacheForNode", reflect.TypeOf((*MockVMSet)(nil).DeleteCacheForNode), nodeName)
 }
 
-// GetNodeNameByProviderID mocks base method.
-func (m *MockVMSet) GetNodeNameByProviderID(providerID string) (types.NodeName, error) {
+// DetachDisk mocks base method.
+func (m *MockVMSet) DetachDisk(nodeName types.NodeName, diskMap map[string]string) (*azure.Future, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetNodeNameByProviderID", providerID)
-	ret0, _ := ret[0].(types.NodeName)
+	ret := m.ctrl.Call(m, "DetachDisk", nodeName, diskMap)
+	ret0, _ := ret[0].(*azure.Future)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-// GetNodeNameByProviderID indicates an expected call of GetNodeNameByProviderID.
-func (mr *MockVMSetMockRecorder) GetNodeNameByProviderID(providerID interface{}) *gomock.Call {
+// DetachDisk indicates an expected call of DetachDisk.
+func (mr *MockVMSetMockRecorder) DetachDisk(nodeName, diskMap interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetNodeNameByProviderID", reflect.TypeOf((*MockVMSet)(nil).GetNodeNameByProviderID), providerID)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DetachDisk", reflect.TypeOf((*MockVMSet)(nil).DetachDisk), nodeName, diskMap)
 }
 
-// GetZoneByNodeName mocks base method.
-func (m *MockVMSet) GetZoneByNodeName(name string) (cloud_provider.Zone, error) {
+// EnsureBackendPoolDeleted mocks base method.
+func (m *MockVMSet) EnsureBackendPoolDeleted(service *v1.Service, backendPoolID, vmSetName string, backendAddressPools *[]network.BackendAddressPool, deleteFromVMSet bool) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetZoneByNodeName", name)
-	ret0, _ := ret[0].(cloud_provider.Zone)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
+	ret := m.ctrl.Call(m, "EnsureBackendPoolDeleted", service, backendPoolID, vmSetName, backendAddressPools, deleteFromVMSet)
+	ret0, _ := ret[0].(error)
+	return ret0
 }
 
-// GetZoneByNodeName indicates an expected call of GetZoneByNodeName.
-func (mr *MockVMSetMockRecorder) GetZoneByNodeName(name interface{}) *gomock.Call {
+// EnsureBackendPoolDeleted indicates an expected call of EnsureBackendPoolDeleted.
+func (mr *MockVMSetMockRecorder) EnsureBackendPoolDeleted(service, backendPoolID, vmSetName, backendAddressPools, deleteFromVMSet interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetZoneByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetZoneByNodeName), name)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EnsureBackendPoolDeleted", reflect.TypeOf((*MockVMSet)(nil).EnsureBackendPoolDeleted), service, backendPoolID, vmSetName, backendAddressPools, deleteFromVMSet)
 }
 
-// GetPrimaryVMSetName mocks base method.
-func (m *MockVMSet) GetPrimaryVMSetName() string {
+// EnsureBackendPoolDeletedFromVMSets mocks base method.
+func (m *MockVMSet) EnsureBackendPoolDeletedFromVMSets(vmSetNamesMap map[string]bool, backendPoolID string) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetPrimaryVMSetName")
-	ret0, _ := ret[0].(string)
+	ret := m.ctrl.Call(m, "EnsureBackendPoolDeletedFromVMSets", vmSetNamesMap, backendPoolID)
+	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-// GetPrimaryVMSetName indicates an expected call of GetPrimaryVMSetName.
-func (mr *MockVMSetMockRecorder) GetPrimaryVMSetName() *gomock.Call {
+// EnsureBackendPoolDeletedFromVMSets indicates an expected call of EnsureBackendPoolDeletedFromVMSets.
+func (mr *MockVMSetMockRecorder) EnsureBackendPoolDeletedFromVMSets(vmSetNamesMap, backendPoolID interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPrimaryVMSetName", reflect.TypeOf((*MockVMSet)(nil).GetPrimaryVMSetName))
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EnsureBackendPoolDeletedFromVMSets", reflect.TypeOf((*MockVMSet)(nil).EnsureBackendPoolDeletedFromVMSets), vmSetNamesMap, backendPoolID)
 }
 
-// GetVMSetNames mocks base method.
-func (m *MockVMSet) GetVMSetNames(service *v1.Service, nodes []*v1.Node) (*[]string, error) {
+// EnsureHostInPool mocks base method.
+func (m *MockVMSet) EnsureHostInPool(service *v1.Service, nodeName types.NodeName, backendPoolID, vmSetName string, isInternal bool) (string, string, string, *compute.VirtualMachineScaleSetVM, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetVMSetNames", service, nodes)
-	ret0, _ := ret[0].(*[]string)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
+	ret := m.ctrl.Call(m, "EnsureHostInPool", service, nodeName, backendPoolID, vmSetName, isInternal)
+	ret0, _ := ret[0].(string)
+	ret1, _ := ret[1].(string)
+	ret2, _ := ret[2].(string)
+	ret3, _ := ret[3].(*compute.VirtualMachineScaleSetVM)
+	ret4, _ := ret[4].(error)
+	return ret0, ret1, ret2, ret3, ret4
 }
 
-// GetVMSetNames indicates an expected call of GetVMSetNames.
-func (mr *MockVMSetMockRecorder) GetVMSetNames(service, nodes interface{}) *gomock.Call {
+// EnsureHostInPool indicates an expected call of EnsureHostInPool.
+func (mr *MockVMSetMockRecorder) EnsureHostInPool(service, nodeName, backendPoolID, vmSetName, isInternal interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetVMSetNames", reflect.TypeOf((*MockVMSet)(nil).GetVMSetNames), service, nodes)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EnsureHostInPool", reflect.TypeOf((*MockVMSet)(nil).EnsureHostInPool), service, nodeName, backendPoolID, vmSetName, isInternal)
 }
 
 // EnsureHostsInPool mocks base method.
@@ -187,124 +157,128 @@ func (mr *MockVMSetMockRecorder) EnsureHostsInPool(service, nodes, backendPoolID
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EnsureHostsInPool", reflect.TypeOf((*MockVMSet)(nil).EnsureHostsInPool), service, nodes, backendPoolID, vmSetName, isInternal)
 }
 
-// EnsureHostInPool mocks base method.
-func (m *MockVMSet) EnsureHostInPool(service *v1.Service, nodeName types.NodeName, backendPoolID, vmSetName string, isInternal bool) (string, string, string, *compute.VirtualMachineScaleSetVM, error) {
+// GetAgentPoolVMSetNames mocks base method.
+func (m *MockVMSet) GetAgentPoolVMSetNames(nodes []*v1.Node) (*[]string, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "EnsureHostInPool", service, nodeName, backendPoolID, vmSetName, isInternal)
-	ret0, _ := ret[0].(string)
-	ret1, _ := ret[1].(string)
-	ret2, _ := ret[2].(string)
-	ret3, _ := ret[3].(*compute.VirtualMachineScaleSetVM)
-	ret4, _ := ret[4].(error)
-	return ret0, ret1, ret2, ret3, ret4
+	ret := m.ctrl.Call(m, "GetAgentPoolVMSetNames", nodes)
+	ret0, _ := ret[0].(*[]string)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
 }
 
-// EnsureHostInPool indicates an expected call of EnsureHostInPool.
-func (mr *MockVMSetMockRecorder) EnsureHostInPool(service, nodeName, backendPoolID, vmSetName, isInternal interface{}) *gomock.Call {
+// GetAgentPoolVMSetNames indicates an expected call of GetAgentPoolVMSetNames.
+func (mr *MockVMSetMockRecorder) GetAgentPoolVMSetNames(nodes interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EnsureHostInPool", reflect.TypeOf((*MockVMSet)(nil).EnsureHostInPool), service, nodeName, backendPoolID, vmSetName, isInternal)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentPoolVMSetNames", reflect.TypeOf((*MockVMSet)(nil).GetAgentPoolVMSetNames), nodes)
 }
 
-// EnsureBackendPoolDeleted mocks base method.
-func (m *MockVMSet) EnsureBackendPoolDeleted(service *v1.Service, backendPoolID, vmSetName string, backendAddressPools *[]network.BackendAddressPool, deleteFromVMSet bool) error {
+// GetDataDisks mocks base method.
+func (m *MockVMSet) GetDataDisks(nodeName types.NodeName, crt cache.AzureCacheReadType) ([]compute.DataDisk, *string, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "EnsureBackendPoolDeleted", service, backendPoolID, vmSetName, backendAddressPools, deleteFromVMSet)
-	ret0, _ := ret[0].(error)
-	return ret0
+	ret := m.ctrl.Call(m, "GetDataDisks", nodeName, crt)
+	ret0, _ := ret[0].([]compute.DataDisk)
+	ret1, _ := ret[1].(*string)
+	ret2, _ := ret[2].(error)
+	return ret0, ret1, ret2
 }
 
-// EnsureBackendPoolDeleted indicates an expected call of EnsureBackendPoolDeleted.
-func (mr *MockVMSetMockRecorder) EnsureBackendPoolDeleted(service, backendPoolID, vmSetName, backendAddressPools, deleteFromVMSet interface{}) *gomock.Call {
+// GetDataDisks indicates an expected call of GetDataDisks.
+func (mr *MockVMSetMockRecorder) GetDataDisks(nodeName, crt interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EnsureBackendPoolDeleted", reflect.TypeOf((*MockVMSet)(nil).EnsureBackendPoolDeleted), service, backendPoolID, vmSetName, backendAddressPools, deleteFromVMSet)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDataDisks", reflect.TypeOf((*MockVMSet)(nil).GetDataDisks), nodeName, crt)
 }
 
-// EnsureBackendPoolDeletedFromVMSets mocks base method.
-func (m *MockVMSet) EnsureBackendPoolDeletedFromVMSets(vmSetNamesMap map[string]bool, backendPoolID string) error {
+// GetIPByNodeName mocks base method.
+func (m *MockVMSet) GetIPByNodeName(name string) (string, string, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "EnsureBackendPoolDeletedFromVMSets", vmSetNamesMap, backendPoolID)
-	ret0, _ := ret[0].(error)
-	return ret0
+	ret := m.ctrl.Call(m, "GetIPByNodeName", name)
+	ret0, _ := ret[0].(string)
+	ret1, _ := ret[1].(string)
+	ret2, _ := ret[2].(error)
+	return ret0, ret1, ret2
 }
 
-// EnsureBackendPoolDeletedFromVMSets indicates an expected call of EnsureBackendPoolDeletedFromVMSets.
-func (mr *MockVMSetMockRecorder) EnsureBackendPoolDeletedFromVMSets(vmSetNamesMap, backendPoolID interface{}) *gomock.Call {
+// GetIPByNodeName indicates an expected call of GetIPByNodeName.
+func (mr *MockVMSetMockRecorder) GetIPByNodeName(name interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EnsureBackendPoolDeletedFromVMSets", reflect.TypeOf((*MockVMSet)(nil).EnsureBackendPoolDeletedFromVMSets), vmSetNamesMap, backendPoolID)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetIPByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetIPByNodeName), name)
 }
 
-// AttachDisk mocks base method.
-func (m *MockVMSet) AttachDisk(nodeName types.NodeName, diskMap map[string]*AttachDiskOptions) (*azure.Future, error) {
+// GetInstanceIDByNodeName mocks base method.
+func (m *MockVMSet) GetInstanceIDByNodeName(name string) (string, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "AttachDisk", nodeName, diskMap)
-	ret0, _ := ret[0].(*azure.Future)
+	ret := m.ctrl.Call(m, "GetInstanceIDByNodeName", name)
+	ret0, _ := ret[0].(string)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-// AttachDisk indicates an expected call of AttachDisk.
-func (mr *MockVMSetMockRecorder) AttachDisk(nodeName, diskMap interface{}) *gomock.Call {
+// GetInstanceIDByNodeName indicates an expected call of GetInstanceIDByNodeName.
+func (mr *MockVMSetMockRecorder) GetInstanceIDByNodeName(name interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AttachDisk", reflect.TypeOf((*MockVMSet)(nil).AttachDisk), nodeName, diskMap)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetInstanceIDByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetInstanceIDByNodeName), name)
 }
 
-// DetachDisk mocks base method.
-func (m *MockVMSet) DetachDisk(nodeName types.NodeName, diskMap map[string]string) (*azure.Future, error) {
+// GetInstanceTypeByNodeName mocks base method.
+func (m *MockVMSet) GetInstanceTypeByNodeName(name string) (string, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "DetachDisk", nodeName, diskMap)
-	ret0, _ := ret[0].(*azure.Future)
+	ret := m.ctrl.Call(m, "GetInstanceTypeByNodeName", name)
+	ret0, _ := ret[0].(string)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-// DetachDisk indicates an expected call of DetachDisk.
-func (mr *MockVMSetMockRecorder) DetachDisk(nodeName, diskMap interface{}) *gomock.Call {
+// GetInstanceTypeByNodeName indicates an expected call of GetInstanceTypeByNodeName.
+func (mr *MockVMSetMockRecorder) GetInstanceTypeByNodeName(name interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DetachDisk", reflect.TypeOf((*MockVMSet)(nil).DetachDisk), nodeName, diskMap)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetInstanceTypeByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetInstanceTypeByNodeName), name)
 }
 
-// WaitForUpdateResult mocks base method.
-func (m *MockVMSet) WaitForUpdateResult(ctx context.Context, future *azure.Future, resourceGroupName, source string) error {
+// GetNodeCIDRMasksByProviderID mocks base method.
+func (m *MockVMSet) GetNodeCIDRMasksByProviderID(providerID string) (int, int, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "WaitForUpdateResult", ctx, future, resourceGroupName, source)
-	ret0, _ := ret[0].(error)
-	return ret0
+	ret := m.ctrl.Call(m, "GetNodeCIDRMasksByProviderID", providerID)
+	ret0, _ := ret[0].(int)
+	ret1, _ := ret[1].(int)
+	ret2, _ := ret[2].(error)
+	return ret0, ret1, ret2
 }
 
-// WaitForUpdateResult indicates an expected call of WaitForUpdateResult.
-func (mr *MockVMSetMockRecorder) WaitForUpdateResult(ctx, future, resourceGroupName, source interface{}) *gomock.Call {
+// GetNodeCIDRMasksByProviderID indicates an expected call of GetNodeCIDRMasksByProviderID.
+func (mr *MockVMSetMockRecorder) GetNodeCIDRMasksByProviderID(providerID interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WaitForUpdateResult", reflect.TypeOf((*MockVMSet)(nil).WaitForUpdateResult), ctx, future, resourceGroupName, source)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetNodeCIDRMasksByProviderID", reflect.TypeOf((*MockVMSet)(nil).GetNodeCIDRMasksByProviderID), providerID)
 }
 
-// GetDataDisks mocks base method.
-func (m *MockVMSet) GetDataDisks(nodeName types.NodeName, crt cache.AzureCacheReadType) ([]compute.DataDisk, *string, error) {
+// GetNodeNameByIPConfigurationID mocks base method.
+func (m *MockVMSet) GetNodeNameByIPConfigurationID(ipConfigurationID string) (string, string, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetDataDisks", nodeName, crt)
-	ret0, _ := ret[0].([]compute.DataDisk)
-	ret1, _ := ret[1].(*string)
+	ret := m.ctrl.Call(m, "GetNodeNameByIPConfigurationID", ipConfigurationID)
+	ret0, _ := ret[0].(string)
+	ret1, _ := ret[1].(string)
 	ret2, _ := ret[2].(error)
 	return ret0, ret1, ret2
 }
 
-// GetDataDisks indicates an expected call of GetDataDisks.
-func (mr *MockVMSetMockRecorder) GetDataDisks(nodeName, crt interface{}) *gomock.Call {
+// GetNodeNameByIPConfigurationID indicates an expected call of GetNodeNameByIPConfigurationID.
+func (mr *MockVMSetMockRecorder) GetNodeNameByIPConfigurationID(ipConfigurationID interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDataDisks", reflect.TypeOf((*MockVMSet)(nil).GetDataDisks), nodeName, crt)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetNodeNameByIPConfigurationID", reflect.TypeOf((*MockVMSet)(nil).GetNodeNameByIPConfigurationID), ipConfigurationID)
 }
 
-// UpdateVM mocks base method.
-func (m *MockVMSet) UpdateVM(nodeName types.NodeName) error {
+// GetNodeNameByProviderID mocks base method.
+func (m *MockVMSet) GetNodeNameByProviderID(providerID string) (types.NodeName, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "UpdateVM", nodeName)
-	ret0, _ := ret[0].(error)
-	return ret0
+	ret := m.ctrl.Call(m, "GetNodeNameByProviderID", providerID)
+	ret0, _ := ret[0].(types.NodeName)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
 }
 
-// UpdateVM indicates an expected call of UpdateVM.
-func (mr *MockVMSetMockRecorder) UpdateVM(nodeName interface{}) *gomock.Call {
+// GetNodeNameByProviderID indicates an expected call of GetNodeNameByProviderID.
+func (mr *MockVMSetMockRecorder) GetNodeNameByProviderID(providerID interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateVM", reflect.TypeOf((*MockVMSet)(nil).UpdateVM), nodeName)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetNodeNameByProviderID", reflect.TypeOf((*MockVMSet)(nil).GetNodeNameByProviderID), providerID)
 }
 
 // GetPowerStatusByNodeName mocks base method.
@@ -322,19 +296,33 @@ func (mr *MockVMSetMockRecorder) GetPowerStatusByNodeName(name interface{}) *gom
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPowerStatusByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetPowerStatusByNodeName), name)
 }
 
-// GetProvisioningStateByNodeName mocks base method.
-func (m *MockVMSet) GetProvisioningStateByNodeName(name string) (string, error) {
+// GetPrimaryInterface mocks base method.
+func (m *MockVMSet) GetPrimaryInterface(nodeName string) (network.Interface, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetProvisioningStateByNodeName", name)
-	ret0, _ := ret[0].(string)
+	ret := m.ctrl.Call(m, "GetPrimaryInterface", nodeName)
+	ret0, _ := ret[0].(network.Interface)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-// GetProvisioningStateByNodeName indicates an expected call of GetProvisioningStateByNodeName.
-func (mr *MockVMSetMockRecorder) GetProvisioningStateByNodeName(name interface{}) *gomock.Call {
+// GetPrimaryInterface indicates an expected call of GetPrimaryInterface.
+func (mr *MockVMSetMockRecorder) GetPrimaryInterface(nodeName interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetProvisioningStateByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetProvisioningStateByNodeName), name)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPrimaryInterface", reflect.TypeOf((*MockVMSet)(nil).GetPrimaryInterface), nodeName)
+}
+
+// GetPrimaryVMSetName mocks base method.
+func (m *MockVMSet) GetPrimaryVMSetName() string {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetPrimaryVMSetName")
+	ret0, _ := ret[0].(string)
+	return ret0
+}
+
+// GetPrimaryVMSetName indicates an expected call of GetPrimaryVMSetName.
+func (mr *MockVMSetMockRecorder) GetPrimaryVMSetName() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPrimaryVMSetName", reflect.TypeOf((*MockVMSet)(nil).GetPrimaryVMSetName))
 }
 
 // GetPrivateIPsByNodeName mocks base method.
@@ -352,49 +340,75 @@ func (mr *MockVMSetMockRecorder) GetPrivateIPsByNodeName(name interface{}) *gomo
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPrivateIPsByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetPrivateIPsByNodeName), name)
 }
 
-// GetNodeNameByIPConfigurationID mocks base method.
-func (m *MockVMSet) GetNodeNameByIPConfigurationID(ipConfigurationID string) (string, string, error) {
+// GetProvisioningStateByNodeName mocks base method.
+func (m *MockVMSet) GetProvisioningStateByNodeName(name string) (string, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetNodeNameByIPConfigurationID", ipConfigurationID)
+	ret := m.ctrl.Call(m, "GetProvisioningStateByNodeName", name)
 	ret0, _ := ret[0].(string)
-	ret1, _ := ret[1].(string)
-	ret2, _ := ret[2].(error)
-	return ret0, ret1, ret2
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
 }
 
-// GetNodeNameByIPConfigurationID indicates an expected call of GetNodeNameByIPConfigurationID.
-func (mr *MockVMSetMockRecorder) GetNodeNameByIPConfigurationID(ipConfigurationID interface{}) *gomock.Call {
+// GetProvisioningStateByNodeName indicates an expected call of GetProvisioningStateByNodeName.
+func (mr *MockVMSetMockRecorder) GetProvisioningStateByNodeName(name interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetNodeNameByIPConfigurationID", reflect.TypeOf((*MockVMSet)(nil).GetNodeNameByIPConfigurationID), ipConfigurationID)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetProvisioningStateByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetProvisioningStateByNodeName), name)
 }
 
-// GetNodeCIDRMasksByProviderID mocks base method.
-func (m *MockVMSet) GetNodeCIDRMasksByProviderID(providerID string) (int, int, error) {
+// GetVMSetNames mocks base method.
+func (m *MockVMSet) GetVMSetNames(service *v1.Service, nodes []*v1.Node) (*[]string, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetNodeCIDRMasksByProviderID", providerID)
-	ret0, _ := ret[0].(int)
-	ret1, _ := ret[1].(int)
-	ret2, _ := ret[2].(error)
-	return ret0, ret1, ret2
+	ret := m.ctrl.Call(m, "GetVMSetNames", service, nodes)
+	ret0, _ := ret[0].(*[]string)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
 }
 
-// GetNodeCIDRMasksByProviderID indicates an expected call of GetNodeCIDRMasksByProviderID.
-func (mr *MockVMSetMockRecorder) GetNodeCIDRMasksByProviderID(providerID interface{}) *gomock.Call {
+// GetVMSetNames indicates an expected call of GetVMSetNames.
+func (mr *MockVMSetMockRecorder) GetVMSetNames(service, nodes interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetNodeCIDRMasksByProviderID", reflect.TypeOf((*MockVMSet)(nil).GetNodeCIDRMasksByProviderID), providerID)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetVMSetNames", reflect.TypeOf((*MockVMSet)(nil).GetVMSetNames), service, nodes)
 }
 
-// GetAgentPoolVMSetNames mocks base method.
-func (m *MockVMSet) GetAgentPoolVMSetNames(nodes []*v1.Node) (*[]string, error) {
+// GetZoneByNodeName mocks base method.
+func (m *MockVMSet) GetZoneByNodeName(name string) (cloud_provider.Zone, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetAgentPoolVMSetNames", nodes)
-	ret0, _ := ret[0].(*[]string)
+	ret := m.ctrl.Call(m, "GetZoneByNodeName", name)
+	ret0, _ := ret[0].(cloud_provider.Zone)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-// GetAgentPoolVMSetNames indicates an expected call of GetAgentPoolVMSetNames.
-func (mr *MockVMSetMockRecorder) GetAgentPoolVMSetNames(nodes interface{}) *gomock.Call {
+// GetZoneByNodeName indicates an expected call of GetZoneByNodeName.
+func (mr *MockVMSetMockRecorder) GetZoneByNodeName(name interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentPoolVMSetNames", reflect.TypeOf((*MockVMSet)(nil).GetAgentPoolVMSetNames), nodes)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetZoneByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetZoneByNodeName), name)
+}
+
+// UpdateVM mocks base method.
+func (m *MockVMSet) UpdateVM(nodeName types.NodeName) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "UpdateVM", nodeName)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// UpdateVM indicates an expected call of UpdateVM.
+func (mr *MockVMSetMockRecorder) UpdateVM(nodeName interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateVM", reflect.TypeOf((*MockVMSet)(nil).UpdateVM), nodeName)
+}
+
+// WaitForUpdateResult mocks base method.
+func (m *MockVMSet) WaitForUpdateResult(ctx context.Context, future *azure.Future, resourceGroupName, source string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "WaitForUpdateResult", ctx, future, resourceGroupName, source)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// WaitForUpdateResult indicates an expected call of WaitForUpdateResult.
+func (mr *MockVMSetMockRecorder) WaitForUpdateResult(ctx, future, resourceGroupName, source interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WaitForUpdateResult", reflect.TypeOf((*MockVMSet)(nil).WaitForUpdateResult), ctx, future, resourceGroupName, source)
 }
diff --git a/pkg/provider/azure_vmss_cache_test.go b/pkg/provider/azure_vmss_cache_test.go
index 4478bfa2b..172ed6bd6 100644
--- a/pkg/provider/azure_vmss_cache_test.go
+++ b/pkg/provider/azure_vmss_cache_test.go
@@ -28,11 +28,14 @@ import (
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 
+	"k8s.io/apimachinery/pkg/util/sets"
 	cloudprovider "k8s.io/cloud-provider"
 
+	"sigs.k8s.io/cloud-provider-azure/pkg/azureclients/vmclient/mockvmclient"
 	"sigs.k8s.io/cloud-provider-azure/pkg/azureclients/vmssclient/mockvmssclient"
 	"sigs.k8s.io/cloud-provider-azure/pkg/azureclients/vmssvmclient/mockvmssvmclient"
 	azcache "sigs.k8s.io/cloud-provider-azure/pkg/cache"
+	"sigs.k8s.io/cloud-provider-azure/pkg/consts"
 	"sigs.k8s.io/cloud-provider-azure/pkg/retry"
 )
 
@@ -85,8 +88,11 @@ func TestVMSSVMCache(t *testing.T) {
 	defer ctrl.Finish()
 
 	vmList := []string{"vmssee6c2000000", "vmssee6c2000001", "vmssee6c2000002"}
-	ss, err := NewTestScaleSet(ctrl)
+	c := GetTestCloud(ctrl)
+	c.DisableAvailabilitySetNodes = true
+	vmSet, err := newScaleSet(c)
 	assert.NoError(t, err)
+	ss := vmSet.(*ScaleSet)
 
 	mockVMSSClient := mockvmssclient.NewMockInterface(ctrl)
 	mockVMSSVMClient := mockvmssvmclient.NewMockInterface(ctrl)
@@ -110,13 +116,13 @@ func TestVMSSVMCache(t *testing.T) {
 		assert.Equal(t, &vm, realVM)
 	}
 
-	// validate deleteCacheForNode().
+	// validate DeleteCacheForNode().
 	vm := expectedVMs[0]
 	vmName := to.String(vm.OsProfile.ComputerName)
-	err = ss.deleteCacheForNode(vmName)
+	err = ss.DeleteCacheForNode(vmName)
 	assert.NoError(t, err)
 
-	// the VM should be removed from cache after deleteCacheForNode().
+	// the VM should be removed from cache after DeleteCacheForNode().
 	cacheKey, cache, err := ss.getVMSSVMCache("rg", testVMSSName)
 	assert.NoError(t, err)
 	cached, err := cache.Get(cacheKey, azcache.CacheReadTypeDefault)
@@ -133,6 +139,32 @@ func TestVMSSVMCache(t *testing.T) {
 	assert.Equal(t, &vm, realVM)
 }
 
+func TestDeleteCacheForAvailabilitySetNodeInVMSS(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	ss, err := NewTestScaleSet(ctrl)
+	assert.NoError(t, err)
+
+	mockVMClient := mockvmclient.NewMockInterface(ctrl)
+	ss.cloud.VirtualMachinesClient = mockVMClient
+
+	mockVMClient.EXPECT().List(gomock.Any(), gomock.Any()).Return([]compute.VirtualMachine{
+		{Name: to.StringPtr("vm1")},
+	}, nil).AnyTimes()
+	ss.cloud.nodeNames = sets.NewString("vm1")
+
+	err = ss.DeleteCacheForNode("vm1")
+	assert.NoError(t, err)
+
+	cached, err := ss.availabilitySetNodesCache.Get(consts.AvailabilitySetNodesKey, azcache.CacheReadTypeUnsafe)
+	assert.NoError(t, err)
+	entry := cached.(*availabilitySetNodeEntry)
+	assert.Equal(t, 0, len(entry.nodeNames))
+	assert.Equal(t, 0, len(entry.vmNames))
+	assert.Equal(t, 0, len(entry.vms))
+}
+
 func TestVMSSVMCacheWithDeletingNodes(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
diff --git a/pkg/provider/azure_vmss_test.go b/pkg/provider/azure_vmss_test.go
index 0d9dbd508..d60bf59f3 100644
--- a/pkg/provider/azure_vmss_test.go
+++ b/pkg/provider/azure_vmss_test.go
@@ -2507,6 +2507,9 @@ func TestEnsureBackendPoolDeleted(t *testing.T) {
 		mockVMSSVMClient.EXPECT().List(gomock.Any(), ss.ResourceGroup, testVMSSName, gomock.Any()).Return(expectedVMSSVMs, nil).AnyTimes()
 		mockVMSSVMClient.EXPECT().UpdateVMs(gomock.Any(), ss.ResourceGroup, testVMSSName, gomock.Any(), gomock.Any(), gomock.Any()).Return(test.vmClientErr).Times(test.expectedVMSSVMPutTimes)
 
+		mockVMClient := ss.cloud.VirtualMachinesClient.(*mockvmclient.MockInterface)
+		mockVMClient.EXPECT().List(gomock.Any(), ss.ResourceGroup).Return([]compute.VirtualMachine{}, nil).AnyTimes()
+
 		err = ss.EnsureBackendPoolDeleted(&v1.Service{}, test.backendpoolID, testVMSSName, test.backendAddressPools, true)
 		assert.Equal(t, test.expectedErr, err != nil, test.description+", but an error occurs")
 	}
@@ -2567,6 +2570,9 @@ func TestEnsureBackendPoolDeletedConcurrently(t *testing.T) {
 		*lbBackendpools = append(*lbBackendpools, compute.SubResource{ID: to.StringPtr(testLBBackendpoolID1)})
 	}
 
+	mockVMClient := ss.cloud.VirtualMachinesClient.(*mockvmclient.MockInterface)
+	mockVMClient.EXPECT().List(gomock.Any(), ss.ResourceGroup).Return([]compute.VirtualMachine{}, nil).AnyTimes()
+
 	mockVMSSClient := ss.cloud.VirtualMachineScaleSetsClient.(*mockvmssclient.MockInterface)
 	mockVMSSClient.EXPECT().List(gomock.Any(), ss.ResourceGroup).Return([]compute.VirtualMachineScaleSet{vmss0, vmss1}, nil).AnyTimes()
 	mockVMSSClient.EXPECT().List(gomock.Any(), "rg1").Return(nil, nil).AnyTimes()