Skip to content

feature(rbg): support role instance index env #194

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cheyang merged 2 commits into from
Mar 13, 2026
Merged

feature(rbg): support role instance index env #194

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
848e6ec
feature(rbg): support role instance index env
Mar 12, 2026
fc3f888
fix comments
Mar 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion internal/controller/workloads/rolebasedgroupscalingadapter_controller.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -427,7 +427,7 @@ func (r *RoleBasedGroupScalingAdapterReconciler) extractLabelSelectorDefault(
}

if kind == "RoleInstanceSet" && role.IsLeaderWorkerPattern() {
selectorStr += fmt.Sprintf(",%s=0", constants.RBGComponentIndexLabelKey)
selectorStr += fmt.Sprintf(",%s=0", constants.ComponentIndexLabelKey)
}

return selectorStr, nil
Expand Down
30 changes: 3 additions & 27 deletions pkg/constants/label.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,9 @@ const (
// GroupRevisionLabelKey is the labels key used to store the revision hash of the
// RoleBasedGroup Roles's template.
GroupRevisionLabelKey = RBGPrefix + "group-revision"

// GroupUniqueHashLabelKey is used for pod affinity rules in exclusive topology
GroupUniqueHashLabelKey = RBGPrefix + "group-unique-hash"
)

// Role level labels
Expand Down Expand Up @@ -90,30 +93,3 @@ const (
// - worker's ComponentIndexLabelKey = ComponentIDKey.value + 1
ComponentIndexLabelKey = RBGPrefix + "component-index"
)

// Unique Hash label
const (
// GroupUniqueHashLabelKey is used for pod affinity rules in exclusive topology
GroupUniqueHashLabelKey = RBGPrefix + "group-unique-hash"
)

// RoleInstance component labels (aliases for compatibility)
const (
// RoleInstanceComponentNameKey is the name of the component within a RoleInstance
RoleInstanceComponentNameKey = RBGPrefix + "role-instance-component-name"

// RoleInstanceComponentIDKey is the id of the component within a RoleInstance
RoleInstanceComponentIDKey = RBGPrefix + "role-instance-component-id"
)

// LWS (LeaderWorkerSet) related labels
const (
// RBGComponentIndexLabelKey identifies the component instance index
// for LWS workloads this maps to:
// - leader's RBGComponentIndex = 0
// - worker's RBGComponentIndex = ComponentIDKey.value + 1
RBGComponentIndexLabelKey = RBGPrefix + "component-index"

// RBGComponentSizeLabelKey identifies the component size
RBGComponentSizeLabelKey = RBGPrefix + "component-size"
)
10 changes: 5 additions & 5 deletions pkg/discovery/env_builder.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,15 +41,15 @@ func (b *EnvBuilder) BuildLwsEnv(svcName string) []corev1.EnvVar {
Name: constants.EnvRBGIndex,
ValueFrom: &corev1.EnvVarSource{
FieldRef: &corev1.ObjectFieldSelector{
FieldPath: fmt.Sprintf("metadata.labels['%s']", constants.RBGComponentIndexLabelKey),
FieldPath: fmt.Sprintf("metadata.labels['%s']", constants.ComponentIndexLabelKey),
},
},
},
{
Name: constants.EnvRBGSize,
ValueFrom: &corev1.EnvVarSource{
FieldRef: &corev1.ObjectFieldSelector{
FieldPath: fmt.Sprintf("metadata.labels['%s']", constants.RBGComponentSizeLabelKey),
FieldPath: fmt.Sprintf("metadata.labels['%s']", constants.ComponentSizeLabelKey),
},
},
},
Expand All @@ -72,7 +72,7 @@ func (g *EnvBuilder) buildLocalRoleVars() []corev1.EnvVar {
},
}

