Skip to content

fix(rbg): use SSA Apply with distinct field managers on RBG writes #309

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

Open
sebest wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(rbg): use SSA Apply with distinct field managers on RBG writes #309

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
61 changes: 32 additions & 29 deletions internal/controller/workloads/discovery_config_mode_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,17 +37,16 @@ import (

func TestEnsureDiscoveryConfigMode(t *testing.T) {
type testCase struct {
name string
mutateRBG func(*workloadsv1alpha2.RoleBasedGroup)
buildExtraObjects func(*workloadsv1alpha2.RoleBasedGroup) []runtime.Object
wantRequeue bool
wantMode constants.DiscoveryConfigMode
wantModeAnnotation bool
name string
mutateRBG func(*workloadsv1alpha2.RoleBasedGroup)
buildExtraObjects func(*workloadsv1alpha2.RoleBasedGroup) []runtime.Object
wantRequeue bool
wantMode constants.DiscoveryConfigMode
}

tests := []testCase{
{
name: "missing annotation with legacy role configmap should set legacy mode without requeue",
name: "missing annotation with legacy role configmap sets legacy mode",
buildExtraObjects: func(rbg *workloadsv1alpha2.RoleBasedGroup) []runtime.Object {
return []runtime.Object{
&corev1.ConfigMap{
Expand All @@ -58,33 +57,37 @@ func TestEnsureDiscoveryConfigMode(t *testing.T) {
},
}
},
wantRequeue: false,
wantMode: constants.LegacyDiscoveryConfigMode,
wantModeAnnotation: true,
wantRequeue: false,
wantMode: constants.LegacyDiscoveryConfigMode,
},
{
name: "missing annotation without legacy signal should set refine mode without requeue",
wantRequeue: false,
wantMode: constants.RefineDiscoveryConfigMode,
wantModeAnnotation: true,
name: "missing annotation with observed generation sets legacy mode",
mutateRBG: func(rbg *workloadsv1alpha2.RoleBasedGroup) {
rbg.Status.ObservedGeneration = 1
},
wantRequeue: false,
wantMode: constants.LegacyDiscoveryConfigMode,
},
{
name: "missing annotation without legacy signal sets refine mode",
wantRequeue: false,
wantMode: constants.RefineDiscoveryConfigMode,
},
{
name: "existing legacy annotation should not requeue",
name: "existing legacy annotation is a no-op",
mutateRBG: func(rbg *workloadsv1alpha2.RoleBasedGroup) {
rbg.SetDiscoveryConfigMode(constants.LegacyDiscoveryConfigMode)
},
wantRequeue: false,
wantMode: constants.LegacyDiscoveryConfigMode,
wantModeAnnotation: true,
wantRequeue: false,
wantMode: constants.LegacyDiscoveryConfigMode,
},
{
name: "existing refine annotation should not requeue",
name: "existing refine annotation is a no-op",
mutateRBG: func(rbg *workloadsv1alpha2.RoleBasedGroup) {
rbg.SetDiscoveryConfigMode(constants.RefineDiscoveryConfigMode)
},
wantRequeue: false,
wantMode: constants.RefineDiscoveryConfigMode,
wantModeAnnotation: true,
wantRequeue: false,
wantMode: constants.RefineDiscoveryConfigMode,
},
}

Expand All @@ -111,7 +114,7 @@ func TestEnsureDiscoveryConfigMode(t *testing.T) {
current := &workloadsv1alpha2.RoleBasedGroup{}
key := types.NamespacedName{Name: rbg.Name, Namespace: rbg.Namespace}
if err := client.Get(context.Background(), key, current); err != nil {
t.Fatalf("get rbg error: %v", err)
t.Fatalf("get rbg: %v", err)
}

requeue, err := reconciler.ensureDiscoveryConfigMode(context.Background(), current)
Expand All @@ -122,16 +125,16 @@ func TestEnsureDiscoveryConfigMode(t *testing.T) {
t.Fatalf("requeue = %v, want %v", requeue, tt.wantRequeue)
}

if got := current.GetDiscoveryConfigMode(); got != tt.wantMode {
t.Fatalf("in-memory mode = %s, want %s", got, tt.wantMode)
}

persisted := &workloadsv1alpha2.RoleBasedGroup{}
if err := client.Get(context.Background(), key, persisted); err != nil {
t.Fatalf("get persisted rbg error: %v", err)
t.Fatalf("get persisted rbg: %v", err)
}
if got := persisted.GetDiscoveryConfigMode(); got != tt.wantMode {
t.Fatalf("mode = %s, want %s", got, tt.wantMode)
}
_, hasModeAnnotation := persisted.Annotations[constants.DiscoveryConfigModeAnnotationKey]
if hasModeAnnotation != tt.wantModeAnnotation {
t.Fatalf("has discovery-config-mode annotation = %v, want %v", hasModeAnnotation, tt.wantModeAnnotation)
t.Fatalf("persisted mode = %s, want %s", got, tt.wantMode)
}
})
}
Expand Down
15 changes: 12 additions & 3 deletions internal/controller/workloads/rolebasedgroup_controller.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ import (
lwsv1 "sigs.k8s.io/lws/api/leaderworkerset/v1"
"sigs.k8s.io/rbgs/api/workloads/constants"
workloadsv1alpha2 "sigs.k8s.io/rbgs/api/workloads/v1alpha2"
applyconfiguration "sigs.k8s.io/rbgs/client-go/applyconfiguration/workloads/v1alpha2"
"sigs.k8s.io/rbgs/pkg/coordination/coordinationscaling"
"sigs.k8s.io/rbgs/pkg/dependency"
"sigs.k8s.io/rbgs/pkg/discovery"
Expand Down Expand Up @@ -339,11 +340,19 @@ func (r *RoleBasedGroupReconciler) ensureDiscoveryConfigMode(
mode = constants.RefineDiscoveryConfigMode
}

old := rbg.DeepCopy()
rbg.SetDiscoveryConfigMode(mode)
if err := r.client.Patch(ctx, rbg, client.MergeFrom(old)); err != nil {
gvk := utils.GetRbgGVK()
applyCfg := applyconfiguration.RoleBasedGroup(rbg.Name, rbg.Namespace).
WithKind(gvk.Kind).
WithAPIVersion(gvk.GroupVersion().String()).
WithAnnotations(map[string]string{
constants.DiscoveryConfigModeAnnotationKey: string(mode),
})
Comment thread
sebest marked this conversation as resolved.
if err := utils.PatchObjectApplyConfigurationWithFieldManager(
ctx, r.client, applyCfg, utils.PatchSpec, utils.RBGDiscoveryFieldManager,
); err != nil {
return true, err
}
rbg.SetDiscoveryConfigMode(mode)

log.FromContext(ctx).Info("Initialized discovery config mode", "mode", mode)
// Don't requeue here - continue to reconcile ConfigMap and workloads in the same loop
Expand Down
43 changes: 19 additions & 24 deletions internal/controller/workloads/rolebasedgroupscalingadapter_controller.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,15 +24,13 @@ import (

"github.com/pkg/errors"
corev1 "k8s.io/api/core/v1"
apierrors "k8s.io/apimachinery/pkg/api/errors"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
"k8s.io/apimachinery/pkg/runtime"
"k8s.io/apimachinery/pkg/runtime/schema"
"k8s.io/apimachinery/pkg/types"
metaapplyv1 "k8s.io/client-go/applyconfigurations/meta/v1"
"k8s.io/client-go/tools/record"
"k8s.io/client-go/util/retry"
"k8s.io/klog/v2"
ctrl "sigs.k8s.io/controller-runtime"
"sigs.k8s.io/controller-runtime/pkg/builder"
Expand Down Expand Up @@ -498,28 +496,25 @@ func (r *RoleBasedGroupScalingAdapterReconciler) GetTargetRbgFromAdapter(
func (r *RoleBasedGroupScalingAdapterReconciler) updateRoleReplicas(

@cheyang cheyang Jun 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

All calls to updateRoleReplicas share the single rbg-replicas field manager. Under SSA list-map merge semantics, applying only one role releases the manager's claim on any previously-owned roles. If two RBGSAs target different roles of the same RBG, the second apply releases ownership of the first role's replicas field.

Is there a guarantee that only one RBGSA at a time updates replicas for a given RBG? If multiple RBGSAs can exist per RBG, a per-role field manager like rbg-replicas-<roleName> would preserve independent ownership.

ctx context.Context, rbg *workloadsv1alpha2.RoleBasedGroup, targetRoleName string, newReplicas *int32,
) error {
return retry.RetryOnConflict(
retry.DefaultBackoff, func() error {
for index, role := range rbg.Spec.Roles {
if role.Name == targetRoleName {
role.Replicas = newReplicas
rbg.Spec.Roles[index] = role
break
}
}
if err := r.client.Update(ctx, rbg); err != nil {
if apierrors.IsConflict(err) {
if err := r.client.Get(
ctx, types.NamespacedName{Name: rbg.Name, Namespace: rbg.Namespace}, rbg,
); err != nil {
return err
}
}
return err
}
return nil
},
)
gvk := utils.GetRbgGVK()
applyCfg := applyconfiguration.RoleBasedGroup(rbg.Name, rbg.Namespace).
WithKind(gvk.Kind).
WithAPIVersion(gvk.GroupVersion().String()).
WithSpec(
applyconfiguration.RoleBasedGroupSpec().
WithRoles(
applyconfiguration.RoleSpec().
WithName(targetRoleName).
WithReplicas(*newReplicas),

@cheyang cheyang Jun 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

The old code assigned role.Replicas = newReplicas (pointer-safe even when nil), but the new code dereferences with *newReplicas on line 508. The current caller guards against nil before reaching this point, so this is not a live bug today, but a defensive nil check (or accepting the pointer directly via WithReplicas) would prevent a panic if a future caller omits the guard.

Non-blocking -- just something to keep in mind.

),
)

if err := utils.PatchObjectApplyConfigurationWithFieldManager(
ctx, r.client, applyCfg, utils.PatchSpec, utils.RBGReplicaFieldManager,
); err != nil {
return fmt.Errorf("apply replica update for role %q: %w", targetRoleName, err)
}
return nil
}

// extractLabelSelectorDefault extracts a LabelSelector string from the given role's scale subresource.
Expand Down
39 changes: 39 additions & 0 deletions internal/controller/workloads/rolebasedgroupscalingadapter_controller_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1110,3 +1110,42 @@ func TestRBGRoleStatusPredicate(t *testing.T) {
}))
})
}

