[YUNIKORN-2057] Optimize FindQueueByAppID (apache#1037)

Closes: apache#1037

Signed-off-by: Peter Bacsko <[email protected]>

Author: adrian-lin-1-0-0

pkg/scheduler/objects/app_queue_mapping.go

@@ -0,0 +1,51 @@
+/*
+Licensed to the Apache Software Foundation (ASF) under one
+or more contributor license agreements.  See the NOTICE file
+distributed with this work for additional information
+regarding copyright ownership.  The ASF licenses this file
+to you under the Apache License, Version 2.0 (the
+"License"); you may not use this file except in compliance
+with the License.  You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package objects
+
+import "github.com/apache/yunikorn-core/pkg/locking"
+
+// AppQueueMapping maintains a mapping between application IDs and their corresponding queues.
+type AppQueueMapping struct {
+	byAppID map[string]*Queue
+	locking.RWMutex
+}
+
+func NewAppQueueMapping() *AppQueueMapping {
+	return &AppQueueMapping{
+		byAppID: make(map[string]*Queue),
+	}
+}
+
+func (aqm *AppQueueMapping) AddAppQueueMapping(appID string, queue *Queue) {
+	aqm.Lock()
+	defer aqm.Unlock()
+	aqm.byAppID[appID] = queue
+}
+
+func (aqm *AppQueueMapping) GetQueueByAppId(appID string) *Queue {
+	aqm.RLock()
+	defer aqm.RUnlock()
+	return aqm.byAppID[appID]
+}
+
+func (aqm *AppQueueMapping) RemoveAppQueueMapping(appID string) {
+	aqm.Lock()
+	defer aqm.Unlock()
+	delete(aqm.byAppID, appID)
+}

pkg/scheduler/objects/app_queue_mapping_test.go

@@ -0,0 +1,48 @@
+/*
+Licensed to the Apache Software Foundation (ASF) under one
+or more contributor license agreements.  See the NOTICE file
+distributed with this work for additional information
+regarding copyright ownership.  The ASF licenses this file
+to you under the Apache License, Version 2.0 (the
+"License"); you may not use this file except in compliance
+with the License.  You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package objects
+
+import (
+	"testing"
+
+	"gotest.tools/v3/assert"
+)
+
+func TestNewAppQueueMapping(t *testing.T) {
+	aqm := NewAppQueueMapping()
+	assert.Assert(t, aqm != nil, "expected non-nil AppQueueMapping")
+	assert.Equal(t, 0, len(aqm.byAppID), "expected empty byAppID map")
+}
+func TestAppQueueMappingOperations(t *testing.T) {
+	aqm := NewAppQueueMapping()
+	queue := &Queue{}
+	appID := "app-1234"
+
+	// Test AddAppQueueMapping
+	aqm.AddAppQueueMapping(appID, queue)
+	assert.Equal(t, 1, len(aqm.byAppID), "expected 1 entry in byAppID map")
+
+	// Test FindQueueByAppID
+	foundQueue := aqm.GetQueueByAppId(appID)
+	assert.Equal(t, foundQueue, queue, "expected to find the correct queue for appID %s", appID)
+
+	// Test RemoveAppQueueMapping
+	aqm.RemoveAppQueueMapping(appID)
+	assert.Equal(t, 0, len(aqm.byAppID), "expected 0 entries in byAppID map after removal")
+}

pkg/scheduler/objects/application.go

@@ -1395,7 +1395,7 @@ func (sa *Application) tryRequiredNodePreemption(reserve *reservation, ask *Allo
 			zap.String("allocation name", ask.GetAllocationName()),
 			zap.Int("no.of victims", len(victims)))
 		for _, victim := range victims {
-			if victimQueue := sa.queue.FindQueueByAppID(victim.GetApplicationID()); victimQueue != nil {
+			if victimQueue := sa.queue.GetQueueByAppID(victim.GetApplicationID()); victimQueue != nil {
 				victimQueue.IncPreemptingResource(victim.GetAllocatedResource())
 			}
 			victim.MarkPreempted()

pkg/scheduler/objects/application_test.go

@@ -1040,7 +1040,7 @@ func TestUpdateAllocationResourcePending(t *testing.T) {
 	app := newApplication(appID1, "default", "root.a")
 	root, err := createRootQueue(nil)
 	assert.NilError(t, err, "failed to create root queue")
-	queue, err := createDynamicQueue(root, "test", false)
+	queue, err := createDynamicQueue(root, "test", false, nil)
 	assert.NilError(t, err, "failed to create test queue")
 	app.SetQueue(queue)
 
@@ -1088,7 +1088,7 @@ func TestUpdateAllocationResourceAllocated(t *testing.T) {
 	app := newApplication(appID1, "default", "root.a")
 	root, err := createRootQueue(nil)
 	assert.NilError(t, err, "failed to create root queue")
-	queue, err := createDynamicQueue(root, "test", false)
+	queue, err := createDynamicQueue(root, "test", false, nil)
 	assert.NilError(t, err, "failed to create test queue")
 	app.SetQueue(queue)
 
@@ -1138,7 +1138,7 @@ func TestQueueUpdate(t *testing.T) {
 
 	root, err := createRootQueue(nil)
 	assert.NilError(t, err, "failed to create root queue")
-	queue, err := createDynamicQueue(root, "test", false)
+	queue, err := createDynamicQueue(root, "test", false, nil)
 	assert.NilError(t, err, "failed to create test queue")
 	app.SetQueue(queue)
 	assert.Equal(t, app.GetQueuePath(), "root.test")
@@ -2195,19 +2195,21 @@ func TestTryAllocatePreemptQueue(t *testing.T) {
 	getNode := func(nodeID string) *Node {
 		return nodeMap[nodeID]
 	}
+	appQueueMapping := NewAppQueueMapping()
 
 	rootQ, err := createRootQueue(map[string]string{"first": "20"})
 	assert.NilError(t, err)
-	parentQ, err := createManagedQueueGuaranteed(rootQ, "parent", true, map[string]string{"first": "20"}, map[string]string{"first": "10"})
+	parentQ, err := createManagedQueueGuaranteed(rootQ, "parent", true, map[string]string{"first": "20"}, map[string]string{"first": "10"}, appQueueMapping)
 	assert.NilError(t, err)
-	childQ1, err := createManagedQueueGuaranteed(parentQ, "child1", false, nil, map[string]string{"first": "5"})
+	childQ1, err := createManagedQueueGuaranteed(parentQ, "child1", false, nil, map[string]string{"first": "5"}, appQueueMapping)
 	assert.NilError(t, err)
-	childQ2, err := createManagedQueueGuaranteed(parentQ, "child2", false, nil, map[string]string{"first": "5"})
+	childQ2, err := createManagedQueueGuaranteed(parentQ, "child2", false, nil, map[string]string{"first": "5"}, appQueueMapping)
 	assert.NilError(t, err)
 
 	app1 := newApplication(appID1, "default", "root.parent.child1")
 	app1.SetQueue(childQ1)
-	childQ1.applications[appID1] = app1
+	childQ1.AddApplication(app1)
+	appQueueMapping.AddAppQueueMapping(appID1, childQ1)
 	ask1 := newAllocationAsk("alloc1", appID1, resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5}))
 	err = app1.AddAllocationAsk(ask1)
 	assert.NilError(t, err)
@@ -2217,7 +2219,8 @@ func TestTryAllocatePreemptQueue(t *testing.T) {
 
 	app2 := newApplication(appID2, "default", "root.parent.child2")
 	app2.SetQueue(childQ2)
-	childQ2.applications[appID2] = app2
+	childQ2.AddApplication(app2)
+	appQueueMapping.AddAppQueueMapping(appID2, childQ2)
 	ask3 := newAllocationAsk("alloc3", appID2, resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5}))
 	ask3.allowPreemptOther = true
 	err = app2.AddAllocationAsk(ask3)
@@ -2276,20 +2279,23 @@ func createPreemptNodeTestSetup(t *testing.T) (func() NodeIterator, func(NodeID
 		return nodeMap[nodeID]
 	}
 
+	appQueueMapping := NewAppQueueMapping()
+
 	rootQ, err := createRootQueue(map[string]string{"first": "40"})
 	assert.NilError(t, err)
-	parentQ, err := createManagedQueueGuaranteed(rootQ, "parent", true, map[string]string{"first": "20"}, map[string]string{"first": "10"})
+	parentQ, err := createManagedQueueGuaranteed(rootQ, "parent", true, map[string]string{"first": "20"}, map[string]string{"first": "10"}, appQueueMapping)
 	assert.NilError(t, err)
-	unlimitedQ, err := createManagedQueueGuaranteed(rootQ, "unlimited", false, nil, nil)
+	unlimitedQ, err := createManagedQueueGuaranteed(rootQ, "unlimited", false, nil, nil, appQueueMapping)
 	assert.NilError(t, err)
-	childQ1, err := createManagedQueueGuaranteed(parentQ, "child1", false, nil, map[string]string{"first": "5"})
+	childQ1, err := createManagedQueueGuaranteed(parentQ, "child1", false, nil, map[string]string{"first": "5"}, appQueueMapping)
 	assert.NilError(t, err)
-	childQ2, err := createManagedQueueGuaranteed(parentQ, "child2", false, nil, map[string]string{"first": "5"})
+	childQ2, err := createManagedQueueGuaranteed(parentQ, "child2", false, nil, map[string]string{"first": "5"}, appQueueMapping)
 	assert.NilError(t, err)
 
 	app0 := newApplication(appID0, "default", "root.unlimited")
 	app0.SetQueue(unlimitedQ)
-	unlimitedQ.applications[appID0] = app0
+	unlimitedQ.AddApplication(app0)
+	appQueueMapping.AddAppQueueMapping(appID0, unlimitedQ)
 	ask00 := newAllocationAsk("alloc0-0", appID0, resources.NewResourceFromMap(map[string]resources.Quantity{"first": 11}))
 	err = app0.AddAllocationAsk(ask00)
 	assert.NilError(t, err)
@@ -2299,7 +2305,8 @@ func createPreemptNodeTestSetup(t *testing.T) (func() NodeIterator, func(NodeID
 
 	app1 := newApplication(appID1, "default", "root.parent.child1")
 	app1.SetQueue(childQ1)
-	childQ1.applications[appID1] = app1
+	childQ1.AddApplication(app1)
+	appQueueMapping.AddAppQueueMapping(appID1, childQ1)
 	ask1 := newAllocationAsk("alloc1", appID1, resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5}))
 	err = app1.AddAllocationAsk(ask1)
 	assert.NilError(t, err)
@@ -2309,7 +2316,8 @@ func createPreemptNodeTestSetup(t *testing.T) (func() NodeIterator, func(NodeID
 
 	app2 := newApplication(appID2, "default", "root.parent.child2")
 	app2.SetQueue(childQ2)
-	childQ2.applications[appID2] = app2
+	childQ2.AddApplication(app2)
+	appQueueMapping.AddAppQueueMapping(appID2, childQ2)
 	ask3 := newAllocationAsk("alloc3", appID2, resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5}))
 	ask3.allowPreemptOther = true
 	err = app2.AddAllocationAsk(ask3)
@@ -2358,7 +2366,7 @@ func createPreemptNodeWithReservationsTestSetup(t *testing.T) (func() NodeIterat
 
 	app3 := newApplication(appID3, "default", "root.parent.child2")
 	app3.SetQueue(childQ2)
-	childQ2.applications[appID3] = app3
+	childQ2.AddApplication(app3)
 	ask4 := newAllocationAsk("alloc4", appID3, resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5}))
 	ask4.allowPreemptOther = true
 	ask4.priority = math.MaxInt32
@@ -3257,13 +3265,15 @@ func TestRequiredNodePreemption(t *testing.T) {
 	getNode := func(nodeID string) *Node {
 		return node
 	}
+	appQueueMapping := NewAppQueueMapping()
 
 	// set queue
 	rootQ, err := createRootQueue(map[string]string{"first": "20"})
 	assert.NilError(t, err)
-	childQ, err := createManagedQueue(rootQ, "default", false, map[string]string{"first": "20"})
+	childQ, err := createManagedQueueWithAppQueueMapping(rootQ, "default", false, map[string]string{"first": "20"}, appQueueMapping)
 	assert.NilError(t, err)
 	app.SetQueue(childQ)
+	appQueueMapping.AddAppQueueMapping(app.ApplicationID, childQ)
 
 	// add an ask
 	mockEvents := mock.NewEventSystem()

pkg/scheduler/objects/preemption.go

@@ -649,7 +649,7 @@ func (p *Preemptor) TryPreemption() (*AllocationResult, bool) {
 
 	// preempt the victims
 	for _, victim := range finalVictims {
-		if victimQueue := p.queue.FindQueueByAppID(victim.GetApplicationID()); victimQueue != nil {
+		if victimQueue := p.queue.GetQueueByAppID(victim.GetApplicationID()); victimQueue != nil {
 			victimQueue.IncPreemptingResource(victim.GetAllocatedResource())
 			victim.MarkPreempted()
 			log.Log(log.SchedPreemption).Info("Preempting task",