From 4834b19be615034361b4bea6b80de443114d148b Mon Sep 17 00:00:00 2001
From: Peter Bacsko <pbacsko@cloudera.com>
Date: Mon, 24 Jun 2024 15:19:31 +0200
Subject: [PATCH] [YUNIKORN-2661] Fix hard-coded boolean in setLimit (#883)

Closes: #883

Signed-off-by: Peter Bacsko <pbacsko@cloudera.com>
(cherry picked from commit 1aafed0f3444acaeaee0b05d6c98620e88eb0d95)
---
 pkg/scheduler/ugm/queue_tracker.go      |  2 +-
 pkg/scheduler/ugm/queue_tracker_test.go | 55 +++++++++++++++++++++++++
 pkg/scheduler/ugm/user_tracker_test.go  |  4 +-
 3 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/pkg/scheduler/ugm/queue_tracker.go b/pkg/scheduler/ugm/queue_tracker.go
index 830610fa6..aea321c9e 100644
--- a/pkg/scheduler/ugm/queue_tracker.go
+++ b/pkg/scheduler/ugm/queue_tracker.go
@@ -188,7 +188,7 @@ func (qt *QueueTracker) setLimit(hierarchy []string, maxResource *resources.Reso
 		if qt.childQueueTrackers[childName] == nil {
 			qt.childQueueTrackers[childName] = newQueueTracker(qt.queuePath, childName, trackType)
 		}
-		qt.childQueueTrackers[childName].setLimit(hierarchy[1:], maxResource, maxApps, useWildCard, trackType, false)
+		qt.childQueueTrackers[childName].setLimit(hierarchy[1:], maxResource, maxApps, useWildCard, trackType, doWildCardCheck)
 	} else if len(hierarchy) == 1 {
 		// don't override named user/group specific limits with wild card limits
 		if doWildCardCheck && !qt.useWildCard {
diff --git a/pkg/scheduler/ugm/queue_tracker_test.go b/pkg/scheduler/ugm/queue_tracker_test.go
index 929b77fd0..90f312496 100644
--- a/pkg/scheduler/ugm/queue_tracker_test.go
+++ b/pkg/scheduler/ugm/queue_tracker_test.go
@@ -346,6 +346,61 @@ func TestNewQueueTracker(t *testing.T) {
 	assert.Assert(t, resources.IsZero(parent.resourceUsage))
 }
 
+func TestSetLimit(t *testing.T) {
+	manager := GetUserManager()
+	defer manager.ClearConfigLimits()
+	manager.userWildCardLimitsConfig = map[string]*LimitConfig{
+		path1: {
+			maxApplications: 3,
+			maxResources:    resources.NewResource(),
+		},
+	}
+	root := newRootQueueTracker(user)
+	assert.Assert(t, !root.useWildCard)
+	assert.Equal(t, uint64(0), root.maxRunningApps)
+	assert.Assert(t, root.maxResources == nil)
+
+	// create tracker hierarchy
+	limit := resources.NewResourceFromMap(map[string]resources.Quantity{
+		"mem":   10,
+		"vcore": 10})
+	root.setLimit(strings.Split(queuePath1, configs.DOT), limit.Clone(), 9, true, user, true)
+
+	// check settings
+	parentQ := root.childQueueTrackers["parent"]
+	assert.Assert(t, parentQ.maxResources == nil)
+	assert.Equal(t, uint64(0), parentQ.maxRunningApps)
+	childQ := parentQ.childQueueTrackers["child1"]
+	assert.Assert(t, parentQ.maxResources == nil)
+	assert.Equal(t, uint64(9), childQ.maxRunningApps)
+	assert.Assert(t, resources.Equals(limit, childQ.maxResources))
+
+	// check if settings are overridden
+	newLimit := resources.NewResourceFromMap(map[string]resources.Quantity{
+		"mem":   20,
+		"vcore": 20})
+	root.setLimit(strings.Split(queuePath1, configs.DOT), newLimit.Clone(), 3, false, user, true) // override
+	assert.Assert(t, resources.Equals(newLimit, childQ.maxResources))
+	assert.Assert(t, !childQ.useWildCard)
+	newLimit2 := resources.NewResourceFromMap(map[string]resources.Quantity{
+		"mem":   30,
+		"vcore": 30})
+	root.setLimit(strings.Split(queuePath1, configs.DOT), newLimit2.Clone(), 2, true, user, true) // no override
+	assert.Assert(t, !childQ.useWildCard)
+	assert.Assert(t, resources.Equals(newLimit, childQ.maxResources))
+	assert.Equal(t, uint64(3), childQ.maxRunningApps)
+
+	root.setLimit(strings.Split(queuePath1, configs.DOT), newLimit2.Clone(), 4, true, user, false) // override -> changes qt.doWildCardCheck
+	assert.Assert(t, childQ.useWildCard)
+	assert.Assert(t, resources.Equals(newLimit2, childQ.maxResources))
+	assert.Equal(t, uint64(4), childQ.maxRunningApps)
+
+	root.setLimit(strings.Split(queuePath1, configs.DOT), newLimit.Clone(), 5, false, user, false) // override
+	assert.Assert(t, !childQ.useWildCard)
+	assert.Assert(t, resources.Equals(newLimit, childQ.maxResources))
+	assert.Equal(t, uint64(5), childQ.maxRunningApps)
+}
+
 func getQTResource(qt *QueueTracker) map[string]*resources.Resource {
 	resources := make(map[string]*resources.Resource)
 	usage := qt.getResourceUsageDAOInfo("")
diff --git a/pkg/scheduler/ugm/user_tracker_test.go b/pkg/scheduler/ugm/user_tracker_test.go
index 28db413fc..ae0d6a885 100644
--- a/pkg/scheduler/ugm/user_tracker_test.go
+++ b/pkg/scheduler/ugm/user_tracker_test.go
@@ -285,12 +285,12 @@ func TestSetAndClearMaxLimits(t *testing.T) {
 
 	// clear limits
 	eventSystem.Reset()
-	userTracker.clearLimits(path1, true)
+	userTracker.clearLimits(path1, false)
 	assert.Assert(t, resources.Equals(userTracker.headroom(hierarchy1), lowerParentHeadroom))
 	assert.Assert(t, resources.Equals(userTracker.headroom(hierarchy5), lowerParentHeadroom))
 	assert.Assert(t, !userTracker.canRunApp(hierarchy1, TestApp4))
 	assert.Assert(t, !userTracker.canRunApp(hierarchy5, TestApp4))
-	userTracker.clearLimits(path5, true)
+	userTracker.clearLimits(path5, false)
 	assert.Assert(t, userTracker.headroom(hierarchy1) == nil)
 	assert.Assert(t, userTracker.headroom(hierarchy5) == nil)
 	assert.Assert(t, userTracker.canRunApp(hierarchy1, TestApp4))