tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, -1);
RackResolver.init(tezConf);
TaskSchedulerAppCallback mockApp = mock(TaskSchedulerAppCallback.class);
CapturingEventHandler eventHandler = new CapturingEventHandler();
TezDAGID dagID1 = TezDAGID.getInstance("0", 1, 0);
AMRMClient<CookieContainerRequest> rmClientCore = new AMRMClientForTest();
TezAMRMClientAsync<CookieContainerRequest> rmClient = spy(new AMRMClientAsyncForTest(rmClientCore, 100));
String appUrl = "url";
String appMsg = "success";
AppFinalStatus finalStatus =
new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
doReturn(finalStatus).when(mockApp).getFinalAppStatus();
AppContext appContext = mock(AppContext.class);
ChangingDAGIDAnswer dagIDAnswer = new ChangingDAGIDAnswer(dagID1);
AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class),
mock(TaskAttemptListener.class), new ContainerContextMatcher(), appContext);
AMNodeMap amNodeMap = new AMNodeMap(eventHandler, appContext);
doReturn(amContainerMap).when(appContext).getAllContainers();
doReturn(amNodeMap).when(appContext).getAllNodes();
doReturn(DAGAppMasterState.RUNNING).when(appContext).getAMState();
doReturn(true).when(appContext).isSession();
doAnswer(dagIDAnswer).when(appContext).getCurrentDAGID();
doReturn(mock(ClusterInfo.class)).when(appContext).getClusterInfo();
TaskSchedulerEventHandler taskSchedulerEventHandlerReal = new TaskSchedulerEventHandlerForTest(appContext, eventHandler, rmClient, new AlwaysMatchesContainerMatcher());
TaskSchedulerEventHandler taskSchedulerEventHandler = spy(taskSchedulerEventHandlerReal);
taskSchedulerEventHandler.init(tezConf);
taskSchedulerEventHandler.start();
TaskSchedulerWithDrainableAppCallback taskScheduler = (TaskSchedulerWithDrainableAppCallback) ((TaskSchedulerEventHandlerForTest) taskSchedulerEventHandler)
.getSpyTaskScheduler();
TaskSchedulerAppCallbackDrainable drainableAppCallback = taskScheduler.getDrainableAppCallback();
AtomicBoolean drainNotifier = new AtomicBoolean(false);
taskScheduler.delayedContainerManager.drainedDelayedContainersForTest = drainNotifier;
Resource resource1 = Resource.newInstance(1024, 1);
String[] host1 = {"host1"};
String []racks = {"/default-rack"};
Priority priority1 = Priority.newInstance(1);
String rsrc1 = "rsrc1";
String rsrc2 = "rsrc2";
String rsrc3 = "rsrc3";
LocalResource lr1 = mock(LocalResource.class);
LocalResource lr2 = mock(LocalResource.class);
LocalResource lr3 = mock(LocalResource.class);
AMContainerEventAssignTA assignEvent = null;
Map<String, LocalResource> dag1LRs = Maps.newHashMap();
dag1LRs.put(rsrc1, lr1);
TezVertexID vertexID11 = TezVertexID.getInstance(dagID1, 1);
//Vertex 1, Task 1, Attempt 1, host1, lr1
TezTaskAttemptID taID111 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID11, 1), 1);
TaskAttempt ta111 = mock(TaskAttempt.class);
AMSchedulerEventTALaunchRequest lrEvent11 = createLaunchRequestEvent(taID111, ta111, resource1, host1, racks, priority1, dag1LRs);
//Vertex 1, Task 2, Attempt 1, host1, lr1
TezTaskAttemptID taID112 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID11, 2), 1);
TaskAttempt ta112 = mock(TaskAttempt.class);
AMSchedulerEventTALaunchRequest lrEvent12 = createLaunchRequestEvent(taID112, ta112, resource1, host1, racks, priority1, dag1LRs);
taskSchedulerEventHandler.handleEvent(lrEvent11);
taskSchedulerEventHandler.handleEvent(lrEvent12);
Container container1 = createContainer(1, "host1", resource1, priority1);
// One container allocated.
drainNotifier.set(false);
taskScheduler.onContainersAllocated(Collections.singletonList(container1));
TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
drainableAppCallback.drain();
verify(taskSchedulerEventHandler).taskAllocated(eq(ta111), any(Object.class), eq(container1));
assignEvent = (AMContainerEventAssignTA) eventHandler.verifyInvocation(AMContainerEventAssignTA.class);
assertEquals(1, assignEvent.getRemoteTaskLocalResources().size());
// Task assigned to container completed successfully. Container should be re-used.
taskSchedulerEventHandler.handleEvent(new AMSchedulerEventTAEnded(ta111, container1.getId(), TaskAttemptState.SUCCEEDED));
drainableAppCallback.drain();
verify(taskScheduler).deallocateTask(eq(ta111), eq(true));
verify(taskSchedulerEventHandler).taskAllocated(eq(ta112), any(Object.class), eq(container1));
verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));
eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
assignEvent = (AMContainerEventAssignTA) eventHandler.verifyInvocation(AMContainerEventAssignTA.class);
assertEquals(1, assignEvent.getRemoteTaskLocalResources().size());
eventHandler.reset();
// Task assigned to container completed successfully.
// Verify reuse across hosts.
taskSchedulerEventHandler.handleEvent(new AMSchedulerEventTAEnded(ta112, container1.getId(), TaskAttemptState.SUCCEEDED));
drainableAppCallback.drain();
verify(taskScheduler).deallocateTask(eq(ta112), eq(true));
verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));
eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
eventHandler.reset();
// Setup DAG2 with additional resources. Make sure the container, even without all resources, is reused.
TezDAGID dagID2 = TezDAGID.getInstance("0", 2, 0);
dagIDAnswer.setDAGID(dagID2);
Map<String, LocalResource> dag2LRs = Maps.newHashMap();
dag2LRs.put(rsrc2, lr2);
dag2LRs.put(rsrc3, lr3);