From d7d1408c6395d2d4d058cbe22a7f15205a28e69a Mon Sep 17 00:00:00 2001 From: Wilfred Spiegelenburg Date: Wed, 13 Nov 2024 14:09:23 -0600 Subject: [PATCH] [YUNIKORN-2972] Remove internal object from user and group REST info (#993) Remove the Resource object from the REST response for the user and group trackers. Removed the usage of the DAO objects in the internal unit tests. The unit tests use new functions which directly expose the same structures without using the DAO. Cleanup: - remove export from internal functions (UserTracker, GroupTracker) - renamed functions to match their results (Manager) - update handler to call renamed functions in Manager - reimplemented test functions to not use DAO objects (Manager, UserTracker, GroupTracker, QueueTracker) - assert function changes to use new functions removing DAO usage (scheduler/utilites_test, scheduler/objects/utilites_test) Closes: #993 Signed-off-by: Craig Condit --- pkg/scheduler/objects/utilities_test.go | 24 ++++--- pkg/scheduler/ugm/group_tracker.go | 65 ++++++++++++----- pkg/scheduler/ugm/group_tracker_test.go | 13 +--- pkg/scheduler/ugm/manager.go | 38 +++++----- pkg/scheduler/ugm/manager_test.go | 47 ++++++------ pkg/scheduler/ugm/queue_tracker.go | 96 ++++++++++++++++++------- pkg/scheduler/ugm/queue_tracker_test.go | 34 ++++----- pkg/scheduler/ugm/user_tracker.go | 65 ++++++++++++----- pkg/scheduler/ugm/user_tracker_test.go | 14 ++-- pkg/scheduler/ugm/utilities_test.go | 13 ---- pkg/scheduler/utilities_test.go | 48 ++++--------- pkg/webservice/dao/ugm_info.go | 6 +- pkg/webservice/handlers.go | 24 +++---- pkg/webservice/handlers_test.go | 18 ++--- 14 files changed, 286 insertions(+), 219 deletions(-) diff --git a/pkg/scheduler/objects/utilities_test.go b/pkg/scheduler/objects/utilities_test.go index f1369a7ea..97be13c24 100644 --- a/pkg/scheduler/objects/utilities_test.go +++ b/pkg/scheduler/objects/utilities_test.go @@ -284,19 +284,23 @@ func getUserGroup(userName string, groupNameList []string) security.UserGroup { } func assertUserGroupResource(t *testing.T, userGroup security.UserGroup, expected *resources.Resource) { - ugm := ugm.GetUserManager() - userResource := ugm.GetUserResources(userGroup) - groupResource := ugm.GetGroupResources(userGroup.Groups[0]) - assert.Equal(t, resources.Equals(userResource, expected), true) - assert.Equal(t, resources.Equals(groupResource, nil), true) + assertUserResourcesAndGroupResources(t, userGroup, expected, nil, 0) } func assertUserResourcesAndGroupResources(t *testing.T, userGroup security.UserGroup, expectedUserResources *resources.Resource, expectedGroupResources *resources.Resource, i int) { - ugm := ugm.GetUserManager() - userResource := ugm.GetUserResources(userGroup) - groupResource := ugm.GetGroupResources(userGroup.Groups[i]) - assert.Equal(t, resources.Equals(userResource, expectedUserResources), true) - assert.Equal(t, resources.Equals(groupResource, expectedGroupResources), true) + m := ugm.GetUserManager() + userResource := m.GetUserResources(userGroup.User) + if expectedUserResources == nil { + assert.Assert(t, userResource.IsEmpty(), "expected empty resource in user tracker") + } else { + assert.Assert(t, resources.Equals(userResource, expectedUserResources), "user value '%s' not equal to expected '%s'", userResource.String(), expectedUserResources.String()) + } + groupResource := m.GetGroupResources(userGroup.Groups[i]) + if expectedGroupResources == nil { + assert.Assert(t, groupResource.IsEmpty(), "expected empty resource in group tracker") + } else { + assert.Assert(t, resources.Equals(groupResource, expectedGroupResources), "group value '%s' not equal to expected '%s'", groupResource.String(), expectedGroupResources.String()) + } } func assertAllocationLog(t *testing.T, ask *Allocation) { diff --git a/pkg/scheduler/ugm/group_tracker.go b/pkg/scheduler/ugm/group_tracker.go index d6ec777fd..29fa15e68 100644 --- a/pkg/scheduler/ugm/group_tracker.go +++ b/pkg/scheduler/ugm/group_tracker.go @@ -21,7 +21,6 @@ package ugm import ( "strings" - "github.com/apache/yunikorn-core/pkg/common" "github.com/apache/yunikorn-core/pkg/common/configs" "github.com/apache/yunikorn-core/pkg/common/resources" "github.com/apache/yunikorn-core/pkg/locking" @@ -92,43 +91,46 @@ func (gt *GroupTracker) clearLimits(queuePath string) { gt.queueTracker.setLimit(strings.Split(queuePath, configs.DOT), nil, 0, false, group, false) } -// Note: headroom of queue tracker is not read-only, it also traverses the queue hierarchy and creates childQueueTracker if it does not exist. +// headroom calculate the resource headroom for the group in the hierarchy defined +// Note: headroom of queue tracker is not read-only. +// It traverses the queue hierarchy and creates a childQueueTracker if it does not exist. func (gt *GroupTracker) headroom(hierarchy []string) *resources.Resource { gt.Lock() defer gt.Unlock() return gt.queueTracker.headroom(hierarchy, group) } -func (gt *GroupTracker) GetGroupResourceUsageDAOInfo() *dao.GroupResourceUsageDAOInfo { +// GetResourceUsageDAOInfo returns the DAO object used in the REST API for this group tracker +func (gt *GroupTracker) GetResourceUsageDAOInfo() *dao.GroupResourceUsageDAOInfo { gt.RLock() defer gt.RUnlock() - groupResourceUsage := &dao.GroupResourceUsageDAOInfo{ - Applications: []string{}, - } - groupResourceUsage.GroupName = gt.groupName + var apps []string for app := range gt.applications { - groupResourceUsage.Applications = append(groupResourceUsage.Applications, app) + apps = append(apps, app) + } + return &dao.GroupResourceUsageDAOInfo{ + Applications: apps, + GroupName: gt.groupName, + Queues: gt.queueTracker.getResourceUsageDAOInfo(), } - groupResourceUsage.Queues = gt.queueTracker.getResourceUsageDAOInfo(common.Empty) - return groupResourceUsage } -func (gt *GroupTracker) IsQueuePathTrackedCompletely(hierarchy []string) bool { +func (gt *GroupTracker) isQueuePathTrackedCompletely(hierarchy []string) bool { gt.RLock() defer gt.RUnlock() - return gt.queueTracker.IsQueuePathTrackedCompletely(hierarchy) + return gt.queueTracker.isQueuePathTrackedCompletely(hierarchy) } -func (gt *GroupTracker) IsUnlinkRequired(hierarchy []string) bool { +func (gt *GroupTracker) isUnlinkRequired(hierarchy []string) bool { gt.RLock() defer gt.RUnlock() - return gt.queueTracker.IsUnlinkRequired(hierarchy) + return gt.queueTracker.isUnlinkRequired(hierarchy) } -func (gt *GroupTracker) UnlinkQT(hierarchy []string) bool { +func (gt *GroupTracker) unlinkQT(hierarchy []string) bool { gt.Lock() defer gt.Unlock() - return gt.queueTracker.UnlinkQT(hierarchy) + return gt.queueTracker.unlink(hierarchy) } func (gt *GroupTracker) canBeRemoved() bool { @@ -153,9 +155,38 @@ func (gt *GroupTracker) decreaseAllTrackedResourceUsage(hierarchy []string) map[ return removedApplications } -// Note: canRunApp of queue tracker is not read-only, it also traverses the queue hierarchy and creates a childQueueTracker if it does not exist. +// canRunApp checks if the group is allowed to run the application in the queue defined in hierarchy. +// Note: canRunApp of queue tracker is not read-only, +// It traverses the queue hierarchy and creates a childQueueTracker if it does not exist. func (gt *GroupTracker) canRunApp(hierarchy []string, applicationID string) bool { gt.Lock() defer gt.Unlock() return gt.queueTracker.canRunApp(hierarchy, applicationID, group) } + +// GetMaxResources returns a map of the maxResources for all queues registered under this group tracker. +// The key into the map is the queue path. +// This should only be used in test +func (gt *GroupTracker) GetMaxResources() map[string]*resources.Resource { + gt.RLock() + defer gt.RUnlock() + return gt.queueTracker.getMaxResources() +} + +// GetMaxApplications returns a map of the maxRunningApps for all queues registered under this group tracker. +// The key into the map is the queue path. +// This should only be used in test +func (gt *GroupTracker) GetMaxApplications() map[string]uint64 { + gt.RLock() + defer gt.RUnlock() + return gt.queueTracker.getMaxApplications() +} + +// getUsedResources returns a map of the usedResources for all queues registered under this group tracker. +// The key into the map is the queue path. +// This should only be used in test +func (gt *GroupTracker) getUsedResources() map[string]*resources.Resource { + gt.RLock() + defer gt.RUnlock() + return gt.queueTracker.getUsedResources() +} diff --git a/pkg/scheduler/ugm/group_tracker_test.go b/pkg/scheduler/ugm/group_tracker_test.go index 245e76131..833be072f 100644 --- a/pkg/scheduler/ugm/group_tracker_test.go +++ b/pkg/scheduler/ugm/group_tracker_test.go @@ -69,8 +69,8 @@ func TestGTIncreaseTrackedResource(t *testing.T) { t.Errorf("new resource create returned error or wrong resource: error %t, res %v", err, usage3) } groupTracker.increaseTrackedResource(path4, TestApp4, usage4, user.User) - actualResources := getGroupResource(groupTracker) + actualResources := groupTracker.queueTracker.getUsedResources() assert.Equal(t, "map[mem:80000000 vcore:80000]", actualResources["root"].String(), "wrong resource") assert.Equal(t, "map[mem:80000000 vcore:80000]", actualResources["root.parent"].String(), "wrong resource") assert.Equal(t, "map[mem:40000000 vcore:40000]", actualResources["root.parent.child1"].String(), "wrong resource") @@ -104,9 +104,9 @@ func TestGTDecreaseTrackedResource(t *testing.T) { t.Errorf("new resource create returned error or wrong resource: error %t, res %v", err, usage2) } groupTracker.increaseTrackedResource(path2, TestApp2, usage2, user.User) - actualResources := getGroupResource(groupTracker) assert.Equal(t, 2, len(groupTracker.getTrackedApplications())) + actualResources := groupTracker.getUsedResources() assert.Equal(t, "map[mem:90000000 vcore:90000]", actualResources["root"].String(), "wrong resource") assert.Equal(t, "map[mem:90000000 vcore:90000]", actualResources["root.parent"].String(), "wrong resource") assert.Equal(t, "map[mem:70000000 vcore:70000]", actualResources["root.parent.child1"].String(), "wrong resource") @@ -126,8 +126,7 @@ func TestGTDecreaseTrackedResource(t *testing.T) { removeQT = groupTracker.decreaseTrackedResource(path2, TestApp2, usage3, false) assert.Equal(t, removeQT, false, "wrong remove queue tracker value") - actualResources1 := getGroupResource(groupTracker) - + actualResources1 := groupTracker.getUsedResources() assert.Equal(t, "map[mem:70000000 vcore:70000]", actualResources1["root"].String(), "wrong resource") assert.Equal(t, "map[mem:70000000 vcore:70000]", actualResources1["root.parent"].String(), "wrong resource") assert.Equal(t, "map[mem:60000000 vcore:60000]", actualResources1["root.parent.child1"].String(), "wrong resource") @@ -282,9 +281,3 @@ func TestGTCanRunApp(t *testing.T) { assert.Assert(t, groupTracker.canRunApp(hierarchy1, TestApp1)) assert.Assert(t, !groupTracker.canRunApp(hierarchy1, TestApp2)) } - -func getGroupResource(gt *GroupTracker) map[string]*resources.Resource { - resources := make(map[string]*resources.Resource) - usage := gt.GetGroupResourceUsageDAOInfo() - return internalGetResource(usage.Queues, resources) -} diff --git a/pkg/scheduler/ugm/manager.go b/pkg/scheduler/ugm/manager.go index 517751a16..e235f36b0 100644 --- a/pkg/scheduler/ugm/manager.go +++ b/pkg/scheduler/ugm/manager.go @@ -188,7 +188,7 @@ func (m *Manager) DecreaseTrackedResource(queuePath, applicationID string, usage } } -func (m *Manager) GetUsersResources() []*UserTracker { +func (m *Manager) GetUserTrackers() []*UserTracker { m.RLock() defer m.RUnlock() var userTrackers []*UserTracker @@ -205,7 +205,7 @@ func (m *Manager) GetUserTracker(user string) *UserTracker { return m.userTrackers[user] } -func (m *Manager) GetGroupsResources() []*GroupTracker { +func (m *Manager) GetGroupTrackers() []*GroupTracker { m.RLock() defer m.RUnlock() var groupTrackers []*GroupTracker @@ -496,14 +496,14 @@ func (m *Manager) clearEarlierSetUserLimits(newUserLimits map[string]map[string] func (m *Manager) resetUserEarlierUsage(ut *UserTracker, queuePath string) { // Is this user already tracked for the queue path? hierarchy := strings.Split(queuePath, configs.DOT) - if ut.IsQueuePathTrackedCompletely(hierarchy) { + if ut.isQueuePathTrackedCompletely(hierarchy) { log.Log(log.SchedUGM).Debug("Need to clear earlier set configs for user", zap.String("user", ut.userName), zap.Strings("queue path", hierarchy)) ut.clearLimits(queuePath, false) // Is there any running applications in end queue of this queue path? If not, then remove the linkage between end queue and its immediate parent - if ut.IsUnlinkRequired(hierarchy) { - ut.UnlinkQT(hierarchy) + if ut.isUnlinkRequired(hierarchy) { + ut.unlinkQT(hierarchy) } log.Log(log.SchedUGM).Debug("Cleared earlier set limit configs for user", zap.String("user", ut.userName), @@ -544,7 +544,7 @@ func (m *Manager) clearEarlierSetGroupLimits(newGroupLimits map[string]map[strin // eventually remove group tracker object itself from ugm if it can be removed. func (m *Manager) resetGroupEarlierUsage(gt *GroupTracker, queuePath string) { hierarchy := strings.Split(queuePath, configs.DOT) - if gt.IsQueuePathTrackedCompletely(hierarchy) { + if gt.isQueuePathTrackedCompletely(hierarchy) { log.Log(log.SchedUGM).Debug("Need to clear earlier set configs for group", zap.String("group", gt.groupName), zap.Strings("queue path", hierarchy)) @@ -555,8 +555,8 @@ func (m *Manager) resetGroupEarlierUsage(gt *GroupTracker, queuePath string) { } gt.clearLimits(queuePath) // Is there any running applications in end queue of this queue path? If not, then remove the linkage between end queue and its immediate parent - if gt.IsUnlinkRequired(hierarchy) { - gt.UnlinkQT(hierarchy) + if gt.isUnlinkRequired(hierarchy) { + gt.unlinkQT(hierarchy) } log.Log(log.SchedUGM).Debug("Cleared earlier set limit configs for group", zap.String("group", gt.groupName), @@ -710,24 +710,26 @@ func (m *Manager) ClearConfigLimits() { m.groupLimits = make(map[string]map[string]*LimitConfig) } -// GetUserResources only for tests -func (m *Manager) GetUserResources(user security.UserGroup) *resources.Resource { +// GetUserResources returns the root queue maxResources for the user +// Should only be used in tests +func (m *Manager) GetUserResources(user string) *resources.Resource { m.RLock() defer m.RUnlock() - ut := m.userTrackers[user.User] - if ut != nil && ut.GetUserResourceUsageDAOInfo().Queues.ResourceUsage != nil && len(ut.GetUserResourceUsageDAOInfo().Queues.ResourceUsage.Resources) > 0 { - return ut.GetUserResourceUsageDAOInfo().Queues.ResourceUsage + ut := m.userTrackers[user] + if ut == nil { + return nil } - return nil + return ut.queueTracker.resourceUsage.Clone() } -// GetGroupResources only for tests +// GetGroupResources returns the root queue maxResources +// Should only be used in tests func (m *Manager) GetGroupResources(group string) *resources.Resource { m.RLock() defer m.RUnlock() gt := m.groupTrackers[group] - if gt != nil && gt.GetGroupResourceUsageDAOInfo().Queues.ResourceUsage != nil && len(gt.GetGroupResourceUsageDAOInfo().Queues.ResourceUsage.Resources) > 0 { - return gt.GetGroupResourceUsageDAOInfo().Queues.ResourceUsage + if gt == nil { + return nil } - return nil + return gt.queueTracker.resourceUsage.Clone() } diff --git a/pkg/scheduler/ugm/manager_test.go b/pkg/scheduler/ugm/manager_test.go index 91d702570..4edc6ce49 100644 --- a/pkg/scheduler/ugm/manager_test.go +++ b/pkg/scheduler/ugm/manager_test.go @@ -194,7 +194,7 @@ func TestAddRemoveUserAndGroups(t *testing.T) { manager.IncreaseTrackedResource("", "", usage1, user) manager.IncreaseTrackedResource(queuePath1, TestApp1, usage1, user) - groupTrackers := manager.GetGroupsResources() + groupTrackers := manager.GetGroupTrackers() assert.Equal(t, len(groupTrackers), 0) assertUGM(t, user, usage1, 1) assert.Equal(t, user.User, manager.GetUserTracker(user.User).userName) @@ -213,19 +213,19 @@ func TestAddRemoveUserAndGroups(t *testing.T) { assert.Equal(t, user.User, manager.GetUserTracker(user.User).userName) assert.Equal(t, user1.User, manager.GetUserTracker(user1.User).userName) - assert.Equal(t, true, manager.GetUserTracker(user.User).queueTracker.IsQueuePathTrackedCompletely(strings.Split(queuePath1, configs.DOT))) - assert.Equal(t, true, manager.GetUserTracker(user1.User).queueTracker.IsQueuePathTrackedCompletely(strings.Split(queuePath2, configs.DOT))) - assert.Equal(t, false, manager.GetUserTracker(user1.User).queueTracker.IsQueuePathTrackedCompletely(strings.Split(queuePath1, configs.DOT))) - assert.Equal(t, false, manager.GetUserTracker(user.User).queueTracker.IsQueuePathTrackedCompletely(strings.Split(queuePath2, configs.DOT))) - assert.Equal(t, false, manager.GetUserTracker(user.User).queueTracker.IsQueuePathTrackedCompletely(strings.Split(queuePath3, configs.DOT))) - assert.Equal(t, false, manager.GetUserTracker(user.User).queueTracker.IsQueuePathTrackedCompletely(strings.Split(queuePath4, configs.DOT))) + assert.Equal(t, true, manager.GetUserTracker(user.User).queueTracker.isQueuePathTrackedCompletely(strings.Split(queuePath1, configs.DOT))) + assert.Equal(t, true, manager.GetUserTracker(user1.User).queueTracker.isQueuePathTrackedCompletely(strings.Split(queuePath2, configs.DOT))) + assert.Equal(t, false, manager.GetUserTracker(user1.User).queueTracker.isQueuePathTrackedCompletely(strings.Split(queuePath1, configs.DOT))) + assert.Equal(t, false, manager.GetUserTracker(user.User).queueTracker.isQueuePathTrackedCompletely(strings.Split(queuePath2, configs.DOT))) + assert.Equal(t, false, manager.GetUserTracker(user.User).queueTracker.isQueuePathTrackedCompletely(strings.Split(queuePath3, configs.DOT))) + assert.Equal(t, false, manager.GetUserTracker(user.User).queueTracker.isQueuePathTrackedCompletely(strings.Split(queuePath4, configs.DOT))) - assert.Equal(t, true, manager.GetUserTracker(user.Groups[0]).queueTracker.IsQueuePathTrackedCompletely(strings.Split(queuePath1, configs.DOT))) - assert.Equal(t, true, manager.GetUserTracker(user1.Groups[0]).queueTracker.IsQueuePathTrackedCompletely(strings.Split(queuePath2, configs.DOT))) - assert.Equal(t, false, manager.GetUserTracker(user1.Groups[0]).queueTracker.IsQueuePathTrackedCompletely(strings.Split(queuePath1, configs.DOT))) - assert.Equal(t, false, manager.GetUserTracker(user.Groups[0]).queueTracker.IsQueuePathTrackedCompletely(strings.Split(queuePath2, configs.DOT))) - assert.Equal(t, false, manager.GetUserTracker(user.Groups[0]).queueTracker.IsQueuePathTrackedCompletely(strings.Split(queuePath3, configs.DOT))) - assert.Equal(t, false, manager.GetUserTracker(user.Groups[0]).queueTracker.IsQueuePathTrackedCompletely(strings.Split(queuePath4, configs.DOT))) + assert.Equal(t, true, manager.GetUserTracker(user.Groups[0]).queueTracker.isQueuePathTrackedCompletely(strings.Split(queuePath1, configs.DOT))) + assert.Equal(t, true, manager.GetUserTracker(user1.Groups[0]).queueTracker.isQueuePathTrackedCompletely(strings.Split(queuePath2, configs.DOT))) + assert.Equal(t, false, manager.GetUserTracker(user1.Groups[0]).queueTracker.isQueuePathTrackedCompletely(strings.Split(queuePath1, configs.DOT))) + assert.Equal(t, false, manager.GetUserTracker(user.Groups[0]).queueTracker.isQueuePathTrackedCompletely(strings.Split(queuePath2, configs.DOT))) + assert.Equal(t, false, manager.GetUserTracker(user.Groups[0]).queueTracker.isQueuePathTrackedCompletely(strings.Split(queuePath3, configs.DOT))) + assert.Equal(t, false, manager.GetUserTracker(user.Groups[0]).queueTracker.isQueuePathTrackedCompletely(strings.Split(queuePath4, configs.DOT))) usage3, err := resources.NewResourceFromConf(map[string]string{"mem": "5M", "vcore": "5"}) if err != nil { @@ -237,12 +237,12 @@ func TestAddRemoveUserAndGroups(t *testing.T) { assertUGM(t, user, usage1, 2) manager.DecreaseTrackedResource(queuePath1, TestApp1, usage3, user, true) - assert.Equal(t, 1, len(manager.GetUsersResources()), "userTrackers count should be 1") - assert.Equal(t, 0, len(manager.GetGroupsResources()), "groupTrackers count should be 0") + assert.Equal(t, 1, len(manager.GetUserTrackers()), "userTrackers count should be 1") + assert.Equal(t, 0, len(manager.GetGroupTrackers()), "groupTrackers count should be 0") manager.DecreaseTrackedResource(queuePath2, TestApp2, usage2, user1, true) - assert.Equal(t, 0, len(manager.GetUsersResources()), "userTrackers count should be 0") - assert.Equal(t, 0, len(manager.GetGroupsResources()), "groupTrackers count should be 0") + assert.Equal(t, 0, len(manager.GetUserTrackers()), "userTrackers count should be 0") + assert.Equal(t, 0, len(manager.GetGroupTrackers()), "groupTrackers count should be 0") assert.Assert(t, manager.GetUserTracker(user.User) == nil) assert.Assert(t, manager.GetGroupTracker(user.Groups[0]) == nil) @@ -1979,12 +1979,13 @@ func setupUGM() { func assertUGM(t *testing.T, userGroup security.UserGroup, expected *resources.Resource, usersCount int) { manager := GetUserManager() - assert.Equal(t, usersCount, len(manager.GetUsersResources()), "userTrackers count should be "+strconv.Itoa(usersCount)) - assert.Equal(t, 0, len(manager.GetGroupsResources()), "groupTrackers count should be "+strconv.Itoa(0)) - userRes := manager.GetUserResources(userGroup) - assert.Equal(t, resources.Equals(userRes, expected), true) - groupRes := manager.GetGroupResources(userGroup.Groups[0]) - assert.Equal(t, resources.Equals(groupRes, nil), true) + assert.Equal(t, usersCount, len(manager.GetUserTrackers()), "userTrackers count not as expected") + assert.Equal(t, 0, len(manager.GetGroupTrackers()), "groupTrackers count should be 0") + userTR := manager.GetUserTracker(userGroup.User) + assert.Assert(t, userTR != nil, "user tracker should be defined") + assert.Assert(t, resources.Equals(userTR.queueTracker.resourceUsage, expected), "user max resource for root not correct") + groupTR := manager.GetGroupTracker(userGroup.Groups[0]) + assert.Assert(t, groupTR == nil, "group tracker should not be defined") } func assertMaxLimits(t *testing.T, userGroup security.UserGroup, expectedResource *resources.Resource, expectedMaxApps int) { diff --git a/pkg/scheduler/ugm/queue_tracker.go b/pkg/scheduler/ugm/queue_tracker.go index 36ddba259..c53519477 100644 --- a/pkg/scheduler/ugm/queue_tracker.go +++ b/pkg/scheduler/ugm/queue_tracker.go @@ -20,6 +20,7 @@ package ugm import ( "go.uber.org/zap" + "golang.org/x/exp/maps" "github.com/apache/yunikorn-core/pkg/common" "github.com/apache/yunikorn-core/pkg/common/configs" @@ -230,41 +231,86 @@ func (qt *QueueTracker) headroom(hierarchy []string, trackType trackingType) *re return resources.ComponentWiseMin(headroom, childHeadroom) } +// getResourceUsageDAOInfo returns the REST representation of the queue tracker // Note: Lock free call. The RLock of the linked tracker (UserTracker and GroupTracker) should be held before calling this function. -func (qt *QueueTracker) getResourceUsageDAOInfo(parentQueuePath string) *dao.ResourceUsageDAOInfo { +func (qt *QueueTracker) getResourceUsageDAOInfo() *dao.ResourceUsageDAOInfo { if qt == nil { return &dao.ResourceUsageDAOInfo{} } - fullQueuePath := parentQueuePath + "." + qt.queueName - if parentQueuePath == common.Empty { - fullQueuePath = qt.queueName + apps := make([]string, len(qt.runningApplications)) + i := 0 + for app := range qt.runningApplications { + apps[i] = app + i++ } - usage := &dao.ResourceUsageDAOInfo{ - QueuePath: fullQueuePath, - ResourceUsage: qt.resourceUsage.Clone(), + children := make([]*dao.ResourceUsageDAOInfo, len(qt.childQueueTrackers)) + i = 0 + for _, cqt := range qt.childQueueTrackers { + children[i] = cqt.getResourceUsageDAOInfo() + i++ } - for app := range qt.runningApplications { - usage.RunningApplications = append(usage.RunningApplications, app) + return &dao.ResourceUsageDAOInfo{ + QueuePath: qt.queuePath, + ResourceUsage: qt.resourceUsage.DAOMap(), + MaxResources: qt.maxResources.DAOMap(), + MaxApplications: qt.maxRunningApps, + RunningApplications: apps, + Children: children, + } +} + +// getMaxResources returns a map of all maxResources defined in the queue hierarchy. +// Note: Lock free call. The RLock of the linked tracker (UserTracker and GroupTracker) should be held before calling this function. +func (qt *QueueTracker) getMaxResources() map[string]*resources.Resource { + if qt == nil { + return nil + } + maxRes := map[string]*resources.Resource{qt.queuePath: qt.maxResources} + for _, cqt := range qt.childQueueTrackers { + childUsage := cqt.getMaxResources() + maps.Copy(maxRes, childUsage) + } + return maxRes +} + +// getMaxApplications returns a map of all maxRunningApps defined in the queue hierarchy. +// Note: Lock free call. The RLock of the linked tracker (UserTracker and GroupTracker) should be held before calling this function. +func (qt *QueueTracker) getMaxApplications() map[string]uint64 { + if qt == nil { + return nil + } + maxApps := map[string]uint64{qt.queuePath: qt.maxRunningApps} + for _, cqt := range qt.childQueueTrackers { + childApps := cqt.getMaxApplications() + maps.Copy(maxApps, childApps) + } + return maxApps +} + +// getUsedResources returns a map of all resourceUsage defined in the queue hierarchy. +// Note: Lock free call. The RLock of the linked tracker (UserTracker and GroupTracker) should be held before calling this function. +func (qt *QueueTracker) getUsedResources() map[string]*resources.Resource { + if qt == nil { + return nil } - usage.MaxResources = qt.maxResources - usage.MaxApplications = qt.maxRunningApps + maxRes := map[string]*resources.Resource{qt.queuePath: qt.resourceUsage} for _, cqt := range qt.childQueueTrackers { - childUsage := cqt.getResourceUsageDAOInfo(fullQueuePath) - usage.Children = append(usage.Children, childUsage) + childUsage := cqt.getUsedResources() + maps.Copy(maxRes, childUsage) } - return usage + return maxRes } -// IsQueuePathTrackedCompletely Traverse queue path upto the end queue through its linkage +// isQueuePathTrackedCompletely Traverse queue path upto the end queue through its linkage // to confirm entire queuePath has been tracked completely or not // Note: Lock free call. The RLock of the linked tracker (UserTracker and GroupTracker) should be held before calling this function. -func (qt *QueueTracker) IsQueuePathTrackedCompletely(hierarchy []string) bool { +func (qt *QueueTracker) isQueuePathTrackedCompletely(hierarchy []string) bool { // depth first: all the way to the leaf, ignore if not exists // more than 1 in the slice means we need to recurse down if len(hierarchy) > 1 { childName := hierarchy[1] if qt.childQueueTrackers[childName] != nil { - return qt.childQueueTrackers[childName].IsQueuePathTrackedCompletely(hierarchy[1:]) + return qt.childQueueTrackers[childName].isQueuePathTrackedCompletely(hierarchy[1:]) } } else if len(hierarchy) == 1 { // reach end of hierarchy @@ -275,18 +321,18 @@ func (qt *QueueTracker) IsQueuePathTrackedCompletely(hierarchy []string) bool { return false } -// IsUnlinkRequired Traverse queue path upto the leaf queue and decide whether +// isUnlinkRequired Traverse queue path upto the leaf queue and decide whether // linkage needs to be removed or not based on the running applications. // If there are any running applications in end leaf queue, we should remove the linkage between -// the leaf and its parent queue using UnlinkQT method. Otherwise, we should leave as it is. +// the leaf and its parent queue using unlink method. Otherwise, we should leave as it is. // Note: Lock free call. The RLock of the linked tracker (UserTracker and GroupTracker) should be held before calling this function. -func (qt *QueueTracker) IsUnlinkRequired(hierarchy []string) bool { +func (qt *QueueTracker) isUnlinkRequired(hierarchy []string) bool { // depth first: all the way to the leaf, ignore if not exists // more than 1 in the slice means we need to recurse down if len(hierarchy) > 1 { childName := hierarchy[1] if qt.childQueueTrackers[childName] != nil { - return qt.childQueueTrackers[childName].IsUnlinkRequired(hierarchy[1:]) + return qt.childQueueTrackers[childName].isUnlinkRequired(hierarchy[1:]) } } else if len(hierarchy) == 1 { // reach end of hierarchy @@ -302,10 +348,10 @@ func (qt *QueueTracker) IsUnlinkRequired(hierarchy []string) bool { return false } -// UnlinkQT Traverse queue path upto the end queue. If end queue has any more child queue trackers, +// unlink Traverse queue path upto the end queue. If end queue has any more child queue trackers, // then goes upto each child queue and removes the linkage with its immediate parent // Note: Lock free call. The Lock of the linked tracker (UserTracker and GroupTracker) should be held before calling this function. -func (qt *QueueTracker) UnlinkQT(hierarchy []string) bool { +func (qt *QueueTracker) unlink(hierarchy []string) bool { log.Log(log.SchedUGM).Debug("Unlinking current queue tracker from its parent", zap.String("current queue ", qt.queueName), zap.String("queue path", qt.queuePath), @@ -316,7 +362,7 @@ func (qt *QueueTracker) UnlinkQT(hierarchy []string) bool { if len(hierarchy) > 1 { childName := hierarchy[1] if qt.childQueueTrackers[childName] != nil { - if qt.childQueueTrackers[childName].UnlinkQT(hierarchy[1:]) { + if qt.childQueueTrackers[childName].unlink(hierarchy[1:]) { delete(qt.childQueueTrackers, childName) // returning false, so that it comes out when end queue detach itself from its immediate parent. // i.e., once leaf detached from root.parent for root.parent.leaf queue path. @@ -327,7 +373,7 @@ func (qt *QueueTracker) UnlinkQT(hierarchy []string) bool { } else if len(hierarchy) <= 1 { // reach end of hierarchy, unlink all queues under this queue for childName, childQT := range qt.childQueueTrackers { - if childQT.UnlinkQT([]string{childName}) { + if childQT.unlink([]string{childName}) { delete(qt.childQueueTrackers, childName) } } diff --git a/pkg/scheduler/ugm/queue_tracker_test.go b/pkg/scheduler/ugm/queue_tracker_test.go index 52f7c2164..da1eb7507 100644 --- a/pkg/scheduler/ugm/queue_tracker_test.go +++ b/pkg/scheduler/ugm/queue_tracker_test.go @@ -60,7 +60,7 @@ func TestQTIncreaseTrackedResource(t *testing.T) { t.Errorf("new resource create returned error or wrong resource: error %t, res %v", err, usage3) } queueTracker.increaseTrackedResource(strings.Split(queuePath4, configs.DOT), TestApp4, user, usage4) - actualResources := getQTResource(queueTracker) + actualResources := queueTracker.getUsedResources() assert.Equal(t, "map[mem:80000000 vcore:80000]", actualResources["root"].String(), "wrong resource") assert.Equal(t, "map[mem:80000000 vcore:80000]", actualResources["root.parent"].String(), "wrong resource") @@ -90,7 +90,7 @@ func TestQTDecreaseTrackedResource(t *testing.T) { t.Errorf("new resource create returned error or wrong resource: error %t, res %v", err, usage2) } queueTracker.increaseTrackedResource(strings.Split(queuePath2, configs.DOT), TestApp2, user, usage2) - actualResources := getQTResource(queueTracker) + actualResources := queueTracker.getUsedResources() assert.Equal(t, 2, len(queueTracker.runningApplications)) assert.Equal(t, "map[mem:90000000 vcore:90000]", actualResources["root"].String(), "wrong resource") @@ -107,7 +107,7 @@ func TestQTDecreaseTrackedResource(t *testing.T) { assert.Equal(t, removeQT, false, "wrong remove queue tracker value") removeQT = queueTracker.decreaseTrackedResource(strings.Split(queuePath2, configs.DOT), TestApp2, usage3, false) - actualResources1 := getQTResource(queueTracker) + actualResources1 := queueTracker.getUsedResources() assert.Equal(t, removeQT, false, "wrong remove queue tracker value") assert.Equal(t, "map[mem:70000000 vcore:70000]", actualResources1["root"].String(), "wrong resource") @@ -361,33 +361,39 @@ func TestGetResourceUsageDAOInfo(t *testing.T) { childQ.maxResources = maxRes.Clone() parentQ.maxRunningApps = 3 - rootDao := root.getResourceUsageDAOInfo("") - assert.Assert(t, resources.Equals(usage1, rootDao.ResourceUsage)) + rootDao := root.getResourceUsageDAOInfo() + assert.DeepEqual(t, usage1.DAOMap(), rootDao.ResourceUsage) assert.Equal(t, "root", rootDao.QueuePath) assert.Equal(t, 1, len(rootDao.RunningApplications)) assert.Equal(t, TestApp1, rootDao.RunningApplications[0]) assert.Equal(t, 1, len(rootDao.Children)) assert.Equal(t, uint64(0), rootDao.MaxApplications) - assert.Assert(t, rootDao.MaxResources == nil) + assert.Assert(t, len(rootDao.MaxResources) == 0, "expected empty max resource") parentDao := rootDao.Children[0] - assert.Assert(t, resources.Equals(usage1, parentDao.ResourceUsage)) + assert.DeepEqual(t, usage1.DAOMap(), parentDao.ResourceUsage) assert.Equal(t, "root.parent", parentDao.QueuePath) assert.Equal(t, 1, len(parentDao.RunningApplications)) assert.Equal(t, TestApp1, parentDao.RunningApplications[0]) assert.Equal(t, uint64(3), parentDao.MaxApplications) - assert.Assert(t, parentDao.MaxResources == nil) + assert.Assert(t, len(parentDao.MaxResources) == 0, "expected empty max resource") assert.Equal(t, 1, len(parentDao.Children)) childDao := parentDao.Children[0] - assert.Assert(t, resources.Equals(usage1, childDao.ResourceUsage)) + assert.DeepEqual(t, usage1.DAOMap(), childDao.ResourceUsage) assert.Equal(t, "root.parent.child1", childDao.QueuePath) assert.Equal(t, 1, len(childDao.RunningApplications)) assert.Equal(t, TestApp1, childDao.RunningApplications[0]) assert.Equal(t, uint64(2), childDao.MaxApplications) - assert.Assert(t, resources.Equals(maxRes, childDao.MaxResources)) + assert.DeepEqual(t, maxRes.DAOMap(), childDao.MaxResources) assert.Equal(t, 0, len(childDao.Children)) + // final nil check for receiver + defer func() { + if r := recover(); r != nil { + t.Fatal("getResourceUsageDAOInfo panic on nil receiver") + } + }() root = nil - rootDao = root.getResourceUsageDAOInfo("") + rootDao = root.getResourceUsageDAOInfo() assert.DeepEqual(t, rootDao, &dao.ResourceUsageDAOInfo{}) } @@ -445,9 +451,3 @@ func TestSetLimit(t *testing.T) { 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("") - return internalGetResource(usage, resources) -} diff --git a/pkg/scheduler/ugm/user_tracker.go b/pkg/scheduler/ugm/user_tracker.go index 56996464c..df7de0725 100644 --- a/pkg/scheduler/ugm/user_tracker.go +++ b/pkg/scheduler/ugm/user_tracker.go @@ -121,45 +121,49 @@ func (ut *UserTracker) clearLimits(queuePath string, doWildCardCheck bool) { ut.queueTracker.setLimit(strings.Split(queuePath, configs.DOT), nil, 0, false, user, doWildCardCheck) } -// Note: headroom of queue tracker is not read-only, it also traverses the queue hierarchy and creates childQueueTracker if it does not exist. +// headroom calculate the resource headroom for the user in the hierarchy defined +// Note: headroom of queue tracker is not read-only. +// It traverses the queue hierarchy and creates a childQueueTracker if it does not exist. func (ut *UserTracker) headroom(hierarchy []string) *resources.Resource { ut.Lock() defer ut.Unlock() return ut.queueTracker.headroom(hierarchy, user) } -func (ut *UserTracker) GetUserResourceUsageDAOInfo() *dao.UserResourceUsageDAOInfo { +// GetResourceUsageDAOInfo returns the DAO object used in the REST API for this user tracker +func (ut *UserTracker) GetResourceUsageDAOInfo() *dao.UserResourceUsageDAOInfo { ut.RLock() defer ut.RUnlock() - userResourceUsage := &dao.UserResourceUsageDAOInfo{ - Groups: make(map[string]string), - } - userResourceUsage.UserName = ut.userName + groups := make(map[string]string, len(ut.appGroupTrackers)) for app, gt := range ut.appGroupTrackers { if gt != nil { - userResourceUsage.Groups[app] = gt.groupName + groups[app] = gt.groupName } } - userResourceUsage.Queues = ut.queueTracker.getResourceUsageDAOInfo(common.Empty) - return userResourceUsage + + return &dao.UserResourceUsageDAOInfo{ + Groups: groups, + UserName: ut.userName, + Queues: ut.queueTracker.getResourceUsageDAOInfo(), + } } -func (ut *UserTracker) IsQueuePathTrackedCompletely(hierarchy []string) bool { +func (ut *UserTracker) isQueuePathTrackedCompletely(hierarchy []string) bool { ut.RLock() defer ut.RUnlock() - return ut.queueTracker.IsQueuePathTrackedCompletely(hierarchy) + return ut.queueTracker.isQueuePathTrackedCompletely(hierarchy) } -func (ut *UserTracker) IsUnlinkRequired(hierarchy []string) bool { +func (ut *UserTracker) isUnlinkRequired(hierarchy []string) bool { ut.RLock() defer ut.RUnlock() - return ut.queueTracker.IsUnlinkRequired(hierarchy) + return ut.queueTracker.isUnlinkRequired(hierarchy) } -func (ut *UserTracker) UnlinkQT(hierarchy []string) bool { +func (ut *UserTracker) unlinkQT(hierarchy []string) bool { ut.Lock() defer ut.Unlock() - return ut.queueTracker.UnlinkQT(hierarchy) + return ut.queueTracker.unlink(hierarchy) } func (ut *UserTracker) canBeRemoved() bool { @@ -168,9 +172,38 @@ func (ut *UserTracker) canBeRemoved() bool { return ut.queueTracker.canBeRemoved() } -// Note: canRunApp of queue tracker is not read-only, it also traverses the queue hierarchy and creates a childQueueTracker if it does not exist. +// canRunApp checks if the user is allowed to run the application in the queue defined in hierarchy. +// Note: canRunApp of queue tracker is not read-only. +// It traverses the queue hierarchy and creates a childQueueTracker if it does not exist. func (ut *UserTracker) canRunApp(hierarchy []string, applicationID string) bool { ut.Lock() defer ut.Unlock() return ut.queueTracker.canRunApp(hierarchy, applicationID, user) } + +// GetMaxResources returns a map of the maxResources for all queues registered under this user tracker. +// The key into the map is the queue path. +// This should only be used in test +func (ut *UserTracker) GetMaxResources() map[string]*resources.Resource { + ut.RLock() + defer ut.RUnlock() + return ut.queueTracker.getMaxResources() +} + +// GetMaxApplications returns a map of the maxRunningApps for all queues registered under this user tracker. +// The key into the map is the queue path. +// This should only be used in test +func (ut *UserTracker) GetMaxApplications() map[string]uint64 { + ut.RLock() + defer ut.RUnlock() + return ut.queueTracker.getMaxApplications() +} + +// getUsedResources returns a map of the usedResources for all queues registered under this user tracker. +// The key into the map is the queue path. +// This should only be used in test +func (ut *UserTracker) getUsedResources() map[string]*resources.Resource { + ut.RLock() + defer ut.RUnlock() + return ut.queueTracker.getUsedResources() +} diff --git a/pkg/scheduler/ugm/user_tracker_test.go b/pkg/scheduler/ugm/user_tracker_test.go index 1b860561c..1fe819012 100644 --- a/pkg/scheduler/ugm/user_tracker_test.go +++ b/pkg/scheduler/ugm/user_tracker_test.go @@ -99,7 +99,7 @@ func TestIncreaseTrackedResource(t *testing.T) { userTracker.increaseTrackedResource(path4, TestApp4, usage4) userTracker.setGroupForApp(TestApp4, groupTracker) - actualResources := getUserResource(userTracker) + actualResources := userTracker.getUsedResources() assert.Equal(t, "map[mem:80000000 vcore:80000]", actualResources["root"].String(), "wrong resource") assert.Equal(t, "map[mem:80000000 vcore:80000]", actualResources["root.parent"].String(), "wrong resource") assert.Equal(t, "map[mem:40000000 vcore:40000]", actualResources["root.parent.child1"].String(), "wrong resource") @@ -134,8 +134,8 @@ func TestDecreaseTrackedResource(t *testing.T) { } userTracker.increaseTrackedResource(path2, TestApp2, usage2) userTracker.setGroupForApp(TestApp2, groupTracker) - actualResources := getUserResource(userTracker) + actualResources := userTracker.getUsedResources() assert.Equal(t, 2, len(userTracker.getTrackedApplications())) assert.Equal(t, "map[mem:90000000 vcore:90000]", actualResources["root"].String(), "wrong resource") assert.Equal(t, "map[mem:90000000 vcore:90000]", actualResources["root.parent"].String(), "wrong resource") @@ -153,9 +153,9 @@ func TestDecreaseTrackedResource(t *testing.T) { assert.Equal(t, si.EventRecord_REMOVE, eventSystem.Events[0].EventChangeType) removeQT = userTracker.decreaseTrackedResource(path2, TestApp2, usage3, false) - actualResources1 := getUserResource(userTracker) - assert.Equal(t, removeQT, false, "wrong remove queue tracker value") + + actualResources1 := userTracker.getUsedResources() assert.Equal(t, "map[mem:70000000 vcore:70000]", actualResources1["root"].String(), "wrong resource") assert.Equal(t, "map[mem:70000000 vcore:70000]", actualResources1["root.parent"].String(), "wrong resource") assert.Equal(t, "map[mem:60000000 vcore:60000]", actualResources1["root.parent.child1"].String(), "wrong resource") @@ -289,9 +289,3 @@ func TestUTCanRunApp(t *testing.T) { assert.Assert(t, userTracker.canRunApp(hierarchy1, TestApp1)) assert.Assert(t, !userTracker.canRunApp(hierarchy1, TestApp2)) } - -func getUserResource(ut *UserTracker) map[string]*resources.Resource { - resources := make(map[string]*resources.Resource) - usage := ut.GetUserResourceUsageDAOInfo() - return internalGetResource(usage.Queues, resources) -} diff --git a/pkg/scheduler/ugm/utilities_test.go b/pkg/scheduler/ugm/utilities_test.go index 09481e1e2..ec666b46f 100644 --- a/pkg/scheduler/ugm/utilities_test.go +++ b/pkg/scheduler/ugm/utilities_test.go @@ -22,21 +22,8 @@ import ( "testing" "gotest.tools/v3/assert" - - "github.com/apache/yunikorn-core/pkg/common/resources" - "github.com/apache/yunikorn-core/pkg/webservice/dao" ) -func internalGetResource(usage *dao.ResourceUsageDAOInfo, resources map[string]*resources.Resource) map[string]*resources.Resource { - resources[usage.QueuePath] = usage.ResourceUsage - if len(usage.Children) > 0 { - for _, resourceUsage := range usage.Children { - internalGetResource(resourceUsage, resources) - } - } - return resources -} - func TestGetParentQueuePath(t *testing.T) { assert.Equal(t, getParentPath(""), "") assert.Equal(t, getParentPath("root"), "") diff --git a/pkg/scheduler/utilities_test.go b/pkg/scheduler/utilities_test.go index d2f9df80d..c3bd6c9ac 100644 --- a/pkg/scheduler/utilities_test.go +++ b/pkg/scheduler/utilities_test.go @@ -30,7 +30,6 @@ import ( "github.com/apache/yunikorn-core/pkg/rmproxy" "github.com/apache/yunikorn-core/pkg/scheduler/objects" "github.com/apache/yunikorn-core/pkg/scheduler/ugm" - "github.com/apache/yunikorn-core/pkg/webservice/dao" siCommon "github.com/apache/yunikorn-scheduler-interface/lib/go/common" "github.com/apache/yunikorn-scheduler-interface/lib/go/si" ) @@ -714,20 +713,24 @@ func assertLimits(t *testing.T, userGroup security.UserGroup, expected *resource func assertUserGroupResourceMaxLimits(t *testing.T, userGroup security.UserGroup, expected *resources.Resource, expectedQueuesMaxLimits map[string]map[string]interface{}) { manager := ugm.GetUserManager() - userResource := manager.GetUserResources(userGroup) + userResource := manager.GetUserResources(userGroup.User) groupResource := manager.GetGroupResources(userGroup.Groups[0]) + if expected == nil { + assert.Assert(t, userResource.IsEmpty(), "expected empty resource in user tracker") + assert.Assert(t, groupResource.IsEmpty(), "expected empty resource in group tracker") + } else { + assert.Assert(t, resources.Equals(userResource, expected), "user value '%s' not equal to expected '%s'", userResource.String(), expected.String()) + assert.Assert(t, resources.Equals(groupResource, expected), "group value '%s' not equal to expected '%s'", groupResource.String(), expected.String()) + } ut := manager.GetUserTracker(userGroup.User) if ut != nil { - maxResources := make(map[string]*resources.Resource) - usage := ut.GetUserResourceUsageDAOInfo() - getMaxResource(usage.Queues, maxResources) + maxResources := ut.GetMaxResources() for q, qMaxLimits := range expectedQueuesMaxLimits { if qRes, ok := maxResources[q]; ok { assert.Equal(t, resources.Equals(qRes, qMaxLimits[maxresources].(*resources.Resource)), true) } } - maxApplications := make(map[string]uint64) - getMaxApplications(usage.Queues, maxApplications) + maxApplications := ut.GetMaxApplications() for q, qMaxLimits := range expectedQueuesMaxLimits { if qApps, ok := maxApplications[q]; ok { assert.Equal(t, qApps, qMaxLimits[maxapplications].(uint64), "queue path is "+q+" actual: "+strconv.Itoa(int(qApps))+", expected: "+strconv.Itoa(int(qMaxLimits[maxapplications].(uint64)))) @@ -735,44 +738,19 @@ func assertUserGroupResourceMaxLimits(t *testing.T, userGroup security.UserGroup } } - gt := manager.GetUserTracker(userGroup.User) + gt := manager.GetGroupTracker(userGroup.Groups[0]) if gt != nil { - gMaxResources := make(map[string]*resources.Resource) - gUsage := gt.GetUserResourceUsageDAOInfo() - getMaxResource(gUsage.Queues, gMaxResources) + gMaxResources := gt.GetMaxResources() for q, qMaxLimits := range expectedQueuesMaxLimits { if qRes, ok := gMaxResources[q]; ok { assert.Equal(t, resources.Equals(qRes, qMaxLimits[maxresources].(*resources.Resource)), true) } } - gMaxApps := make(map[string]uint64) - getMaxApplications(gUsage.Queues, gMaxApps) + gMaxApps := gt.GetMaxApplications() for q, qMaxLimits := range expectedQueuesMaxLimits { if qApps, ok := gMaxApps[q]; ok { assert.Equal(t, qApps, qMaxLimits[maxapplications].(uint64)) } } } - assert.Equal(t, resources.Equals(userResource, expected), true) - assert.Equal(t, resources.Equals(groupResource, expected), true) -} - -func getMaxResource(usage *dao.ResourceUsageDAOInfo, maxResources map[string]*resources.Resource) map[string]*resources.Resource { - maxResources[usage.QueuePath] = usage.MaxResources - if len(usage.Children) > 0 { - for _, resourceUsage := range usage.Children { - getMaxResource(resourceUsage, maxResources) - } - } - return maxResources -} - -func getMaxApplications(usage *dao.ResourceUsageDAOInfo, maxApplications map[string]uint64) map[string]uint64 { - maxApplications[usage.QueuePath] = usage.MaxApplications - if len(usage.Children) > 0 { - for _, resourceUsage := range usage.Children { - getMaxApplications(resourceUsage, maxApplications) - } - } - return maxApplications } diff --git a/pkg/webservice/dao/ugm_info.go b/pkg/webservice/dao/ugm_info.go index e7b26cf2d..4ac5dee1f 100644 --- a/pkg/webservice/dao/ugm_info.go +++ b/pkg/webservice/dao/ugm_info.go @@ -18,8 +18,6 @@ package dao -import "github.com/apache/yunikorn-core/pkg/common/resources" - type UserResourceUsageDAOInfo struct { UserName string `json:"userName"` // no omitempty, user name should not be empty Groups map[string]string `json:"groups,omitempty"` @@ -34,9 +32,9 @@ type GroupResourceUsageDAOInfo struct { type ResourceUsageDAOInfo struct { QueuePath string `json:"queuePath"` // no omitempty, queue path should not be empty - ResourceUsage *resources.Resource `json:"resourceUsage,omitempty"` + ResourceUsage map[string]int64 `json:"resourceUsage,omitempty"` RunningApplications []string `json:"runningApplications,omitempty"` - MaxResources *resources.Resource `json:"maxResources,omitempty"` + MaxResources map[string]int64 `json:"maxResources,omitempty"` MaxApplications uint64 `json:"maxApplications,omitempty"` Children []*ResourceUsageDAOInfo `json:"children,omitempty"` } diff --git a/pkg/webservice/handlers.go b/pkg/webservice/handlers.go index d06d62340..934685165 100644 --- a/pkg/webservice/handlers.go +++ b/pkg/webservice/handlers.go @@ -1183,10 +1183,10 @@ func getMetrics(w http.ResponseWriter, r *http.Request) { func getUsersResourceUsage(w http.ResponseWriter, _ *http.Request) { writeHeaders(w) userManager := ugm.GetUserManager() - usersResources := userManager.GetUsersResources() - result := make([]*dao.UserResourceUsageDAOInfo, len(usersResources)) - for i, tracker := range usersResources { - result[i] = tracker.GetUserResourceUsageDAOInfo() + trackers := userManager.GetUserTrackers() + result := make([]*dao.UserResourceUsageDAOInfo, len(trackers)) + for i, tracker := range trackers { + result[i] = tracker.GetResourceUsageDAOInfo() } if err := json.NewEncoder(w).Encode(result); err != nil { buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError) @@ -1215,8 +1215,8 @@ func getUserResourceUsage(w http.ResponseWriter, r *http.Request) { buildJSONErrorResponse(w, UserDoesNotExists, http.StatusNotFound) return } - var result = userTracker.GetUserResourceUsageDAOInfo() - if err := json.NewEncoder(w).Encode(result); err != nil { + result := userTracker.GetResourceUsageDAOInfo() + if err = json.NewEncoder(w).Encode(result); err != nil { buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError) } } @@ -1224,10 +1224,10 @@ func getUserResourceUsage(w http.ResponseWriter, r *http.Request) { func getGroupsResourceUsage(w http.ResponseWriter, r *http.Request) { writeHeaders(w) userManager := ugm.GetUserManager() - groupsResources := userManager.GetGroupsResources() - result := make([]*dao.GroupResourceUsageDAOInfo, len(groupsResources)) - for i, tracker := range groupsResources { - result[i] = tracker.GetGroupResourceUsageDAOInfo() + trackers := userManager.GetGroupTrackers() + result := make([]*dao.GroupResourceUsageDAOInfo, len(trackers)) + for i, tracker := range trackers { + result[i] = tracker.GetResourceUsageDAOInfo() } if err := json.NewEncoder(w).Encode(result); err != nil { buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError) @@ -1256,8 +1256,8 @@ func getGroupResourceUsage(w http.ResponseWriter, r *http.Request) { buildJSONErrorResponse(w, GroupDoesNotExists, http.StatusNotFound) return } - var result = groupTracker.GetGroupResourceUsageDAOInfo() - if err := json.NewEncoder(w).Encode(result); err != nil { + result := groupTracker.GetResourceUsageDAOInfo() + if err = json.NewEncoder(w).Encode(result); err != nil { buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError) } } diff --git a/pkg/webservice/handlers_test.go b/pkg/webservice/handlers_test.go index 4ffb231cf..f262140b1 100644 --- a/pkg/webservice/handlers_test.go +++ b/pkg/webservice/handlers_test.go @@ -2149,12 +2149,12 @@ func TestSpecificUserResourceUsage(t *testing.T) { Groups: map[string]string{"app-1": "testgroup"}, Queues: &dao.ResourceUsageDAOInfo{ QueuePath: "root", - ResourceUsage: resources.NewResourceFromMap(map[string]resources.Quantity{"vcore": 1}), + ResourceUsage: map[string]int64{"vcore": 1}, RunningApplications: []string{"app-1"}, Children: []*dao.ResourceUsageDAOInfo{ { QueuePath: "root.default", - ResourceUsage: resources.NewResourceFromMap(map[string]resources.Quantity{"vcore": 1}), + ResourceUsage: map[string]int64{"vcore": 1}, RunningApplications: []string{"app-1"}, }, }, @@ -2220,13 +2220,13 @@ func TestSpecificGroupResourceUsage(t *testing.T) { Applications: []string{"app-1"}, Queues: &dao.ResourceUsageDAOInfo{ QueuePath: "root", - ResourceUsage: resources.NewResourceFromMap(map[string]resources.Quantity{"vcore": 1}), + ResourceUsage: map[string]int64{"vcore": 1}, RunningApplications: []string{"app-1"}, Children: []*dao.ResourceUsageDAOInfo{ { QueuePath: "root.default", - ResourceUsage: resources.NewResourceFromMap(map[string]resources.Quantity{"vcore": 1}), - MaxResources: resources.NewResourceFromMap(map[string]resources.Quantity{"cpu": 200}), + ResourceUsage: map[string]int64{"vcore": 1}, + MaxResources: map[string]int64{"cpu": 200}, RunningApplications: []string{"app-1"}, }, }, @@ -2283,8 +2283,8 @@ func TestUsersAndGroupsResourceUsage(t *testing.T) { getUsersResourceUsage(resp, req) err = json.Unmarshal(resp.outputBytes, &usersResourceUsageDao) assert.NilError(t, err, unmarshalError) - assert.Equal(t, usersResourceUsageDao[0].Queues.ResourceUsage.String(), - resources.NewResourceFromMap(map[string]resources.Quantity{siCommon.CPU: 1}).String()) + assert.DeepEqual(t, usersResourceUsageDao[0].Queues.ResourceUsage, + resources.NewResourceFromMap(map[string]resources.Quantity{siCommon.CPU: 1}).DAOMap()) // Assert existing users assert.Equal(t, len(usersResourceUsageDao), 1) @@ -2298,8 +2298,8 @@ func TestUsersAndGroupsResourceUsage(t *testing.T) { getGroupsResourceUsage(resp, req) err = json.Unmarshal(resp.outputBytes, &groupsResourceUsageDao) assert.NilError(t, err, unmarshalError) - assert.Equal(t, groupsResourceUsageDao[0].Queues.ResourceUsage.String(), - resources.NewResourceFromMap(map[string]resources.Quantity{siCommon.CPU: 1}).String()) + assert.DeepEqual(t, groupsResourceUsageDao[0].Queues.ResourceUsage, + resources.NewResourceFromMap(map[string]resources.Quantity{siCommon.CPU: 1}).DAOMap()) // Assert existing groups assert.Equal(t, len(groupsResourceUsageDao), 1)