From 6559050ee15e87c4fdb6712476dfd6881bc877e3 Mon Sep 17 00:00:00 2001
From: Gunju Kim <gjkim042@gmail.com>
Date: Mon, 26 Sep 2022 22:11:07 +0900
Subject: [PATCH] StatefulSet: Cleanup the complex defer function updating the
 status

In the long term, the complex defer function makes the code harder to
maintain as code after it should take that into account. This removes
the complex defer function updating the status of a statefulset.
---
 .../statefulset/stateful_set_control.go       | 43 +++++------
 .../statefulset/stateful_set_control_test.go  | 72 +++++++++++++++++++
 2 files changed, 95 insertions(+), 20 deletions(-)

diff --git a/pkg/controller/statefulset/stateful_set_control.go b/pkg/controller/statefulset/stateful_set_control.go
index cc281c0a710..30d4d626871 100644
--- a/pkg/controller/statefulset/stateful_set_control.go
+++ b/pkg/controller/statefulset/stateful_set_control.go
@@ -105,8 +105,28 @@ func (ssc *defaultStatefulSetControl) performUpdate(
 
 	// perform the main update function and get the status
 	currentStatus, err = ssc.updateStatefulSet(ctx, set, currentRevision, updateRevision, collisionCount, pods)
-	if err != nil {
+	if err != nil && currentStatus == nil {
+		return currentRevision, updateRevision, nil, err
+	}
+
+	// make sure to update the latest status even if there is an error with non-nil currentStatus
+	statusErr := ssc.updateStatefulSetStatus(ctx, set, currentStatus)
+	if statusErr == nil {
+		klog.V(4).InfoS("Updated status", "statefulSet", klog.KObj(set),
+			"replicas", currentStatus.Replicas,
+			"readyReplicas", currentStatus.ReadyReplicas,
+			"currentReplicas", currentStatus.CurrentReplicas,
+			"updatedReplicas", currentStatus.UpdatedReplicas)
+	}
+
+	switch {
+	case err != nil && statusErr != nil:
+		klog.ErrorS(statusErr, "Could not update status", "statefulSet", klog.KObj(set))
 		return currentRevision, updateRevision, currentStatus, err
+	case err != nil:
+		return currentRevision, updateRevision, currentStatus, err
+	case statusErr != nil:
+		return currentRevision, updateRevision, currentStatus, statusErr
 	}
 
 	klog.V(4).InfoS("StatefulSet revisions", "statefulSet", klog.KObj(set),
@@ -263,7 +283,7 @@ func (ssc *defaultStatefulSetControl) updateStatefulSet(
 	currentRevision *apps.ControllerRevision,
 	updateRevision *apps.ControllerRevision,
 	collisionCount int32,
-	pods []*v1.Pod) (statefulSetStatus *apps.StatefulSetStatus, updateErr error) {
+	pods []*v1.Pod) (*apps.StatefulSetStatus, error) {
 	// get the current and update revisions of the set.
 	currentSet, err := ApplyRevision(set, currentRevision)
 	if err != nil {
@@ -325,23 +345,6 @@ func (ssc *defaultStatefulSetControl) updateStatefulSet(
 		// If the ordinal could not be parsed (ord < 0), ignore the Pod.
 	}
 
-	// make sure to update the latest status even if there is an error later
-	defer func() {
-		// update the set's status
-		statusErr := ssc.updateStatefulSetStatus(ctx, set, &status)
-		if statusErr == nil {
-			klog.V(4).InfoS("Updated status", "statefulSet", klog.KObj(set),
-				"replicas", status.Replicas,
-				"readyReplicas", status.ReadyReplicas,
-				"currentReplicas", status.CurrentReplicas,
-				"updatedReplicas", status.UpdatedReplicas)
-		} else if updateErr == nil {
-			updateErr = statusErr
-		} else {
-			klog.V(4).InfoS("Could not update status", "statefulSet", klog.KObj(set), "err", statusErr)
-		}
-	}()
-
 	// for any empty indices in the sequence [0,set.Spec.Replicas) create a new Pod at the correct revision
 	for ord := 0; ord < replicaCount; ord++ {
 		if replicas[ord] == nil {
@@ -647,7 +650,7 @@ func updateStatefulSetAfterInvariantEstablished(
 				replicas[target].Name)
 			if err := ssc.podControl.DeleteStatefulPod(set, replicas[target]); err != nil {
 				if !errors.IsNotFound(err) {
-					return nil, err
+					return &status, err
 				}
 			}
 			deletedPods++
diff --git a/pkg/controller/statefulset/stateful_set_control_test.go b/pkg/controller/statefulset/stateful_set_control_test.go
index 43e06388df7..8ad7b586ff6 100644
--- a/pkg/controller/statefulset/stateful_set_control_test.go
+++ b/pkg/controller/statefulset/stateful_set_control_test.go
@@ -2106,6 +2106,78 @@ func TestStatefulSetAvailability(t *testing.T) {
 	}
 }
 
+func TestStatefulSetStatusUpdate(t *testing.T) {
+	var (
+		syncErr   = fmt.Errorf("sync error")
+		statusErr = fmt.Errorf("status error")
+	)
+
+	testCases := []struct {
+		desc string
+
+		hasSyncErr   bool
+		hasStatusErr bool
+
+		expectedErr error
+	}{
+		{
+			desc:         "no error",
+			hasSyncErr:   false,
+			hasStatusErr: false,
+			expectedErr:  nil,
+		},
+		{
+			desc:         "sync error",
+			hasSyncErr:   true,
+			hasStatusErr: false,
+			expectedErr:  syncErr,
+		},
+		{
+			desc:         "status error",
+			hasSyncErr:   false,
+			hasStatusErr: true,
+			expectedErr:  statusErr,
+		},
+		{
+			desc:         "sync and status error",
+			hasSyncErr:   true,
+			hasStatusErr: true,
+			expectedErr:  syncErr,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			set := newStatefulSet(3)
+			client := fake.NewSimpleClientset(set)
+			om, ssu, ssc := setupController(client)
+
+			if tc.hasSyncErr {
+				om.SetCreateStatefulPodError(syncErr, 0)
+			}
+			if tc.hasStatusErr {
+				ssu.SetUpdateStatefulSetStatusError(statusErr, 0)
+			}
+
+			selector, err := metav1.LabelSelectorAsSelector(set.Spec.Selector)
+			if err != nil {
+				t.Error(err)
+			}
+			pods, err := om.podsLister.Pods(set.Namespace).List(selector)
+			if err != nil {
+				t.Error(err)
+			}
+			_, err = ssc.UpdateStatefulSet(context.TODO(), set, pods)
+			if ssu.updateStatusTracker.requests != 1 {
+				t.Errorf("Did not update status")
+			}
+			if !errors.Is(err, tc.expectedErr) {
+				t.Errorf("Expected error: %v, got: %v", tc.expectedErr, err)
+			}
+		})
+	}
+}
+
 type requestTracker struct {
 	requests int
 	err      error