Java tutorial
/** * 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.tez.dag.app.rm; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import static org.mockito.Matchers.any; import static org.mockito.Matchers.eq; import static org.mockito.Mockito.RETURNS_DEEP_STUBS; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.reset; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import javax.annotation.Nullable; import java.io.IOException; import java.lang.reflect.Method; import java.net.InetSocketAddress; import java.nio.ByteBuffer; import java.util.HashMap; import java.util.HashSet; import java.util.LinkedList; import java.util.List; import java.util.Set; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.security.Credentials; import org.apache.hadoop.yarn.api.records.ApplicationId; import org.apache.hadoop.yarn.api.records.Container; import org.apache.hadoop.yarn.api.records.ContainerExitStatus; import org.apache.hadoop.yarn.api.records.ContainerId; import org.apache.hadoop.yarn.api.records.ContainerStatus; import org.apache.hadoop.yarn.api.records.LocalResource; import org.apache.hadoop.yarn.api.records.NodeId; import org.apache.hadoop.yarn.api.records.Priority; import org.apache.hadoop.yarn.api.records.Resource; import org.apache.hadoop.yarn.event.Event; import org.apache.hadoop.yarn.event.EventHandler; import org.apache.tez.common.ContainerSignatureMatcher; import org.apache.tez.common.TezUtils; import org.apache.tez.dag.api.NamedEntityDescriptor; import org.apache.tez.dag.api.TaskLocationHint; import org.apache.tez.dag.api.TezConfiguration; import org.apache.tez.dag.api.TezConstants; import org.apache.tez.dag.api.TezException; import org.apache.tez.dag.api.UserPayload; import org.apache.tez.dag.api.client.DAGClientServer; import org.apache.tez.dag.app.AppContext; import org.apache.tez.dag.app.ContainerContext; import org.apache.tez.dag.app.ServicePluginLifecycleAbstractService; import org.apache.tez.dag.app.dag.DAG; import org.apache.tez.dag.app.dag.TaskAttempt; import org.apache.tez.dag.app.dag.event.DAGAppMasterEventType; import org.apache.tez.dag.app.dag.event.DAGAppMasterEventUserServiceFatalError; import org.apache.tez.dag.app.dag.event.DAGEventTerminateDag; import org.apache.tez.dag.app.dag.impl.TaskAttemptImpl; import org.apache.tez.dag.app.dag.impl.TaskImpl; import org.apache.tez.dag.app.dag.impl.VertexImpl; import org.apache.tez.dag.app.rm.container.AMContainer; import org.apache.tez.dag.app.rm.container.AMContainerEventAssignTA; import org.apache.tez.dag.app.rm.container.AMContainerEventCompleted; import org.apache.tez.dag.app.rm.container.AMContainerEventType; import org.apache.tez.dag.app.rm.container.AMContainerMap; import org.apache.tez.dag.app.rm.container.AMContainerState; import org.apache.tez.dag.app.web.WebUIService; import org.apache.tez.dag.helpers.DagInfoImplForTest; import org.apache.tez.dag.records.TaskAttemptTerminationCause; import org.apache.tez.dag.records.TezDAGID; import org.apache.tez.dag.records.TezTaskAttemptID; import org.apache.tez.dag.records.TezTaskID; import org.apache.tez.dag.records.TezVertexID; import org.apache.tez.runtime.api.impl.TaskSpec; import org.apache.tez.serviceplugins.api.ServicePluginErrorDefaults; import org.apache.tez.serviceplugins.api.ServicePluginException; import org.apache.tez.serviceplugins.api.TaskAttemptEndReason; import org.apache.tez.serviceplugins.api.TaskScheduler; import org.apache.tez.serviceplugins.api.TaskSchedulerContext; import org.apache.tez.serviceplugins.api.TaskSchedulerDescriptor; import org.junit.Assert; import org.junit.Before; import org.junit.Test; import com.google.common.collect.Lists; import org.mockito.ArgumentCaptor; import org.mockito.invocation.InvocationOnMock; import org.mockito.stubbing.Answer; @SuppressWarnings("rawtypes") public class TestTaskSchedulerManager { class TestEventHandler implements EventHandler { List<Event> events = Lists.newLinkedList(); @Override public void handle(Event event) { events.add(event); } } class MockTaskSchedulerManager extends TaskSchedulerManager { final AtomicBoolean notify = new AtomicBoolean(false); public MockTaskSchedulerManager(AppContext appContext, DAGClientServer clientService, EventHandler eventHandler, ContainerSignatureMatcher containerSignatureMatcher, WebUIService webUI) { super(appContext, clientService, eventHandler, containerSignatureMatcher, webUI, Lists.newArrayList(new NamedEntityDescriptor("FakeDescriptor", null)), false); } @Override protected void instantiateSchedulers(String host, int port, String trackingUrl, AppContext appContext) { taskSchedulers[0] = new TaskSchedulerWrapper(mockTaskScheduler); taskSchedulerServiceWrappers[0] = new ServicePluginLifecycleAbstractService<>( taskSchedulers[0].getTaskScheduler()); } @Override protected void notifyForTest() { synchronized (notify) { notify.set(true); notify.notifyAll(); } } } AppContext mockAppContext; DAGClientServer mockClientService; TestEventHandler mockEventHandler; ContainerSignatureMatcher mockSigMatcher; MockTaskSchedulerManager schedulerHandler; TaskScheduler mockTaskScheduler; AMContainerMap mockAMContainerMap; WebUIService mockWebUIService; @Before public void setup() { mockAppContext = mock(AppContext.class, RETURNS_DEEP_STUBS); doReturn(new Configuration(false)).when(mockAppContext).getAMConf(); mockClientService = mock(DAGClientServer.class); mockEventHandler = new TestEventHandler(); mockSigMatcher = mock(ContainerSignatureMatcher.class); mockTaskScheduler = mock(TaskScheduler.class); mockAMContainerMap = mock(AMContainerMap.class); mockWebUIService = mock(WebUIService.class); when(mockAppContext.getAllContainers()).thenReturn(mockAMContainerMap); when(mockClientService.getBindAddress()).thenReturn(new InetSocketAddress(10000)); schedulerHandler = new MockTaskSchedulerManager(mockAppContext, mockClientService, mockEventHandler, mockSigMatcher, mockWebUIService); } @Test(timeout = 5000) public void testSimpleAllocate() throws Exception { Configuration conf = new Configuration(false); schedulerHandler.init(conf); schedulerHandler.start(); TaskAttemptImpl mockTaskAttempt = mock(TaskAttemptImpl.class); TezTaskAttemptID mockAttemptId = mock(TezTaskAttemptID.class); when(mockAttemptId.getId()).thenReturn(0); when(mockTaskAttempt.getID()).thenReturn(mockAttemptId); Resource resource = Resource.newInstance(1024, 1); ContainerContext containerContext = new ContainerContext(new HashMap<String, LocalResource>(), new Credentials(), new HashMap<String, String>(), ""); int priority = 10; TaskLocationHint locHint = TaskLocationHint.createTaskLocationHint(new HashSet<String>(), null); ContainerId mockCId = mock(ContainerId.class); Container container = mock(Container.class); when(container.getId()).thenReturn(mockCId); AMContainer mockAMContainer = mock(AMContainer.class); when(mockAMContainer.getContainerId()).thenReturn(mockCId); when(mockAMContainer.getState()).thenReturn(AMContainerState.IDLE); when(mockAMContainerMap.get(mockCId)).thenReturn(mockAMContainer); AMSchedulerEventTALaunchRequest lr = new AMSchedulerEventTALaunchRequest(mockAttemptId, resource, null, mockTaskAttempt, locHint, priority, containerContext, 0, 0, 0); schedulerHandler.taskAllocated(0, mockTaskAttempt, lr, container); assertEquals(1, mockEventHandler.events.size()); assertTrue(mockEventHandler.events.get(0) instanceof AMContainerEventAssignTA); AMContainerEventAssignTA assignEvent = (AMContainerEventAssignTA) mockEventHandler.events.get(0); assertEquals(priority, assignEvent.getPriority()); assertEquals(mockAttemptId, assignEvent.getTaskAttemptId()); } @Test(timeout = 5000) public void testTaskBasedAffinity() throws Exception { Configuration conf = new Configuration(false); schedulerHandler.init(conf); schedulerHandler.start(); TaskAttemptImpl mockTaskAttempt = mock(TaskAttemptImpl.class); TezTaskAttemptID taId = mock(TezTaskAttemptID.class); String affVertexName = "srcVertex"; int affTaskIndex = 1; TaskLocationHint locHint = TaskLocationHint.createTaskLocationHint(affVertexName, affTaskIndex); VertexImpl affVertex = mock(VertexImpl.class); TaskImpl affTask = mock(TaskImpl.class); TaskAttemptImpl affAttempt = mock(TaskAttemptImpl.class); ContainerId affCId = mock(ContainerId.class); when(affVertex.getTotalTasks()).thenReturn(2); when(affVertex.getTask(affTaskIndex)).thenReturn(affTask); when(affTask.getSuccessfulAttempt()).thenReturn(affAttempt); when(affAttempt.getAssignedContainerID()).thenReturn(affCId); when(mockAppContext.getCurrentDAG().getVertex(affVertexName)).thenReturn(affVertex); Resource resource = Resource.newInstance(100, 1); AMSchedulerEventTALaunchRequest event = new AMSchedulerEventTALaunchRequest(taId, resource, null, mockTaskAttempt, locHint, 3, null, 0, 0, 0); schedulerHandler.notify.set(false); schedulerHandler.handle(event); synchronized (schedulerHandler.notify) { while (!schedulerHandler.notify.get()) { schedulerHandler.notify.wait(); } } // verify mockTaskAttempt affinitized to expected affCId verify(mockTaskScheduler, times(1)).allocateTask(mockTaskAttempt, resource, affCId, Priority.newInstance(3), null, event); schedulerHandler.stop(); schedulerHandler.close(); } @Test(timeout = 5000) public void testContainerPreempted() throws IOException { Configuration conf = new Configuration(false); schedulerHandler.init(conf); schedulerHandler.start(); String diagnostics = "Container preempted by RM."; TaskAttemptImpl mockTask = mock(TaskAttemptImpl.class); ContainerStatus mockStatus = mock(ContainerStatus.class); ContainerId mockCId = mock(ContainerId.class); AMContainer mockAMContainer = mock(AMContainer.class); when(mockAMContainerMap.get(mockCId)).thenReturn(mockAMContainer); when(mockAMContainer.getContainerId()).thenReturn(mockCId); when(mockStatus.getContainerId()).thenReturn(mockCId); when(mockStatus.getDiagnostics()).thenReturn(diagnostics); when(mockStatus.getExitStatus()).thenReturn(ContainerExitStatus.PREEMPTED); schedulerHandler.containerCompleted(0, mockTask, mockStatus); assertEquals(1, mockEventHandler.events.size()); Event event = mockEventHandler.events.get(0); assertEquals(AMContainerEventType.C_COMPLETED, event.getType()); AMContainerEventCompleted completedEvent = (AMContainerEventCompleted) event; assertEquals(mockCId, completedEvent.getContainerId()); assertEquals("Container preempted externally. Container preempted by RM.", completedEvent.getDiagnostics()); assertTrue(completedEvent.isPreempted()); assertEquals(TaskAttemptTerminationCause.EXTERNAL_PREEMPTION, completedEvent.getTerminationCause()); Assert.assertFalse(completedEvent.isDiskFailed()); schedulerHandler.stop(); schedulerHandler.close(); } @Test(timeout = 5000) public void testContainerInternalPreempted() throws IOException, ServicePluginException { Configuration conf = new Configuration(false); schedulerHandler.init(conf); schedulerHandler.start(); AMContainer mockAmContainer = mock(AMContainer.class); when(mockAmContainer.getTaskSchedulerIdentifier()).thenReturn(0); when(mockAmContainer.getContainerLauncherIdentifier()).thenReturn(0); when(mockAmContainer.getTaskCommunicatorIdentifier()).thenReturn(0); ContainerId mockCId = mock(ContainerId.class); verify(mockTaskScheduler, times(0)).deallocateContainer((ContainerId) any()); when(mockAMContainerMap.get(mockCId)).thenReturn(mockAmContainer); schedulerHandler.preemptContainer(0, mockCId); verify(mockTaskScheduler, times(1)).deallocateContainer(mockCId); assertEquals(1, mockEventHandler.events.size()); Event event = mockEventHandler.events.get(0); assertEquals(AMContainerEventType.C_COMPLETED, event.getType()); AMContainerEventCompleted completedEvent = (AMContainerEventCompleted) event; assertEquals(mockCId, completedEvent.getContainerId()); assertEquals("Container preempted internally", completedEvent.getDiagnostics()); assertTrue(completedEvent.isPreempted()); Assert.assertFalse(completedEvent.isDiskFailed()); assertEquals(TaskAttemptTerminationCause.INTERNAL_PREEMPTION, completedEvent.getTerminationCause()); schedulerHandler.stop(); schedulerHandler.close(); } @Test(timeout = 5000) public void testContainerDiskFailed() throws IOException { Configuration conf = new Configuration(false); schedulerHandler.init(conf); schedulerHandler.start(); String diagnostics = "NM disk failed."; TaskAttemptImpl mockTask = mock(TaskAttemptImpl.class); ContainerStatus mockStatus = mock(ContainerStatus.class); ContainerId mockCId = mock(ContainerId.class); AMContainer mockAMContainer = mock(AMContainer.class); when(mockAMContainerMap.get(mockCId)).thenReturn(mockAMContainer); when(mockAMContainer.getContainerId()).thenReturn(mockCId); when(mockStatus.getContainerId()).thenReturn(mockCId); when(mockStatus.getDiagnostics()).thenReturn(diagnostics); when(mockStatus.getExitStatus()).thenReturn(ContainerExitStatus.DISKS_FAILED); schedulerHandler.containerCompleted(0, mockTask, mockStatus); assertEquals(1, mockEventHandler.events.size()); Event event = mockEventHandler.events.get(0); assertEquals(AMContainerEventType.C_COMPLETED, event.getType()); AMContainerEventCompleted completedEvent = (AMContainerEventCompleted) event; assertEquals(mockCId, completedEvent.getContainerId()); assertEquals("Container disk failed. NM disk failed.", completedEvent.getDiagnostics()); Assert.assertFalse(completedEvent.isPreempted()); assertTrue(completedEvent.isDiskFailed()); assertEquals(TaskAttemptTerminationCause.NODE_DISK_ERROR, completedEvent.getTerminationCause()); schedulerHandler.stop(); schedulerHandler.close(); } @Test(timeout = 5000) public void testContainerExceededPMem() throws IOException { Configuration conf = new Configuration(false); schedulerHandler.init(conf); schedulerHandler.start(); String diagnostics = "Exceeded Physical Memory"; TaskAttemptImpl mockTask = mock(TaskAttemptImpl.class); ContainerStatus mockStatus = mock(ContainerStatus.class); ContainerId mockCId = mock(ContainerId.class); AMContainer mockAMContainer = mock(AMContainer.class); when(mockAMContainerMap.get(mockCId)).thenReturn(mockAMContainer); when(mockAMContainer.getContainerId()).thenReturn(mockCId); when(mockStatus.getContainerId()).thenReturn(mockCId); when(mockStatus.getDiagnostics()).thenReturn(diagnostics); // use -104 rather than ContainerExitStatus.KILLED_EXCEEDED_PMEM because // ContainerExitStatus.KILLED_EXCEEDED_PMEM is only available after hadoop-2.5 when(mockStatus.getExitStatus()).thenReturn(-104); schedulerHandler.containerCompleted(0, mockTask, mockStatus); assertEquals(1, mockEventHandler.events.size()); Event event = mockEventHandler.events.get(0); assertEquals(AMContainerEventType.C_COMPLETED, event.getType()); AMContainerEventCompleted completedEvent = (AMContainerEventCompleted) event; assertEquals(mockCId, completedEvent.getContainerId()); assertEquals("Container failed, exitCode=-104. Exceeded Physical Memory", completedEvent.getDiagnostics()); Assert.assertFalse(completedEvent.isPreempted()); Assert.assertFalse(completedEvent.isDiskFailed()); assertEquals(TaskAttemptTerminationCause.CONTAINER_EXITED, completedEvent.getTerminationCause()); schedulerHandler.stop(); schedulerHandler.close(); } @Test(timeout = 5000) public void testHistoryUrlConf() throws Exception { Configuration conf = schedulerHandler.appContext.getAMConf(); final ApplicationId mockApplicationId = mock(ApplicationId.class); doReturn("TEST_APP_ID").when(mockApplicationId).toString(); doReturn(mockApplicationId).when(mockAppContext).getApplicationID(); // ensure history url is empty when timeline server is not the logging class conf.set(TezConfiguration.TEZ_HISTORY_URL_BASE, "http://ui-host:9999"); assertTrue("http://ui-host:9999/#/tez-app/TEST_APP_ID".equals(schedulerHandler.getHistoryUrl())); // ensure the trailing / in history url is handled conf.set(TezConfiguration.TEZ_HISTORY_URL_BASE, "http://ui-host:9998/"); assertTrue("http://ui-host:9998/#/tez-app/TEST_APP_ID".equals(schedulerHandler.getHistoryUrl())); // ensure missing scheme in history url is handled conf.set(TezConfiguration.TEZ_HISTORY_URL_BASE, "ui-host:9998/"); Assert.assertTrue("http://ui-host:9998/#/tez-app/TEST_APP_ID".equals(schedulerHandler.getHistoryUrl())); // handle bad template ex without begining / conf.set(TezConfiguration.TEZ_AM_TEZ_UI_HISTORY_URL_TEMPLATE, "__HISTORY_URL_BASE__#/somepath"); assertTrue("http://ui-host:9998/#/somepath".equals(schedulerHandler.getHistoryUrl())); conf.set(TezConfiguration.TEZ_AM_TEZ_UI_HISTORY_URL_TEMPLATE, "__HISTORY_URL_BASE__?viewPath=tez-app/__APPLICATION_ID__"); conf.set(TezConfiguration.TEZ_HISTORY_URL_BASE, "http://localhost/ui/tez"); assertTrue("http://localhost/ui/tez?viewPath=tez-app/TEST_APP_ID".equals(schedulerHandler.getHistoryUrl())); } @Test(timeout = 5000) public void testNoSchedulerSpecified() throws IOException { try { new TSEHForMultipleSchedulersTest(mockAppContext, mockClientService, mockEventHandler, mockSigMatcher, mockWebUIService, null, false); fail("Expecting an IllegalStateException with no schedulers specified"); } catch (IllegalArgumentException e) { } } // Verified via statics @Test(timeout = 5000) public void testCustomTaskSchedulerSetup() throws IOException { Configuration conf = new Configuration(false); conf.set("testkey", "testval"); UserPayload defaultPayload = TezUtils.createUserPayloadFromConf(conf); String customSchedulerName = "fakeScheduler"; List<NamedEntityDescriptor> taskSchedulers = new LinkedList<>(); ByteBuffer bb = ByteBuffer.allocate(4); bb.putInt(0, 3); UserPayload userPayload = UserPayload.create(bb); taskSchedulers.add(new NamedEntityDescriptor(customSchedulerName, FakeTaskScheduler.class.getName()) .setUserPayload(userPayload)); taskSchedulers.add(new NamedEntityDescriptor(TezConstants.getTezYarnServicePluginName(), null) .setUserPayload(defaultPayload)); TSEHForMultipleSchedulersTest tseh = new TSEHForMultipleSchedulersTest(mockAppContext, mockClientService, mockEventHandler, mockSigMatcher, mockWebUIService, taskSchedulers, false); tseh.init(conf); tseh.start(); // Verify that the YARN task scheduler is installed by default assertTrue(tseh.getYarnSchedulerCreated()); assertFalse(tseh.getUberSchedulerCreated()); assertEquals(2, tseh.getNumCreateInvocations()); // Verify the order of the schedulers assertEquals(customSchedulerName, tseh.getTaskSchedulerName(0)); assertEquals(TezConstants.getTezYarnServicePluginName(), tseh.getTaskSchedulerName(1)); // Verify the payload setup for the custom task scheduler assertNotNull(tseh.getTaskSchedulerContext(0)); assertEquals(bb, tseh.getTaskSchedulerContext(0).getInitialUserPayload().getPayload()); // Verify the payload on the yarn scheduler assertNotNull(tseh.getTaskSchedulerContext(1)); Configuration parsed = TezUtils .createConfFromUserPayload(tseh.getTaskSchedulerContext(1).getInitialUserPayload()); assertEquals("testval", parsed.get("testkey")); } @Test(timeout = 5000) public void testTaskSchedulerRouting() throws Exception { Configuration conf = new Configuration(false); UserPayload defaultPayload = TezUtils.createUserPayloadFromConf(conf); String customSchedulerName = "fakeScheduler"; List<NamedEntityDescriptor> taskSchedulers = new LinkedList<>(); ByteBuffer bb = ByteBuffer.allocate(4); bb.putInt(0, 3); UserPayload userPayload = UserPayload.create(bb); taskSchedulers.add(new NamedEntityDescriptor(customSchedulerName, FakeTaskScheduler.class.getName()) .setUserPayload(userPayload)); taskSchedulers.add(new NamedEntityDescriptor(TezConstants.getTezYarnServicePluginName(), null) .setUserPayload(defaultPayload)); TSEHForMultipleSchedulersTest tseh = new TSEHForMultipleSchedulersTest(mockAppContext, mockClientService, mockEventHandler, mockSigMatcher, mockWebUIService, taskSchedulers, false); tseh.init(conf); tseh.start(); // Verify that the YARN task scheduler is installed by default assertTrue(tseh.getYarnSchedulerCreated()); assertFalse(tseh.getUberSchedulerCreated()); assertEquals(2, tseh.getNumCreateInvocations()); // Verify the order of the schedulers assertEquals(customSchedulerName, tseh.getTaskSchedulerName(0)); assertEquals(TezConstants.getTezYarnServicePluginName(), tseh.getTaskSchedulerName(1)); verify(tseh.getTestTaskScheduler(0)).initialize(); verify(tseh.getTestTaskScheduler(0)).start(); ApplicationId appId = ApplicationId.newInstance(1000, 1); TezDAGID dagId = TezDAGID.getInstance(appId, 1); TezVertexID vertexID = TezVertexID.getInstance(dagId, 1); TezTaskID taskId1 = TezTaskID.getInstance(vertexID, 1); TezTaskAttemptID attemptId11 = TezTaskAttemptID.getInstance(taskId1, 1); TezTaskID taskId2 = TezTaskID.getInstance(vertexID, 2); TezTaskAttemptID attemptId21 = TezTaskAttemptID.getInstance(taskId2, 1); Resource resource = Resource.newInstance(1024, 1); TaskAttempt mockTaskAttempt1 = mock(TaskAttempt.class); TaskAttempt mockTaskAttempt2 = mock(TaskAttempt.class); AMSchedulerEventTALaunchRequest launchRequest1 = new AMSchedulerEventTALaunchRequest(attemptId11, resource, mock(TaskSpec.class), mockTaskAttempt1, mock(TaskLocationHint.class), 1, mock(ContainerContext.class), 0, 0, 0); tseh.handle(launchRequest1); verify(tseh.getTestTaskScheduler(0)).allocateTask(eq(mockTaskAttempt1), eq(resource), any(String[].class), any(String[].class), any(Priority.class), any(Object.class), eq(launchRequest1)); AMSchedulerEventTALaunchRequest launchRequest2 = new AMSchedulerEventTALaunchRequest(attemptId21, resource, mock(TaskSpec.class), mockTaskAttempt2, mock(TaskLocationHint.class), 1, mock(ContainerContext.class), 1, 0, 0); tseh.handle(launchRequest2); verify(tseh.getTestTaskScheduler(1)).allocateTask(eq(mockTaskAttempt2), eq(resource), any(String[].class), any(String[].class), any(Priority.class), any(Object.class), eq(launchRequest2)); } @SuppressWarnings("unchecked") @Test(timeout = 5000) public void testReportFailureFromTaskScheduler() { String dagName = DAG_NAME; Configuration conf = new TezConfiguration(); String taskSchedulerName = "testTaskScheduler"; String expIdentifier = "[0:" + taskSchedulerName + "]"; EventHandler eventHandler = mock(EventHandler.class); AppContext appContext = mock(AppContext.class, RETURNS_DEEP_STUBS); doReturn(taskSchedulerName).when(appContext).getTaskSchedulerName(0); doReturn(eventHandler).when(appContext).getEventHandler(); doReturn(conf).when(appContext).getAMConf(); InetSocketAddress address = new InetSocketAddress("host", 55000); DAGClientServer dagClientServer = mock(DAGClientServer.class); doReturn(address).when(dagClientServer).getBindAddress(); DAG dag = mock(DAG.class); TezDAGID dagId = TezDAGID.getInstance(ApplicationId.newInstance(1, 0), DAG_INDEX); doReturn(dagName).when(dag).getName(); doReturn(dagId).when(dag).getID(); doReturn(dag).when(appContext).getCurrentDAG(); NamedEntityDescriptor<TaskSchedulerDescriptor> namedEntityDescriptor = new NamedEntityDescriptor<>( taskSchedulerName, TaskSchedulerForFailureTest.class.getName()); List<NamedEntityDescriptor> list = new LinkedList<>(); list.add(namedEntityDescriptor); TaskSchedulerManager taskSchedulerManager = new TaskSchedulerManager(appContext, dagClientServer, eventHandler, mock(ContainerSignatureMatcher.class), mock(WebUIService.class), list, false) { @Override TaskSchedulerContext wrapTaskSchedulerContext(TaskSchedulerContext rawContext) { // Avoid wrapping in threads return rawContext; } }; try { taskSchedulerManager.init(new TezConfiguration()); taskSchedulerManager.start(); taskSchedulerManager.getTotalResources(0); ArgumentCaptor<Event> argumentCaptor = ArgumentCaptor.forClass(Event.class); verify(eventHandler, times(1)).handle(argumentCaptor.capture()); Event rawEvent = argumentCaptor.getValue(); assertTrue(rawEvent instanceof DAGEventTerminateDag); DAGEventTerminateDag killEvent = (DAGEventTerminateDag) rawEvent; assertTrue(killEvent.getDiagnosticInfo().contains("ReportError")); assertTrue( killEvent.getDiagnosticInfo().contains(ServicePluginErrorDefaults.SERVICE_UNAVAILABLE.name())); assertTrue(killEvent.getDiagnosticInfo().contains(expIdentifier)); reset(eventHandler); taskSchedulerManager.getAvailableResources(0); argumentCaptor = ArgumentCaptor.forClass(Event.class); verify(eventHandler, times(1)).handle(argumentCaptor.capture()); rawEvent = argumentCaptor.getValue(); assertTrue(rawEvent instanceof DAGAppMasterEventUserServiceFatalError); DAGAppMasterEventUserServiceFatalError event = (DAGAppMasterEventUserServiceFatalError) rawEvent; assertEquals(DAGAppMasterEventType.TASK_SCHEDULER_SERVICE_FATAL_ERROR, event.getType()); assertTrue(event.getDiagnosticInfo().contains("ReportedFatalError")); assertTrue(event.getDiagnosticInfo().contains(ServicePluginErrorDefaults.INCONSISTENT_STATE.name())); assertTrue(event.getDiagnosticInfo().contains(expIdentifier)); } finally { taskSchedulerManager.stop(); } } @SuppressWarnings("unchecked") @Test(timeout = 5000) public void testTaskSchedulerUserError() { TaskScheduler taskScheduler = mock(TaskScheduler.class, new ExceptionAnswer()); EventHandler eventHandler = mock(EventHandler.class); AppContext appContext = mock(AppContext.class, RETURNS_DEEP_STUBS); when(appContext.getEventHandler()).thenReturn(eventHandler); doReturn("testTaskScheduler").when(appContext).getTaskSchedulerName(0); String expectedId = "[0:testTaskScheduler]"; Configuration conf = new Configuration(false); InetSocketAddress address = new InetSocketAddress(15222); DAGClientServer mockClientService = mock(DAGClientServer.class); doReturn(address).when(mockClientService).getBindAddress(); TaskSchedulerManager taskSchedulerManager = new TaskSchedulerManager(taskScheduler, appContext, mock(ContainerSignatureMatcher.class), mockClientService, Executors.newFixedThreadPool(1)) { @Override protected void instantiateSchedulers(String host, int port, String trackingUrl, AppContext appContext) throws TezException { // Stubbed out since these are setup up front in the constructor used for testing } }; try { taskSchedulerManager.init(conf); taskSchedulerManager.start(); // Invoking a couple of random methods AMSchedulerEventTALaunchRequest launchRequest = new AMSchedulerEventTALaunchRequest( mock(TezTaskAttemptID.class), mock(Resource.class), mock(TaskSpec.class), mock(TaskAttempt.class), mock(TaskLocationHint.class), 0, mock(ContainerContext.class), 0, 0, 0); taskSchedulerManager.handleEvent(launchRequest); ArgumentCaptor<Event> argumentCaptor = ArgumentCaptor.forClass(Event.class); verify(eventHandler, times(1)).handle(argumentCaptor.capture()); Event rawEvent = argumentCaptor.getValue(); assertTrue(rawEvent instanceof DAGAppMasterEventUserServiceFatalError); DAGAppMasterEventUserServiceFatalError event = (DAGAppMasterEventUserServiceFatalError) rawEvent; assertEquals(DAGAppMasterEventType.TASK_SCHEDULER_SERVICE_FATAL_ERROR, event.getType()); assertTrue(event.getError().getMessage().contains("TestException_" + "allocateTask")); assertTrue(event.getDiagnosticInfo().contains("Task Allocation")); assertTrue(event.getDiagnosticInfo().contains(expectedId)); taskSchedulerManager.dagCompleted(); argumentCaptor = ArgumentCaptor.forClass(Event.class); verify(eventHandler, times(2)).handle(argumentCaptor.capture()); rawEvent = argumentCaptor.getAllValues().get(1); assertTrue(rawEvent instanceof DAGAppMasterEventUserServiceFatalError); event = (DAGAppMasterEventUserServiceFatalError) rawEvent; assertEquals(DAGAppMasterEventType.TASK_SCHEDULER_SERVICE_FATAL_ERROR, event.getType()); assertTrue(event.getError().getMessage().contains("TestException_" + "dagComplete")); assertTrue(event.getDiagnosticInfo().contains("Dag Completion")); assertTrue(event.getDiagnosticInfo().contains(expectedId)); } finally { taskSchedulerManager.stop(); } } private static class ExceptionAnswer implements Answer { @Override public Object answer(InvocationOnMock invocation) throws Throwable { Method method = invocation.getMethod(); if (method.getDeclaringClass().equals(TaskScheduler.class) && !method.getName().equals("getContext") && !method.getName().equals("initialize") && !method.getName().equals("start") && !method.getName().equals("shutdown")) { throw new RuntimeException("TestException_" + method.getName()); } else { return invocation.callRealMethod(); } } } public static class TSEHForMultipleSchedulersTest extends TaskSchedulerManager { private final TaskScheduler yarnTaskScheduler; private final TaskScheduler uberTaskScheduler; private final AtomicBoolean uberSchedulerCreated = new AtomicBoolean(false); private final AtomicBoolean yarnSchedulerCreated = new AtomicBoolean(false); private final AtomicInteger numCreateInvocations = new AtomicInteger(0); private final Set<Integer> seenSchedulers = new HashSet<>(); private final List<TaskSchedulerContext> taskSchedulerContexts = new LinkedList<>(); private final List<String> taskSchedulerNames = new LinkedList<>(); private final List<TaskScheduler> testTaskSchedulers = new LinkedList<>(); public TSEHForMultipleSchedulersTest(AppContext appContext, DAGClientServer clientService, EventHandler eventHandler, ContainerSignatureMatcher containerSignatureMatcher, WebUIService webUI, List<NamedEntityDescriptor> schedulerDescriptors, boolean isPureLocalMode) { super(appContext, clientService, eventHandler, containerSignatureMatcher, webUI, schedulerDescriptors, isPureLocalMode); yarnTaskScheduler = mock(TaskScheduler.class); uberTaskScheduler = mock(TaskScheduler.class); } @Override TaskScheduler createTaskScheduler(String host, int port, String trackingUrl, AppContext appContext, NamedEntityDescriptor taskSchedulerDescriptor, long customAppIdIdentifier, int schedulerId) throws TezException { numCreateInvocations.incrementAndGet(); boolean added = seenSchedulers.add(schedulerId); assertTrue("Cannot add multiple schedulers with the same schedulerId", added); taskSchedulerNames.add(taskSchedulerDescriptor.getEntityName()); return super.createTaskScheduler(host, port, trackingUrl, appContext, taskSchedulerDescriptor, customAppIdIdentifier, schedulerId); } @Override TaskSchedulerContext wrapTaskSchedulerContext(TaskSchedulerContext rawContext) { // Avoid wrapping in threads return rawContext; } @Override TaskScheduler createYarnTaskScheduler(TaskSchedulerContext taskSchedulerContext, int schedulerId) { taskSchedulerContexts.add(taskSchedulerContext); testTaskSchedulers.add(yarnTaskScheduler); yarnSchedulerCreated.set(true); return yarnTaskScheduler; } @Override TaskScheduler createUberTaskScheduler(TaskSchedulerContext taskSchedulerContext, int schedulerId) { taskSchedulerContexts.add(taskSchedulerContext); uberSchedulerCreated.set(true); testTaskSchedulers.add(yarnTaskScheduler); return uberTaskScheduler; } @Override TaskScheduler createCustomTaskScheduler(TaskSchedulerContext taskSchedulerContext, NamedEntityDescriptor taskSchedulerDescriptor, int schedulerId) throws TezException { taskSchedulerContexts.add(taskSchedulerContext); TaskScheduler taskScheduler = spy( super.createCustomTaskScheduler(taskSchedulerContext, taskSchedulerDescriptor, schedulerId)); testTaskSchedulers.add(taskScheduler); return taskScheduler; } @Override // Inline handling of events. public void handle(AMSchedulerEvent event) { handleEvent(event); } public boolean getUberSchedulerCreated() { return uberSchedulerCreated.get(); } public boolean getYarnSchedulerCreated() { return yarnSchedulerCreated.get(); } public int getNumCreateInvocations() { return numCreateInvocations.get(); } public TaskSchedulerContext getTaskSchedulerContext(int schedulerId) { return taskSchedulerContexts.get(schedulerId); } public String getTaskSchedulerName(int schedulerId) { return taskSchedulerNames.get(schedulerId); } public TaskScheduler getTestTaskScheduler(int schedulerId) { return testTaskSchedulers.get(schedulerId); } } public static class FakeTaskScheduler extends TaskScheduler { public FakeTaskScheduler(TaskSchedulerContext taskSchedulerContext) { super(taskSchedulerContext); } @Override public Resource getAvailableResources() { return null; } @Override public int getClusterNodeCount() { return 0; } @Override public void dagComplete() { } @Override public Resource getTotalResources() { return null; } @Override public void blacklistNode(NodeId nodeId) { } @Override public void unblacklistNode(NodeId nodeId) { } @Override public void allocateTask(Object task, Resource capability, String[] hosts, String[] racks, Priority priority, Object containerSignature, Object clientCookie) { } @Override public void allocateTask(Object task, Resource capability, ContainerId containerId, Priority priority, Object containerSignature, Object clientCookie) { } @Override public boolean deallocateTask(Object task, boolean taskSucceeded, TaskAttemptEndReason endReason, String diagnostics) { return false; } @Override public Object deallocateContainer(ContainerId containerId) { return null; } @Override public void setShouldUnregister() { } @Override public boolean hasUnregistered() { return false; } } private static final String DAG_NAME = "dagName"; private static final int DAG_INDEX = 1; public static class TaskSchedulerForFailureTest extends TaskScheduler { public TaskSchedulerForFailureTest(TaskSchedulerContext taskSchedulerContext) { super(taskSchedulerContext); } @Override public Resource getAvailableResources() throws ServicePluginException { getContext().reportError(ServicePluginErrorDefaults.INCONSISTENT_STATE, "ReportedFatalError", null); return Resource.newInstance(1024, 1); } @Override public Resource getTotalResources() throws ServicePluginException { getContext().reportError(ServicePluginErrorDefaults.SERVICE_UNAVAILABLE, "ReportError", new DagInfoImplForTest(DAG_INDEX, DAG_NAME)); return Resource.newInstance(1024, 1); } @Override public int getClusterNodeCount() throws ServicePluginException { return 0; } @Override public void blacklistNode(NodeId nodeId) throws ServicePluginException { } @Override public void unblacklistNode(NodeId nodeId) throws ServicePluginException { } @Override public void allocateTask(Object task, Resource capability, String[] hosts, String[] racks, Priority priority, Object containerSignature, Object clientCookie) throws ServicePluginException { } @Override public void allocateTask(Object task, Resource capability, ContainerId containerId, Priority priority, Object containerSignature, Object clientCookie) throws ServicePluginException { } @Override public boolean deallocateTask(Object task, boolean taskSucceeded, TaskAttemptEndReason endReason, @Nullable String diagnostics) throws ServicePluginException { return false; } @Override public Object deallocateContainer(ContainerId containerId) throws ServicePluginException { return null; } @Override public void setShouldUnregister() throws ServicePluginException { } @Override public boolean hasUnregistered() throws ServicePluginException { return false; } @Override public void dagComplete() throws ServicePluginException { } } }