From eba0fa2f954e93b129440a129232cb31b07b658a Mon Sep 17 00:00:00 2001
From: Beata Skiba <bskiba@google.com>
Date: Mon, 11 Sep 2017 17:23:59 +0200
Subject: [PATCH] Remove nodes that are not in the cluster from
 unremovableNodes

---
 cluster-autoscaler/core/scale_down.go      | 24 ++++++++++++++++++++++
 cluster-autoscaler/core/scale_down_test.go |  4 ++++
 2 files changed, 28 insertions(+)

diff --git a/cluster-autoscaler/core/scale_down.go b/cluster-autoscaler/core/scale_down.go
index ec1b7da429..3c775d167d 100644
--- a/cluster-autoscaler/core/scale_down.go
+++ b/cluster-autoscaler/core/scale_down.go
@@ -154,6 +154,7 @@ func (sd *ScaleDown) UpdateUnneededNodes(
 	nodeNameToNodeInfo := schedulercache.CreateNodeNameToInfoMap(pods, nodes)
 	utilizationMap := make(map[string]float64)
 
+	sd.updateUnremovableNodes(nodes)
 	// Filter out nodes that were recently checked
 	filteredNodesToCheck := make([]*apiv1.Node, 0)
 	for _, node := range nodesToCheck {
@@ -298,6 +299,29 @@ func (sd *ScaleDown) UpdateUnneededNodes(
 	return nil
 }
 
+// updateUnremovableNodes updates unremovableNodes map according to current
+// state of the cluster. Removes from the map nodes that are no longer in the
+// nodes list.
+func (sd *ScaleDown) updateUnremovableNodes(nodes []*apiv1.Node) {
+	if len(sd.unremovableNodes) <= 0 {
+		return
+	}
+	// A set of nodes to delete from unremovableNodes map.
+	nodesToDelete := make(map[string]struct{}, len(sd.unremovableNodes))
+	for name := range sd.unremovableNodes {
+		nodesToDelete[name] = struct{}{}
+	}
+	// Nodes that are in the cluster should not be deleted.
+	for _, node := range nodes {
+		if _, ok := nodesToDelete[node.Name]; ok {
+			delete(nodesToDelete, node.Name)
+		}
+	}
+	for nodeName := range nodesToDelete {
+		delete(sd.unremovableNodes, nodeName)
+	}
+}
+
 // markSimulationError indicates a simulation error by clearing  relevant scale
 // down state and returning an apropriate error.
 func (sd *ScaleDown) markSimulationError(simulatorErr errors.AutoscalerError,
diff --git a/cluster-autoscaler/core/scale_down_test.go b/cluster-autoscaler/core/scale_down_test.go
index ca35331181..02efb390fa 100644
--- a/cluster-autoscaler/core/scale_down_test.go
+++ b/cluster-autoscaler/core/scale_down_test.go
@@ -126,10 +126,14 @@ func TestFindUnneededNodes(t *testing.T) {
 	// Node n1 is unneeded, but should be skipped because it has just recently been found to be unremovable
 	sd.UpdateUnneededNodes([]*apiv1.Node{n1}, []*apiv1.Node{n1}, []*apiv1.Pod{}, time.Now(), nil)
 	assert.Equal(t, 0, len(sd.unneededNodes))
+	// Verify that no other nodes are in unremovable map.
+	assert.Equal(t, 1, len(sd.unremovableNodes))
 
 	// But it should be checked after timeout
 	sd.UpdateUnneededNodes([]*apiv1.Node{n1}, []*apiv1.Node{n1}, []*apiv1.Pod{}, time.Now().Add(UnremovableNodeRecheckTimeout+time.Second), nil)
 	assert.Equal(t, 1, len(sd.unneededNodes))
+	// Verify that nodes that are no longer unremovable are removed.
+	assert.Equal(t, 0, len(sd.unremovableNodes))
 }
 
 func TestFindUnneededMaxCandidates(t *testing.T) {