scheduler.allocateTask(mockTask1, mockCapability, hosts,
racks, mockPriority1, null, mockCookie1);
drainableAppCallback.drain();
verify(mockRMClient, times(1)).
addContainerRequest(requestCaptor.capture());
CookieContainerRequest request1 = requestCaptor.getValue();
scheduler.allocateTask(mockTask2, mockCapability, hosts,
racks, mockPriority2, null, mockCookie2);
drainableAppCallback.drain();
verify(mockRMClient, times(2)).
addContainerRequest(requestCaptor.capture());
CookieContainerRequest request2 = requestCaptor.getValue();
scheduler.allocateTask(mockTask3, mockCapability, hosts,
racks, mockPriority3, null, mockCookie3);
drainableAppCallback.drain();
verify(mockRMClient, times(3)).
addContainerRequest(requestCaptor.capture());
CookieContainerRequest request3 = requestCaptor.getValue();
List<Container> containers = new ArrayList<Container>();
// sending lower priority container first to make sure its not matched
Container mockContainer4 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer4.getNodeId().getHost()).thenReturn("host4");
when(mockContainer4.toString()).thenReturn("container4");
when(mockContainer4.getPriority()).thenReturn(mockPriority4);
ContainerId mockCId4 = mock(ContainerId.class);
when(mockContainer4.getId()).thenReturn(mockCId4);
when(mockCId4.toString()).thenReturn("container4");
containers.add(mockContainer4);
Container mockContainer1 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer1.getNodeId().getHost()).thenReturn("host1");
when(mockContainer1.getPriority()).thenReturn(mockPriority1);
when(mockContainer1.toString()).thenReturn("container1");
ContainerId mockCId1 = mock(ContainerId.class);
when(mockContainer1.getId()).thenReturn(mockCId1);
when(mockCId1.toString()).thenReturn("container1");
containers.add(mockContainer1);
Container mockContainer2 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer2.getNodeId().getHost()).thenReturn("host2");
when(mockContainer2.getPriority()).thenReturn(mockPriority2);
when(mockContainer2.toString()).thenReturn("container2");
ContainerId mockCId2 = mock(ContainerId.class);
when(mockContainer2.getId()).thenReturn(mockCId2);
when(mockCId2.toString()).thenReturn("container2");
containers.add(mockContainer2);
Container mockContainer3 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer3.getNodeId().getHost()).thenReturn("host3");
when(mockContainer3.getPriority()).thenReturn(mockPriority3);
when(mockContainer3.toString()).thenReturn("container3");
ContainerId mockCId3 = mock(ContainerId.class);
when(mockContainer3.getId()).thenReturn(mockCId3);
when(mockCId3.toString()).thenReturn("container3");
containers.add(mockContainer3);
ArrayList<CookieContainerRequest> hostContainers =
new ArrayList<CookieContainerRequest>();
hostContainers.add(request1);
ArrayList<CookieContainerRequest> rackContainers =
new ArrayList<CookieContainerRequest>();
rackContainers.add(request2);
ArrayList<CookieContainerRequest> anyContainers =
new ArrayList<CookieContainerRequest>();
anyContainers.add(request3);
final List<ArrayList<CookieContainerRequest>> hostList =
new LinkedList<ArrayList<CookieContainerRequest>>();
hostList.add(hostContainers);
final List<ArrayList<CookieContainerRequest>> rackList =
new LinkedList<ArrayList<CookieContainerRequest>>();
rackList.add(rackContainers);
final List<ArrayList<CookieContainerRequest>> anyList =
new LinkedList<ArrayList<CookieContainerRequest>>();
anyList.add(anyContainers);
final List<ArrayList<CookieContainerRequest>> emptyList =
new LinkedList<ArrayList<CookieContainerRequest>>();
// return pri1 requests for host1
when(
mockRMClient.getMatchingRequestsForTopPriority(eq("host1"),
(Resource) any())).thenAnswer(
new Answer<List<? extends Collection<CookieContainerRequest>>>() {
@Override
public List<? extends Collection<CookieContainerRequest>> answer(
InvocationOnMock invocation) throws Throwable {
return hostList;
}
});
// second request matched to rack. RackResolver by default puts hosts in
// /default-rack. We need to workaround by returning rack matches only once
when(
mockRMClient.getMatchingRequestsForTopPriority(eq("/default-rack"),
(Resource) any())).thenAnswer(
new Answer<List<? extends Collection<CookieContainerRequest>>>() {
@Override
public List<? extends Collection<CookieContainerRequest>> answer(
InvocationOnMock invocation) throws Throwable {
return rackList;
}
}).thenAnswer(
new Answer<List<? extends Collection<CookieContainerRequest>>>() {
@Override
public List<? extends Collection<CookieContainerRequest>> answer(
InvocationOnMock invocation) throws Throwable {
return emptyList;
}
});
// third request matched to ANY
when(
mockRMClient.getMatchingRequestsForTopPriority(
eq(ResourceRequest.ANY), (Resource) any())).thenAnswer(
new Answer<List<? extends Collection<CookieContainerRequest>>>() {
@Override
public List<? extends Collection<CookieContainerRequest>> answer(
InvocationOnMock invocation) throws Throwable {
return anyList;
}
}).thenAnswer(
new Answer<List<? extends Collection<CookieContainerRequest>>>() {
@Override
public List<? extends Collection<CookieContainerRequest>> answer(
InvocationOnMock invocation) throws Throwable {
return emptyList;
}
});
when(mockRMClient.getTopPriority()).then(
new Answer<Priority>() {
@Override
public Priority answer(
InvocationOnMock invocation) throws Throwable {
int allocations = drainableAppCallback.count.get();
if (allocations == 0) {
return mockPriority1;
}
if (allocations == 1) {
return mockPriority2;
}
if (allocations == 2) {
return mockPriority3;
}
if (allocations == 3) {
return mockPriority4;
}
return null;
}
});
AtomicBoolean drainNotifier = new AtomicBoolean(false);
scheduler.delayedContainerManager.drainedDelayedContainers = drainNotifier;
scheduler.onContainersAllocated(containers);
TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
drainableAppCallback.drain();
// exact number allocations returned
verify(mockApp, times(3)).taskAllocated(any(), any(), (Container) any());
// first container allocated
verify(mockApp).taskAllocated(mockTask1, mockCookie1, mockContainer1);
verify(mockApp).taskAllocated(mockTask2, mockCookie2, mockContainer2);
verify(mockApp).taskAllocated(mockTask3, mockCookie3, mockContainer3);
verify(mockRMClient).removeContainerRequest(request1);
verify(mockRMClient).removeContainerRequest(request2);
verify(mockRMClient).removeContainerRequest(request3);
// verify unwanted container released
verify(mockRMClient).releaseAssignedContainer(mockCId4);
// deallocate allocated task
assertTrue(scheduler.deallocateTask(mockTask1, true));
drainableAppCallback.drain();
verify(mockApp).containerBeingReleased(mockCId1);
verify(mockRMClient).releaseAssignedContainer(mockCId1);
// deallocate allocated container
Assert.assertEquals(mockTask2, scheduler.deallocateContainer(mockCId2));
drainableAppCallback.drain();
verify(mockRMClient).releaseAssignedContainer(mockCId2);
verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any());
List<ContainerStatus> statuses = new ArrayList<ContainerStatus>();
ContainerStatus mockStatus1 = mock(ContainerStatus.class);
when(mockStatus1.getContainerId()).thenReturn(mockCId1);
statuses.add(mockStatus1);
ContainerStatus mockStatus2 = mock(ContainerStatus.class);
when(mockStatus2.getContainerId()).thenReturn(mockCId2);
statuses.add(mockStatus2);
ContainerStatus mockStatus3 = mock(ContainerStatus.class);
when(mockStatus3.getContainerId()).thenReturn(mockCId3);
statuses.add(mockStatus3);
ContainerStatus mockStatus4 = mock(ContainerStatus.class);
when(mockStatus4.getContainerId()).thenReturn(mockCId4);
statuses.add(mockStatus4);
scheduler.onContainersCompleted(statuses);
drainableAppCallback.drain();
// released container status returned
verify(mockApp).containerCompleted(mockTask1, mockStatus1);
verify(mockApp).containerCompleted(mockTask2, mockStatus2);
// currently allocated container status returned and not released
verify(mockApp).containerCompleted(mockTask3, mockStatus3);
// no other statuses returned
verify(mockApp, times(3)).containerCompleted(any(), (ContainerStatus) any());
verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any());
// verify blacklisting
verify(mockRMClient, times(0)).addNodeToBlacklist((NodeId)any());
String badHost = "host6";
NodeId badNodeId = mock(NodeId.class);
when(badNodeId.getHost()).thenReturn(badHost);
scheduler.blacklistNode(badNodeId);
verify(mockRMClient, times(1)).addNodeToBlacklist(badNodeId);
Object mockTask4 = mock(Object.class);
when(mockTask4.toString()).thenReturn("task4");
Object mockCookie4 = mock(Object.class);
scheduler.allocateTask(mockTask4, mockCapability, null,
null, mockPriority4, null, mockCookie4);
drainableAppCallback.drain();
verify(mockRMClient, times(4)).addContainerRequest(requestCaptor.capture());
CookieContainerRequest request4 = requestCaptor.getValue();
anyContainers.clear();
anyContainers.add(request4);
Container mockContainer5 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer5.getNodeId().getHost()).thenReturn(badHost);
when(mockContainer5.getNodeId()).thenReturn(badNodeId);
ContainerId mockCId5 = mock(ContainerId.class);
when(mockContainer5.toString()).thenReturn("container5");
when(mockCId5.toString()).thenReturn("container5");
when(mockContainer5.getId()).thenReturn(mockCId5);
when(mockContainer5.getPriority()).thenReturn(mockPriority4);
containers.clear();
containers.add(mockContainer5);
when(
mockRMClient.getMatchingRequestsForTopPriority(
eq(ResourceRequest.ANY), (Resource) any())).thenAnswer(
new Answer<List<? extends Collection<CookieContainerRequest>>>() {
@Override
public List<? extends Collection<CookieContainerRequest>> answer(
InvocationOnMock invocation) throws Throwable {
return anyList;
}
}).thenAnswer(
new Answer<List<? extends Collection<CookieContainerRequest>>>() {
@Override
public List<? extends Collection<CookieContainerRequest>> answer(
InvocationOnMock invocation) throws Throwable {
return emptyList;
}
});
drainNotifier.set(false);
scheduler.onContainersAllocated(containers);
TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
drainableAppCallback.drain();
// no new allocation
verify(mockApp, times(3)).taskAllocated(any(), any(), (Container) any());
// verify blacklisted container released
verify(mockRMClient).releaseAssignedContainer(mockCId5);
verify(mockRMClient, times(4)).releaseAssignedContainer((ContainerId) any());
// verify request added back
verify(mockRMClient, times(5)).addContainerRequest(requestCaptor.capture());
CookieContainerRequest request5 = requestCaptor.getValue();
anyContainers.clear();
anyContainers.add(request5);
Container mockContainer6 = mock(Container.class, RETURNS_DEEP_STUBS);
when(mockContainer6.getNodeId().getHost()).thenReturn("host7");
ContainerId mockCId6 = mock(ContainerId.class);