YARN-10408. Extract MockQueueHierarchyBuilder to a separate class. Contributed by Gergely Pollak
This commit is contained in:
Szilard Nemeth
parent
6e3d7050a9
commit
44542863f4
|
|
@ -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);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -24,11 +24,7 @@ import static org.mockito.Mockito.when;
|
||||||
import static org.mockito.Mockito.isNull;
|
import static org.mockito.Mockito.isNull;
|
||||||
|
|
||||||
import java.util.Arrays;
|
import java.util.Arrays;
|
||||||
import java.util.List;
|
|
||||||
import java.util.Map;
|
|
||||||
|
|
||||||
import com.google.common.collect.Maps;
|
|
||||||
import org.apache.commons.compress.utils.Lists;
|
|
||||||
import org.apache.hadoop.fs.CommonConfigurationKeys;
|
import org.apache.hadoop.fs.CommonConfigurationKeys;
|
||||||
import org.apache.hadoop.security.GroupMappingServiceProvider;
|
import org.apache.hadoop.security.GroupMappingServiceProvider;
|
||||||
import org.apache.hadoop.security.Groups;
|
import org.apache.hadoop.security.Groups;
|
||||||
|
|
@ -39,11 +35,7 @@ import org.apache.hadoop.yarn.exceptions.YarnException;
|
||||||
import org.apache.hadoop.yarn.server.resourcemanager.placement.QueueMapping.MappingType;
|
import org.apache.hadoop.yarn.server.resourcemanager.placement.QueueMapping.MappingType;
|
||||||
import org.apache.hadoop.yarn.server.resourcemanager.placement.QueueMapping.QueueMappingBuilder;
|
import org.apache.hadoop.yarn.server.resourcemanager.placement.QueueMapping.QueueMappingBuilder;
|
||||||
import org.apache.hadoop.yarn.server.resourcemanager.placement.TestUserGroupMappingPlacementRule.QueueMappingTestData.QueueMappingTestDataBuilder;
|
import org.apache.hadoop.yarn.server.resourcemanager.placement.TestUserGroupMappingPlacementRule.QueueMappingTestData.QueueMappingTestDataBuilder;
|
||||||
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.AbstractCSQueue;
|
|
||||||
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CapacitySchedulerQueueManager;
|
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CapacitySchedulerQueueManager;
|
||||||
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.LeafQueue;
|
|
||||||
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.ManagedParentQueue;
|
|
||||||
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.ParentQueue;
|
|
||||||
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.PrimaryGroupMapping;
|
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.PrimaryGroupMapping;
|
||||||
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.SimpleGroupsMapping;
|
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.SimpleGroupsMapping;
|
||||||
import org.apache.hadoop.yarn.util.Records;
|
import org.apache.hadoop.yarn.util.Records;
|
||||||
|
|
@ -57,141 +49,6 @@ public class TestUserGroupMappingPlacementRule {
|
||||||
private static final Logger LOG =
|
private static final Logger LOG =
|
||||||
LoggerFactory.getLogger(TestUserGroupMappingPlacementRule.class);
|
LoggerFactory.getLogger(TestUserGroupMappingPlacementRule.class);
|
||||||
|
|
||||||
private static 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) {
|
|
||||||
LOG.info("Processing queue path: " + queuePath);
|
|
||||||
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);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
YarnConfiguration conf = new YarnConfiguration();
|
YarnConfiguration conf = new YarnConfiguration();
|
||||||
|
|
||||||
@Before
|
@Before
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue