-
Notifications
You must be signed in to change notification settings - Fork 8.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
YARN-10408. Extract MockQueueHierarchyBuilder to a separate class. Co…
…ntributed by Gergely Pollak
- Loading branch information
1 parent
6e3d705
commit 4454286
Showing
2 changed files
with
163 additions
and
143 deletions.
There are no files selected for viewing
163 changes: 163 additions & 0 deletions
163
...va/org/apache/hadoop/yarn/server/resourcemanager/placement/MockQueueHierarchyBuilder.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,163 @@ | ||
| /** | ||
| * 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 org.apache.hadoop.yarn.server.resourcemanager.placement; | ||
|
|
||
| import com.google.common.collect.Maps; | ||
| import org.apache.commons.compress.utils.Lists; | ||
| import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.*; | ||
|
|
||
| import java.util.List; | ||
| import java.util.Map; | ||
|
|
||
| import static org.mockito.Mockito.mock; | ||
| import static org.mockito.Mockito.when; | ||
|
|
||
| class MockQueueHierarchyBuilder { | ||
| private static final String ROOT = "root"; | ||
| private static final String QUEUE_SEP = "."; | ||
| private List<String> queuePaths = Lists.newArrayList(); | ||
| private List<String> managedParentQueues = Lists.newArrayList(); | ||
| private CapacitySchedulerQueueManager queueManager; | ||
|
|
||
| public static MockQueueHierarchyBuilder create() { | ||
| return new MockQueueHierarchyBuilder(); | ||
| } | ||
|
|
||
| public MockQueueHierarchyBuilder withQueueManager( | ||
| CapacitySchedulerQueueManager queueManager) { | ||
| this.queueManager = queueManager; | ||
| return this; | ||
| } | ||
|
|
||
| public MockQueueHierarchyBuilder withQueue(String queue) { | ||
| this.queuePaths.add(queue); | ||
| return this; | ||
| } | ||
|
|
||
| public MockQueueHierarchyBuilder withManagedParentQueue( | ||
| String managedQueue) { | ||
| this.managedParentQueues.add(managedQueue); | ||
| return this; | ||
| } | ||
|
|
||
| public void build() { | ||
| if (this.queueManager == null) { | ||
| throw new IllegalStateException( | ||
| "QueueManager instance is not provided!"); | ||
| } | ||
|
|
||
| for (String managedParentQueue : managedParentQueues) { | ||
| if (!queuePaths.contains(managedParentQueue)) { | ||
| queuePaths.add(managedParentQueue); | ||
| } else { | ||
| throw new IllegalStateException("Cannot add a managed parent " + | ||
| "and a simple queue with the same path"); | ||
| } | ||
| } | ||
|
|
||
| Map<String, AbstractCSQueue> queues = Maps.newHashMap(); | ||
| for (String queuePath : queuePaths) { | ||
| addQueues(queues, queuePath); | ||
| } | ||
| } | ||
|
|
||
| private void addQueues(Map<String, AbstractCSQueue> queues, | ||
| String queuePath) { | ||
| final String[] pathComponents = queuePath.split("\\" + QUEUE_SEP); | ||
|
|
||
| String currentQueuePath = ""; | ||
| for (int i = 0; i < pathComponents.length; ++i) { | ||
| boolean isLeaf = i == pathComponents.length - 1; | ||
| String queueName = pathComponents[i]; | ||
| String parentPath = currentQueuePath; | ||
| currentQueuePath += currentQueuePath.equals("") ? | ||
| queueName : QUEUE_SEP + queueName; | ||
|
|
||
| if (managedParentQueues.contains(parentPath) && !isLeaf) { | ||
| throw new IllegalStateException("Cannot add a queue under " + | ||
| "managed parent"); | ||
| } | ||
| if (!queues.containsKey(currentQueuePath)) { | ||
| ParentQueue parentQueue = (ParentQueue) queues.get(parentPath); | ||
| AbstractCSQueue queue = createQueue(parentQueue, queueName, | ||
| currentQueuePath, isLeaf); | ||
| queues.put(currentQueuePath, queue); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| private AbstractCSQueue createQueue(ParentQueue parentQueue, | ||
| String queueName, String currentQueuePath, boolean isLeaf) { | ||
| if (queueName.equals(ROOT)) { | ||
| return createRootQueue(ROOT); | ||
| } else if (managedParentQueues.contains(currentQueuePath)) { | ||
| return addManagedParentQueueAsChildOf(parentQueue, queueName); | ||
| } else if (isLeaf) { | ||
| return addLeafQueueAsChildOf(parentQueue, queueName); | ||
| } else { | ||
| return addParentQueueAsChildOf(parentQueue, queueName); | ||
| } | ||
| } | ||
|
|
||
| private AbstractCSQueue createRootQueue(String rootQueueName) { | ||
| ParentQueue root = mock(ParentQueue.class); | ||
| when(root.getQueuePath()).thenReturn(rootQueueName); | ||
| when(queueManager.getQueue(rootQueueName)).thenReturn(root); | ||
| when(queueManager.getQueueByFullName(rootQueueName)).thenReturn(root); | ||
| return root; | ||
| } | ||
|
|
||
| private AbstractCSQueue addParentQueueAsChildOf(ParentQueue parent, | ||
| String queueName) { | ||
| ParentQueue queue = mock(ParentQueue.class); | ||
| setQueueFields(parent, queue, queueName); | ||
| return queue; | ||
| } | ||
|
|
||
| private AbstractCSQueue addManagedParentQueueAsChildOf(ParentQueue parent, | ||
| String queueName) { | ||
| ManagedParentQueue queue = mock(ManagedParentQueue.class); | ||
| setQueueFields(parent, queue, queueName); | ||
| return queue; | ||
| } | ||
|
|
||
| private AbstractCSQueue addLeafQueueAsChildOf(ParentQueue parent, | ||
| String queueName) { | ||
| LeafQueue queue = mock(LeafQueue.class); | ||
| setQueueFields(parent, queue, queueName); | ||
| return queue; | ||
| } | ||
|
|
||
| private void setQueueFields(ParentQueue parent, AbstractCSQueue newQueue, | ||
| String queueName) { | ||
| String fullPathOfQueue = parent.getQueuePath() + QUEUE_SEP + queueName; | ||
| addQueueToQueueManager(queueName, newQueue, fullPathOfQueue); | ||
|
|
||
| when(newQueue.getParent()).thenReturn(parent); | ||
| when(newQueue.getQueuePath()).thenReturn(fullPathOfQueue); | ||
| when(newQueue.getQueueName()).thenReturn(queueName); | ||
| } | ||
|
|
||
| private void addQueueToQueueManager(String queueName, AbstractCSQueue queue, | ||
| String fullPathOfQueue) { | ||
| when(queueManager.getQueue(queueName)).thenReturn(queue); | ||
| when(queueManager.getQueue(fullPathOfQueue)).thenReturn(queue); | ||
| when(queueManager.getQueueByFullName(fullPathOfQueue)).thenReturn(queue); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters