Use WorkloadRequestUseMergePatch on UpdateStatus(). (#8029)

* Use WorkloadRequestUseMergePatch on UpdateStatus().

* Add unit tests.

Author: mbobrovskyi

pkg/controller/jobframework/reconciler.go

@@ -516,8 +516,7 @@ func (r *JobReconciler) ReconcileGenericJob(ctx context.Context, req ctrl.Reques
 		condition := generatePodsReadyCondition(ctx, job, wl, r.clock)
 		if !workload.HasConditionWithTypeAndReason(wl, &condition) {
 			log.V(3).Info("Updating the PodsReady condition", "reason", condition.Reason, "status", condition.Status)
-			apimeta.SetStatusCondition(&wl.Status.Conditions, condition)
-			err := workload.UpdateStatus(ctx, r.client, wl, condition.Type, condition.Status, condition.Reason, condition.Message, constants.JobControllerName, r.clock)
+			err := workload.SetConditionAndUpdate(ctx, r.client, wl, condition.Type, condition.Status, condition.Reason, condition.Message, constants.JobControllerName, r.clock)
 			if err != nil {
 				log.Error(err, "Updating workload status")
 				return ctrl.Result{}, client.IgnoreNotFound(err)

pkg/workload/workload.go

@@ -561,29 +561,60 @@ func totalRequestsFromAdmission(wl *kueue.Workload) []PodSetResources {
 	return res
 }
 
-// UpdateStatus updates the condition of a workload with ssa,
-// fieldManager being set to managerPrefix + "-" + conditionType
-func UpdateStatus(ctx context.Context,
+// SetConditionAndUpdate sets (or replaces) a single condition in a Workload's status.
+//
+// Behaviour depends on the feature gate WorkloadRequestUseMergePatch:
+//
+//   - Enabled → uses merge-patch via PatchStatus (preserves other conditions,
+//     safe for concurrent controllers).
+//
+//   - Disabled → uses server-side apply with field manager
+//     "<managerPrefix>-<conditionType>" (legacy path; only the written condition
+//     is managed by this controller).
+//
+// The condition gets:
+//   - ObservedGeneration = wl.Generation
+//   - LastTransitionTime = clock.Now()
+//   - Message truncated to the allowed length
+func SetConditionAndUpdate(ctx context.Context,
 	c client.Client,
 	wl *kueue.Workload,
 	conditionType string,
 	conditionStatus metav1.ConditionStatus,
 	reason, message string,
 	managerPrefix string,
-	clock clock.Clock) error {
-	now := metav1.NewTime(clock.Now())
+	clock clock.Clock,
+) error {
 	condition := metav1.Condition{
 		Type:               conditionType,
 		Status:             conditionStatus,
-		LastTransitionTime: now,
+		ObservedGeneration: wl.Generation,
+		LastTransitionTime: metav1.NewTime(clock.Now()),
 		Reason:             reason,
 		Message:            api.TruncateConditionMessage(message),
-		ObservedGeneration: wl.Generation,
 	}
 
-	newWl := BaseSSAWorkload(wl, false)
-	newWl.Status.Conditions = []metav1.Condition{condition}
-	return c.Status().Patch(ctx, newWl, client.Apply, client.FieldOwner(managerPrefix+"-"+condition.Type))
+	var (
+		wlCopy *kueue.Workload
+		err    error
+	)
+
+	if features.Enabled(features.WorkloadRequestUseMergePatch) {
+		wlCopy = wl.DeepCopy()
+		err = clientutil.PatchStatus(ctx, c, wlCopy, func() (bool, error) {
+			return apimeta.SetStatusCondition(&wlCopy.Status.Conditions, condition), nil
+		})
+	} else {
+		wlCopy = BaseSSAWorkload(wl, true)
+		wlCopy.Status.Conditions = []metav1.Condition{condition}
+		err = c.Status().Patch(ctx, wlCopy, client.Apply, client.FieldOwner(managerPrefix+"-"+condition.Type))
+	}
+	if err != nil {
+		return err
+	}
+
+	wlCopy.DeepCopyInto(wl)
+	return nil
 }
 
 // UnsetQuotaReservationWithCondition sets the QuotaReserved condition to false, clears
@@ -1223,7 +1254,7 @@ func CreatePodsReadyCondition(status metav1.ConditionStatus, reason, message str
 		Reason:             reason,
 		Message:            message,
 		LastTransitionTime: metav1.NewTime(clock.Now()),
-		// ObservedGeneration is added via workload.UpdateStatus
+		// ObservedGeneration is added via workload.SetConditionAndUpdate
 	}
 }
 

pkg/workload/workload_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package workload
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"testing"
@@ -32,6 +33,7 @@ import (
 	testingclock "k8s.io/utils/clock/testing"
 	"k8s.io/utils/ptr"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/client/interceptor"
 
 	config "sigs.k8s.io/kueue/apis/config/v1beta2"
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta2"
@@ -44,6 +46,10 @@ import (
 	utiltestingapi "sigs.k8s.io/kueue/pkg/util/testing/v1beta2"
 )
 
+var (
+	errTest = errors.New("test error")
+)
+
 func TestNewInfo(t *testing.T) {
 	cases := map[string]struct {
 		workload     kueue.Workload
@@ -450,7 +456,7 @@ func TestNewInfo(t *testing.T) {
 	}
 }
 
-func TestUpdateWorkloadStatus(t *testing.T) {
+func TestSetConditionAndUpdate(t *testing.T) {
 	now := time.Now()
 	fakeClock := testingclock.NewFakeClock(now)
 	cases := map[string]struct {
@@ -459,7 +465,9 @@ func TestUpdateWorkloadStatus(t *testing.T) {
 		condStatus metav1.ConditionStatus
 		reason     string
 		message    string
+		err        error
 		wantStatus kueue.WorkloadStatus
+		wantErr    error
 	}{
 		"initial empty": {
 			condType:   kueue.WorkloadQuotaReserved,
@@ -478,6 +486,15 @@ func TestUpdateWorkloadStatus(t *testing.T) {
 				},
 			},
 		},
+		"initial empty with error": {
+			condType:   kueue.WorkloadQuotaReserved,
+			condStatus: metav1.ConditionFalse,
+			reason:     "Pending",
+			message:    "didn't fit",
+			err:        errTest,
+			wantStatus: kueue.WorkloadStatus{},
+			wantErr:    errTest,
+		},
 		"same condition type": {
 			oldStatus: kueue.WorkloadStatus{
 				Conditions: []metav1.Condition{
@@ -505,27 +522,47 @@ func TestUpdateWorkloadStatus(t *testing.T) {
 		},
 	}
 	for name, tc := range cases {
-		t.Run(name, func(t *testing.T) {
-			ctx, _ := utiltesting.ContextWithLog(t)
-			workload := utiltestingapi.MakeWorkload("foo", "bar").Generation(1).Obj()
-			workload.Status = tc.oldStatus
-			cl := utiltesting.NewFakeClientSSAAsSM(workload)
-			err := UpdateStatus(ctx, cl, workload, tc.condType, tc.condStatus, tc.reason, tc.message, "manager-prefix", fakeClock)
-			if err != nil {
-				t.Fatalf("Failed updating status: %v", err)
-			}
-			var updatedWl kueue.Workload
-			if err := cl.Get(ctx, client.ObjectKeyFromObject(workload), &updatedWl); err != nil {
-				t.Fatalf("Failed obtaining updated object: %v", err)
-			}
-			if diff := cmp.Diff(
-				tc.wantStatus,
-				updatedWl.Status,
-				cmpopts.IgnoreFields(metav1.Condition{}, "LastTransitionTime"),
-			); diff != "" {
-				t.Errorf("Unexpected status after updating (-want,+got):\n%s", diff)
-			}
-		})
+		for _, useMergePatch := range []bool{false, true} {
+			t.Run(name, func(t *testing.T) {
+				features.SetFeatureGateDuringTest(t, features.WorkloadRequestUseMergePatch, useMergePatch)
+
+				ctx, _ := utiltesting.ContextWithLog(t)
+
+				workload := utiltestingapi.MakeWorkload("foo", "bar").Generation(1).Obj()
+				workload.Status = tc.oldStatus
+
+				cl := utiltesting.NewClientBuilder().
+					WithObjects(workload).
+					WithStatusSubresource(&kueue.Workload{}).
+					WithInterceptorFuncs(interceptor.Funcs{
+						SubResourcePatch: func(ctx context.Context, c client.Client, subResourceName string, obj client.Object, patch client.Patch, opts ...client.SubResourcePatchOption) error {
+							if tc.err != nil {
+								return tc.err
+							}
+							return utiltesting.TreatSSAAsStrategicMerge(ctx, c, subResourceName, obj, patch, opts...)
+						},
+					}).
+					Build()
+
+				err := SetConditionAndUpdate(ctx, cl, workload, tc.condType, tc.condStatus, tc.reason, tc.message, "manager-prefix", fakeClock)
+				if diff := cmp.Diff(tc.wantErr, err, cmpopts.EquateErrors()); diff != "" {
+					t.Errorf("Unexpected error (-want,+got):\n%s", diff)
+				}
+
+				var updatedWl kueue.Workload
+				if err := cl.Get(ctx, client.ObjectKeyFromObject(workload), &updatedWl); err != nil {
+					t.Fatalf("Failed obtaining updated object: %v", err)
+				}
+
+				if diff := cmp.Diff(
+					tc.wantStatus,
+					updatedWl.Status,
+					cmpopts.IgnoreFields(metav1.Condition{}, "LastTransitionTime"),
+				); diff != "" {
+					t.Errorf("Unexpected status after updating (-want,+got):\n%s", diff)
+				}
+			})
+		}
 	}
 }
 

test/integration/singlecluster/controller/jobs/appwrapper/appwrapper_controller_test.go

@@ -283,7 +283,7 @@ var _ = ginkgo.Describe("AppWrapper controller", ginkgo.Ordered, ginkgo.Continue
 			ginkgo.By("preempt the workload", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sClient.Get(ctx, wlLookupKey, createdWorkload)).To(gomega.Succeed())
-					g.Expect(workload.UpdateStatus(ctx, k8sClient, createdWorkload, kueue.WorkloadEvicted, metav1.ConditionTrue, kueue.WorkloadEvictedByPreemption, "By test", "evict", util.RealClock)).To(gomega.Succeed())
+					g.Expect(workload.SetConditionAndUpdate(ctx, k8sClient, createdWorkload, kueue.WorkloadEvicted, metav1.ConditionTrue, kueue.WorkloadEvictedByPreemption, "By test", "evict", util.RealClock)).To(gomega.Succeed())
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
 			})
 

test/integration/singlecluster/controller/jobs/job/job_controller_test.go

@@ -630,7 +630,7 @@ var _ = ginkgo.Describe("Job controller", ginkgo.Ordered, ginkgo.ContinueOnFailu
 		ginkgo.By("preempt the workload", func() {
 			gomega.Eventually(func(g gomega.Gomega) {
 				g.Expect(k8sClient.Get(ctx, wlLookupKey, wl)).To(gomega.Succeed())
-				g.Expect(workload.UpdateStatus(
+				g.Expect(workload.SetConditionAndUpdate(
 					ctx,
 					k8sClient,
 					wl,

test/integration/singlecluster/controller/jobs/jobset/jobset_controller_test.go

@@ -352,7 +352,7 @@ var _ = ginkgo.Describe("JobSet controller", ginkgo.Ordered, ginkgo.ContinueOnFa
 			ginkgo.By("preempt the workload", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sClient.Get(ctx, wlLookupKey, createdWorkload)).To(gomega.Succeed())
-					g.Expect(workload.UpdateStatus(ctx, k8sClient, createdWorkload, kueue.WorkloadEvicted, metav1.ConditionTrue, kueue.WorkloadEvictedByPreemption, "By test", "evict", util.RealClock)).To(gomega.Succeed())
+					g.Expect(workload.SetConditionAndUpdate(ctx, k8sClient, createdWorkload, kueue.WorkloadEvicted, metav1.ConditionTrue, kueue.WorkloadEvictedByPreemption, "By test", "evict", util.RealClock)).To(gomega.Succeed())
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
 			})
 

test/integration/singlecluster/controller/jobs/pod/pod_controller_test.go

@@ -329,7 +329,7 @@ var _ = ginkgo.Describe("Pod controller", ginkgo.Ordered, ginkgo.ContinueOnFailu
 					ginkgo.By("checking that pod is stopped when workload is evicted")
 
 					gomega.Expect(
-						workload.UpdateStatus(ctx, k8sClient, createdWorkload, kueue.WorkloadEvicted, metav1.ConditionTrue,
+						workload.SetConditionAndUpdate(ctx, k8sClient, createdWorkload, kueue.WorkloadEvicted, metav1.ConditionTrue,
 							kueue.WorkloadEvictedByPreemption, "By test", "evict", util.RealClock),
 					).Should(gomega.Succeed())
 					util.FinishEvictionForWorkloads(ctx, k8sClient, createdWorkload)

test/integration/singlecluster/controller/jobs/trainjob/trainjob_controller_test.go

@@ -259,7 +259,7 @@ var _ = ginkgo.Describe("Trainjob controller", ginkgo.Ordered, ginkgo.ContinueOn
 			ginkgo.By("preempt the workload", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sClient.Get(ctx, wlLookupKey, createdWorkload)).To(gomega.Succeed())
-					g.Expect(workload.UpdateStatus(ctx, k8sClient, createdWorkload, kueue.WorkloadEvicted, metav1.ConditionTrue, kueue.WorkloadEvictedByPreemption, "By test", "evict", util.RealClock)).To(gomega.Succeed())
+					g.Expect(workload.SetConditionAndUpdate(ctx, k8sClient, createdWorkload, kueue.WorkloadEvicted, metav1.ConditionTrue, kueue.WorkloadEvictedByPreemption, "By test", "evict", util.RealClock)).To(gomega.Succeed())
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
 			})
 

test/performance/scheduler/runner/controller/controller.go

@@ -143,7 +143,7 @@ func (r *reconciler) Reconcile(ctx context.Context, req reconcile.Request) (reco
 		if remaining > 0 {
 			return reconcile.Result{RequeueAfter: remaining}, nil
 		} else {
-			err := workload.UpdateStatus(ctx, r.client, &wl, kueue.WorkloadFinished, metav1.ConditionTrue, "ByTest", "By test runner", constants.JobControllerName, r.clock)
+			err := workload.SetConditionAndUpdate(ctx, r.client, &wl, kueue.WorkloadFinished, metav1.ConditionTrue, "ByTest", "By test runner", constants.JobControllerName, r.clock)
 			if err == nil {
 				log.V(5).Info("Finish Workload")
 			}