From 1a234b5aa93dcf807297d64ba8737d38b5d68707 Mon Sep 17 00:00:00 2001
From: Danny Thomson <dthomson25@gmail.com>
Date: Mon, 7 Oct 2019 15:51:30 -0700
Subject: [PATCH] Remove scale down annot after scaling down

---
 rollout/bluegreen_test.go             |  1 +
 rollout/sync.go                       |  4 +++
 utils/annotations/annotations.go      |  4 +++
 utils/annotations/annotations_test.go | 48 ++++++++++++++++-----------
 4 files changed, 37 insertions(+), 20 deletions(-)

diff --git a/rollout/bluegreen_test.go b/rollout/bluegreen_test.go
index 70050b25f8..97666b3eee 100644
--- a/rollout/bluegreen_test.go
+++ b/rollout/bluegreen_test.go
@@ -911,6 +911,7 @@ func TestBlueGreenReadyToScaleDownOldReplica(t *testing.T) {
 	f.run(getKey(r2, t))
 	updatedRS := f.getUpdatedReplicaSet(updatedRSIndex)
 	assert.Equal(t, int32(0), *updatedRS.Spec.Replicas)
+	assert.Equal(t, "", updatedRS.Annotations[v1alpha1.DefaultReplicaSetScaleDownDeadlineAnnotationKey])
 
 	patch := f.getPatchedRollout(patchIndex)
 	expectedPatch := calculatePatch(r2, OnlyObservedGenerationPatch)
diff --git a/rollout/sync.go b/rollout/sync.go
index 43d3ee69c2..39f4bceaa8 100644
--- a/rollout/sync.go
+++ b/rollout/sync.go
@@ -272,6 +272,7 @@ func (c *RolloutController) scaleReplicaSetAndRecordEvent(rs *appsv1.ReplicaSet,
 func (c *RolloutController) scaleReplicaSet(rs *appsv1.ReplicaSet, newScale int32, rollout *v1alpha1.Rollout, scalingOperation string) (bool, *appsv1.ReplicaSet, error) {
 
 	sizeNeedsUpdate := *(rs.Spec.Replicas) != newScale
+	fullScaleDown := newScale == int32(0)
 	rolloutReplicas := defaults.GetRolloutReplicasOrDefault(rollout)
 	annotationsNeedUpdate := annotations.ReplicasAnnotationsNeedUpdate(rs, rolloutReplicas)
 
@@ -281,6 +282,9 @@ func (c *RolloutController) scaleReplicaSet(rs *appsv1.ReplicaSet, newScale int3
 		rsCopy := rs.DeepCopy()
 		*(rsCopy.Spec.Replicas) = newScale
 		annotations.SetReplicasAnnotations(rsCopy, rolloutReplicas)
+		if fullScaleDown {
+			delete(rsCopy.Annotations, v1alpha1.DefaultReplicaSetScaleDownDeadlineAnnotationKey)
+		}
 		rs, err = c.kubeclientset.AppsV1().ReplicaSets(rsCopy.Namespace).Update(rsCopy)
 		if err == nil && sizeNeedsUpdate {
 			scaled = true
diff --git a/utils/annotations/annotations.go b/utils/annotations/annotations.go
index 0820543ac4..e5d3824553 100644
--- a/utils/annotations/annotations.go
+++ b/utils/annotations/annotations.go
@@ -82,6 +82,10 @@ func ReplicasAnnotationsNeedUpdate(rs *appsv1.ReplicaSet, desiredReplicas int32)
 	if hasString := rs.Annotations[DesiredReplicasAnnotation]; hasString != desiredString {
 		return true
 	}
+	hasScaleDownDelay := rs.Annotations[v1alpha1.DefaultReplicaSetScaleDownDeadlineAnnotationKey]
+	if desiredReplicas == int32(0) && hasScaleDownDelay != "" {
+		return true
+	}
 	return false
 }
 
diff --git a/utils/annotations/annotations_test.go b/utils/annotations/annotations_test.go
index bf106e5828..b3c69cda34 100644
--- a/utils/annotations/annotations_test.go
+++ b/utils/annotations/annotations_test.go
@@ -280,34 +280,41 @@ func TestAnnotationUtils(t *testing.T) {
 func TestReplicasAnnotationsNeedUpdate(t *testing.T) {
 
 	desiredReplicas := fmt.Sprintf("%d", int32(10))
+	zeroDesiredReplicas := fmt.Sprintf("%d", int32(0))
 	tests := []struct {
-		name       string
-		replicaSet *appsv1.ReplicaSet
-		expected   bool
+		name            string
+		replicaSet      *appsv1.ReplicaSet
+		desiredReplicas int32
+		expected        bool
 	}{
+		// {
+		// 	name: "test Annotations nil",
+		// 	replicaSet: &appsv1.ReplicaSet{
+		// 		ObjectMeta: metav1.ObjectMeta{Name: "hello", Namespace: "test"},
+		// 		Spec: appsv1.ReplicaSetSpec{
+		// 			Selector: &metav1.LabelSelector{MatchLabels: map[string]string{"foo": "bar"}},
+		// 		},
+		// 	},
+		// 	desiredReplicas: 10,
+		// 	expected:        true,
+		// },
 		{
-			name: "test Annotations nil",
-			replicaSet: &appsv1.ReplicaSet{
-				ObjectMeta: metav1.ObjectMeta{Name: "hello", Namespace: "test"},
-				Spec: appsv1.ReplicaSetSpec{
-					Selector: &metav1.LabelSelector{MatchLabels: map[string]string{"foo": "bar"}},
-				},
-			},
-			expected: true,
-		},
-		{
-			name: "test desiredReplicas update",
+			name: "test remove scale-down-delay",
 			replicaSet: &appsv1.ReplicaSet{
 				ObjectMeta: metav1.ObjectMeta{
-					Name:        "hello",
-					Namespace:   "test",
-					Annotations: map[string]string{DesiredReplicasAnnotation: "8"},
+					Name:      "hello",
+					Namespace: "test",
+					Annotations: map[string]string{
+						DesiredReplicasAnnotation:                                zeroDesiredReplicas,
+						v1alpha1.DefaultReplicaSetScaleDownDeadlineAnnotationKey: "set-to-something",
+					},
 				},
 				Spec: appsv1.ReplicaSetSpec{
 					Selector: &metav1.LabelSelector{MatchLabels: map[string]string{"foo": "bar"}},
 				},
 			},
-			expected: true,
+			desiredReplicas: 0,
+			expected:        true,
 		},
 		{
 			name: "test needn't update",
@@ -321,13 +328,14 @@ func TestReplicasAnnotationsNeedUpdate(t *testing.T) {
 					Selector: &metav1.LabelSelector{MatchLabels: map[string]string{"foo": "bar"}},
 				},
 			},
-			expected: false,
+			desiredReplicas: 10,
+			expected:        false,
 		},
 	}
 
 	for i, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			result := ReplicasAnnotationsNeedUpdate(test.replicaSet, 10)
+			result := ReplicasAnnotationsNeedUpdate(test.replicaSet, test.desiredReplicas)
 			if result != test.expected {
 				t.Errorf("case[%d]:%s Expected %v, Got: %v", i, test.name, test.expected, result)
 			}