Repository: Azure/aad-pod-identity
Pull request #219 (2019-05-17 20:44:49)

Performance  Improvement - reduce identity assignment time
resolves 
https://github.com/Azure/aad-pod-identity/issues/145
https://github.com/Azure/aad-pod-identity/issues/181
https://github.com/Azure/aad-pod-identity/issues/217

- Batch process id assignments to node on each sync cycle (pooling delete and add ids, issuing one CreateOrUpdate call for each sync cycle)
- Reduces identity assignment for 44 pods from 29m to ~3m

Add states to the assigned identities
1. Created - assigned identity will be in created state when the crd is created.
2. Assigned - once the identity has been successfully assigned to underlying node, the state of the assigned identity will be changed to Assigned.
3. Unassigned - the assigned identity will be in this state when the identity has been successfully removed from the underlying node.

**Note**
1. Assigned identities will be created and then the identity will be tried to be assigned to the node. If assignment fails, then the identity will be deleted. No change in this behavior compared to previous version.
2. Assigned identities will be deleted only after the identity has been successfully removed from the node. This is change in behavior from previous version to ensure we successfully remove the identities from the underlying node.

kkmsft: Thank you @aramase . A good e2e could be validating that #145 is resolved and the identities are assigned in a reasonable amount of time.
aramase: @kkmsft I was able to validate https://github.com/Azure/aad-pod-identity/issues/145 with my changes. Created close to 44 pods with 2 assigned identities. With version 1.3 it took about 29m to create all assigned identities and assign msi to node. With my changes, it takes only about a minute since it's batching all the updates with a single call. I've a lot of refactoring of my PR still though.
kkmsft: Thank you @aramas. Good to know that it's moving in the right direction. After our initial discussion on the approach have not looked deeply into the changes. Please let me know when you want me to start looking at the changes. 

Also, any chance (I know I am being greedy here ;-)) you can convert that test to an e2e test so that we can track any regression in future. 
aramase: @kkmsft Sure, I'll add that as an e2e. Still working on testing a couple of things. I'll need to refactor the PR. I'll let you know when it's ready for review.
aramase: Tests passed - 

```
Ran 8 of 9 Specs in 1339.441 seconds
SUCCESS! -- 8 Passed | 0 Failed | 0 Pending | 1 Skipped
--- PASS: TestAADPodIdentity (1339.44s)
PASS
ok  	github.com/Azure/aad-pod-identity/test/e2e	1339.459s
```

@kkmsft 
aramase: <img width="408" alt="Screen Shot 2019-06-25 at 12 40 56 PM" src="https://user-images.githubusercontent.com/10608509/60128015-823c3700-9746-11e9-9a69-f478f2fd70ef.png">

aramase: Ran 10 of 10 Specs in 2353.337 seconds
SUCCESS! -- 10 Passed | 0 Failed | 0 Pending | 0 Skipped
--- PASS: TestAADPodIdentity (2353.34s)
PASS
ok  	github.com/Azure/aad-pod-identity/test/e2e	2353.353s

Patch:
diff --git a/pkg/cloudprovider/cloudprovider.go b/pkg/cloudprovider/cloudprovider.go
index 922a084d..af6af1a2 100644
--- a/pkg/cloudprovider/cloudprovider.go
+++ b/pkg/cloudprovider/cloudprovider.go
@@ -27,9 +27,12 @@ type Client struct {
 	Config     config.AzureConfig
 }
 
+// ClientInt client interface
 type ClientInt interface {
 	RemoveUserMSI(userAssignedMSIID string, node *corev1.Node) error
 	AssignUserMSI(userAssignedMSIID string, node *corev1.Node) error
+	UpdateUserMSI(addUserAssignedMSIIDs []string, removeUserAssignedMSIIDs []string, node *corev1.Node) error
+	GetUserMSIs(node *corev1.Node) ([]string, error)
 }
 
 // NewCloudProvider returns a azure cloud provider client
@@ -112,6 +115,59 @@ func withInspection() autorest.PrepareDecorator {
 	}
 }
 
+// GetUserMSIs will return a list of all identities on the node
+func (c *Client) GetUserMSIs(node *corev1.Node) ([]string, error) {
+	idH, _, err := c.getIdentityResource(node)
+	if err != nil {
+		glog.Errorf("GetUserMSIs: get identity resource failed with error %v", err)
+		return nil, err
+	}
+	info := idH.IdentityInfo()
+	if info == nil {
+		return nil, fmt.Errorf("identity info is nil")
+	}
+	idList := info.GetUserIdentityList()
+	return idList, nil
+}
+
+// UpdateUserMSI will batch process the removal and addition of ids
+func (c *Client) UpdateUserMSI(addUserAssignedMSIIDs []string, removeUserAssignedMSIIDs []string, node *corev1.Node) error {
+	idH, updateFunc, err := c.getIdentityResource(node)
+	if err != nil {
+		return err
+	}
+
+	info := idH.IdentityInfo()
+	if info == nil {
+		info = idH.ResetIdentity()
+	}
+
+	requiresUpdate := false
+	// remove msi ids from the list
+	for _, userAssignedMSIID := range removeUserAssignedMSIIDs {
+		requiresUpdate = true
+		if err := info.RemoveUserIdentity(userAssignedMSIID); err != nil {
+			return fmt.Errorf("could not remove identity from node %s: %v", node.Name, err)
+		}
+	}
+	// add new ids to the list
+	for _, userAssignedMSIID := range addUserAssignedMSIIDs {
+		addedToList := info.AppendUserIdentity(userAssignedMSIID)
+		if !addedToList {
+			glog.V(6).Infof("Identity %s already assigned to node %s. Skipping assignment.", userAssignedMSIID, node.Name)
+		}
+		requiresUpdate = requiresUpdate || addedToList
+	}
+	if requiresUpdate {
+		timeStarted := time.Now()
+		if err := updateFunc(); err != nil {
+			return err
+		}
+		glog.V(6).Infof("UpdateUserMSI of %s completed in %s", node.Name, time.Since(timeStarted))
+	}
+	return nil
+}
+
 //RemoveUserMSI - Use the underlying cloud api calls and remove the given user assigned MSI from the vm.
 func (c *Client) RemoveUserMSI(userAssignedMSIID string, node *corev1.Node) error {
 	idH, updateFunc, err := c.getIdentityResource(node)
@@ -137,6 +193,7 @@ func (c *Client) RemoveUserMSI(userAssignedMSIID string, node *corev1.Node) erro
 	return nil
 }
 
