diff options
Diffstat (limited to 'exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/TestResourcePoolTree.java')
-rw-r--r-- | exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/TestResourcePoolTree.java | 277 |
1 files changed, 277 insertions, 0 deletions
diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/TestResourcePoolTree.java b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/TestResourcePoolTree.java new file mode 100644 index 000000000..fea132d41 --- /dev/null +++ b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/TestResourcePoolTree.java @@ -0,0 +1,277 @@ +/* + * 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.drill.exec.resourcemgr; + +import com.typesafe.config.Config; +import com.typesafe.config.ConfigFactory; +import com.typesafe.config.ConfigValueFactory; +import org.apache.drill.categories.ResourceManagerTest; +import org.apache.drill.exec.ExecConstants; +import org.apache.drill.exec.ops.QueryContext; +import org.apache.drill.exec.resourcemgr.config.QueueAssignmentResult; +import org.apache.drill.exec.resourcemgr.config.ResourcePool; +import org.apache.drill.exec.resourcemgr.config.ResourcePoolImpl; +import org.apache.drill.exec.resourcemgr.config.ResourcePoolTree; +import org.apache.drill.exec.resourcemgr.config.ResourcePoolTreeImpl; +import org.apache.drill.exec.resourcemgr.config.exception.RMConfigException; +import org.apache.drill.exec.server.options.OptionValue; +import org.junit.After; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.apache.drill.exec.resourcemgr.config.RMCommonDefaults.ROOT_POOL_DEFAULT_QUEUE_SELECTION_POLICY; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Category(ResourceManagerTest.class) +public final class TestResourcePoolTree { + + private static final Map<String, Object> poolTreeConfig = new HashMap<>(); + + private static final Map<String, Object> pool1 = new HashMap<>(); + + private static final Map<String, Object> pool2 = new HashMap<>(); + + private static final Map<String, Object> queue1 = new HashMap<>(); + + private static final List<Object> childResourcePools = new ArrayList<>(); + + private static final Map<String, Object> tagSelectorConfig1 = new HashMap<>(); + + private static final Map<String, Object> tagSelectorConfig2 = new HashMap<>(); + + private static final QueryContext mockContext = mock(QueryContext.class); + + @BeforeClass + public static void testSuiteSetup() { + pool1.put(ResourcePoolImpl.POOL_NAME_KEY, "dev"); + pool1.put(ResourcePoolImpl.POOL_MEMORY_SHARE_KEY, 0.80); + + pool2.put(ResourcePoolImpl.POOL_NAME_KEY, "qa"); + pool2.put(ResourcePoolImpl.POOL_MEMORY_SHARE_KEY, 0.20); + + queue1.put("max_query_memory_per_node", 5534); + + tagSelectorConfig1.put("tag", "small"); + tagSelectorConfig2.put("tag", "large"); + } + + @After + public void afterTestCleanup() { + // cleanup resource tree + poolTreeConfig.clear(); + + // cleanup pools + pool1.remove(ResourcePoolImpl.POOL_QUEUE_KEY); + pool1.remove(ResourcePoolImpl.POOL_SELECTOR_KEY); + + pool2.remove(ResourcePoolImpl.POOL_QUEUE_KEY); + pool2.remove(ResourcePoolImpl.POOL_SELECTOR_KEY); + + childResourcePools.clear(); + } + + private ResourcePoolTree getPoolTreeConfig() throws RMConfigException { + poolTreeConfig.put(ResourcePoolImpl.POOL_NAME_KEY, "drill"); + poolTreeConfig.put(ResourcePoolImpl.POOL_CHILDREN_POOLS_KEY, childResourcePools); + + Config rmConfig = ConfigFactory.empty() + .withValue("drill.exec.rm", ConfigValueFactory.fromMap(poolTreeConfig)); + return new ResourcePoolTreeImpl(rmConfig, 10000, 10, 2); + } + + private boolean checkExpectedVsActualPools(List<ResourcePool> actual, List<String> expectedNames) { + if (actual.size() != expectedNames.size()) { + return false; + } + + for (ResourcePool pool : actual) { + if (!expectedNames.contains(pool.getPoolName())) { + return false; + } + } + return true; + } + + @Test + public void testTreeWith2LeafPool() throws Exception { + // pool with tag selector + pool1.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + pool1.put(ResourcePoolImpl.POOL_SELECTOR_KEY, tagSelectorConfig1); + + // pool with default selector + pool2.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + + childResourcePools.add(pool1); + childResourcePools.add(pool2); + + ResourcePoolTree configTree = getPoolTreeConfig(); + + // get all leaf queues names + Set<String> expectedLeafQueue = new HashSet<>(); + expectedLeafQueue.add((String)pool1.get("pool_name")); + expectedLeafQueue.add((String)pool2.get("pool_name")); + + assertEquals("Root pool is different than expected", "drill", configTree.getRootPool().getPoolName()); + assertEquals("Expected and actual leaf queue names are different", expectedLeafQueue, + configTree.getAllLeafQueues().keySet()); + assertEquals("Unexpected Selection policy is in use", ROOT_POOL_DEFAULT_QUEUE_SELECTION_POLICY, + configTree.getSelectionPolicyInUse().getSelectionPolicy()); + } + + @Test(expected = RMConfigException.class) + public void testDuplicateLeafPool() throws Exception { + // leaf pool + pool1.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + childResourcePools.add(pool1); + childResourcePools.add(pool1); + + getPoolTreeConfig(); + } + + @Test(expected = RMConfigException.class) + public void testMissingQueueAtLeafPool() throws Exception { + // leaf pool with queue + pool1.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + pool1.put(ResourcePoolImpl.POOL_SELECTOR_KEY, tagSelectorConfig1); + childResourcePools.add(pool1); + childResourcePools.add(pool2); + + getPoolTreeConfig(); + } + + @Test(expected = RMConfigException.class) + public void testInvalidQueueAtLeafPool() throws Exception { + // leaf pool with invalid queue + int initialValue = (Integer)queue1.remove("max_query_memory_per_node"); + + try { + pool1.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + pool1.put(ResourcePoolImpl.POOL_SELECTOR_KEY, tagSelectorConfig1); + childResourcePools.add(pool1); + + getPoolTreeConfig(); + } finally { + queue1.put("max_query_memory_per_node", initialValue); + } + } + + @Test + public void testRootPoolAsLeaf() throws Exception { + // leaf pool with queue + poolTreeConfig.put(ResourcePoolImpl.POOL_NAME_KEY, "drill"); + poolTreeConfig.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + poolTreeConfig.put(ResourcePoolImpl.POOL_SELECTOR_KEY, tagSelectorConfig1); + + Config rmConfig = ConfigFactory.empty() + .withValue("drill.exec.rm", ConfigValueFactory.fromMap(poolTreeConfig)); + ResourcePoolTree poolTree = new ResourcePoolTreeImpl(rmConfig, 10000, 10, 2); + + assertTrue("Root pool is not a leaf pool", poolTree.getRootPool().isLeafPool()); + assertEquals("Root pool name is not drill", "drill", poolTree.getRootPool().getPoolName()); + assertTrue("Root pool is not the only leaf pool", poolTree.getAllLeafQueues().size() == 1); + assertTrue("Root pool name is not same as leaf pool name", poolTree.getAllLeafQueues().containsKey("drill")); + assertFalse("Root pool should not be a default pool", poolTree.getRootPool().isDefaultPool()); + } + + @Test + public void testTreeWithLeafAndIntermediatePool() throws Exception { + // left leaf pool1 with tag selector + pool1.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + pool1.put(ResourcePoolImpl.POOL_SELECTOR_KEY, tagSelectorConfig1); + + // left leaf pool2 with default selector + pool2.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + + // intermediate left pool1 with 2 leaf pools (pool1, pool2) + Map<String, Object> interPool1 = new HashMap<>(); + List<Object> childPools1 = new ArrayList<>(); + childPools1.add(pool1); + childPools1.add(pool2); + interPool1.put(ResourcePoolImpl.POOL_NAME_KEY, "eng"); + interPool1.put(ResourcePoolImpl.POOL_MEMORY_SHARE_KEY, 0.9); + interPool1.put(ResourcePoolImpl.POOL_CHILDREN_POOLS_KEY, childPools1); + + // right leaf pool + Map<String, Object> rightLeafPool = new HashMap<>(); + rightLeafPool.put(ResourcePoolImpl.POOL_NAME_KEY, "marketing"); + rightLeafPool.put(ResourcePoolImpl.POOL_MEMORY_SHARE_KEY, 0.1); + rightLeafPool.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + rightLeafPool.put(ResourcePoolImpl.POOL_SELECTOR_KEY, tagSelectorConfig2); + + childResourcePools.add(interPool1); + childResourcePools.add(rightLeafPool); + ResourcePoolTree configTree = getPoolTreeConfig(); + + // Test successful selection of all leaf pools + OptionValue testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, ExecConstants + .RM_QUERY_TAGS_KEY, "small,large", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + QueueAssignmentResult assignmentResult = configTree.selectAllQueues(mockContext); + List<ResourcePool> selectedPools = assignmentResult.getSelectedLeafPools(); + List<String> expectedPools = new ArrayList<>(); + expectedPools.add("dev"); + expectedPools.add("qa"); + expectedPools.add("marketing"); + + assertTrue("All leaf pools are not selected", selectedPools.size() == 3); + assertTrue("Selected leaf pools and expected pools are different", + checkExpectedVsActualPools(selectedPools, expectedPools)); + + // Test successful selection of multiple leaf pools + expectedPools.clear(); + testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, ExecConstants + .RM_QUERY_TAGS_KEY, "small", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assignmentResult = configTree.selectAllQueues(mockContext); + selectedPools = assignmentResult.getSelectedLeafPools(); + expectedPools.add("qa"); + expectedPools.add("dev"); + assertTrue("Expected 2 pools to be selected", selectedPools.size() == 2); + assertTrue("Selected leaf pools and expected pools are different", + checkExpectedVsActualPools(selectedPools, expectedPools)); + + // Test successful selection of only left default pool + expectedPools.clear(); + testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, ExecConstants + .RM_QUERY_TAGS_KEY, "medium", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assignmentResult = configTree.selectAllQueues(mockContext); + selectedPools = assignmentResult.getSelectedLeafPools(); + expectedPools.add("qa"); + assertTrue("More than one leaf pool is selected", selectedPools.size() == 1); + assertTrue("Selected leaf pools and expected pools are different", + checkExpectedVsActualPools(selectedPools, expectedPools)); + + // cleanup + interPool1.clear(); + rightLeafPool.clear(); + expectedPools.clear(); + } +} |