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);
TezVertexID vertexID21 = TezVertexID.getInstance(dagID2, 1);
//Vertex 2, Task 1, Attempt 1, host1, lr2
TezTaskAttemptID taID211 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID21, 1), 1);
TaskAttempt ta211 = mock(TaskAttempt.class);
AMSchedulerEventTALaunchRequest lrEvent21 = createLaunchRequestEvent(taID211, ta211, resource1, host1, racks, priority1, dag2LRs);
//Vertex 2, Task 2, Attempt 1, host1, lr2
TezTaskAttemptID taID212 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID21, 2), 1);
TaskAttempt ta212 = mock(TaskAttempt.class);
AMSchedulerEventTALaunchRequest lrEvent22 = createLaunchRequestEvent(taID212, ta212, resource1, host1, racks, priority1, dag2LRs);
taskSchedulerEventHandler.handleEvent(lrEvent21);
taskSchedulerEventHandler.handleEvent(lrEvent22);
drainableAppCallback.drain();
// TODO This is terrible, need a better way to ensure the scheduling loop has run
LOG.info("Sleeping to ensure that the scheduling loop runs");
Thread.sleep(6000l);
verify(taskSchedulerEventHandler).taskAllocated(eq(ta211), any(Object.class), eq(container1));
verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));
eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
assignEvent = (AMContainerEventAssignTA) eventHandler.verifyInvocation(AMContainerEventAssignTA.class);
assertEquals(2, assignEvent.getRemoteTaskLocalResources().size());
eventHandler.reset();
taskSchedulerEventHandler.handleEvent(new AMSchedulerEventTAEnded(ta211, container1.getId(), TaskAttemptState.SUCCEEDED));
drainableAppCallback.drain();
verify(taskScheduler).deallocateTask(eq(ta211), eq(true));
verify(taskSchedulerEventHandler).taskAllocated(eq(ta212), any(Object.class), eq(container1));
verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));
eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
assignEvent = (AMContainerEventAssignTA) eventHandler.verifyInvocation(AMContainerEventAssignTA.class);
assertEquals(2, assignEvent.getRemoteTaskLocalResources().size());
eventHandler.reset();
taskScheduler.close();
taskSchedulerEventHandler.close();
}