From 8aafccea1a3675558b6b3cd42666d33dd20e8f38 Mon Sep 17 00:00:00 2001
From: "dom.bozzuto" <dom.bozzuto@datadoghq.com>
Date: Fri, 14 Jun 2024 13:30:40 -0400
Subject: [PATCH] Prevent autoscaler from performing multiple deletes

---
 .../cloudprovider/azure/azure_scale_set.go    | 54 +++++++++++++++++--
 1 file changed, 51 insertions(+), 3 deletions(-)

diff --git a/cluster-autoscaler/cloudprovider/azure/azure_scale_set.go b/cluster-autoscaler/cloudprovider/azure/azure_scale_set.go
index 3b7978c5868b..f19b2bdfc5be 100644
--- a/cluster-autoscaler/cloudprovider/azure/azure_scale_set.go
+++ b/cluster-autoscaler/cloudprovider/azure/azure_scale_set.go
@@ -63,6 +63,9 @@ type ScaleSet struct {
 	instanceMutex       sync.Mutex
 	instanceCache       []cloudprovider.Instance
 	lastInstanceRefresh time.Time
+
+	deletionMutex       sync.Mutex
+	deletionsInProgress map[string]struct{}
 }
 
 // NewScaleSet creates a new NewScaleSet.
@@ -78,6 +81,7 @@ func NewScaleSet(spec *dynamic.NodeGroupSpec, az *AzureManager, curSize int64) (
 		sizeRefreshPeriod:         az.azureCache.refreshInterval,
 		enableDynamicInstanceList: az.config.EnableDynamicInstanceList,
 		instancesRefreshJitter:    az.config.VmssVmsCacheJitter,
+		deletionsInProgress:       make(map[string]struct{}),
 	}
 
 	if az.config.VmssVmsCacheTTL != 0 {
@@ -183,6 +187,9 @@ func (scaleSet *ScaleSet) waitForDeleteInstances(future *azure.Future, requiredI
 
 	klog.V(3).Infof("Calling virtualMachineScaleSetsClient.WaitForDeleteInstancesResult(%v) for %s", requiredIds.InstanceIds, scaleSet.Name)
 	httpResponse, err := scaleSet.manager.azClient.virtualMachineScaleSetsClient.WaitForDeleteInstancesResult(ctx, future, scaleSet.manager.config.ResourceGroup)
+	for _, instanceID := range *requiredIds.InstanceIds {
+		scaleSet.trackDeletionEnd(instanceID)
+	}
 	isSuccess, err := isSuccessHTTPResponse(httpResponse, err)
 	if isSuccess {
 		klog.V(3).Infof("virtualMachineScaleSetsClient.WaitForDeleteInstancesResult(%v) for %s success", requiredIds.InstanceIds, scaleSet.Name)
@@ -395,6 +402,7 @@ func (scaleSet *ScaleSet) DeleteInstances(instances []*azureRef, hasUnregistered
 			klog.V(3).Infof("Skipping deleting instance %s as its current state is deleting", instance.Name)
 			continue
 		}
+
 		instancesToDelete = append(instancesToDelete, instance)
 	}
 
@@ -414,18 +422,24 @@ func (scaleSet *ScaleSet) DeleteInstances(instances []*azureRef, hasUnregistered
 		instanceIDs = append(instanceIDs, instanceID)
 	}
 
+	var instanceIDsToDelete []string
+	for _, instanceID := range instanceIDs {
+		if scaleSet.isBeingDeleted(instanceID) {
+			continue
+		}
+		instanceIDsToDelete = append(instanceIDsToDelete, instanceID)
+	}
+
 	requiredIds := &compute.VirtualMachineScaleSetVMInstanceRequiredIDs{
-		InstanceIds: &instanceIDs,
+		InstanceIds: &instanceIDsToDelete,
 	}
 
 	ctx, cancel := getContextWithTimeout(vmssContextTimeout)
 	defer cancel()
 	resourceGroup := scaleSet.manager.config.ResourceGroup
 
-	scaleSet.instanceMutex.Lock()
 	klog.V(3).Infof("Calling virtualMachineScaleSetsClient.DeleteInstancesAsync(%v)", requiredIds.InstanceIds)
 	future, rerr := scaleSet.manager.azClient.virtualMachineScaleSetsClient.DeleteInstancesAsync(ctx, resourceGroup, commonAsg.Id(), *requiredIds, false)
-	scaleSet.instanceMutex.Unlock()
 	if rerr != nil {
 		klog.Errorf("virtualMachineScaleSetsClient.DeleteInstancesAsync for instances %v failed: %v", requiredIds.InstanceIds, rerr)
 		return rerr.Error()
@@ -440,6 +454,9 @@ func (scaleSet *ScaleSet) DeleteInstances(instances []*azureRef, hasUnregistered
 		scaleSet.lastSizeRefresh = time.Now()
 		scaleSet.sizeMutex.Unlock()
 	}
+	for _, instanceID := range *requiredIds.InstanceIds {
+		scaleSet.trackDeletionStart(instanceID)
+	}
 
 	// Proactively set the status of the instances to be deleted in cache
 	for _, instance := range instancesToDelete {
@@ -732,3 +749,34 @@ func (scaleSet *ScaleSet) getOrchestrationMode() (compute.OrchestrationMode, err
 	}
 	return vmss.OrchestrationMode, nil
 }
+
+func (scaleSet *ScaleSet) trackDeletionStart(instanceID string) {
+	scaleSet.deletionMutex.Lock()
+	defer scaleSet.deletionMutex.Unlock()
+	if scaleSet.deletionsInProgress == nil {
+		scaleSet.deletionsInProgress = make(map[string]struct{})
+	}
+
+	scaleSet.deletionsInProgress[instanceID] = struct{}{}
+}
+
+func (scaleSet *ScaleSet) trackDeletionEnd(instanceID string) {
+	scaleSet.deletionMutex.Lock()
+	defer scaleSet.deletionMutex.Unlock()
+	if scaleSet.deletionsInProgress == nil {
+		scaleSet.deletionsInProgress = make(map[string]struct{})
+	}
+
+	delete(scaleSet.deletionsInProgress, instanceID)
+}
+
+func (scaleSet *ScaleSet) isBeingDeleted(instanceID string) bool {
+	scaleSet.deletionMutex.Lock()
+	defer scaleSet.deletionMutex.Unlock()
+	if scaleSet.deletionsInProgress == nil {
+		scaleSet.deletionsInProgress = make(map[string]struct{})
+	}
+
+	_, ok := scaleSet.deletionsInProgress[instanceID]
+	return ok
+}