Introduce workload.PatchStatus function to commonize workload status updates

pkg/workload/workload.go

@@ -594,28 +594,9 @@ func SetConditionAndUpdate(ctx context.Context,
 		Reason:             reason,
 		Message:            api.TruncateConditionMessage(message),
 	}
-
-	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
+	return PatchStatus(ctx, c, wl, client.FieldOwner(managerPrefix+"-"+condition.Type), func(wl *kueue.Workload) (bool, error) {
+		return apimeta.SetStatusCondition(&wl.Status.Conditions, condition), nil
+	})
 }
 
 // UnsetQuotaReservationWithCondition sets the QuotaReserved condition to false, clears
@@ -965,12 +946,14 @@ func PrepareWorkloadPatch(w *kueue.Workload, strict bool, clk clock.Clock) *kueu
 	return wlCopy
 }
 
-// PatchAdmissionStatusOption defines a functional option for customizing PatchAdmissionStatusOptions.
+type UpdateFunc func(*kueue.Workload) (bool, error)
+
+// PatchStatusOption defines a functional option for customizing PatchStatusOptions.
 // It follows the functional options pattern, allowing callers to configure
-// patch behavior at call sites without directly manipulating PatchAdmissionStatusOptions.
-type PatchAdmissionStatusOption func(*PatchAdmissionStatusOptions)
+// patch behavior at call sites without directly manipulating PatchStatusOptions.
+type PatchStatusOption func(*PatchStatusOptions)
 
-// PatchAdmissionStatusOptions contains configuration parameters that control how patches
+// PatchStatusOptions contains configuration parameters that control how patches
 // are generated and applied.
 //
 // Fields:
@@ -983,28 +966,28 @@ type PatchAdmissionStatusOption func(*PatchAdmissionStatusOptions)
 //     patch. Defaults to true. Setting StrictPatch=false preserves the current
 //     ResourceVersion.
 //
