Repository: kubernetes-sigs/kueue
Pull request #295 (2022-07-11 15:25:33)

Track number of active running workloads per queue
#### What type of PR is this?

/kind feature

#### What this PR does / why we need it:

This PR doesn't have any user visible changes. The information can be used to populate the information into a Queue status, to improve observability of kueue.

This is a spin off from #291

#### Which issue(s) this PR fixes:

Part of #259

#### Special notes for your reviewer:



k8s-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/kubernetes-sigs/kueue/pull/295#" title="Author self-approved">alculquicondor</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=kubernetes-sigs%2Fkueue).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/kubernetes-sigs/kueue/blob/main/OWNERS)~~ [alculquicondor]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
alculquicondor: Part of #259
ahg-g: /label tide/merge-method-squash
ahg-g: /lgtm

Patch:
diff --git a/pkg/cache/cache.go b/pkg/cache/cache.go
index 63d333bb9..94644e5e4 100644
--- a/pkg/cache/cache.go
+++ b/pkg/cache/cache.go
@@ -33,9 +33,13 @@ import (
 	"sigs.k8s.io/kueue/pkg/workload"
 )
 
-const workloadClusterQueueKey = "spec.admission.clusterQueue"
+const (
+	workloadClusterQueueKey = "spec.admission.clusterQueue"
+	queueClusterQueueKey    = "spec.clusterQueue"
+)
 
 var (
+	errQueueAlreadyExists  = errors.New("queue already exists")
 	errCqNotFound          = errors.New("cluster queue not found")
 	errWorkloadNotAdmitted = errors.New("workload not admitted by a ClusterQueue")
 )
@@ -107,6 +111,10 @@ type ClusterQueue struct {
 	// that can be matched against the flavors.
 	LabelKeys map[corev1.ResourceName]sets.String
 	Status    ClusterQueueStatus
+
+	// The following fields are not populated in a snapshot.
+
+	admittedWorkloadsPerQueue map[string]int
 }
 
 // FlavorLimits holds a processed ClusterQueue flavor quota.
@@ -118,8 +126,9 @@ type FlavorLimits struct {
 
 func (c *Cache) newClusterQueue(cq *kueue.ClusterQueue) (*ClusterQueue, error) {
 	cqImpl := &ClusterQueue{
-		Name:      cq.Name,
-		Workloads: map[string]*workload.Info{},
+		Name:                      cq.Name,
+		Workloads:                 make(map[string]*workload.Info),
+		admittedWorkloadsPerQueue: make(map[string]int),
 	}
 	if err := cqImpl.update(cq, c.resourceFlavors); err != nil {
 		return nil, err
@@ -232,6 +241,32 @@ func (c *ClusterQueue) updateWorkloadUsage(wi *workload.Info, m int64) {
 			}
 		}
 	}
+	qKey := workload.QueueKey(wi.Obj)
+	if _, ok := c.admittedWorkloadsPerQueue[qKey]; ok {
+		c.admittedWorkloadsPerQueue[qKey] += int(m)
+	}
+}
+
+func (c *ClusterQueue) addQueue(q *kueue.Queue) error {
+	qKey := queueKey(q)
+	if _, ok := c.admittedWorkloadsPerQueue[qKey]; ok {
+		return errQueueAlreadyExists
+	}
+	// We need to count the workloads, because they could have been added before
+	// receiving the queue add event.
+	workloads := 0
+	for _, wl := range c.Workloads {
+		if workloadBelongsToQueue(wl.Obj, q) {
+			workloads++
+		}
+	}
+	c.admittedWorkloadsPerQueue[qKey] = workloads
+	return nil
+}
+
+func (c *ClusterQueue) deleteQueue(q *kueue.Queue) {
+	qKey := queueKey(q)
+	delete(c.admittedWorkloadsPerQueue, qKey)
 }
 
 func (c *ClusterQueue) flavorInUse(flavor string) bool {
@@ -301,7 +336,18 @@ func (c *Cache) AddClusterQueue(ctx context.Context, cq *kueue.ClusterQueue) err
 	c.clusterQueues[cq.Name] = cqImpl
 
 	// On controller restart, an add ClusterQueue event may come after
-	// add workload events, and so here we explicitly list and add existing workloads.
+	// add queue and workload, so here we explicitly list and add existing queues
+	// and workloads.
+	var queues kueue.QueueList
+	if err := c.client.List(ctx, &queues, client.MatchingFields{queueClusterQueueKey: cq.Name}); err != nil {
+		return fmt.Errorf("listing queues that match the clusterQueue: %w", err)
+	}
+	for _, q := range queues.Items {
+		// Checking ClusterQueue name again because the field index is not available in tests.
+		if string(q.Spec.ClusterQueue) == cq.Name {
+			cqImpl.admittedWorkloadsPerQueue[queueKey(&q)] = 0
+		}
+	}
 	var workloads kueue.WorkloadList
 	if err := c.client.List(ctx, &workloads, client.MatchingFields{workloadClusterQueueKey: cq.Name}); err != nil {
 		return fmt.Errorf("listing workloads that match the queue: %w", err)
@@ -312,6 +358,9 @@ func (c *Cache) AddClusterQueue(ctx context.Context, cq *kueue.ClusterQueue) err
 			continue
 		}
 		c.addOrUpdateWorkload(&workloads.Items[i])
+		if _, ok := cqImpl.admittedWorkloadsPerQueue[w.Spec.QueueName]; ok {
+			cqImpl.admittedWorkloadsPerQueue[w.Spec.QueueName]++
+		}
 	}
 
 	return nil
