Repository: Azure/aad-pod-identity
Pull request #408 (2019-10-03 15:37:28)

improvements for scale
<!-- 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? -->
- Pod informer modified to use node based listing in NMI
- 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 label. This reduces the payload passed between the etcd, api-server and NMI. With this we can scale to large number of nodes.
- `CREATE/UPDATE/DELETE` operations for each node/VMSS in parallel based on '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.
- Add flags for setting `QPS`, `burst` values for kube client.
- Use maps instead of for loops for efficient computation.
- Update stats.

**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..d6e81a7a 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", 20, "Per node/VMSS create/delete batches")
+
+	// Client QPS is used to configure the client-go QPS throttling and bursting.
+	flag.Float64Var(&clientQPS, "clientQps", 5, "Client QPS used for throttling of calls to kube-api 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/cloudprovider/vm.go b/pkg/cloudprovider/vm.go
index 4d991f17..97070c11 100644
--- a/pkg/cloudprovider/vm.go
+++ b/pkg/cloudprovider/vm.go
@@ -63,6 +63,7 @@ func (c *VMClient) CreateOrUpdate(rg string, nodeName string, vm compute.Virtual
 		glog.Error(err)
 		return err
 	}
+	stats.UpdateCount(stats.TotalPutCalls, 1)
 	stats.Update(stats.CloudPut, time.Since(begin))
 	return nil
 }
@@ -75,6 +76,7 @@ func (c *VMClient) Get(rgName string, nodeName string) (compute.VirtualMachine,
 		glog.Error(err)
 		return vm, err
 	}
+	stats.UpdateCount(stats.TotalGetCalls, 1)
 	stats.Update(stats.CloudGet, time.Since(beginGetTime))
 	return vm, nil
 }
diff --git a/pkg/cloudprovider/vmss.go b/pkg/cloudprovider/vmss.go
index b13f911d..69bef7f9 100644
--- a/pkg/cloudprovider/vmss.go
+++ b/pkg/cloudprovider/vmss.go
@@ -69,6 +69,7 @@ func (c *VMSSClient) CreateOrUpdate(rg string, vmssName string, vm compute.Virtu
 		glog.Error(err)
 		return err
 	}
+	stats.UpdateCount(stats.TotalPutCalls, 1)
 	stats.Update(stats.CloudPut, time.Since(begin))
 	return nil
 }
@@ -82,6 +83,7 @@ func (c *VMSSClient) Get(rgName string, vmssName string) (ret compute.VirtualMac
 		glog.Error(err)
 		return vm, err
 	}
+	stats.UpdateCount(stats.TotalGetCalls, 1)
 	stats.Update(stats.CloudGet, time.Since(beginGetTime))
 	return vm, nil
 }
diff --git a/pkg/crd/crd.go b/pkg/crd/crd.go
index 10e7c656..39b432a8 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,26 +36,35 @@ 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
 	ListBindings() (res *[]aadpodid.AzureIdentityBinding, err error)
 	ListAssignedIDs() (res *[]aadpodid.AzureAssignedIdentity, err error)
+	ListAssignedIDsInMap() (res map[string]aadpodid.AzureAssignedIdentity, err error)
 	ListIds() (res *[]aadpodid.AzureIdentity, err error)
 	ListPodIds(podns, podname string) (map[string][]aadpodid.AzureIdentity, error)
 	ListPodIdentityExceptions(ns string) (res *[]aadpodid.AzurePodIdentityException, err error)
 }
 
 // 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 +212,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 +264,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 +273,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 +304,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 +315,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 +323,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
@@ -346,6 +385,37 @@ func (c *Client) ListAssignedIDs() (res *[]aadpodid.AzureAssignedIdentity, err e
 	return &resList, nil
 }
 
