CLOUDP-59612: Configure StatefulSet UpdatePolicy (#31)

Author: chatton

pkg/apis/mongodb/v1/mongodb_types.go

@@ -19,6 +19,12 @@ const (
 	Running Phase = "Running"
 )
 
+const (
+	// LastVersionAnnotationKey should indicate which version of MongoDB was last
+	// configured
+	LastVersionAnnotationKey = "lastVersion"
+)
+
 // MongoDBSpec defines the desired state of MongoDB
 type MongoDBSpec struct {
 	// Members is the number of members in the replica set
@@ -59,6 +65,13 @@ func (m *MongoDB) UpdateSuccess() {
 	m.Status.Phase = Running
 }
 
+func (m MongoDB) ChangingVersion() bool {
+	if lastVersion, ok := m.Annotations[LastVersionAnnotationKey]; ok {
+		return (m.Spec.Version != lastVersion) && lastVersion != ""
+	}
+	return false
+}
+
 // MongoURI returns a mongo uri which can be used to connect to this deployment
 func (m MongoDB) MongoURI() string {
 	members := make([]string, m.Spec.Members)

pkg/controller/mongodb/mongodb_controller.go

@@ -8,6 +8,8 @@ import (
 	"os"
 	"time"
 
+	"github.com/mongodb/mongodb-kubernetes-operator/pkg/controller/predicates"
+
 	mdbv1 "github.com/mongodb/mongodb-kubernetes-operator/pkg/apis/mongodb/v1"
 	"github.com/mongodb/mongodb-kubernetes-operator/pkg/automationconfig"
 	mdbClient "github.com/mongodb/mongodb-kubernetes-operator/pkg/kube/client"
@@ -57,6 +59,7 @@ func newReconciler(mgr manager.Manager, manifestProvider ManifestProvider) recon
 		client:           mdbClient.NewClient(mgrClient),
 		scheme:           mgr.GetScheme(),
 		manifestProvider: manifestProvider,
+		log:              zap.S(),
 	}
 }
 
@@ -69,7 +72,7 @@ func add(mgr manager.Manager, r reconcile.Reconciler) error {
 	}
 
 	// Watch for changes to primary resource MongoDB
-	err = c.Watch(&source.Kind{Type: &mdbv1.MongoDB{}}, &handler.EnqueueRequestForObject{})
+	err = c.Watch(&source.Kind{Type: &mdbv1.MongoDB{}}, &handler.EnqueueRequestForObject{}, predicates.OnlyOnSpecChange())
 	if err != nil {
 		return err
 	}
@@ -86,6 +89,7 @@ type ReplicaSetReconciler struct {
 	client           mdbClient.Client
 	scheme           *runtime.Scheme
 	manifestProvider func() (automationconfig.VersionManifest, error)
+	log              *zap.SugaredLogger
 }
 
 // Reconcile reads that state of the cluster for a MongoDB object and makes changes based on the state read
@@ -94,8 +98,8 @@ type ReplicaSetReconciler struct {
 // The Controller will requeue the Request to be processed again if the returned error is non-nil or
 // Result.Requeue is true, otherwise upon completion it will remove the work from the queue.
 func (r *ReplicaSetReconciler) Reconcile(request reconcile.Request) (reconcile.Result, error) {
-	log := zap.S().With("ReplicaSet", request.NamespacedName)
-	log.Info("Reconciling MongoDB")
+	r.log = zap.S().With("ReplicaSet", request.NamespacedName)
+	r.log.Info("Reconciling MongoDB")
 
 	// TODO: generalize preparation for resource
 	// Fetch the MongoDB instance
@@ -108,61 +112,116 @@ func (r *ReplicaSetReconciler) Reconcile(request reconcile.Request) (reconcile.R
 			// Return and don't requeue
 			return reconcile.Result{}, nil
 		}
-		log.Errorf("error reconciling MongoDB resource: %s", err)
+		r.log.Errorf("error reconciling MongoDB resource: %s", err)
 		// Error reading the object - requeue the request.
 		return reconcile.Result{}, err
 	}
 
 	// TODO: Read current automation config version from config map
 	if err := r.ensureAutomationConfig(mdb); err != nil {
-		log.Warnf("error creating automation config config map: %s", err)
+		r.log.Infof("error creating automation config config map: %s", err)
 		return reconcile.Result{}, err
 	}
 
 	svc := buildService(mdb)
 	if err = r.client.CreateOrUpdate(&svc); err != nil {
-		log.Warnf("The service already exists... moving forward: %s", err)
+		r.log.Infof("The service already exists... moving forward: %s", err)
 	}
 
-	sts, err := buildStatefulSet(mdb)
-	if err != nil {
-		log.Infof("Error building StatefulSet: %s", err)
-		return reconcile.Result{}, nil
+	if err := r.createOrUpdateStatefulSet(mdb); err != nil {
+		r.log.Infof("Error creating/updating StatefulSet: %+v", err)
+		return reconcile.Result{}, err
 	}
 
-	if err = r.client.CreateOrUpdate(&sts); err != nil {
-		log.Infof("Error creating/updating StatefulSet: %s", err)
+	if ready, err := r.isStatefulSetReady(mdb); err != nil {
+		r.log.Infof("error checking StatefulSet status: %+v", err)
 		return reconcile.Result{}, err
-	} else {
-		log.Infof("StatefulSet successfully Created/Updated")
+	} else if !ready {
+		r.log.Infof("StatefulSet %s/%s is not yet ready, retrying in 10 seconds", mdb.Namespace, mdb.Name)
+		return reconcile.Result{RequeueAfter: time.Second * 10}, nil
 	}
 
-	log.Debugf("waiting for StatefulSet %s/%s to reach ready state", mdb.Namespace, mdb.Name)
-	set := appsv1.StatefulSet{}
-	if err := r.client.Get(context.TODO(), types.NamespacedName{Name: mdb.Name, Namespace: mdb.Namespace}, &set); err != nil {
-		log.Infof("Error getting StatefulSet: %s", err)
+	if err := r.resetStatefulSetUpdateStrategy(mdb); err != nil {
+		r.log.Infof("error resetting StatefulSet UpdateStrategyType: %+v", err)
 		return reconcile.Result{}, err
 	}
 
-	if !statefulset.IsReady(set) {
-		log.Infof("Stateful Set has not yet reached the ready state, requeuing reconciliation")
-		return reconcile.Result{RequeueAfter: time.Second * 10}, nil
+	if err := r.setAnnotation(types.NamespacedName{Name: mdb.Name, Namespace: mdb.Namespace}, mdbv1.LastVersionAnnotationKey, mdb.Spec.Version); err != nil {
+		r.log.Infof("Error setting annotation: %+v", err)
+		return reconcile.Result{}, err
 	}
 
-	log.Infof("Stateful Set reached ready state!")
-
 	if err := r.updateStatusSuccess(&mdb); err != nil {
-		log.Infof("Error updating the status of the MongoDB resource: %+v", err)
-		return reconcile.Result{}, nil
+		r.log.Infof("Error updating the status of the MongoDB resource: %+v", err)
+		return reconcile.Result{}, err
 	}
 
-	log.Info("Successfully finished reconciliation", "MongoDB.Spec:", mdb.Spec, "MongoDB.Status", mdb.Status)
+	r.log.Info("Successfully finished reconciliation", "MongoDB.Spec:", mdb.Spec, "MongoDB.Status", mdb.Status)
 	return reconcile.Result{}, nil
 }
 
+// resetStatefulSetUpdateStrategy ensures the stateful set is configured back to using RollingUpdateStatefulSetStrategyType
+// and does not keep using OnDelete
+func (r *ReplicaSetReconciler) resetStatefulSetUpdateStrategy(mdb mdbv1.MongoDB) error {
+	if !mdb.ChangingVersion() {
+		return nil
+	}
+	// if we changed the version, we need to reset the UpdatePolicy back to OnUpdate
+	sts := &appsv1.StatefulSet{}
+	return r.client.GetAndUpdate(types.NamespacedName{Name: mdb.Name, Namespace: mdb.Namespace}, sts, func() {
+		sts.Spec.UpdateStrategy.Type = appsv1.RollingUpdateStatefulSetStrategyType
+	})
+}
+
+// isStatefulSetReady checks to see if the stateful set corresponding to the given MongoDB resource
+// is currently in the ready state
+func (r *ReplicaSetReconciler) isStatefulSetReady(mdb mdbv1.MongoDB) (bool, error) {
+	set := appsv1.StatefulSet{}
+	if err := r.client.Get(context.TODO(), types.NamespacedName{Name: mdb.Name, Namespace: mdb.Namespace}, &set); err != nil {
+		return false, fmt.Errorf("error getting StatefulSet: %s", err)
+	}
+	return statefulset.IsReady(set), nil
+}
+
+func (r *ReplicaSetReconciler) createOrUpdateStatefulSet(mdb mdbv1.MongoDB) error {
+	sts, err := buildStatefulSet(mdb)
+	if err != nil {
+		return fmt.Errorf("error building StatefulSet: %s", err)
+	}
+	if err = r.client.CreateOrUpdate(&sts); err != nil {
+		return fmt.Errorf("error creating/updating StatefulSet: %s", err)
+	}
+
+	r.log.Debugf("Waiting for StatefulSet %s/%s to reach ready state", mdb.Namespace, mdb.Name)
+	set := appsv1.StatefulSet{}
+	if err := r.client.Get(context.TODO(), types.NamespacedName{Name: mdb.Name, Namespace: mdb.Namespace}, &set); err != nil {
+		return fmt.Errorf("error getting StatefulSet: %s", err)
+	}
+	return nil
+}
+
+// setAnnotation updates the monogdb resource with the given namespaced name and sets the annotation
+// "key" with the provided value "val"
+func (r ReplicaSetReconciler) setAnnotation(nsName types.NamespacedName, key, val string) error {
+	mdb := mdbv1.MongoDB{}
+	return r.client.GetAndUpdate(nsName, &mdb, func() {
+		if mdb.Annotations == nil {
+			mdb.Annotations = map[string]string{}
+		}
+		mdb.Annotations[key] = val
+	})
+}
+
+// updateStatusSuccess should be called after a successful reconciliation
+// the resource's status is updated to reflect to the state, and any other cleanup
+// operators should be performed here
 func (r ReplicaSetReconciler) updateStatusSuccess(mdb *mdbv1.MongoDB) error {
-	mdb.UpdateSuccess()
-	if err := r.client.Status().Update(context.TODO(), mdb); err != nil {
+	newMdb := &mdbv1.MongoDB{}
+	if err := r.client.Get(context.TODO(), types.NamespacedName{Name: mdb.Name, Namespace: mdb.Namespace}, newMdb); err != nil {
+		return fmt.Errorf("error getting resource: %+v", err)
+	}
+	newMdb.UpdateSuccess()
+	if err := r.client.Status().Update(context.TODO(), newMdb); err != nil {
 		return fmt.Errorf("error updating status: %+v", err)
 	}
 	return nil
@@ -293,6 +352,15 @@ func defaultReadinessProbe() corev1.Probe {
 	}
 }
 
+// getUpdateStrategyType returns the type of RollingUpgradeStrategy that the StatefulSet
+// should be configured with
+func getUpdateStrategyType(mdb mdbv1.MongoDB) appsv1.StatefulSetUpdateStrategyType {
+	if !mdb.ChangingVersion() {
+		return appsv1.RollingUpdateStatefulSetStrategyType
+	}
+	return appsv1.OnDeleteStatefulSetStrategyType
+}
+
 // buildStatefulSet takes a MongoDB resource and converts it into
 // the corresponding stateful set
 func buildStatefulSet(mdb mdbv1.MongoDB) (appsv1.StatefulSet, error) {
@@ -316,7 +384,8 @@ func buildStatefulSet(mdb mdbv1.MongoDB) (appsv1.StatefulSet, error) {
 		SetReplicas(mdb.Spec.Members).
 		SetLabels(labels).
 		SetMatchLabels(labels).
-		SetServiceName(mdb.ServiceName())
+		SetServiceName(mdb.ServiceName()).
+		SetUpdateStrategy(getUpdateStrategyType(mdb))
 
 	// TODO: Add this section to architecture document.
 	// The design of the multi-container and the different volumes mounted to them is as follows:

pkg/controller/mongodb/replicaset_controller_test.go

@@ -26,8 +26,9 @@ func init() {
 func newTestReplicaSet() mdbv1.MongoDB {
 	return mdbv1.MongoDB{
 		ObjectMeta: metav1.ObjectMeta{
-			Name:      "my-rs",
-			Namespace: "my-ns",
+			Name:        "my-rs",
+			Namespace:   "my-ns",
+			Annotations: map[string]string{},
 		},
 		Spec: mdbv1.MongoDBSpec{
 			Members: 3,
@@ -102,3 +103,62 @@ func TestStatefulSet_IsCorrectlyConfigured(t *testing.T) {
 
 	assert.Equal(t, resourcerequirements.Defaults(), agentContainer.Resources)
 }
+
+func TestChangingVersion_ResultsInRollingUpdateStrategyType(t *testing.T) {
+	mdb := newTestReplicaSet()
+	mgr := client.NewManager(&mdb)
+	mgrClient := mgr.GetClient()
+	r := newReconciler(mgr, mockManifestProvider(mdb.Spec.Version))
+	res, err := r.Reconcile(reconcile.Request{NamespacedName: types.NamespacedName{Namespace: mdb.Namespace, Name: mdb.Name}})
+	assertReconciliationSuccessful(t, res, err)
+
+	// fetch updated resource after first reconciliation
+	_ = mgrClient.Get(context.TODO(), types.NamespacedName{Name: mdb.Name, Namespace: mdb.Namespace}, &mdb)
+
+	sts := appsv1.StatefulSet{}
+	err = mgrClient.Get(context.TODO(), types.NamespacedName{Name: mdb.Name, Namespace: mdb.Namespace}, &sts)
+	assert.NoError(t, err)
+	assert.Equal(t, appsv1.RollingUpdateStatefulSetStrategyType, sts.Spec.UpdateStrategy.Type)
+
+	mdbRef := &mdb
+	mdbRef.Spec.Version = "4.2.3"
+
+	_ = mgrClient.Update(context.TODO(), &mdb)
+
+	res, err = r.Reconcile(reconcile.Request{NamespacedName: types.NamespacedName{Namespace: mdb.Namespace, Name: mdb.Name}})
+	assertReconciliationSuccessful(t, res, err)
+
+	sts = appsv1.StatefulSet{}
+	err = mgrClient.Get(context.TODO(), types.NamespacedName{Name: mdb.Name, Namespace: mdb.Namespace}, &sts)
+	assert.NoError(t, err)
+
+	assert.Equal(t, appsv1.RollingUpdateStatefulSetStrategyType, sts.Spec.UpdateStrategy.Type,
+		"The StatefulSet should have be re-configured to use RollingUpdates after it reached the ready state")
+}
+
+func TestBuildStatefulSet_ConfiguresUpdateStrategyCorrectly(t *testing.T) {
+	t.Run("On No Version Change, Same Version", func(t *testing.T) {
+		mdb := newTestReplicaSet()
+		mdb.Spec.Version = "4.0.0"
+		mdb.Annotations[mdbv1.LastVersionAnnotationKey] = "4.0.0"
+		sts, err := buildStatefulSet(mdb)
+		assert.NoError(t, err)
+		assert.Equal(t, appsv1.RollingUpdateStatefulSetStrategyType, sts.Spec.UpdateStrategy.Type)
+	})
+	t.Run("On No Version Change, First Version", func(t *testing.T) {
+		mdb := newTestReplicaSet()
+		mdb.Spec.Version = "4.0.0"
+		delete(mdb.Annotations, mdbv1.LastVersionAnnotationKey)
+		sts, err := buildStatefulSet(mdb)
+		assert.NoError(t, err)
+		assert.Equal(t, appsv1.RollingUpdateStatefulSetStrategyType, sts.Spec.UpdateStrategy.Type)
+	})
+	t.Run("On Version Change", func(t *testing.T) {
+		mdb := newTestReplicaSet()
+		mdb.Spec.Version = "4.0.0"
+		mdb.Annotations[mdbv1.LastVersionAnnotationKey] = "4.2.0"
+		sts, err := buildStatefulSet(mdb)
+		assert.NoError(t, err)
+		assert.Equal(t, appsv1.OnDeleteStatefulSetStrategyType, sts.Spec.UpdateStrategy.Type)
+	})
+}

pkg/controller/predicates/predicates.go

@@ -0,0 +1,24 @@
+package predicates
+
+import (
+	"reflect"
+
+	mdbv1 "github.com/mongodb/mongodb-kubernetes-operator/pkg/apis/mongodb/v1"
+	"sigs.k8s.io/controller-runtime/pkg/event"
+	"sigs.k8s.io/controller-runtime/pkg/predicate"
+)
+
+// OnlyOnSpecChange returns a set of predicates indicating
+// that reconciliations should only happen on changes to the Spec of the resource.
+// any other changes won't trigger a reconciliation. This allows us to freely update the annotations
+// of the resource without triggering unintentional reconciliations.
+func OnlyOnSpecChange() predicate.Funcs {
+	return predicate.Funcs{
+		UpdateFunc: func(e event.UpdateEvent) bool {
+			oldResource := e.ObjectOld.(*mdbv1.MongoDB)
+			newResource := e.ObjectNew.(*mdbv1.MongoDB)
+			specChanged := !reflect.DeepEqual(oldResource.Spec, newResource.Spec)
+			return specChanged
+		},
+	}
+}

pkg/kube/client/client.go

@@ -19,6 +19,7 @@ func NewClient(c k8sClient.Client) Client {
 type Client interface {
 	k8sClient.Client
 	CreateOrUpdate(obj runtime.Object) error
+	GetAndUpdate(nsName types.NamespacedName, obj runtime.Object, updateFunc func()) error
 }
 
 type client struct {
@@ -39,6 +40,19 @@ func (c client) CreateOrUpdate(obj runtime.Object) error {
 	return c.Update(context.TODO(), obj)
 }
 
+// GetAndUpdate fetches the most recent version of the runtime.Object with the provided
+// nsName and applies the update function. The update function should update "obj" from
+// an outer scope
+func (c client) GetAndUpdate(nsName types.NamespacedName, obj runtime.Object, updateFunc func()) error {
+	err := c.Get(context.TODO(), nsName, obj)
+	if err != nil {
+		return err
+	}
+	// apply the function on the most recent version of the resource
+	updateFunc()
+	return c.Update(context.TODO(), obj)
+}
+
 func namespacedNameFromObject(obj runtime.Object) types.NamespacedName {
 	ns := reflect.ValueOf(obj).Elem().FieldByName("Namespace").String()
 	name := reflect.ValueOf(obj).Elem().FieldByName("Name").String()

pkg/kube/client/mocked_client.go

@@ -59,15 +59,15 @@ func (m *mockedClient) Create(_ context.Context, obj runtime.Object, _ ...k8sCli
 
 	switch v := obj.(type) {
 	case *appsv1.StatefulSet:
-		onStatefulsetUpdate(v)
+		makeStatefulSetReady(v)
 	}
 
 	relevantMap[objKey] = obj
 	return nil
 }
 
-// onStatefulsetUpdate configures the statefulset to be in the running state.
-func onStatefulsetUpdate(set *appsv1.StatefulSet) {
+// makeStatefulSetReady configures the statefulset to be in the running state.
+func makeStatefulSetReady(set *appsv1.StatefulSet) {
 	set.Status.UpdatedReplicas = *set.Spec.Replicas
 	set.Status.ReadyReplicas = *set.Spec.Replicas
 }
@@ -86,6 +86,10 @@ func (m *mockedClient) Update(_ context.Context, obj runtime.Object, _ ...k8sCli
 	if err != nil {
 		return err
 	}
+	switch v := obj.(type) {
+	case *appsv1.StatefulSet:
+		makeStatefulSetReady(v)
+	}
 	relevantMap[objKey] = obj
 	return nil
 }