From d96104b271d273865bb3c201e1c78bd10fec1b34 Mon Sep 17 00:00:00 2001
From: Rob Skillington <rob.skillington@gmail.com>
Date: Fri, 31 May 2019 01:50:11 +0200
Subject: [PATCH] Always create missing statefulsets

---
 pkg/controller/controller.go      | 35 ++++++++++++++++---------------
 pkg/controller/controller_test.go |  4 ++++
 2 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/pkg/controller/controller.go b/pkg/controller/controller.go
index 6a775019..00d62370 100644
--- a/pkg/controller/controller.go
+++ b/pkg/controller/controller.go
@@ -360,7 +360,9 @@ func (c *Controller) handleClusterUpdate(cluster *myspec.M3DBCluster) error {
 		return err
 	}
 
+	childrenSetsByName := make(map[string]*appsv1.StatefulSet)
 	for _, sts := range childrenSets {
+		childrenSetsByName[sts.Name] = sts
 		// if any of the statefulsets aren't ready, wait until they are as we'll get
 		// another event (ready == bootstrapped)
 		if sts.Spec.Replicas != nil && *sts.Spec.Replicas != sts.Status.ReadyReplicas {
@@ -370,26 +372,25 @@ func (c *Controller) handleClusterUpdate(cluster *myspec.M3DBCluster) error {
 		}
 	}
 
-	// At this point all existing statefulsets are bootstrapped.
-
-	if len(childrenSets) != len(isoGroups) {
-		nextID := len(childrenSets)
-		// create a statefulset
+	// Create any missing statefulsets, at this point all existing stateful sets are bootstrapped.
+	for i := 0; i < len(isoGroups); i++ {
+		name := fmt.Sprintf("%s-%d", cluster.Name, i)
+		_, exists := childrenSetsByName[name]
+		if !exists {
+			sts, err := k8sops.GenerateStatefulSet(cluster, isoGroups[i].Name, isoGroups[i].NumInstances)
+			if err != nil {
+				return err
+			}
 
-		name := fmt.Sprintf("%s-%d", cluster.Name, nextID)
-		sts, err := k8sops.GenerateStatefulSet(cluster, isoGroups[nextID].Name, isoGroups[nextID].NumInstances)
-		if err != nil {
-			return err
-		}
+			_, err = c.kubeClient.AppsV1().StatefulSets(cluster.Namespace).Create(sts)
+			if err != nil {
+				c.logger.Error(err.Error())
+				return err
+			}
 
-		_, err = c.kubeClient.AppsV1().StatefulSets(cluster.Namespace).Create(sts)
-		if err != nil {
-			c.logger.Error(err.Error())
-			return err
+			c.logger.Info("created statefulset", zap.String("name", name))
+			return nil
 		}
-
-		c.logger.Info("created statefulset", zap.String("name", name))
-		return nil
 	}
 
 	if err := c.reconcileNamespaces(cluster); err != nil {
diff --git a/pkg/controller/controller_test.go b/pkg/controller/controller_test.go
index 8f5d5b8a..1c6e484a 100644
--- a/pkg/controller/controller_test.go
+++ b/pkg/controller/controller_test.go
@@ -415,3 +415,7 @@ func TestValidateIsolationGroups(t *testing.T) {
 		}
 	}
 }
+
+// TODO: Add test that ensures creates any missing stateful sets
+// not in any particular order (if statefulset is missing, but it's
+// not the last alphabetically it should create it)