-// Typically, PatchAdmissionStatusOptions are constructed via DefaultPatchAdmissionStatusOptions and
-// modified using PatchAdmissionStatusOption functions (e.g., WithLoose).
-type PatchAdmissionStatusOptions struct {
+// Typically, PatchStatusOptions are constructed via DefaultPatchStatusOptions and
+// modified using PatchStatusOption functions (e.g., WithLoose).
+type PatchStatusOptions struct {
 	StrictPatch             bool
 	StrictApply             bool
 	RetryOnConflictForPatch bool
 }
 
-// DefaultPatchAdmissionStatusOptions returns a new PatchAdmissionStatusOptions instance configured with
+// DefaultPatchStatusOptions returns a new PatchStatusOptions instance configured with
 // default settings.
 //
 // By default, StrictPatch and StrictApply is set to true, meaning ResourceVersion is cleared
 // from the original object so it will always be included in the generated
 // patch. This ensures stricter version handling during patch application.
-func DefaultPatchAdmissionStatusOptions() *PatchAdmissionStatusOptions {
-	return &PatchAdmissionStatusOptions{
+func DefaultPatchStatusOptions() *PatchStatusOptions {
+	return &PatchStatusOptions{
 		StrictPatch: true, // default is strict
 		StrictApply: true, // default is strict
 	}
 }
 
-// WithLooseOnApply returns a PatchAdmissionStatusOption that resets the StrictApply field on PatchAdmissionStatusOptions.
+// WithLooseOnApply returns a PatchStatusOption that resets the StrictApply field on PatchStatusOptions.
 //
 // When using Patch Apply, setting StrictApply to false enforces looser
 // version handling only for Patch Apply.
@@ -1016,52 +999,83 @@ func DefaultPatchAdmissionStatusOptions() *PatchAdmissionStatusOptions {
 //	    return updateFn(obj), nil
 //	}, WithLooseOnApply()) // disables strict mode for Patch Apply
 
-func WithLooseOnApply() PatchAdmissionStatusOption {
-	return func(o *PatchAdmissionStatusOptions) {
+func WithLooseOnApply() PatchStatusOption {
+	return func(o *PatchStatusOptions) {
 		o.StrictApply = false
 	}
 }
 
-// WithRetryOnConflictForPatch configures PatchAdmissionStatusOptions to enable retry logic on conflicts.
+// WithRetryOnConflictForPatch configures PatchStatusOptions to enable retry logic on conflicts.
 // Note: This only works with merge patches.
-func WithRetryOnConflictForPatch() PatchAdmissionStatusOption {
-	return func(o *PatchAdmissionStatusOptions) {
+func WithRetryOnConflictForPatch() PatchStatusOption {
+	return func(o *PatchStatusOptions) {
 		o.RetryOnConflictForPatch = true
 	}
 }
 
-// PatchAdmissionStatus updates the admission status of a workload.
-// If the WorkloadRequestUseMergePatch feature is enabled, it uses a Merge Patch with update function.
-// Otherwise, it runs the update function and, if updated, applies the SSA Patch status.
-func PatchAdmissionStatus(ctx context.Context, c client.Client, w *kueue.Workload, clk clock.Clock, update func(*kueue.Workload) (bool, error), options ...PatchAdmissionStatusOption) error {
-	opts := DefaultPatchAdmissionStatusOptions()
+func patchStatusOptions(options []PatchStatusOption) *PatchStatusOptions {
+	opts := DefaultPatchStatusOptions()
 	for _, opt := range options {
 		opt(opts)
 	}
-	var err error
-	wlCopy := w.DeepCopy()
+	return opts
+}
+
+// patchStatus updates the status of a workload.
+// If the WorkloadRequestUseMergePatch feature is enabled, it uses a Merge Patch with update function.
+// Otherwise, it runs the update function and, if updated, applies the SSA Patch status.
+func patchStatus(ctx context.Context, c client.Client, wl *kueue.Workload, owner client.FieldOwner, update UpdateFunc, opts *PatchStatusOptions) error {
+	wlCopy := wl.DeepCopy()
 	if features.Enabled(features.WorkloadRequestUseMergePatch) {
-		var patchOptions []clientutil.PatchOption
+		patchOptions := make([]clientutil.PatchOption, 0, 2)
 		if !opts.StrictPatch {
 			patchOptions = append(patchOptions, clientutil.WithLoose())
 		}
 		if opts.RetryOnConflictForPatch {
 			patchOptions = append(patchOptions, clientutil.WithRetryOnConflict())
 		}
-		err = clientutil.PatchStatus(ctx, c, wlCopy, func() (bool, error) {
+		err := clientutil.PatchStatus(ctx, c, wlCopy, func() (bool, error) {
 			return update(wlCopy)
 		}, patchOptions...)
+		if err != nil {
+			return err
+		}
 	} else {
 		if updated, err := update(wlCopy); err != nil || !updated {
 			return err
 		}
-		wlCopy = PrepareWorkloadPatch(wlCopy, opts.StrictApply, clk)
-		err = c.Status().Patch(ctx, wlCopy, client.Apply, client.FieldOwner(constants.AdmissionName), client.ForceOwnership)
-	}
-	if err == nil {
-		wlCopy.DeepCopyInto(w)
+		err := c.Status().Patch(ctx, wlCopy, client.Apply, owner, client.ForceOwnership)
+		if err != nil {
+			return err
+		}
 	}
-	return err
+	wlCopy.DeepCopyInto(wl)
+	return nil
+}
+
+func PatchStatus(ctx context.Context, c client.Client, wl *kueue.Workload, owner client.FieldOwner, update UpdateFunc, options ...PatchStatusOption) error {
+	opts := patchStatusOptions(options)
+	return patchStatus(ctx, c, wl, owner, func(wl *kueue.Workload) (bool, error) {
+		if !features.Enabled(features.WorkloadRequestUseMergePatch) {
+			wlPatch := BaseSSAWorkload(wl, opts.StrictApply)
+			wlPatch.DeepCopyInto(wl)
+		}
+		return update(wl)
+	}, opts)
+}
+
+func PatchAdmissionStatus(ctx context.Context, c client.Client, wl *kueue.Workload, clk clock.Clock, update UpdateFunc, options ...PatchStatusOption) error {
+	opts := patchStatusOptions(options)
+	return patchStatus(ctx, c, wl, constants.AdmissionName, func(wl *kueue.Workload) (bool, error) {
+		if updated, err := update(wl); err != nil || !updated {
+			return updated, err
+		}
+		if !features.Enabled(features.WorkloadRequestUseMergePatch) {
+			wlPatch := PrepareWorkloadPatch(wl, opts.StrictApply, clk)
+			wlPatch.DeepCopyInto(wl)
+		}
+		return true, nil
+	}, opts)
 }
 
 type Ordering struct {