feature: range reserveOrdinals for AdvancedStatefulSet (#1873)

* feature: range reserveOrdinals for AdvancedStatefulSet Signed-off-by: AiRanthem <[email protected]> * feature: range reserveOrdinals for AdvancedStatefulSet Signed-off-by: AiRanthem <[email protected]> * feature: range reserveOrdinals for AdvancedStatefulSet Signed-off-by: AiRanthem <[email protected]> * feature: range reserveOrdinals for AdvancedStatefulSet Signed-off-by: AiRanthem <[email protected]> * feature: range reserveOrdinals for AdvancedStatefulSet Signed-off-by: AiRanthem <[email protected]> --------- Signed-off-by: AiRanthem <[email protected]>
openkruise · Feb 12, 2025 · 8f727a4 · 8f727a4
1 parent 4183fbc
commit 8f727a4
Show file tree

Hide file tree

Showing 18 changed files with 528 additions and 369 deletions.
diff --git a/apis/apps/v1beta1/statefulset_types.go b/apis/apps/v1beta1/statefulset_types.go
@@ -17,12 +17,11 @@ limitations under the License.
 package v1beta1
 
 import (
+	appspub "github.com/openkruise/kruise/apis/apps/pub"
 	apps "k8s.io/api/apps/v1"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
-
-	appspub "github.com/openkruise/kruise/apis/apps/pub"
 )
 
 const (
@@ -275,7 +274,8 @@ type StatefulSetSpec struct {
 	//   Then controller will delete Pod-1 and create Pod-3 (existing Pods will be [0, 2, 3])
 	// - If you just want to delete Pod-1, you should set spec.reserveOrdinal to [1] and spec.replicas to 2.
 	//   Then controller will delete Pod-1 (existing Pods will be [0, 2])
-	ReserveOrdinals []int `json:"reserveOrdinals,omitempty"`
+	// You can also use ranges along with numbers, such as [1, 3-5], which is a shortcut for [1, 3, 4, 5].
+	ReserveOrdinals []intstr.IntOrString `json:"reserveOrdinals,omitempty"`
 
 	// Lifecycle defines the lifecycle hooks for Pods pre-delete, in-place update.
 	Lifecycle *appspub.Lifecycle `json:"lifecycle,omitempty"`

diff --git a/apis/apps/v1beta1/zz_generated.deepcopy.go b/apis/apps/v1beta1/zz_generated.deepcopy.go
diff --git a/config/crd/bases/apps.kruise.io_statefulsets.yaml b/config/crd/bases/apps.kruise.io_statefulsets.yaml
@@ -652,8 +652,12 @@ spec:
                     Then controller will delete Pod-1 and create Pod-3 (existing Pods will be [0, 2, 3])
                   - If you just want to delete Pod-1, you should set spec.reserveOrdinal to [1] and spec.replicas to 2.
                     Then controller will delete Pod-1 (existing Pods will be [0, 2])
+                  You can also use ranges along with numbers, such as [1, 3-5], which is a shortcut for [1, 3, 4, 5].
                 items:
-                  type: integer
+                  anyOf:
+                  - type: integer
+                  - type: string
+                  x-kubernetes-int-or-string: true
                 type: array
               revisionHistoryLimit:
                 description: |-

diff --git a/config/crd/bases/apps.kruise.io_uniteddeployments.yaml b/config/crd/bases/apps.kruise.io_uniteddeployments.yaml
@@ -275,8 +275,12 @@ spec:
                                 Then controller will delete Pod-1 and create Pod-3 (existing Pods will be [0, 2, 3])
                               - If you just want to delete Pod-1, you should set spec.reserveOrdinal to [1] and spec.replicas to 2.
                                 Then controller will delete Pod-1 (existing Pods will be [0, 2])
+                              You can also use ranges along with numbers, such as [1, 3-5], which is a shortcut for [1, 3, 4, 5].
                             items:
-                              type: integer
+                              anyOf:
+                              - type: integer
+                              - type: string
+                              x-kubernetes-int-or-string: true
                             type: array
                           revisionHistoryLimit:
                             description: |-

diff --git a/pkg/controller/persistentpodstate/persistent_pod_state_controller.go b/pkg/controller/persistentpodstate/persistent_pod_state_controller.go
@@ -24,6 +24,7 @@ import (
 	"strings"
 
 	"github.com/openkruise/kruise/pkg/util/configuration"
+	"k8s.io/utils/ptr"
 
 	ctrlUtil "github.com/openkruise/kruise/pkg/controller/util"
 
@@ -36,7 +37,6 @@ import (
 	"k8s.io/client-go/util/retry"
 	"k8s.io/klog/v2"
 	podutil "k8s.io/kubernetes/pkg/api/v1/pod"
-	utilpointer "k8s.io/utils/pointer"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/controller"
@@ -146,7 +146,7 @@ type innerStatefulset struct {
 	// replicas
 	Replicas int32
 	// kruise statefulset filed
-	ReserveOrdinals   []int
+	ReserveOrdinals   sets.Set[int]
 	DeletionTimestamp *metav1.Time
 }
 
@@ -354,11 +354,10 @@ func (r *ReconcilePersistentPodState) getPodState(pod *corev1.Pod, nodeTopologyK
 }
 
 func isInStatefulSetReplicas(index int, sts *innerStatefulset) bool {
-	reserveOrdinals := sets.NewInt(sts.ReserveOrdinals...)
 	replicas := sets.NewInt()
 	replicaIndex := 0
 	for realReplicaCount := 0; realReplicaCount < int(sts.Replicas); replicaIndex++ {
-		if reserveOrdinals.Has(replicaIndex) {
+		if sts.ReserveOrdinals.Has(replicaIndex) {
 			continue
 		}
 		realReplicaCount++
@@ -457,7 +456,7 @@ func newStatefulSetPersistentPodState(workload *controllerfinder.ScaleAndSelecto
 					APIVersion: workload.APIVersion,
 					Kind:       workload.Kind,
 					Name:       workload.Name,
-					Controller: utilpointer.BoolPtr(true),
+					Controller: ptr.To(true),
 					UID:        workload.UID,
 				},
 			},

diff --git a/pkg/controller/persistentpodstate/persistent_pod_state_controller_test.go b/pkg/controller/persistentpodstate/persistent_pod_state_controller_test.go
@@ -23,7 +23,9 @@ import (
 	"testing"
 	"time"
 
+	"k8s.io/apimachinery/pkg/util/intstr"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
+	"k8s.io/utils/ptr"
 
 	appsv1alpha1 "github.com/openkruise/kruise/apis/apps/v1alpha1"
 	appsv1beta1 "github.com/openkruise/kruise/apis/apps/v1beta1"
@@ -36,7 +38,6 @@ import (
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/kubernetes/pkg/apis/apps"
-	"k8s.io/utils/pointer"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
@@ -56,7 +57,7 @@ var (
 			UID:       "012d18d5-5eb9-449d-b670-3da8fec8852f",
 		},
 		Spec: appsv1beta1.StatefulSetSpec{
-			Replicas: pointer.Int32Ptr(10),
+			Replicas: ptr.To[int32](10),
 			Template: corev1.PodTemplateSpec{
 				ObjectMeta: metav1.ObjectMeta{
 					Annotations: map[string]string{},
@@ -70,7 +71,7 @@ var (
 			Namespace: "ns-test",
 			OwnerReferences: []metav1.OwnerReference{
 				{
-					Controller: pointer.BoolPtr(true),
+					Controller: ptr.To(true),
 				},
 			},
 			Annotations: map[string]string{
@@ -222,7 +223,7 @@ func TestReconcilePersistentPodState(t *testing.T) {
 			name: "kruise statefulset, scale down replicas 10->8, 1 pod deleted, 1 pod running",
 			getSts: func() (*apps.StatefulSet, *appsv1beta1.StatefulSet) {
 				kruise := kruiseStsDemo.DeepCopy()
-				kruise.Spec.Replicas = pointer.Int32Ptr(8)
+				kruise.Spec.Replicas = ptr.To[int32](8)
 				return nil, kruise
 			},
 			getPods: func() []*corev1.Pod {
@@ -316,8 +317,12 @@ func TestReconcilePersistentPodState(t *testing.T) {
 			name: "kruise reserveOrigin statefulset, scale down replicas 10->8, 1 pod deleted, 1 pod running",
 			getSts: func() (*apps.StatefulSet, *appsv1beta1.StatefulSet) {
 				kruise := kruiseStsDemo.DeepCopy()
-				kruise.Spec.Replicas = pointer.Int32Ptr(8)
-				kruise.Spec.ReserveOrdinals = []int{0, 3, 7}
+				kruise.Spec.Replicas = ptr.To[int32](8)
+				kruise.Spec.ReserveOrdinals = []intstr.IntOrString{
+					intstr.FromInt32(0),
+					intstr.FromInt32(3),
+					intstr.FromInt32(7),
+				}
 				return nil, kruise
 			},
 			getPods: func() []*corev1.Pod {

diff --git a/pkg/controller/statefulset/stateful_pod_control_test.go b/pkg/controller/statefulset/stateful_pod_control_test.go
@@ -1,6 +1,5 @@
 /*
 Copyright 2019 The Kruise Authors.
-Copyright 2016 The Kubernetes Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -34,6 +33,7 @@ import (
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/intstr"
 	"k8s.io/client-go/kubernetes/fake"
 	corelisters "k8s.io/client-go/listers/core/v1"
 	storagelisters "k8s.io/client-go/listers/storage/v1"
@@ -1045,7 +1045,10 @@ func TestUpdatePodClaimForRetentionPolicy(t *testing.T) {
 					WhenDeleted: appsv1beta1.DeletePersistentVolumeClaimRetentionPolicyType,
 					WhenScaled:  appsv1beta1.RetainPersistentVolumeClaimRetentionPolicyType,
 				}
-				set.Spec.ReserveOrdinals = []int{2, 4}
+				set.Spec.ReserveOrdinals = []intstr.IntOrString{
+					intstr.FromInt32(2),
+					intstr.FromInt32(4),
+				}
 				return set
 			},
 			getPods: func(set *appsv1beta1.StatefulSet) []*v1.Pod {
@@ -1083,7 +1086,10 @@ func TestUpdatePodClaimForRetentionPolicy(t *testing.T) {
 					WhenDeleted: appsv1beta1.DeletePersistentVolumeClaimRetentionPolicyType,
 					WhenScaled:  appsv1beta1.DeletePersistentVolumeClaimRetentionPolicyType,
 				}
-				set.Spec.ReserveOrdinals = []int{2, 4}
+				set.Spec.ReserveOrdinals = []intstr.IntOrString{
+					intstr.FromInt32(2),
+					intstr.FromInt32(4),
+				}
 				return set
 			},
 			getPods: func(set *appsv1beta1.StatefulSet) []*v1.Pod {
@@ -1121,7 +1127,10 @@ func TestUpdatePodClaimForRetentionPolicy(t *testing.T) {
 					WhenDeleted: appsv1beta1.DeletePersistentVolumeClaimRetentionPolicyType,
 					WhenScaled:  appsv1beta1.DeletePersistentVolumeClaimRetentionPolicyType,
 				}
-				set.Spec.ReserveOrdinals = []int{2, 4}
+				set.Spec.ReserveOrdinals = []intstr.IntOrString{
+					intstr.FromInt32(2),
+					intstr.FromInt32(4),
+				}
 				return set
 			},
 			getPods: func(set *appsv1beta1.StatefulSet) []*v1.Pod {
@@ -1170,7 +1179,10 @@ func TestUpdatePodClaimForRetentionPolicy(t *testing.T) {
 					WhenDeleted: appsv1beta1.DeletePersistentVolumeClaimRetentionPolicyType,
 					WhenScaled:  appsv1beta1.RetainPersistentVolumeClaimRetentionPolicyType,
 				}
-				set.Spec.ReserveOrdinals = []int{2, 4}
+				set.Spec.ReserveOrdinals = []intstr.IntOrString{
+					intstr.FromInt32(2),
+					intstr.FromInt32(4),
+				}
 				return set
 			},
 			getPods: func(set *appsv1beta1.StatefulSet) []*v1.Pod {
@@ -1211,6 +1223,58 @@ func TestUpdatePodClaimForRetentionPolicy(t *testing.T) {
 				}
 			},
 		},
+		{
+			name: "reserveOrdinals is [1,3-5], scaleDown=true, whenScaled=Retain, whenDeleted=Delete",
+			getStatefulSet: func() *appsv1beta1.StatefulSet {
+				set := newStatefulSet(3)
+				set.Spec.PersistentVolumeClaimRetentionPolicy = &appsv1beta1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+					WhenDeleted: appsv1beta1.DeletePersistentVolumeClaimRetentionPolicyType,
+					WhenScaled:  appsv1beta1.RetainPersistentVolumeClaimRetentionPolicyType,
+				}
+				set.Spec.ReserveOrdinals = []intstr.IntOrString{
+					intstr.FromInt32(1),
+					intstr.FromString("3-5"),
+				}
+				return set
+			},
+			getPods: func(set *appsv1beta1.StatefulSet) []*v1.Pod {
+				setClone := set.DeepCopy()
+				setClone.Spec.Replicas = utilpointer.Int32(5)
+				startOrdinal, endOrdinal, reserveOrdinals := getStatefulSetReplicasRange(setClone)
+				pods := make([]*v1.Pod, 0)
+				expectIndex := []int{0, 2, 6, 7, 8}
+				currentIndex := make([]int, 0)
+				for i := startOrdinal; i < endOrdinal; i++ {
+					if reserveOrdinals.Has(i) {
+						continue
+					}
+					currentIndex = append(currentIndex, i)
+					pods = append(pods, newStatefulSetPod(set, i))
+				}
+				if !reflect.DeepEqual(expectIndex, currentIndex) {
+					t.Fatalf("expect(%v), but get(%v)", expectIndex, currentIndex)
+				}
+				return pods
+			},
+			expectPvcOwnerRef: func(pvcName string) metav1.OwnerReference {
+				sIndex1 := strings.Index(pvcName, "-") + 1
+				podName := pvcName[sIndex1:]
+				sIndex2 := strings.LastIndex(pvcName, "-") + 1
+				index, _ := strconv.Atoi(pvcName[sIndex2:])
+				if index < 9 {
+					return metav1.OwnerReference{
+						APIVersion: "apps.kruise.io/v1beta1",
+						Kind:       "StatefulSet",
+						Name:       "foo",
+					}
+				}
+				return metav1.OwnerReference{
+					APIVersion: "v1",
+					Kind:       "Pod",
+					Name:       podName,
+				}
+			},
+		},
 	}
 
 	for _, cs := range cases {

diff --git a/pkg/controller/statefulset/stateful_set_control_test.go b/pkg/controller/statefulset/stateful_set_control_test.go
@@ -1,6 +1,5 @@
 /*
 Copyright 2019 The Kruise Authors.
-Copyright 2016 The Kubernetes Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -4651,7 +4650,8 @@ func TestScaleUpWithMaxUnavailable(t *testing.T) {
 }
 
 func isOrHasInternalError(err error) bool {
-	agg, ok := err.(utilerrors.Aggregate)
+	var agg utilerrors.Aggregate
+	ok := errors.As(err, &agg)
 	return !ok && !apierrors.IsInternalError(err) || ok && len(agg.Errors()) > 0 && !apierrors.IsInternalError(agg.Errors()[0])
 }
 
@@ -4662,10 +4662,12 @@ func emptyInvariants(set *appsv1beta1.StatefulSet, om *fakeObjectManager) error
 func TestStatefulSetControlWithStartOrdinal(t *testing.T) {
 	defer utilfeature.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.StatefulSetStartOrdinal, true)()
 
-	simpleSetFn := func(replicas, startOrdinal int, reservedIds ...int) *appsv1beta1.StatefulSet {
+	simpleSetFn := func(replicas, startOrdinal int, reservedIds ...int32) *appsv1beta1.StatefulSet {
 		statefulSet := newStatefulSet(replicas)
 		statefulSet.Spec.Ordinals = &appsv1beta1.StatefulSetOrdinals{Start: int32(startOrdinal)}
-		statefulSet.Spec.ReserveOrdinals = append([]int{}, reservedIds...)
+		for _, id := range reservedIds {
+			statefulSet.Spec.ReserveOrdinals = append(statefulSet.Spec.ReserveOrdinals, intstr.FromInt32(id))
+		}
 		return statefulSet
 	}
 

diff --git a/pkg/controller/statefulset/stateful_set_status_updater_test.go b/pkg/controller/statefulset/stateful_set_status_updater_test.go
@@ -1,6 +1,5 @@
 /*
 Copyright 2019 The Kruise Authors.
-Copyright 2017 The Kubernetes Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.

diff --git a/pkg/controller/statefulset/stateful_set_utils.go b/pkg/controller/statefulset/stateful_set_utils.go
@@ -26,6 +26,7 @@ import (
 	"strconv"
 	"time"
 
+	apiutil "github.com/openkruise/kruise/pkg/util/api"
 	apps "k8s.io/api/apps/v1"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -92,7 +93,7 @@ func podInOrdinalRange(pod *v1.Pod, set *appsv1beta1.StatefulSet) bool {
 	return podInOrdinalRangeWithParams(pod, startOrdinal, endOrdinal, reserveOrdinals)
 }
 
-func podInOrdinalRangeWithParams(pod *v1.Pod, startOrdinal, endOrdinal int, reserveOrdinals sets.Int) bool {
+func podInOrdinalRangeWithParams(pod *v1.Pod, startOrdinal, endOrdinal int, reserveOrdinals sets.Set[int]) bool {
 	ordinal := getOrdinal(pod)
 	return ordinal >= startOrdinal && ordinal < endOrdinal &&
 		!reserveOrdinals.Has(ordinal)
@@ -791,8 +792,8 @@ func decreaseAndCheckMaxUnavailable(maxUnavailable *int) bool {
 // result is startOrdinal 2(inclusive), endOrdinal 7(exclusive), reserveOrdinals = {1, 3}
 // replicas[endOrdinal - startOrdinal] stores [replica-2, nil(reserveOrdinal 3), replica-4, replica-5, replica-6]
 // todo: maybe we should remove ineffective reserveOrdinals in webhook, reserveOrdinals = {3}
-func getStatefulSetReplicasRange(set *appsv1beta1.StatefulSet) (int, int, sets.Int) {
-	reserveOrdinals := sets.NewInt(set.Spec.ReserveOrdinals...)
+func getStatefulSetReplicasRange(set *appsv1beta1.StatefulSet) (int, int, sets.Set[int]) {
+	reserveOrdinals := apiutil.GetReserveOrdinalIntSet(set.Spec.ReserveOrdinals)
 	replicaMaxOrdinal := getStartOrdinal(set)
 	for realReplicaCount := 0; realReplicaCount < int(*set.Spec.Replicas); replicaMaxOrdinal++ {
 		if reserveOrdinals.Has(replicaMaxOrdinal) {