Repository: Azure/aad-pod-identity
Pull request #400 (2019-10-01 01:44:37)

Node filtered listing for pod/AzureAssignedIdentity, enable pprof, and misc cache sync fixes
<!-- Thank you for helping AAD Pod Identity with a pull request! -->

**Reason for Change**:
<!-- What does this PR improve or fix in AAD Pod Identity? Why is it needed? -->

- Add ability to enable pprof profiling.
-  Add labels to AzureAssignedIdentity which indicates the nodename, pod name and pod namespace. This is used to tune the list/watches in NMI. NMI in 'enableScaleFeatures' mode will list/watch based on nodename in the lable. This reduces the payload passed between the etcd, api-server and NMI. With this we can scale to large number of nodes.     
-  Create/Delete operations for each node/VMSS based on the parallelism set by 'createDeleteBatch' parameter.
-  Parallelize the create/delete of AzureAssignedIdentity operations for individual VMSS. 
- Cache sync fixes.
- Fix to not call CacheSync in every cycle to avoid go-client thread leak issue.

**Issue Fixed**:
<!-- If this PR fixes GitHub issue 1234, add "Fixes #1234" to the next line. -->

**Notes for Reviewers**:



Patch:
diff --git a/cmd/mic/main.go b/cmd/mic/main.go
index e13a3cb1..63caf524 100644
--- a/cmd/mic/main.go
+++ b/cmd/mic/main.go
@@ -2,6 +2,8 @@ package main
 
 import (
 	"flag"
+	"net/http"
+	_ "net/http/pprof"
 	"os"
 	"time"
 
@@ -14,13 +16,17 @@ import (
 )
 
 var (
-	kubeconfig        string
-	cloudconfig       string
-	forceNamespaced   bool
-	versionInfo       bool
-	syncRetryDuration time.Duration
-	leaderElectionCfg mic.LeaderElectionConfig
-	httpProbePort     string
+	kubeconfig          string
+	cloudconfig         string
+	forceNamespaced     bool
+	versionInfo         bool
+	syncRetryDuration   time.Duration
+	leaderElectionCfg   mic.LeaderElectionConfig
+	httpProbePort       string
+	enableProfile       bool
+	enableScaleFeatures bool
+	createDeleteBatch   int64
+	clientQPS           float64
 )
 
 func main() {
@@ -44,6 +50,18 @@ func main() {
 	//Probe port
 	flag.StringVar(&httpProbePort, "http-probe-port", "8080", "http liveliness probe port")
 
+	// Profile
+	flag.BoolVar(&enableProfile, "enableProfile", false, "Enable/Disable pprof profiling")
+
+	// Enable scale features handles the label based azureassignedidentity.
+	flag.BoolVar(&enableScaleFeatures, "enableScaleFeatures", false, "Enable/Disable new features used for clusters at scale")
+
+	// createDeleteBatch can be used for tuning the number of outstanding api server operations we do per node/VMSS.
+	flag.Int64Var(&createDeleteBatch, "createDeleteBatch", 200, "Per node/VMSS create/delete batches")
+
+	// Client QPS is used to configure the client-go QPS throttling and bursting.
+	flag.Float64Var(&clientQPS, "clientQps", 10, "Client QPS used for throttling of calls to server")
+
 	flag.Parse()
 	if versionInfo {
 		version.PrintVersionAndExit()
@@ -55,6 +73,17 @@ func main() {
 	if kubeconfig == "" {
 		glog.Warningf("--kubeconfig not passed will use InClusterConfig")
 	}
+	if enableProfile {
+		profilePort := "6060"
+		glog.Infof("Starting profiling on port %s", profilePort)
+		go func() {
+			glog.Error(http.ListenAndServe("localhost:"+profilePort, nil))
+		}()
+	}
+
+	if enableScaleFeatures {
+		glog.Infof("Enabling features for scale clusters")
+	}
 
 	glog.Infof("kubeconfig (%s) cloudconfig (%s)", kubeconfig, cloudconfig)
 	config, err := buildConfig(kubeconfig)
@@ -65,7 +94,11 @@ func main() {
 
 	forceNamespaced = forceNamespaced || "true" == os.Getenv("FORCENAMESPACED")
 
-	micClient, err := mic.NewMICClient(cloudconfig, config, forceNamespaced, syncRetryDuration, &leaderElectionCfg)
+	config.QPS = float32(clientQPS)
+	config.Burst = int(clientQPS)
+	glog.Infof("Client QPS set to: %v. Burst to: %v", config.QPS, config.Burst)
+
+	micClient, err := mic.NewMICClient(cloudconfig, config, forceNamespaced, syncRetryDuration, &leaderElectionCfg, enableScaleFeatures, createDeleteBatch)
 	if err != nil {
 		glog.Fatalf("Could not get the MIC client: %+v", err)
 	}
diff --git a/cmd/nmi/main.go b/cmd/nmi/main.go
index 2e72ee8e..69856431 100644
--- a/cmd/nmi/main.go
+++ b/cmd/nmi/main.go
@@ -3,6 +3,9 @@ package main
 import (
 	"os"
 
+	"net/http"
+	_ "net/http/pprof"
+
 	"github.com/Azure/aad-pod-identity/pkg/k8s"
 	server "github.com/Azure/aad-pod-identity/pkg/nmi/server"
 	"github.com/Azure/aad-pod-identity/pkg/probes"
@@ -36,6 +39,8 @@ var (
 	retryAttemptsForCreated            = pflag.Int("retry-attempts-for-created", defaultlistPodIDsRetryAttemptsForCreated, "Number of retries in NMI to find assigned identity in CREATED state")
 	retryAttemptsForAssigned           = pflag.Int("retry-attempts-for-assigned", defaultlistPodIDsRetryAttemptsForAssigned, "Number of retries in NMI to find assigned identity in ASSIGNED state")
 	findIdentityRetryIntervalInSeconds = pflag.Int("find-identity-retry-interval", defaultlistPodIDsRetryIntervalInSeconds, "Retry interval to find assigned identities in seconds")
+	enableProfile                      = pflag.Bool("enableProfile", false, "Enable/Disable pprof profiling")
+	enableScaleFeatures                = pflag.Bool("enableScaleFeatures", false, "Enable/Disable features for scale clusters")
 )
 
 func main() {
@@ -49,9 +54,21 @@ func main() {
 		log.SetLevel(log.DebugLevel)
 	}
 	log.Infof("Starting nmi process. Version: %v. Build date: %v. Log level: %s.", version.NMIVersion, version.BuildDate, log.GetLevel())
+
+	if *enableProfile {
+		profilePort := "6060"
+		log.Infof("Starting profiling on port %s", profilePort)
+		go func() {
+			log.Error(http.ListenAndServe("localhost:"+profilePort, nil))
+		}()
+	}
+	if *enableScaleFeatures {
+		log.Infof("Features for scale clusters enabled")
+	}
+
 	logger := &server.Log{}
 
-	client, err := k8s.NewKubeClient(logger)
+	client, err := k8s.NewKubeClient(logger, *nodename, *enableScaleFeatures)
 	if err != nil {
 		log.Fatalf("%+v", err)
 	}
diff --git a/go.mod b/go.mod
index 35c77dd2..eac9dd7d 100644
--- a/go.mod
+++ b/go.mod
@@ -29,6 +29,7 @@ require (
 	github.com/sirupsen/logrus v1.2.0
 	github.com/spf13/pflag v1.0.1
 	go.opencensus.io v0.22.0 // indirect
+	golang.org/x/sync v0.0.0-20190227155943-e225da77a7e6
 	golang.org/x/time v0.0.0-20180412165947-fbb02b2291d2 // indirect
 	gopkg.in/airbrake/gobrake.v2 v2.0.9 // indirect
 	gopkg.in/gemnasium/logrus-airbrake-hook.v2 v2.1.2 // indirect
diff --git a/pkg/crd/crd.go b/pkg/crd/crd.go
index 10e7c656..cdeb15ce 100644
--- a/pkg/crd/crd.go
+++ b/pkg/crd/crd.go
@@ -18,6 +18,7 @@ import (
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/runtime/serializer"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/client-go/informers/internalinterfaces"
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/cache"
 )
@@ -35,7 +36,8 @@ type Client struct {
 // ClientInt ...
 type ClientInt interface {
 	Start(exit <-chan struct{})
-	SyncCache(exit <-chan struct{}, initial bool)
+	SyncCache(exit <-chan struct{})
+	SyncCacheLite(exit <-chan struct{})
 	RemoveAssignedIdentity(assignedIdentity *aadpodid.AzureAssignedIdentity) error
 	CreateAssignedIdentity(assignedIdentity *aadpodid.AzureAssignedIdentity) error
 	UpdateAzureAssignedIdentityStatus(assignedIdentity *aadpodid.AzureAssignedIdentity, status string) error
@@ -47,14 +49,21 @@ type ClientInt interface {
 }
 
 // NewCRDClientLite ...
-func NewCRDClientLite(config *rest.Config, log inlog.Logger) (crdClient *Client, err error) {
+func NewCRDClientLite(config *rest.Config, log inlog.Logger, nodeName string, scale bool) (crdClient *Client, err error) {
 	restClient, err := newRestClient(config)
 	if err != nil {
 		log.Error(err)
 		return nil, err
 	}
 
-	assignedIDListWatch := newAssignedIDListWatch(restClient)
+	var assignedIDListWatch *cache.ListWatch
+
+	if scale {
+		assignedIDListWatch = newAssignedIDNodeListWatch(restClient, nodeName)
+	} else {
+		assignedIDListWatch = newAssignedIDListWatch(restClient)
+	}
+
 	assignedIDListInformer, err := newAssignedIDInformer(assignedIDListWatch)
 	if err != nil {
 		log.Error(err)
@@ -202,6 +211,23 @@ func newIDInformer(r *rest.RESTClient, eventCh chan aadpodid.EventType, lw *cach
 	return azIDInformer, nil
 }
 
+// NodeNameFilter - CRDs do not yet support field selectors. Instead of that we
+// apply labels with node name and then later use the NodeNameFilter to tweak
+// options to filter using nodename label.
+func NodeNameFilter(nodeName string) internalinterfaces.TweakListOptionsFunc {
+	return func(l *v1.ListOptions) {
+		if l == nil {
+			l = &v1.ListOptions{}
+		}
+		l.LabelSelector = l.LabelSelector + "nodename=" + nodeName
+		return
+	}
+}
+
+func newAssignedIDNodeListWatch(r *rest.RESTClient, nodeName string) *cache.ListWatch {
+	return cache.NewFilteredListWatchFromClient(r, aadpodid.AzureAssignedIDResource, v1.NamespaceAll, NodeNameFilter(nodeName))
+}
+
 func newAssignedIDListWatch(r *rest.RESTClient) *cache.ListWatch {
 	return cache.NewListWatchFromClient(r, aadpodid.AzureAssignedIDResource, v1.NamespaceAll, fields.Everything())
 }
@@ -237,7 +263,7 @@ func newPodIdentityExceptionInformer(lw *cache.ListWatch) (cache.SharedInformer,
 func (c *Client) StartLite(exit <-chan struct{}) {
 	go c.AssignedIDInformer.Run(exit)
 	go c.PodIdentityExceptionInformer.Run(exit)
-	c.SyncCache(exit, true)
+	c.SyncCacheLite(exit)
 	c.log.Info("CRD lite informers started ")
 }
 
@@ -246,13 +272,24 @@ func (c *Client) Start(exit <-chan struct{}) {
 	go c.BindingInformer.Run(exit)
 	go c.IDInformer.Run(exit)
 	go c.AssignedIDInformer.Run(exit)
-	c.SyncCache(exit, true)
+	c.SyncCache(exit)
 	c.log.Info("CRD informers started")
 }
 
+func (c *Client) SyncCache(exit <-chan struct{}) {
+	c.syncCache(exit, true, c.BindingInformer.HasSynced,
+		c.IDInformer.HasSynced,
+		c.AssignedIDInformer.HasSynced)
+}
+
+func (c *Client) SyncCacheLite(exit <-chan struct{}) {
+	c.syncCache(exit, true, c.AssignedIDInformer.HasSynced,
+		c.PodIdentityExceptionInformer.HasSynced)
+}
+
 // SyncCache synchronizes cache
-func (c *Client) SyncCache(exit <-chan struct{}, initial bool) {
-	if !cache.WaitForCacheSync(exit) {
+func (c *Client) syncCache(exit <-chan struct{}, initial bool, cacheSyncs ...cache.InformerSynced) {
+	if !cache.WaitForCacheSync(exit, cacheSyncs...) {
 		if !initial {
 			c.log.Errorf("Cache could not be synchronized")
 			return
@@ -266,6 +303,7 @@ func (c *Client) RemoveAssignedIdentity(assignedIdentity *aadpodid.AzureAssigned
 	glog.V(6).Infof("Deletion of assigned id named: %s", assignedIdentity.Name)
 	begin := time.Now()
 	err := c.rest.Delete().Namespace(assignedIdentity.Namespace).Resource("azureassignedidentities").Name(assignedIdentity.Name).Do().Error()
+	glog.V(5).Infof("Deletion %s took: %v", assignedIdentity.Name, time.Since(begin))
 	stats.Update(stats.AssignedIDDel, time.Since(begin))
 	return err
 }
@@ -276,7 +314,6 @@ func (c *Client) CreateAssignedIdentity(assignedIdentity *aadpodid.AzureAssigned
 	begin := time.Now()
 	// Create a new AzureAssignedIdentity which maps the relationship between
 	// id and pod
-	glog.Infof("Creating assigned Id: %s", assignedIdentity.Name)
 	var res aadpodid.AzureAssignedIdentity
 	// TODO: Ensure that the status reflects the corresponding
 	err := c.rest.Post().Namespace(assignedIdentity.Namespace).Resource("azureassignedidentities").Body(assignedIdentity).Do().Into(&res)
@@ -285,6 +322,7 @@ func (c *Client) CreateAssignedIdentity(assignedIdentity *aadpodid.AzureAssigned
 		return err
 	}
 
+	glog.V(5).Infof("Time take to create %s: %v", assignedIdentity.Name, time.Since(begin))
 	stats.Update(stats.AssignedIDAdd, time.Since(begin))
 	//TODO: Update the status of the assign identity to indicate that the node assignment got done.
 	return nil
@@ -441,6 +479,7 @@ func (c *Client) UpdateAzureAssignedIdentityStatus(assignedIdentity *aadpodid.Az
 		return err
 	}
 
+	begin := time.Now()
 	err = c.rest.
 		Patch(types.JSONPatchType).
 		Namespace(assignedIdentity.Namespace).
@@ -449,6 +488,7 @@ func (c *Client) UpdateAzureAssignedIdentityStatus(assignedIdentity *aadpodid.Az
 		Body(patchBytes).
 		Do().
 		Error()
+	glog.V(5).Infof("Patch of %s took: %v", assignedIdentity.Name, time.Since(begin))
 
 	return err
 }
diff --git a/pkg/k8s/client.go b/pkg/k8s/client.go
index 7428fe0a..d6a59851 100644
--- a/pkg/k8s/client.go
+++ b/pkg/k8s/client.go
@@ -8,9 +8,9 @@ import (
 	"time"
 
 	v1 "k8s.io/api/core/v1"
-	"k8s.io/client-go/informers"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/rest"
+	"k8s.io/client-go/tools/cache"
 	"k8s.io/client-go/tools/clientcmd"
 
 	aadpodid "github.com/Azure/aad-pod-identity/pkg/apis/aadpodidentity/v1"
@@ -20,8 +20,8 @@ import (
 	"github.com/golang/glog"
 	log "github.com/sirupsen/logrus"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/apimachinery/pkg/labels"
 	informersv1 "k8s.io/client-go/informers/core/v1"
+	"k8s.io/client-go/informers/internalinterfaces"
 )
 
 const (
@@ -49,12 +49,12 @@ type KubeClient struct {
 	ClientSet kubernetes.Interface
 	// Crd client used to access our CRD resources.
 	CrdClient   *crd.Client
-	PodInformer informersv1.PodInformer
+	PodInformer cache.SharedIndexInformer
 	log         inlog.Logger
 }
 
 // NewKubeClient new kubernetes api client
-func NewKubeClient(log inlog.Logger) (Client, error) {
+func NewKubeClient(log inlog.Logger, nodeName string, scale bool) (Client, error) {
 	config, err := buildConfig()
 	if err != nil {
 		return nil, err
@@ -64,13 +64,14 @@ func NewKubeClient(log inlog.Logger) (Client, error) {
 	if err != nil {
 		return nil, err
 	}
-	crdclient, err := crd.NewCRDClientLite(config, log)
+	crdclient, err := crd.NewCRDClientLite(config, log, nodeName, scale)
 	if err != nil {
 		return nil, err
 	}
 
-	informer := informers.NewSharedInformerFactory(clientset, 30*time.Second)
-	podInformer := informer.Core().V1().Pods()
+	podInformer := informersv1.NewFilteredPodInformer(clientset, v1.NamespaceAll, 10*time.Minute,
+		cache.Indexers{cache.NamespaceIndex: cache.MetaNamespaceIndexFunc},
+		NodeNameFilter(nodeName))
 
 	kubeClient := &KubeClient{
 		CrdClient:   crdclient,
@@ -82,11 +83,18 @@ func NewKubeClient(log inlog.Logger) (Client, error) {
 	return kubeClient, nil
 }
 
+func (c *KubeClient) Sync(exit <-chan struct{}) {
+	if !cache.WaitForCacheSync(exit, c.PodInformer.HasSynced) {
+		c.log.Errorf("Pod cache could not be synchronized")
+	}
+	c.CrdClient.SyncCacheLite(exit)
+}
+
 // Start the corresponding starts
 func (c *KubeClient) Start(exit <-chan struct{}) {
-	go c.PodInformer.Informer().Run(exit)
+	go c.PodInformer.Run(exit)
 	c.CrdClient.StartLite(exit)
-	c.CrdClient.SyncCache(exit, true)
+	c.Sync(exit)
 }
 
 func (c *KubeClient) getReplicasetName(pod v1.Pod) string {
@@ -98,6 +106,18 @@ func (c *KubeClient) getReplicasetName(pod v1.Pod) string {
 	return ""
 }
 
+// NodeNameFilter will tweak the options to include the node name as field
+// selector.
+func NodeNameFilter(nodeName string) internalinterfaces.TweakListOptionsFunc {
+	return func(l *metav1.ListOptions) {
+		if l == nil {
+			l = &metav1.ListOptions{}
+		}
+		l.FieldSelector = l.FieldSelector + "spec.nodeName=" + nodeName
+		return
+	}
+}
+
 // GetPodInfo get pod ns,name from apiserver
 func (c *KubeClient) GetPodInfo(podip string) (podns, poddname, rsName string, labels *metav1.LabelSelector, err error) {
 	if podip == "" {
@@ -123,14 +143,15 @@ func isPhaseValid(p v1.PodPhase) bool {
 }
 
 func (c *KubeClient) getPodList(podip string) ([]*v1.Pod, error) {
-	list, err := c.PodInformer.Lister().List(labels.Everything())
-	if err != nil {
-		glog.Error(err)
-		return nil, err
-	}
-
 	var podList []*v1.Pod
-	for _, pod := range list {
+	list := c.PodInformer.GetStore().List()
+	for _, o := range list {
+		pod, ok := o.(*v1.Pod)
+		if !ok {
+			err := fmt.Errorf("could not cast %T to %s", pod, "v1.Pod")
+			glog.Error(err)
+			return nil, err
+		}
 		if pod.Status.PodIP == podip && isPhaseValid(pod.Status.Phase) {
 			podList = append(podList, pod)
 		}
diff --git a/pkg/mic/mic.go b/pkg/mic/mic.go
index b22a51b4..73d82a34 100644
--- a/pkg/mic/mic.go
+++ b/pkg/mic/mic.go
@@ -1,6 +1,7 @@
 package mic
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"strings"
@@ -15,6 +16,7 @@ import (
 	"github.com/Azure/aad-pod-identity/pkg/stats"
 	"github.com/Azure/aad-pod-identity/version"
 	"github.com/golang/glog"
+	"golang.org/x/sync/semaphore"
 	corev1 "k8s.io/api/core/v1"
 	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/informers"
@@ -49,15 +51,17 @@ type LeaderElectionConfig struct {
 // Client has the required pointers to talk to the api server
 // and interact with the CRD related datastructure.
 type Client struct {
-	CRDClient         crd.ClientInt
-	CloudClient       cloudprovider.ClientInt
-	PodClient         pod.ClientInt
-	EventRecorder     record.EventRecorder
-	EventChannel      chan aadpodid.EventType
-	NodeClient        NodeGetter
-	IsNamespaced      bool
-	SyncLoopStarted   bool
-	syncRetryInterval time.Duration
+	CRDClient           crd.ClientInt
+	CloudClient         cloudprovider.ClientInt
+	PodClient           pod.ClientInt
+	EventRecorder       record.EventRecorder
+	EventChannel        chan aadpodid.EventType
+	NodeClient          NodeGetter
+	IsNamespaced        bool
+	SyncLoopStarted     bool
+	syncRetryInterval   time.Duration
+	enableScaleFeatures bool
+	createDeleteBatch   int64
 
 	syncing int32 // protect against conucrrent sync's
 
@@ -80,7 +84,8 @@ type trackUserAssignedMSIIds struct {
 }
 
 // NewMICClient returnes new mic client
-func NewMICClient(cloudconfig string, config *rest.Config, isNamespaced bool, syncRetryInterval time.Duration, leaderElectionConfig *LeaderElectionConfig) (*Client, error) {
+func NewMICClient(cloudconfig string, config *rest.Config, isNamespaced bool, syncRetryInterval time.Duration,
+	leaderElectionConfig *LeaderElectionConfig, enableScaleFeatures bool, createDeleteBatch int64) (*Client, error) {
 	glog.Infof("Starting to create the pod identity client. Version: %v. Build date: %v", version.MICVersion, version.BuildDate)
 
 	clientSet := kubernetes.NewForConfigOrDie(config)
@@ -109,19 +114,20 @@ func NewMICClient(cloudconfig string, config *rest.Config, isNamespaced bool, sy
 	glog.V(1).Infof("Pod Client initialized")
 
 	eventBroadcaster := record.NewBroadcaster()
-	eventBroadcaster.StartLogging(glog.Infof)
 	eventBroadcaster.StartRecordingToSink(&typedcorev1.EventSinkImpl{Interface: clientSet.CoreV1().Events("")})
 	recorder := eventBroadcaster.NewRecorder(scheme.Scheme, corev1.EventSource{Component: aadpodid.CRDGroup})
 
 	c := &Client{
-		CRDClient:         crdClient,
-		CloudClient:       cloudClient,
-		PodClient:         podClient,
-		EventRecorder:     recorder,
-		EventChannel:      eventCh,
-		NodeClient:        &NodeClient{informer.Core().V1().Nodes()},
-		IsNamespaced:      isNamespaced,
-		syncRetryInterval: syncRetryInterval,
+		CRDClient:           crdClient,
+		CloudClient:         cloudClient,
+		PodClient:           podClient,
+		EventRecorder:       recorder,
+		EventChannel:        eventCh,
+		NodeClient:          &NodeClient{informer.Core().V1().Nodes()},
+		IsNamespaced:        isNamespaced,
+		syncRetryInterval:   syncRetryInterval,
+		enableScaleFeatures: enableScaleFeatures,
+		createDeleteBatch:   createDeleteBatch,
 	}
 	leaderElector, err := c.NewLeaderElector(clientSet, recorder, leaderElectionConfig)
 	if err != nil {
@@ -229,6 +235,8 @@ func (c *Client) Sync(exit <-chan struct{}) {
 	glog.Info("Sync thread started.")
 	c.SyncLoopStarted = true
 	var event aadpodid.EventType
+	totalWorkDoneCycles := 0
+	totalSyncCycles := 0
 	for {
 		select {
 		case <-exit:
@@ -238,19 +246,12 @@ func (c *Client) Sync(exit <-chan struct{}) {
 		case <-ticker.C:
 			glog.V(6).Infof("Running periodic sync loop")
 		}
-
+		totalSyncCycles++
 		stats.Init()
 		// This is the only place where the AzureAssignedIdentity creation is initiated.
 		begin := time.Now()
 		workDone := false
 
-		cacheTime := time.Now()
-
-		// There is a delay in data propogation to cache. It's possible that the creates performed in the previous sync cycle
-		// are not propogated before this sync cycle began. In order to avoid redoing the cycle, we sync cache again.
-		c.CRDClient.SyncCache(exit, false)
-		stats.Put(stats.CacheSync, time.Since(cacheTime))
-
 		// List all pods in all namespaces
 		systemTime := time.Now()
 		listPods, err := c.PodClient.GetPods()
@@ -328,7 +329,10 @@ func (c *Client) Sync(exit <-chan struct{}) {
 
 		wg.Wait()
 
-		if workDone {
+		if workDone || ((totalSyncCycles % 1000) == 0) {
+			if workDone {
+				totalWorkDoneCycles++
+			}
 			idsFound := 0
 			bindingsFound := 0
 			if listIDs != nil {
@@ -337,9 +341,16 @@ func (c *Client) Sync(exit <-chan struct{}) {
 			if listBindings != nil {
 				bindingsFound = len(*listBindings)
 			}
-			glog.Infof("Found %d pods, %d ids, %d bindings", len(listPods), idsFound, bindingsFound)
+			glog.Infof("Work done: %v. Found %d pods, %d ids, %d bindings", workDone, len(listPods), idsFound, bindingsFound)
+			glog.Infof("Total work cycles: %d, out of which work was done in: %d.", totalSyncCycles, totalWorkDoneCycles)
 			stats.Put(stats.Total, time.Since(begin))
 			stats.PrintSync()
+			if workDone {
+				// We need to synchornize the cache inorder to get the latest updates. Sync cache has a bug in the current go client which caused thread leak.
+				// Updating of go client has issues with case sensitivity. Avoid this issue by sleping for 500 milliseconds to reduce the chance
+				// of cache misses for assignedidentities updated in the previous cycle.
+				time.Sleep(time.Millisecond * 200)
+			}
 		}
 	}
 }
@@ -480,14 +491,14 @@ func (c *Client) matchAssignedID(x *aadpodid.AzureAssignedIdentity, y *aadpodid.
 	idX := x.Spec.AzureIdentityRef
 	idY := y.Spec.AzureIdentityRef
 
-	glog.V(6).Infof("assignedidX - %+v\n", x)
-	glog.V(6).Infof("assignedidY - %+v\n", y)
+	glog.V(7).Infof("assignedidX - %+v\n", x)
+	glog.V(7).Infof("assignedidY - %+v\n", y)
 
-	glog.V(6).Infof("bindingX - %+v\n", bindingX)
-	glog.V(6).Infof("bindingY - %+v\n", bindingY)
+	glog.V(7).Infof("bindingX - %+v\n", bindingX)
+	glog.V(7).Infof("bindingY - %+v\n", bindingY)
 
-	glog.V(6).Infof("idX - %+v\n", idX)
-	glog.V(6).Infof("idY - %+v\n", idY)
+	glog.V(7).Infof("idX - %+v\n", idX)
+	glog.V(7).Infof("idY - %+v\n", idY)
 
 	if bindingX.Name == bindingY.Name && bindingX.ResourceVersion == bindingY.ResourceVersion &&
 		idX.Name == idY.Name && idX.ResourceVersion == idY.ResourceVersion &&
@@ -580,10 +591,17 @@ func (c *Client) makeAssignedIDs(azID aadpodid.AzureIdentity, azBinding aadpodid
 	binding := azBinding
 	id := azID
 
+	labels := make(map[string]string)
+	labels["nodename"] = nodeName
+	labels["podnamespace"] = podNameSpace
+	labels["podname"] = podName
+
+	oMeta := v1.ObjectMeta{
+		Name:   c.getAssignedIDName(podName, podNameSpace, azID.Name),
+		Labels: labels,
+	}
 	assignedID := &aadpodid.AzureAssignedIdentity{
-		ObjectMeta: v1.ObjectMeta{
-			Name: c.getAssignedIDName(podName, podNameSpace, azID.Name),
-		},
+		ObjectMeta: oMeta,
 		Spec: aadpodid.AzureAssignedIdentitySpec{
 			AzureIdentityRef: &id,
 			AzureBindingRef:  &binding,
@@ -747,21 +765,41 @@ func (c *Client) updateUserMSI(newAssignedIDs []aadpodid.AzureAssignedIdentity,
 	beginAdding := time.Now()
 	glog.Infof("Processing node %s, add [%d], del [%d]", nodeOrVMSSName, len(nodeTrackList.assignedIDsToCreate), len(nodeTrackList.assignedIDsToDelete))
 
+	ctx := context.TODO()
+	// We have to ensure that we don't overwhelm the API server with too many
+	// requests in flight. We use a token based approach implemented using semaphore to
+	// ensure that only given createDeleteBatch requests are in flight at any point in time.
+	// Note that at this point in the code path, we are doing this in parallel per node/VMSS already.
+	semCreate := semaphore.NewWeighted(c.createDeleteBatch)
+
 	for _, createID := range nodeTrackList.assignedIDsToCreate {
-		if createID.Status.Status == "" {
-			binding := createID.Spec.AzureBindingRef
+		if err := semCreate.Acquire(ctx, 1); err != nil {
+			glog.Errorf("Failed to acquire semaphore in the create loop: %v", err)
+			return
+		}
+		go func(assignedID aadpodid.AzureAssignedIdentity) {
+			defer semCreate.Release(1)
+			if assignedID.Status.Status == "" {
+				binding := assignedID.Spec.AzureBindingRef
 
-			// this is the state when the azure assigned identity is yet to be created
-			glog.V(5).Infof("Initiating assigned id creation for pod - %s, binding - %s", createID.Spec.Pod, binding.Name)
+				// this is the state when the azure assigned identity is yet to be created
+				glog.V(5).Infof("Initiating assigned id creation for pod - %s, binding - %s", assignedID.Spec.Pod, binding.Name)
 
-			createID.Status.Status = aadpodid.AssignedIDCreated
-			err := c.createAssignedIdentity(&createID)
-			if err != nil {
-				c.EventRecorder.Event(binding, corev1.EventTypeWarning, "binding apply error",
-					fmt.Sprintf("Creating assigned identity for pod %s resulted in error %v", createID.Name, err))
-				glog.Error(err)
+				assignedID.Status.Status = aadpodid.AssignedIDCreated
+				err := c.createAssignedIdentity(&assignedID)
+				if err != nil {
+					c.EventRecorder.Event(binding, corev1.EventTypeWarning, "binding apply error",
+						fmt.Sprintf("Creating assigned identity for pod %s resulted in error %v", assignedID.Name, err))
+					glog.Error(err)
+				}
 			}
-		}
+		}(createID)
+	}
+
+	// Ensure that all creates are complete
+	if err := semCreate.Acquire(ctx, c.createDeleteBatch); err != nil {
+		glog.Errorf("Failed to acquire semaphore at the end of creates: %v", err)
+		return
 	}
 	// generate unique list so we don't make multiple calls to assign/remove same id
 	addUserAssignedMSIIDs := c.getUniqueIDs(nodeTrackList.addUserAssignedMSIIDs)
@@ -844,44 +882,74 @@ func (c *Client) updateUserMSI(newAssignedIDs []aadpodid.AzureAssignedIdentity,
 		return
 	}
 
+	semUpdate := semaphore.NewWeighted(c.createDeleteBatch)
+
 	for _, createID := range nodeTrackList.assignedIDsToCreate {
-		binding := createID.Spec.AzureBindingRef
-		// update the status to assigned for assigned identity as identity was successfully assigned to node.
-		err = c.updateAssignedIdentityStatus(&createID, aadpodid.AssignedIDAssigned)
-		if err != nil {
-			message := fmt.Sprintf("Updating assigned identity %s status to %s for pod %s failed with error %v", createID.Name, aadpodid.AssignedIDAssigned, createID.Spec.Pod, err.Error())
-			c.EventRecorder.Event(&createID, corev1.EventTypeWarning, "status update error", message)
-			glog.Error(message)
-			continue
+		if err := semUpdate.Acquire(ctx, 1); err != nil {
+			glog.Errorf("Failed to acquire semaphore in the update loop: %v", err)
+			return
 		}
-		c.EventRecorder.Event(binding, corev1.EventTypeNormal, "binding applied",
-			fmt.Sprintf("Binding %s applied on node %s for pod %s", binding.Name, createID.Spec.NodeName, createID.Name))
+		go func(assignedID aadpodid.AzureAssignedIdentity) {
+			defer semUpdate.Release(1)
+			binding := assignedID.Spec.AzureBindingRef
+			// update the status to assigned for assigned identity as identity was successfully assigned to node.
+			err := c.updateAssignedIdentityStatus(&assignedID, aadpodid.AssignedIDAssigned)
+			if err != nil {
+				message := fmt.Sprintf("Updating assigned identity %s status to %s for pod %s failed with error %v", assignedID.Name, aadpodid.AssignedIDAssigned, assignedID.Spec.Pod, err.Error())
+				c.EventRecorder.Event(&assignedID, corev1.EventTypeWarning, "status update error", message)
+				glog.Error(message)
+				return
+			}
+			c.EventRecorder.Event(binding, corev1.EventTypeNormal, "binding applied",
+				fmt.Sprintf("Binding %s applied on node %s for pod %s", binding.Name, assignedID.Spec.NodeName, assignedID.Name))
+		}(createID)
 	}
 
-	for _, delID := range nodeTrackList.assignedIDsToDelete {
-		removedBinding := delID.Spec.AzureBindingRef
+	// Ensure that all updates are complete
+	if err := semUpdate.Acquire(ctx, c.createDeleteBatch); err != nil {
+		glog.Errorf("Failed to acquire semaphore at the end of updates: %v", err)
+		return
+	}
 
-		// update the status for the assigned identity to Unassigned as the identity has been successfully removed from node.
-		// this will ensure on next sync loop we only try to delete the assigned identity instead of doing everything.
-		err = c.updateAssignedIdentityStatus(&delID, aadpodid.AssignedIDUnAssigned)
-		if err != nil {
-			message := fmt.Sprintf("Updating assigned identity %s status to %s for pod %s failed with error %v", delID.Name, aadpodid.AssignedIDUnAssigned, delID.Spec.Pod, err.Error())
-			c.EventRecorder.Event(&delID, corev1.EventTypeWarning, "status update error", message)
-			glog.Error(message)
-			continue
-		}
-		// remove assigned identity crd from cluster as the identity has successfully been removed from the node
-		err = c.removeAssignedIdentity(&delID)
-		if err != nil {
-			c.EventRecorder.Event(removedBinding, corev1.EventTypeWarning, "binding remove error",
-				fmt.Sprintf("Removing assigned identity binding %s node %s for pod %s resulted in error %v", removedBinding.Name, delID.Spec.NodeName, delID.Name, err))
-			glog.Error(err)
-			continue
+	semDel := semaphore.NewWeighted(c.createDeleteBatch)
+
+	for _, delID := range nodeTrackList.assignedIDsToDelete {
+		if err := semDel.Acquire(ctx, 1); err != nil {
+			glog.Errorf("Failed to acquire semaphore in the delete loop: %v", err)
+			return
 		}
-		// the identity was successfully removed from node
-		c.EventRecorder.Event(removedBinding, corev1.EventTypeNormal, "binding removed",
-			fmt.Sprintf("Binding %s removed from node %s for pod %s", removedBinding.Name, delID.Spec.NodeName, delID.Spec.Pod))
+		go func(assignedID aadpodid.AzureAssignedIdentity) {
+			defer semDel.Release(1)
+			removedBinding := assignedID.Spec.AzureBindingRef
+			// update the status for the assigned identity to Unassigned as the identity has been successfully removed from node.
+			// this will ensure on next sync loop we only try to delete the assigned identity instead of doing everything.
+			err := c.updateAssignedIdentityStatus(&assignedID, aadpodid.AssignedIDUnAssigned)
+			if err != nil {
+				message := fmt.Sprintf("Updating assigned identity %s status to %s for pod %s failed with error %v", assignedID.Name, aadpodid.AssignedIDUnAssigned, assignedID.Spec.Pod, err.Error())
+				c.EventRecorder.Event(&assignedID, corev1.EventTypeWarning, "status update error", message)
+				glog.Error(message)
+				return
+			}
+			// remove assigned identity crd from cluster as the identity has successfully been removed from the node
+			err = c.removeAssignedIdentity(&assignedID)
+			if err != nil {
+				c.EventRecorder.Event(removedBinding, corev1.EventTypeWarning, "binding remove error",
+					fmt.Sprintf("Removing assigned identity binding %s node %s for pod %s resulted in error %v", removedBinding.Name, assignedID.Spec.NodeName, assignedID.Name, err))
+				glog.Error(err)
+				return
+			}
+			// the identity was successfully removed from node
+			c.EventRecorder.Event(removedBinding, corev1.EventTypeNormal, "binding removed",
+				fmt.Sprintf("Binding %s removed from node %s for pod %s", removedBinding.Name, assignedID.Spec.NodeName, assignedID.Spec.Pod))
+		}(delID)
+	}
+
+	// Ensure that all deletes are complete
+	if err := semDel.Acquire(ctx, c.createDeleteBatch); err != nil {
+		glog.Errorf("Failed to acquire semaphore at the end of deletes: %v", err)
+		return
 	}
+
 	stats.Put(stats.TotalCreateOrUpdate, time.Since(beginAdding))
 }
 


Test patch:
diff --git a/pkg/mic/mic_test.go b/pkg/mic/mic_test.go
index b8940018..8c41ab32 100644
--- a/pkg/mic/mic_test.go
+++ b/pkg/mic/mic_test.go
@@ -345,7 +345,11 @@ func NewTestCrdClient(config *rest.Config) *TestCrdClient {
 func (c *TestCrdClient) Start(exit <-chan struct{}) {
 }
 
-func (c *TestCrdClient) SyncCache(exit <-chan struct{}, initial bool) {
+func (c *TestCrdClient) SyncCache(exit <-chan struct{}) {
+
+}
+
+func (c *TestCrdClient) SyncCacheLite(exit <-chan struct{}) {
 
 }
 
@@ -586,7 +590,8 @@ func NewMICTestClient(eventCh chan aadpodid.EventType,
 	crdClient *TestCrdClient,
 	podClient *TestPodClient,
 	nodeClient *TestNodeClient,
-	eventRecorder *TestEventRecorder, isNamespaced bool) *TestMICClient {
+	eventRecorder *TestEventRecorder, isNamespaced bool,
+	createDeleteBatch int64) *TestMICClient {
 
 	realMICClient := &Client{
 		CloudClient:       cpClient,
@@ -597,6 +602,7 @@ func NewMICTestClient(eventCh chan aadpodid.EventType,
 		NodeClient:        nodeClient,
 		syncRetryInterval: 120 * time.Second,
 		IsNamespaced:      isNamespaced,
+		createDeleteBatch: createDeleteBatch,
 	}
 
 	return &TestMICClient{
@@ -703,7 +709,7 @@ func TestSimpleMICClient(t *testing.T) {
 	evtRecorder.lastEvent = new(LastEvent)
 	evtRecorder.eventChannel = make(chan bool, 100)
 
-	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false)
+	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false, 4)
 
 	crdClient.CreateID("test-id", "default", aadpodid.UserAssignedMSI, "test-user-msi-resourceid", "test-user-msi-clientid", nil, "", "", "", "")
 	crdClient.CreateBinding("testbinding", "default", "test-id", "test-select", "")
@@ -824,7 +830,7 @@ func TestAddDelMICClient(t *testing.T) {
 	evtRecorder.lastEvent = new(LastEvent)
 	evtRecorder.eventChannel = make(chan bool, 100)
 
-	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false)
+	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false, 4)
 
 	// Test to add and delete at the same time.
 	// Add a pod, identity and binding.
@@ -916,7 +922,7 @@ func TestMicAddDelVMSS(t *testing.T) {
 	evtRecorder.lastEvent = new(LastEvent)
 	evtRecorder.eventChannel = make(chan bool, 100)
 
-	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false)
+	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false, 4)
 
 	// Test to add and delete at the same time.
 	// Add a pod, identity and binding.
@@ -1019,7 +1025,7 @@ func TestMICStateFlow(t *testing.T) {
 	evtRecorder.lastEvent = new(LastEvent)
 	evtRecorder.eventChannel = make(chan bool, 100)
 
-	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false)
+	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false, 4)
 
 	// Add a pod, identity and binding.
 	crdClient.CreateID("test-id1", "default", aadpodid.UserAssignedMSI, "test-user-msi-resourceid", "test-user-msi-clientid", nil, "", "", "", "")
@@ -1133,7 +1139,7 @@ func TestForceNamespaced(t *testing.T) {
 	evtRecorder.lastEvent = new(LastEvent)
 	evtRecorder.eventChannel = make(chan bool, 100)
 
-	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, true)
+	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, true, 4)
 
 	crdClient.CreateID("test-id1", "default", aadpodid.UserAssignedMSI, "test-user-msi-resourceid", "test-user-msi-clientid", nil, "", "", "", "idrv1")
 	crdClient.CreateBinding("testbinding1", "default", "test-id1", "test-select1", "bindingrv1")
@@ -1197,7 +1203,7 @@ func TestSyncRetryLoop(t *testing.T) {
 	evtRecorder.lastEvent = new(LastEvent)
 	evtRecorder.eventChannel = make(chan bool, 100)
 
-	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false)
+	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false, 4)
 	syncRetryInterval, err := time.ParseDuration("10s")
 	if err != nil {
 		t.Errorf("error parsing duration: %v", err)
@@ -1276,7 +1282,7 @@ func TestSyncNodeNotFound(t *testing.T) {
 	evtRecorder.lastEvent = new(LastEvent)
 	evtRecorder.eventChannel = make(chan bool, 100)
 
-	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false)
+	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false, 4)
 
 	// Add a pod, identity and binding.
 	crdClient.CreateID("test-id1", "default", aadpodid.UserAssignedMSI, "test-user-msi-resourceid", "test-user-msi-clientid", nil, "", "", "", "")
@@ -1348,7 +1354,7 @@ func TestSyncExit(t *testing.T) {
 	evtRecorder.lastEvent = new(LastEvent)
 	evtRecorder.eventChannel = make(chan bool)
 
-	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false)
+	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false, 4)
 
 	micClient.testRunSync()(t)
 }
diff --git a/test/common/k8s/infra/infra.go b/test/common/k8s/infra/infra.go
index ccd783db..d322c1e7 100644
--- a/test/common/k8s/infra/infra.go
+++ b/test/common/k8s/infra/infra.go
@@ -11,7 +11,7 @@ import (
 )
 
 // CreateInfra will deploy all the infrastructure components (nmi and mic) on a Kubernetes cluster
-func CreateInfra(namespace, registry, nmiVersion, micVersion, templateOutputPath string, old bool) error {
+func CreateInfra(namespace, registry, nmiVersion, micVersion, templateOutputPath string, old, enableScaleFeatures bool) error {
 	var err error
 	var t *template.Template
 	if !old {
@@ -37,12 +37,13 @@ func CreateInfra(namespace, registry, nmiVersion, micVersion, templateOutputPath
 	nmiArg = nmiVersion == "1.4"
 
 	deployData := struct {
-		Namespace  string
-		Registry   string
-		NMIVersion string
-		MICVersion string
-		MICArg     bool
-		NMIArg     bool
+		Namespace           string
+		Registry            string
+		NMIVersion          string
+		MICVersion          string
+		MICArg              bool
+		NMIArg              bool
+		EnableScaleFeatures bool
 	}{
 		namespace,
 		registry,
@@ -50,6 +51,7 @@ func CreateInfra(namespace, registry, nmiVersion, micVersion, templateOutputPath
 		micVersion,
 		micArg,
 		nmiArg,
+		enableScaleFeatures,
 	}
 	if err := t.Execute(deployFile, deployData); err != nil {
 		return errors.Wrap(err, "Failed to create a deployment file from deployment-rbac.yaml")
diff --git a/test/e2e/aadpodidentity_test.go b/test/e2e/aadpodidentity_test.go
index 5fb9ce3f..b8a08991 100644
--- a/test/e2e/aadpodidentity_test.go
+++ b/test/e2e/aadpodidentity_test.go
@@ -60,7 +60,7 @@ var _ = BeforeSuite(func() {
 	Expect(err).NotTo(HaveOccurred())
 	cfg = *c
 	fmt.Printf("System MSI enabled: %v\n", cfg.SystemMSICluster)
-	setupInfra(cfg.Registry, cfg.NMIVersion, cfg.MICVersion)
+	setupInfra(cfg.Registry, cfg.NMIVersion, cfg.MICVersion, cfg.EnableScaleFeatures)
 })
 
 var _ = AfterSuite(func() {
@@ -498,7 +498,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		}
 
 		// reset the infra to previous state
-		setupInfra(cfg.Registry, cfg.NMIVersion, cfg.MICVersion)
+		setupInfra(cfg.Registry, cfg.NMIVersion, cfg.MICVersion, cfg.EnableScaleFeatures)
 	})
 
 	It("should not alter the system assigned identity after creating and deleting pod identity", func() {
@@ -603,7 +603,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		Expect(ok).To(Equal(true))
 
 		// update the infra to use latest mic and nmi images
-		setupInfra(cfg.Registry, cfg.NMIVersion, cfg.MICVersion)
+		setupInfra(cfg.Registry, cfg.NMIVersion, cfg.MICVersion, cfg.EnableScaleFeatures)
 
 		ok, err = daemonset.WaitOnReady(nmiDaemonSet)
 		Expect(err).NotTo(HaveOccurred())
@@ -1167,15 +1167,15 @@ func checkInfra() {
 // setupInfra creates the crds, mic, nmi and blocks until iptable entries exist
 func setupInfraOld(registry, nmiVersion, micVersion string) {
 	// Install CRDs and deploy MIC and NMI
-	err := infra.CreateInfra("default", registry, nmiVersion, micVersion, templateOutputPath, true)
+	err := infra.CreateInfra("default", registry, nmiVersion, micVersion, templateOutputPath, true, false)
 	Expect(err).NotTo(HaveOccurred())
 	checkInfra()
 }
 
 // setupInfra creates the crds, mic, nmi and blocks until iptable entries exist
-func setupInfra(registry, nmiVersion, micVersion string) {
+func setupInfra(registry, nmiVersion, micVersion string, enableScaleFeatures bool) {
 	// Install CRDs and deploy MIC and NMI
-	err := infra.CreateInfra("default", registry, nmiVersion, micVersion, templateOutputPath, false)
+	err := infra.CreateInfra("default", registry, nmiVersion, micVersion, templateOutputPath, false, enableScaleFeatures)
 	Expect(err).NotTo(HaveOccurred())
 	checkInfra()
 }
diff --git a/test/e2e/config/config.go b/test/e2e/config/config.go
index 394e9521..3a0c05dc 100644
--- a/test/e2e/config/config.go
+++ b/test/e2e/config/config.go
@@ -15,6 +15,7 @@ type Config struct {
 	Registry                 string `envconfig:"REGISTRY" default:"mcr.microsoft.com/k8s/aad-pod-identity"`
 	IdentityValidatorVersion string `envconfig:"IDENTITY_VALIDATOR_VERSION" default:"1.5.2"`
 	SystemMSICluster         bool   `envconfig:"SYSTEM_MSI_CLUSTER" default:"false"`
+	EnableScaleFeatures      bool   `envconfig:"ENABLE_SCALE_FEATURES" default:"false"`
 }
 
 // ParseConfig will parse needed environment variables for running the tests
diff --git a/test/e2e/template/deployment-rbac.yaml b/test/e2e/template/deployment-rbac.yaml
index 7109b133..e35e83f7 100644
--- a/test/e2e/template/deployment-rbac.yaml
+++ b/test/e2e/template/deployment-rbac.yaml
@@ -116,6 +116,7 @@ spec:
           {{if .NMIArg}}- nmi {{ end }}
           - "--host-ip=$(HOST_IP)"
           - "--node=$(NODE_NAME)"
+          {{if .EnableScaleFeatures}}- "--enableScaleFeatures=true" {{end}}
         env:
           - name: HOST_IP
             valueFrom:
@@ -215,6 +216,7 @@ spec:
           {{if .MICArg}}- mic {{ end }}
           - "--cloudconfig=/etc/kubernetes/azure.json"
           - "--logtostderr"
+          {{if .EnableScaleFeatures}}- "--enableScaleFeatures=true" {{end}}
         resources:
           limits:
             cpu: 200m