Repository: m3db/m3db-operator
Pull request #236 (2020-09-25 19:01:29)

[controller] Add logic to update StatefulSets
This commit to the controller adds support for updating `StatefulSet`'s when their image or config map are out of date with the `ClusterSpec`. After checking that the cluster is healthy, the controller will next check if any `StatefulSet` needs to be updated, and will then perform any needed updates one by one. To ensure that the controller will only perform updates when it is safe to do so, the controller will check for an annotation that must be set by a user to indicate the controller can perform an update.

codecov[bot]: # [Codecov](https://codecov.io/gh/m3db/m3db-operator/pull/236?src=pr&el=h1) Report
> Merging [#236](https://codecov.io/gh/m3db/m3db-operator/pull/236?src=pr&el=desc) into [master](https://codecov.io/gh/m3db/m3db-operator/commit/acaa0f1f682b447cdd214d65c975782111610e2f?el=desc) will **increase** coverage by `0.60%`.
> The diff coverage is `n/a`.


[![Impacted file tree graph](https://codecov.io/gh/m3db/m3db-operator/pull/236/graphs/tree.svg?width=650&height=150&src=pr&token=jtHDyWhsto)](https://codecov.io/gh/m3db/m3db-operator/pull/236?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #236      +/-   ##
==========================================
+ Coverage   75.59%   76.19%   +0.60%     
==========================================
  Files          30       30              
  Lines        2118     2168      +50     
==========================================
+ Hits         1601     1652      +51     
+ Misses        392      381      -11     
- Partials      125      135      +10     
```


------

[Continue to review full report at Codecov](https://codecov.io/gh/m3db/m3db-operator/pull/236?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/m3db/m3db-operator/pull/236?src=pr&el=footer). Last update [acaa0f1...f1db6ed](https://codecov.io/gh/m3db/m3db-operator/pull/236?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

jeromefroe: > Can do this in another PR if you want but would love to have something in `docs` on using this.

Definitely! I tested the changes on a running cluster yesterday using a local build, but want to deploy the changes and do some more thorough testing today. So I'll update the docs after that once I've gotten a better handle on the UX.


Patch:
diff --git a/go.mod b/go.mod
index 059ec47..3c1c3ad 100644
--- a/go.mod
+++ b/go.mod
@@ -77,7 +77,7 @@ require (
 	golang.org/x/crypto v0.0.0-20200220183623-bac4c82f6975 // indirect
 	golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3
 	golang.org/x/net v0.0.0-20200324143707-d3edc9973b7e // indirect
-	golang.org/x/tools v0.0.0-20200102140908-9497f49d5709 // indirect
+	golang.org/x/tools v0.0.0-20200522201501-cb1345f3a375 // indirect
 	google.golang.org/appengine v1.6.2 // indirect
 	gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect
 	gopkg.in/go-ini/ini.v1 v1.51.1 // indirect
diff --git a/go.sum b/go.sum
index b4ef31c..94450c3 100644
--- a/go.sum
+++ b/go.sum
@@ -623,6 +623,7 @@ github.com/xeipuuv/gojsonschema v1.2.0 h1:LhYJRs+L4fBtjZUfuSZIKGeVu0QRy8e5Xi7D17
 github.com/xeipuuv/gojsonschema v1.2.0/go.mod h1:anYRn/JVcOK2ZgGU+IjEV4nwlhoK5sQluxsYJ78Id3Y=
 github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2/go.mod h1:UETIi67q53MR2AWcXfiuqkDkRtnGDLqkBTpCHuJHxtU=
 github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q=
+github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 go.etcd.io/bbolt v1.3.3/go.mod h1:IbVyRI1SCnLcuJnV2u8VeU0CEYM7e686BmAb1XKL+uU=
 go.etcd.io/etcd v0.0.0-20191023171146-3cf2f69b5738/go.mod h1:dnLIgRNXwCJa5e+c6mIZCrds/GIG4ncV9HhK5PX7jPg=
 go.mongodb.org/mongo-driver v1.0.3/go.mod h1:u7ryQJ+DOzQmeO7zB6MHyr8jkEQvC8vH7qLUO4lqsUM=
@@ -667,7 +668,8 @@ golang.org/x/lint v0.0.0-20190301231843-5614ed5bae6f/go.mod h1:UVdnD1Gm6xHRNCYTk
 golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3 h1:XQyxROzUlZH+WIQwySDgnISgOivlhjIEwaQaJEJrrN0=
 golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
 golang.org/x/mobile v0.0.0-20190312151609-d3739f865fa6/go.mod h1:z+o9i4GpDbdi3rU15maQ/Ox0txvL9dWGYEHz965HBQE=
-golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg=
+golang.org/x/mod v0.2.0 h1:KU7oHjnv3XNWfa5COkzUifxZmxp1TyI7ImMXqFxLwvQ=
+golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/net v0.0.0-20170114055629-f2499483f923/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20170915142106-8351a756f30f/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
@@ -692,6 +694,7 @@ golang.org/x/net v0.0.0-20190827160401-ba9fcec4b297 h1:k7pJ2yAPLPgbskkFdhRCsA77k
 golang.org/x/net v0.0.0-20190827160401-ba9fcec4b297/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20191004110552-13f9640d40b9 h1:rjwSpXsdiK0dV8/Naq3kAw9ymfAeJIyd0upUIElB+lI=
 golang.org/x/net v0.0.0-20191004110552-13f9640d40b9/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
+golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20200324143707-d3edc9973b7e h1:3G+cUijn7XD+S4eJFddp53Pv7+slrESplyjG25HgL+k=
 golang.org/x/net v0.0.0-20200324143707-d3edc9973b7e/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
@@ -705,6 +708,8 @@ golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20190227155943-e225da77a7e6/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58 h1:8gQV6CLnAEikrhgkHFbMAEhagSSnXWGV915qUMm9mrU=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e h1:vcxGaoTs7kV8m5Np9uUNQin4BrLOthgV7252N8V+FwY=
+golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20170830134202-bb24a47a89ea/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20171026204733-164713f0dfce/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
@@ -766,8 +771,9 @@ golang.org/x/tools v0.0.0-20190719005602-e377ae9d6386/go.mod h1:jcCCGcm9btYwXyDq
 golang.org/x/tools v0.0.0-20190729092621-ff9f1409240a/go.mod h1:jcCCGcm9btYwXyDqrUWc6MKQKKGJCWEQ3AfLSRIbEuI=
 golang.org/x/tools v0.0.0-20190909030654-5b82db07426d/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.0.0-20190920225731-5eefd052ad72/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
-golang.org/x/tools v0.0.0-20200102140908-9497f49d5709 h1:AfG1EmoRkFK24HWWLxSrRKNg2G+oA3JVOG8GJsHWypQ=
-golang.org/x/tools v0.0.0-20200102140908-9497f49d5709/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28=
+golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
+golang.org/x/tools v0.0.0-20200522201501-cb1345f3a375 h1:SjQ2+AKWgZLc1xej6WSzL+Dfs5Uyd5xcZH1mGC411IA=
+golang.org/x/tools v0.0.0-20200522201501-cb1345f3a375/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898 h1:/atklqdjdhuosWIl6AIbOeHJjicWYPqR9bpxqxYG2pA=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
diff --git a/pkg/controller/controller.go b/pkg/controller/controller.go
index a0a1dab..2d162fe 100644
--- a/pkg/controller/controller.go
+++ b/pkg/controller/controller.go
@@ -24,6 +24,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"reflect"
 	"sort"
 	"sync"
 
@@ -32,6 +33,7 @@ import (
 	clientset "github.com/m3db/m3db-operator/pkg/client/clientset/versioned"
 	samplescheme "github.com/m3db/m3db-operator/pkg/client/clientset/versioned/scheme"
 	clusterlisters "github.com/m3db/m3db-operator/pkg/client/listers/m3dboperator/v1alpha1"
+	"github.com/m3db/m3db-operator/pkg/k8sops/annotations"
 	"github.com/m3db/m3db-operator/pkg/k8sops/labels"
 	"github.com/m3db/m3db-operator/pkg/k8sops/m3db"
 	"github.com/m3db/m3db-operator/pkg/k8sops/podidentity"
@@ -42,6 +44,7 @@ import (
 
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/equality"
 	kerrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	klabels "k8s.io/apimachinery/pkg/labels"
@@ -477,6 +480,34 @@ func (c *M3DBController) handleClusterUpdate(cluster *myspec.M3DBCluster) error
 		}
 	}
 
+	// Now that we know the cluster is healthy, iterate over each isolation group and check
+	// if it should be updated.
+	for i := 0; i < len(isoGroups); i++ {
+		name := m3db.StatefulSetName(cluster.Name, i)
+
+		// This StatefulSet is guaranteed to exist since if it didn't originally it would be
+		// created above when we first iterate over the isolation groups.
+		actual := childrenSetsByName[name]
+		expected, update, err := updatedStatefulSet(actual, cluster, isoGroups[i])
+		if err != nil {
+			c.logger.Error(err.Error())
+			return err
+		}
+
+		if !update {
+			continue
+		}
+
+		_, err = c.kubeClient.AppsV1().StatefulSets(cluster.Namespace).Update(expected)
+		if err != nil {
+			c.logger.Error(err.Error())
+			return err
+		}
+
+		c.logger.Info("updated statefulset", zap.String("name", name))
+		return nil
+	}
+
 	if err := c.reconcileNamespaces(cluster); err != nil {
 		c.recorder.WarningEvent(cluster, eventer.ReasonFailedCreate, "failed to create namespace: %s", err)
 		c.logger.Error("error reconciling namespaces", zap.Error(err))
@@ -919,3 +950,58 @@ func validateIsolationGroups(cluster *myspec.M3DBCluster) error {
 	}
 	return nil
 }
+
+// updatedStatefulSet checks if we should update the actual StatefulSet to match it's
+// expected state. If so, it returns true and the updated StatefulSet and if not, it
+// return false.
+func updatedStatefulSet(
+	actual *appsv1.StatefulSet, cluster *myspec.M3DBCluster, isoGroup myspec.IsolationGroup,
+) (*appsv1.StatefulSet, bool, error) {
+	// The operator will only perform an update if the current StatefulSet has been
+	// annotated to indicate that it is okay to update it.
+	if val, ok := actual.Annotations[annotations.Update]; !ok || val != "enabled" {
+		return nil, false, nil
+	}
+
+	expected, err := m3db.GenerateStatefulSet(cluster, isoGroup.Name, isoGroup.NumInstances)
+	if err != nil {
+		return nil, false, err
+	}
+
+	// The Kubernetes API server sets various default values for fields so instead of comparing
+	// if the desired spec is equal to the actual spec, which may have fail because the desired
+	// spec hasn't yet been transformed by the API server, we use DeepDerivative to only compare
+	// those fields in the desired spec which are actually set.
+	var update bool
+	if !equality.Semantic.DeepDerivative(expected.Spec, actual.Spec) {
+		update = true
+	}
+
+	// We also want to check if the labels in the cluster spec have changed and update the
+	// StatefulSet if so.
+	if !reflect.DeepEqual(expected.Labels, actual.Labels) {
+		update = true
+	}
+
+	if !update {
+		return nil, false, nil
+	}
+
+	// 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
+	// will always have to explicitly trigger updates. Note that this means removing an
+	// annotation will require users to delete the annotation from the StatefulSet as well as
+	// the cluster spec.
+	for k, v := range actual.Annotations {
+		if k == annotations.Update {
+			continue
+		}
+
+		if _, ok := expected.Annotations[k]; !ok {
+			expected.Annotations[k] = v
+		}
+	}
+
+	return expected, true, nil
+}
diff --git a/pkg/k8sops/annotations/annotations.go b/pkg/k8sops/annotations/annotations.go
index b6dc6f4..292ad0a 100644
--- a/pkg/k8sops/annotations/annotations.go
+++ b/pkg/k8sops/annotations/annotations.go
@@ -34,6 +34,9 @@ const (
 	AppM3DB = labels.AppM3DB
 	// Cluster is the label identifying what m3db cluster an object is a part of.
 	Cluster = labels.Cluster
+	// Update is the annotation used by the operator to determine if a StatefulSet is
+	// allowed to be updated.
+	Update = "operator.m3db.io/update"
 )
 
 // BaseAnnotations returns the base annotations we apply to all objects


Test patch:
diff --git a/pkg/controller/controller_test.go b/pkg/controller/controller_test.go
index d2deb09..2dd3977 100644
--- a/pkg/controller/controller_test.go
+++ b/pkg/controller/controller_test.go
@@ -23,6 +23,7 @@ package controller
 import (
 	"fmt"
 	"sort"
+	"strconv"
 	"sync"
 	"testing"
 	"time"
@@ -30,11 +31,14 @@ import (
 	myspec "github.com/m3db/m3db-operator/pkg/apis/m3dboperator/v1alpha1"
 	clientsetfake "github.com/m3db/m3db-operator/pkg/client/clientset/versioned/fake"
 	m3dbinformers "github.com/m3db/m3db-operator/pkg/client/informers/externalversions"
+	"github.com/m3db/m3db-operator/pkg/k8sops/annotations"
 	"github.com/m3db/m3db-operator/pkg/k8sops/labels"
 	"github.com/m3db/m3db-operator/pkg/k8sops/m3db"
 	"github.com/m3db/m3db-operator/pkg/k8sops/podidentity"
 
 	"github.com/m3db/m3/src/cluster/placement"
+	namespacepb "github.com/m3db/m3/src/dbnode/generated/proto/namespace"
+	"github.com/m3db/m3/src/query/generated/proto/admin"
 
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
@@ -55,12 +59,152 @@ import (
 	"go.uber.org/zap"
 )
 
-func newMeta(name string, labels map[string]string) *metav1.ObjectMeta {
+const (
+	defaultTestImage     = "m3db:v1.0.0"
+	defaultConfigMapName = "configMapName"
+)
+
+func newMeta(name string, labels, annotations map[string]string) *metav1.ObjectMeta {
 	return &metav1.ObjectMeta{
-		Name:      name,
-		Labels:    labels,
-		Namespace: "namespace",
+		Name:        name,
+		Labels:      labels,
+		Annotations: annotations,
+		Namespace:   "namespace",
+	}
+}
+
+func setupTestCluster(
+	t *testing.T,
+	clusterMeta metav1.ObjectMeta,
+	sets []*metav1.ObjectMeta,
+	replicationFactor int,
+) (*myspec.M3DBCluster, *testDeps) {
+	cfgMapName := defaultConfigMapName
+	cluster := &myspec.M3DBCluster{
+		ObjectMeta: clusterMeta,
+		Spec: myspec.ClusterSpec{
+			Image:             defaultTestImage,
+			ReplicationFactor: int32(replicationFactor),
+			ConfigMapName:     &cfgMapName,
+		},
 	}
+	cluster.ObjectMeta.UID = "abcd"
+	for i := 0; i < replicationFactor; i++ {
+		group := myspec.IsolationGroup{
+			Name:         fmt.Sprintf("group%d", i),
+			NumInstances: 1,
+		}
+		cluster.Spec.IsolationGroups = append(cluster.Spec.IsolationGroups, group)
+	}
+	cluster.ObjectMeta.Finalizers = []string{labels.EtcdDeletionFinalizer}
+
+	objects := make([]runtime.Object, len(sets))
+	statefulSets := make([]*appsv1.StatefulSet, len(sets))
+	for i, s := range sets {
+		set := &appsv1.StatefulSet{
+			ObjectMeta: *s,
+		}
+		// Apply base labels
+		if set.ObjectMeta.Labels == nil {
+			set.ObjectMeta.Labels = make(map[string]string)
+		}
+		for k, v := range labels.BaseLabels(cluster) {
+			set.ObjectMeta.Labels[k] = v
+		}
+		set.Spec.Template.Spec.Containers = append(set.Spec.Template.Spec.Containers, corev1.Container{
+			Image: defaultTestImage,
+		})
+		set.Spec.Template.Spec.Volumes = append(set.Spec.Template.Spec.Volumes, corev1.Volume{
+			VolumeSource: corev1.VolumeSource{
+				ConfigMap: &corev1.ConfigMapVolumeSource{
+					LocalObjectReference: corev1.LocalObjectReference{Name: defaultConfigMapName},
+				},
+			},
+		})
+		statefulSets[i] = set
+		objects[i] = set
+		set.OwnerReferences = []metav1.OwnerReference{
+			*metav1.NewControllerRef(cluster, schema.GroupVersionKind{
+				Group:   myspec.SchemeGroupVersion.Group,
+				Version: myspec.SchemeGroupVersion.Version,
+				Kind:    "m3dbcluster",
+			}),
+		}
+	}
+
+	deps := newTestDeps(t, &testOpts{
+		crdObjects: []runtime.Object{
+			&myspec.M3DBCluster{
+				ObjectMeta: clusterMeta,
+			},
+		},
+		kubeObjects: objects,
+	})
+
+	return cluster, deps
+}
+
+func waitForStatefulSets(
+	t *testing.T,
+	controller *M3DBController,
+	cluster *myspec.M3DBCluster,
+	verb string,
+	setReadyReplicas bool,
+	expectedStatefulSets []string,
+) bool {
+	var (
+		mu         sync.Mutex
+		actualSets = make(map[string]bool)
+	)
+
+	controller.kubeClient.(*kubefake.Clientset).PrependReactor(verb, "statefulsets", func(action ktesting.Action) (bool, runtime.Object, error) {
+		var sts *appsv1.StatefulSet
+		switch verb {
+		case "create":
+			sts = action.(kubetesting.CreateActionImpl).GetObject().(*appsv1.StatefulSet)
+		case "update":
+			sts = action.(kubetesting.UpdateActionImpl).GetObject().(*appsv1.StatefulSet)
+		default:
+			t.Errorf("verb %s is not supported", verb)
+		}
+
+		if setReadyReplicas && sts.Spec.Replicas != nil {
+			sts.Status.ReadyReplicas = *sts.Spec.Replicas
+		}
+
+		name := sts.Name
+		mu.Lock()
+		actualSets[name] = true
+		mu.Unlock()
+
+		// Return false to indicate that the next reactor in the chain should process this
+		// object as well.
+		return false, nil, nil
+	})
+
+	// Iterate through the expected stateful sets twice to make sure we see all stateful
+	// sets that we expect and also be able to catch any extra stateful sets that we don't.
+	var done bool
+	for i := 0; i < 2*len(expectedStatefulSets); i++ {
+		err := controller.handleClusterUpdate(cluster)
+		require.NoError(t, err)
+
+		mu.Lock()
+		var seen int
+		for _, found := range actualSets {
+			if found {
+				seen++
+			}
+		}
+		mu.Unlock()
+
+		if seen != len(expectedStatefulSets) {
+			time.Sleep(100 * time.Millisecond)
+			continue
+		}
+		done = true
+	}
+	return done
 }
 
 func TestNew(t *testing.T) {
@@ -95,31 +239,31 @@ func TestGetChildStatefulSets(t *testing.T) {
 		expChildren []string
 	}{
 		{
-			cluster: newMeta("cluster1", map[string]string{"foo": "bar"}),
+			cluster: newMeta("cluster1", map[string]string{"foo": "bar"}, nil),
 			sets: []*metav1.ObjectMeta{
-				newMeta("set1", nil),
+				newMeta("set1", nil, nil),
 			},
 			expChildren: []string{},
 		},
 		{
-			cluster: newMeta("cluster1", map[string]string{"foo": "bar"}),
+			cluster: newMeta("cluster1", map[string]string{"foo": "bar"}, nil),
 			sets: []*metav1.ObjectMeta{
 				newMeta("set1", map[string]string{
 					"foo":                      "bar",
 					"operator.m3db.io/app":     "m3db",
 					"operator.m3db.io/cluster": "cluster1",
-				}),
+				}, nil),
 			},
 			expChildren: []string{"set1"},
 		},
 		{
-			cluster: newMeta("cluster1", map[string]string{"foo": "bar"}),
+			cluster: newMeta("cluster1", map[string]string{"foo": "bar"}, nil),
 			sets: []*metav1.ObjectMeta{
 				newMeta("set1", map[string]string{
 					"foo":                      "bar",
 					"operator.m3db.io/app":     "m3db",
 					"operator.m3db.io/cluster": "cluster2",
-				}),
+				}, nil),
 			},
 			expChildren: []string{},
 		},
@@ -437,9 +581,9 @@ func TestHandleUpdateClusterCreatesStatefulSets(t *testing.T) {
 				"foo":                      "bar",
 				"operator.m3db.io/app":     "m3db",
 				"operator.m3db.io/cluster": "cluster1",
-			}),
+			}, nil),
 			sets: []*metav1.ObjectMeta{
-				newMeta("cluster1-rep0", nil),
+				newMeta("cluster1-rep0", nil, nil),
 			},
 			replicationFactor:     3,
 			expCreateStatefulSets: []string{"cluster1-rep1", "cluster1-rep2"},
@@ -450,9 +594,9 @@ func TestHandleUpdateClusterCreatesStatefulSets(t *testing.T) {
 				"foo":                      "bar",
 				"operator.m3db.io/app":     "m3db",
 				"operator.m3db.io/cluster": "cluster1",
-			}),
+			}, nil),
 			sets: []*metav1.ObjectMeta{
-				newMeta("cluster1-rep1", nil),
+				newMeta("cluster1-rep1", nil, nil),
 			},
 			replicationFactor:     3,
 			expCreateStatefulSets: []string{"cluster1-rep0", "cluster1-rep2"},
@@ -463,9 +607,9 @@ func TestHandleUpdateClusterCreatesStatefulSets(t *testing.T) {
 				"foo":                      "bar",
 				"operator.m3db.io/app":     "m3db",
 				"operator.m3db.io/cluster": "cluster1",
-			}),
+			}, nil),
 			sets: []*metav1.ObjectMeta{
-				newMeta("cluster1-rep2", nil),
+				newMeta("cluster1-rep2", nil, nil),
 			},
 			replicationFactor:     3,
 			expCreateStatefulSets: []string{"cluster1-rep0", "cluster1-rep1"},
@@ -474,91 +618,124 @@ func TestHandleUpdateClusterCreatesStatefulSets(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			cfgMapName := "configMapName"
-			cluster := &myspec.M3DBCluster{
-				ObjectMeta: *test.cluster,
-				Spec: myspec.ClusterSpec{
-					ReplicationFactor: int32(test.replicationFactor),
-					ConfigMapName:     &cfgMapName,
-				},
-			}
-			cluster.ObjectMeta.UID = "abcd"
-			for i := 0; i < test.replicationFactor; i++ {
-				group := myspec.IsolationGroup{
-					Name:         fmt.Sprintf("group%d", i),
-					NumInstances: 1,
-				}
-				cluster.Spec.IsolationGroups = append(cluster.Spec.IsolationGroups, group)
-			}
-			cluster.ObjectMeta.Finalizers = []string{labels.EtcdDeletionFinalizer}
-
-			objects := make([]runtime.Object, len(test.sets))
-			statefulSets := make([]*appsv1.StatefulSet, len(test.sets))
-			for i, s := range test.sets {
-				set := &appsv1.StatefulSet{
-					ObjectMeta: *s,
-				}
-				// Apply base labels
-				if set.ObjectMeta.Labels == nil {
-					set.ObjectMeta.Labels = make(map[string]string)
-				}
-				for k, v := range labels.BaseLabels(cluster) {
-					set.ObjectMeta.Labels[k] = v
-				}
-				statefulSets[i] = set
-				objects[i] = set
-				set.OwnerReferences = []metav1.OwnerReference{
-					*metav1.NewControllerRef(cluster, schema.GroupVersionKind{
-						Group:   myspec.SchemeGroupVersion.Group,
-						Version: myspec.SchemeGroupVersion.Version,
-						Kind:    "m3dbcluster",
-					}),
-				}
-			}
+			cluster, deps := setupTestCluster(t, *test.cluster, test.sets, test.replicationFactor)
+			defer deps.cleanup()
+			c := deps.newController(t)
 
-			deps := newTestDeps(t, &testOpts{
-				crdObjects: []runtime.Object{
-					&myspec.M3DBCluster{
-						ObjectMeta: *test.cluster,
-					},
-				},
-				kubeObjects: objects,
-			})
+			done := waitForStatefulSets(t, c, cluster, "create", false, test.expCreateStatefulSets)
+			assert.True(t, done, "expected all sets to be created")
+		})
+	}
+}
+
+func TestHandleUpdateClusterUpdatesStatefulSets(t *testing.T) {
+	tests := []struct {
+		name                  string
+		cluster               *metav1.ObjectMeta
+		sets                  []*metav1.ObjectMeta
+		newImage              string
+		newConfigMap          string
+		expUpdateStatefulSets []string
+	}{
+		{
+			name: "updates stateful sets when image is out of date",
+			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, map[string]string{
+					annotations.Update: "enabled",
+				}),
+				newMeta("cluster1-rep1", nil, map[string]string{
+					annotations.Update: "enabled",
+				}),
+				newMeta("cluster1-rep2", nil, map[string]string{
+					annotations.Update: "enabled",
+				}),
+			},
+			newImage:              "m3db:v2.0.0",
+			expUpdateStatefulSets: []string{"cluster1-rep0", "cluster1-rep1", "cluster1-rep2"},
+		},
+		{
+			name: "updates stateful sets when config map name is out of date",
+			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, map[string]string{
+					annotations.Update: "enabled",
+				}),
+				newMeta("cluster1-rep1", nil, map[string]string{
+					annotations.Update: "enabled",
+				}),
+				newMeta("cluster1-rep2", nil, map[string]string{
+					annotations.Update: "enabled",
+				}),
+			},
+			newConfigMap:          "configMapName-2",
+			expUpdateStatefulSets: []string{"cluster1-rep0", "cluster1-rep1", "cluster1-rep2"},
+		},
+		{
+			name: "only updates stateful sets with update annotation",
+			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, nil),
+				newMeta("cluster1-rep2", nil, map[string]string{
+					annotations.Update: "enabled",
+				}),
+			},
+			newImage:              "m3db:v2.0.0",
+			expUpdateStatefulSets: []string{"cluster1-rep2"},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			const replicas = 3
+			cluster, deps := setupTestCluster(t, *test.cluster, test.sets, replicas)
 			defer deps.cleanup()
 			c := deps.newController(t)
 
-			// Keep running cluster updates until no more stateful sets required
-			var expectedMu sync.Mutex
-			expectedSetsCreated := make(map[string]struct{})
-
-			c.kubeClient.(*kubefake.Clientset).PrependReactor("create", "statefulsets", func(action ktesting.Action) (bool, runtime.Object, error) {
-				cluster := action.(kubetesting.CreateActionImpl).GetObject().(*appsv1.StatefulSet)
-				name := cluster.Name
-				expectedMu.Lock()
-				expectedSetsCreated[name] = struct{}{}
-				expectedMu.Unlock()
-				// Return false to indicate that the next reactor in the chain should process this
-				// object as well.
-				return false, nil, nil
-			})
+			deps.namespaceClient.EXPECT().List().AnyTimes().Return(&admin.NamespaceGetResponse{
+				Registry: &namespacepb.Registry{},
+			}, nil)
+
+			var mockInstances []placement.Instance
+			for i := 0; i < replicas; i++ {
+				inst := placement.NewMockInstance(deps.mockController)
+
+				// Since the controller checks all instances are available in the placement after
+				// checking and performing any updates, and we're only concerned with the latter
+				// in this test, configure the mock instances to be unavailable to reduce the
+				// amount of mocks we need to set up.
+				inst.EXPECT().IsAvailable().AnyTimes().Return(false)
+				inst.EXPECT().ID().AnyTimes().Return(strconv.Itoa(i))
+				mockInstances = append(mockInstances, inst)
+			}
+
+			mockPlacement := placement.NewMockPlacement(deps.mockController)
+			mockPlacement.EXPECT().NumInstances().AnyTimes().Return(len(mockInstances))
+			mockPlacement.EXPECT().Instances().AnyTimes().Return(mockInstances)
+			deps.placementClient.EXPECT().Get().AnyTimes().Return(mockPlacement, nil)
 
-			// Iterate through the expected stateful sets twice to make sure we see all stateful
-			// sets that we expect and also be able to catch any extra stateful sets that we don't.
-			var done bool
-			for i := 0; i < 2*len(test.expCreateStatefulSets); i++ {
-				err := c.handleClusterUpdate(cluster)
-				require.NoError(t, err)
-
-				expectedMu.Lock()
-				created := len(expectedSetsCreated)
-				expectedMu.Unlock()
-				if created != len(test.expCreateStatefulSets) {
-					time.Sleep(100 * time.Millisecond)
-					continue
-				}
-				done = true
+			if test.newImage != "" {
+				cluster.Spec.Image = test.newImage
 			}
-			assert.True(t, done, "expected all sets to be created")
+			if test.newConfigMap != "" {
+				cluster.Spec.ConfigMapName = &test.newConfigMap
+			}
+
+			done := waitForStatefulSets(t, c, cluster, "update", true, test.expUpdateStatefulSets)
+			assert.True(t, done, "expected all sets to be updated")
 		})
 	}
 }