loadshedding for arm clients

Azure · Jul 1, 2024 · 6517178 · 6517178
1 parent c8a70f5
commit 6517178
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 22 deletions.
diff --git a/cmd/kube-egress-gateway-controller/cmd/root.go b/cmd/kube-egress-gateway-controller/cmd/root.go
@@ -14,16 +14,13 @@ import (
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
 	_ "k8s.io/client-go/plugin/pkg/client/auth"
-	"sigs.k8s.io/cloud-provider-azure/pkg/azclient"
 	"sigs.k8s.io/cloud-provider-azure/pkg/azclient/configloader"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/healthz"
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
 	ctrlmetrics "sigs.k8s.io/controller-runtime/pkg/metrics"
 	metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server"
 
-	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
-
 	egressgatewayv1alpha1 "github.com/Azure/kube-egress-gateway/api/v1alpha1"
 	controllers "github.com/Azure/kube-egress-gateway/controllers/manager"
 	"github.com/Azure/kube-egress-gateway/pkg/azmanager"
@@ -152,25 +149,7 @@ func startControllers(cmd *cobra.Command, args []string) {
 		setupLog.Error(err, "cloud configuration is invalid")
 		os.Exit(1)
 	}
-	var authProvider *azclient.AuthProvider
-	authProvider, err = azclient.NewAuthProvider(
-		&cloudConfig.ARMClientConfig,
-		&cloudConfig.AzureAuthConfig)
-	if err != nil {
-		setupLog.Error(err, "unable to create auth provider")
-		os.Exit(1)
-	}
-	var cred azcore.TokenCredential
-	if cloudConfig.UseManagedIdentityExtension {
-		cred = authProvider.ManagedIdentityCredential
-	} else {
-		cred = authProvider.ClientSecretCredential
-	}
-	if cloudConfig.UserAgent == "" {
-		cloudConfig.UserAgent = consts.DefaultUserAgent
-	}
-	var factory azclient.ClientFactory
-	factory, err = azclient.NewClientFactory(&azclient.ClientFactoryConfig{SubscriptionID: cloudConfig.SubscriptionID}, &azclient.ARMClientConfig{Cloud: cloudConfig.Cloud, UserAgent: cloudConfig.UserAgent}, cred)
+	factory, err := azmanager.GetClientFactoryFromConfig(cloudConfig)
 	if err != nil {
 		setupLog.Error(err, "unable to create client factory")
 		os.Exit(1)

diff --git a/pkg/azmanager/azmanager.go b/pkg/azmanager/azmanager.go
@@ -6,6 +6,7 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	compute "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v5"
 	network "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork/v4"
 	"sigs.k8s.io/cloud-provider-azure/pkg/azclient"
@@ -42,6 +43,31 @@ type AzureManager struct {
 	SubnetClient         subnetclient.Interface
 }
 
+func GetClientFactoryFromConfig(cloud *config.CloudConfig) (azclient.ClientFactory, error) {
+	var factory azclient.ClientFactory
+	authProvider, err := azclient.NewAuthProvider(&cloud.ARMClientConfig, &cloud.AzureAuthConfig)
+	if err != nil {
+		return factory, err
+	}
+	var cred azcore.TokenCredential
+	if cloud.UseManagedIdentityExtension {
+		cred = authProvider.ManagedIdentityCredential
+	} else {
+		cred = authProvider.ClientSecretCredential
+	}
+	if cloud.UserAgent == "" {
+		cloud.UserAgent = consts.DefaultUserAgent
+	}
+	factory, err = azclient.NewClientFactory(
+		&azclient.ClientFactoryConfig{SubscriptionID: cloud.SubscriptionID},
+		&azclient.ARMClientConfig{Cloud: cloud.Cloud, UserAgent: cloud.UserAgent},
+		cred)
+	if err != nil {
+		return factory, err
+	}
+	return factory, nil
+}
+
 func CreateAzureManager(cloud *config.CloudConfig, factory azclient.ClientFactory) (*AzureManager, error) {
 	az := AzureManager{
 		CloudConfig: cloud,