+// ListAssignedIDsInMap gets the list of current assigned ids, adds it to a map
+// with assigned identity name as key and assigned identity as value.
+func (c *Client) ListAssignedIDsInMap() (map[string]aadpodid.AzureAssignedIdentity, error) {
+	begin := time.Now()
+
+	result := make(map[string]aadpodid.AzureAssignedIdentity)
+	list := c.AssignedIDInformer.GetStore().List()
+
+	for _, assignedID := range list {
+
+		o, ok := assignedID.(*aadpodid.AzureAssignedIdentity)
+		if !ok {
+			err := fmt.Errorf("could not cast %T to %s", assignedID, aadpodid.AzureAssignedIDResource)
+			c.log.Error(err)
+			return nil, err
+		}
+		// Note: List items returned from cache have empty Kind and API version..
+		// Work around this issue since we need that for event recording to work.
+		o.SetGroupVersionKind(schema.GroupVersionKind{
+			Group:   aadpodid.CRDGroup,
+			Version: aadpodid.CRDVersion,
+			Kind:    reflect.TypeOf(*o).String()})
+
+		// assigned identities names are unique across namespaces as we use pod name-<id ns>-<id name>
+		result[o.Name] = *o
+	}
+
+	stats.Update(stats.AssignedIDList, time.Since(begin))
+	return result, nil
+}
+
 // ListIds returns a list of azureidentities
 func (c *Client) ListIds() (res *[]aadpodid.AzureIdentity, err error) {
 	begin := time.Now()
@@ -441,6 +511,7 @@ func (c *Client) UpdateAzureAssignedIdentityStatus(assignedIdentity *aadpodid.Az
 		return err
 	}
 
+	begin := time.Now()
 	err = c.rest.
 		Patch(types.JSONPatchType).
 		Namespace(assignedIdentity.Namespace).
@@ -449,6 +520,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..df1260ee 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()
@@ -272,7 +273,7 @@ func (c *Client) Sync(exit <-chan struct{}) {
 			continue
 		}
 
-		currentAssignedIDs, err := c.CRDClient.ListAssignedIDs()
+		currentAssignedIDs, err := c.CRDClient.ListAssignedIDsInMap()
 		if err != nil {
 			continue
 		}
@@ -288,12 +289,12 @@ func (c *Client) Sync(exit <-chan struct{}) {
 
 		// Extract add list and delete list based on existing assigned ids in the system (currentAssignedIDs).
 		// and the ones we have arrived at in the volatile list (newAssignedIDs).
-		addList, err := c.getAzureAssignedIDsToCreate(*currentAssignedIDs, newAssignedIDs)
+		addList, err := c.getAzureAssignedIDsToCreate(currentAssignedIDs, newAssignedIDs)
 		if err != nil {
 			glog.Error(err)
 			continue
 		}
-		deleteList, err := c.getAzureAssignedIDsToDelete(currentAssignedIDs, &newAssignedIDs)
+		deleteList, err := c.getAzureAssignedIDsToDelete(currentAssignedIDs, newAssignedIDs)
 		if err != nil {
 			glog.Error(err)
 			continue
@@ -309,13 +310,13 @@ func (c *Client) Sync(exit <-chan struct{}) {
 
 		// process the delete and add list
 		// determine the list of identities that need to updated, create a node to identity list mapping for add and delete
-		if deleteList != nil && len(*deleteList) > 0 {
+		if len(deleteList) > 0 {
 			workDone = true
-			c.getListOfIdsToDelete(*deleteList, newAssignedIDs, nodeMap, nodeRefs)
+			c.getListOfIdsToDelete(deleteList, newAssignedIDs, nodeMap, nodeRefs)
 		}
-		if addList != nil && len(*addList) > 0 {
+		if len(addList) > 0 {
 			workDone = true
-			c.getListOfIdsToAssign(*addList, nodeMap)
+			c.getListOfIdsToAssign(addList, nodeMap)
 		}
 
 		var wg sync.WaitGroup
@@ -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,16 +341,23 @@ 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)
+			}
 		}
 	}
 }
 
-func (c *Client) convertAssignedIDListToMap(addList, deleteList *[]aadpodid.AzureAssignedIdentity, nodeMap map[string]trackUserAssignedMSIIds) {
+func (c *Client) convertAssignedIDListToMap(addList, deleteList map[string]aadpodid.AzureAssignedIdentity, nodeMap map[string]trackUserAssignedMSIIds) {
 	if addList != nil {
-		for _, createID := range *addList {
+		for _, createID := range addList {
 			if trackList, ok := nodeMap[createID.Spec.NodeName]; ok {
 				trackList.assignedIDsToCreate = append(trackList.assignedIDsToCreate, createID)
 				nodeMap[createID.Spec.NodeName] = trackList
@@ -357,7 +368,7 @@ func (c *Client) convertAssignedIDListToMap(addList, deleteList *[]aadpodid.Azur
 	}
 
 	if deleteList != nil {
-		for _, delID := range *deleteList {
+		for _, delID := range deleteList {
 			if trackList, ok := nodeMap[delID.Spec.NodeName]; ok {
 				trackList.assignedIDsToDelete = append(trackList.assignedIDsToDelete, delID)
 				nodeMap[delID.Spec.NodeName] = trackList
@@ -369,13 +380,13 @@ func (c *Client) convertAssignedIDListToMap(addList, deleteList *[]aadpodid.Azur
 }
 
 func (c *Client) createDesiredAssignedIdentityList(
-	listPods []*corev1.Pod, listBindings *[]aadpodid.AzureIdentityBinding, idMap map[string]aadpodid.AzureIdentity) ([]aadpodid.AzureAssignedIdentity, map[string]bool, error) {
+	listPods []*corev1.Pod, listBindings *[]aadpodid.AzureIdentityBinding, idMap map[string]aadpodid.AzureIdentity) (map[string]aadpodid.AzureAssignedIdentity, map[string]bool, error) {
 	//For each pod, check what bindings are matching. For each binding create volatile azure assigned identity.
 	//Compare this list with the current list of azure assigned identities.
 	//For any new assigned identities found in this volatile list, create assigned identity and assign user assigned msis.
 	//For any assigned ids not present the volatile list, proceed with the deletion.
 	nodeRefs := make(map[string]bool)
-	var newAssignedIDs []aadpodid.AzureAssignedIdentity
+	newAssignedIDs := make(map[string]aadpodid.AzureAssignedIdentity)
 
 	for _, pod := range listPods {
 		if pod.Spec.NodeName == "" {
@@ -417,7 +428,7 @@ func (c *Client) createDesiredAssignedIdentityList(
 					glog.Errorf("failed to create assignment for pod %s/%s with identity %s/%s with error %v", pod.Namespace, pod.Name, azureID.Namespace, azureID.Name, err.Error())
 					continue
 				}
-				newAssignedIDs = append(newAssignedIDs, *assignedID)
+				newAssignedIDs[assignedID.Name] = *assignedID
 			} else {
 				// This is the case where the identity has been deleted.
 				// In such a case, we will skip it from matching binding.
@@ -432,7 +443,10 @@ func (c *Client) createDesiredAssignedIdentityList(
 
 // getListOfIdsToDelete will go over the delete list to determine if the id is required to be deleted
 // only user assigned identity not in use are added to the remove list for the node
-func (c *Client) getListOfIdsToDelete(deleteList, newAssignedIDs []aadpodid.AzureAssignedIdentity, nodeMap map[string]trackUserAssignedMSIIds, nodeRefs map[string]bool) {
+func (c *Client) getListOfIdsToDelete(deleteList map[string]aadpodid.AzureAssignedIdentity,
+	newAssignedIDs map[string]aadpodid.AzureAssignedIdentity,
+	nodeMap map[string]trackUserAssignedMSIIds,
+	nodeRefs map[string]bool) {
 	vmssGroups, err := getVMSSGroups(c.NodeClient, nodeRefs)
 	if err != nil {
 		glog.Error(err)
@@ -459,7 +473,7 @@ func (c *Client) getListOfIdsToDelete(deleteList, newAssignedIDs []aadpodid.Azur
 }
 
 // getListOfIdsToAssign will add the id to the append list for node if it's user assigned identity
-func (c *Client) getListOfIdsToAssign(addList []aadpodid.AzureAssignedIdentity, nodeMap map[string]trackUserAssignedMSIIds) {
+func (c *Client) getListOfIdsToAssign(addList map[string]aadpodid.AzureAssignedIdentity, nodeMap map[string]trackUserAssignedMSIIds) {
 	for _, createID := range addList {
 		id := createID.Spec.AzureIdentityRef
 		isUserAssignedMSI := c.checkIfUserAssignedMSI(id)
@@ -473,117 +487,110 @@ func (c *Client) getListOfIdsToAssign(addList []aadpodid.AzureAssignedIdentity,
 	}
 }
 
-func (c *Client) matchAssignedID(x *aadpodid.AzureAssignedIdentity, y *aadpodid.AzureAssignedIdentity) (ret bool, err error) {
+func (c *Client) matchAssignedID(x aadpodid.AzureAssignedIdentity, y aadpodid.AzureAssignedIdentity) (ret bool) {
 	bindingX := x.Spec.AzureBindingRef
 	bindingY := y.Spec.AzureBindingRef
 
 	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 &&
-		x.Spec.Pod == y.Spec.Pod && x.Spec.PodNamespace == y.Spec.PodNamespace && x.Spec.NodeName == y.Spec.NodeName {
-		return true, nil
-	}
-	return false, nil
+	return bindingX.Name == bindingY.Name &&
+		bindingX.ResourceVersion == bindingY.ResourceVersion &&
+		idX.Name == idY.Name &&
+		idX.ResourceVersion == idY.ResourceVersion &&
+		x.Spec.Pod == y.Spec.Pod &&
+		x.Spec.PodNamespace == y.Spec.PodNamespace &&
+		x.Spec.NodeName == y.Spec.NodeName
 }
 
-func (c *Client) getAzureAssignedIDsToCreate(old, new []aadpodid.AzureAssignedIdentity) (*[]aadpodid.AzureAssignedIdentity, error) {
+func (c *Client) getAzureAssignedIDsToCreate(old, new map[string]aadpodid.AzureAssignedIdentity) (map[string]aadpodid.AzureAssignedIdentity, error) {
 	// everything in new needs to be created
-	if old == nil || len(old) == 0 {
-		return &new, nil
+	if len(old) == 0 {
+		return new, nil
 	}
 
-	create := make([]aadpodid.AzureAssignedIdentity, 0)
-	var err error
-	idMatch := false
+	create := make(map[string]aadpodid.AzureAssignedIdentity)
 	begin := time.Now()
-	// TODO: We should be able to optimize the many for loops.
-	for _, newAssignedID := range new {
-		idMatch = false
-		for _, oldAssignedID := range old {
-			idMatch, err = c.matchAssignedID(&newAssignedID, &oldAssignedID)
-			if err != nil {
-				glog.Error(err)
-				continue
-			}
-			if idMatch {
-				// if the old assigned id is in created state, then the identity assignment to the node
-				// is not done. Adding to the list will ensure we retry identity assignment to node for
-				// this assigned identity.
-				if oldAssignedID.Status.Status == aadpodid.AssignedIDCreated {
-					create = append(create, oldAssignedID)
-				}
-				break
-			}
+
+	for assignedIDName, newAssignedID := range new {
+		oldAssignedID, exists := old[assignedIDName]
+		idMatch := false
+		if exists {
+			idMatch = c.matchAssignedID(oldAssignedID, newAssignedID)
+		}
+		if idMatch && oldAssignedID.Status.Status == aadpodid.AssignedIDCreated {
+			// if the old assigned id is in created state, then the identity assignment to the node
+			// is not done. Adding to the list will ensure we retry identity assignment to node for
+			// this assigned identity.
+			glog.V(5).Infof("ok: %v, Create added: %s", idMatch, assignedIDName)
+			create[assignedIDName] = oldAssignedID
 		}
 		if !idMatch {
-			glog.V(5).Infof("ok: %v, Create added: %s", idMatch, newAssignedID.Name)
 			// We are done checking that this new id is not present in the old
 			// list. So we will add it to the create list.
-			create = append(create, newAssignedID)
+			glog.V(5).Infof("ok: %v, Create added: %s", idMatch, assignedIDName)
+			create[assignedIDName] = newAssignedID
 		}
 	}
 	stats.Put(stats.FindAssignedIDCreate, time.Since(begin))
-	return &create, nil
+	return create, nil
 }
 
-func (c *Client) getAzureAssignedIDsToDelete(old *[]aadpodid.AzureAssignedIdentity, new *[]aadpodid.AzureAssignedIdentity) (*[]aadpodid.AzureAssignedIdentity, error) {
+func (c *Client) getAzureAssignedIDsToDelete(old, new map[string]aadpodid.AzureAssignedIdentity) (map[string]aadpodid.AzureAssignedIdentity, error) {
+	delete := make(map[string]aadpodid.AzureAssignedIdentity)
 	// nothing to delete
-	if old == nil || len(*old) == 0 {
-		return nil, nil
+	if len(old) == 0 {
+		return delete, nil
 	}
 	// delete everything as nothing in new
-	if new == nil || len(*new) == 0 {
+	if len(new) == 0 {
 		return old, nil
 	}
 
-	delete := make([]aadpodid.AzureAssignedIdentity, 0)
-	var err error
-	idMatch := false
-
 	begin := time.Now()
-	for _, oldAssignedID := range *old {
-		idMatch = false
-		for _, newAssignedID := range *new {
-			idMatch, err = c.matchAssignedID(&newAssignedID, &oldAssignedID)
-			if err != nil {
-				glog.Error(err)
-				continue
-			}
-			//glog.Infof("Match %s %s %v", newAssignedID.Name, oldAssignedID.Name, idMatch)
-			if idMatch {
-				break
-			}
-		}
-		if !idMatch {
-			glog.V(5).Infof("ok: %v, Delete added: %s", idMatch, oldAssignedID.Name)
-			// We are done checking that this old id is not present in the new
-			// list. So we will add it to the delete list.
-			delete = append(delete, oldAssignedID)
+	for assignedIDName, oldAssignedID := range old {
+		newAssignedID, exists := new[assignedIDName]
+		idMatch := false
+		if exists {
+			idMatch = c.matchAssignedID(oldAssignedID, newAssignedID)
+		}
+		// assigned identity exists in the desired list too which means
+		// it should not be deleted
+		if exists && idMatch {
+			continue
 		}
+		// We are done checking that this old id is not present in the new
+		// list. So we will add it to the delete list.
+		delete[assignedIDName] = oldAssignedID
 	}
 	stats.Put(stats.FindAssignedIDDel, time.Since(begin))
-	return &delete, nil
+	return delete, nil
 }
 
 func (c *Client) makeAssignedIDs(azID aadpodid.AzureIdentity, azBinding aadpodid.AzureIdentityBinding, podName, podNameSpace, nodeName string) (res *aadpodid.AzureAssignedIdentity, err error) {
 	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,
@@ -677,8 +684,8 @@ func (c *Client) convertIDListToMap(arr []aadpodid.AzureIdentity) (m map[string]
 	return m, nil
 }
 
-func (c *Client) checkIfInUse(checkAssignedID aadpodid.AzureAssignedIdentity, arr []aadpodid.AzureAssignedIdentity, vmssGroups *vmssGroupList) (bool, error) {
-	for _, assignedID := range arr {
+func (c *Client) checkIfInUse(checkAssignedID aadpodid.AzureAssignedIdentity, assignedIDMap map[string]aadpodid.AzureAssignedIdentity, vmssGroups *vmssGroupList) (bool, error) {
+	for _, assignedID := range assignedIDMap {
 		checkID := checkAssignedID.Spec.AzureIdentityRef
 		id := assignedID.Spec.AzureIdentityRef
 		// If they have the same client id, reside on the same node but the pod name is different, then the
@@ -735,33 +742,53 @@ func (c *Client) updateAssignedIdentityStatus(assignedID *aadpodid.AzureAssigned
 	return c.CRDClient.UpdateAzureAssignedIdentityStatus(assignedID, status)
 }
 
-func (c *Client) updateNodeAndDeps(newAssignedIDs []aadpodid.AzureAssignedIdentity, nodeMap map[string]trackUserAssignedMSIIds, nodeRefs map[string]bool, wg *sync.WaitGroup) {
+func (c *Client) updateNodeAndDeps(newAssignedIDs map[string]aadpodid.AzureAssignedIdentity, nodeMap map[string]trackUserAssignedMSIIds, nodeRefs map[string]bool, wg *sync.WaitGroup) {
 	for nodeName, nodeTrackList := range nodeMap {
 		wg.Add(1)
 		go c.updateUserMSI(newAssignedIDs, nodeName, nodeTrackList, nodeRefs, wg)
 	}
 }
 
-func (c *Client) updateUserMSI(newAssignedIDs []aadpodid.AzureAssignedIdentity, nodeOrVMSSName string, nodeTrackList trackUserAssignedMSIIds, nodeRefs map[string]bool, wg *sync.WaitGroup) {
+func (c *Client) updateUserMSI(newAssignedIDs map[string]aadpodid.AzureAssignedIdentity, nodeOrVMSSName string, nodeTrackList trackUserAssignedMSIIds, nodeRefs map[string]bool, wg *sync.WaitGroup) {
 	defer wg.Done()
 	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 +871,76 @@ 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.UpdateCount(stats.TotalAssignedIDsCreated, len(nodeTrackList.assignedIDsToCreate))
+	stats.UpdateCount(stats.TotalAssignedIDsDeleted, len(nodeTrackList.assignedIDsToDelete))
 	stats.Put(stats.TotalCreateOrUpdate, time.Since(beginAdding))
 }
 
diff --git a/pkg/stats/stats.go b/pkg/stats/stats.go
index 4802976b..47baae6d 100644
--- a/pkg/stats/stats.go
+++ b/pkg/stats/stats.go
@@ -8,37 +8,44 @@ import (
 )
 
 var GlobalStats map[StatsType]time.Duration
+var CountStats map[StatsType]int
+
 var Mutex *sync.RWMutex
 
 type StatsType string
 
 const (
-	Total                StatsType = "Total"
-	System               StatsType = "System"
-	CacheSync            StatsType = "CacheSync"
-	CurrentState         StatsType = "Gather current state"
-	PodList              StatsType = "Pod listing"
-	BindingList          StatsType = "Binding listing"
-	IDList               StatsType = "ID listing"
-	ExceptionList        StatsType = "Pod Identity Exception listing"
-	AssignedIDList       StatsType = "Assigned ID listing"
-	CloudGet             StatsType = "Cloud provider get"
-	CloudPut             StatsType = "Cloud provider put"
-	K8sGet               StatsType = "K8s get"
-	K8sPut               StatsType = "K8s put"
-	FindAssignedIDDel    StatsType = "Find assigned ids to delete"
-	FindAssignedIDCreate StatsType = "Find assigned ids to create"
-	AssignedIDDel        StatsType = "Assigned ID deletion"
-	AssignedIDAdd        StatsType = "Assigned ID addition"
-	TotalIDDel           StatsType = "Total time to delete assigned IDs"
-	TotalIDAdd           StatsType = "Total time to add assigned IDs"
-	TotalCreateOrUpdate  StatsType = "Total time to assign or remove IDs"
+	Total                   StatsType = "Total"
+	System                  StatsType = "System"
+	CacheSync               StatsType = "CacheSync"
+	CurrentState            StatsType = "Gather current state"
+	PodList                 StatsType = "Pod listing"
+	BindingList             StatsType = "Binding listing"
+	IDList                  StatsType = "ID listing"
+	ExceptionList           StatsType = "Pod Identity Exception listing"
+	AssignedIDList          StatsType = "Assigned ID listing"
+	CloudGet                StatsType = "Cloud provider get"
+	CloudPut                StatsType = "Cloud provider put"
+	TotalPutCalls           StatsType = "Number of cloud provider PUT"
+	TotalGetCalls           StatsType = "Number of cloud provider GET"
+	TotalAssignedIDsCreated StatsType = "Number of assigned ids created in this sync cycle"
+	TotalAssignedIDsDeleted StatsType = "Number of assigned ids deleted in this sync cycle"
+	K8sGet                  StatsType = "K8s get"
+	K8sPut                  StatsType = "K8s put"
+	FindAssignedIDDel       StatsType = "Find assigned ids to delete"
+	FindAssignedIDCreate    StatsType = "Find assigned ids to create"
+	AssignedIDDel           StatsType = "Assigned ID deletion"
+	AssignedIDAdd           StatsType = "Assigned ID addition"
+	TotalIDDel              StatsType = "Total time to delete assigned IDs"
+	TotalIDAdd              StatsType = "Total time to add assigned IDs"
+	TotalCreateOrUpdate     StatsType = "Total time to assign or remove IDs"
 
 	EventRecord StatsType = "Event recording"
 )
 
 func Init() {
 	GlobalStats = make(map[StatsType]time.Duration)
+	CountStats = make(map[StatsType]int)
 	Mutex = &sync.RWMutex{}
 }
 
@@ -70,9 +77,23 @@ func Update(key StatsType, val time.Duration) {
 func Print(key StatsType) {
 	Mutex.RLock()
 	defer Mutex.RUnlock()
+
 	glog.Infof("%s: %s", key, GlobalStats[key])
 }
 
+func PrintCount(key StatsType) {
+	Mutex.RLock()
+	defer Mutex.RUnlock()
+
+	glog.Infof("%s: %d", key, CountStats[key])
+}
+
+func UpdateCount(key StatsType, val int) {
+	Mutex.Lock()
+	defer Mutex.Unlock()
+	CountStats[key] = CountStats[key] + val
+}
+
 func PrintSync() {
 	glog.Infof("** Stats collected **")
 	if GlobalStats != nil {
@@ -89,6 +110,12 @@ func PrintSync() {
 		Print(AssignedIDAdd)
 		Print(AssignedIDDel)
 
+		PrintCount(TotalPutCalls)
+		PrintCount(TotalGetCalls)
+
+		PrintCount(TotalAssignedIDsCreated)
+		PrintCount(TotalAssignedIDsDeleted)
+
 		Print(FindAssignedIDCreate)
 		Print(FindAssignedIDDel)
 


Test patch:
diff --git a/pkg/mic/mic_test.go b/pkg/mic/mic_test.go
index b8940018..b7c72f95 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{}) {
 
 }
 
@@ -451,6 +455,16 @@ func (c *TestCrdClient) ListAssignedIDs() (res *[]aadpodid.AzureAssignedIdentity
 	return &assignedIDList, nil
 }
 
+func (c *TestCrdClient) ListAssignedIDsInMap() (res map[string]aadpodid.AzureAssignedIdentity, err error) {
+	assignedIDMap := make(map[string]aadpodid.AzureAssignedIdentity)
+	c.mu.Lock()
+	for k, v := range c.assignedIDMap {
+		assignedIDMap[k] = *v
+	}
+	c.mu.Unlock()
+	return assignedIDMap, nil
+}
+
 func (c *Client) ListPodIds(podns, podname string) (map[string][]aadpodid.AzureIdentity, error) {
 	return map[string][]aadpodid.AzureIdentity{}, nil
 }
@@ -586,7 +600,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 +612,7 @@ func NewMICTestClient(eventCh chan aadpodid.EventType,
 		NodeClient:        nodeClient,
 		syncRetryInterval: 120 * time.Second,
 		IsNamespaced:      isNamespaced,
+		createDeleteBatch: createDeleteBatch,
 	}
 
 	return &TestMICClient{
@@ -703,7 +719,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 +840,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 +932,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 +1035,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 +1149,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 +1213,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 +1292,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, "", "", "", "")
@@ -1338,6 +1354,61 @@ func TestSyncNodeNotFound(t *testing.T) {
 	}
 }
 
+func TestProcessingTimeForScale(t *testing.T) {
+	eventCh := make(chan aadpodid.EventType, 20000)
+	cloudClient := NewTestCloudClient(config.AzureConfig{})
+	crdClient := NewTestCrdClient(nil)
+	podClient := NewTestPodClient()
+	nodeClient := NewTestNodeClient()
+	var evtRecorder TestEventRecorder
+	evtRecorder.lastEvent = new(LastEvent)
+	evtRecorder.eventChannel = make(chan bool, 20000)
+
+	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, "", "", "", "")
+	crdClient.CreateBinding("testbinding1", "default", "test-id1", "test-select1", "")
+
+	nodeClient.AddNode("test-node1")
+	for i := 0; i < 20000; i++ {
+		podClient.AddPod(fmt.Sprintf("test-pod%d", i), "default", "test-node1", "test-select1")
+	}
+	eventCh <- aadpodid.PodCreated
+
+	defer micClient.testRunSync()(t)
+
+	if !evtRecorder.WaitForEvents(20000) {
+		t.Fatalf("Timeout waiting for mic sync cycles")
+	}
+
+	listAssignedIDs, err := crdClient.ListAssignedIDs()
+	if err != nil {
+		glog.Error(err)
+		t.Errorf("list assigned failed")
+	}
+	if !(len(*listAssignedIDs) == 20000) {
+		t.Fatalf("expected assigned identities len: %d, got: %d", 20000, len(*listAssignedIDs))
+	}
+
+	for i := 10000; i < 20000; i++ {
+		podClient.DeletePod(fmt.Sprintf("test-pod%d", i), "default")
+	}
+	eventCh <- aadpodid.PodDeleted
+
+	if !evtRecorder.WaitForEvents(10000) {
+		t.Fatalf("Timeout waiting for mic sync cycles")
+	}
+	listAssignedIDs, err = crdClient.ListAssignedIDs()
+	if err != nil {
+		glog.Error(err)
+		t.Errorf("list assigned failed")
+	}
+	if !(len(*listAssignedIDs) == 10000) {
+		t.Fatalf("expected assigned identities len: %d, got: %d", 10000, len(*listAssignedIDs))
+	}
+}
+
 func TestSyncExit(t *testing.T) {
 	eventCh := make(chan aadpodid.EventType)
 	cloudClient := NewTestCloudClient(config.AzureConfig{VMType: "vmss"})
@@ -1348,7 +1419,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