From 7ed9d21f42862d1ad627333b53a5e86bd31d36a3 Mon Sep 17 00:00:00 2001
From: Jerome Froelich <jeromefroelich@hotmail.com>
Date: Tue, 29 Sep 2020 16:02:23 -0400
Subject: [PATCH] [controller] Always remove update annotation after processing
 StatefulSet

---
 pkg/controller/controller.go      | 14 +++++++++++---
 pkg/controller/controller_test.go | 16 ++++++++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/pkg/controller/controller.go b/pkg/controller/controller.go
index dc503926..3cc177a3 100644
--- a/pkg/controller/controller.go
+++ b/pkg/controller/controller.go
@@ -987,10 +987,20 @@ func updatedStatefulSet(
 		update = true
 	}
 
+	// If we don't need to perform an update to the StatefulSet's spec, but the StatefulSet
+	// has the update annotation, we'll still update the StatefulSet to remove the update
+	// annotation. This ensures that users can always set the update annotation and then
+	// wait for it to be removed to know if the operator has processed a StatefulSet.
 	if !update {
-		return nil, false, nil
+		delete(actual.Annotations, annotations.Update)
+		return actual, true, nil
 	}
 
+	copyAnnotations(expected, actual)
+	return expected, true, nil
+}
+
+func copyAnnotations(expected, actual *appsv1.StatefulSet) {
 	// It's okay for users to add annotations to a StatefulSet after it has been created so
 	// we'll want to copy over any that exist on the actual StatefulSet but not the expected
 	// one. The only exception is we don't want to copy over the update annotation so users
@@ -1006,6 +1016,4 @@ func updatedStatefulSet(
 			expected.Annotations[k] = v
 		}
 	}
-
-	return expected, true, nil
 }
diff --git a/pkg/controller/controller_test.go b/pkg/controller/controller_test.go
index 2b501ebf..3b223719 100644
--- a/pkg/controller/controller_test.go
+++ b/pkg/controller/controller_test.go
@@ -702,6 +702,22 @@ func TestHandleUpdateClusterUpdatesStatefulSets(t *testing.T) {
 			newImage:              "m3db:v2.0.0",
 			expUpdateStatefulSets: []string{"cluster1-rep2"},
 		},
+		{
+			name: "removes update annotation even if stateful set doesn't change",
+			cluster: newMeta("cluster1", map[string]string{
+				"foo":                      "bar",
+				"operator.m3db.io/app":     "m3db",
+				"operator.m3db.io/cluster": "cluster1",
+			}, nil),
+			sets: []*metav1.ObjectMeta{
+				newMeta("cluster1-rep0", nil, nil),
+				newMeta("cluster1-rep1", nil, map[string]string{
+					annotations.Update: "enabled",
+				}),
+				newMeta("cluster1-rep2", nil, nil),
+			},
+			expUpdateStatefulSets: []string{"cluster1-rep1"},
+		},
 		{
 			name: "doesn't call update for replica changes",
 			cluster: newMeta("cluster1", map[string]string{