if g.role.Workload.Kind == "StatefulSet" || g.role.Workload.Kind == "LeaderWorkerSet" {
if workloadsv1alpha2.IsStatefulRole(g.role) {
envVars = append(envVars,
Comment on lines +75 to 76
Copy link

Copilot AI Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

EnvBuilder.buildLocalRoleVars now uses workloadsv1alpha2.IsStatefulRole(g.role), which changes behavior for RoleInstanceSet roles (e.g., injecting RBG_ROLE_INDEX for stateful RoleInstanceSet patterns). There are existing tests for EnvBuilder in this package, but they don’t cover RoleInstanceSet stateful vs stateless patterns, so this behavior change is currently untested. Please add/extend unit tests in env_builder_test.go to cover RoleInstanceSet roles (stateful + stateless) and assert whether RBG_ROLE_INDEX is included as expected.

Copilot uses AI. Check for mistakes.
corev1.EnvVar{
Name: constants.EnvRBGRoleIndex,
Expand All @@ -98,15 +98,15 @@ func (g *EnvBuilder) buildLocalRoleVars() []corev1.EnvVar {
Name: constants.EnvRBGComponentName,
ValueFrom: &corev1.EnvVarSource{
FieldRef: &corev1.ObjectFieldSelector{
FieldPath: fmt.Sprintf("metadata.labels['%s']", constants.RoleInstanceComponentNameKey),
FieldPath: fmt.Sprintf("metadata.labels['%s']", constants.ComponentNameLabelKey),
},
},
},
corev1.EnvVar{
Name: constants.EnvRBGComponentIndex,
ValueFrom: &corev1.EnvVarSource{
FieldRef: &corev1.ObjectFieldSelector{
FieldPath: fmt.Sprintf("metadata.labels['%s']", constants.RoleInstanceComponentIDKey),
FieldPath: fmt.Sprintf("metadata.labels['%s']", constants.ComponentIDLabelKey),
},
},
},
Expand Down
2 changes: 1 addition & 1 deletion pkg/inplace/instance/instance_util.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ func MatchComponent(pod *corev1.Pod, componentName string) bool {
if pod.Labels == nil {
return false
}
return pod.Labels[constants.RoleInstanceComponentNameKey] == componentName
return pod.Labels[constants.ComponentNameLabelKey] == componentName
}

// IsRoleInstanceReady checks whether the role instance is ready
Expand Down
4 changes: 2 additions & 2 deletions pkg/inplace/instance/instance_util_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,15 +74,15 @@ func TestIsMatchComponent(t *testing.T) {
{
name: "Pod matches component",
podLabels: map[string]string{
constants.RoleInstanceComponentNameKey: "web",
constants.ComponentNameLabelKey: "web",
},
componentName: "web",
expected: true,
},
{
name: "Pod does not match component",
podLabels: map[string]string{
constants.RoleInstanceComponentNameKey: "db",
constants.ComponentNameLabelKey: "db",
},
componentName: "web",
expected: false,
Expand Down
12 changes: 6 additions & 6 deletions pkg/reconciler/roleinstance/utils/instance_utils.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,13 +75,13 @@ func FormatComponentPodName(instanceName, componentName string, id int32, roleTe

func InitComponentPodLabels(instanceName, componentName string, id int32, roleTemplateType constants.RoleTemplateType) map[string]string {
l := GetSelectorMatchLabels(instanceName)
l[constants.RoleInstanceComponentNameKey] = componentName
l[constants.RoleInstanceComponentIDKey] = fmt.Sprintf("%d", id)
l[constants.ComponentNameLabelKey] = componentName
l[constants.ComponentIDLabelKey] = fmt.Sprintf("%d", id)
if roleTemplateType == constants.LeaderWorkerSetTemplateType {
l[constants.RBGComponentIndexLabelKey] = fmt.Sprintf("%d", id)
l[constants.ComponentIndexLabelKey] = fmt.Sprintf("%d", id)
// when roleTemplateType is LWS, component name will be controlled by RBG-controller
if componentName == "worker" {
l[constants.RBGComponentIndexLabelKey] = fmt.Sprintf("%d", id+1)
l[constants.ComponentIndexLabelKey] = fmt.Sprintf("%d", id+1)
}
}

Expand Down Expand Up @@ -142,7 +142,7 @@ func NextRevision(revisions []*apps.ControllerRevision) int64 {
}

func GetPodComponentName(pod *v1.Pod) string {
componentName := pod.Labels[constants.RoleInstanceComponentNameKey]
componentName := pod.Labels[constants.ComponentNameLabelKey]
if len(componentName) != 0 {
return componentName
}
Expand All @@ -154,7 +154,7 @@ func GetPodComponentName(pod *v1.Pod) string {
}

func GetPodComponentID(pod *v1.Pod) int32 {
componentId := pod.Labels[constants.RoleInstanceComponentIDKey]
componentId := pod.Labels[constants.ComponentIDLabelKey]
if len(componentId) != 0 {
id, _ := strconv.Atoi(componentId)
return int32(id)
Expand Down
37 changes: 33 additions & 4 deletions pkg/reconciler/roleinstanceset/statefulmode/stateful_instance_set_utils.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ import (
intstrutil "k8s.io/apimachinery/pkg/util/intstr"
"k8s.io/apimachinery/pkg/util/sets"
workloadsv1alpha2 "sigs.k8s.io/rbgs/api/workloads/v1alpha2"
"sigs.k8s.io/rbgs/pkg/constants"
)

// statefulInstanceRegex is a regular expression that extracts the parent InstanceSet and ordinal from the Name of an Instance
Expand Down Expand Up @@ -96,8 +97,28 @@ func updateIdentity(set *workloadsv1alpha2.RoleInstanceSet, instance *workloadsv
if instance.Labels == nil {
instance.Labels = make(map[string]string)
}
instance.Labels[apps.StatefulSetPodNameLabel] = instance.Name
instance.Labels[apps.PodIndexLabel] = strconv.Itoa(ordinal)
identityLabels := map[string]string{
apps.StatefulSetPodNameLabel: instance.Name,
constants.RoleInstanceIndexLabelKey: strconv.Itoa(ordinal),
}
for k, v := range identityLabels {
instance.Labels[k] = v
}
injectIdentityLabelsIntoComponents(instance, identityLabels)
Copy link

Copilot AI Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When propagating identityLabels into component pod templates, apps.StatefulSetPodNameLabel ("statefulset.kubernetes.io/pod-name") will be set to the RoleInstance name for all pods created from the template. That label’s expected semantic is the pod name (and cannot be correctly set at template level for multi-pod components), so this will produce incorrect/misleading labels on the pods. Consider excluding apps.StatefulSetPodNameLabel from template injection and only propagate labels that are valid at the RoleInstance level (e.g., apps.PodIndexLabel / RoleInstanceIndexLabelKey), or set a dedicated RoleInstance name label instead.

Suggested change
injectIdentityLabelsIntoComponents(instance, identityLabels)
// Only inject identity labels that are valid at the pod-template level.
templateIdentityLabels := make(map[string]string, len(identityLabels))
for k, v := range identityLabels {
if k == apps.StatefulSetPodNameLabel {
// apps.StatefulSetPodNameLabel is per-pod and must not be set on pod templates.
continue
}
templateIdentityLabels[k] = v
}
injectIdentityLabelsIntoComponents(instance, templateIdentityLabels)

Copilot uses AI. Check for mistakes.
}

// injectIdentityLabelsIntoComponents writes the given identity labels into each
// component's pod template metadata, so that pods created from the template
// also carry the instance's ordinal identity labels.
func injectIdentityLabelsIntoComponents(instance *workloadsv1alpha2.RoleInstance, identityLabels map[string]string) {
for i := range instance.Spec.Components {
if instance.Spec.Components[i].Template.Labels == nil {
instance.Spec.Components[i].Template.Labels = make(map[string]string)
}
for k, v := range identityLabels {
instance.Spec.Components[i].Template.Labels[k] = v
}
}
}

// getInstanceSetReplicasRange returns the start ordinal, end ordinal (exclusive), and set of reserved ordinals for a InstanceSet.
Expand Down Expand Up @@ -181,15 +202,23 @@ func newVersionedInstance(
instance.Labels[k] = v
}
}
instance.Labels[apps.StatefulSetPodNameLabel] = instance.Name

// Set revision label
if useUpdateRevision {
instance.Labels[apps.ControllerRevisionHashLabelKey] = updateRevision
} else {
instance.Labels[apps.ControllerRevisionHashLabelKey] = currentRevision
}

// Set identity labels and propagate them into each component's pod template.
identityLabels := map[string]string{
apps.StatefulSetPodNameLabel: instance.Name,
constants.RoleInstanceIndexLabelKey: strconv.Itoa(ordinal),
}
for k, v := range identityLabels {
instance.Labels[k] = v
}
injectIdentityLabelsIntoComponents(instance, identityLabels)
Comment on lines +213 to +220
Copy link

@gemini-code-assist gemini-code-assist bot Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This block of code for setting and injecting identity labels is duplicated from the updateIdentity function (lines 100-108). To improve maintainability and reduce redundancy, consider extracting this logic into a new helper function.

For example, you could create a function setAndInjectIdentityLabels(instance *workloadsv1alpha2.RoleInstance, ordinal int) that encapsulates this logic, and then call it from both updateIdentity and newVersionedInstance.

func setAndInjectIdentityLabels(instance *workloadsv1alpha2.RoleInstance, ordinal int) {
	if instance.Labels == nil {
		instance.Labels = make(map[string]string)
	}
	identityLabels := map[string]string{
		apps.StatefulSetPodNameLabel:        instance.Name,
		apps.PodIndexLabel:                  strconv.Itoa(ordinal),
		constants.RoleInstanceIndexLabelKey: strconv.Itoa(ordinal),
	}
	for k, v := range identityLabels {
		instance.Labels[k] = v
	}
	injectIdentityLabelsIntoComponents(instance, identityLabels)
}


Comment on lines +212 to +221
Copy link

Copilot AI Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same issue as above: injecting apps.StatefulSetPodNameLabel into component pod templates will set "statefulset.kubernetes.io/pod-name" to the RoleInstance name for every pod created from the template. Since pod names differ per component replica, this label will be incorrect on the resulting pods. Consider not including apps.StatefulSetPodNameLabel in the set of labels propagated into component templates.

Copilot uses AI. Check for mistakes.
// Set owner reference
instance.OwnerReferences = []metav1.OwnerReference{
*metav1.NewControllerRef(currentSet, controllerKind),
Expand Down
2 changes: 1 addition & 1 deletion pkg/reconciler/roleinstanceset_reconciler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -303,7 +303,7 @@ func (r *RoleInstanceSetReconciler) constructRoleInstanceTemplateByLeaderWorkerP

workerPodReconciler := NewPodReconciler(r.scheme, r.client)
// workerTemplate do not need to inject sidecar
workerPodReconciler.SetInjectors([]string{"config", "common_env", "lws_env"})
workerPodReconciler.SetInjectors([]string{"config", "common_env", "lwp_env"})
workerTemplateApplyCfg, err := workerPodReconciler.ConstructPodTemplateSpecApplyConfiguration(
ctx, rbg, role, matchLabels, *workerTemp,
)
Expand Down
Loading