@@ -351,6 +400,43 @@ func (c *Cache) DeleteClusterQueue(cq *kueue.ClusterQueue) {
 	delete(c.clusterQueues, cq.Name)
 }
 
+func (c *Cache) AddQueue(q *kueue.Queue) error {
+	c.Lock()
+	defer c.Unlock()
+	cq, ok := c.clusterQueues[string(q.Spec.ClusterQueue)]
+	if !ok {
+		return nil
+	}
+	return cq.addQueue(q)
+}
+
+func (c *Cache) DeleteQueue(q *kueue.Queue) {
+	c.Lock()
+	defer c.Unlock()
+	cq, ok := c.clusterQueues[string(q.Spec.ClusterQueue)]
+	if !ok {
+		return
+	}
+	cq.deleteQueue(q)
+}
+
+func (c *Cache) UpdateQueue(oldQ, newQ *kueue.Queue) error {
+	if oldQ.Spec.ClusterQueue == newQ.Spec.ClusterQueue {
+		return nil
+	}
+	c.Lock()
+	defer c.Unlock()
+	cq, ok := c.clusterQueues[string(oldQ.Spec.ClusterQueue)]
+	if ok {
+		cq.deleteQueue(oldQ)
+	}
+	cq, ok = c.clusterQueues[string(newQ.Spec.ClusterQueue)]
+	if ok {
+		return cq.addQueue(newQ)
+	}
+	return nil
+}
+
 func (c *Cache) AddOrUpdateWorkload(w *kueue.Workload) bool {
 	c.Lock()
 	defer c.Unlock()
@@ -574,3 +660,12 @@ func SetupIndexes(indexer client.FieldIndexer) error {
 		return []string{string(wl.Spec.Admission.ClusterQueue)}
 	})
 }
