feat: add pod index label

Signed-off-by: Se7en <[email protected]>
openkruise · Jul 21, 2024 · 1aa682f · 1aa682f
1 parent bfb70a1
commit 1aa682f
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 1 deletion.
diff --git a/pkg/controller/statefulset/stateful_set_utils.go b/pkg/controller/statefulset/stateful_set_utils.go
@@ -40,6 +40,8 @@ import (
 
 	appspub "github.com/openkruise/kruise/apis/apps/pub"
 	appsv1beta1 "github.com/openkruise/kruise/apis/apps/v1beta1"
+	"github.com/openkruise/kruise/pkg/features"
+	utilfeature "github.com/openkruise/kruise/pkg/util/feature"
 	"github.com/openkruise/kruise/pkg/util/lifecycle"
 	"github.com/openkruise/kruise/pkg/util/revision"
 )
@@ -371,12 +373,16 @@ func initIdentity(set *appsv1beta1.StatefulSet, pod *v1.Pod) {
 // updateIdentity updates pod's name, hostname, and subdomain, and StatefulSetPodNameLabel to conform to set's name
 // and headless service.
 func updateIdentity(set *appsv1beta1.StatefulSet, pod *v1.Pod) {
-	pod.Name = getPodName(set, getOrdinal(pod))
+	ordinal := getOrdinal(pod)
+	pod.Name = getPodName(set, ordinal)
 	pod.Namespace = set.Namespace
 	if pod.Labels == nil {
 		pod.Labels = make(map[string]string)
 	}
 	pod.Labels[apps.StatefulSetPodNameLabel] = pod.Name
+	if utilfeature.DefaultFeatureGate.Enabled(features.PodIndexLabel) {
+		pod.Labels[apps.PodIndexLabel] = strconv.Itoa(ordinal)
+	}
 }
 
 // isRunningAndAvailable returns true if pod is in the PodRunning Phase,

diff --git a/pkg/controller/statefulset/stateful_set_utils_test.go b/pkg/controller/statefulset/stateful_set_utils_test.go
@@ -206,6 +206,23 @@ func TestUpdateIdentity(t *testing.T) {
 	}
 }
 
+func TestUpdateIdentityWithPodIndexLabel(t *testing.T) {
+	defer utilfeature.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.PodIndexLabel, true)()
+
+	set := newStatefulSet(3)
+	pod := newStatefulSetPod(set, 1)
+	updateIdentity(set, pod)
+
+	podIndexFromLabel, exists := pod.Labels[apps.PodIndexLabel]
+	if !exists {
+		t.Errorf("Missing pod index label: %s", apps.PodIndexLabel)
+	}
+	podIndexFromName := strconv.Itoa(getOrdinal(pod))
+	if podIndexFromLabel != podIndexFromName {
+		t.Errorf("Pod index label value (%s) does not match pod index in pod name (%s)", podIndexFromLabel, podIndexFromName)
+	}
+}
+
 func TestUpdateStorage(t *testing.T) {
 	set := newStatefulSet(3)
 	pod := newStatefulSetPod(set, 1)

diff --git a/pkg/features/kruise_features.go b/pkg/features/kruise_features.go
@@ -122,6 +122,9 @@ const (
 
 	// Enables a StatefulSet to start from an arbitrary non zero ordinal
 	StatefulSetStartOrdinal featuregate.Feature = "StatefulSetStartOrdinal"
+
+	// Set pod completion index as a pod label for Indexed Jobs.
+	PodIndexLabel featuregate.Feature = "PodIndexLabel"
 )
 
 var defaultFeatureGates = map[featuregate.Feature]featuregate.FeatureSpec{
@@ -154,6 +157,7 @@ var defaultFeatureGates = map[featuregate.Feature]featuregate.FeatureSpec{
 	EnhancedLivenessProbeGate:              {Default: false, PreRelease: featuregate.Alpha},
 	RecreatePodWhenChangeVCTInCloneSetGate: {Default: false, PreRelease: featuregate.Alpha},
 	StatefulSetStartOrdinal:                {Default: false, PreRelease: featuregate.Alpha},
+	PodIndexLabel:                          {Default: true, PreRelease: featuregate.Beta},
 }
 
 func init() {