From a2d05b038baa475ca3fd55d00715f3e86e8047ed Mon Sep 17 00:00:00 2001
From: Mykhailo Bobrovskyi <mikhail.bobrovsky@gmail.com>
Date: Fri, 31 Jan 2025 14:08:01 +0200
Subject: [PATCH] Update workload AdmissionChecks on AdmissionRequest creation
 error.

---
 .../provisioning/controller.go                | 23 +++++++++++----
 .../provisioning/controller_test.go           | 28 ++++++++++++++-----
 2 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/pkg/controller/admissionchecks/provisioning/controller.go b/pkg/controller/admissionchecks/provisioning/controller.go
index 308473aa5e..083f5cd191 100644
--- a/pkg/controller/admissionchecks/provisioning/controller.go
+++ b/pkg/controller/admissionchecks/provisioning/controller.go
@@ -306,12 +306,9 @@ func (c *Controller) syncOwnedProvisionRequest(
 			}
 
 			if err := c.client.Create(ctx, req); err != nil {
-				msg := fmt.Sprintf("Error creating ProvisioningRequest %q: %v", requestName, err)
-				ac.Message = api.TruncateConditionMessage(msg)
-				workload.SetAdmissionCheckState(&wl.Status.AdmissionChecks, *ac, c.clock)
-
-				c.record.Eventf(wl, corev1.EventTypeWarning, "FailedCreate", api.TruncateEventMessage(msg))
-				return nil, err
+				msg := api.TruncateEventMessage(fmt.Sprintf("Error creating ProvisioningRequest %q: %v", requestName, err))
+				c.record.Eventf(wl, corev1.EventTypeWarning, "FailedCreate", msg)
+				return nil, c.handleError(ctx, wl, ac, msg, err)
 			}
 			c.record.Eventf(wl, corev1.EventTypeNormal, "ProvisioningRequestCreated", "Created ProvisioningRequest: %q", req.Name)
 			activeOrLastPRForChecks[checkName] = req
@@ -343,6 +340,20 @@ func (c *Controller) remainingTimeToRetry(pr *autoscaling.ProvisioningRequest, f
 	return backoffDuration - timeElapsedSinceLastFailure
 }
 
+func (c *Controller) handleError(ctx context.Context, wl *kueue.Workload, ac *kueue.AdmissionCheckState, msg string, err error) error {
+	ac.Message = msg
+	wlPatch := workload.BaseSSAWorkload(wl)
+	workload.SetAdmissionCheckState(&wlPatch.Status.AdmissionChecks, *ac, c.clock)
+
+	patchErr := c.client.Status().Patch(
+		ctx, wlPatch, client.Apply,
+		client.FieldOwner(kueue.ProvisioningRequestControllerName),
+		client.ForceOwnership,
+	)
+
+	return errors.Join(err, patchErr)
+}
+
 func (c *Controller) syncProvisionRequestsPodTemplates(ctx context.Context, wl *kueue.Workload, prName string, prc *kueue.ProvisioningRequestConfig) error {
 	request := &autoscaling.ProvisioningRequest{}
 	requestKey := types.NamespacedName{
diff --git a/pkg/controller/admissionchecks/provisioning/controller_test.go b/pkg/controller/admissionchecks/provisioning/controller_test.go
index 8b3324371e..48157c05a7 100644
--- a/pkg/controller/admissionchecks/provisioning/controller_test.go
+++ b/pkg/controller/admissionchecks/provisioning/controller_test.go
@@ -1151,6 +1151,20 @@ func TestReconcile(t *testing.T) {
 			checks:             []kueue.AdmissionCheck{*baseCheck.DeepCopy()},
 			configs:            []kueue.ProvisioningRequestConfig{*utiltesting.MakeProvisioningRequestConfig("config1").Obj()},
 			wantReconcileError: errInvalidProvisioningRequest,
+			wantWorkloads: map[string]*kueue.Workload{
+				"wl": utiltesting.MakeWorkload("wl", TestNamespace).
+					Annotations(map[string]string{
+						"provreq.kueue.x-k8s.io/ValidUntilSeconds": "0",
+						"invalid-provreq-prefix/Foo1":              "Bar1",
+						"another-invalid-provreq-prefix/Foo2":      "Bar2"}).
+					AdmissionChecks(kueue.AdmissionCheckState{
+						Name:    "check1",
+						State:   kueue.CheckStatePending,
+						Message: "Error creating ProvisioningRequest \"wl-check1-1\": invalid ProvisioningRequest error",
+					}).
+					ReserveQuota(utiltesting.MakeAdmission("q1").Obj()).
+					Obj(),
+			},
 			wantEvents: []utiltesting.EventRecord{
 				{
 					Key:       client.ObjectKeyFromObject(baseWorkload),
@@ -1167,16 +1181,16 @@ func TestReconcile(t *testing.T) {
 			for _, gate := range tc.enableGates {
 				features.SetFeatureGateDuringTest(t, gate, true)
 			}
-			builder, ctx := getClientBuilder()
-			builder = builder.WithInterceptorFuncs(interceptor.Funcs{SubResourcePatch: utiltesting.TreatSSAAsStrategicMerge})
 
+			interceptorFuncs := interceptor.Funcs{SubResourcePatch: utiltesting.TreatSSAAsStrategicMerge}
 			if tc.wantReconcileError != nil {
-				builder = builder.WithInterceptorFuncs(
-					interceptor.Funcs{
-						Create: func(ctx context.Context, client client.WithWatch, obj client.Object, opts ...client.CreateOption) error {
-							return tc.wantReconcileError
-						}})
+				interceptorFuncs.Create = func(ctx context.Context, client client.WithWatch, obj client.Object, opts ...client.CreateOption) error {
+					return tc.wantReconcileError
+				}
 			}
+
+			builder, ctx := getClientBuilder()
+			builder = builder.WithInterceptorFuncs(interceptorFuncs)
 			builder = builder.WithObjects(tc.workload)
 			builder = builder.WithStatusSubresource(tc.workload)
 			builder = builder.WithLists(