+// AssignUserMSI - Use the underlying cloud api call and add the given user assigned MSI to the vm
 func (c *Client) AssignUserMSI(userAssignedMSIID string, node *corev1.Node) error {
 	// Get the vm using the VmClient
 	// Update the assigned identity into the VM using the CreateOrUpdate
diff --git a/pkg/cloudprovider/identity.go b/pkg/cloudprovider/identity.go
index 3163a6ff..1b0f7349 100644
--- a/pkg/cloudprovider/identity.go
+++ b/pkg/cloudprovider/identity.go
@@ -21,6 +21,7 @@ type IdentityHolder interface {
 type IdentityInfo interface {
 	AppendUserIdentity(id string) bool
 	RemoveUserIdentity(id string) error
+	GetUserIdentityList() []string
 }
 
 var (
@@ -35,7 +36,7 @@ func filterUserIdentity(idType *compute.ResourceIdentityType, idList *[]string,
 	case compute.ResourceIdentityTypeUserAssigned,
 		compute.ResourceIdentityTypeSystemAssignedUserAssigned:
 	default:
-		return errNotFound
+		return nil
 	}
 
 	origLen := len(*idList)
@@ -78,10 +79,8 @@ func appendUserIdentity(idType *compute.ResourceIdentityType, idList *[]string,
 	switch *idType {
 	case compute.ResourceIdentityTypeUserAssigned, compute.ResourceIdentityTypeSystemAssignedUserAssigned:
 		// check if this ID is already in the list
-		for _, id := range *idList {
-			if id == newID {
-				return false
-			}
+		if checkIfIDInList(*idList, newID) {
+			return false
 		}
 	case compute.ResourceIdentityTypeSystemAssigned:
 		*idType = compute.ResourceIdentityTypeSystemAssignedUserAssigned
@@ -92,3 +91,12 @@ func appendUserIdentity(idType *compute.ResourceIdentityType, idList *[]string,
 	*idList = append(*idList, newID)
 	return true
 }
+
+func checkIfIDInList(idList []string, desiredID string) bool {
+	for _, id := range idList {
+		if id == desiredID {
+			return true
+		}
+	}
+	return false
+}
diff --git a/pkg/cloudprovider/vm.go b/pkg/cloudprovider/vm.go
index 56760d8a..43a7849d 100644
--- a/pkg/cloudprovider/vm.go
+++ b/pkg/cloudprovider/vm.go
@@ -115,3 +115,7 @@ func (i *vmIdentityInfo) AppendUserIdentity(id string) bool {
 	}
 	return appendUserIdentity(&i.info.Type, i.info.IdentityIds, id)
 }
+
+func (i *vmIdentityInfo) GetUserIdentityList() []string {
+	return *i.info.IdentityIds
+}
diff --git a/pkg/cloudprovider/vmss.go b/pkg/cloudprovider/vmss.go
index 567aabfc..94edf577 100644
--- a/pkg/cloudprovider/vmss.go
+++ b/pkg/cloudprovider/vmss.go
@@ -121,3 +121,7 @@ func (i *vmssIdentityInfo) AppendUserIdentity(id string) bool {
 	}
 	return appendUserIdentity(&i.info.Type, i.info.IdentityIds, id)
 }
+
+func (i *vmssIdentityInfo) GetUserIdentityList() []string {
+	return *i.info.IdentityIds
+}
diff --git a/pkg/crd/crd.go b/pkg/crd/crd.go
index 114992b6..7a792ddf 100644
--- a/pkg/crd/crd.go
+++ b/pkg/crd/crd.go
@@ -1,18 +1,21 @@
 package crd
 
 import (
+	"encoding/json"
 	"fmt"
 	"time"
 
+	"github.com/golang/glog"
+
 	aadpodid "github.com/Azure/aad-pod-identity/pkg/apis/aadpodidentity/v1"
 	"github.com/Azure/aad-pod-identity/pkg/stats"
 
-	"github.com/golang/glog"
 	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/fields"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/runtime/serializer"
+	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/cache"
 )
@@ -31,6 +34,7 @@ type ClientInt interface {
 	SyncCache(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)
 	ListIds() (res *[]aadpodid.AzureIdentity, err error)
@@ -190,16 +194,18 @@ func (c *Client) SyncCache(exit <-chan struct{}) {
 	}
 }
 
+// RemoveAssignedIdentity removes the assigned identity
 func (c *Client) RemoveAssignedIdentity(assignedIdentity *aadpodid.AzureAssignedIdentity) error {
-	glog.V(6).Infof("Deletion of id named: %s", assignedIdentity.Name)
+	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()
 	stats.Update(stats.AssignedIDDel, time.Since(begin))
 	return err
 }
 
+// CreateAssignedIdentity creates new assigned identity
 func (c *Client) CreateAssignedIdentity(assignedIdentity *aadpodid.AzureAssignedIdentity) error {
-	glog.Infof("Got id %s to assign", assignedIdentity.Name)
+	glog.Infof("Got assigned id %s to assign", assignedIdentity.Name)
 	begin := time.Now()
 	// Create a new AzureAssignedIdentity which maps the relationship between
 	// id and pod
@@ -268,3 +274,35 @@ func (c *Client) ListPodIds(podns, podname string) (*[]aadpodid.AzureIdentity, e
 
 	return &matchedIds, nil
 }
+
+type patchStatusOps struct {
+	Op    string      `json:"op"`
+	Path  string      `json:"path"`
+	Value interface{} `json:"value"`
+}
+
+// UpdateAzureAssignedIdentityStatus updates the status field in AzureAssignedIdentity to indicate current status
+func (c *Client) UpdateAzureAssignedIdentityStatus(assignedIdentity *aadpodid.AzureAssignedIdentity, status string) error {
+	glog.Infof("Updating assigned identity %s/%s status to %s", assignedIdentity.Namespace, assignedIdentity.Name, status)
+
+	ops := make([]patchStatusOps, 1)
+	ops[0].Op = "replace"
+	ops[0].Path = "/Status/status"
+	ops[0].Value = status
+
+	patchBytes, err := json.Marshal(ops)
+	if err != nil {
+		return err
+	}
+
+	err = c.rest.
+		Patch(types.JSONPatchType).
+		Namespace(assignedIdentity.Namespace).
+		Resource("azureassignedidentities").
+		Name(assignedIdentity.Name).
+		Body(patchBytes).
+		Do().
+		Error()
+
+	return err
+}
diff --git a/pkg/mic/mic.go b/pkg/mic/mic.go
index 31526e65..3825d7ab 100644
--- a/pkg/mic/mic.go
+++ b/pkg/mic/mic.go
@@ -26,8 +26,15 @@ import (
 const (
 	stopped = int32(0)
 	running = int32(1)
+	// IdentityCreated status indicates azure assigned identity is created
+	IdentityCreated = "Created"
+	// IdentityAssigned status indicates identity has been assigned to the node
+	IdentityAssigned = "Assigned"
+	// IdentityUnassigned status indicates identity has been unassigned from the node
+	IdentityUnassigned = "Unassigned"
 )
 
+// NodeGetter ...
 type NodeGetter interface {
 	Get(name string) (*corev1.Node, error)
 	Start(<-chan struct{})
@@ -47,12 +54,20 @@ type Client struct {
 	syncing int32 // protect against conucrrent sync's
 }
 
+// ClientInt ...
 type ClientInt interface {
 	Start(exit <-chan struct{})
-
 	Sync(exit <-chan struct{})
 }
 
+type trackUserAssignedMSIIds struct {
+	addUserAssignedMSIIDs    []string
+	removeUserAssignedMSIIDs []string
+	assignedIDsToCreate      []aadpodid.AzureAssignedIdentity
+	assignedIDsToDelete      []aadpodid.AzureAssignedIdentity
+}
+
+// NewMICClient returnes new mic client
 func NewMICClient(cloudconfig string, config *rest.Config, isNamespaced bool) (*Client, error) {
 	glog.Infof("Starting to create the pod identity client. Version: %v. Build date: %v", version.MICVersion, version.BuildDate)
 
@@ -91,6 +106,7 @@ func NewMICClient(cloudconfig string, config *rest.Config, isNamespaced bool) (*
 	}, nil
 }
 
+// Start ...
 func (c *Client) Start(exit <-chan struct{}) {
 	glog.V(6).Infof("MIC client starting..")
 
@@ -129,6 +145,7 @@ func (c *Client) setStopped() {
 	atomic.StoreInt32(&c.syncing, stopped)
 }
 
+// Sync ...
 func (c *Client) Sync(exit <-chan struct{}) {
 	if !c.canSync() {
 		panic("concurrent syncs")
@@ -136,9 +153,7 @@ func (c *Client) Sync(exit <-chan struct{}) {
 	defer c.setStopped()
 
 	glog.Info("Sync thread started.")
-
 	var event aadpodid.EventType
-
 	for {
 		select {
 		case <-exit:
@@ -178,148 +193,48 @@ func (c *Client) Sync(exit <-chan struct{}) {
 		}
 		stats.Put(stats.System, time.Since(systemTime))
 
-		nodeRefs := make(map[string]bool)
-
-		var newAssignedIDs []aadpodid.AzureAssignedIdentity
 		beginNewListTime := time.Now()
-		//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.
-		for _, pod := range listPods {
-			//Node is not yet allocated. In that case skip the pod
-			if pod.Spec.NodeName == "" {
-				glog.V(2).Infof("Pod %s/%s has no assigned node yet. it will be ignored", pod.Name, pod.Namespace)
-				continue
-			}
-			crdPodLabelVal := pod.Labels[aadpodid.CRDLabelKey]
-			if crdPodLabelVal == "" {
-				//No binding mentioned in the label. Just continue to the next pod
-				glog.V(2).Infof("Pod %s/%s has correct %s label but with no value. it will be ignored", pod.Name, pod.Namespace, aadpodid.CRDLabelKey)
-				continue
-			}
-			//glog.Infof("Found label with our CRDKey %s for pod: %s", crdPodLabelVal, pod.Name)
-			var matchedBindings []aadpodid.AzureIdentityBinding
-			for _, allBinding := range *listBindings {
-				if allBinding.Spec.Selector == crdPodLabelVal {
-					glog.V(5).Infof("Found binding match for pod %s/%s with binding %s", pod.Name, pod.Namespace, allBinding.Name)
-					matchedBindings = append(matchedBindings, allBinding)
-					nodeRefs[pod.Spec.NodeName] = true
-				}
-			}
-
-			for _, binding := range matchedBindings {
-				glog.V(5).Infof("Looking up id map: %v", binding.Spec.AzureIdentity)
-				if azureID, idPresent := idMap[binding.Spec.AzureIdentity]; idPresent {
-					// working in Namespaced mode or this specific identity is namespaced
-					if c.IsNamespaced || aadpodid.IsNamespacedIdentity(&azureID) {
-						// They have to match all
-						if !(azureID.Namespace == binding.Namespace && binding.Namespace == pod.Namespace) {
-							glog.V(5).Infof("identity %s/%s was matched via binding %s/%s to %s/%s but namespaced identity is enforced, so it will be ignored",
-								azureID.Namespace, azureID.Name, binding.Namespace, binding.Name, pod.Namespace, pod.Name)
-							continue
-						}
-					}
-					glog.V(5).Infof("identity %s/%s assigned to %s/%s via %s/%s", azureID.Namespace, azureID.Name, pod.Namespace, pod.Name, binding.Namespace, binding.Namespace)
-					assignedID, err := c.makeAssignedIDs(&azureID, &binding, pod.Name, pod.Namespace, pod.Spec.NodeName)
-
-					if err != nil {
-						glog.Errorf("failed to create assignment for pod %s/%s with identity %s/%s with error %v", pod.Name, pod.Namespace, azureID.Namespace, azureID.Name, err.Error())
-						continue
-					}
-					newAssignedIDs = append(newAssignedIDs, *assignedID)
-				} else {
-					// This is the case where the identity has been deleted.
-					// In such a case, we will skip it from matching binding.
-					// This will ensure that the new assigned ids created will not have the
-					// one associated with this azure identity.
-					glog.V(5).Infof("%s identity not found when using %s binding", binding.Spec.AzureIdentity, binding.Name)
-				}
-			}
+		newAssignedIDs, nodeRefs, err := c.createDesiredAssignedIdentityList(listPods, listBindings, idMap)
+		if err != nil {
+			glog.Error(err)
+			continue
 		}
 		stats.Put(stats.CurrentState, time.Since(beginNewListTime))
 
 		// 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).
-		// TODO: Separate this into two methods.
-		addList, deleteList, err := c.splitAzureAssignedIDs(currentAssignedIDs, &newAssignedIDs)
+		addList, err := c.getAzureAssignedIDsToCreate(*currentAssignedIDs, newAssignedIDs)
+		if err != nil {
+			glog.Error(err)
+			continue
+		}
+		deleteList, err := c.getAzureAssignedIDsToDelete(currentAssignedIDs, &newAssignedIDs)
 		if err != nil {
 			glog.Error(err)
 			continue
 		}
-
 		glog.V(5).Infof("del: %v, add: %v", deleteList, addList)
 
-		if deleteList != nil && len(*deleteList) > 0 {
-			beginDeletion := time.Now()
-			workDone = true
+		nodeMap := make(map[string]trackUserAssignedMSIIds)
 
-			vmssGroups, err := getVMSSGroups(c.NodeClient, nodeRefs)
-			if err != nil {
-				glog.Error(err)
-				continue
-			}
+		// seperate the add and delete list per node
+		c.convertAssignedIDListToMap(addList, deleteList, nodeMap)
 
-			for _, delID := range *deleteList {
-				glog.V(5).Infof("Deletion of id: %s", delID.Name)
-				inUse, err := c.checkIfInUse(delID, newAssignedIDs, vmssGroups)
-				if err != nil {
-					glog.Error(err)
-					continue
-				}
-				removedBinding := delID.Spec.AzureBindingRef
-				// The inUse here checks if there are pods which are using the MSI in the newAssignedIDs.
-				err = c.removeAssignedIDsWithDeps(&delID, inUse)
-				if err != nil {
-					// Since k8s event has only Info and Warning, using Warning.
-					message := fmt.Sprintf("Binding %s removal from node %s for pod %s resulted in error %v", removedBinding.Name, delID.Spec.NodeName, delID.Spec.Pod, err.Error())
-					c.EventRecorder.Event(removedBinding, corev1.EventTypeWarning, "binding remove error", message)
-					glog.Error(message)
-					continue
-				}
-				eventRecordStart := time.Now()
-				glog.V(5).Infof("Binding removed: %+v", removedBinding)
-				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))
-				stats.Update(stats.EventRecord, time.Since(eventRecordStart))
-			}
-			stats.Update(stats.TotalIDDel, time.Since(beginDeletion))
+		// 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 {
+			workDone = true
+			c.getListOfIdsToDelete(*deleteList, newAssignedIDs, nodeMap, nodeRefs)
 		}
-
 		if addList != nil && len(*addList) > 0 {
-			beginAdding := time.Now()
 			workDone = true
-			for _, createID := range *addList {
-				id := createID.Spec.AzureIdentityRef
-				binding := createID.Spec.AzureBindingRef
-
-				node, err := c.NodeClient.Get(createID.Spec.NodeName)
-				if err != nil {
-					c.EventRecorder.Event(binding, corev1.EventTypeWarning, "get node error",
-						fmt.Sprintf("Lookup of node %s for pod %s resulted in error %v", createID.Spec.NodeName, createID.Name, err))
-					continue
-				}
-
-				glog.V(5).Infof("Initiating assigned id creation for pod - %s, binding - %s", createID.Spec.Pod, binding.Name)
-
-				err = c.createAssignedIdentityDeps(&createID, id, node)
-				if err != nil {
-					// Since k8s event has only Info and Warning, using Warning.
-					c.EventRecorder.Event(binding, corev1.EventTypeWarning, "binding apply error",
-						fmt.Sprintf("Applying binding %s node %s for pod %s resulted in error %v", binding.Name, createID.Spec.NodeName, createID.Name, err))
-					glog.Error(err)
-					continue
-				}
-				appliedBinding := createID.Spec.AzureBindingRef
-				eventRecordStart := time.Now()
-				glog.V(5).Infof("Binding applied: %+v", appliedBinding)
-				c.EventRecorder.Event(appliedBinding, corev1.EventTypeNormal, "binding applied",
-					fmt.Sprintf("Binding %s applied on node %s for pod %s", appliedBinding.Name, createID.Spec.NodeName, createID.Name))
-				stats.Update(stats.EventRecord, time.Since(eventRecordStart))
-			}
-			stats.Put(stats.TotalIDAdd, time.Since(beginAdding))
+			c.getListOfIdsToAssign(*addList, nodeMap)
 		}
 
+		// one final createorupdate to each node in the map
+		// TODO: parallelize this process
+		c.updateNodeAndDeps(newAssignedIDs, nodeMap, nodeRefs)
+
 		if workDone {
 			idsFound := 0
 			bindingsFound := 0
@@ -336,6 +251,134 @@ func (c *Client) Sync(exit <-chan struct{}) {
 	}
 }
 
+func (c *Client) convertAssignedIDListToMap(addList, deleteList *[]aadpodid.AzureAssignedIdentity, nodeMap map[string]trackUserAssignedMSIIds) {
+	if addList != nil {
+		for _, createID := range *addList {
+			if trackList, ok := nodeMap[createID.Spec.NodeName]; ok {
+				trackList.assignedIDsToCreate = append(trackList.assignedIDsToCreate, createID)
+				nodeMap[createID.Spec.NodeName] = trackList
+				continue
+			}
+			nodeMap[createID.Spec.NodeName] = trackUserAssignedMSIIds{assignedIDsToCreate: []aadpodid.AzureAssignedIdentity{createID}}
+		}
+	}
+	if deleteList != nil {
+		for _, delID := range *deleteList {
+			if trackList, ok := nodeMap[delID.Spec.NodeName]; ok {
+				trackList.assignedIDsToDelete = append(trackList.assignedIDsToDelete, delID)
+				nodeMap[delID.Spec.NodeName] = trackList
+				continue
+			}
+			nodeMap[delID.Spec.NodeName] = trackUserAssignedMSIIds{assignedIDsToDelete: []aadpodid.AzureAssignedIdentity{delID}}
+		}
+	}
+}
+
+func (c *Client) createDesiredAssignedIdentityList(
+	listPods []*corev1.Pod, listBindings *[]aadpodid.AzureIdentityBinding, idMap map[string]aadpodid.AzureIdentity) ([]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
+
+	for _, pod := range listPods {
+		if pod.Spec.NodeName == "" {
+			//Node is not yet allocated. In that case skip the pod
+			glog.V(2).Infof("Pod %s/%s has no assigned node yet. it will be ignored", pod.Namespace, pod.Name)
+			continue
+		}
+		crdPodLabelVal := pod.Labels[aadpodid.CRDLabelKey]
+		if crdPodLabelVal == "" {
+			//No binding mentioned in the label. Just continue to the next pod
+			glog.V(2).Infof("Pod %s/%s has correct %s label but with no value. it will be ignored", pod.Namespace, pod.Name, aadpodid.CRDLabelKey)
+			continue
+		}
+		var matchedBindings []aadpodid.AzureIdentityBinding
+		for _, allBinding := range *listBindings {
+			if allBinding.Spec.Selector == crdPodLabelVal {
+				glog.V(5).Infof("Found binding match for pod %s/%s with binding %s", pod.Namespace, pod.Name, allBinding.Name)
+				matchedBindings = append(matchedBindings, allBinding)
+				nodeRefs[pod.Spec.NodeName] = true
+			}
+		}
+
+		for _, binding := range matchedBindings {
+			glog.V(5).Infof("Looking up id map: %v", binding.Spec.AzureIdentity)
+			if azureID, idPresent := idMap[binding.Spec.AzureIdentity]; idPresent {
+				// working in Namespaced mode or this specific identity is namespaced
+				if c.IsNamespaced || aadpodid.IsNamespacedIdentity(&azureID) {
+					// They have to match all
+					if !(azureID.Namespace == binding.Namespace && binding.Namespace == pod.Namespace) {
+						glog.V(5).Infof("identity %s/%s was matched via binding %s/%s to %s/%s but namespaced identity is enforced, so it will be ignored",
+							azureID.Namespace, azureID.Name, binding.Namespace, binding.Name, pod.Name, pod.Namespace)
+						continue
+					}
+				}
+				glog.V(5).Infof("identity %s/%s assigned to %s/%s via %s/%s", azureID.Namespace, azureID.Name, pod.Namespace, pod.Name, binding.Namespace, binding.Namespace)
+				assignedID, err := c.makeAssignedIDs(&azureID, &binding, pod.Name, pod.Namespace, pod.Spec.NodeName)
+
+				if err != nil {
+					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)
+			} else {
+				// This is the case where the identity has been deleted.
+				// In such a case, we will skip it from matching binding.
+				// This will ensure that the new assigned ids created will not have the
+				// one associated with this azure identity.
+				glog.V(5).Infof("%s identity not found when using %s binding", binding.Spec.AzureIdentity, binding.Name)
+			}
+		}
+	}
+	return newAssignedIDs, nodeRefs, nil
+}
+
+// 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) {
+	vmssGroups, err := getVMSSGroups(c.NodeClient, nodeRefs)
+	if err != nil {
+		glog.Error(err)
+		return
+	}
+	for _, delID := range deleteList {
+		glog.V(5).Infof("Deletion of id: %s", delID.Name)
+		inUse, err := c.checkIfInUse(delID, newAssignedIDs, vmssGroups)
+		if err != nil {
+			glog.Error(err)
+			continue
+		}
+		id := delID.Spec.AzureIdentityRef
+		isUserAssignedMSI := c.checkIfUserAssignedMSI(id)
+
+		// this case includes Assigned state and empty state to ensure backward compatability
+		if delID.Status.Status == IdentityAssigned || delID.Status.Status == "" {
+			if !inUse && isUserAssignedMSI {
+				c.appendToRemoveListForNode(id.Spec.ResourceID, delID.Spec.NodeName, nodeMap)
+			}
+		}
+		glog.V(5).Infof("Binding removed: %+v", delID.Spec.AzureBindingRef)
+	}
+}
+
+// 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) {
+	for _, createID := range addList {
+		id := createID.Spec.AzureIdentityRef
+		isUserAssignedMSI := c.checkIfUserAssignedMSI(id)
+
+		if createID.Status.Status == "" || createID.Status.Status == IdentityCreated {
+			if isUserAssignedMSI {
+				c.appendToAddListForNode(id.Spec.ResourceID, createID.Spec.NodeName, nodeMap)
+			}
+		}
+		glog.V(5).Infof("Binding applied: %+v", createID.Spec.AzureBindingRef)
+	}
+}
+
 func (c *Client) matchAssignedID(x *aadpodid.AzureAssignedIdentity, y *aadpodid.AzureAssignedIdentity) (ret bool, err error) {
 	bindingX := x.Spec.AzureBindingRef
 	bindingY := y.Spec.AzureBindingRef
@@ -351,28 +394,32 @@ func (c *Client) matchAssignedID(x *aadpodid.AzureAssignedIdentity, y *aadpodid.
 	return false, nil
 }
 
-func (c *Client) splitAzureAssignedIDs(old *[]aadpodid.AzureAssignedIdentity, new *[]aadpodid.AzureAssignedIdentity) (retCreate *[]aadpodid.AzureAssignedIdentity, retDelete *[]aadpodid.AzureAssignedIdentity, err error) {
-
-	if old == nil || len(*old) == 0 {
-		return new, nil, nil
+func (c *Client) getAzureAssignedIDsToCreate(old, new []aadpodid.AzureAssignedIdentity) (*[]aadpodid.AzureAssignedIdentity, error) {
+	// everything in new needs to be created
+	if old == nil || len(old) == 0 {
+		return &new, nil
 	}
 
 	create := make([]aadpodid.AzureAssignedIdentity, 0)
-	delete := make([]aadpodid.AzureAssignedIdentity, 0)
-
+	var err error
 	idMatch := false
 	begin := time.Now()
 	// TODO: We should be able to optimize the many for loops.
-	for _, newAssignedID := range *new {
+	for _, newAssignedID := range new {
 		idMatch = false
-		for _, oldAssignedID := range *old {
+		for _, oldAssignedID := range old {
 			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 {
+				// 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 == IdentityCreated {
+					create = append(create, oldAssignedID)
+				}
 				break
 			}
 		}
@@ -384,8 +431,24 @@ func (c *Client) splitAzureAssignedIDs(old *[]aadpodid.AzureAssignedIdentity, ne
 		}
 	}
 	stats.Put(stats.FindAssignedIDCreate, time.Since(begin))
+	return &create, nil
+}
+
+func (c *Client) getAzureAssignedIDsToDelete(old *[]aadpodid.AzureAssignedIdentity, new *[]aadpodid.AzureAssignedIdentity) (*[]aadpodid.AzureAssignedIdentity, error) {
+	// nothing to delete
+	if old == nil || len(*old) == 0 {
+		return nil, nil
+	}
+	// delete everything as nothing in new
+	if new == nil || len(*new) == 0 {
+		return old, nil
+	}
 
-	begin = time.Now()
+	delete := make([]aadpodid.AzureAssignedIdentity, 0)
+	var err error
+	idMatch := false
+
+	begin := time.Now()
 	for _, oldAssignedID := range *old {
 		idMatch = false
 		for _, newAssignedID := range *new {
@@ -407,12 +470,10 @@ func (c *Client) splitAzureAssignedIDs(old *[]aadpodid.AzureAssignedIdentity, ne
 		}
 	}
 	stats.Put(stats.FindAssignedIDDel, time.Since(begin))
-
-	//	glog.Info("Time taken to split create/delete list: %s", time.Since(begin).String())
-	return &create, &delete, nil
+	return &delete, nil
 }
 
-func (c *Client) makeAssignedIDs(azID *aadpodid.AzureIdentity, azBinding *aadpodid.AzureIdentityBinding, podName string, podNameSpace string, nodeName string) (res *aadpodid.AzureAssignedIdentity, err error) {
+func (c *Client) makeAssignedIDs(azID *aadpodid.AzureIdentity, azBinding *aadpodid.AzureIdentityBinding, podName, podNameSpace, nodeName string) (res *aadpodid.AzureAssignedIdentity, err error) {
 	assignedID := &aadpodid.AzureAssignedIdentity{
 		ObjectMeta: v1.ObjectMeta{
 			Name: c.getAssignedIDName(podName, podNameSpace, azID.Name),
@@ -432,7 +493,7 @@ func (c *Client) makeAssignedIDs(azID *aadpodid.AzureIdentity, azBinding *aadpod
 	if c.IsNamespaced || aadpodid.IsNamespacedIdentity(azID) {
 		assignedID.Namespace = azID.Namespace
 	} else {
-		// evantually this should be identity namespace
+		// eventually this should be identity namespace
 		// but to maintain back compat we will use existing
 		// behavior
 		assignedID.Namespace = "default"
@@ -442,57 +503,61 @@ func (c *Client) makeAssignedIDs(azID *aadpodid.AzureIdentity, azBinding *aadpod
 	return assignedID, nil
 }
 
-func (c *Client) createAssignedIdentityDeps(assignedID *aadpodid.AzureAssignedIdentity, id *aadpodid.AzureIdentity, node *corev1.Node) error {
+func (c *Client) createAssignedIdentity(assignedID *aadpodid.AzureAssignedIdentity) error {
 	err := c.CRDClient.CreateAssignedIdentity(assignedID)
 	if err != nil {
-		glog.Error(err)
 		return err
 	}
-
-	if id.Spec.Type == aadpodid.UserAssignedMSI {
-
-		err = c.CloudClient.AssignUserMSI(id.Spec.ResourceID, node)
-		if err != nil {
-			glog.Error(err)
-			newErr := c.CRDClient.RemoveAssignedIdentity(assignedID)
-			if newErr != nil {
-				glog.Errorf("Error when removing assigned identity in create error path err: %v", newErr)
-			}
-			return err
-		}
-	}
 	return nil
 }
 
-func (c *Client) removeAssignedIDsWithDeps(assignedID *aadpodid.AzureAssignedIdentity, inUse bool) error {
+func (c *Client) removeAssignedIdentity(assignedID *aadpodid.AzureAssignedIdentity) error {
 	err := c.CRDClient.RemoveAssignedIdentity(assignedID)
 	if err != nil {
-		glog.Error(err)
-		return nil
+		return err
 	}
-	if !inUse {
-		id := assignedID.Spec.AzureIdentityRef
-		if id.Spec.Type == aadpodid.UserAssignedMSI {
+	return nil
+}
 
-			node, err := c.NodeClient.Get(assignedID.Spec.NodeName)
-			if err != nil {
-				return err
-			}
+func (c *Client) appendToRemoveListForNode(resourceID, nodeName string, nodeMap map[string]trackUserAssignedMSIIds) {
+	if trackList, ok := nodeMap[nodeName]; ok {
+		trackList.removeUserAssignedMSIIDs = append(trackList.removeUserAssignedMSIIDs, resourceID)
+		nodeMap[nodeName] = trackList
+		return
+	}
+	nodeMap[nodeName] = trackUserAssignedMSIIds{removeUserAssignedMSIIDs: []string{resourceID}}
+}
 
-			err = c.CloudClient.RemoveUserMSI(id.Spec.ResourceID, node)
-			if err != nil {
-				glog.Error(err)
-				return err
-			}
-		}
+func (c *Client) appendToAddListForNode(resourceID, nodeName string, nodeMap map[string]trackUserAssignedMSIIds) {
+	if trackList, ok := nodeMap[nodeName]; ok {
+		trackList.addUserAssignedMSIIDs = append(trackList.addUserAssignedMSIIDs, resourceID)
+		nodeMap[nodeName] = trackList
+		return
 	}
-	return nil
+	nodeMap[nodeName] = trackUserAssignedMSIIds{addUserAssignedMSIIDs: []string{resourceID}}
+}
+
+func (c *Client) checkIfUserAssignedMSI(id *aadpodid.AzureIdentity) bool {
+	return id.Spec.Type == aadpodid.UserAssignedMSI
 }
 
-func (c *Client) getAssignedIDName(podName string, podNameSpace string, idName string) string {
+func (c *Client) getAssignedIDName(podName, podNameSpace, idName string) string {
 	return podName + "-" + podNameSpace + "-" + idName
 }
 
+func (c *Client) checkIfMSIExistsOnNode(id *aadpodid.AzureIdentity, nodeName string, nodeMSIList []string) bool {
+	for _, userAssignedMSI := range nodeMSIList {
+		if userAssignedMSI == id.Spec.ResourceID {
+			return true
+		}
+	}
+	return false
+}
+
+func (c *Client) getUserMSIListForNode(node *corev1.Node) ([]string, error) {
+	return c.CloudClient.GetUserMSIs(node)
+}
+
 func (c *Client) convertIDListToMap(arr []aadpodid.AzureIdentity) (m map[string]aadpodid.AzureIdentity, err error) {
 	m = make(map[string]aadpodid.AzureIdentity, len(arr))
 	for _, element := range arr {
@@ -541,3 +606,168 @@ func (c *Client) checkIfInUse(checkAssignedID aadpodid.AzureAssignedIdentity, ar
 
 	return false, nil
 }
+
+func (c *Client) getUniqueIDs(idList []string) []string {
+	idSet := make(map[string]struct{})
+	var uniqueList []string
+
+	for _, id := range idList {
+		idSet[id] = struct{}{}
+	}
+	for id := range idSet {
+		uniqueList = append(uniqueList, id)
+	}
+	return uniqueList
+}
+
+func (c *Client) updateAssignedIdentityStatus(assignedID *aadpodid.AzureAssignedIdentity, status string) error {
+	return c.CRDClient.UpdateAzureAssignedIdentityStatus(assignedID, status)
+}
+
+func (c *Client) updateNodeAndDeps(newAssignedIDs []aadpodid.AzureAssignedIdentity, nodeMap map[string]trackUserAssignedMSIIds, nodeRefs map[string]bool) {
+	for nodeName, nodeTrackList := range nodeMap {
+		c.updateUserMSI(newAssignedIDs, nodeName, nodeTrackList, nodeRefs)
+	}
+}
+
+func (c *Client) updateUserMSI(newAssignedIDs []aadpodid.AzureAssignedIdentity, nodeName string, nodeTrackList trackUserAssignedMSIIds, nodeRefs map[string]bool) {
+	beginAdding := time.Now()
+	glog.V(5).Infof("Processing node %s", nodeName)
+
+	for _, createID := range nodeTrackList.assignedIDsToCreate {
+		if createID.Status.Status == "" {
+			binding := createID.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)
+
+			createID.Status.Status = IdentityCreated
+			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)
+			}
+		}
+	}
+	// generate unique list so we don't make multiple calls to assign/remove same id
+	addUserAssignedMSIIDs := c.getUniqueIDs(nodeTrackList.addUserAssignedMSIIDs)
+	removeUserAssignedMSIIDs := c.getUniqueIDs(nodeTrackList.removeUserAssignedMSIIDs)
+
+	node, err := c.NodeClient.Get(nodeName)
+	if err != nil {
+		glog.Errorf("Lookup of node %s resulted in error %v", nodeName, err)
+		return
+	}
+
+	err = c.CloudClient.UpdateUserMSI(addUserAssignedMSIIDs, removeUserAssignedMSIIDs, node)
+	if err != nil {
+		idList, getErr := c.getUserMSIListForNode(node)
+		if getErr != nil {
+			glog.Errorf("Getting list of msi's from node %s resulted in error %v", nodeName, getErr)
+			return
+		}
+
+		for _, createID := range nodeTrackList.assignedIDsToCreate {
+			id := createID.Spec.AzureIdentityRef
+			binding := createID.Spec.AzureBindingRef
+
+			isUserAssignedMSI := c.checkIfUserAssignedMSI(id)
+			idExistsOnNode := c.checkIfMSIExistsOnNode(id, createID.Spec.NodeName, idList)
+
+			if isUserAssignedMSI && !idExistsOnNode {
+				message := fmt.Sprintf("Applying binding %s node %s for pod %s resulted in error %v", binding.Name, createID.Spec.NodeName, createID.Name, err.Error())
+				c.EventRecorder.Event(binding, corev1.EventTypeWarning, "binding apply error", message)
+				glog.Error(message)
+				continue
+			}
+			// the identity was successfully assigned to the node
+			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))
+
+			// Identity is successfully assigned to node, so update the status of assigned identity to assigned
+			if updateErr := c.updateAssignedIdentityStatus(&createID, IdentityAssigned); updateErr != nil {
+				message := fmt.Sprintf("Updating assigned identity %s status to %s for pod %s failed with error %v", createID.Name, IdentityAssigned, createID.Spec.Pod, updateErr.Error())
+				c.EventRecorder.Event(&createID, corev1.EventTypeWarning, "status update error", message)
+				glog.Error(message)
+			}
+		}
+
+		for _, delID := range nodeTrackList.assignedIDsToDelete {
+			id := delID.Spec.AzureIdentityRef
+			removedBinding := delID.Spec.AzureBindingRef
+			isUserAssignedMSI := c.checkIfUserAssignedMSI(id)
+			idExistsOnNode := c.checkIfMSIExistsOnNode(id, delID.Spec.NodeName, idList)
+			vmssGroups, getErr := getVMSSGroups(c.NodeClient, nodeRefs)
+			if getErr != nil {
+				glog.Error(getErr)
+				continue
+			}
+			inUse, checkErr := c.checkIfInUse(delID, newAssignedIDs, vmssGroups)
+			if checkErr != nil {
+				glog.Error(checkErr)
+				continue
+			}
+			// the identity still exists on node, which means removing the identity from the node failed
+			if isUserAssignedMSI && !inUse && idExistsOnNode {
+				message := fmt.Sprintf("Binding %s removal from node %s for pod %s resulted in error %v", removedBinding.Name, delID.Spec.NodeName, delID.Spec.Pod, err.Error())
+				c.EventRecorder.Event(removedBinding, corev1.EventTypeWarning, "binding remove 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.Error()))
+				glog.Error(err)
+				continue
+			}
+			// 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))
+		}
+		stats.Put(stats.TotalCreateOrUpdate, time.Since(beginAdding))
+		return
+	}
+
+	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, IdentityAssigned)
+		if err != nil {
+			message := fmt.Sprintf("Updating assigned identity %s status to %s for pod %s failed with error %v", createID.Name, IdentityAssigned, createID.Spec.Pod, err.Error())
+			c.EventRecorder.Event(&createID, corev1.EventTypeWarning, "status update error", message)
+			glog.Error(message)
+			continue
+		}
+		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))
+	}
+
+	for _, delID := range nodeTrackList.assignedIDsToDelete {
+		removedBinding := delID.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(&delID, IdentityUnassigned)
+		if err != nil {
+			message := fmt.Sprintf("Updating assigned identity %s status to %s for pod %s failed with error %v", delID.Name, IdentityUnassigned, 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
+		}
+		// 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))
+	}
+	stats.Put(stats.TotalCreateOrUpdate, time.Since(beginAdding))
+}
diff --git a/pkg/stats/stats.go b/pkg/stats/stats.go
index 3a0449b6..d418c2fb 100644
--- a/pkg/stats/stats.go
+++ b/pkg/stats/stats.go
@@ -28,6 +28,7 @@ const (
 	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"
 )
@@ -77,8 +78,7 @@ func PrintSync() {
 		Print(FindAssignedIDCreate)
 		Print(FindAssignedIDDel)
 
-		Print(TotalIDAdd)
-		Print(TotalIDDel)
+		Print(TotalCreateOrUpdate)
 
 		Print(EventRecord)
 		Print(Total)


Test patch:
diff --git a/pkg/cloudprovider/cloudprovider_test.go b/pkg/cloudprovider/cloudprovider_test.go
index a138e497..32c03582 100644
--- a/pkg/cloudprovider/cloudprovider_test.go
+++ b/pkg/cloudprovider/cloudprovider_test.go
@@ -118,6 +118,28 @@ func TestSimple(t *testing.T) {
 				cloudClient.PrintMSI(t)
 				t.Error("MSI mismatch")
 			}
+
+			// test the UpdateUserMSI interface
+			cloudClient.UpdateUserMSI([]string{"ID1", "ID2", "ID3"}, []string{"ID0again"}, node0)
+			testMSI = []string{"ID1", "ID2", "ID3"}
+			if !cloudClient.CompareMSI(node0, testMSI) {
+				cloudClient.PrintMSI(t)
+				t.Error("MSI mismatch")
+			}
+
+			cloudClient.UpdateUserMSI(nil, []string{"ID3"}, node3)
+			testMSI = []string{}
+			if !cloudClient.CompareMSI(node3, testMSI) {
+				cloudClient.PrintMSI(t)
+				t.Error("MSI mismatch")
+			}
+
+			cloudClient.UpdateUserMSI([]string{"ID3"}, nil, node4)
+			testMSI = []string{"ID4", "ID3"}
+			if !cloudClient.CompareMSI(node4, testMSI) {
+				cloudClient.PrintMSI(t)
+				t.Error("MSI mismatch")
+			}
 		})
 	}
 }
diff --git a/pkg/cloudprovider/identity_test.go b/pkg/cloudprovider/identity_test.go
index ab3d05bb..9296bda1 100644
--- a/pkg/cloudprovider/identity_test.go
+++ b/pkg/cloudprovider/identity_test.go
@@ -9,8 +9,8 @@ import (
 func TestFilterIdentity(t *testing.T) {
 	idList := []string{}
 	idType := compute.ResourceIdentityTypeNone
-	if err := filterUserIdentity(&idType, &idList, "A"); err == nil || err != errNotFound {
-		t.Fatalf("expected error %q, got: %v", errNotFound, err)
+	if err := filterUserIdentity(&idType, &idList, "A"); err != nil {
+		t.Fatalf("expected nil error, got: %v", err)
 	}
 
 	idType = compute.ResourceIdentityTypeUserAssigned
diff --git a/pkg/mic/mic_test.go b/pkg/mic/mic_test.go
index 76ebf6a3..bf3ae042 100644
--- a/pkg/mic/mic_test.go
+++ b/pkg/mic/mic_test.go
@@ -33,9 +33,10 @@ type TestCloudClient struct {
 type TestVMClient struct {
 	*cp.VMClient
 
-	mu      sync.Mutex
-	nodeMap map[string]*compute.VirtualMachine
-	err     *error
+	mu       sync.Mutex
+	nodeMap  map[string]*compute.VirtualMachine
+	err      *error
+	identity *compute.VirtualMachineIdentity
 }
 
 func (c *TestVMClient) SetError(err error) {
@@ -64,9 +65,9 @@ func (c *TestVMClient) CreateOrUpdate(rg string, nodeName string, vm compute.Vir
 	defer c.mu.Unlock()
 
 	if c.err != nil {
+		c.nodeMap[nodeName].Identity = c.identity
 		return *c.err
 	}
-
 	c.nodeMap[nodeName] = &vm
 	return nil
 }
@@ -78,7 +79,9 @@ func (c *TestVMClient) ListMSI() (ret map[string]*[]string) {
 	defer c.mu.Unlock()
 
 	for key, val := range c.nodeMap {
-		ret[key] = val.Identity.IdentityIds
+		if val.Identity != nil {
+			ret[key] = val.Identity.IdentityIds
+		}
 	}
 	return ret
 }
@@ -105,9 +108,10 @@ func (c *TestVMClient) CompareMSI(nodeName string, userIDs []string) bool {
 type TestVMSSClient struct {
 	*cp.VMSSClient
 
-	mu      sync.Mutex
-	nodeMap map[string]*compute.VirtualMachineScaleSet
-	err     *error
+	mu       sync.Mutex
+	nodeMap  map[string]*compute.VirtualMachineScaleSet
+	err      *error
+	identity *compute.VirtualMachineScaleSetIdentity
 }
 
 func (c *TestVMSSClient) SetError(err error) {
@@ -136,6 +140,7 @@ func (c *TestVMSSClient) CreateOrUpdate(rg string, nodeName string, vm compute.V
 	defer c.mu.Unlock()
 
 	if c.err != nil {
+		c.nodeMap[nodeName].Identity = c.identity
 		return *c.err
 	}
 	c.nodeMap[nodeName] = &vm
@@ -209,20 +214,24 @@ func (c *TestCloudClient) UnSetError() {
 func NewTestVMClient() *TestVMClient {
 	nodeMap := make(map[string]*compute.VirtualMachine)
 	vmClient := &cp.VMClient{}
+	identity := &compute.VirtualMachineIdentity{IdentityIds: &[]string{}}
 
 	return &TestVMClient{
 		VMClient: vmClient,
 		nodeMap:  nodeMap,
+		identity: identity,
 	}
 }
 
 func NewTestVMSSClient() *TestVMSSClient {
 	nodeMap := make(map[string]*compute.VirtualMachineScaleSet)
 	vmssClient := &cp.VMSSClient{}
+	identity := &compute.VirtualMachineScaleSetIdentity{IdentityIds: &[]string{}}
 
 	return &TestVMSSClient{
 		VMSSClient: vmssClient,
 		nodeMap:    nodeMap,
+		identity:   identity,
 	}
 }
 
@@ -270,7 +279,7 @@ func (c *TestPodClient) GetPods() ([]*corev1.Pod, error) {
 	return pods, nil
 }
 
-func (c *TestPodClient) AddPod(podName string, podNs string, nodeName string, binding string) {
+func (c *TestPodClient) AddPod(podName, podNs, nodeName, binding string) {
 	labels := make(map[string]string)
 	labels[aadpodid.CRDLabelKey] = binding
 	pod := &corev1.Pod{
@@ -317,6 +326,7 @@ type TestCrdClient struct {
 	assignedIDMap map[string]*aadpodid.AzureAssignedIdentity
 	bindingMap    map[string]*aadpodid.AzureIdentityBinding
 	idMap         map[string]*aadpodid.AzureIdentity
+	err           *error
 }
 
 func NewTestCrdClient(config *rest.Config) *TestCrdClient {
@@ -340,8 +350,12 @@ func (c *TestCrdClient) CreateCrdWatchers(eventCh chan aadpodid.EventType) (err
 
 func (c *TestCrdClient) RemoveAssignedIdentity(assignedIdentity *aadpodid.AzureAssignedIdentity) error {
 	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if c.err != nil {
+		return *c.err
+	}
 	delete(c.assignedIDMap, assignedIdentity.Name)
-	c.mu.Unlock()
 	return nil
 }
 
@@ -355,6 +369,13 @@ func (c *TestCrdClient) CreateAssignedIdentity(assignedIdentity *aadpodid.AzureA
 	return nil
 }
 
+func (c *TestCrdClient) UpdateAzureAssignedIdentityStatus(assignedIdentity *aadpodid.AzureAssignedIdentity, status string) error {
+	assignedIdentity.Status.Status = status
+	assignedIdentityToStore := *assignedIdentity //Make a copy to store in the map.
+	c.assignedIDMap[assignedIdentity.Name] = &assignedIdentityToStore
+	return nil
+}
+
 func (c *TestCrdClient) CreateBinding(bindingName string, idName string, selector string) {
 	binding := &aadpodid.AzureIdentityBinding{
 		ObjectMeta: v1.ObjectMeta{
@@ -419,10 +440,19 @@ func (c *TestCrdClient) ListAssignedIDs() (res *[]aadpodid.AzureAssignedIdentity
 	c.mu.Unlock()
 	return &assignedIDList, nil
 }
+
 func (c *Client) ListPodIds(podns, podname string) (*[]aadpodid.AzureIdentity, error) {
 	return &[]aadpodid.AzureIdentity{}, nil
 }
 
+func (c *TestCrdClient) SetError(err error) {
+	c.err = &err
+}
+
+func (c *TestCrdClient) UnSetError() {
+	c.err = nil
+}
+
 /************************ NODE MOCK *************************************/
 
 type TestNodeClient struct {
@@ -701,6 +731,7 @@ func TestSimpleMICClient(t *testing.T) {
 	cloudClient.SetError(err)
 
 	podClient.AddPod("test-pod", "default", "test-node", "test-select")
+
 	eventCh <- aadpodid.PodCreated
 	evtRecorder.WaitForEvents(1)
 
@@ -710,8 +741,8 @@ func TestSimpleMICClient(t *testing.T) {
 		t.Fatalf("list assigned failed")
 	}
 
-	if len(*listAssignedIDs) != 0 {
-		t.Fatalf("ID assigned")
+	if (*listAssignedIDs)[0].Status.Status != IdentityCreated {
+		t.Fatalf("expected status to be %s, got: %s", IdentityCreated, (*listAssignedIDs)[0].Status.Status)
 	}
 
 	/*
@@ -723,7 +754,7 @@ func TestSimpleMICClient(t *testing.T) {
 		} */
 
 	// Test4: Removal error event test
-	//Reset the state to add the id.
+	// Reset the state to add the id.
 	cloudClient.UnSetError()
 
 	//podClient.AddPod("test-pod", "default", "test-node", "test-select")
@@ -793,7 +824,7 @@ func TestAddDelMICClient(t *testing.T) {
 		t.Fatalf("Add and delete id at same time mismatch")
 	}
 
-	//Delete the pod
+	// Delete the pod
 	podClient.DeletePod("test-pod2", "default")
 	podClient.DeletePod("test-pod4", "default")
 
@@ -877,6 +908,14 @@ func TestMicAddDelVMSS(t *testing.T) {
 	if !evtRecorder.WaitForEvents(3) {
 		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) == 3) {
+		t.Fatalf("expected assigned identities len: %d, got: %d", 3, len(*listAssignedIDs))
+	}
 
 	if !cloudClient.CompareMSI("testvmss1", []string{"test-user-msi-resourceid"}) {
 		t.Fatalf("missing identity: %+v", cloudClient.ListMSI()["testvmss1"])
@@ -891,6 +930,14 @@ func TestMicAddDelVMSS(t *testing.T) {
 	if !evtRecorder.WaitForEvents(1) {
 		t.Fatal("Timeout waiting for mic sync cycles")
 	}
+	listAssignedIDs, err = crdClient.ListAssignedIDs()
+	if err != nil {
+		glog.Error(err)
+		t.Errorf("list assigned failed")
+	}
+	if !(len(*listAssignedIDs) == 2) {
+		t.Fatalf("expected assigned identities len: %d, got: %d", 2, len(*listAssignedIDs))
+	}
 
 	if !cloudClient.CompareMSI("testvmss1", []string{"test-user-msi-resourceid"}) {
 		t.Fatalf("missing identity: %+v", cloudClient.ListMSI()["testvmss1"])
@@ -900,10 +947,20 @@ func TestMicAddDelVMSS(t *testing.T) {
 	}
 
 	podClient.DeletePod("test-pod2", "default")
+
 	eventCh <- aadpodid.PodDeleted
+
 	if !evtRecorder.WaitForEvents(1) {
 		t.Fatal("Timeout waiting for mic sync cycles")
 	}
+	listAssignedIDs, err = crdClient.ListAssignedIDs()
+	if err != nil {
+		glog.Error(err)
+		t.Errorf("list assigned failed")
+	}
+	if !(len(*listAssignedIDs) == 1) {
+		t.Fatalf("expected assigned identities len: %d, got: %d", 1, len(*listAssignedIDs))
+	}
 
 	if !cloudClient.CompareMSI("testvmss1", []string{}) {
 		t.Fatalf("missing identity: %+v", cloudClient.ListMSI()["testvmss1"])
@@ -913,6 +970,115 @@ func TestMicAddDelVMSS(t *testing.T) {
 	}
 }
 
+func TestMICStateFlow(t *testing.T) {
+	eventCh := make(chan aadpodid.EventType, 100)
+	cloudClient := NewTestCloudClient(config.AzureConfig{})
+	crdClient := NewTestCrdClient(nil)
+	podClient := NewTestPodClient()
+	nodeClient := NewTestNodeClient()
+	var evtRecorder TestEventRecorder
+	evtRecorder.lastEvent = new(LastEvent)
+	evtRecorder.eventChannel = make(chan bool, 100)
+
+	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder)
+
+	// Add a pod, identity and binding.
+	crdClient.CreateID("test-id1", aadpodid.UserAssignedMSI, "test-user-msi-resourceid", "test-user-msi-clientid", nil, "", "", "")
+	crdClient.CreateBinding("testbinding1", "test-id1", "test-select1")
+
+	nodeClient.AddNode("test-node1")
+	podClient.AddPod("test-pod1", "default", "test-node1", "test-select1")
+
+	eventCh <- aadpodid.PodCreated
+	defer micClient.testRunSync()(t)
+
+	if !evtRecorder.WaitForEvents(1) {
+		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) == 1) {
+		t.Fatalf("expected assigned identities len: %d, got: %d", 1, len(*listAssignedIDs))
+	}
+	if !((*listAssignedIDs)[0].Status.Status == IdentityAssigned) {
+		t.Fatalf("expected status to be %s, got: %s", IdentityCreated, (*listAssignedIDs)[0].Status.Status)
+	}
+
+	// delete the pod, simulate failure in cloud calls on trying to un-assign identity from node
+	podClient.DeletePod("test-pod1", "default")
+	cloudClient.SetError(errors.New("error removing identity from node"))
+	cloudClient.testVMClient.identity = nil
+
+	eventCh <- aadpodid.PodDeleted
+
+	listAssignedIDs, err = crdClient.ListAssignedIDs()
+	if err != nil {
+		glog.Error(err)
+		t.Errorf("list assigned failed")
+	}
+	if !(len(*listAssignedIDs) == 1) {
+		t.Fatalf("expected assigned identities len: %d, got: %d", 1, len(*listAssignedIDs))
+	}
+	if !((*listAssignedIDs)[0].Status.Status == IdentityAssigned) {
+		t.Fatalf("expected status to be %s, got: %s", IdentityAssigned, (*listAssignedIDs)[0].Status.Status)
+	}
+
+	cloudClient.UnSetError()
+	crdClient.SetError(errors.New("error from crd client"))
+
+	// add new pod, this time the old assigned identity which is in Assigned state should be tried to delete
+	// simulate failure on kube api call to delete crd
+	crdClient.CreateID("test-id2", aadpodid.UserAssignedMSI, "test-user-msi-resourceid2", "test-user-msi-clientid2", nil, "", "", "")
+	crdClient.CreateBinding("testbinding2", "test-id2", "test-select2")
+
+	nodeClient.AddNode("test-node2")
+	podClient.AddPod("test-pod2", "default", "test-node2", "test-select2")
+
+	eventCh <- aadpodid.PodCreated
+	if !evtRecorder.WaitForEvents(2) {
+		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) == 2) {
+		t.Fatalf("expected assigned identities len: %d, got: %d", 2, len(*listAssignedIDs))
+	}
+	for _, assignedID := range *listAssignedIDs {
+		if assignedID.Spec.Pod == "test-pod1" {
+			if assignedID.Status.Status != IdentityUnassigned {
+				t.Fatalf("Expected status to be: %s. Got: %s", IdentityUnassigned, assignedID.Status.Status)
+			}
+		}
+		if assignedID.Spec.Pod == "test-pod2" {
+			if assignedID.Status.Status != IdentityAssigned {
+				t.Fatalf("Expected status to be: %s. Got: %s", IdentityAssigned, assignedID.Status.Status)
+			}
+		}
+	}
+	crdClient.UnSetError()
+
+	// delete pod2 and everything should be cleaned up now
+	podClient.DeletePod("test-pod2", "default")
+	eventCh <- aadpodid.PodDeleted
+	if !evtRecorder.WaitForEvents(2) {
+		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) == 0) {
+		t.Fatalf("expected assigned identities len: %d, got: %d", 0, len(*listAssignedIDs))
+	}
+}
+
 func TestSyncExit(t *testing.T) {
 	eventCh := make(chan aadpodid.EventType)
 	cloudClient := NewTestCloudClient(config.AzureConfig{VMType: "vmss"})
diff --git a/test/common/k8s/azureassignedidentity/azureassignedidentity.go b/test/common/k8s/azureassignedidentity/azureassignedidentity.go
index a25b8995..f0dd38b9 100644
--- a/test/common/k8s/azureassignedidentity/azureassignedidentity.go
+++ b/test/common/k8s/azureassignedidentity/azureassignedidentity.go
@@ -46,10 +46,32 @@ func GetByPrefix(prefix string) (aadpodid.AzureAssignedIdentity, error) {
 	return aadpodid.AzureAssignedIdentity{}, errors.Errorf("No AzureAssignedIdentity has a name prefix with '%s'", prefix)
 }
 
+// GetAllByPrefix will return all AzureAssignedIdentity with matched prefix
+func GetAllByPrefix(prefix string) ([]aadpodid.AzureAssignedIdentity, error) {
+	var aaidList []aadpodid.AzureAssignedIdentity
+	list, err := GetAll()
+	if err != nil {
+		return aaidList, err
+	}
+
+	for _, azureAssignedIdentity := range list.Items {
+		if strings.HasPrefix(azureAssignedIdentity.ObjectMeta.Name, prefix) {
+			aaidList = append(aaidList, azureAssignedIdentity)
+		}
+	}
+
+	if len(aaidList) < 1 {
+		err = errors.Errorf("No AzureAssignedIdentity has a name prefix with '%s'", prefix)
+	}
+
+	return aaidList, err
+}
+
 // WaitOnLengthMatched will block until the number of Azure Assigned Identity matches the target
 func WaitOnLengthMatched(target int) (bool, error) {
 	successChannel, errorChannel := make(chan bool, 1), make(chan error)
-	duration, sleep := 70*time.Second, 10*time.Second
+	// defining ~2 mins as an acceptable timeframe for ids to be assigned to node
+	duration, sleep := 120*time.Second, 10*time.Second
 	ctx, cancel := context.WithTimeout(context.Background(), duration)
 	defer cancel()
 
diff --git a/test/common/k8s/infra/infra.go b/test/common/k8s/infra/infra.go
index 89dd3cc3..59b5435f 100644
--- a/test/common/k8s/infra/infra.go
+++ b/test/common/k8s/infra/infra.go
@@ -60,13 +60,13 @@ func CreateInfra(namespace, registry, nmiVersion, micVersion, templateOutputPath
 }
 
 // CreateIdentityValidator will create an identity validator deployment on a Kubernetes cluster
-func CreateIdentityValidator(subscriptionID, resourceGroup, registryName, name, identityBinding, identityValidatorVersion, templateOutputPath string) error {
+func CreateIdentityValidator(subscriptionID, resourceGroup, registryName, deploymentName, identityBinding, identityValidatorVersion, templateOutputPath, replicas string) error {
 	t, err := template.New("deployment.yaml").ParseFiles(path.Join("template", "deployment.yaml"))
 	if err != nil {
 		return errors.Wrap(err, "Failed to parse deployment.yaml")
 	}
 
-	deployFilePath := path.Join(templateOutputPath, name+"-deployment.yaml")
+	deployFilePath := path.Join(templateOutputPath, deploymentName+"-deployment.yaml")
 	deployFile, err := os.Create(deployFilePath)
 	if err != nil {
 		return errors.Wrap(err, "Failed to create a deployment file from deployment.yaml")
@@ -79,11 +79,13 @@ func CreateIdentityValidator(subscriptionID, resourceGroup, registryName, name,
 		IdentityBinding          string
 		Registry                 string
 		IdentityValidatorVersion string
+		Replicas                 string
 	}{
-		name,
+		deploymentName,
 		identityBinding,
 		registryName,
 		identityValidatorVersion,
+		replicas,
 	}
 	if err := t.Execute(deployFile, deployData); err != nil {
 		return errors.Wrap(err, "Failed to create a deployment file from deployment.yaml")
diff --git a/test/e2e/aadpodidentity_test.go b/test/e2e/aadpodidentity_test.go
index 4c26a06f..fb383dfe 100644
--- a/test/e2e/aadpodidentity_test.go
+++ b/test/e2e/aadpodidentity_test.go
@@ -124,7 +124,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 	})
 
 	It("should pass the identity validating test", func() {
-		setUpIdentityAndDeployment(keyvaultIdentity, "")
+		setUpIdentityAndDeployment(keyvaultIdentity, "", "1")
 
 		ok, err := azureassignedidentity.WaitOnLengthMatched(1)
 		Expect(err).NotTo(HaveOccurred())
@@ -137,7 +137,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 	})
 
 	It("should not pass the identity validating test if the AzureIdentity is deleted", func() {
-		setUpIdentityAndDeployment(keyvaultIdentity, "")
+		setUpIdentityAndDeployment(keyvaultIdentity, "", "1")
 
 		err := azureidentity.DeleteOnCluster(keyvaultIdentity, templateOutputPath)
 		Expect(err).NotTo(HaveOccurred())
@@ -159,7 +159,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 	})
 
 	It("should not pass the identity validating test if the AzureIdentityBinding is deleted", func() {
-		setUpIdentityAndDeployment(keyvaultIdentity, "")
+		setUpIdentityAndDeployment(keyvaultIdentity, "", "1")
 
 		err := azureidentitybinding.Delete(keyvaultIdentity, templateOutputPath)
 		Expect(err).NotTo(HaveOccurred())
@@ -181,7 +181,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 	})
 
 	It("should delete the AzureAssignedIdentity if the deployment is deleted", func() {
-		setUpIdentityAndDeployment(keyvaultIdentity, "")
+		setUpIdentityAndDeployment(keyvaultIdentity, "", "1")
 
 		waitForDeployDeletion(identityValidator)
 
@@ -191,7 +191,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 	})
 
 	It("should establish a new AzureAssignedIdentity and remove the old one when draining the node containing identity validator", func() {
-		setUpIdentityAndDeployment(keyvaultIdentity, "")
+		setUpIdentityAndDeployment(keyvaultIdentity, "", "1")
 
 		podName, err := pod.GetNameByPrefix(identityValidator)
 		Expect(err).NotTo(HaveOccurred())
@@ -230,8 +230,11 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 			identityName := fmt.Sprintf("%s-%d", keyvaultIdentity, i)
 			identityValidatorName := fmt.Sprintf("identity-validator-%d", i)
 
-			setUpIdentityAndDeployment(keyvaultIdentity, fmt.Sprintf("%d", i))
-			time.Sleep(5 * time.Second)
+			setUpIdentityAndDeployment(keyvaultIdentity, fmt.Sprintf("%d", i), "1")
+
+			ok, err := azureassignedidentity.WaitOnLengthMatched(i + 1)
+			Expect(ok).To(Equal(true))
+			Expect(err).NotTo(HaveOccurred())
 
 			identityClientID, err := azure.GetIdentityClientID(cfg.ResourceGroup, identityName)
 			Expect(err).NotTo(HaveOccurred())
@@ -339,7 +342,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 
 		enableUserAssignedIdentityOnCluster(nodeList, clusterIdentity)
 
-		setUpIdentityAndDeployment(keyvaultIdentity, "")
+		setUpIdentityAndDeployment(keyvaultIdentity, "", "1")
 
 		podName, err := pod.GetNameByPrefix(identityValidator)
 		Expect(err).NotTo(HaveOccurred())
@@ -489,7 +492,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 
 		enableSystemAssignedIdentityOnCluster(nodeList)
 
-		setUpIdentityAndDeployment(keyvaultIdentity, "")
+		setUpIdentityAndDeployment(keyvaultIdentity, "", "1")
 
 		podName, err := pod.GetNameByPrefix(identityValidator)
 		Expect(err).NotTo(HaveOccurred())
@@ -524,12 +527,42 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 
 		removeSystemAssignedIdentityOnCluster(nodeList)
 	})
+
+	It("should create azureassignedidentities for 40 pods within ~2mins", func() {
+		// setup all the 40 pods in a loop to ensure mic handles
+		// scale out efficiently
+		for i := 0; i < 5; i++ {
+			setUpIdentityAndDeployment(keyvaultIdentity, fmt.Sprintf("%d", i), "8")
+		}
+
+		// WaitOnLengthMatched waits for 2 mins, so this will ensure we are performant at high scale
+		ok, err := azureassignedidentity.WaitOnLengthMatched(40)
+		Expect(ok).To(Equal(true))
+		Expect(err).NotTo(HaveOccurred())
+
+		for i := 0; i < 5; i++ {
+			identityName := fmt.Sprintf("%s-%d", keyvaultIdentity, i)
+			identityValidatorName := fmt.Sprintf("identity-validator-%d", i)
+
+			identityClientID, err := azure.GetIdentityClientID(cfg.ResourceGroup, identityName)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(identityClientID).NotTo(Equal(""))
+
+			azureAssignedIdentities, err := azureassignedidentity.GetAllByPrefix(identityValidatorName)
+			Expect(err).NotTo(HaveOccurred())
+
+			for _, azureAssignedIdentity := range azureAssignedIdentities {
+				validateAzureAssignedIdentity(azureAssignedIdentity, identityName)
+			}
+		}
+	})
 })
 
 // setupInfra creates the crds, mic, nmi and blocks until iptable entries exist
 func setupInfra() {
 	// Install CRDs and deploy MIC and NMI
 	err := infra.CreateInfra("default", cfg.Registry, cfg.NMIVersion, cfg.MICVersion, templateOutputPath)
+
 	Expect(err).NotTo(HaveOccurred())
 
 	ok, err := daemonset.WaitOnReady(nmiDaemonSet)
@@ -566,7 +599,8 @@ func setupInfra() {
 
 // setUpIdentityAndDeployment will deploy AzureIdentity, AzureIdentityBinding, and an identity validator
 // Suffix will give the tests the option to add a suffix to the end of the identity name, useful for scale tests
-func setUpIdentityAndDeployment(azureIdentityName, suffix string) {
+// replicas to indicate the number of replicas for the deployment
+func setUpIdentityAndDeployment(azureIdentityName, suffix, replicas string) {
 	identityValidatorName := identityValidator
 
 	if suffix != "" {
@@ -580,7 +614,7 @@ func setUpIdentityAndDeployment(azureIdentityName, suffix string) {
 	err = azureidentitybinding.Create(azureIdentityName, templateOutputPath)
 	Expect(err).NotTo(HaveOccurred())
 
-	err = infra.CreateIdentityValidator(cfg.SubscriptionID, cfg.ResourceGroup, cfg.Registry, identityValidatorName, azureIdentityName, cfg.IdentityValidatorVersion, templateOutputPath)
+	err = infra.CreateIdentityValidator(cfg.SubscriptionID, cfg.ResourceGroup, cfg.Registry, identityValidatorName, azureIdentityName, cfg.IdentityValidatorVersion, templateOutputPath, replicas)
 	Expect(err).NotTo(HaveOccurred())
 
 	ok, err := deploy.WaitOnReady(identityValidatorName)
@@ -592,6 +626,8 @@ func setUpIdentityAndDeployment(azureIdentityName, suffix string) {
 	ok, err = daemonset.WaitOnReady(nmiDaemonSet)
 	Expect(err).NotTo(HaveOccurred())
 	Expect(ok).To(Equal(true))
+
+	time.Sleep(30 * time.Second)
 }
 
 // validateAzureAssignedIdentity will make sure a given AzureAssignedIdentity has the correct properties
diff --git a/test/e2e/template/deployment.yaml b/test/e2e/template/deployment.yaml
index dc1017b4..ee4eb5a1 100644
--- a/test/e2e/template/deployment.yaml
+++ b/test/e2e/template/deployment.yaml
@@ -7,6 +7,7 @@ metadata:
   name: {{.Name}}
   namespace: default
 spec:
+  replicas: {{.Replicas}}
   template:
     metadata:
       labels: