/*
 * 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 com.google.common.collect.Sets;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.yarn.api.protocolrecords.RegisterApplicationMasterResponse;
import org.apache.hadoop.yarn.api.records.ApplicationAccessType;
import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.api.records.Container;
import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.hadoop.yarn.api.records.ContainerStatus;
import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;
import org.apache.hadoop.yarn.api.records.NodeId;
import org.apache.hadoop.yarn.api.records.NodeReport;
import org.apache.hadoop.yarn.api.records.Priority;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.client.api.async.AMRMClientAsync;
import org.apache.hadoop.yarn.client.api.impl.AMRMClientImpl;
import org.apache.hadoop.yarn.exceptions.YarnException;
import org.apache.hadoop.yarn.proto.YarnServiceProtos.SchedulerResourceTypes;
import org.apache.hadoop.yarn.util.resource.Resources;
import org.apache.tez.common.MockDNSToSwitchMapping;
import org.apache.tez.dag.api.TezConfiguration;
import org.apache.tez.dag.app.MockClock;
import org.apache.tez.dag.app.dag.Task;
import org.apache.tez.dag.app.dag.TaskAttempt;
import org.apache.tez.dag.app.rm.DagAwareYarnTaskScheduler.AMRMClientAsyncWrapper;
import org.apache.tez.dag.app.rm.DagAwareYarnTaskScheduler.HeldContainer;
import org.apache.tez.dag.app.rm.DagAwareYarnTaskScheduler.TaskRequest;
import org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerContextDrainable;
import org.apache.tez.serviceplugins.api.DagInfo;
import org.apache.tez.serviceplugins.api.TaskScheduler;
import org.apache.tez.serviceplugins.api.TaskSchedulerContext;
import org.apache.tez.serviceplugins.api.TaskSchedulerContext.AppFinalStatus;
import org.apache.tez.test.ControlledScheduledExecutorService;
import org.junit.After;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;
import org.mockito.ArgumentCaptor;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.BitSet;
import java.util.Collections;
import java.util.EnumSet;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;

import static org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.createCountingExecutingService;
import static org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.setupMockTaskSchedulerContext;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyInt;
import static org.mockito.Matchers.anyListOf;
import static org.mockito.Matchers.eq;
import static org.mockito.Matchers.isNull;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

public class TestDagAwareYarnTaskScheduler {
  private ExecutorService contextCallbackExecutor;

  @BeforeClass
  public static void beforeClass() {

    MockDNSToSwitchMapping.initializeMockRackResolver();
  }

  @Before
  public void preTest() {
    contextCallbackExecutor = Executors.newSingleThreadExecutor(
        new ThreadFactoryBuilder().setNameFormat("TaskSchedulerAppCallbackExecutor #%d")
            .setDaemon(true)
            .build());
  }

  @After
  public void postTest() {
    contextCallbackExecutor.shutdownNow();
  }

  private TaskSchedulerContextDrainable createDrainableContext(
      TaskSchedulerContext taskSchedulerContext) {
    TaskSchedulerContextImplWrapper wrapper =
        new TaskSchedulerContextImplWrapper(taskSchedulerContext,
            createCountingExecutingService(contextCallbackExecutor));
    return new TaskSchedulerContextDrainable(wrapper);
  }

  @SuppressWarnings({ "unchecked" })
  @Test(timeout=30000)
  public void testNoReuse() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback,
        mockRMClient, clock);

    scheduler.initialize();
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
        regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
        regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Object mockTask1 = new MockTask("task1");
    Object mockCookie1 = new Object();
    Resource mockCapability = Resources.createResource(1024, 1);
    String[] hosts = {"host1", "host5"};
    String[] racks = {"/default-rack", "/default-rack"};
    Priority mockPriority = Priority.newInstance(1);
    ArgumentCaptor<TaskRequest> requestCaptor =
        ArgumentCaptor.forClass(TaskRequest.class);
    // allocate task
    scheduler.allocateTask(mockTask1, mockCapability, hosts,
        racks, mockPriority, null, mockCookie1);
    drainableAppCallback.drain();
    verify(mockRMClient, times(1)).
        addContainerRequest(any(TaskRequest.class));

    // returned from task requests before allocation happens
    assertFalse(scheduler.deallocateTask(mockTask1, true, null, null));
    verify(mockApp, times(0)).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, times(1)).
        removeContainerRequest(any(TaskRequest.class));
    verify(mockRMClient, times(0)).
        releaseAssignedContainer((ContainerId) any());

    // deallocating unknown task
    assertFalse(scheduler.deallocateTask(mockTask1, true, null, null));
    verify(mockApp, times(0)).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, times(1)).
        removeContainerRequest(any(TaskRequest.class));
    verify(mockRMClient, times(0)).
        releaseAssignedContainer((ContainerId) any());

    // allocate tasks
    Object mockTask2 = new MockTask("task2");
    Object mockCookie2 = new Object();
    Object mockTask3 = new MockTask("task3");
    Object mockCookie3 = new Object();
    scheduler.allocateTask(mockTask1, mockCapability, hosts,
        racks, mockPriority, null, mockCookie1);
    drainableAppCallback.drain();
    verify(mockRMClient, times(2)).
        addContainerRequest(requestCaptor.capture());
    TaskRequest request1 = requestCaptor.getValue();
    scheduler.allocateTask(mockTask2, mockCapability, hosts,
        racks, mockPriority, null, mockCookie2);
    drainableAppCallback.drain();
    verify(mockRMClient, times(3)).
        addContainerRequest(requestCaptor.capture());
    TaskRequest request2 = requestCaptor.getValue();
    scheduler.allocateTask(mockTask3, mockCapability, hosts,
        racks, mockPriority, null, mockCookie3);
    drainableAppCallback.drain();
    verify(mockRMClient, times(4)).
        addContainerRequest(requestCaptor.capture());
    TaskRequest request3 = requestCaptor.getValue();

    NodeId host1 = NodeId.newInstance("host1", 1);
    NodeId host2 = NodeId.newInstance("host2", 2);
    NodeId host3 = NodeId.newInstance("host3", 3);
    NodeId host4 = NodeId.newInstance("host4", 4);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId mockCId1 = ContainerId.newContainerId(attemptId, 1);
    Container mockContainer1 = Container.newInstance(mockCId1, host1, null, mockCapability, mockPriority, null);
    ContainerId mockCId2 = ContainerId.newContainerId(attemptId, 2);
    Container mockContainer2 = Container.newInstance(mockCId2, host2, null, mockCapability, mockPriority, null);
    ContainerId mockCId3 = ContainerId.newContainerId(attemptId, 3);
    Container mockContainer3 = Container.newInstance(mockCId3, host3, null, mockCapability, mockPriority, null);
    ContainerId mockCId4 = ContainerId.newContainerId(attemptId, 4);
    Container mockContainer4 = Container.newInstance(mockCId4, host4, null, mockCapability, mockPriority, null);
    List<Container> containers = new ArrayList<>();
    containers.add(mockContainer1);
    containers.add(mockContainer2);
    containers.add(mockContainer3);
    containers.add(mockContainer4);
    scheduler.onContainersAllocated(containers);
    drainableAppCallback.drain();
    // first container allocated
    verify(mockApp).taskAllocated(mockTask1, mockCookie1, mockContainer1);
    verify(mockApp).taskAllocated(mockTask2, mockCookie2, mockContainer2);
    verify(mockApp).taskAllocated(mockTask3, mockCookie3, mockContainer3);
    // no other allocations returned
    verify(mockApp, times(3)).taskAllocated(any(), any(), (Container) any());
    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, null, null));
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(mockCId1);
    verify(mockRMClient).releaseAssignedContainer(mockCId1);
    // deallocate allocated container
    assertEquals(mockTask2, scheduler.deallocateContainer(mockCId2));
    drainableAppCallback.drain();
    verify(mockRMClient).releaseAssignedContainer(mockCId2);
    verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any());

    List<ContainerStatus> statuses = new ArrayList<>();
    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)).updateBlacklist(anyListOf(String.class), anyListOf(String.class));
    String badHost = "host6";
    NodeId badNodeId = NodeId.newInstance(badHost, 1);
    scheduler.blacklistNode(badNodeId);
    List<String> badNodeList = Collections.singletonList(badHost);
    verify(mockRMClient, times(1)).updateBlacklist(eq(badNodeList), isNull(List.class));
    Object mockTask4 = new MockTask("task4");
    Object mockCookie4 = new Object();
    scheduler.allocateTask(mockTask4, mockCapability, null,
        null, mockPriority, null, mockCookie4);
    drainableAppCallback.drain();
    verify(mockRMClient, times(5)).addContainerRequest(requestCaptor.capture());
    ContainerId mockCId5 = ContainerId.newContainerId(attemptId, 5);
    Container mockContainer5 = Container.newInstance(mockCId5, badNodeId, null, mockCapability, mockPriority, null);
    containers.clear();
    containers.add(mockContainer5);
    scheduler.onContainersAllocated(containers);
    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(6)).addContainerRequest(requestCaptor.capture());
    NodeId host6 = NodeId.newInstance("host6", 6);
    ContainerId mockCId6 = ContainerId.newContainerId(attemptId, 6);
    Container mockContainer6 = Container.newInstance(mockCId6, host6, null, mockCapability, mockPriority, null);
    containers.clear();
    containers.add(mockContainer6);
    scheduler.onContainersAllocated(containers);
    drainableAppCallback.drain();
    // new allocation
    verify(mockApp, times(4)).taskAllocated(any(), any(), (Container) any());
    verify(mockApp).taskAllocated(mockTask4, mockCookie4, mockContainer6);
    // deallocate allocated task
    assertTrue(scheduler.deallocateTask(mockTask4, true, null, null));
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(mockCId6);
    verify(mockRMClient).releaseAssignedContainer(mockCId6);
    verify(mockRMClient, times(5)).releaseAssignedContainer((ContainerId) any());
    // test unblacklist
    scheduler.unblacklistNode(badNodeId);
    verify(mockRMClient, times(1)).updateBlacklist(isNull(List.class), eq(badNodeList));
    assertEquals(0, scheduler.getNumBlacklistedNodes());

    float progress = 0.5f;
    when(mockApp.getProgress()).thenReturn(progress);
    assertEquals(progress, scheduler.getProgress(), 0);

    // check duplicate allocation request
    scheduler.allocateTask(mockTask1, mockCapability, hosts, racks,
        mockPriority, null, mockCookie1);
    drainableAppCallback.drain();
    verify(mockRMClient, times(7)).addContainerRequest(any(TaskRequest.class));
    verify(mockRMClient, times(6)).
        removeContainerRequest(any(TaskRequest.class));
    scheduler.allocateTask(mockTask1, mockCapability, hosts, racks,
        mockPriority, null, mockCookie1);
    drainableAppCallback.drain();
    // old request removed and new one added
    verify(mockRMClient, times(7)).
        removeContainerRequest(any(TaskRequest.class));
    verify(mockRMClient, times(8)).addContainerRequest(any(TaskRequest.class));
    assertFalse(scheduler.deallocateTask(mockTask1, true, null, null));

    // test speculative node adjustment
    String speculativeNode = "host8";
    NodeId speculativeNodeId = mock(NodeId.class);
    when(speculativeNodeId.getHost()).thenReturn(speculativeNode);
    TaskAttempt mockTask5 = mock(TaskAttempt.class);
    Task task = mock(Task.class);
    when(mockTask5.getTask()).thenReturn(task);
    when(task.getNodesWithRunningAttempts()).thenReturn(Sets.newHashSet(speculativeNodeId));
    Object mockCookie5 = new Object();
    scheduler.allocateTask(mockTask5, mockCapability, hosts, racks,
        mockPriority, null, mockCookie5);
    drainableAppCallback.drain();
    // no new allocation
    verify(mockApp, times(4)).taskAllocated(any(), any(), (Container) any());
    // verify container released
    verify(mockRMClient, times(5)).releaseAssignedContainer((ContainerId) any());
    // verify request added back
    verify(mockRMClient, times(9)).addContainerRequest(requestCaptor.capture());

    List<NodeReport> mockUpdatedNodes = mock(List.class);
    scheduler.onNodesUpdated(mockUpdatedNodes);
    drainableAppCallback.drain();
    verify(mockApp).nodesUpdated(mockUpdatedNodes);

    ArgumentCaptor<String> argumentCaptor = ArgumentCaptor.forClass(String.class);
    Exception mockException = new IOException("mockexception");
    scheduler.onError(mockException);
    drainableAppCallback.drain();
    verify(mockApp)
        .reportError(eq(YarnTaskSchedulerServiceError.RESOURCEMANAGER_ERROR), argumentCaptor.capture(),
            any(DagInfo.class));
    assertTrue(argumentCaptor.getValue().contains("mockexception"));

    scheduler.onShutdownRequest();
    drainableAppCallback.drain();
    verify(mockApp).appShutdownRequested();

    String appMsg = "success";
    AppFinalStatus finalStatus =
        new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).
        unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED,
            appMsg, appUrl);
    verify(mockRMClient).stop();
  }

  @Test(timeout=30000)
  public void testSimpleReuseLocalMatching() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, false);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback,
        mockRMClient, clock);

    scheduler.initialize();
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
        regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
        regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv1 = Priority.newInstance(2);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, "host3");
    MockTaskInfo taskv0t2 = new MockTaskInfo("taskv0t2", priorityv0, hostsv0t0);
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, hostsv0t0);
    MockTaskInfo taskv1t1 = new MockTaskInfo("taskv1t1", priorityv1, hostsv0t0);

    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient,
        scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    taskRequestCaptor.scheduleTask(taskv0t1);
    TaskRequest reqv0t2 = taskRequestCaptor.scheduleTask(taskv0t2);
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0);
    taskRequestCaptor.scheduleTask(taskv1t1);

    NodeId host1 = NodeId.newInstance("host1", 1);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    Container container1 = Container.newInstance(cid1, host1, null, taskv0t0.capability, priorityv0, null);

    // allocate one container at v0 priority
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);

    // finish v0t0 successfully, verify v0t1 is skipped and v0t2 instead is assigned to the container
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv0t2.task, taskv0t2.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t2);

    // finish v0t2 successfully, verify v1t0 is assigned to the same container
    assertTrue(scheduler.deallocateTask(taskv0t2.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv1t0.task, taskv1t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv1t0);

    // fail v1t0 and verify container is released instead of reused for v1t1
    assertTrue(scheduler.deallocateTask(taskv1t0.task, false, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(cid1);
    verify(mockRMClient).releaseAssignedContainer(cid1);

    String appMsg = "success";
    AppFinalStatus finalStatus =
        new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).
        unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED,
            appMsg, appUrl);
    verify(mockRMClient).stop();
  }

  @Test(timeout=30000)
  public void testSimpleReuseRackMatching() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback,
        mockRMClient, clock);

    scheduler.initialize();
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
        regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
        regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv1 = Priority.newInstance(2);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, "host2");
    MockTaskInfo taskv0t2 = new MockTaskInfo("taskv0t2", priorityv0, "host4", "/somerack");
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, "host1");
    MockTaskInfo taskv1t1 = new MockTaskInfo("taskv1t1", priorityv1, "host5");

    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient,
        scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    TaskRequest reqv0t1 = taskRequestCaptor.scheduleTask(taskv0t1);
    taskRequestCaptor.scheduleTask(taskv0t2);
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0);
    taskRequestCaptor.scheduleTask(taskv1t1);

    NodeId host1 = NodeId.newInstance("host1", 1);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    Container container1 = Container.newInstance(cid1, host1, null, taskv0t0.capability, priorityv0, null);

    // allocate one container at v0 priority
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);

    // finish v0t0 successfully, verify v0t1 is skipped and v1t0 assigned instead
    // since host locality is preferred to rack locality and lower priority vertex
    // is not blocked by higher priority vertex
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv1t0.task, taskv1t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv1t0);

    // finish v1t0 successfully, verify v0t1 is assigned
    assertTrue(scheduler.deallocateTask(taskv1t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv0t1.task, taskv0t1.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t1);

    // fail v0t1 and verify container is released instead of reused for v1t1
    assertTrue(scheduler.deallocateTask(taskv0t1.task, false, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(cid1);
    verify(mockRMClient).releaseAssignedContainer(cid1);

    String appMsg = "success";
    AppFinalStatus finalStatus =
        new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).
        unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED,
            appMsg, appUrl);
    verify(mockRMClient).stop();
  }

  @Test(timeout=30000)
  public void testSimpleReuseAnyMatching() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback,
        mockRMClient, clock);

    scheduler.initialize();
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
        regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
        regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv1 = Priority.newInstance(2);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, "host2");
    MockTaskInfo taskv0t2 = new MockTaskInfo("taskv0t2", priorityv0, "host4", "/rack4");
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, "host1");
    MockTaskInfo taskv1t1 = new MockTaskInfo("taskv1t1", priorityv1, "host6", "/rack6");

    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient,
        scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    TaskRequest reqv0t1 = taskRequestCaptor.scheduleTask(taskv0t1);
    TaskRequest reqv0t2 = taskRequestCaptor.scheduleTask(taskv0t2);
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0);
    taskRequestCaptor.scheduleTask(taskv1t1);

    NodeId host1 = NodeId.newInstance("host1", 1);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    Container container1 = Container.newInstance(cid1, host1, null, taskv0t0.capability, priorityv0, null);

    // allocate one container at v0 priority
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);

    // finish v0t0 successfully, verify v0t1 is skipped and v1t0 assigned instead
    // since host locality is preferred to rack locality and lower priority vertex
    // is not blocked by higher priority vertex
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv1t0.task, taskv1t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv1t0);

    // finish v1t0 successfully, verify v0t1 is assigned
    assertTrue(scheduler.deallocateTask(taskv1t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv0t1.task, taskv0t1.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t1);

    // finish v0t1 successfully, verify v0t2 is assigned
    assertTrue(scheduler.deallocateTask(taskv0t1.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv0t2.task, taskv0t2.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t2);

    // fail v0t2 and verify container is released instead of reused for v1t1
    assertTrue(scheduler.deallocateTask(taskv0t2.task, false, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(cid1);
    verify(mockRMClient).releaseAssignedContainer(cid1);

    String appMsg = "success";
    AppFinalStatus finalStatus =
        new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).
        unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED,
            appMsg, appUrl);
    verify(mockRMClient).stop();
  }

  @Test(timeout=30000)
  public void testReuseWithAffinity() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback,
        mockRMClient, clock);

    scheduler.initialize();
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
        regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
        regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv1 = Priority.newInstance(2);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, hostsv0t0);

    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient,
        scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    taskRequestCaptor.scheduleTask(taskv0t1);

    NodeId host1 = NodeId.newInstance("host1", 1);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    Container container1 = Container.newInstance(cid1, host1, null, taskv0t0.capability, priorityv0, null);

    // allocate one container at v0 priority
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);

    // add a new request for this container
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, "host1");
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0, cid1);

    // finish v0t0 successfully, verify v0t1 is skipped even though it is node-local
    // and v1t0 assigned instead for affinity
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv1t0.task, taskv1t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv1t0);

    String appMsg = "success";
    AppFinalStatus finalStatus =
        new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).
        unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED,
            appMsg, appUrl);
    verify(mockRMClient).stop();
  }

  @Test(timeout=30000)
  public void testReuseVertexDescendants() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    // vertex 0 and vertex 2 are root vertices and vertex 1 is a child of vertex 0
    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(3);
    when(mockDagInfo.getVertexDescendants(0)).thenReturn(BitSet.valueOf(new long[] { 0x2 }));
    when(mockDagInfo.getVertexDescendants(1)).thenReturn(new BitSet());
    when(mockDagInfo.getVertexDescendants(2)).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback,
        mockRMClient, clock);

    scheduler.initialize();
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
        regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
        regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv1 = Priority.newInstance(2);
    Priority priorityv2 = Priority.newInstance(3);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv0t0.task)).thenReturn(0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, "host3");
    when(mockApp.getVertexIndexForTask(taskv0t1.task)).thenReturn(0);
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv1t0.task)).thenReturn(1);
    MockTaskInfo taskv2t0 = new MockTaskInfo("taskv2t0", priorityv2, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv2t0.task)).thenReturn(2);
    MockTaskInfo taskv2t1 = new MockTaskInfo("taskv2t1", priorityv2, "host3");
    when(mockApp.getVertexIndexForTask(taskv2t1.task)).thenReturn(2);

    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient,
        scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    TaskRequest reqv0t1 = taskRequestCaptor.scheduleTask(taskv0t1);
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0);
    TaskRequest reqv2t0 = taskRequestCaptor.scheduleTask(taskv2t0);
    taskRequestCaptor.scheduleTask(taskv2t1);

    NodeId host1 = NodeId.newInstance("host1", 1);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    Container container1 = Container.newInstance(cid1, host1, null, taskv0t0.capability, priorityv0, null);

    // allocate one container at v0 priority
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);

    // finish v0t0 successfully, verify v1t0 is skipped and v2t0 assigned instead
    // since host locality beats rack locality for unblocked vertex v2 and
    // v1 is blocked by pending v0 request
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv2t0.task, taskv2t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv2t0);

    // finish v2t0 successfully, verify v0t1 is assigned since it is higher
    // priority than v2
    assertTrue(scheduler.deallocateTask(taskv2t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv0t1.task, taskv0t1.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t1);

    // finish v2t0 successfully, verify v1t0 is assigned since it is now unblocked
    assertTrue(scheduler.deallocateTask(taskv0t1.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv1t0.task, taskv1t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv1t0);

    // fail v1t0 and verify container is released instead of reused for v2t1
    assertTrue(scheduler.deallocateTask(taskv1t0.task, false, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(cid1);
    verify(mockRMClient).releaseAssignedContainer(cid1);

    String appMsg = "success";
    AppFinalStatus finalStatus =
        new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).
        unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED,
            appMsg, appUrl);
    verify(mockRMClient).stop();
  }

  @Test(timeout=30000)
  public void testSessionContainers() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 4000);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 5000);
    conf.setInt(TezConfiguration.TEZ_AM_SESSION_MIN_HELD_CONTAINERS, 5);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    when(mockApp.isSession()).thenReturn(true);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback,
        mockRMClient, clock);

    scheduler.initialize();
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
        regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
        regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    final String rack1 = "/r1";
    final String rack2 = "/r2";
    final String rack3 = "/r3";
    final String node1Rack1 = "n1r1";
    final String node2Rack1 = "n2r1";
    final String node1Rack2 = "n1r2";
    final String node2Rack2 = "n2r2";
    final String node1Rack3 = "n1r3";
    MockDNSToSwitchMapping.addRackMapping(node1Rack1, rack1);
    MockDNSToSwitchMapping.addRackMapping(node2Rack1, rack1);
    MockDNSToSwitchMapping.addRackMapping(node1Rack2, rack2);
    MockDNSToSwitchMapping.addRackMapping(node2Rack2, rack2);
    MockDNSToSwitchMapping.addRackMapping(node1Rack3, rack3);

    Priority priorityv0 = Priority.newInstance(1);
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, node1Rack1, rack1);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, node2Rack1, rack1);
    MockTaskInfo taskv0t2 = new MockTaskInfo("taskv0t2", priorityv0, node1Rack1, rack1);
    MockTaskInfo taskv0t3 = new MockTaskInfo("taskv0t3", priorityv0, node2Rack1, rack1);
    MockTaskInfo taskv0t4 = new MockTaskInfo("taskv0t4", priorityv0, node1Rack2, rack2);
    MockTaskInfo taskv0t5 = new MockTaskInfo("taskv0t5", priorityv0, node2Rack2, rack2);
    MockTaskInfo taskv0t6 = new MockTaskInfo("taskv0t6", priorityv0, node1Rack3, rack3);

    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient,
        scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    TaskRequest reqv0t1 = taskRequestCaptor.scheduleTask(taskv0t1);
    TaskRequest reqv0t2 = taskRequestCaptor.scheduleTask(taskv0t2);
    TaskRequest reqv0t3 = taskRequestCaptor.scheduleTask(taskv0t3);
    TaskRequest reqv0t4 = taskRequestCaptor.scheduleTask(taskv0t4);
    TaskRequest reqv0t5 = taskRequestCaptor.scheduleTask(taskv0t5);
    TaskRequest reqv0t6 = taskRequestCaptor.scheduleTask(taskv0t6);

    List<Container> containers = new ArrayList<>();
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    NodeId n1r1 = NodeId.newInstance(node1Rack1, 1);
    Container container1 = Container.newInstance(cid1, n1r1, null, taskv0t0.capability, priorityv0, null);
    containers.add(container1);
    ContainerId cid2 = ContainerId.newContainerId(attemptId, 2);
    NodeId n2r1 = NodeId.newInstance(node2Rack1, 1);
    Container container2 = Container.newInstance(cid2, n2r1, null, taskv0t1.capability, priorityv0, null);
    containers.add(container2);
    ContainerId cid3 = ContainerId.newContainerId(attemptId, 3);
    Container container3 = Container.newInstance(cid3, n1r1, null, taskv0t2.capability, priorityv0, null);
    containers.add(container3);
    ContainerId cid4 = ContainerId.newContainerId(attemptId, 4);
    Container container4 = Container.newInstance(cid4, n2r1, null, taskv0t3.capability, priorityv0, null);
    containers.add(container4);
    ContainerId cid5 = ContainerId.newContainerId(attemptId, 5);
    NodeId n1r2 = NodeId.newInstance(node1Rack2, 1);
    Container container5 = Container.newInstance(cid5, n1r2, null, taskv0t4.capability, priorityv0, null);
    containers.add(container5);
    ContainerId cid6 = ContainerId.newContainerId(attemptId, 6);
    NodeId n2r2 = NodeId.newInstance(node2Rack2, 1);
    Container container6 = Container.newInstance(cid6, n2r2, null, taskv0t5.capability, priorityv0, null);
    containers.add(container6);
    ContainerId cid7 = ContainerId.newContainerId(attemptId, 7);
    NodeId n1r3 = NodeId.newInstance(node1Rack3, 1);
    Container container7 = Container.newInstance(cid7, n1r3, null, taskv0t6.capability, priorityv0, null);
    containers.add(container7);

    scheduler.onContainersAllocated(containers);
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);
    verify(mockApp).taskAllocated(taskv0t1.task, taskv0t1.cookie, container2);
    verify(mockRMClient).removeContainerRequest(reqv0t1);
    verify(mockApp).taskAllocated(taskv0t2.task, taskv0t2.cookie, container3);
    verify(mockRMClient).removeContainerRequest(reqv0t2);
    verify(mockApp).taskAllocated(taskv0t3.task, taskv0t3.cookie, container4);
    verify(mockRMClient).removeContainerRequest(reqv0t3);
    verify(mockApp).taskAllocated(taskv0t4.task, taskv0t4.cookie, container5);
    verify(mockRMClient).removeContainerRequest(reqv0t4);
    verify(mockApp).taskAllocated(taskv0t5.task, taskv0t5.cookie, container6);
    verify(mockRMClient).removeContainerRequest(reqv0t5);
    verify(mockApp).taskAllocated(taskv0t6.task, taskv0t6.cookie, container7);
    verify(mockRMClient).removeContainerRequest(reqv0t6);

    clock.incrementTime(10000);
    drainableAppCallback.drain();
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t1.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t2.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t3.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t4.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t5.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t6.task, true, null, null));
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));

    // verify only two of the containers were released after idle expiration
    // and the rest were spread across the nodes and racks
    clock.incrementTime(5000);
    drainableAppCallback.drain();
    verify(mockApp, times(2)).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, times(2)).releaseAssignedContainer(any(ContainerId.class));
    Set<String> hosts = new HashSet<>();
    Set<String> racks = new HashSet<>();
    for (HeldContainer hc : scheduler.getSessionContainers()) {
      hosts.add(hc.getHost());
      racks.add(hc.getRack());
    }
    assertEquals(5, hosts.size());
    assertEquals(3, racks.size());
    assertTrue(hosts.contains(node1Rack1));
    assertTrue(hosts.contains(node2Rack1));
    assertTrue(hosts.contains(node1Rack2));
    assertTrue(hosts.contains(node2Rack2));
    assertTrue(hosts.contains(node1Rack3));
    assertTrue(racks.contains(rack1));
    assertTrue(racks.contains(rack2));
    assertTrue(racks.contains(rack3));

    String appMsg = "success";
    AppFinalStatus finalStatus =
        new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).
        unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED,
            appMsg, appUrl);
    verify(mockRMClient).stop();
  }

  @Test(timeout=50000)
  public void testPreemptionNoHeadroom() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);
    conf.setInt(TezConfiguration.TEZ_AM_PREEMPTION_PERCENTAGE, 10);
    conf.setInt(TezConfiguration.TEZ_AM_PREEMPTION_HEARTBEATS_BETWEEN_PREEMPTIONS, 3);
    conf.setInt(TezConfiguration.TEZ_AM_PREEMPTION_MAX_WAIT_TIME_MS, 60 * 1000);

    // vertex 0 and vertex 2 are root vertices and vertex 1 is a child of vertex 0
    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(3);
    when(mockDagInfo.getVertexDescendants(0)).thenReturn(BitSet.valueOf(new long[] { 0x2 }));
    when(mockDagInfo.getVertexDescendants(1)).thenReturn(new BitSet());
    when(mockDagInfo.getVertexDescendants(2)).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback,
        mockRMClient, clock);

    scheduler.initialize();
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
        regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
        regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv1 = Priority.newInstance(2);
    Priority priorityv2 = Priority.newInstance(3);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv0t0.task)).thenReturn(0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv0t1.task)).thenReturn(0);
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv1t0.task)).thenReturn(1);
    MockTaskInfo taskv1t1 = new MockTaskInfo("taskv1t1", priorityv1, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv1t1.task)).thenReturn(1);
    MockTaskInfo taskv2t0 = new MockTaskInfo("taskv2t0", priorityv2, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv2t0.task)).thenReturn(2);

    // asks for two tasks for vertex 1 and start running one of them
    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient,
        scheduler, drainableAppCallback);
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0);
    TaskRequest reqv1t1 = taskRequestCaptor.scheduleTask(taskv1t1);
    NodeId host1 = NodeId.newInstance("host1", 1);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    Container container1 = Container.newInstance(cid1, host1, null, taskv1t0.capability, priorityv1, null);
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv1t0.task, taskv1t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv1t0);

    // start running the other task for vertex 1 a bit later
    clock.incrementTime(1000);
    ContainerId cid2 = ContainerId.newContainerId(attemptId, 2);
    Container container2 = Container.newInstance(cid2, host1, null, taskv1t0.capability, priorityv1, null);
    scheduler.onContainersAllocated(Collections.singletonList(container2));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv1t1.task, taskv1t1.cookie, container2);
    verify(mockRMClient).removeContainerRequest(reqv1t1);

    // add a request for vertex 0 but there is no headroom
    when(mockRMClient.getAvailableResources()).thenReturn(Resources.none());
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);

    // should preempt after enough heartbeats to get past preemption interval
    // only the youngest container should be preempted to meet the demand
    scheduler.getProgress();
    scheduler.getProgress();
    scheduler.getProgress();
    drainableAppCallback.drain();
    verify(mockApp, times(1)).preemptContainer(any(ContainerId.class));
    verify(mockApp).preemptContainer(cid2);
    assertEquals(taskv1t1.task, scheduler.deallocateContainer(cid2));
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(cid2);
    verify(mockRMClient).releaseAssignedContainer(cid2);
    verify(mockApp, never()).containerBeingReleased(cid1);
    verify(mockRMClient, never()).releaseAssignedContainer(cid1);

    // add a request for vertex 2 and allocate another container
    clock.incrementTime(1000);
    taskRequestCaptor.scheduleTask(taskv2t0);
    ContainerId cid3 = ContainerId.newContainerId(attemptId, 3);
    Container container3 = Container.newInstance(cid3, host1, null, taskv0t0.capability, priorityv0, null);
    scheduler.onContainersAllocated(Collections.singletonList(container3));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container3);
    verify(mockRMClient).removeContainerRequest(reqv0t0);

    // no more preemptions since v1 is not a descendant of v2
    scheduler.getProgress();
    scheduler.getProgress();
    scheduler.getProgress();
    drainableAppCallback.drain();
    verify(mockApp, times(1)).preemptContainer(any(ContainerId.class));

    // adding request for v0 should trigger preemption on next heartbeat
    taskRequestCaptor.scheduleTask(taskv0t1);
    scheduler.getProgress();
    drainableAppCallback.drain();
    verify(mockApp, times(2)).preemptContainer(any(ContainerId.class));
    verify(mockApp).preemptContainer(cid1);
    assertEquals(taskv1t0.task, scheduler.deallocateContainer(cid1));
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(cid1);
    verify(mockRMClient).releaseAssignedContainer(cid1);

    String appMsg = "success";
    AppFinalStatus finalStatus =
        new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).
        unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED,
            appMsg, appUrl);
    verify(mockRMClient).stop();
  }

  @Test (timeout = 50000L)
  public void testPreemptionWhenBlocked() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);
    conf.setInt(TezConfiguration.TEZ_AM_PREEMPTION_PERCENTAGE, 10);
    conf.setInt(TezConfiguration.TEZ_AM_PREEMPTION_HEARTBEATS_BETWEEN_PREEMPTIONS, 3);
    conf.setInt(TezConfiguration.TEZ_AM_PREEMPTION_MAX_WAIT_TIME_MS, 60 * 1000);


    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(3);
    when(mockDagInfo.getVertexDescendants(0)).thenReturn(BitSet.valueOf(new long[] { 0x6 }));
    when(mockDagInfo.getVertexDescendants(1)).thenReturn(BitSet.valueOf(new long[] { 0x2 }));
    when(mockDagInfo.getVertexDescendants(2)).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback,
        mockRMClient, clock);

    scheduler.initialize();
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
        regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
        regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv2 = Priority.newInstance(3);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv0t0.task)).thenReturn(0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv0t1.task)).thenReturn(0);
    MockTaskInfo taskv2t0 = new MockTaskInfo("taskv2t0", priorityv2, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv2t0.task)).thenReturn(2);
    MockTaskInfo taskv2t1 = new MockTaskInfo("taskv2t1", priorityv2, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv2t1.task)).thenReturn(2);
    when(mockApp.getVertexIndexForTask(taskv2t0.task)).thenReturn(2);

    // asks for one task for vertex 2 and start running
    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient,
        scheduler, drainableAppCallback);
    TaskRequest reqv2t0 = taskRequestCaptor.scheduleTask(taskv2t0);
    NodeId host1 = NodeId.newInstance("host1", 1);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    Container container1 = Container.newInstance(cid1, host1, null, taskv2t0.capability, priorityv2, null);
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv2t0.task, taskv2t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv2t0);
    clock.incrementTime(1000);

    when(mockRMClient.getAvailableResources()).thenReturn(Resources.none());
    scheduler.getProgress();
    scheduler.getProgress();
    scheduler.getProgress();
    drainableAppCallback.drain();
    //ask another task for v2
    TaskRequest reqv2t1 = taskRequestCaptor.scheduleTask(taskv2t1);
    scheduler.getProgress();
    scheduler.getProgress();
    scheduler.getProgress();
    drainableAppCallback.drain();

    clock.incrementTime(1000);
    // add a request for vertex 0 but there is no headroom, this should preempt
    when(mockRMClient.getAvailableResources()).thenReturn(Resources.none());
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);

    // should preempt after enough heartbeats to get past preemption interval
    scheduler.getProgress();
    scheduler.getProgress();
    scheduler.getProgress();
    drainableAppCallback.drain();
    verify(mockApp, times(1)).preemptContainer(any(ContainerId.class));
    verify(mockApp).preemptContainer(cid1);
    String appMsg = "success";
    AppFinalStatus finalStatus =
        new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).
        unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED,
            appMsg, appUrl);
    verify(mockRMClient).stop();
  }

  @Test(timeout=50000)
  public void testContainerAssignmentReleaseNewContainers() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, false);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NEW_CONTAINERS_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 4000);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 5000);
    conf.setInt(TezConfiguration.TEZ_AM_SESSION_MIN_HELD_CONTAINERS, 5);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    when(mockApp.isSession()).thenReturn(true);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback,
        mockRMClient, clock);

    scheduler.initialize();
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
        regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
        regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    final String rack1 = "/r1";
    final String rack2 = "/r2";
    final String node1Rack1 = "n1r1";
    final String node2Rack1 = "n2r1";
    final String node1Rack2 = "n1r2";
    MockDNSToSwitchMapping.addRackMapping(node1Rack1, rack1);
    MockDNSToSwitchMapping.addRackMapping(node2Rack1, rack1);
    MockDNSToSwitchMapping.addRackMapping(node1Rack2, rack2);

    Priority priorityv0 = Priority.newInstance(1);
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, node1Rack1, rack1);

    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    NodeId n2r1 = NodeId.newInstance(node2Rack1, 1);
    Container container1 = Container.newInstance(cid1, n2r1, null, taskv0t0.capability, priorityv0, null);

    // verify new container is released is not immediately allocated
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    // app is not notified of the container being released since it never launched
    verify(mockApp, never()).containerBeingReleased(cid1);
    verify(mockRMClient).releaseAssignedContainer(eq(cid1));
  }

  @Test(timeout=50000)
  public void testIdleContainerAssignmentReuseNewContainers() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, false);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NEW_CONTAINERS_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 4000);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 5000);
    conf.setInt(TezConfiguration.TEZ_AM_SESSION_MIN_HELD_CONTAINERS, 5);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    when(mockApp.isSession()).thenReturn(true);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback,
        mockRMClient, clock);

    scheduler.initialize();
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
        regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
        regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    final String rack1 = "/r1";
    final String rack2 = "/r2";
    final String node1Rack1 = "n1r1";
    final String node2Rack1 = "n2r1";
    final String node1Rack2 = "n1r2";
    MockDNSToSwitchMapping.addRackMapping(node1Rack1, rack1);
    MockDNSToSwitchMapping.addRackMapping(node2Rack1, rack1);
    MockDNSToSwitchMapping.addRackMapping(node1Rack2, rack2);

    Priority priorityv0 = Priority.newInstance(1);
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, node1Rack1, rack1);

    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    NodeId n2r1 = NodeId.newInstance(node2Rack1, 1);
    Container container1 = Container.newInstance(cid1, n2r1, null, taskv0t0.capability, priorityv0, null);

    // verify idle container is kept for now
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    clock.incrementTime(2000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(cid1);
    verify(mockRMClient, never()).releaseAssignedContainer(cid1);

    // verify idle container is released without being assigned to a task because rack-local reuse is
    // disabled
    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient,
        scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient, never()).removeContainerRequest(reqv0t0);
    verify(mockApp, never()).containerBeingReleased(cid1);
    verify(mockRMClient).releaseAssignedContainer(cid1);

    // cancel the task request
    assertFalse(scheduler.deallocateTask(taskv0t0.task, false, null, null));

    // allocate another container that's node-local
    ContainerId cid2 = ContainerId.newContainerId(attemptId, 2);
    NodeId n1r1 = NodeId.newInstance(node1Rack1, 1);
    Container container2 = Container.newInstance(cid2, n1r1, null, taskv0t0.capability, priorityv0, null);
    scheduler.onContainersAllocated(Collections.singletonList(container2));
    clock.incrementTime(2000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(cid2);
    verify(mockRMClient, never()).releaseAssignedContainer(cid2);

    // reschedule the task, verify it's now scheduled without a container request
    // since node-local idle container is available
    reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0, false);
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container2);
    verify(mockRMClient).removeContainerRequest(reqv0t0);
  }

  @Test
  public void testMinMaxContainerIdleMillisAreEqual() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = new AMRMClientAsyncWrapperForTest();
    Configuration conf = new Configuration();
    conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 10000);
    conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 10000);

    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext("host", 0, "url", conf);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);
    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback, mockRMClient, clock);
    scheduler.initialize();

    NodeId host1 = NodeId.newInstance("host1", 1);
    Container container1 = Container.newInstance(null, host1, null, null, null, null);
    HeldContainer heldContainer = scheduler.new HeldContainer(container1);
    long now = clock.getTime();
    assertEquals(now + 10000, heldContainer.getIdleExpirationTimestamp(now));
  }

  static class AMRMClientAsyncWrapperForTest extends AMRMClientAsyncWrapper {
    AMRMClientAsyncWrapperForTest() {
      super(new MockAMRMClient(), 10000, null);
    }

    RegisterApplicationMasterResponse getRegistrationResponse() {
      return ((MockAMRMClient) client).getRegistrationResponse();
    }

    @Override
    public RegisterApplicationMasterResponse registerApplicationMaster(String appHostName, int appHostPort,
        String appTrackingUrl) throws YarnException, IOException {
      return client.registerApplicationMaster(appHostName, appHostPort, appTrackingUrl);
    }

    @Override
    protected void serviceStart() {
    }

    @Override
    protected void serviceStop() {
    }
  }

  static class MockAMRMClient extends AMRMClientImpl<TaskRequest> {
    private RegisterApplicationMasterResponse mockRegResponse;

    MockAMRMClient() {
      super();
      this.clusterAvailableResources = Resource.newInstance(4000, 4);
      this.clusterNodeCount = 5;
    }

    @Override
    protected void serviceStart() {
    }

    @Override
    protected void serviceStop() {
    }

    @Override
    public RegisterApplicationMasterResponse registerApplicationMaster(String appHostName, int appHostPort,
        String appTrackingUrl) {
      mockRegResponse = mock(RegisterApplicationMasterResponse.class);
      Resource mockMaxResource = Resources.createResource(1024*1024, 1024);
      Map<ApplicationAccessType, String> mockAcls = Collections.emptyMap();
      when(mockRegResponse.getMaximumResourceCapability()).thenReturn(
          mockMaxResource);
      when(mockRegResponse.getApplicationACLs()).thenReturn(mockAcls);
      when(mockRegResponse.getSchedulerResourceTypes()).thenReturn(
          EnumSet.of(SchedulerResourceTypes.MEMORY, SchedulerResourceTypes.CPU));
      return mockRegResponse;
    }

    @Override
    public void unregisterApplicationMaster(FinalApplicationStatus appStatus,
        String appMessage, String appTrackingUrl) {
    }

    RegisterApplicationMasterResponse getRegistrationResponse() {
      return mockRegResponse;
    }
  }

  static class MockTask {
    final String name;

    MockTask(String name) {
      this.name = name;
    }

    @Override
    public String toString() {
      return name;
    }
  }

  static class MockTaskInfo {
    final static Object DEFAULT_SIGNATURE = new Object();

    final MockTask task;
    final Object cookie = new Object();
    final Object signature = DEFAULT_SIGNATURE;
    final String[] hosts;
    final String[] racks;
    final Priority priority;
    final Resource capability;

    MockTaskInfo(String name, Priority priority, String host) {
      this(name, priority, host == null ? null : new String[] { host });
    }

    MockTaskInfo(String name, Priority priority, String[] hosts) {
      this(name, priority, hosts, buildDefaultRacks(hosts));
    }

    MockTaskInfo(String name, Priority priority, String host, String rack) {
      this(name, priority, host == null ? null : new String[] { host },
          rack == null ? null : new String[] { rack });
    }

    MockTaskInfo(String name, Priority priority, String[] hosts, String[] racks) {
      this.task = new MockTask(name);
      this.hosts = hosts;
      this.racks = racks;
      this.priority = priority;
      this.capability = Resource.newInstance(1024, 1);
    }

    static String[] buildDefaultRacks(String[] hosts) {
      if (hosts == null) {
        return null;
      }
      String[] racks = new String[hosts.length];
      Arrays.fill(racks, "/default-rack");
      return racks;
    }
  }

  static class TaskRequestCaptor {
    final AMRMClientAsync<TaskRequest> client;
    final TaskScheduler scheduler;
    final TaskSchedulerContextDrainable drainableAppCallback;
    final ArgumentCaptor<TaskRequest> captor = ArgumentCaptor.forClass(TaskRequest.class);
    int invocationCount = 0;

    TaskRequestCaptor(AMRMClientAsync<TaskRequest> client, TaskScheduler scheduler,
        TaskSchedulerContextDrainable drainableAppCallback) {
      this.client = client;
      this.scheduler = scheduler;
      this.drainableAppCallback = drainableAppCallback;
    }

    TaskRequest scheduleTask(MockTaskInfo taskInfo) throws Exception {
      return scheduleTask(taskInfo, true);
    }

    TaskRequest scheduleTask(MockTaskInfo taskInfo, boolean expectContainerRequest) throws Exception {
      scheduler.allocateTask(taskInfo.task, taskInfo.capability, taskInfo.hosts, taskInfo.racks,
          taskInfo.priority, taskInfo.signature, taskInfo.cookie);
      drainableAppCallback.drain();
      if (expectContainerRequest) {
        ++invocationCount;
      }
      verify(client, times(invocationCount)).addContainerRequest(captor.capture());
      TaskRequest request = captor.getValue();
      assertEquals(request.getTask(), taskInfo.task);
      assertEquals(request.getCookie(), taskInfo.cookie);
      return request;
    }

    TaskRequest scheduleTask(MockTaskInfo taskInfo, ContainerId affinity) throws Exception {
      scheduler.allocateTask(taskInfo.task, taskInfo.capability, affinity, taskInfo.priority,
          taskInfo.signature, taskInfo.cookie);
      drainableAppCallback.drain();
      verify(client, times(++invocationCount)).addContainerRequest(captor.capture());
      TaskRequest request = captor.getValue();
      assertEquals(request.getTask(), taskInfo.task);
      assertEquals(request.getCookie(), taskInfo.cookie);
      return request;
    }
  }

  static class NewTaskSchedulerForTest extends DagAwareYarnTaskScheduler {
    final AMRMClientAsyncWrapper mockClient;
    final MockClock clock;

    NewTaskSchedulerForTest(
        TaskSchedulerContextDrainable appClient,
        AMRMClientAsyncWrapper client, MockClock clock) {
      super(appClient);
      this.mockClient = client;
      this.clock = clock;
      setShouldUnregister();
    }

    @Override
    public void initialize() throws Exception {
      initialize(mockClient);
    }

    @Override
    protected ScheduledExecutorService createExecutor() {
      return new ControlledScheduledExecutorService(clock);
    }

    @Override
    protected long now() {
      return clock.getTime();
    }
  }
}