fix incorrect quota management when lendingLimit enabled in preemption (

kubernetes-sigs#1770)

Signed-off-by: B1F030 <b1fzhang@gmail.com>
Co-authored-by: B1F030 <b1fzhang@gmail.com>

pkg/cache/clusterqueue.go

@@ -484,6 +484,29 @@ func updateUsage(wi *workload.Info, flvUsage FlavorResourceQuantities, m int64)
 	}
 }
 
+func updateCohortUsage(wi *workload.Info, cq *ClusterQueue, m int64) {
+	for _, ps := range wi.TotalRequests {
+		for wlRes, wlResFlv := range ps.Flavors {
+			v, wlResExist := ps.Requests[wlRes]
+			flv, flvExist := cq.Cohort.Usage[wlResFlv]
+			if flvExist && wlResExist {
+				if _, exists := flv[wlRes]; exists {
+					after := cq.Usage[wlResFlv][wlRes] - cq.guaranteedQuota(wlResFlv, wlRes)
+					// rollback update cq.Usage
+					before := after - v*m
+					if before > 0 {
+						flv[wlRes] -= before
+					}
+					// simulate updating cq.Usage
+					if after > 0 {
+						flv[wlRes] += after
+					}
+				}
+			}
+		}
+	}
+}
+
 func (c *ClusterQueue) addLocalQueue(q *kueue.LocalQueue) error {
 	qKey := queueKey(q)
 	if _, ok := c.localQueues[qKey]; ok {

pkg/cache/snapshot.go

@@ -43,18 +43,26 @@ func (s *Snapshot) RemoveWorkload(wl *workload.Info) {
 	delete(cq.Workloads, workload.Key(wl.Obj))
 	updateUsage(wl, cq.Usage, -1)
 	if cq.Cohort != nil {
-		updateUsage(wl, cq.Cohort.Usage, -1)
+		if features.Enabled(features.LendingLimit) {
+			updateCohortUsage(wl, cq, -1)
+		} else {
+			updateUsage(wl, cq.Cohort.Usage, -1)
+		}
 	}
 }
 
-// AddWorkload removes a workload from its corresponding ClusterQueue and
+// AddWorkload adds a workload from its corresponding ClusterQueue and
 // updates resource usage.
 func (s *Snapshot) AddWorkload(wl *workload.Info) {
 	cq := s.ClusterQueues[wl.ClusterQueue]
 	cq.Workloads[workload.Key(wl.Obj)] = wl
 	updateUsage(wl, cq.Usage, 1)
 	if cq.Cohort != nil {
-		updateUsage(wl, cq.Cohort.Usage, 1)
+		if features.Enabled(features.LendingLimit) {
+			updateCohortUsage(wl, cq, 1)
+		} else {
+			updateUsage(wl, cq.Cohort.Usage, 1)
+		}
 	}
 }