+
+func workloadBelongsToQueue(wl *kueue.Workload, q *kueue.Queue) bool {
+	return wl.Namespace == q.Namespace && wl.Spec.QueueName == q.Name
+}
+
+// Key is the key used to index the queue.
+func queueKey(q *kueue.Queue) string {
+	return fmt.Sprintf("%s/%s", q.Namespace, q.Name)
+}
diff --git a/pkg/controller/core/core.go b/pkg/controller/core/core.go
index 99795593a..00166de74 100644
--- a/pkg/controller/core/core.go
+++ b/pkg/controller/core/core.go
@@ -32,7 +32,7 @@ func SetupControllers(mgr ctrl.Manager, qManager *queue.Manager, cc *cache.Cache
 	if err := rfRec.SetupWithManager(mgr); err != nil {
 		return "ResourceFlavor", err
 	}
-	qRec := NewQueueReconciler(mgr.GetClient(), qManager)
+	qRec := NewQueueReconciler(mgr.GetClient(), qManager, cc)
 	if err := qRec.SetupWithManager(mgr); err != nil {
 		return "Queue", err
 	}
diff --git a/pkg/controller/core/queue_controller.go b/pkg/controller/core/queue_controller.go
index df38c8811..8a6576f23 100644
--- a/pkg/controller/core/queue_controller.go
+++ b/pkg/controller/core/queue_controller.go
@@ -31,6 +31,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/source"
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
+	"sigs.k8s.io/kueue/pkg/cache"
 	"sigs.k8s.io/kueue/pkg/constants"
 	"sigs.k8s.io/kueue/pkg/queue"
 )
@@ -40,13 +41,15 @@ type QueueReconciler struct {
 	client     client.Client
 	log        logr.Logger
 	queues     *queue.Manager
+	cache      *cache.Cache
 	wlUpdateCh chan event.GenericEvent
 }
 
