[YUNIKORN-2024] clear limits only upon the right queue (apache#669)

Closes: apache#669 Signed-off-by: Manikandan R <[email protected]>
doupache · Nov 3, 2023 · 0a05068 · 0a05068
1 parent dee60d1
commit 0a05068
Showing 1 changed file with 97 additions and 0 deletions.
diff --git a/pkg/scheduler/ugm/manager_test.go b/pkg/scheduler/ugm/manager_test.go
@@ -1070,6 +1070,103 @@ func TestUserGroupLimit(t *testing.T) { //nolint:funlen
 	}
 }
 
+func TestUserGroupLimitChange(t *testing.T) { //nolint:funlen
+	testCases := []struct {
+		name      string
+		user      security.UserGroup
+		limits    []configs.Limit
+		newLimits []configs.Limit
+	}{
+		{
+			name: "maxresources with an updated specific group limit",
+			user: security.UserGroup{User: "user1", Groups: []string{"group1"}},
+			limits: []configs.Limit{
+				createLimit(nil, []string{"group1"}, largeResource, 2),
+			},
+			newLimits: []configs.Limit{
+				createLimit(nil, []string{"group1"}, mediumResource, 2),
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			setupUGM()
+
+			manager := GetUserManager()
+			conf := createConfigWithLimits(tc.limits)
+
+			assert.NilError(t, manager.UpdateConfig(conf.Queues[0], "root"))
+
+			usage, err := resources.NewResourceFromConf(mediumResource)
+			if err != nil {
+				t.Errorf("new resource create returned error or wrong resource: error %t, res %v", err, usage)
+			}
+
+			increased := manager.IncreaseTrackedResource(queuePathParent, TestApp1, usage, tc.user)
+			assert.Equal(t, increased, true, "unable to increase tracked resource: queuepath "+queuePathParent+", app "+TestApp1+", res "+usage.String())
+
+			increased = manager.IncreaseTrackedResource(queuePathParent, TestApp2, usage, tc.user)
+			assert.Equal(t, increased, true, "unable to increase tracked resource: queuepath "+queuePathParent+", app "+TestApp2+", res "+usage.String())
+
+			decreased := manager.DecreaseTrackedResource(queuePathParent, TestApp2, usage, tc.user, true)
+			assert.Equal(t, decreased, true, "unable to decreased tracked resource: queuepath "+queuePathParent+", app "+TestApp2+", res "+usage.String())
+
+			conf.Queues[0].Queues[0].Limits = tc.newLimits
+			assert.NilError(t, manager.UpdateConfig(conf.Queues[0], "root"))
+
+			increased = manager.IncreaseTrackedResource(queuePathParent, TestApp2, usage, tc.user)
+			assert.Equal(t, increased, false, "should not increase tracked resource: queuepath "+queuePathParent+", app "+TestApp2+", res "+usage.String())
+		})
+	}
+}
+
+func TestMultipleGroupLimitChange(t *testing.T) {
+	setupUGM()
+
+	manager := GetUserManager()
+	conf := createConfigWithLimits([]configs.Limit{
+		createLimit(nil, []string{"group1", "group2"}, largeResource, 2),
+		createLimit(nil, []string{"*"}, mediumResource, 1),
+	})
+	assert.NilError(t, manager.UpdateConfig(conf.Queues[0], "root"))
+
+	user1 := security.UserGroup{User: "user1", Groups: []string{"group1"}}
+	user2 := security.UserGroup{User: "user2", Groups: []string{"group2"}}
+	user3 := security.UserGroup{User: "user3", Groups: []string{"group3"}}
+
+	usage, err := resources.NewResourceFromConf(mediumResource)
+	if err != nil {
+		t.Errorf("new resource create returned error or wrong resource: error %t, res %v", err, usage)
+	}
+
+	// all users can increate usage within the quota
+	increased := manager.IncreaseTrackedResource(queuePathParent, "test-app-1-1", usage, user1)
+	assert.Equal(t, increased, true, "unable to increase tracked resource: queuepath "+queuePathParent+", app test-app-1-1, res "+usage.String())
+
+	increased = manager.IncreaseTrackedResource(queuePathParent, "test-app-2-1", usage, user2)
+	assert.Equal(t, increased, true, "unable to increase tracked resource: queuepath "+queuePathParent+", app test-app-2-1, res "+usage.String())
+
+	increased = manager.IncreaseTrackedResource(queuePathParent, "test-app-3-1", usage, user3)
+	assert.Equal(t, increased, true, "unable to increase tracked resource: queuepath "+queuePathParent+", app test-app-3-1, res "+usage.String())
+
+	// remove group2 from the specific group
+	conf.Queues[0].Queues[0].Limits[0].Groups = []string{"group1"}
+	assert.NilError(t, manager.UpdateConfig(conf.Queues[0], "root"))
+
+	// user1 still can increase usage within the quota
+	increased = manager.IncreaseTrackedResource(queuePathParent, "test-app-1-2", usage, user1)
+	assert.Equal(t, increased, true, "unable to increase tracked resource: queuepath "+queuePathParent+", app test-app-1-2, res "+usage.String())
+
+	// user2 can't increase usage more than wildcard limit
+	increased = manager.IncreaseTrackedResource(queuePathParent, "test-app-2-2", usage, user2)
+	assert.Equal(t, increased, false, "should not increase tracked resource: queuepath "+queuePathParent+", app test-app-2-2, res "+usage.String())
+
+	// user3 can't increase usage more than wildcard limit
+	increased = manager.IncreaseTrackedResource(queuePathParent, "test-app-3-2", usage, user3)
+	assert.Equal(t, increased, false, "should not increase tracked resource: queuepath "+queuePathParent+", app test-app-3-2, res "+usage.String())
+}
+
 func createLimit(users, groups []string, maxResources map[string]string, maxApps uint64) configs.Limit {
 	return configs.Limit{
 		Users:           users,