Java tutorial
/* * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0 * (the "License"). You may not use this work except in compliance with the License, which is * available at www.apache.org/licenses/LICENSE-2.0 * * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, * either express or implied, as more fully set forth in the License. * * See the NOTICE file distributed with this work for information regarding copyright ownership. */ package alluxio.yarn; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import static org.mockito.Matchers.any; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import alluxio.exception.ExceptionMessage; import org.apache.hadoop.yarn.api.records.Container; import org.apache.hadoop.yarn.api.records.NodeId; import org.apache.hadoop.yarn.api.records.NodeReport; import org.apache.hadoop.yarn.api.records.NodeState; import org.apache.hadoop.yarn.api.records.Resource; import org.apache.hadoop.yarn.client.api.AMRMClient.ContainerRequest; import org.apache.hadoop.yarn.client.api.YarnClient; import org.apache.hadoop.yarn.client.api.async.AMRMClientAsync; import org.apache.hadoop.yarn.util.Records; import org.junit.Before; import org.junit.Rule; import org.junit.Test; import org.junit.rules.ExpectedException; import org.mockito.ArgumentMatcher; import org.mockito.Matchers; import org.mockito.invocation.InvocationOnMock; import org.mockito.stubbing.Answer; import java.util.ArrayList; import java.util.HashSet; import java.util.List; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; /** * Unit tests for {@link ContainerAllocator}. */ public final class ContainerAllocatorTest { private static final String CONTAINER_NAME = "test"; private Resource mResource; private YarnClient mYarnClient; private AMRMClientAsync<ContainerRequest> mRMClient; @Rule public ExpectedException mThrown = ExpectedException.none(); @Before @SuppressWarnings("unchecked") public void before() { mResource = mock(Resource.class); mYarnClient = mock(YarnClient.class); mRMClient = (AMRMClientAsync<ContainerRequest>) mock(AMRMClientAsync.class); } @Test(timeout = 10000) public void oneContainerPerHostFullAllocation() throws Exception { int numHosts = 10; int maxContainersPerHost = 1; testFullAllocation(numHosts, maxContainersPerHost); } @Test(timeout = 10000) public void fiveContainersPerHostFullAllocation() throws Exception { int numHosts = 10; int maxContainersPerHost = 5; testFullAllocation(numHosts, maxContainersPerHost); } @Test(timeout = 10000) public void fiveContainersPerHostHalfAllocation() throws Exception { int numHosts = 10; int maxContainersPerHost = 5; int numContainers = numHosts * maxContainersPerHost / 2; ContainerAllocator containerAllocator = setup(numHosts, maxContainersPerHost, numContainers); List<Container> containers = containerAllocator.allocateContainers(); assertEquals(numContainers, containers.size()); checkMaxHostsLimitNotExceeded(containers, maxContainersPerHost); } @Test(timeout = 10000) public void notEnoughHosts() throws Exception { int numHosts = 10; int maxContainersPerHost = 5; int numContainers = numHosts * maxContainersPerHost + 1; // one container too many ContainerAllocator containerAllocator = setup(numHosts, maxContainersPerHost, numContainers); mThrown.expect(RuntimeException.class); mThrown.expectMessage( ExceptionMessage.YARN_NOT_ENOUGH_HOSTS.getMessage(numContainers, CONTAINER_NAME, numHosts)); containerAllocator.allocateContainers(); } @Test(timeout = 1000) public void allocateMasterInAnyHost() throws Exception { ContainerAllocator containerAllocator = new ContainerAllocator(CONTAINER_NAME, 1, 1, mResource, mYarnClient, mRMClient, "any"); doAnswer(allocateFirstHostAnswer(containerAllocator)).when(mRMClient) .addContainerRequest(Matchers.argThat(new ArgumentMatcher<ContainerRequest>() { @Override public boolean matches(Object o) { ContainerRequest request = (ContainerRequest) o; if (request.getRelaxLocality() == true && request.getNodes().size() == 1 && request.getNodes().get(0).equals("any")) { return true; } return false; } })); containerAllocator.allocateContainers(); } /* * Creates a container allocator for allocating the specified numContainers with the specified * maxContainersPerHost. * * The yarn client is mocked to make it look like there are numHosts different hosts in the * system, and the resource manager client is mocked to allocate containers when they are * requested. */ private ContainerAllocator setup(int numHosts, int maxContainersPerHost, int numContainers) throws Exception { ContainerAllocator containerAllocator = new ContainerAllocator(CONTAINER_NAME, numContainers, maxContainersPerHost, mResource, mYarnClient, mRMClient); List<NodeReport> nodeReports = new ArrayList<>(); for (int i = 0; i < numHosts; i++) { NodeReport nodeReport = Records.newRecord(NodeReport.class); nodeReport.setNodeId(NodeId.newInstance("host" + i, 0)); nodeReports.add(nodeReport); } when(mYarnClient.getNodeReports(Matchers.<NodeState[]>anyVararg())).thenReturn(nodeReports); doAnswer(allocateFirstHostAnswer(containerAllocator)).when(mRMClient) .addContainerRequest(any(ContainerRequest.class)); return containerAllocator; } /* * Tests that when there are the specified number of hosts and the specified max containers per * host, a ContainerAllocator can fully allocate all hosts so that every host has * maxContainersPerHost containers. */ private void testFullAllocation(int numHosts, int maxContainersPerHost) throws Exception { int numContainers = numHosts * maxContainersPerHost; ContainerAllocator containerAllocator = setup(numHosts, maxContainersPerHost, numContainers); List<Container> containers = containerAllocator.allocateContainers(); Set<String> containerHosts = new HashSet<>(); for (Container container : containers) { containerHosts.add(container.getNodeId().getHost()); } assertEquals("All hosts are allocated", numHosts, containerHosts.size()); assertEquals("All containers are allocated", numContainers, containers.size()); checkMaxHostsLimitNotExceeded(containers, maxContainersPerHost); } /* * Creates an Answer to an addContainerRequest method. The Answer picks the first node requested * and allocates a container on it, sending a callback to the specified container allocator */ private Answer<Void> allocateFirstHostAnswer(final ContainerAllocator containerAllocator) { return new Answer<Void>() { @Override public Void answer(InvocationOnMock invocation) throws Throwable { ContainerRequest containerRequest = invocation.getArgumentAt(0, ContainerRequest.class); Container container = Records.newRecord(Container.class); container.setNodeId(NodeId.newInstance(containerRequest.getNodes().get(0), 0)); containerAllocator.allocateContainer(container); return null; } }; } private void checkMaxHostsLimitNotExceeded(List<Container> containers, int maxContainersPerHost) { ConcurrentHashMap<String, Integer> counts = new ConcurrentHashMap<>(); for (Container container : containers) { String host = container.getNodeId().getHost(); counts.putIfAbsent(host, 0); int newCount = counts.get(host) + 1; assertTrue(newCount <= maxContainersPerHost); counts.put(host, newCount); } } }