feat: support pause reconcile

client/apis/apps/v1/helper/helper.go

@@ -23,6 +23,13 @@ import (
 
 const (
 	DeleteSlotsAnn = "delete-slots"
+
+	// PausedReconcileAnn is the annotation key for the paused reconcile.
+	// If the value is "true", the controller will not reconcile the statefulset.
+	// The controller will not delete the pods, update the status, etc.
+	// We use an annotation instead of a field in the status
+	// so that we can convert between the K8s built-in StatefulSet and ours.
+	PausedReconcileAnn = "paused-reconcile"
 )
 
 func GetDeleteSlots(set metav1.Object) (deleteSlots sets.Int32) {
@@ -87,6 +94,31 @@ func GetMaxReplicaCountAndDeleteSlots(replicas int32, deleteSlots sets.Int32) (i
 	return replicaCount, deleteSlotsCopy
 }
 
+func SetPausedReconcile(set metav1.Object, paused bool) {
+	annotations := set.GetAnnotations()
+	if annotations == nil {
+		annotations = make(map[string]string)
+	}
+	if paused {
+		annotations[PausedReconcileAnn] = "true"
+	} else {
+		delete(annotations, PausedReconcileAnn)
+	}
+	set.SetAnnotations(annotations)
+}
+
+func GetPausedReconcile(set metav1.Object) bool {
+	annotations := set.GetAnnotations()
+	if annotations == nil {
+		return false
+	}
+	value, ok := annotations[PausedReconcileAnn]
+	if !ok {
+		return false
+	}
+	return value == "true"
+}
+
 func GetPodOrdinals(replicas int32, set metav1.Object) sets.Int32 {
 	return GetPodOrdinalsFromReplicasAndDeleteSlots(replicas, GetDeleteSlots(set))
 }

client/apis/apps/v1/helper/helper_test.go

@@ -181,6 +181,127 @@ func TestSetDeleteSlots(t *testing.T) {
 	}
 }
 
+func TestSetPausedReconcile(t *testing.T) {
+	tests := []struct {
+		name string
+		set  bool
+		sts  asappsv1.StatefulSet
+		want asappsv1.StatefulSet
+	}{
+		{
+			name: "set paused reconcile",
+			set:  true,
+			sts: asappsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{},
+			},
+			want: asappsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{
+						PausedReconcileAnn: "true",
+					},
+				},
+			},
+		},
+		{
+			name: "unset paused reconcile",
+			set:  false,
+			sts: asappsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{
+						PausedReconcileAnn: "true",
+					},
+				},
+			},
+			want: asappsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{},
+				},
+			},
+		},
+		{
+			name: "toggle paused reconcile to false",
+			set:  false,
+			sts: asappsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{
+						PausedReconcileAnn: "true",
+					},
+				},
+			},
+			want: asappsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{},
+				},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			SetPausedReconcile(&tt.sts, tt.set)
+			if diff := cmp.Diff(tt.want, tt.sts); diff != "" {
+				t.Errorf("unexpected result (-want, +got): %s", diff)
+			}
+		})
+	}
+}
+
+func TestGetPausedReconcile(t *testing.T) {
+	tests := []struct {
+		name string
+		sts  asappsv1.StatefulSet
+		want bool
+	}{
+		{
+			name: "no annotation",
+			sts: asappsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{},
+			},
+			want: false,
+		},
+		{
+			name: "empty annotation",
+			sts: asappsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{},
+				},
+			},
+			want: false,
+		},
+		{
+			name: "paused reconcile",
+			sts: asappsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{
+						PausedReconcileAnn: "true",
+					},
+				},
+			},
+			want: true,
+		},
+		{
+			name: "not paused reconcile",
+			sts: asappsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{
+						PausedReconcileAnn: "false",
+					},
+				},
+			},
+			want: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := GetPausedReconcile(&tt.sts)
+			if got != tt.want {
+				t.Errorf("GetPausedReconcile want %v got %v", tt.want, got)
+			}
+		})
+	}
+}
+
 func int32ptr(i int32) *int32 {
 	return &i
 }

pkg/controller/statefulset/stateful_set.go

@@ -474,6 +474,12 @@ func (ssc *StatefulSetController) sync(key string) error {
 		return err
 	}
 
+	// If the StatefulSet is paused, don't do anything.
+	if helper.GetPausedReconcile(set) {
+		klog.V(4).Infof("StatefulSet %v/%v is paused, skipping", set.Namespace, set.Name)
+		return nil
+	}
+
 	selector, err := metav1.LabelSelectorAsSelector(set.Spec.Selector)
 	if err != nil {
 		utilruntime.HandleError(fmt.Errorf("error converting StatefulSet %v selector: %v", key, err))