-func NewQueueReconciler(client client.Client, queues *queue.Manager) *QueueReconciler {
+func NewQueueReconciler(client client.Client, queues *queue.Manager, cache *cache.Cache) *QueueReconciler {
 	return &QueueReconciler{
 		log:        ctrl.Log.WithName("queue-reconciler"),
 		queues:     queues,
+		cache:      cache,
 		client:     client,
 		wlUpdateCh: make(chan event.GenericEvent, updateChBuffer),
 	}
@@ -98,7 +101,10 @@ func (r *QueueReconciler) Create(e event.CreateEvent) bool {
 	log.V(2).Info("Queue create event")
 	ctx := logr.NewContext(context.Background(), log)
 	if err := r.queues.AddQueue(ctx, q); err != nil {
-		log.Error(err, "Failed to add queue to system")
+		log.Error(err, "Failed to add queue to the queueing system")
+	}
+	if err := r.cache.AddQueue(q); err != nil {
+		log.Error(err, "Failed to add queue to the cache")
 	}
 	return true
 }
@@ -111,6 +117,7 @@ func (r *QueueReconciler) Delete(e event.DeleteEvent) bool {
 	}
 	r.log.V(2).Info("Queue delete event", "queue", klog.KObj(q))
 	r.queues.DeleteQueue(q)
+	r.cache.DeleteQueue(q)
 	return true
 }
 
@@ -123,7 +130,11 @@ func (r *QueueReconciler) Update(e event.UpdateEvent) bool {
 	log := r.log.WithValues("queue", klog.KObj(q))
 	log.V(2).Info("Queue update event")
 	if err := r.queues.UpdateQueue(q); err != nil {
-		log.Error(err, "Failed to update queue in system")
+		log.Error(err, "Failed to update queue in the queueing system")
+	}
+	oldQ := e.ObjectOld.(*kueue.Queue)
+	if err := r.cache.UpdateQueue(oldQ, q); err != nil {
+		log.Error(err, "Failed to update queue in the cache")
 	}
 	return true
 }
diff --git a/pkg/metrics/metrics.go b/pkg/metrics/metrics.go
index d50f5a173..1b785b5e3 100644
--- a/pkg/metrics/metrics.go
+++ b/pkg/metrics/metrics.go
@@ -54,7 +54,8 @@ var (
 			Subsystem: subsystemName,
 			Name:      "pending_workloads",
 			Help:      "Number of pending workloads, per cluster_queue.",
-		}, []string{"cluster_queue"})
+		}, []string{"cluster_queue"},
+	)
 )
 
 func AdmissionAttempt(result AdmissionResult, duration time.Duration) {
diff --git a/pkg/queue/manager.go b/pkg/queue/manager.go
index a716da90e..1796a686f 100644
--- a/pkg/queue/manager.go
+++ b/pkg/queue/manager.go
@@ -240,7 +240,7 @@ func (m *Manager) Pending(cq *kueue.ClusterQueue) int32 {
 func (m *Manager) QueueForWorkloadExists(wl *kueue.Workload) bool {
 	m.RLock()
 	defer m.RUnlock()
-	_, ok := m.queues[queueKeyForWorkload(wl)]
+	_, ok := m.queues[workload.QueueKey(wl)]
 	return ok
 
 }
@@ -251,7 +251,7 @@ func (m *Manager) QueueForWorkloadExists(wl *kueue.Workload) bool {
 func (m *Manager) ClusterQueueForWorkload(wl *kueue.Workload) (string, bool) {
 	m.RLock()
 	defer m.RUnlock()
-	q, ok := m.queues[queueKeyForWorkload(wl)]
+	q, ok := m.queues[workload.QueueKey(wl)]
 	if !ok {
 		return "", false
 	}
@@ -268,7 +268,7 @@ func (m *Manager) AddOrUpdateWorkload(w *kueue.Workload) bool {
 }
 
 func (m *Manager) addOrUpdateWorkload(w *kueue.Workload) bool {
-	qKey := queueKeyForWorkload(w)
+	qKey := workload.QueueKey(w)
 	q := m.queues[qKey]
 	if q == nil {
 		return false
@@ -300,7 +300,7 @@ func (m *Manager) RequeueWorkload(ctx context.Context, info *workload.Info, imme
 		return false
 	}
 
-	q := m.queues[queueKeyForWorkload(&w)]
+	q := m.queues[workload.QueueKey(&w)]
 	if q == nil {
 		return false
 	}
@@ -321,7 +321,7 @@ func (m *Manager) RequeueWorkload(ctx context.Context, info *workload.Info, imme
 
 func (m *Manager) DeleteWorkload(w *kueue.Workload) {
 	m.Lock()
-	m.deleteWorkloadFromQueueAndClusterQueue(w, queueKeyForWorkload(w))
+	m.deleteWorkloadFromQueueAndClusterQueue(w, workload.QueueKey(w))
 	m.Unlock()
 }
 
@@ -345,7 +345,7 @@ func (m *Manager) QueueAssociatedInadmissibleWorkloads(w *kueue.Workload) {
 	m.Lock()
 	defer m.Unlock()
 
-	q := m.queues[queueKeyForWorkload(w)]
+	q := m.queues[workload.QueueKey(w)]
 	if q == nil {
 		return
 	}
@@ -417,7 +417,7 @@ func (m *Manager) UpdateWorkload(oldW, w *kueue.Workload) bool {
 	m.Lock()
 	defer m.Unlock()
 	if oldW.Spec.QueueName != w.Spec.QueueName {
-		m.deleteWorkloadFromQueueAndClusterQueue(w, queueKeyForWorkload(oldW))
+		m.deleteWorkloadFromQueueAndClusterQueue(w, workload.QueueKey(oldW))
 	}
 	return m.addOrUpdateWorkload(w)
 }
@@ -486,7 +486,7 @@ func (m *Manager) heads() []workload.Info {
 		wlCopy := *wl
 		wlCopy.ClusterQueue = cqName
 		workloads = append(workloads, wlCopy)
-		q := m.queues[queueKeyForWorkload(wl.Obj)]
+		q := m.queues[workload.QueueKey(wl.Obj)]
 		delete(q.items, workload.Key(wl.Obj))
 	}
 	return workloads
diff --git a/pkg/queue/queue.go b/pkg/queue/queue.go
index 0582cf750..86e0b0932 100644
--- a/pkg/queue/queue.go
+++ b/pkg/queue/queue.go
@@ -33,11 +33,6 @@ func Key(q *kueue.Queue) string {
 	return fmt.Sprintf("%s/%s", q.Namespace, q.Name)
 }
 
-// queueKeyForWorkload is the key to find the queue for the workload in the index.
-func queueKeyForWorkload(w *kueue.Workload) string {
-	return fmt.Sprintf("%s/%s", w.Namespace, w.Spec.QueueName)
-}
-
 // Queue is the internal implementation of kueue.Queue.
 type Queue struct {
 	Key          string
diff --git a/pkg/workload/workload.go b/pkg/workload/workload.go
index b5d97b666..6868f9791 100644
--- a/pkg/workload/workload.go
+++ b/pkg/workload/workload.go
@@ -59,6 +59,10 @@ func Key(w *kueue.Workload) string {
 	return fmt.Sprintf("%s/%s", w.Namespace, w.Name)
 }
 
+func QueueKey(w *kueue.Workload) string {
+	return fmt.Sprintf("%s/%s", w.Namespace, w.Spec.QueueName)
+}
+
 func totalRequests(spec *kueue.WorkloadSpec) []PodSetResources {
 	if len(spec.PodSets) == 0 {
 		return nil


Test patch:
diff --git a/pkg/cache/cache_test.go b/pkg/cache/cache_test.go
index 123f7cc64..646c158c2 100644
--- a/pkg/cache/cache_test.go
+++ b/pkg/cache/cache_test.go
@@ -29,6 +29,7 @@ import (
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/util/sets"
+	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
@@ -461,7 +462,7 @@ func TestCacheClusterQueueOperations(t *testing.T) {
 			cache := New(fake.NewClientBuilder().WithScheme(scheme).Build())
 			tc.operation(cache)
 			if diff := cmp.Diff(tc.wantClusterQueues, cache.clusterQueues,
-				cmpopts.IgnoreFields(ClusterQueue{}, "Cohort", "Workloads")); diff != "" {
+				cmpopts.IgnoreFields(ClusterQueue{}, "Cohort", "Workloads"), cmpopts.IgnoreUnexported(ClusterQueue{})); diff != "" {
 				t.Errorf("Unexpected clusterQueues (-want,+got):\n%s", diff)
 			}
 
@@ -547,7 +548,7 @@ func TestCacheWorkloadOperations(t *testing.T) {
 	if err := kueue.AddToScheme(scheme); err != nil {
 		t.Fatalf("Failed adding kueue scheme: %v", err)
 	}
-	client := fake.NewClientBuilder().WithScheme(scheme).WithObjects(
+	cl := fake.NewClientBuilder().WithScheme(scheme).WithObjects(
 		utiltesting.MakeWorkload("a", "").PodSets(podSets).Admit(&kueue.Admission{
 			ClusterQueue:  "one",
 			PodSetFlavors: podSetFlavors,
@@ -583,11 +584,10 @@ func TestCacheWorkloadOperations(t *testing.T) {
 					utiltesting.MakeWorkload("d", "").Admit(&kueue.Admission{
 						ClusterQueue: "two",
 					}).Obj(),
+					utiltesting.MakeWorkload("pending", "").Obj(),
 				}
 				for i := range workloads {
-					if !cache.AddOrUpdateWorkload(workloads[i]) {
-						return fmt.Errorf("failed to add workload")
-					}
+					cache.AddOrUpdateWorkload(workloads[i])
 				}
 				return nil
 			},
@@ -956,7 +956,7 @@ func TestCacheWorkloadOperations(t *testing.T) {
 	}
 	for _, step := range steps {
 		t.Run(step.name, func(t *testing.T) {
-			cache := New(client)
+			cache := New(cl)
 
 			for _, c := range clusterQueues {
 				if err := cache.AddClusterQueue(context.Background(), &c); err != nil {
@@ -1161,6 +1161,254 @@ func TestClusterQueueUsage(t *testing.T) {
 	}
 }
 
+func TestCacheQueueOperations(t *testing.T) {
+	cqs := []*kueue.ClusterQueue{
+		utiltesting.MakeClusterQueue("foo").Obj(),
+		utiltesting.MakeClusterQueue("bar").Obj(),
+	}
+	queues := []*kueue.Queue{
+		utiltesting.MakeQueue("alpha", "ns1").ClusterQueue("foo").Obj(),
+		utiltesting.MakeQueue("beta", "ns2").ClusterQueue("foo").Obj(),
+		utiltesting.MakeQueue("gamma", "ns1").ClusterQueue("bar").Obj(),
+	}
+	workloads := []*kueue.Workload{
+		utiltesting.MakeWorkload("job1", "ns1").Queue("alpha").Admit(utiltesting.MakeAdmission("foo").Obj()).Obj(),
+		utiltesting.MakeWorkload("job2", "ns2").Queue("beta").Admit(utiltesting.MakeAdmission("foo").Obj()).Obj(),
+		utiltesting.MakeWorkload("job3", "ns1").Queue("gamma").Admit(utiltesting.MakeAdmission("bar").Obj()).Obj(),
+		utiltesting.MakeWorkload("job4", "ns2").Queue("beta").Admit(utiltesting.MakeAdmission("foo").Obj()).Obj(),
+	}
+	insertAllClusterQueues := func(ctx context.Context, cl client.Client, cache *Cache) error {
+		for _, cq := range cqs {
+			cq := cq.DeepCopy()
+			if err := cl.Create(ctx, cq); err != nil {
+				return err
+			}
+			if err := cache.AddClusterQueue(ctx, cq); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+	insertAllQueues := func(ctx context.Context, cl client.Client, cache *Cache) error {
+		for _, q := range queues {
+			q := q.DeepCopy()
+			if err := cl.Create(ctx, q.DeepCopy()); err != nil {
+				return err
+			}
+			if err := cache.AddQueue(q); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+	insertAllWorkloads := func(ctx context.Context, cl client.Client, cache *Cache) error {
+		for _, wl := range workloads {
+			wl := wl.DeepCopy()
+			if err := cl.Create(ctx, wl); err != nil {
+				return err
+			}
+			cache.AddOrUpdateWorkload(wl)
+		}
+		return nil
+	}
+	cases := map[string]struct {
+		ops             []func(context.Context, client.Client, *Cache) error
+		wantQueueCounts map[string]map[string]int
+	}{
+		"insert cqs, queues, workloads": {
+			ops: []func(ctx context.Context, cl client.Client, cache *Cache) error{
+				insertAllClusterQueues,
+				insertAllQueues,
+				insertAllWorkloads,
+			},
+			wantQueueCounts: map[string]map[string]int{
+				"foo": {
+					"ns1/alpha": 1,
+					"ns2/beta":  2,
+				},
+				"bar": {
+					"ns1/gamma": 1,
+				},
+			},
+		},
+		"insert cqs, workloads but no queues": {
+			ops: []func(context.Context, client.Client, *Cache) error{
+				insertAllClusterQueues,
+				insertAllWorkloads,
+			},
+			wantQueueCounts: map[string]map[string]int{
+				"foo": {},
+				"bar": {},
+			},
+		},
+		"insert queues, workloads but no cqs": {
+			ops: []func(context.Context, client.Client, *Cache) error{
+				insertAllQueues,
+				insertAllWorkloads,
+			},
+			wantQueueCounts: map[string]map[string]int{},
+		},
+		"insert queues last": {
+			ops: []func(context.Context, client.Client, *Cache) error{
+				insertAllClusterQueues,
+				insertAllWorkloads,
+				insertAllQueues,
+			},
+			wantQueueCounts: map[string]map[string]int{
+				"foo": {
+					"ns1/alpha": 1,
+					"ns2/beta":  2,
+				},
+				"bar": {
+					"ns1/gamma": 1,
+				},
+			},
+		},
+		"insert cqs last": {
+			ops: []func(context.Context, client.Client, *Cache) error{
+				insertAllQueues,
+				insertAllWorkloads,
+				insertAllClusterQueues,
+			},
+			wantQueueCounts: map[string]map[string]int{
+				"foo": {
+					"ns1/alpha": 1,
+					"ns2/beta":  2,
+				},
+				"bar": {
+					"ns1/gamma": 1,
+				},
+			},
+		},
+		"assume": {
+			ops: []func(context.Context, client.Client, *Cache) error{
+				insertAllClusterQueues,
+				insertAllQueues,
+				func(ctx context.Context, cl client.Client, cache *Cache) error {
+					wl := workloads[0].DeepCopy()
+					if err := cl.Create(ctx, wl); err != nil {
+						return err
+					}
+					return cache.AssumeWorkload(wl)
+				},
+			},
+			wantQueueCounts: map[string]map[string]int{
+				"foo": {
+					"ns1/alpha": 1,
+					"ns2/beta":  0,
+				},
+				"bar": {
+					"ns1/gamma": 0,
+				},
+			},
+		},
+		"assume and forget": {
+			ops: []func(context.Context, client.Client, *Cache) error{
+				insertAllClusterQueues,
+				insertAllQueues,
+				func(ctx context.Context, cl client.Client, cache *Cache) error {
+					wl := workloads[0].DeepCopy()
+					if err := cl.Create(ctx, wl); err != nil {
+						return err
+					}
+					if err := cache.AssumeWorkload(wl); err != nil {
+						return err
+					}
+					return cache.ForgetWorkload(wl)
+				},
+			},
+			wantQueueCounts: map[string]map[string]int{
+				"foo": {
+					"ns1/alpha": 0,
+					"ns2/beta":  0,
+				},
+				"bar": {
+					"ns1/gamma": 0,
+				},
+			},
+		},
+		"delete workload": {
+			ops: []func(ctx context.Context, cl client.Client, cache *Cache) error{
+				insertAllClusterQueues,
+				insertAllQueues,
+				insertAllWorkloads,
+				func(ctx context.Context, cl client.Client, cache *Cache) error {
+					return cache.DeleteWorkload(workloads[0])
+				},
+			},
+			wantQueueCounts: map[string]map[string]int{
+				"foo": {
+					"ns1/alpha": 0,
+					"ns2/beta":  2,
+				},
+				"bar": {
+					"ns1/gamma": 1,
+				},
+			},
+		},
+		"delete cq": {
+			ops: []func(ctx context.Context, cl client.Client, cache *Cache) error{
+				insertAllClusterQueues,
+				insertAllQueues,
+				insertAllWorkloads,
+				func(ctx context.Context, cl client.Client, cache *Cache) error {
+					cache.DeleteClusterQueue(cqs[0])
+					return nil
+				},
+			},
+			wantQueueCounts: map[string]map[string]int{
+				"bar": {
+					"ns1/gamma": 1,
+				},
+			},
+		},
+		"delete queue": {
+			ops: []func(ctx context.Context, cl client.Client, cache *Cache) error{
+				insertAllClusterQueues,
+				insertAllQueues,
+				insertAllWorkloads,
+				func(ctx context.Context, cl client.Client, cache *Cache) error {
+					cache.DeleteQueue(queues[0])
+					return nil
+				},
+			},
+			wantQueueCounts: map[string]map[string]int{
+				"foo": {
+					"ns2/beta": 2,
+				},
+				"bar": {
+					"ns1/gamma": 1,
+				},
+			},
+		},
+		// Not tested: changing a workload's queue and changing a queue's cluster queue.
+		// These operations should not be allowed by the webhook.
+	}
+	for name, tc := range cases {
+		t.Run(name, func(t *testing.T) {
+			scheme := runtime.NewScheme()
+			if err := kueue.AddToScheme(scheme); err != nil {
+				t.Fatalf("Failed adding kueue scheme: %v", err)
+			}
+			cl := fake.NewClientBuilder().WithScheme(scheme).Build()
+			cache := New(cl)
+			ctx := context.Background()
+			for i, op := range tc.ops {
+				if err := op(ctx, cl, cache); err != nil {
+					t.Fatalf("Running op %d: %v", i, err)
+				}
+			}
+			qCounts := make(map[string]map[string]int)
+			for _, cq := range cache.clusterQueues {
+				qCounts[cq.Name] = cq.admittedWorkloadsPerQueue
+			}
+			if diff := cmp.Diff(tc.wantQueueCounts, qCounts); diff != "" {
+				t.Errorf("Wrong active workloads counters for queues (-want,+got):\n%s", diff)
+			}
+		})
+	}
+}
+
 func TestFlavorInUse(t *testing.T) {
 	rf := utiltesting.MakeResourceFlavor("x86").Obj()
 	flavor := utiltesting.MakeFlavor(rf.Name, "5").Obj()
diff --git a/pkg/cache/snapshot_test.go b/pkg/cache/snapshot_test.go
index 0e1234d1c..608f4159c 100644
--- a/pkg/cache/snapshot_test.go
+++ b/pkg/cache/snapshot_test.go
@@ -386,7 +386,7 @@ func TestSnapshot(t *testing.T) {
 		},
 		InactiveClusterQueueSets: sets.String{"flavor-nonexistent-cq": {}},
 	}
-	if diff := cmp.Diff(wantSnapshot, snapshot, cmpopts.IgnoreUnexported(Cohort{})); diff != "" {
+	if diff := cmp.Diff(wantSnapshot, snapshot, cmpopts.IgnoreUnexported(Cohort{}, ClusterQueue{})); diff != "" {
 		t.Errorf("Unexpected Snapshot (-want,+got):\n%s", diff)
 	}
 }
diff --git a/pkg/queue/manager_test.go b/pkg/queue/manager_test.go
index 2379c8f96..9cf4578c2 100644
--- a/pkg/queue/manager_test.go
+++ b/pkg/queue/manager_test.go
@@ -572,7 +572,7 @@ func TestUpdateWorkload(t *testing.T) {
 			if updated := manager.UpdateWorkload(tc.workloads[0], wl); updated != tc.wantUpdated {
 				t.Errorf("UpdatedWorkload returned %t, want %t", updated, tc.wantUpdated)
 			}
-			q := manager.queues[queueKeyForWorkload(wl)]
+			q := manager.queues[workload.QueueKey(wl)]
 			if q != nil {
 				key := workload.Key(wl)
 				item := q.items[key]
diff --git a/test/integration/controller/core/queue_controller_test.go b/test/integration/controller/core/queue_controller_test.go
index 75ebbe75e..b13c37978 100644
--- a/test/integration/controller/core/queue_controller_test.go
+++ b/test/integration/controller/core/queue_controller_test.go
@@ -47,10 +47,11 @@ var _ = ginkgo.Describe("Queue controller", func() {
 	})
 
 	ginkgo.BeforeEach(func() {
-		clusterQueue = testing.MakeClusterQueue("clusterQueue_queue-controller").
+		clusterQueue = testing.MakeClusterQueue("cluster-queue.queue-controller").
 			Resource(testing.MakeResource(resourceGPU).
 				Flavor(testing.MakeFlavor(flavorModelA, "5").Max("10").Obj()).
 				Flavor(testing.MakeFlavor(flavorModelB, "5").Max("10").Obj()).Obj()).Obj()
+		gomega.Expect(k8sClient.Create(ctx, clusterQueue)).To(gomega.Succeed())
 		queue = testing.MakeQueue("queue", ns.Name).ClusterQueue(clusterQueue.Name).Obj()
 		gomega.Expect(k8sClient.Create(ctx, queue)).To(gomega.Succeed())
 	})