// TestUpdateRoleReplicas only covers behaviors the controller-runtime fake
// client can faithfully model: that the apply config is accepted and the
// targeted role's replicas reflect the desired value. SSA list-merge by
// name and sibling-field preservation are real API-server behaviors and
// must be verified via envtest or kind.
func TestUpdateRoleReplicas(t *testing.T) {
scheme := runtime.NewScheme()
_ = workloadsv1alpha2.AddToScheme(scheme)
_ = corev1.AddToScheme(scheme)

rbg := wrappersv2.BuildBasicRoleBasedGroup("test-rbg", "default").
WithRoles([]workloadsv1alpha2.RoleSpec{
wrappersv2.BuildStandaloneRole("worker").WithReplicas(1).Obj(),
}).Obj()

cl := fake.NewClientBuilder().WithScheme(scheme).WithRuntimeObjects(rbg).Build()
reconciler := &RoleBasedGroupScalingAdapterReconciler{
client: cl,
recorder: record.NewFakeRecorder(10),
}
ctx := ctrl.LoggerInto(context.Background(), zap.New().WithValues("env", "unit-test"))

if err := reconciler.updateRoleReplicas(ctx, rbg, "worker", ptr.To(int32(5))); err != nil {
t.Fatalf("updateRoleReplicas: %v", err)
}

updated := &workloadsv1alpha2.RoleBasedGroup{}
if err := cl.Get(ctx, types.NamespacedName{Name: rbg.Name, Namespace: rbg.Namespace}, updated); err != nil {
t.Fatalf("get updated RBG: %v", err)
}
role, err := updated.GetRole("worker")
if err != nil {
t.Fatalf("worker role missing after update: %v", err)
}
if role.Replicas == nil || *role.Replicas != 5 {
t.Errorf("worker replicas = %v, want 5", role.Replicas)
}
}
Loading
Loading