Repository: kubernetes-sigs/cloud-provider-azure
Pull request #7242 (2024-10-09 04:25:11)

[release-1.31] Introduce context.Context to vmset interface
This is an automated cherry-pick of #7231

/assign MartinForReal

```release-note
NONE
```

k8s-ci-robot: [APPROVALNOTIFIER] This PR is **NOT APPROVED**

This pull-request has been approved by: *<a href="https://github.com/kubernetes-sigs/cloud-provider-azure/pull/7242#" title="Author self-approved">k8s-infra-cherrypick-robot</a>*
**Once this PR has been reviewed and has the lgtm label**, please assign [lzhecheng](https://github.com/lzhecheng) for approval. For more information see [the Kubernetes Code Review Process](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process).

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=kubernetes-sigs%2Fcloud-provider-azure).

<details open>
Needs approval from an approver in each of these files:

- **[OWNERS](https://github.com/kubernetes-sigs/cloud-provider-azure/blob/release-1.31/OWNERS)**

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":["lzhecheng"]} -->
MartinForReal: /lgtm
/approve
k8s-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/kubernetes-sigs/cloud-provider-azure/pull/7242#" title="Author self-approved">k8s-infra-cherrypick-robot</a>*, *<a href="https://github.com/kubernetes-sigs/cloud-provider-azure/pull/7242#issuecomment-2401613785" title="Approved">MartinForReal</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=kubernetes-sigs%2Fcloud-provider-azure).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/kubernetes-sigs/cloud-provider-azure/blob/release-1.31/OWNERS)~~ [MartinForReal]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/pkg/provider/azure_instances_v1.go b/pkg/provider/azure_instances_v1.go
index 2e5c9ebbf..59a99ea76 100644
--- a/pkg/provider/azure_instances_v1.go
+++ b/pkg/provider/azure_instances_v1.go
@@ -44,8 +44,8 @@ var (
 	errNodeNotInitialized = fmt.Errorf("providerID is empty, the node is not initialized yet")
 )
 
-func (az *Cloud) addressGetter(nodeName types.NodeName) ([]v1.NodeAddress, error) {
-	ip, publicIP, err := az.getIPForMachine(nodeName)
+func (az *Cloud) addressGetter(ctx context.Context, nodeName types.NodeName) ([]v1.NodeAddress, error) {
+	ip, publicIP, err := az.getIPForMachine(ctx, nodeName)
 	if err != nil {
 		klog.V(2).Infof("NodeAddresses(%s) abort backoff: %v", nodeName, err)
 		return nil, err
@@ -65,7 +65,7 @@ func (az *Cloud) addressGetter(nodeName types.NodeName) ([]v1.NodeAddress, error
 }
 
 // NodeAddresses returns the addresses of the specified instance.
-func (az *Cloud) NodeAddresses(_ context.Context, name types.NodeName) ([]v1.NodeAddress, error) {
+func (az *Cloud) NodeAddresses(ctx context.Context, name types.NodeName) ([]v1.NodeAddress, error) {
 	// Returns nil for unmanaged nodes because azure cloud provider couldn't fetch information for them.
 	unmanaged, err := az.IsNodeUnmanaged(string(name))
 	if err != nil {
@@ -94,7 +94,7 @@ func (az *Cloud) NodeAddresses(_ context.Context, name types.NodeName) ([]v1.Nod
 		// Not local instance, get addresses from Azure ARM API.
 		if !isLocalInstance {
 			if az.VMSet != nil {
-				return az.addressGetter(name)
+				return az.addressGetter(ctx, name)
 			}
 
 			// vmSet == nil indicates credentials are not provided.
@@ -104,7 +104,7 @@ func (az *Cloud) NodeAddresses(_ context.Context, name types.NodeName) ([]v1.Nod
 		return az.getLocalInstanceNodeAddresses(metadata.Network.Interface, string(name))
 	}
 
-	return az.addressGetter(name)
+	return az.addressGetter(ctx, name)
 }
 
 func (az *Cloud) getLocalInstanceNodeAddresses(netInterfaces []NetworkInterface, nodeName string) ([]v1.NodeAddress, error) {
@@ -287,7 +287,7 @@ func (az *Cloud) isCurrentInstance(name types.NodeName, metadataVMName string) (
 
 // InstanceID returns the cloud provider ID of the specified instance.
 // Note that if the instance does not exist or is no longer running, we must return ("", cloudprovider.InstanceNotFound)
-func (az *Cloud) InstanceID(_ context.Context, name types.NodeName) (string, error) {
+func (az *Cloud) InstanceID(ctx context.Context, name types.NodeName) (string, error) {
 	nodeName := mapNodeNameToVMName(name)
 	unmanaged, err := az.IsNodeUnmanaged(nodeName)
 	if err != nil {
@@ -317,7 +317,7 @@ func (az *Cloud) InstanceID(_ context.Context, name types.NodeName) (string, err
 		// Not local instance, get instanceID from Azure ARM API.
 		if !isLocalInstance {
 			if az.VMSet != nil {
-				return az.VMSet.GetInstanceIDByNodeName(nodeName)
+				return az.VMSet.GetInstanceIDByNodeName(ctx, nodeName)
 			}
 
 			// vmSet == nil indicates credentials are not provided.
@@ -326,7 +326,7 @@ func (az *Cloud) InstanceID(_ context.Context, name types.NodeName) (string, err
 		return az.getLocalInstanceProviderID(metadata, nodeName)
 	}
 
-	return az.VMSet.GetInstanceIDByNodeName(nodeName)
+	return az.VMSet.GetInstanceIDByNodeName(ctx, nodeName)
 }
 
 func (az *Cloud) getLocalInstanceProviderID(metadata *InstanceMetadata, _ string) (string, error) {
@@ -377,7 +377,7 @@ func (az *Cloud) InstanceTypeByProviderID(ctx context.Context, providerID string
 // Note that if the instance does not exist or is no longer running, we must return ("", cloudprovider.InstanceNotFound)
 // (Implementer Note): This is used by kubelet. Kubelet will label the node. Real log from kubelet:
 // Adding node label from cloud provider: beta.kubernetes.io/instance-type=[value]
-func (az *Cloud) InstanceType(_ context.Context, name types.NodeName) (string, error) {
+func (az *Cloud) InstanceType(ctx context.Context, name types.NodeName) (string, error) {
 	// Returns "" for unmanaged nodes because azure cloud provider couldn't fetch information for them.
 	unmanaged, err := az.IsNodeUnmanaged(string(name))
 	if err != nil {
@@ -404,7 +404,7 @@ func (az *Cloud) InstanceType(_ context.Context, name types.NodeName) (string, e
 		}
 		if !isLocalInstance {
 			if az.VMSet != nil {
-				return az.VMSet.GetInstanceTypeByNodeName(string(name))
+				return az.VMSet.GetInstanceTypeByNodeName(ctx, string(name))
 			}
 
 			// vmSet == nil indicates credentials are not provided.
@@ -421,7 +421,7 @@ func (az *Cloud) InstanceType(_ context.Context, name types.NodeName) (string, e
 		return "", fmt.Errorf("no credentials provided for Azure cloud provider")
 	}
 
-	return az.VMSet.GetInstanceTypeByNodeName(string(name))
+	return az.VMSet.GetInstanceTypeByNodeName(ctx, string(name))
 }
 
 // AddSSHKeyToAllInstances adds an SSH public key as a legal identity for all instances
diff --git a/pkg/provider/azure_loadbalancer.go b/pkg/provider/azure_loadbalancer.go
index 90400b4ca..cd9572e9b 100644
--- a/pkg/provider/azure_loadbalancer.go
+++ b/pkg/provider/azure_loadbalancer.go
@@ -1867,7 +1867,7 @@ func (az *Cloud) reconcileLoadBalancer(ctx context.Context, clusterName string,
 		if az.useMultipleStandardLoadBalancers() {
 			lbToReconcile = *existingLBs
 		}
-		lb, err = az.reconcileBackendPoolHosts(lb, lbToReconcile, service, nodes, clusterName, vmSetName, lbBackendPoolIDs)
+		lb, err = az.reconcileBackendPoolHosts(ctx, lb, lbToReconcile, service, nodes, clusterName, vmSetName, lbBackendPoolIDs)
 		if err != nil {
 			return nil, err
 		}
@@ -1882,6 +1882,7 @@ func (az *Cloud) reconcileLoadBalancer(ctx context.Context, clusterName string,
 }
 
 func (az *Cloud) reconcileBackendPoolHosts(
+	ctx context.Context,
 	currentLB *network.LoadBalancer,
 	lbs []network.LoadBalancer,
 	service *v1.Service,
@@ -1899,6 +1900,7 @@ func (az *Cloud) reconcileBackendPoolHosts(
 				isIPv6 := isBackendPoolIPv6(ptr.Deref(backendPool.Name, ""))
 				if strings.EqualFold(ptr.Deref(backendPool.Name, ""), az.getBackendPoolNameForService(service, clusterName, isIPv6)) {
 					if err := az.LoadBalancerBackendPool.EnsureHostsInPool(
+						ctx,
 						service,
 						nodes,
 						lbBackendPoolIDs[isIPv6],
diff --git a/pkg/provider/azure_loadbalancer_backendpool.go b/pkg/provider/azure_loadbalancer_backendpool.go
index 2cf104761..b094dc3a5 100644
--- a/pkg/provider/azure_loadbalancer_backendpool.go
+++ b/pkg/provider/azure_loadbalancer_backendpool.go
@@ -19,6 +19,7 @@ package provider
 //go:generate sh -c "mockgen -destination=$GOPATH/src/sigs.k8s.io/cloud-provider-azure/pkg/provider/azure_mock_loadbalancer_backendpool.go -source=$GOPATH/src/sigs.k8s.io/cloud-provider-azure/pkg/provider/azure_loadbalancer_backendpool.go -package=provider BackendPool"
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"strings"
@@ -39,7 +40,7 @@ import (
 
 type BackendPool interface {
 	// EnsureHostsInPool ensures the nodes join the backend pool of the load balancer
-	EnsureHostsInPool(service *v1.Service, nodes []*v1.Node, backendPoolID, vmSetName, clusterName, lbName string, backendPool network.BackendAddressPool) error
+	EnsureHostsInPool(ctx context.Context, service *v1.Service, nodes []*v1.Node, backendPoolID, vmSetName, clusterName, lbName string, backendPool network.BackendAddressPool) error
 
 	// CleanupVMSetFromBackendPoolByCondition removes nodes of the unwanted vmSet from the lb backend pool.
 	// This is needed in two scenarios:
@@ -66,8 +67,8 @@ func newBackendPoolTypeNodeIPConfig(c *Cloud) BackendPool {
 	return &backendPoolTypeNodeIPConfig{c}
 }
 
-func (bc *backendPoolTypeNodeIPConfig) EnsureHostsInPool(service *v1.Service, nodes []*v1.Node, backendPoolID, vmSetName, _, _ string, _ network.BackendAddressPool) error {
-	return bc.VMSet.EnsureHostsInPool(service, nodes, backendPoolID, vmSetName)
+func (bc *backendPoolTypeNodeIPConfig) EnsureHostsInPool(ctx context.Context, service *v1.Service, nodes []*v1.Node, backendPoolID, vmSetName, _, _ string, _ network.BackendAddressPool) error {
+	return bc.VMSet.EnsureHostsInPool(ctx, service, nodes, backendPoolID, vmSetName)
 }
 
 func isLBBackendPoolsExisting(lbBackendPoolNames map[bool]string, bpName *string) (found, isIPv6 bool) {
@@ -401,7 +402,7 @@ func (az *Cloud) getVnetResourceID() string {
 	)
 }
 
-func (bi *backendPoolTypeNodeIP) EnsureHostsInPool(service *v1.Service, nodes []*v1.Node, _, _, clusterName, lbName string, backendPool network.BackendAddressPool) error {
+func (bi *backendPoolTypeNodeIP) EnsureHostsInPool(_ context.Context, service *v1.Service, nodes []*v1.Node, _, _, clusterName, lbName string, backendPool network.BackendAddressPool) error {
 	isIPv6 := isBackendPoolIPv6(ptr.Deref(backendPool.Name, ""))
 
 	var (
diff --git a/pkg/provider/azure_mock_loadbalancer_backendpool.go b/pkg/provider/azure_mock_loadbalancer_backendpool.go
index 5735ed7f1..0a7393a42 100644
--- a/pkg/provider/azure_mock_loadbalancer_backendpool.go
+++ b/pkg/provider/azure_mock_loadbalancer_backendpool.go
@@ -1,22 +1,33 @@
-/*
-Copyright 2021 The Kubernetes Authors.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-*/
-
+// /*
+// Copyright The Kubernetes Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// */
+//
+
+// Code generated by MockGen. DO NOT EDIT.
+// Source: azure_loadbalancer_backendpool.go
+//
+// Generated by this command:
+//
+//	mockgen -package provider -source azure_loadbalancer_backendpool.go -self_package sigs.k8s.io/cloud-provider-azure/pkg/provider -copyright_file ../../hack/boilerplate/boilerplate.generatego.txt BackendPool
+//
+
+// Package provider is a generated GoMock package.
 package provider
 
 import (
+	context "context"
 	reflect "reflect"
 
 	network "github.com/Azure/azure-sdk-for-go/services/network/mgmt/2022-07-01/network"
@@ -57,23 +68,23 @@ func (m *MockBackendPool) CleanupVMSetFromBackendPoolByCondition(slb *network.Lo
 }
 
 // CleanupVMSetFromBackendPoolByCondition indicates an expected call of CleanupVMSetFromBackendPoolByCondition.
-func (mr *MockBackendPoolMockRecorder) CleanupVMSetFromBackendPoolByCondition(slb, service, nodes, clusterName, shouldRemoveVMSetFromSLB interface{}) *gomock.Call {
+func (mr *MockBackendPoolMockRecorder) CleanupVMSetFromBackendPoolByCondition(slb, service, nodes, clusterName, shouldRemoveVMSetFromSLB any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CleanupVMSetFromBackendPoolByCondition", reflect.TypeOf((*MockBackendPool)(nil).CleanupVMSetFromBackendPoolByCondition), slb, service, nodes, clusterName, shouldRemoveVMSetFromSLB)
 }
 
 // EnsureHostsInPool mocks base method.
-func (m *MockBackendPool) EnsureHostsInPool(service *v1.Service, nodes []*v1.Node, backendPoolID, vmSetName, clusterName, lbName string, backendPool network.BackendAddressPool) error {
+func (m *MockBackendPool) EnsureHostsInPool(ctx context.Context, service *v1.Service, nodes []*v1.Node, backendPoolID, vmSetName, clusterName, lbName string, backendPool network.BackendAddressPool) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "EnsureHostsInPool", service, nodes, backendPoolID, vmSetName, clusterName, lbName, backendPool)
+	ret := m.ctrl.Call(m, "EnsureHostsInPool", ctx, service, nodes, backendPoolID, vmSetName, clusterName, lbName, backendPool)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
 // EnsureHostsInPool indicates an expected call of EnsureHostsInPool.
-func (mr *MockBackendPoolMockRecorder) EnsureHostsInPool(service, nodes, backendPoolID, vmSetName, clusterName, lbName, backendPool interface{}) *gomock.Call {
+func (mr *MockBackendPoolMockRecorder) EnsureHostsInPool(ctx, service, nodes, backendPoolID, vmSetName, clusterName, lbName, backendPool any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EnsureHostsInPool", reflect.TypeOf((*MockBackendPool)(nil).EnsureHostsInPool), service, nodes, backendPoolID, vmSetName, clusterName, lbName, backendPool)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EnsureHostsInPool", reflect.TypeOf((*MockBackendPool)(nil).EnsureHostsInPool), ctx, service, nodes, backendPoolID, vmSetName, clusterName, lbName, backendPool)
 }
 
 // GetBackendPrivateIPs mocks base method.
@@ -86,7 +97,7 @@ func (m *MockBackendPool) GetBackendPrivateIPs(clusterName string, service *v1.S
 }
 
 // GetBackendPrivateIPs indicates an expected call of GetBackendPrivateIPs.
-func (mr *MockBackendPoolMockRecorder) GetBackendPrivateIPs(clusterName, service, lb interface{}) *gomock.Call {
+func (mr *MockBackendPoolMockRecorder) GetBackendPrivateIPs(clusterName, service, lb any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetBackendPrivateIPs", reflect.TypeOf((*MockBackendPool)(nil).GetBackendPrivateIPs), clusterName, service, lb)
 }
@@ -103,7 +114,7 @@ func (m *MockBackendPool) ReconcileBackendPools(clusterName string, service *v1.
 }
 
 // ReconcileBackendPools indicates an expected call of ReconcileBackendPools.
-func (mr *MockBackendPoolMockRecorder) ReconcileBackendPools(clusterName, service, lb interface{}) *gomock.Call {
+func (mr *MockBackendPoolMockRecorder) ReconcileBackendPools(clusterName, service, lb any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReconcileBackendPools", reflect.TypeOf((*MockBackendPool)(nil).ReconcileBackendPools), clusterName, service, lb)
 }
diff --git a/pkg/provider/azure_mock_vmsets.go b/pkg/provider/azure_mock_vmsets.go
index 608332b9e..da0fc950d 100644
--- a/pkg/provider/azure_mock_vmsets.go
+++ b/pkg/provider/azure_mock_vmsets.go
@@ -30,7 +30,7 @@ import (
 	context "context"
 	reflect "reflect"
 
-	v5 "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v6"
+	v6 "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v6"
 	compute "github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2022-08-01/compute"
 	network "github.com/Azure/azure-sdk-for-go/services/network/mgmt/2022-07-01/network"
 	gomock "go.uber.org/mock/gomock"
@@ -94,15 +94,15 @@ func (mr *MockVMSetMockRecorder) DeleteCacheForNode(nodeName any) *gomock.Call {
 // DetachDisk mocks base method.
 func (m *MockVMSet) DetachDisk(ctx context.Context, nodeName types.NodeName, diskMap map[string]string, forceDetach bool) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "DetachDisk", ctx, nodeName, diskMap)
+	ret := m.ctrl.Call(m, "DetachDisk", ctx, nodeName, diskMap, forceDetach)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
 // DetachDisk indicates an expected call of DetachDisk.
-func (mr *MockVMSetMockRecorder) DetachDisk(ctx, nodeName, diskMap any) *gomock.Call {
+func (mr *MockVMSetMockRecorder) DetachDisk(ctx, nodeName, diskMap, forceDetach any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DetachDisk", reflect.TypeOf((*MockVMSet)(nil).DetachDisk), ctx, nodeName, diskMap)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DetachDisk", reflect.TypeOf((*MockVMSet)(nil).DetachDisk), ctx, nodeName, diskMap, forceDetach)
 }
 
 // EnsureBackendPoolDeleted mocks base method.
@@ -135,9 +135,9 @@ func (mr *MockVMSetMockRecorder) EnsureBackendPoolDeletedFromVMSets(vmSetNamesMa
 }
 
 // EnsureHostInPool mocks base method.
-func (m *MockVMSet) EnsureHostInPool(service *v1.Service, nodeName types.NodeName, backendPoolID, vmSetName string) (string, string, string, *compute.VirtualMachineScaleSetVM, error) {
+func (m *MockVMSet) EnsureHostInPool(ctx context.Context, service *v1.Service, nodeName types.NodeName, backendPoolID, vmSetName string) (string, string, string, *compute.VirtualMachineScaleSetVM, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "EnsureHostInPool", service, nodeName, backendPoolID, vmSetName)
+	ret := m.ctrl.Call(m, "EnsureHostInPool", ctx, service, nodeName, backendPoolID, vmSetName)
 	ret0, _ := ret[0].(string)
 	ret1, _ := ret[1].(string)
 	ret2, _ := ret[2].(string)
@@ -147,23 +147,23 @@ func (m *MockVMSet) EnsureHostInPool(service *v1.Service, nodeName types.NodeNam
 }
 
 // EnsureHostInPool indicates an expected call of EnsureHostInPool.
-func (mr *MockVMSetMockRecorder) EnsureHostInPool(service, nodeName, backendPoolID, vmSetName any) *gomock.Call {
+func (mr *MockVMSetMockRecorder) EnsureHostInPool(ctx, service, nodeName, backendPoolID, vmSetName any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EnsureHostInPool", reflect.TypeOf((*MockVMSet)(nil).EnsureHostInPool), service, nodeName, backendPoolID, vmSetName)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EnsureHostInPool", reflect.TypeOf((*MockVMSet)(nil).EnsureHostInPool), ctx, service, nodeName, backendPoolID, vmSetName)
 }
 
 // EnsureHostsInPool mocks base method.
-func (m *MockVMSet) EnsureHostsInPool(service *v1.Service, nodes []*v1.Node, backendPoolID, vmSetName string) error {
+func (m *MockVMSet) EnsureHostsInPool(ctx context.Context, service *v1.Service, nodes []*v1.Node, backendPoolID, vmSetName string) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "EnsureHostsInPool", service, nodes, backendPoolID, vmSetName)
+	ret := m.ctrl.Call(m, "EnsureHostsInPool", ctx, service, nodes, backendPoolID, vmSetName)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
 // EnsureHostsInPool indicates an expected call of EnsureHostsInPool.
-func (mr *MockVMSetMockRecorder) EnsureHostsInPool(service, nodes, backendPoolID, vmSetName any) *gomock.Call {
+func (mr *MockVMSetMockRecorder) EnsureHostsInPool(ctx, service, nodes, backendPoolID, vmSetName any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EnsureHostsInPool", reflect.TypeOf((*MockVMSet)(nil).EnsureHostsInPool), service, nodes, backendPoolID, vmSetName)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EnsureHostsInPool", reflect.TypeOf((*MockVMSet)(nil).EnsureHostsInPool), ctx, service, nodes, backendPoolID, vmSetName)
 }
 
 // GetAgentPoolVMSetNames mocks base method.
@@ -182,10 +182,10 @@ func (mr *MockVMSetMockRecorder) GetAgentPoolVMSetNames(nodes any) *gomock.Call
 }
 
 // GetDataDisks mocks base method.
-func (m *MockVMSet) GetDataDisks(nodeName types.NodeName, crt cache.AzureCacheReadType) ([]*v5.DataDisk, *string, error) {
+func (m *MockVMSet) GetDataDisks(nodeName types.NodeName, crt cache.AzureCacheReadType) ([]*v6.DataDisk, *string, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "GetDataDisks", nodeName, crt)
-	ret0, _ := ret[0].([]*v5.DataDisk)
+	ret0, _ := ret[0].([]*v6.DataDisk)
 	ret1, _ := ret[1].(*string)
 	ret2, _ := ret[2].(error)
 	return ret0, ret1, ret2
@@ -198,9 +198,9 @@ func (mr *MockVMSetMockRecorder) GetDataDisks(nodeName, crt any) *gomock.Call {
 }
 
 // GetIPByNodeName mocks base method.
-func (m *MockVMSet) GetIPByNodeName(name string) (string, string, error) {
+func (m *MockVMSet) GetIPByNodeName(ctx context.Context, name string) (string, string, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetIPByNodeName", name)
+	ret := m.ctrl.Call(m, "GetIPByNodeName", ctx, name)
 	ret0, _ := ret[0].(string)
 	ret1, _ := ret[1].(string)
 	ret2, _ := ret[2].(error)
@@ -208,39 +208,39 @@ func (m *MockVMSet) GetIPByNodeName(name string) (string, string, error) {
 }
 
 // GetIPByNodeName indicates an expected call of GetIPByNodeName.
-func (mr *MockVMSetMockRecorder) GetIPByNodeName(name any) *gomock.Call {
+func (mr *MockVMSetMockRecorder) GetIPByNodeName(ctx, name any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetIPByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetIPByNodeName), name)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetIPByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetIPByNodeName), ctx, name)
 }
 
 // GetInstanceIDByNodeName mocks base method.
-func (m *MockVMSet) GetInstanceIDByNodeName(name string) (string, error) {
+func (m *MockVMSet) GetInstanceIDByNodeName(ctx context.Context, name string) (string, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetInstanceIDByNodeName", name)
+	ret := m.ctrl.Call(m, "GetInstanceIDByNodeName", ctx, name)
 	ret0, _ := ret[0].(string)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
 // GetInstanceIDByNodeName indicates an expected call of GetInstanceIDByNodeName.
-func (mr *MockVMSetMockRecorder) GetInstanceIDByNodeName(name any) *gomock.Call {
+func (mr *MockVMSetMockRecorder) GetInstanceIDByNodeName(ctx, name any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetInstanceIDByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetInstanceIDByNodeName), name)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetInstanceIDByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetInstanceIDByNodeName), ctx, name)
 }
 
 // GetInstanceTypeByNodeName mocks base method.
-func (m *MockVMSet) GetInstanceTypeByNodeName(name string) (string, error) {
+func (m *MockVMSet) GetInstanceTypeByNodeName(ctx context.Context, name string) (string, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetInstanceTypeByNodeName", name)
+	ret := m.ctrl.Call(m, "GetInstanceTypeByNodeName", ctx, name)
 	ret0, _ := ret[0].(string)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
 // GetInstanceTypeByNodeName indicates an expected call of GetInstanceTypeByNodeName.
-func (mr *MockVMSetMockRecorder) GetInstanceTypeByNodeName(name any) *gomock.Call {
+func (mr *MockVMSetMockRecorder) GetInstanceTypeByNodeName(ctx, name any) *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, "GetInstanceTypeByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetInstanceTypeByNodeName), ctx, name)
 }
 
 // GetNodeCIDRMasksByProviderID mocks base method.
@@ -321,18 +321,18 @@ func (mr *MockVMSetMockRecorder) GetPowerStatusByNodeName(name any) *gomock.Call
 }
 
 // GetPrimaryInterface mocks base method.
-func (m *MockVMSet) GetPrimaryInterface(nodeName string) (network.Interface, error) {
+func (m *MockVMSet) GetPrimaryInterface(ctx context.Context, nodeName string) (network.Interface, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetPrimaryInterface", nodeName)
+	ret := m.ctrl.Call(m, "GetPrimaryInterface", ctx, nodeName)
 	ret0, _ := ret[0].(network.Interface)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
 // GetPrimaryInterface indicates an expected call of GetPrimaryInterface.
-func (mr *MockVMSetMockRecorder) GetPrimaryInterface(nodeName any) *gomock.Call {
+func (mr *MockVMSetMockRecorder) GetPrimaryInterface(ctx, nodeName any) *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, "GetPrimaryInterface", reflect.TypeOf((*MockVMSet)(nil).GetPrimaryInterface), ctx, nodeName)
 }
 
 // GetPrimaryVMSetName mocks base method.
@@ -350,18 +350,18 @@ func (mr *MockVMSetMockRecorder) GetPrimaryVMSetName() *gomock.Call {
 }
 
 // GetPrivateIPsByNodeName mocks base method.
-func (m *MockVMSet) GetPrivateIPsByNodeName(name string) ([]string, error) {
+func (m *MockVMSet) GetPrivateIPsByNodeName(ctx context.Context, name string) ([]string, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetPrivateIPsByNodeName", name)
+	ret := m.ctrl.Call(m, "GetPrivateIPsByNodeName", ctx, name)
 	ret0, _ := ret[0].([]string)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
 // GetPrivateIPsByNodeName indicates an expected call of GetPrivateIPsByNodeName.
-func (mr *MockVMSetMockRecorder) GetPrivateIPsByNodeName(name any) *gomock.Call {
+func (mr *MockVMSetMockRecorder) GetPrivateIPsByNodeName(ctx, name any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPrivateIPsByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetPrivateIPsByNodeName), name)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPrivateIPsByNodeName", reflect.TypeOf((*MockVMSet)(nil).GetPrivateIPsByNodeName), ctx, name)
 }
 
 // GetProvisioningStateByNodeName mocks base method.
diff --git a/pkg/provider/azure_routes.go b/pkg/provider/azure_routes.go
index a5b7c6c41..93fa18a44 100644
--- a/pkg/provider/azure_routes.go
+++ b/pkg/provider/azure_routes.go
@@ -371,7 +371,7 @@ func (az *Cloud) createRouteTable() error {
 // route.Name will be ignored, although the cloud-provider may use nameHint
 // to create a more user-meaningful name.
 // implements cloudprovider.Routes.CreateRoute
-func (az *Cloud) CreateRoute(_ context.Context, clusterName string, _ string, kubeRoute *cloudprovider.Route) error {
+func (az *Cloud) CreateRoute(ctx context.Context, clusterName string, _ string, kubeRoute *cloudprovider.Route) error {
 	mc := metrics.NewMetricContext("routes", "create_route", az.ResourceGroup, az.getNetworkResourceSubscriptionID(), string(kubeRoute.TargetNode))
 	isOperationSucceeded := false
 	defer func() {
@@ -398,7 +398,7 @@ func (az *Cloud) CreateRoute(_ context.Context, clusterName string, _ string, ku
 	// single stack IPv6 is supported on dual stack host. So the IPv6 IP is secondary IP for both single stack IPv6 and dual stack
 	// Get all private IPs for the machine and find the first one that matches the IPv6 family
 	if !az.ipv6DualStackEnabled && !CIDRv6 {
-		targetIP, _, err = az.getIPForMachine(kubeRoute.TargetNode)
+		targetIP, _, err = az.getIPForMachine(ctx, kubeRoute.TargetNode)
 		if err != nil {
 			return err
 		}
@@ -406,7 +406,7 @@ func (az *Cloud) CreateRoute(_ context.Context, clusterName string, _ string, ku
 		// for dual stack and single stack IPv6 we need to select
 		// a private ip that matches family of the cidr
 		klog.V(4).Infof("CreateRoute: create route instance=%q cidr=%q is in dual stack mode", kubeRoute.TargetNode, kubeRoute.DestinationCIDR)
-		nodePrivateIPs, err := az.getPrivateIPsForMachine(kubeRoute.TargetNode)
+		nodePrivateIPs, err := az.getPrivateIPsForMachine(ctx, kubeRoute.TargetNode)
 		if nil != err {
 			klog.V(3).Infof("CreateRoute: create route: failed(GetPrivateIPsByNodeName) instance=%q cidr=%q with error=%v", kubeRoute.TargetNode, kubeRoute.DestinationCIDR, err)
 			return err
diff --git a/pkg/provider/azure_standard.go b/pkg/provider/azure_standard.go
index 151cbc633..c58d0f9ca 100644
--- a/pkg/provider/azure_standard.go
+++ b/pkg/provider/azure_standard.go
@@ -464,7 +464,7 @@ func newAvailabilitySet(az *Cloud) (VMSet, error) {
 // GetInstanceIDByNodeName gets the cloud provider ID by node name.
 // It must return ("", cloudprovider.InstanceNotFound) if the instance does
 // not exist or is no longer running.
-func (as *availabilitySet) GetInstanceIDByNodeName(name string) (string, error) {
+func (as *availabilitySet) GetInstanceIDByNodeName(_ context.Context, name string) (string, error) {
 	var machine compute.VirtualMachine
 	var err error
 
@@ -537,7 +537,7 @@ func (as *availabilitySet) GetNodeNameByProviderID(providerID string) (types.Nod
 }
 
 // GetInstanceTypeByNodeName gets the instance type by node name.
-func (as *availabilitySet) GetInstanceTypeByNodeName(name string) (string, error) {
+func (as *availabilitySet) GetInstanceTypeByNodeName(_ context.Context, name string) (string, error) {
 	machine, err := as.getVirtualMachine(types.NodeName(name), azcache.CacheReadTypeUnsafe)
 	if err != nil {
 		klog.Errorf("as.GetInstanceTypeByNodeName(%s) failed: as.getVirtualMachine(%s) err=%v", name, name, err)
@@ -588,8 +588,8 @@ func (as *availabilitySet) GetPrimaryVMSetName() string {
 }
 
 // GetIPByNodeName gets machine private IP and public IP by node name.
-func (as *availabilitySet) GetIPByNodeName(name string) (string, string, error) {
-	nic, err := as.GetPrimaryInterface(name)
+func (as *availabilitySet) GetIPByNodeName(ctx context.Context, name string) (string, string, error) {
+	nic, err := as.GetPrimaryInterface(ctx, name)
 	if err != nil {
 		return "", "", err
 	}
@@ -623,9 +623,9 @@ func (as *availabilitySet) GetIPByNodeName(name string) (string, string, error)
 // returns a list of private ips assigned to node
 // TODO (khenidak): This should read all nics, not just the primary
 // allowing users to split ipv4/v6 on multiple nics
-func (as *availabilitySet) GetPrivateIPsByNodeName(name string) ([]string, error) {
+func (as *availabilitySet) GetPrivateIPsByNodeName(ctx context.Context, name string) ([]string, error) {
 	ips := make([]string, 0)
-	nic, err := as.GetPrimaryInterface(name)
+	nic, err := as.GetPrimaryInterface(ctx, name)
 	if err != nil {
 		return ips, err
 	}
@@ -770,7 +770,7 @@ func (as *availabilitySet) GetNodeVMSetName(node *v1.Node) (string, error) {
 }
 
 // GetPrimaryInterface gets machine primary network interface by node name.
-func (as *availabilitySet) GetPrimaryInterface(nodeName string) (network.Interface, error) {
+func (as *availabilitySet) GetPrimaryInterface(_ context.Context, nodeName string) (network.Interface, error) {
 	nic, _, err := as.getPrimaryInterfaceWithVMSet(nodeName, "")
 	return nic, err
 }
@@ -851,7 +851,7 @@ func (as *availabilitySet) getPrimaryInterfaceWithVMSet(nodeName, vmSetName stri
 
 // EnsureHostInPool ensures the given VM's Primary NIC's Primary IP Configuration is
 // participating in the specified LoadBalancer Backend Pool.
-func (as *availabilitySet) EnsureHostInPool(service *v1.Service, nodeName types.NodeName, backendPoolID string, vmSetName string) (string, string, string, *compute.VirtualMachineScaleSetVM, error) {
+func (as *availabilitySet) EnsureHostInPool(_ context.Context, service *v1.Service, nodeName types.NodeName, backendPoolID string, vmSetName string) (string, string, string, *compute.VirtualMachineScaleSetVM, error) {
 	vmName := mapNodeNameToVMName(nodeName)
 	serviceName := getServiceName(service)
 	nic, _, err := as.getPrimaryInterfaceWithVMSet(vmName, vmSetName)
@@ -936,7 +936,7 @@ func (as *availabilitySet) EnsureHostInPool(service *v1.Service, nodeName types.
 
 // EnsureHostsInPool ensures the given Node's primary IP configurations are
 // participating in the specified LoadBalancer Backend Pool.
-func (as *availabilitySet) EnsureHostsInPool(service *v1.Service, nodes []*v1.Node, backendPoolID string, vmSetName string) error {
+func (as *availabilitySet) EnsureHostsInPool(ctx context.Context, service *v1.Service, nodes []*v1.Node, backendPoolID string, vmSetName string) error {
 	mc := metrics.NewMetricContext("services", "vmas_ensure_hosts_in_pool", as.ResourceGroup, as.SubscriptionID, getServiceName(service))
 	isOperationSucceeded := false
 	defer func() {
@@ -962,7 +962,7 @@ func (as *availabilitySet) EnsureHostsInPool(service *v1.Service, nodes []*v1.No
 		}
 
 		f := func() error {
-			_, _, _, _, err := as.EnsureHostInPool(service, types.NodeName(localNodeName), backendPoolID, vmSetName)
+			_, _, _, _, err := as.EnsureHostInPool(ctx, service, types.NodeName(localNodeName), backendPoolID, vmSetName)
 			if err != nil {
 				return fmt.Errorf("ensure(%s): backendPoolID(%s) - failed to ensure host in pool: %w", getServiceName(service), backendPoolID, err)
 			}
diff --git a/pkg/provider/azure_vmsets.go b/pkg/provider/azure_vmsets.go
index 929a0ca9d..03535a0da 100644
--- a/pkg/provider/azure_vmsets.go
+++ b/pkg/provider/azure_vmsets.go
@@ -40,13 +40,13 @@ type VMSet interface {
 	// GetInstanceIDByNodeName gets the cloud provider ID by node name.
 	// It must return ("", cloudprovider.InstanceNotFound) if the instance does
 	// not exist or is no longer running.
-	GetInstanceIDByNodeName(name string) (string, error)
+	GetInstanceIDByNodeName(ctx context.Context, name string) (string, error)
 	// GetInstanceTypeByNodeName gets the instance type by node name.
-	GetInstanceTypeByNodeName(name string) (string, error)
+	GetInstanceTypeByNodeName(ctx context.Context, name string) (string, error)
 	// GetIPByNodeName gets machine private IP and public IP by node name.
-	GetIPByNodeName(name string) (string, string, error)
+	GetIPByNodeName(ctx context.Context, name string) (string, string, error)
 	// GetPrimaryInterface gets machine primary network interface by node name.
-	GetPrimaryInterface(nodeName string) (network.Interface, error)
+	GetPrimaryInterface(ctx context.Context, nodeName string) (network.Interface, error)
 	// GetNodeNameByProviderID gets the node name by provider ID.
 	GetNodeNameByProviderID(providerID string) (types.NodeName, error)
 
@@ -66,10 +66,10 @@ type VMSet interface {
 	GetNodeVMSetName(node *v1.Node) (string, error)
 	// EnsureHostsInPool ensures the given Node's primary IP configurations are
 	// participating in the specified LoadBalancer Backend Pool.
-	EnsureHostsInPool(service *v1.Service, nodes []*v1.Node, backendPoolID string, vmSetName string) error
+	EnsureHostsInPool(ctx context.Context, service *v1.Service, nodes []*v1.Node, backendPoolID string, vmSetName string) error
 	// EnsureHostInPool ensures the given VM's Primary NIC's Primary IP Configuration is
 	// participating in the specified LoadBalancer Backend Pool.
-	EnsureHostInPool(service *v1.Service, nodeName types.NodeName, backendPoolID string, vmSetName string) (string, string, string, *compute.VirtualMachineScaleSetVM, error)
+	EnsureHostInPool(ctx context.Context, service *v1.Service, nodeName types.NodeName, backendPoolID string, vmSetName string) (string, string, string, *compute.VirtualMachineScaleSetVM, error)
 	// EnsureBackendPoolDeleted ensures the loadBalancer backendAddressPools deleted from the specified nodes.
 	EnsureBackendPoolDeleted(service *v1.Service, backendPoolIDs []string, vmSetName string, backendAddressPools *[]network.BackendAddressPool, deleteFromVMSet bool) (bool, error)
 	//EnsureBackendPoolDeletedFromVMSets ensures the loadBalancer backendAddressPools deleted from the specified VMSS/VMAS
@@ -94,7 +94,7 @@ type VMSet interface {
 	GetProvisioningStateByNodeName(name string) (string, error)
 
 	// GetPrivateIPsByNodeName returns a slice of all private ips assigned to node (ipv6 and ipv4)
-	GetPrivateIPsByNodeName(name string) ([]string, error)
+	GetPrivateIPsByNodeName(ctx context.Context, name string) ([]string, error)
 
 	// GetNodeNameByIPConfigurationID gets the nodeName and vmSetName by IP configuration ID.
 	GetNodeNameByIPConfigurationID(ipConfigurationID string) (string, string, error)
diff --git a/pkg/provider/azure_vmsets_repo.go b/pkg/provider/azure_vmsets_repo.go
index bbd79c9c8..25c767dc8 100644
--- a/pkg/provider/azure_vmsets_repo.go
+++ b/pkg/provider/azure_vmsets_repo.go
@@ -17,6 +17,7 @@ limitations under the License.
 package provider
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"strings"
@@ -73,15 +74,15 @@ func (az *Cloud) ListVirtualMachines(resourceGroup string) ([]compute.VirtualMac
 
 // getPrivateIPsForMachine is wrapper for optional backoff getting private ips
 // list of a node by name
-func (az *Cloud) getPrivateIPsForMachine(nodeName types.NodeName) ([]string, error) {
-	return az.getPrivateIPsForMachineWithRetry(nodeName)
+func (az *Cloud) getPrivateIPsForMachine(ctx context.Context, nodeName types.NodeName) ([]string, error) {
+	return az.getPrivateIPsForMachineWithRetry(ctx, nodeName)
 }
 
-func (az *Cloud) getPrivateIPsForMachineWithRetry(nodeName types.NodeName) ([]string, error) {
+func (az *Cloud) getPrivateIPsForMachineWithRetry(ctx context.Context, nodeName types.NodeName) ([]string, error) {
 	var privateIPs []string
 	err := wait.ExponentialBackoff(az.RequestBackoff(), func() (bool, error) {
 		var retryErr error
-		privateIPs, retryErr = az.VMSet.GetPrivateIPsByNodeName(string(nodeName))
+		privateIPs, retryErr = az.VMSet.GetPrivateIPsByNodeName(ctx, string(nodeName))
 		if retryErr != nil {
 			// won't retry since the instance doesn't exist on Azure.
 			if errors.Is(retryErr, cloudprovider.InstanceNotFound) {
@@ -96,16 +97,16 @@ func (az *Cloud) getPrivateIPsForMachineWithRetry(nodeName types.NodeName) ([]st
 	return privateIPs, err
 }
 
-func (az *Cloud) getIPForMachine(nodeName types.NodeName) (string, string, error) {
-	return az.GetIPForMachineWithRetry(nodeName)
+func (az *Cloud) getIPForMachine(ctx context.Context, nodeName types.NodeName) (string, string, error) {
+	return az.GetIPForMachineWithRetry(ctx, nodeName)
 }
 
 // GetIPForMachineWithRetry invokes az.getIPForMachine with exponential backoff retry
-func (az *Cloud) GetIPForMachineWithRetry(name types.NodeName) (string, string, error) {
+func (az *Cloud) GetIPForMachineWithRetry(ctx context.Context, name types.NodeName) (string, string, error) {
 	var ip, publicIP string
-	err := wait.ExponentialBackoff(az.RequestBackoff(), func() (bool, error) {
+	err := wait.ExponentialBackoffWithContext(ctx, az.RequestBackoff(), func(ctx context.Context) (bool, error) {
 		var retryErr error
-		ip, publicIP, retryErr = az.VMSet.GetIPByNodeName(string(name))
+		ip, publicIP, retryErr = az.VMSet.GetIPByNodeName(ctx, string(name))
 		if retryErr != nil {
 			klog.Errorf("GetIPForMachineWithRetry(%s): backoff failure, will retry,err=%v", name, retryErr)
 			return false, nil
diff --git a/pkg/provider/azure_vmss.go b/pkg/provider/azure_vmss.go
index 8b3aa5818..70034ffd8 100644
--- a/pkg/provider/azure_vmss.go
+++ b/pkg/provider/azure_vmss.go
@@ -381,7 +381,7 @@ func (ss *ScaleSet) getVmssVMByInstanceID(resourceGroup, scaleSetName, instanceI
 // GetInstanceIDByNodeName gets the cloud provider ID by node name.
 // It must return ("", cloudprovider.InstanceNotFound) if the instance does
 // not exist or is no longer running.
-func (ss *ScaleSet) GetInstanceIDByNodeName(name string) (string, error) {
+func (ss *ScaleSet) GetInstanceIDByNodeName(ctx context.Context, name string) (string, error) {
 	vmManagementType, err := ss.getVMManagementTypeByNodeName(name, azcache.CacheReadTypeUnsafe)
 	if err != nil {
 		klog.Errorf("Failed to check VM management type: %v", err)
@@ -390,11 +390,11 @@ func (ss *ScaleSet) GetInstanceIDByNodeName(name string) (string, error) {
 
 	if vmManagementType == ManagedByAvSet {
 		// vm is managed by availability set.
-		return ss.availabilitySet.GetInstanceIDByNodeName(name)
+		return ss.availabilitySet.GetInstanceIDByNodeName(ctx, name)
 	}
 	if vmManagementType == ManagedByVmssFlex {
 		// vm is managed by vmss flex.
-		return ss.flexScaleSet.GetInstanceIDByNodeName(name)
+		return ss.flexScaleSet.GetInstanceIDByNodeName(ctx, name)
 	}
 
 	vm, err := ss.getVmssVM(name, azcache.CacheReadTypeUnsafe)
@@ -480,7 +480,7 @@ func (ss *ScaleSet) GetNodeNameByProviderID(providerID string) (types.NodeName,
 }
 
 // GetInstanceTypeByNodeName gets the instance type by node name.
-func (ss *ScaleSet) GetInstanceTypeByNodeName(name string) (string, error) {
+func (ss *ScaleSet) GetInstanceTypeByNodeName(ctx context.Context, name string) (string, error) {
 	vmManagementType, err := ss.getVMManagementTypeByNodeName(name, azcache.CacheReadTypeUnsafe)
 	if err != nil {
 		klog.Errorf("Failed to check VM management type: %v", err)
@@ -489,11 +489,11 @@ func (ss *ScaleSet) GetInstanceTypeByNodeName(name string) (string, error) {
 
 	if vmManagementType == ManagedByAvSet {
 		// vm is managed by availability set.
-		return ss.availabilitySet.GetInstanceTypeByNodeName(name)
+		return ss.availabilitySet.GetInstanceTypeByNodeName(ctx, name)
 	}
 	if vmManagementType == ManagedByVmssFlex {
 		// vm is managed by vmss flex.
-		return ss.flexScaleSet.GetInstanceTypeByNodeName(name)
+		return ss.flexScaleSet.GetInstanceTypeByNodeName(ctx, name)
 	}
 
 	vm, err := ss.getVmssVM(name, azcache.CacheReadTypeUnsafe)
@@ -569,7 +569,7 @@ func (ss *ScaleSet) GetPrimaryVMSetName() string {
 }
 
 // GetIPByNodeName gets machine private IP and public IP by node name.
-func (ss *ScaleSet) GetIPByNodeName(nodeName string) (string, string, error) {
+func (ss *ScaleSet) GetIPByNodeName(ctx context.Context, nodeName string) (string, string, error) {
 	vmManagementType, err := ss.getVMManagementTypeByNodeName(nodeName, azcache.CacheReadTypeUnsafe)
 	if err != nil {
 		klog.Errorf("Failed to check VM management type: %v", err)
@@ -578,14 +578,14 @@ func (ss *ScaleSet) GetIPByNodeName(nodeName string) (string, string, error) {
 
 	if vmManagementType == ManagedByAvSet {
 		// vm is managed by availability set.
-		return ss.availabilitySet.GetIPByNodeName(nodeName)
+		return ss.availabilitySet.GetIPByNodeName(ctx, nodeName)
 	}
 	if vmManagementType == ManagedByVmssFlex {
 		// vm is managed by vmss flex.
-		return ss.flexScaleSet.GetIPByNodeName(nodeName)
+		return ss.flexScaleSet.GetIPByNodeName(ctx, nodeName)
 	}
 
-	nic, err := ss.GetPrimaryInterface(nodeName)
+	nic, err := ss.GetPrimaryInterface(ctx, nodeName)
 	if err != nil {
 		klog.Errorf("error: ss.GetIPByNodeName(%s), GetPrimaryInterface(%q), err=%v", nodeName, nodeName, err)
 		return "", "", err
@@ -646,7 +646,7 @@ func (ss *ScaleSet) getVMSSPublicIPAddress(resourceGroupName string, virtualMach
 // returns a list of private ips assigned to node
 // TODO (khenidak): This should read all nics, not just the primary
 // allowing users to split ipv4/v6 on multiple nics
-func (ss *ScaleSet) GetPrivateIPsByNodeName(nodeName string) ([]string, error) {
+func (ss *ScaleSet) GetPrivateIPsByNodeName(ctx context.Context, nodeName string) ([]string, error) {
 	ips := make([]string, 0)
 	vmManagementType, err := ss.getVMManagementTypeByNodeName(nodeName, azcache.CacheReadTypeUnsafe)
 	if err != nil {
@@ -656,14 +656,14 @@ func (ss *ScaleSet) GetPrivateIPsByNodeName(nodeName string) ([]string, error) {
 
 	if vmManagementType == ManagedByAvSet {
 		// vm is managed by availability set.
-		return ss.availabilitySet.GetPrivateIPsByNodeName(nodeName)
+		return ss.availabilitySet.GetPrivateIPsByNodeName(ctx, nodeName)
 	}
 	if vmManagementType == ManagedByVmssFlex {
 		// vm is managed by vmss flex.
-		return ss.flexScaleSet.GetPrivateIPsByNodeName(nodeName)
+		return ss.flexScaleSet.GetPrivateIPsByNodeName(ctx, nodeName)
 	}
 
-	nic, err := ss.GetPrimaryInterface(nodeName)
+	nic, err := ss.GetPrimaryInterface(ctx, nodeName)
 	if err != nil {
 		klog.Errorf("error: ss.GetIPByNodeName(%s), GetPrimaryInterface(%q), err=%v", nodeName, nodeName, err)
 		return ips, err
@@ -905,7 +905,7 @@ func extractResourceGroupByVMSSNicID(nicID string) (string, error) {
 }
 
 // GetPrimaryInterface gets machine primary network interface by node name and vmSet.
-func (ss *ScaleSet) GetPrimaryInterface(nodeName string) (network.Interface, error) {
+func (ss *ScaleSet) GetPrimaryInterface(ctx context.Context, nodeName string) (network.Interface, error) {
 	vmManagementType, err := ss.getVMManagementTypeByNodeName(nodeName, azcache.CacheReadTypeUnsafe)
 	if err != nil {
 		klog.Errorf("Failed to check VM management type: %v", err)
@@ -914,18 +914,18 @@ func (ss *ScaleSet) GetPrimaryInterface(nodeName string) (network.Interface, err
 
 	if vmManagementType == ManagedByAvSet {
 		// vm is managed by availability set.
-		return ss.availabilitySet.GetPrimaryInterface(nodeName)
+		return ss.availabilitySet.GetPrimaryInterface(ctx, nodeName)
 	}
 	if vmManagementType == ManagedByVmssFlex {
 		// vm is managed by vmss flex.
-		return ss.flexScaleSet.GetPrimaryInterface(nodeName)
+		return ss.flexScaleSet.GetPrimaryInterface(ctx, nodeName)
 	}
 
 	vm, err := ss.getVmssVM(nodeName, azcache.CacheReadTypeDefault)
 	if err != nil {
 		// VM is availability set, but not cached yet in availabilitySetNodesCache.
 		if errors.Is(err, ErrorNotVmssInstance) {
-			return ss.availabilitySet.GetPrimaryInterface(nodeName)
+			return ss.availabilitySet.GetPrimaryInterface(ctx, nodeName)
 		}
 
 		klog.Errorf("error: ss.GetPrimaryInterface(%s), ss.getVmssVM(%s), err=%v", nodeName, nodeName, err)
@@ -1023,7 +1023,7 @@ func getPrimaryIPConfigFromVMSSNetworkConfig(config *compute.VirtualMachineScale
 
 // EnsureHostInPool ensures the given VM's Primary NIC's Primary IP Configuration is
 // participating in the specified LoadBalancer Backend Pool, which returns (resourceGroup, vmasName, instanceID, vmssVM, error).
-func (ss *ScaleSet) EnsureHostInPool(_ *v1.Service, nodeName types.NodeName, backendPoolID string, vmSetNameOfLB string) (string, string, string, *compute.VirtualMachineScaleSetVM, error) {
+func (ss *ScaleSet) EnsureHostInPool(_ context.Context, _ *v1.Service, nodeName types.NodeName, backendPoolID string, vmSetNameOfLB string) (string, string, string, *compute.VirtualMachineScaleSetVM, error) {
 	logger := klog.Background().WithName("EnsureHostInPool").
 		WithValues("nodeName", nodeName, "backendPoolID", backendPoolID, "vmSetNameOfLB", vmSetNameOfLB)
 	vmName := mapNodeNameToVMName(nodeName)
@@ -1356,7 +1356,7 @@ func isWindows2019(vmss *compute.VirtualMachineScaleSet) bool {
 	return osVersion[0] == consts.Windows2019OSBuildVersion
 }
 
-func (ss *ScaleSet) ensureHostsInPool(service *v1.Service, nodes []*v1.Node, backendPoolID string, vmSetNameOfLB string) error {
+func (ss *ScaleSet) ensureHostsInPool(ctx context.Context, service *v1.Service, nodes []*v1.Node, backendPoolID string, vmSetNameOfLB string) error {
 	mc := metrics.NewMetricContext("services", "vmss_ensure_hosts_in_pool", ss.ResourceGroup, ss.SubscriptionID, getServiceName(service))
 	isOperationSucceeded := false
 	defer func() {
@@ -1384,7 +1384,7 @@ func (ss *ScaleSet) ensureHostsInPool(service *v1.Service, nodes []*v1.Node, bac
 			continue
 		}
 
-		nodeResourceGroup, nodeVMSS, nodeInstanceID, nodeVMSSVM, err := ss.EnsureHostInPool(service, types.NodeName(localNodeName), backendPoolID, vmSetNameOfLB)
+		nodeResourceGroup, nodeVMSS, nodeInstanceID, nodeVMSSVM, err := ss.EnsureHostInPool(ctx, service, types.NodeName(localNodeName), backendPoolID, vmSetNameOfLB)
 		if err != nil {
 			klog.Errorf("EnsureHostInPool(%s): backendPoolID(%s) - failed to ensure host in pool: %q", getServiceName(service), backendPoolID, err)
 			errors = append(errors, err)
@@ -1466,9 +1466,9 @@ func (ss *ScaleSet) ensureHostsInPool(service *v1.Service, nodes []*v1.Node, bac
 
 // EnsureHostsInPool ensures the given Node's primary IP configurations are
 // participating in the specified LoadBalancer Backend Pool.
-func (ss *ScaleSet) EnsureHostsInPool(service *v1.Service, nodes []*v1.Node, backendPoolID string, vmSetNameOfLB string) error {
+func (ss *ScaleSet) EnsureHostsInPool(ctx context.Context, service *v1.Service, nodes []*v1.Node, backendPoolID string, vmSetNameOfLB string) error {
 	if ss.DisableAvailabilitySetNodes && !ss.EnableVmssFlexNodes {
-		return ss.ensureHostsInPool(service, nodes, backendPoolID, vmSetNameOfLB)
+		return ss.ensureHostsInPool(ctx, service, nodes, backendPoolID, vmSetNameOfLB)
 	}
 	vmssUniformNodes := make([]*v1.Node, 0)
 	vmssFlexNodes := make([]*v1.Node, 0)
@@ -1522,17 +1522,17 @@ func (ss *ScaleSet) EnsureHostsInPool(service *v1.Service, nodes []*v1.Node, bac
 	}
 
 	if len(vmssFlexNodes) > 0 {
-		vmssFlexError := ss.flexScaleSet.EnsureHostsInPool(service, vmssFlexNodes, backendPoolID, vmSetNameOfLB)
+		vmssFlexError := ss.flexScaleSet.EnsureHostsInPool(ctx, service, vmssFlexNodes, backendPoolID, vmSetNameOfLB)
 		errors = append(errors, vmssFlexError)
 	}
 
 	if len(vmasNodes) > 0 {
-		vmasError := ss.availabilitySet.EnsureHostsInPool(service, vmasNodes, backendPoolID, vmSetNameOfLB)
+		vmasError := ss.availabilitySet.EnsureHostsInPool(ctx, service, vmasNodes, backendPoolID, vmSetNameOfLB)
 		errors = append(errors, vmasError)
 	}
 
 	if len(vmssUniformNodes) > 0 {
-		vmssUniformError := ss.ensureHostsInPool(service, vmssUniformNodes, backendPoolID, vmSetNameOfLB)
+		vmssUniformError := ss.ensureHostsInPool(ctx, service, vmssUniformNodes, backendPoolID, vmSetNameOfLB)
 		errors = append(errors, vmssUniformError)
 	}
 
diff --git a/pkg/provider/azure_vmssflex.go b/pkg/provider/azure_vmssflex.go
index 8bbe3d950..9b1244cbc 100644
--- a/pkg/provider/azure_vmssflex.go
+++ b/pkg/provider/azure_vmssflex.go
@@ -180,7 +180,7 @@ func (fs *FlexScaleSet) GetNodeNameByProviderID(providerID string) (types.NodeNa
 // GetInstanceIDByNodeName gets the cloud provider ID by node name.
 // It must return ("", cloudprovider.InstanceNotFound) if the instance does
 // not exist or is no longer running.
-func (fs *FlexScaleSet) GetInstanceIDByNodeName(name string) (string, error) {
+func (fs *FlexScaleSet) GetInstanceIDByNodeName(_ context.Context, name string) (string, error) {
 	machine, err := fs.getVmssFlexVM(name, azcache.CacheReadTypeUnsafe)
 	if err != nil {
 		return "", err
@@ -199,7 +199,7 @@ func (fs *FlexScaleSet) GetInstanceIDByNodeName(name string) (string, error) {
 }
 
 // GetInstanceTypeByNodeName gets the instance type by node name.
-func (fs *FlexScaleSet) GetInstanceTypeByNodeName(name string) (string, error) {
+func (fs *FlexScaleSet) GetInstanceTypeByNodeName(_ context.Context, name string) (string, error) {
 	machine, err := fs.getVmssFlexVM(name, azcache.CacheReadTypeUnsafe)
 	if err != nil {
 		klog.Errorf("fs.GetInstanceTypeByNodeName(%s) failed: fs.getVmssFlexVMWithoutInstanceView(%s) err=%v", name, name, err)
@@ -279,7 +279,7 @@ func (fs *FlexScaleSet) GetPowerStatusByNodeName(name string) (powerState string
 }
 
 // GetPrimaryInterface gets machine primary network interface by node name.
-func (fs *FlexScaleSet) GetPrimaryInterface(nodeName string) (network.Interface, error) {
+func (fs *FlexScaleSet) GetPrimaryInterface(ctx context.Context, nodeName string) (network.Interface, error) {
 	machine, err := fs.getVmssFlexVM(nodeName, azcache.CacheReadTypeDefault)
 	if err != nil {
 		klog.Errorf("fs.GetInstanceTypeByNodeName(%s) failed: fs.getVmssFlexVMWithoutInstanceView(%s) err=%v", nodeName, nodeName, err)
@@ -300,8 +300,6 @@ func (fs *FlexScaleSet) GetPrimaryInterface(nodeName string) (network.Interface,
 		return network.Interface{}, err
 	}
 
-	ctx, cancel := getContextWithCancel()
-	defer cancel()
 	nic, rerr := fs.InterfacesClient.Get(ctx, nicResourceGroup, nicName, "")
 	if rerr != nil {
 		return network.Interface{}, rerr.Error()
@@ -311,8 +309,8 @@ func (fs *FlexScaleSet) GetPrimaryInterface(nodeName string) (network.Interface,
 }
 
 // GetIPByNodeName gets machine private IP and public IP by node name.
-func (fs *FlexScaleSet) GetIPByNodeName(name string) (string, string, error) {
-	nic, err := fs.GetPrimaryInterface(name)
+func (fs *FlexScaleSet) GetIPByNodeName(ctx context.Context, name string) (string, string, error) {
+	nic, err := fs.GetPrimaryInterface(ctx, name)
 	if err != nil {
 		return "", "", err
 	}
@@ -347,9 +345,9 @@ func (fs *FlexScaleSet) GetIPByNodeName(name string) (string, string, error) {
 // GetPrivateIPsByNodeName returns a slice of all private ips assigned to node (ipv6 and ipv4)
 // TODO (khenidak): This should read all nics, not just the primary
 // allowing users to split ipv4/v6 on multiple nics
-func (fs *FlexScaleSet) GetPrivateIPsByNodeName(name string) ([]string, error) {
+func (fs *FlexScaleSet) GetPrivateIPsByNodeName(ctx context.Context, name string) ([]string, error) {
 	ips := make([]string, 0)
-	nic, err := fs.GetPrimaryInterface(name)
+	nic, err := fs.GetPrimaryInterface(ctx, name)
 	if err != nil {
 		return ips, err
 	}
@@ -441,7 +439,7 @@ func (fs *FlexScaleSet) GetNodeCIDRMasksByProviderID(providerID string) (int, in
 
 // EnsureHostInPool ensures the given VM's Primary NIC's Primary IP Configuration is
 // participating in the specified LoadBalancer Backend Pool, which returns (resourceGroup, vmasName, instanceID, vmssVM, error).
-func (fs *FlexScaleSet) EnsureHostInPool(service *v1.Service, nodeName types.NodeName, backendPoolID string, vmSetNameOfLB string) (string, string, string, *compute.VirtualMachineScaleSetVM, error) {
+func (fs *FlexScaleSet) EnsureHostInPool(ctx context.Context, service *v1.Service, nodeName types.NodeName, backendPoolID string, vmSetNameOfLB string) (string, string, string, *compute.VirtualMachineScaleSetVM, error) {
 	serviceName := getServiceName(service)
 	name := mapNodeNameToVMName(nodeName)
 	vmssFlexName, err := fs.getNodeVmssFlexName(name)
@@ -464,7 +462,7 @@ func (fs *FlexScaleSet) EnsureHostInPool(service *v1.Service, nodeName types.Nod
 		return "", "", "", nil, errNotInVMSet
 	}
 
-	nic, err := fs.GetPrimaryInterface(name)
+	nic, err := fs.GetPrimaryInterface(ctx, name)
 	if err != nil {
 		klog.Errorf("error: fs.EnsureHostInPool(%s), s.GetPrimaryInterface(%s), vmSetNameOfLB: %s, err=%v", name, name, vmSetNameOfLB, err)
 		return "", "", "", nil, err
@@ -702,7 +700,7 @@ func (fs *FlexScaleSet) ensureVMSSFlexInPool(_ *v1.Service, nodes []*v1.Node, ba
 
 // EnsureHostsInPool ensures the given Node's primary IP configurations are
 // participating in the specified LoadBalancer Backend Pool.
-func (fs *FlexScaleSet) EnsureHostsInPool(service *v1.Service, nodes []*v1.Node, backendPoolID string, vmSetNameOfLB string) error {
+func (fs *FlexScaleSet) EnsureHostsInPool(ctx context.Context, service *v1.Service, nodes []*v1.Node, backendPoolID string, vmSetNameOfLB string) error {
 	mc := metrics.NewMetricContext("services", "vmssflex_ensure_hosts_in_pool", fs.ResourceGroup, fs.SubscriptionID, getServiceName(service))
 	isOperationSucceeded := false
 	defer func() {
@@ -728,7 +726,7 @@ func (fs *FlexScaleSet) EnsureHostsInPool(service *v1.Service, nodes []*v1.Node,
 		}
 
 		f := func() error {
-			_, _, _, _, err := fs.EnsureHostInPool(service, types.NodeName(localNodeName), backendPoolID, vmSetNameOfLB)
+			_, _, _, _, err := fs.EnsureHostInPool(ctx, service, types.NodeName(localNodeName), backendPoolID, vmSetNameOfLB)
 			if err != nil {
 				return fmt.Errorf("ensure(%s): backendPoolID(%s) - failed to ensure host in pool: %w", getServiceName(service), backendPoolID, err)
 			}


Test patch:
diff --git a/pkg/provider/azure_instances_test.go b/pkg/provider/azure_instances_test.go
index 49b1a7b2a..68f856575 100644
--- a/pkg/provider/azure_instances_test.go
+++ b/pkg/provider/azure_instances_test.go
@@ -968,7 +968,7 @@ func TestCloud_InstanceExists(t *testing.T) {
 			ObjectMeta: metav1.ObjectMeta{Name: "foo"},
 		}
 
-		cloud.VMSet.(*MockVMSet).EXPECT().GetInstanceIDByNodeName("foo").Return("", cloudprovider.InstanceNotFound)
+		cloud.VMSet.(*MockVMSet).EXPECT().GetInstanceIDByNodeName(gomock.Any(), "foo").Return("", cloudprovider.InstanceNotFound)
 
 		exist, err := cloud.InstanceExists(ctx, node)
 		assert.False(t, exist)
diff --git a/pkg/provider/azure_loadbalancer_backendpool_test.go b/pkg/provider/azure_loadbalancer_backendpool_test.go
index 8b30c27c9..e5727db9f 100644
--- a/pkg/provider/azure_loadbalancer_backendpool_test.go
+++ b/pkg/provider/azure_loadbalancer_backendpool_test.go
@@ -17,6 +17,7 @@ limitations under the License.
 package provider
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"strings"
@@ -433,7 +434,7 @@ func TestEnsureHostsInPoolNodeIP(t *testing.T) {
 			if tc.namespace != "" {
 				service.Namespace = tc.namespace
 			}
-			err := bi.EnsureHostsInPool(&service, nodes, "", "", "kubernetes", "kubernetes", tc.backendPool)
+			err := bi.EnsureHostsInPool(context.Background(), &service, nodes, "", "", "kubernetes", "kubernetes", tc.backendPool)
 			assert.NoError(t, err)
 			assert.Equal(t, tc.expectedBackendPool, tc.backendPool)
 		})
diff --git a/pkg/provider/azure_loadbalancer_test.go b/pkg/provider/azure_loadbalancer_test.go
index 0fb4d8a6f..706bef73d 100644
--- a/pkg/provider/azure_loadbalancer_test.go
+++ b/pkg/provider/azure_loadbalancer_test.go
@@ -738,7 +738,7 @@ func TestEnsureLoadBalancerDeleted(t *testing.T) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 	mockLBBackendPool.EXPECT().GetBackendPrivateIPs(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes()
 
 	clusterResources, expectedInterfaces, expectedVirtualMachines := getClusterResources(az, vmCount, availabilitySetCount)
@@ -3902,7 +3902,7 @@ func TestReconcileLoadBalancerCommon(t *testing.T) {
 			mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 				return false, false, lb, nil
 			}).AnyTimes()
-			mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+			mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 			lb, rerr := az.reconcileLoadBalancer(context.TODO(), "testCluster", &service, clusterResources.nodes, test.wantLb)
 			assert.Equal(t, test.expectedError, rerr)
@@ -8668,22 +8668,22 @@ func TestReconcileBackendPoolHosts(t *testing.T) {
 
 	cloud := GetTestCloud(ctrl)
 	mockLBBackendPool := NewMockBackendPool(ctrl)
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), bp1).DoAndReturn(fakeEnsureHostsInPool())
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), bp2).Return(nil)
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), bp1).DoAndReturn(fakeEnsureHostsInPool())
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), bp2).Return(nil)
 	cloud.LoadBalancerBackendPool = mockLBBackendPool
 
 	var err error
-	lb1, err = cloud.reconcileBackendPoolHosts(lb1, existingLBs, &svc, []*v1.Node{}, clusterName, "vmss", lbBackendPoolIDs)
+	lb1, err = cloud.reconcileBackendPoolHosts(context.Background(), lb1, existingLBs, &svc, []*v1.Node{}, clusterName, "vmss", lbBackendPoolIDs)
 	assert.NoError(t, err)
 	assert.Equal(t, expectedLB, lb1)
 
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(errors.New("error"))
-	_, err = cloud.reconcileBackendPoolHosts(lb1, existingLBs, &svc, []*v1.Node{}, clusterName, "vmss", lbBackendPoolIDs)
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(errors.New("error"))
+	_, err = cloud.reconcileBackendPoolHosts(context.Background(), lb1, existingLBs, &svc, []*v1.Node{}, clusterName, "vmss", lbBackendPoolIDs)
 	assert.Equal(t, errors.New("error"), err)
 }
 
-func fakeEnsureHostsInPool() func(*v1.Service, []*v1.Node, string, string, string, string, network.BackendAddressPool) error {
-	return func(_ *v1.Service, _ []*v1.Node, _, _, _, _ string, backendPool network.BackendAddressPool) error {
+func fakeEnsureHostsInPool() func(context.Context, *v1.Service, []*v1.Node, string, string, string, string, network.BackendAddressPool) error {
+	return func(_ context.Context, _ *v1.Service, _ []*v1.Node, _, _, _, _ string, backendPool network.BackendAddressPool) error {
 		backendPool.LoadBalancerBackendAddresses = &[]network.LoadBalancerBackendAddress{
 			{
 				LoadBalancerBackendAddressPropertiesFormat: &network.LoadBalancerBackendAddressPropertiesFormat{
diff --git a/pkg/provider/azure_routes_test.go b/pkg/provider/azure_routes_test.go
index 9c3692761..b95165892 100644
--- a/pkg/provider/azure_routes_test.go
+++ b/pkg/provider/azure_routes_test.go
@@ -362,8 +362,8 @@ func TestCreateRoute(t *testing.T) {
 			cloud.nodeInformerSynced = func() bool { return true }
 		}
 
-		mockVMSet.EXPECT().GetIPByNodeName(gomock.Any()).Return(nodePrivateIP, "", test.getIPError).MaxTimes(1)
-		mockVMSet.EXPECT().GetPrivateIPsByNodeName("node").Return([]string{nodePrivateIP, "10.10.10.10"}, nil).MaxTimes(1)
+		mockVMSet.EXPECT().GetIPByNodeName(gomock.Any(), gomock.Any()).Return(nodePrivateIP, "", test.getIPError).MaxTimes(1)
+		mockVMSet.EXPECT().GetPrivateIPsByNodeName(gomock.Any(), "node").Return([]string{nodePrivateIP, "10.10.10.10"}, nil).MaxTimes(1)
 		routeTableClient.EXPECT().Get(gomock.Any(), cloud.RouteTableResourceGroup, cloud.RouteTableName, "").Return(initialTable, test.getErr).MaxTimes(1)
 		routeTableClient.EXPECT().CreateOrUpdate(gomock.Any(), cloud.RouteTableResourceGroup, cloud.RouteTableName, updatedTable, "").Return(test.createOrUpdateErr).MaxTimes(1)
 
diff --git a/pkg/provider/azure_standard_test.go b/pkg/provider/azure_standard_test.go
index 17d54ee63..a3dedf815 100644
--- a/pkg/provider/azure_standard_test.go
+++ b/pkg/provider/azure_standard_test.go
@@ -17,6 +17,7 @@ limitations under the License.
 package provider
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net/http"
@@ -993,7 +994,7 @@ func TestGetStandardInstanceIDByNodeName(t *testing.T) {
 			ID:   ptr.To(invalidResouceID),
 		}, nil).AnyTimes()
 
-		instanceID, err := cloud.VMSet.GetInstanceIDByNodeName(test.nodeName)
+		instanceID, err := cloud.VMSet.GetInstanceIDByNodeName(context.Background(), test.nodeName)
 		if test.expectedErrMsg != nil {
 			assert.EqualError(t, test.expectedErrMsg, err.Error(), test.name)
 		}
@@ -1472,7 +1473,7 @@ func TestStandardEnsureHostInPool(t *testing.T) {
 		mockInterfaceClient.EXPECT().Get(gomock.Any(), cloud.ResourceGroup, test.nicName, gomock.Any()).Return(testNIC, nil).AnyTimes()
 		mockInterfaceClient.EXPECT().CreateOrUpdate(gomock.Any(), cloud.ResourceGroup, gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
-		_, _, _, vm, err := cloud.VMSet.EnsureHostInPool(test.service, test.nodeName, test.backendPoolID, test.vmSetName)
+		_, _, _, vm, err := cloud.VMSet.EnsureHostInPool(context.Background(), test.service, test.nodeName, test.backendPoolID, test.vmSetName)
 		assert.Equal(t, test.expectedErrMsg, err, test.name)
 		assert.Nil(t, vm, test.name)
 	}
@@ -1587,7 +1588,7 @@ func TestStandardEnsureHostsInPool(t *testing.T) {
 			mockInterfaceClient.EXPECT().Get(gomock.Any(), cloud.ResourceGroup, test.nicName, gomock.Any()).Return(testNIC, nil).AnyTimes()
 			mockInterfaceClient.EXPECT().CreateOrUpdate(gomock.Any(), cloud.ResourceGroup, gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
-			err := cloud.VMSet.EnsureHostsInPool(test.service, test.nodes, test.backendPoolID, test.vmSetName)
+			err := cloud.VMSet.EnsureHostsInPool(context.Background(), test.service, test.nodes, test.backendPoolID, test.vmSetName)
 			if test.expectedErr {
 				assert.EqualError(t, test.expectedErrMsg, err.Error(), test.name)
 			} else {
diff --git a/pkg/provider/azure_test.go b/pkg/provider/azure_test.go
index 83d5d5b56..df663dfc8 100644
--- a/pkg/provider/azure_test.go
+++ b/pkg/provider/azure_test.go
@@ -130,7 +130,7 @@ func TestAddPort(t *testing.T) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 	lb, err := az.reconcileLoadBalancer(context.TODO(), testClusterName, &svc, clusterResources.nodes, true /* wantLb */)
 	assert.Nil(t, err)
@@ -523,7 +523,7 @@ func testLoadBalancerServiceDefaultModeSelection(t *testing.T, isInternal bool)
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 	mockLBBackendPool.EXPECT().GetBackendPrivateIPs(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes()
 
 	expectedLBs := make([]network.LoadBalancer, 0)
@@ -594,7 +594,7 @@ func testLoadBalancerServiceAutoModeSelection(t *testing.T, isInternal bool) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 	mockLBBackendPool.EXPECT().GetBackendPrivateIPs(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes()
 
 	for index := 1; index <= serviceCount; index++ {
@@ -671,7 +671,7 @@ func testLoadBalancerServicesSpecifiedSelection(t *testing.T, isInternal bool) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 	mockLBBackendPool.EXPECT().GetBackendPrivateIPs(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes()
 
 	for index := 1; index <= serviceCount; index++ {
@@ -732,7 +732,7 @@ func testLoadBalancerMaxRulesServices(t *testing.T, isInternal bool) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 	mockLBBackendPool.EXPECT().GetBackendPrivateIPs(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes()
 
 	for index := 1; index <= az.Config.MaximumLoadBalancerRuleCount; index++ {
@@ -819,7 +819,7 @@ func testLoadBalancerServiceAutoModeDeleteSelection(t *testing.T, isInternal boo
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 	mockLBBackendPool.EXPECT().GetBackendPrivateIPs(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes()
 
 	for index := 1; index <= serviceCount; index++ {
@@ -915,7 +915,7 @@ func TestReconcileLoadBalancerAddServiceOnInternalSubnet(t *testing.T) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 	lb, err := az.reconcileLoadBalancer(context.TODO(), testClusterName, &svc, clusterResources.nodes, true /* wantLb */)
 	assert.Nil(t, err)
@@ -944,7 +944,7 @@ func TestReconcileLoadBalancerAddServicesOnMultipleSubnets(t *testing.T) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 	// svc1 is using LB without "-internal" suffix
 	lb, err := az.reconcileLoadBalancer(context.TODO(), testClusterName, &svc1, clusterResources.nodes, true /* wantLb */)
@@ -994,7 +994,7 @@ func TestReconcileLoadBalancerEditServiceSubnet(t *testing.T) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 	expectedPLS := make([]network.PrivateLinkService, 0)
 	mockPLSClient := az.PrivateLinkServiceClient.(*mockprivatelinkserviceclient.MockInterface)
@@ -1043,7 +1043,7 @@ func TestReconcileLoadBalancerNodeHealth(t *testing.T) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 	lb, err := az.reconcileLoadBalancer(context.TODO(), testClusterName, &svc, clusterResources.nodes, true /* wantLb */)
 	assert.Nil(t, err)
@@ -1072,7 +1072,7 @@ func TestReconcileLoadBalancerRemoveService(t *testing.T) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 	_, err := az.reconcileLoadBalancer(context.TODO(), testClusterName, &svc, clusterResources.nodes, true /* wantLb */)
 	assert.Nil(t, err)
@@ -1112,7 +1112,7 @@ func TestReconcileLoadBalancerRemoveAllPortsRemovesFrontendConfig(t *testing.T)
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 	lb, err := az.reconcileLoadBalancer(context.TODO(), testClusterName, &svc, clusterResources.nodes, true /* wantLb */)
 	assert.Nil(t, err)
@@ -1150,7 +1150,7 @@ func TestReconcileLoadBalancerRemovesPort(t *testing.T) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 	expectedLBs := make([]network.LoadBalancer, 0)
 	setMockLBsDualStack(az, ctrl, &expectedLBs, "service", 1, 1, false)
@@ -1186,7 +1186,7 @@ func TestReconcileLoadBalancerMultipleServices(t *testing.T) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 	expectedPLS := make([]network.PrivateLinkService, 0)
 	mockPLSClient := az.PrivateLinkServiceClient.(*mockprivatelinkserviceclient.MockInterface)
@@ -1250,7 +1250,7 @@ func TestServiceDefaultsToNoSessionPersistence(t *testing.T) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 	lb, err := az.reconcileLoadBalancer(context.TODO(), testClusterName, &svc, clusterResources.nodes, true /* wantLb */)
 	if err != nil {
@@ -1311,7 +1311,7 @@ func TestServiceRespectsNoSessionAffinity(t *testing.T) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 	lb, err := az.reconcileLoadBalancer(context.TODO(), testClusterName, &svc, clusterResources.nodes, true /* wantLb */)
 	if err != nil {
@@ -1374,7 +1374,7 @@ func TestServiceRespectsClientIPSessionAffinity(t *testing.T) {
 	mockLBBackendPool.EXPECT().ReconcileBackendPools(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ string, _ *v1.Service, lb *network.LoadBalancer) (bool, bool, *network.LoadBalancer, error) {
 		return false, false, lb, nil
 	}).AnyTimes()
-	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockLBBackendPool.EXPECT().EnsureHostsInPool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 	lb, err := az.reconcileLoadBalancer(context.TODO(), testClusterName, &svc, clusterResources.nodes, true /* wantLb */)
 	if err != nil {
@@ -1943,7 +1943,7 @@ func validatePublicIP(t *testing.T, publicIP *network.PublicIPAddress, service *
 		t.Fatal("Expected publicIP resource exists, when it is not an internal service")
 	}
 
-	if publicIP.Tags == nil || publicIP.Tags[consts.ServiceTagKey] == nil {
+	if _, ok := publicIP.Tags[consts.ServiceTagKey]; !ok {
 		t.Fatalf("Expected publicIP resource does not have tags[%s]", consts.ServiceTagKey)
 	}
 
diff --git a/pkg/provider/azure_vmss_repo_test.go b/pkg/provider/azure_vmss_repo_test.go
index 0f2c31d05..c9f48c4ce 100644
--- a/pkg/provider/azure_vmss_repo_test.go
+++ b/pkg/provider/azure_vmss_repo_test.go
@@ -17,6 +17,7 @@ limitations under the License.
 package provider
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"testing"
@@ -172,7 +173,7 @@ func TestGetPrivateIPsForMachine(t *testing.T) {
 		mockInterfaceClient := az.InterfacesClient.(*mockinterfaceclient.MockInterface)
 		mockInterfaceClient.EXPECT().Get(gomock.Any(), az.ResourceGroup, "nic", gomock.Any()).Return(expectedInterface, nil).MaxTimes(1)
 
-		privateIPs, err := az.getPrivateIPsForMachine("vm")
+		privateIPs, err := az.getPrivateIPsForMachine(context.Background(), "vm")
 		assert.Equal(t, test.expectedErr, err)
 		assert.Equal(t, test.expectedPrivateIPs, privateIPs)
 	}
@@ -247,7 +248,7 @@ func TestGetIPForMachineWithRetry(t *testing.T) {
 		mockPIPClient := az.PublicIPAddressesClient.(*mockpublicipclient.MockInterface)
 		mockPIPClient.EXPECT().List(gomock.Any(), az.ResourceGroup).Return([]network.PublicIPAddress{expectedPIP}, nil).MaxTimes(1)
 
-		privateIP, publicIP, err := az.GetIPForMachineWithRetry("vm")
+		privateIP, publicIP, err := az.GetIPForMachineWithRetry(context.Background(), "vm")
 		assert.Equal(t, test.expectedErr, err)
 		assert.Equal(t, test.expectedPrivateIP, privateIP)
 		assert.Equal(t, test.expectedPublicIP, publicIP)
diff --git a/pkg/provider/azure_vmss_test.go b/pkg/provider/azure_vmss_test.go
index 5369874e1..5b044bd11 100644
--- a/pkg/provider/azure_vmss_test.go
+++ b/pkg/provider/azure_vmss_test.go
@@ -17,6 +17,7 @@ limitations under the License.
 package provider
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"strings"
@@ -431,7 +432,7 @@ func TestGetInstanceIDByNodeName(t *testing.T) {
 		mockVMsClient := ss.VirtualMachinesClient.(*mockvmclient.MockInterface)
 		mockVMsClient.EXPECT().List(gomock.Any(), gomock.Any()).Return([]compute.VirtualMachine{}, nil).AnyTimes()
 
-		realValue, err := ss.GetInstanceIDByNodeName(test.nodeName)
+		realValue, err := ss.GetInstanceIDByNodeName(context.Background(), test.nodeName)
 		if test.expectError {
 			assert.Error(t, err, test.description)
 			continue
@@ -580,7 +581,7 @@ func TestGetIPByNodeName(t *testing.T) {
 		mockVMsClient := ss.VirtualMachinesClient.(*mockvmclient.MockInterface)
 		mockVMsClient.EXPECT().List(gomock.Any(), gomock.Any()).Return([]compute.VirtualMachine{}, nil).AnyTimes()
 
-		privateIP, publicIP, err := ss.GetIPByNodeName(test.nodeName)
+		privateIP, publicIP, err := ss.GetIPByNodeName(context.Background(), test.nodeName)
 		if test.expectError {
 			assert.Error(t, err, test.description)
 			continue
@@ -1064,7 +1065,7 @@ func TestGetInstanceTypeByNodeName(t *testing.T) {
 		mockVMClient := ss.VirtualMachinesClient.(*mockvmclient.MockInterface)
 		mockVMClient.EXPECT().List(gomock.Any(), gomock.Any()).Return(nil, test.vmClientErr).AnyTimes()
 
-		sku, err := ss.GetInstanceTypeByNodeName("vmss-vm-000000")
+		sku, err := ss.GetInstanceTypeByNodeName(context.Background(), "vmss-vm-000000")
 		if test.expectedErr != nil {
 			assert.EqualError(t, err, test.expectedErr.Error(), test.description)
 		}
@@ -1249,7 +1250,7 @@ func TestGetPrimaryInterface(t *testing.T) {
 			expectedInterface = network.Interface{}
 		}
 
-		nic, err := ss.GetPrimaryInterface(test.nodeName)
+		nic, err := ss.GetPrimaryInterface(context.Background(), test.nodeName)
 		if test.expectedErr != nil {
 			assert.EqualError(t, err, test.expectedErr.Error(), test.description)
 		}
@@ -1362,7 +1363,7 @@ func TestGetPrivateIPsByNodeName(t *testing.T) {
 		mockInterfaceClient := ss.InterfacesClient.(*mockinterfaceclient.MockInterface)
 		mockInterfaceClient.EXPECT().GetVirtualMachineScaleSetNetworkInterface(gomock.Any(), ss.ResourceGroup, testVMSSName, "0", test.nodeName, gomock.Any()).Return(expectedInterface, nil).AnyTimes()
 
-		privateIPs, err := ss.GetPrivateIPsByNodeName(test.nodeName)
+		privateIPs, err := ss.GetPrivateIPsByNodeName(context.Background(), test.nodeName)
 		if test.expectedErr != nil {
 			assert.EqualError(t, err, test.expectedErr.Error(), test.description)
 		}
@@ -2264,7 +2265,7 @@ func TestEnsureHostInPool(t *testing.T) {
 			gomock.Any(),
 		).Return(expectedVMSSVMs, test.vmssVMListError).AnyTimes()
 
-		nodeResourceGroup, ssName, instanceID, vm, err := ss.EnsureHostInPool(test.service, test.nodeName, test.backendPoolID, test.vmSetName)
+		nodeResourceGroup, ssName, instanceID, vm, err := ss.EnsureHostInPool(context.Background(), test.service, test.nodeName, test.backendPoolID, test.vmSetName)
 		assert.Equal(t, test.expectedErr, err, test.description+errMsgSuffix)
 		assert.Equal(t, test.expectedNodeResourceGroup, nodeResourceGroup, test.description)
 		assert.Equal(t, test.expectedVMSSName, ssName, test.description)
@@ -2574,7 +2575,7 @@ func TestEnsureVMSSInPool(t *testing.T) {
 
 			if test.expectedGetInstanceID != "" {
 				mockVMSet := NewMockVMSet(ctrl)
-				mockVMSet.EXPECT().GetInstanceIDByNodeName(gomock.Any()).Return(test.expectedGetInstanceID, test.getInstanceIDErr)
+				mockVMSet.EXPECT().GetInstanceIDByNodeName(gomock.Any(), gomock.Any()).Return(test.expectedGetInstanceID, test.getInstanceIDErr)
 				ss.VMSet = mockVMSet
 			}
 
@@ -2670,7 +2671,7 @@ func TestEnsureHostsInPool(t *testing.T) {
 		mockVMClient := ss.VirtualMachinesClient.(*mockvmclient.MockInterface)
 		mockVMClient.EXPECT().List(gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes()
 
-		err = ss.EnsureHostsInPool(&v1.Service{}, test.nodes, test.backendpoolID, test.vmSetName)
+		err = ss.EnsureHostsInPool(context.Background(), &v1.Service{}, test.nodes, test.backendpoolID, test.vmSetName)
 		assert.Equal(t, test.expectedErr, err != nil, test.description+errMsgSuffix)
 	}
 }
diff --git a/pkg/provider/azure_vmssflex_test.go b/pkg/provider/azure_vmssflex_test.go
index ad3cefa66..82204ecdd 100644
--- a/pkg/provider/azure_vmssflex_test.go
+++ b/pkg/provider/azure_vmssflex_test.go
@@ -17,6 +17,7 @@ limitations under the License.
 package provider
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"testing"
@@ -331,7 +332,7 @@ func TestGetInstanceIDByNodeNameVmssFlex(t *testing.T) {
 		mockVMClient.EXPECT().ListVmssFlexVMsWithoutInstanceView(gomock.Any(), gomock.Any()).Return(tc.testVMListWithoutInstanceView, tc.vmListErr).AnyTimes()
 		mockVMClient.EXPECT().ListVmssFlexVMsWithOnlyInstanceView(gomock.Any(), gomock.Any()).Return(tc.testVMListWithOnlyInstanceView, tc.vmListErr).AnyTimes()
 
-		instanceID, err := fs.GetInstanceIDByNodeName(tc.nodeName)
+		instanceID, err := fs.GetInstanceIDByNodeName(context.Background(), tc.nodeName)
 		assert.Equal(t, tc.expectedInstanceID, instanceID)
 		assert.Equal(t, tc.expectedErr, err)
 	}
@@ -381,7 +382,7 @@ func TestGetInstanceTypeByNodeNameVmssFlex(t *testing.T) {
 		mockVMClient.EXPECT().ListVmssFlexVMsWithoutInstanceView(gomock.Any(), gomock.Any()).Return(tc.testVMListWithoutInstanceView, tc.vmListErr).AnyTimes()
 		mockVMClient.EXPECT().ListVmssFlexVMsWithOnlyInstanceView(gomock.Any(), gomock.Any()).Return(tc.testVMListWithOnlyInstanceView, tc.vmListErr).AnyTimes()
 
-		instanceType, err := fs.GetInstanceTypeByNodeName(tc.nodeName)
+		instanceType, err := fs.GetInstanceTypeByNodeName(context.Background(), tc.nodeName)
 		assert.Equal(t, tc.expectedInstanceType, instanceType)
 		assert.Equal(t, tc.expectedErr, err)
 	}
@@ -646,7 +647,7 @@ func TestGetPrimaryInterfaceVmssFlex(t *testing.T) {
 		mockInterfacesClient := fs.InterfacesClient.(*mockinterfaceclient.MockInterface)
 		mockInterfacesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(tc.nic, tc.nicGetErr).AnyTimes()
 
-		nic, err := fs.GetPrimaryInterface(tc.nodeName)
+		nic, err := fs.GetPrimaryInterface(context.Background(), tc.nodeName)
 		assert.Equal(t, tc.expectedNeworkInterface, nic, tc.description)
 		if tc.expectedErr != nil {
 			assert.EqualError(t, err, tc.expectedErr.Error(), tc.description)
@@ -698,7 +699,7 @@ func TestGetIPByNodeNameVmssFlex(t *testing.T) {
 		mockInterfacesClient := fs.InterfacesClient.(*mockinterfaceclient.MockInterface)
 		mockInterfacesClient.EXPECT().Get(gomock.Any(), gomock.Any(), "testvm1-nic", gomock.Any()).Return(tc.nic, tc.nicGetErr).AnyTimes()
 
-		privateIP, publicIP, err := fs.GetIPByNodeName(tc.nodeName)
+		privateIP, publicIP, err := fs.GetIPByNodeName(context.Background(), tc.nodeName)
 		assert.Equal(t, tc.expectedPrivateIP, privateIP)
 		assert.Equal(t, tc.expectedPublicIP, publicIP)
 		assert.Equal(t, tc.expectedErr, err)
@@ -759,7 +760,7 @@ func TestGetPrivateIPsByNodeNameVmssFlex(t *testing.T) {
 		mockInterfacesClient := fs.InterfacesClient.(*mockinterfaceclient.MockInterface)
 		mockInterfacesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(tc.nic, tc.nicGetErr).AnyTimes()
 
-		ips, err := fs.GetPrivateIPsByNodeName(tc.nodeName)
+		ips, err := fs.GetPrivateIPsByNodeName(context.Background(), tc.nodeName)
 		assert.Equal(t, tc.expectedPrivateIPs, ips)
 		assert.Equal(t, tc.expectedErr, err)
 	}
@@ -1112,7 +1113,7 @@ func TestEnsureHostInPoolVmssFlex(t *testing.T) {
 		mockInterfacesClient.EXPECT().Get(gomock.Any(), gomock.Any(), "testvm1-nic", gomock.Any()).Return(tc.nic, tc.nicGetErr).AnyTimes()
 		mockInterfacesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(tc.nicPutErr).AnyTimes()
 
-		rg, vmSetName, nodeName, _, err := fs.EnsureHostInPool(tc.service, tc.nodeName, tc.backendPoolID, tc.vmSetNameOfLB)
+		rg, vmSetName, nodeName, _, err := fs.EnsureHostInPool(context.Background(), tc.service, tc.nodeName, tc.backendPoolID, tc.vmSetNameOfLB)
 		assert.Equal(t, tc.expectedNodeResourceGroup, rg, tc.description)
 		assert.Equal(t, tc.expectedVMSetName, vmSetName, tc.description)
 		assert.Equal(t, tc.expectedNodeName, nodeName, tc.description)
@@ -1378,7 +1379,7 @@ func TestEnsureHostsInPoolVmssFlex(t *testing.T) {
 		mockInterfacesClient.EXPECT().Get(gomock.Any(), gomock.Any(), "testvm1-nic", gomock.Any()).Return(tc.nic, tc.nicGetErr).AnyTimes()
 		mockInterfacesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
-		err = fs.EnsureHostsInPool(tc.service, tc.nodes, tc.backendPoolID, tc.vmSetNameOfLB)
+		err = fs.EnsureHostsInPool(context.Background(), tc.service, tc.nodes, tc.backendPoolID, tc.vmSetNameOfLB)
 
 		if tc.expectedErr != nil {
 			assert.EqualError(t, err, tc.expectedErr.Error(), tc.description)