From ccc1b3a6d6b1d81f6cd7f19364c30601cf9ad98b Mon Sep 17 00:00:00 2001
From: Patryk Bundyra <pbundyra@google.com>
Date: Tue, 16 Apr 2024 08:50:18 +0000
Subject: [PATCH] Fix a bug on assigning AdmissionChecks for a Workload

---
 pkg/cache/clusterqueue.go                  |  2 +-
 pkg/cache/snapshot.go                      |  4 +---
 pkg/controller/core/workload_controller.go |  1 +
 pkg/util/maps/maps.go                      | 11 +++++++++++
 pkg/workload/workload.go                   | 14 ++++++++++++++
 5 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/pkg/cache/clusterqueue.go b/pkg/cache/clusterqueue.go
index b4ff4025ed..0264883409 100644
--- a/pkg/cache/clusterqueue.go
+++ b/pkg/cache/clusterqueue.go
@@ -418,8 +418,8 @@ func (c *ClusterQueue) updateLabelKeys(flavors map[kueue.ResourceFlavorReference
 // updateWithAdmissionChecks updates a ClusterQueue based on the passed AdmissionChecks set.
 func (c *ClusterQueue) updateWithAdmissionChecks(checks map[string]AdmissionCheck) {
 	hasMissing := false
-	singleInstanceControllers := sets.New[string]()
 	checksPerController := make(map[string]int, len(c.AdmissionChecks))
+	singleInstanceControllers := sets.New[string]()
 	for acName, _ := range c.AdmissionChecks {
 		if ac, found := checks[acName]; !found {
 			hasMissing = true
diff --git a/pkg/cache/snapshot.go b/pkg/cache/snapshot.go
index 64e145d257..c5cdc3ff9e 100644
--- a/pkg/cache/snapshot.go
+++ b/pkg/cache/snapshot.go
@@ -143,10 +143,8 @@ func (c *ClusterQueue) snapshot() *ClusterQueue {
 		Preemption:                    c.Preemption,
 		NamespaceSelector:             c.NamespaceSelector,
 		Status:                        c.Status,
+		AdmissionChecks:               utilmaps.DeepCopySets[string](c.AdmissionChecks),
 	}
-	admissionChecks := make(map[string]sets.Set[string], len(c.AdmissionChecks))
-	maps.Copy(admissionChecks, c.AdmissionChecks)
-	c.AdmissionChecks = admissionChecks
 
 	for fName, rUsage := range c.Usage {
 		cc.Usage[fName] = maps.Clone(rUsage)
diff --git a/pkg/controller/core/workload_controller.go b/pkg/controller/core/workload_controller.go
index 0a964ce7b1..39913b44f7 100644
--- a/pkg/controller/core/workload_controller.go
+++ b/pkg/controller/core/workload_controller.go
@@ -266,6 +266,7 @@ func (r *WorkloadReconciler) reconcileSyncAdmissionChecks(ctx context.Context, w
 	}
 
 	admissionChecks := workload.AdmissionChecksForWorkload(log, wl, utilac.NewAdmissionChecks(&cq))
+	log.V(2).Info("PATRYK created ac", "ac", admissionChecks)
 	newChecks, shouldUpdate := syncAdmissionCheckConditions(wl.Status.AdmissionChecks, admissionChecks)
 	if shouldUpdate {
 		log.V(3).Info("The workload needs admission checks updates", "clusterQueue", klog.KRef("", cqName), "admissionChecks", admissionChecks)
diff --git a/pkg/util/maps/maps.go b/pkg/util/maps/maps.go
index 675d72fa50..5eb5edaca5 100644
--- a/pkg/util/maps/maps.go
+++ b/pkg/util/maps/maps.go
@@ -21,6 +21,8 @@ package maps
 import (
 	"fmt"
 	"maps"
+
+	"k8s.io/apimachinery/pkg/util/sets"
 )
 
 // Merge merges a and b while resolving the conflicts by calling commonKeyValue
@@ -108,3 +110,12 @@ func FilterKeys[K comparable, V any, M ~map[K]V](m M, k []K) M {
 	}
 	return ret
 }
+
+// DeepCopySets create a deep copy of map[string]Set which would otherwise be referenced
+func DeepCopySets[T comparable](src map[string]sets.Set[T]) map[string]sets.Set[T] {
+	copy := make(map[string]sets.Set[T], len(src))
+	for key, set := range src {
+		copy[key] = set.Clone()
+	}
+	return copy
+}
diff --git a/pkg/workload/workload.go b/pkg/workload/workload.go
index 58a535e301..57d5b2e30f 100644
--- a/pkg/workload/workload.go
+++ b/pkg/workload/workload.go
@@ -39,6 +39,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/util/api"
 	"sigs.k8s.io/kueue/pkg/util/limitrange"
+	utilmaps "sigs.k8s.io/kueue/pkg/util/maps"
 )
 
 var (
@@ -550,6 +551,19 @@ func RemoveFinalizer(ctx context.Context, c client.Client, wl *kueue.Workload) e
 // AdmissionChecksForWorkload returns AdmissionChecks that should be assigned to a specific Workload based on
 // ClusterQueue configuration and ResourceFlavors
 func AdmissionChecksForWorkload(log logr.Logger, wl *kueue.Workload, admissionChecks map[string]sets.Set[string]) sets.Set[string] {
+	// If all admissionChecks should be run for all flavors we don't need to wait for Workload's Admission to be set.
+	// This is also the case if admissionChecks are specified with ClusterQueue.Spec.AdmissionChecks instead of
+	// ClusterQueue.Spec.AdmissionCheckStrategy
+	allFlavors := true
+	for _, flavors := range admissionChecks {
+		if len(flavors) != 0 {
+			allFlavors = false
+		}
+	}
+	if allFlavors {
+		return sets.New(utilmaps.Keys(admissionChecks)...)
+	}
+
 	// Kueue sets AdmissionChecks first based on ClusterQueue configuration and at this point Workload has no
 	// ResourceFlavors assigned, so we cannot match AdmissionChecks to ResourceFlavor.
 	// After Quota is reserved, another reconciliation happens and we can match AdmissionChecks to ResourceFlavors