/*
 * Licensed 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.hadoop.hive.llap.tezplugins;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.reset;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.DelayQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hive.conf.HiveConf.ConfVars;
import org.apache.hadoop.hive.llap.registry.ServiceInstance;
import org.apache.hadoop.hive.llap.registry.ServiceInstanceSet;
import org.apache.hadoop.hive.llap.registry.impl.InactiveServiceInstance;
import org.apache.hadoop.hive.llap.registry.impl.LlapFixedRegistryImpl;
import org.apache.hadoop.hive.llap.testhelpers.ControlledClock;
import org.apache.hadoop.hive.llap.tezplugins.helpers.MonotonicClock;
import org.apache.hadoop.util.Time;
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.Priority;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.util.Clock;
import org.apache.tez.common.TezUtils;
import org.apache.tez.dag.api.UserPayload;
import org.apache.tez.serviceplugins.api.TaskAttemptEndReason;
import org.apache.tez.serviceplugins.api.TaskSchedulerContext;
import org.junit.Test;
import org.mockito.ArgumentCaptor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.ImmutableSet;

public class TestLlapTaskSchedulerService {

  private static final Logger LOG = LoggerFactory.getLogger(TestLlapTaskSchedulerService.class);

  private static final String HOST1 = "host1";
  private static final String HOST2 = "host2";
  private static final String HOST3 = "host3";
  private static final String HOST4 = "host4";

  @Test(timeout = 10000)
  public void testSimpleLocalAllocation() throws IOException, InterruptedException {

    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper();

    try {
      Priority priority1 = Priority.newInstance(1);
      String[] hosts1 = new String[]{HOST1};

      Object task1 = new Object();
      Object clientCookie1 = new Object();

      tsWrapper.controlScheduler(true);
      tsWrapper.allocateTask(task1, hosts1, priority1, clientCookie1);

      tsWrapper.awaitLocalTaskAllocations(1);

      verify(tsWrapper.mockAppCallback).taskAllocated(eq(task1), eq(clientCookie1), any(Container.class));
      assertEquals(1, tsWrapper.ts.dagStats.numLocalAllocations);
      assertEquals(1, tsWrapper.ts.dagStats.numAllocationsPerHost.get(HOST1).get());
    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testSimpleNoLocalityAllocation() throws IOException, InterruptedException {
    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper();

    try {
      Priority priority1 = Priority.newInstance(1);

      Object task1 = new Object();
      Object clientCookie1 = new Object();
      tsWrapper.controlScheduler(true);
      tsWrapper.allocateTask(task1, null, priority1, clientCookie1);
      tsWrapper.awaitTotalTaskAllocations(1);
      verify(tsWrapper.mockAppCallback).taskAllocated(eq(task1), eq(clientCookie1), any(Container.class));
      assertEquals(1, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);
    } finally {
      tsWrapper.shutdown();
    }
  }


  @Test(timeout = 10000)
  public void testPreemption() throws InterruptedException, IOException {

    Priority priority1 = Priority.newInstance(1);
    Priority priority2 = Priority.newInstance(2);
    String [] hosts = new String[] {HOST1};
    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1);
    try {

      Object task1 = "task1";
      Object clientCookie1 = "cookie1";
      Object task2 = "task2";
      Object clientCookie2 = "cookie2";
      Object task3 = "task3";
      Object clientCookie3 = "cookie3";
      Object task4 = "task4";
      Object clientCookie4 = "cookie4";

      tsWrapper.controlScheduler(true);
      tsWrapper.allocateTask(task1, hosts, priority2, clientCookie1);
      tsWrapper.allocateTask(task2, hosts, priority2, clientCookie2);
      tsWrapper.allocateTask(task3, hosts, priority2, clientCookie3);
      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numLocalAllocations == 2) {
          break;
        }
      }
      verify(tsWrapper.mockAppCallback, times(2)).taskAllocated(any(Object.class),
          any(Object.class), any(Container.class));
      assertEquals(2, tsWrapper.ts.dagStats.numLocalAllocations);
      assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);

      reset(tsWrapper.mockAppCallback);

      tsWrapper.allocateTask(task4, hosts, priority1, clientCookie4);
      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numPreemptedTasks == 1) {
          break;
        }
      }
      verify(tsWrapper.mockAppCallback).preemptContainer(any(ContainerId.class));


      tsWrapper.deallocateTask(task2, false, TaskAttemptEndReason.INTERNAL_PREEMPTION);

      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numTotalAllocations == 3) {
          break;
        }
      }
      verify(tsWrapper.mockAppCallback, times(1)).taskAllocated(eq(task4),
          eq(clientCookie4), any(Container.class));

    } finally {
      tsWrapper.shutdown();
    }

  }

  @Test(timeout = 10000)
  public void testNodeDisabled() throws IOException, InterruptedException {
    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(10000l);
    try {
      Priority priority1 = Priority.newInstance(1);
      String[] hosts1 = new String[]{HOST1};
      Object task1 = new Object();
      Object clientCookie1 = new Object();
      tsWrapper.controlScheduler(true);
      tsWrapper.allocateTask(task1, hosts1, priority1, clientCookie1);
      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numTotalAllocations == 1) {
          break;
        }
      }
      verify(tsWrapper.mockAppCallback).taskAllocated(eq(task1), eq(clientCookie1),
          any(Container.class));
      assertEquals(1, tsWrapper.ts.dagStats.numLocalAllocations);
      assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);
      assertEquals(0, tsWrapper.ts.dagStats.numNonLocalAllocations);
      assertEquals(1, tsWrapper.ts.dagStats.numTotalAllocations);

      tsWrapper.resetAppCallback();

      tsWrapper.clock.setTime(10000l);
      tsWrapper.rejectExecution(task1);

      // Verify that the node is blacklisted
      assertEquals(1, tsWrapper.ts.dagStats.numRejectedTasks);
      assertEquals(3, tsWrapper.ts.instanceToNodeMap.size());
      LlapTaskSchedulerService.NodeInfo disabledNodeInfo = tsWrapper.ts.disabledNodesQueue.peek();
      assertNotNull(disabledNodeInfo);
      assertEquals(HOST1, disabledNodeInfo.getHost());
      assertEquals((10000l), disabledNodeInfo.getDelay(TimeUnit.MILLISECONDS));
      assertEquals((10000l + 10000l), disabledNodeInfo.expireTimeMillis);

      Object task2 = new Object();
      Object clientCookie2 = new Object();
      tsWrapper.allocateTask(task2, hosts1, priority1, clientCookie2);
      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numTotalAllocations == 2) {
          break;
        }
      }
      verify(tsWrapper.mockAppCallback).taskAllocated(eq(task2), eq(clientCookie2), any(Container.class));
      assertEquals(1, tsWrapper.ts.dagStats.numLocalAllocations);
      assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);
      assertEquals(1, tsWrapper.ts.dagStats.numNonLocalAllocations);
      assertEquals(2, tsWrapper.ts.dagStats.numTotalAllocations);

    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testNodeReEnabled() throws InterruptedException, IOException {
    // Based on actual timing.
    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(1000l);
    try {
      Priority priority1 = Priority.newInstance(1);
      String[] hosts1 = new String[]{HOST1};
      String[] hosts2 = new String[]{HOST2};
      String[] hosts3 = new String[]{HOST3};

      Object task1 = new Object();
      Object clientCookie1 = new Object();
      Object task2 = new Object();
      Object clientCookie2 = new Object();
      Object task3 = new Object();
      Object clientCookie3 = new Object();

      tsWrapper.controlScheduler(true);
      tsWrapper.allocateTask(task1, hosts1, priority1, clientCookie1);
      tsWrapper.allocateTask(task2, hosts2, priority1, clientCookie2);
      tsWrapper.allocateTask(task3, hosts3, priority1, clientCookie3);
      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numTotalAllocations == 3) {
          break;
        }
      }
      verify(tsWrapper.mockAppCallback, times(3)).taskAllocated(any(Object.class), any(Object.class), any(Container.class));
      assertEquals(3, tsWrapper.ts.dagStats.numLocalAllocations);
      assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);
      assertEquals(3, tsWrapper.ts.dagStats.numTotalAllocations);

      tsWrapper.resetAppCallback();

      tsWrapper.rejectExecution(task1);
      tsWrapper.rejectExecution(task2);
      tsWrapper.rejectExecution(task3);

      // Verify that the node is blacklisted
      assertEquals(3, tsWrapper.ts.dagStats.numRejectedTasks);
      assertEquals(3, tsWrapper.ts.instanceToNodeMap.size());
      assertEquals(3, tsWrapper.ts.disabledNodesQueue.size());


      Object task4 = new Object();
      Object clientCookie4 = new Object();
      Object task5 = new Object();
      Object clientCookie5 = new Object();
      Object task6 = new Object();
      Object clientCookie6 = new Object();
      tsWrapper.allocateTask(task4, hosts1, priority1, clientCookie4);
      tsWrapper.allocateTask(task5, hosts2, priority1, clientCookie5);
      tsWrapper.allocateTask(task6, hosts3, priority1, clientCookie6);
      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numTotalAllocations == 6) {
          break;
        }
      }

      ArgumentCaptor<Container> argumentCaptor = ArgumentCaptor.forClass(Container.class);
      verify(tsWrapper.mockAppCallback, times(3)).taskAllocated(any(Object.class), any(Object.class), argumentCaptor.capture());

      // which affects the locality matching
      assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);
      assertEquals(6, tsWrapper.ts.dagStats.numTotalAllocations);

    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testForceLocalityTest1() throws IOException, InterruptedException {
    // 2 hosts. 2 per host. 5 requests at the same priority.
    // First 3 on host1, Next at host2, Last with no host.
    // Third request on host1 should not be allocated immediately.
    forceLocalityTest1(true);

  }

  @Test(timeout = 10000)
  public void testNoForceLocalityCounterTest1() throws IOException, InterruptedException {
    // 2 hosts. 2 per host. 5 requests at the same priority.
    // First 3 on host1, Next at host2, Last with no host.
    // Third should allocate on host2, 4th on host2, 5th will wait.

    forceLocalityTest1(false);
  }

  private void forceLocalityTest1(boolean forceLocality) throws IOException, InterruptedException {
    Priority priority1 = Priority.newInstance(1);

    String[] hosts = new String[] {HOST1, HOST2};

    String[] hostsH1 = new String[] {HOST1};
    String[] hostsH2 = new String[] {HOST2};

    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1, (forceLocality ? -1l : 0l));

    try {
      Object task1 = "task1";
      Object clientCookie1 = "cookie1";
      Object task2 = "task2";
      Object clientCookie2 = "cookie2";
      Object task3 = "task3";
      Object clientCookie3 = "cookie3";
      Object task4 = "task4";
      Object clientCookie4 = "cookie4";
      Object task5 = "task5";
      Object clientCookie5 = "cookie5";

      tsWrapper.controlScheduler(true);
      //H1 - should allocate
      tsWrapper.allocateTask(task1, hostsH1, priority1, clientCookie1);
      //H1 - should allocate
      tsWrapper.allocateTask(task2, hostsH1, priority1, clientCookie2);
      //H1 - no capacity if force, should allocate otherwise
      tsWrapper.allocateTask(task3, hostsH1, priority1, clientCookie3);
      //H2 - should allocate
      tsWrapper.allocateTask(task4, hostsH2, priority1, clientCookie4);
      //No location - should allocate if force, no capacity otherwise
      tsWrapper.allocateTask(task5, null, priority1, clientCookie5);

      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numTotalAllocations == 4) {
          break;
        }
      }

      // Verify no preemption requests - since everything is at the same priority
      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
      verify(tsWrapper.mockAppCallback, times(4)).taskAllocated(argumentCaptor.capture(), any(Object.class), any(Container.class));
      assertEquals(4, argumentCaptor.getAllValues().size());
      assertEquals(task1, argumentCaptor.getAllValues().get(0));
      assertEquals(task2, argumentCaptor.getAllValues().get(1));
      if (forceLocality) {
        // task3 not allocated
        assertEquals(task4, argumentCaptor.getAllValues().get(2));
        assertEquals(task5, argumentCaptor.getAllValues().get(3));
      } else {
        assertEquals(task3, argumentCaptor.getAllValues().get(2));
        assertEquals(task4, argumentCaptor.getAllValues().get(3));
      }

      //Complete one task on host1.
      tsWrapper.deallocateTask(task1, true, null);

      reset(tsWrapper.mockAppCallback);

      // Try scheduling again.
      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numTotalAllocations == 5) {
          break;
        }
      }
      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      argumentCaptor = ArgumentCaptor.forClass(Object.class);
      verify(tsWrapper.mockAppCallback, times(1)).taskAllocated(argumentCaptor.capture(), any(Object.class), any(Container.class));
      assertEquals(1, argumentCaptor.getAllValues().size());
      if (forceLocality) {
        assertEquals(task3, argumentCaptor.getAllValues().get(0));
      } else {
        assertEquals(task5, argumentCaptor.getAllValues().get(0));
      }

    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testForcedLocalityUnknownHost() throws IOException, InterruptedException {
    Priority priority1 = Priority.newInstance(1);

    String[] hostsKnown = new String[]{HOST1};
    String[] hostsUnknown = new String[]{HOST2};

    TestTaskSchedulerServiceWrapper tsWrapper =
        new TestTaskSchedulerServiceWrapper(2000, hostsKnown, 1, 1, -1l);
    try {
      Object task1 = "task1";
      Object clientCookie1 = "cookie1";

      Object task2 = "task2";
      Object clientCookie2 = "cookie2";

      tsWrapper.controlScheduler(true);
      // Should allocate since H2 is not known.
      tsWrapper.allocateTask(task1, hostsUnknown, priority1, clientCookie1);
      tsWrapper.allocateTask(task2, hostsKnown, priority1, clientCookie2);


      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numTotalAllocations == 2) {
          break;
        }
      }

      ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
      verify(tsWrapper.mockAppCallback, times(2))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), any(Container.class));
      assertEquals(2, argumentCaptor.getAllValues().size());
      assertEquals(task1, argumentCaptor.getAllValues().get(0));
      assertEquals(task2, argumentCaptor.getAllValues().get(1));


    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testHostPreferenceUnknownAndNotSpecified() throws IOException, InterruptedException {
    Priority priority1 = Priority.newInstance(1);

    String[] hostsKnown = new String[]{HOST1, HOST2};
    String[] hostsUnknown = new String[]{HOST3};
    String[] noHosts = new String[]{};

    TestTaskSchedulerServiceWrapper tsWrapper =
      new TestTaskSchedulerServiceWrapper(2000, hostsKnown, 1, 1, -1l);
    try {
      Object task1 = "task1";
      Object clientCookie1 = "cookie1";

      Object task2 = "task2";
      Object clientCookie2 = "cookie2";

      Object task3 = "task3";
      Object clientCookie3 = "cookie3";

      Object task4 = "task4";
      Object clientCookie4 = "cookie4";

      tsWrapper.controlScheduler(true);
      tsWrapper.allocateTask(task1, hostsKnown, priority1, clientCookie1);
      tsWrapper.allocateTask(task2, hostsKnown, priority1, clientCookie2);
      tsWrapper.allocateTask(task3, hostsUnknown, priority1, clientCookie3);
      tsWrapper.allocateTask(task4, noHosts, priority1, clientCookie4);

      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numTotalAllocations == 4) {
          break;
        }
      }

      ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
      ArgumentCaptor<Container> argumentCaptor2 = ArgumentCaptor.forClass(Container.class);
      verify(tsWrapper.mockAppCallback, times(4))
        .taskAllocated(argumentCaptor.capture(), any(Object.class), argumentCaptor2.capture());
      assertEquals(4, argumentCaptor.getAllValues().size());
      assertEquals(task1, argumentCaptor.getAllValues().get(0));
      assertEquals(task2, argumentCaptor.getAllValues().get(1));
      assertEquals(task3, argumentCaptor.getAllValues().get(2));
      assertEquals(task4, argumentCaptor.getAllValues().get(3));
      // 1st task requested host1, got host1
      assertEquals(HOST1, argumentCaptor2.getAllValues().get(0).getNodeId().getHost());
      // 2nd task requested host1, got host1
      assertEquals(HOST1, argumentCaptor2.getAllValues().get(1).getNodeId().getHost());
      // 3rd task requested unknown host, got host2 since host1 is full and only host2 is left in random pool
      assertEquals(HOST2, argumentCaptor2.getAllValues().get(2).getNodeId().getHost());
      // 4rd task provided no location preference, got host2 since host1 is full and only host2 is left in random pool
      assertEquals(HOST2, argumentCaptor2.getAllValues().get(3).getNodeId().getHost());

      assertEquals(1, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);
      assertEquals(2, tsWrapper.ts.dagStats.numLocalAllocations);
      assertEquals(1, tsWrapper.ts.dagStats.numNonLocalAllocations);
    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testHostPreferenceMissesConsistentRollover() throws IOException, InterruptedException {
    Priority priority1 = Priority.newInstance(1);

    String[] hostsKnown = new String[]{HOST1, HOST2, HOST3};
    String[] hostsLive = new String[]{HOST1, HOST2, HOST3};
    String[] hostsH2 = new String[]{HOST2};
    TestTaskSchedulerServiceWrapper tsWrapper =
      new TestTaskSchedulerServiceWrapper(2000, hostsKnown, 1, 0, 0l, false, hostsLive, true);
    try {
      Object task1 = "task1";
      Object clientCookie1 = "cookie1";

      Object task2 = "task2";
      Object clientCookie2 = "cookie2";

      Object task3 = "task3";
      Object clientCookie3 = "cookie3";

      tsWrapper.controlScheduler(true);
      tsWrapper.allocateTask(task1, hostsH2, priority1, clientCookie1);
      tsWrapper.allocateTask(task2, hostsH2, priority1, clientCookie2);
      tsWrapper.allocateTask(task3, hostsH2, priority1, clientCookie3);

      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numTotalAllocations == 3) {
          break;
        }
      }

      ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
      ArgumentCaptor<Container> argumentCaptor2 = ArgumentCaptor.forClass(Container.class);
      verify(tsWrapper.mockAppCallback, times(3))
        .taskAllocated(argumentCaptor.capture(), any(Object.class), argumentCaptor2.capture());
      assertEquals(3, argumentCaptor.getAllValues().size());
      assertEquals(task1, argumentCaptor.getAllValues().get(0));
      assertEquals(task2, argumentCaptor.getAllValues().get(1));
      assertEquals(task3, argumentCaptor.getAllValues().get(2));
      // 1st task requested host2, got host2
      assertEquals(HOST2, argumentCaptor2.getAllValues().get(0).getNodeId().getHost());
      // 2nd task requested host2, got host3 as host2 is full
      assertEquals(HOST3, argumentCaptor2.getAllValues().get(1).getNodeId().getHost());
      // 3rd task requested host2, got host1 as host2 and host3 are full
      assertEquals(HOST1, argumentCaptor2.getAllValues().get(2).getNodeId().getHost());

      verify(tsWrapper.mockServiceInstanceSet, times(2)).getAllInstancesOrdered(true);

      assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);
      assertEquals(1, tsWrapper.ts.dagStats.numLocalAllocations);
      assertEquals(2, tsWrapper.ts.dagStats.numNonLocalAllocations);
    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testHostPreferenceMissesConsistentPartialAlive() throws IOException, InterruptedException {
    Priority priority1 = Priority.newInstance(1);

    String[] hostsKnown = new String[]{HOST1, HOST2, HOST3, HOST4};
    String[] hostsLive = new String[]{HOST1, HOST2, null, HOST4}; // host3 dead before scheduling
    String[] hostsH2 = new String[]{HOST2};
    String[] hostsH3 = new String[]{HOST3};
    TestTaskSchedulerServiceWrapper tsWrapper =
      new TestTaskSchedulerServiceWrapper(2000, hostsKnown, 1, 0, 0l, false, hostsLive, true);
    try {
      Object task1 = "task1";
      Object clientCookie1 = "cookie1";

      Object task2 = "task2";
      Object clientCookie2 = "cookie2";

      Object task3 = "task3";
      Object clientCookie3 = "cookie3";

      tsWrapper.controlScheduler(true);
      tsWrapper.allocateTask(task1, hostsH2, priority1, clientCookie1);
      tsWrapper.allocateTask(task2, hostsH2, priority1, clientCookie2);
      tsWrapper.allocateTask(task3, hostsH3, priority1, clientCookie3);

      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numTotalAllocations == 3) {
          break;
        }
      }

      ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
      ArgumentCaptor<Container> argumentCaptor2 = ArgumentCaptor.forClass(Container.class);
      verify(tsWrapper.mockAppCallback, times(3))
        .taskAllocated(argumentCaptor.capture(), any(Object.class), argumentCaptor2.capture());
      assertEquals(3, argumentCaptor.getAllValues().size());
      assertEquals(task1, argumentCaptor.getAllValues().get(0));
      assertEquals(task2, argumentCaptor.getAllValues().get(1));
      assertEquals(task3, argumentCaptor.getAllValues().get(2));

      // 1st task requested host2, got host2
      assertEquals(HOST2, argumentCaptor2.getAllValues().get(0).getNodeId().getHost());
      // 2nd task requested host2, got host4 since host3 is dead and host2 is full
      assertEquals(HOST4, argumentCaptor2.getAllValues().get(1).getNodeId().getHost());
      // 3rd task requested host3, got host1 since host3 is dead and host4 is full
      assertEquals(HOST1, argumentCaptor2.getAllValues().get(2).getNodeId().getHost());

      verify(tsWrapper.mockServiceInstanceSet, times(2)).getAllInstancesOrdered(true);

      assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);
      assertEquals(1, tsWrapper.ts.dagStats.numLocalAllocations);
      assertEquals(2, tsWrapper.ts.dagStats.numNonLocalAllocations);
    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testForcedLocalityPreemption() throws IOException, InterruptedException {
    Priority priority1 = Priority.newInstance(1);
    Priority priority2 = Priority.newInstance(2);
    String [] hosts = new String[] {HOST1, HOST2};

    String [] hostsH1 = new String[] {HOST1};

    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(2000, hosts, 2, 0, -1l);

    // Fill up host1 with p2 tasks.
    // Leave host2 empty
    // Try running p1 task on host1 - should preempt

    try {
      Object task1 = "task1";
      Object clientCookie1 = "cookie1";
      Object task2 = "task2";
      Object clientCookie2 = "cookie2";
      Object task3 = "task3";
      Object clientCookie3 = "cookie3";
      Object task4 = "task4";
      Object clientCookie4 = "cookie4";

      tsWrapper.controlScheduler(true);
      tsWrapper.allocateTask(task1, hostsH1, priority2, clientCookie1);
      tsWrapper.allocateTask(task2, hostsH1, priority2, clientCookie2);
      // This request at a lower priority should not affect anything.
      tsWrapper.allocateTask(task3, hostsH1, priority2, clientCookie3);
      tsWrapper.awaitLocalTaskAllocations(2);

      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
      verify(tsWrapper.mockAppCallback, times(2))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), any(Container.class));
      assertEquals(2, argumentCaptor.getAllValues().size());
      assertEquals(task1, argumentCaptor.getAllValues().get(0));
      assertEquals(task2, argumentCaptor.getAllValues().get(1));

      reset(tsWrapper.mockAppCallback);
      // Allocate t4 at higher priority. t3 should not be allocated,
      // and a preemption should be attempted on host1, despite host2 having available capacity
      tsWrapper.allocateTask(task4, hostsH1, priority1, clientCookie4);

      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numPreemptedTasks == 1) {
          break;
        }
      }
      verify(tsWrapper.mockAppCallback).preemptContainer(any(ContainerId.class));

      tsWrapper.deallocateTask(task1, false, TaskAttemptEndReason.INTERNAL_PREEMPTION);

      tsWrapper.awaitLocalTaskAllocations(3);

      verify(tsWrapper.mockAppCallback, times(1)).taskAllocated(eq(task4),
          eq(clientCookie4), any(Container.class));

    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testPreemptionChoiceTimeOrdering() throws IOException, InterruptedException {

    Priority priority1 = Priority.newInstance(1);
    Priority priority2 = Priority.newInstance(2);

    String [] hosts = new String[] {HOST1, HOST2};

    String [] hostsH1 = new String[] {HOST1};

    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1, -1l);

    try {
      Object task1 = "task1";
      Object clientCookie1 = "cookie1";
      Object task2 = "task2";
      Object clientCookie2 = "cookie2";
      Object task3 = "task3";
      Object clientCookie3 = "cookie3";

      tsWrapper.controlScheduler(true);

      ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
      ArgumentCaptor<Container> cArgCaptor = ArgumentCaptor.forClass(Container.class);

      // Request task1
      tsWrapper.getClock().setTime(10000l);
      tsWrapper.allocateTask(task1, hostsH1, priority2, clientCookie1);
      tsWrapper.awaitLocalTaskAllocations(1);
      verify(tsWrapper.mockAppCallback, times(1))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), cArgCaptor.capture());
      ContainerId t1Cid = cArgCaptor.getValue().getId();

      reset(tsWrapper.mockAppCallback);
      // Move clock backwards (so that t1 allocation is after t2 allocation)
      // Request task2 (task1 already started at previously set time)
      tsWrapper.getClock().setTime(tsWrapper.getClock().getTime() - 1000);
      tsWrapper.allocateTask(task2, hostsH1, priority2, clientCookie2);
      tsWrapper.awaitLocalTaskAllocations(2);
      verify(tsWrapper.mockAppCallback, times(1))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), cArgCaptor.capture());


      reset(tsWrapper.mockAppCallback);
      // Move clock forward, and request a task at p=1
      tsWrapper.getClock().setTime(tsWrapper.getClock().getTime() + 2000);

      tsWrapper.allocateTask(task3, hostsH1, priority1, clientCookie3);
      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numPreemptedTasks == 1) {
          break;
        }
      }
      // Ensure task1 is preempted based on time (match it's allocated containerId)
      ArgumentCaptor<ContainerId> cIdArgCaptor = ArgumentCaptor.forClass(ContainerId.class);
      verify(tsWrapper.mockAppCallback).preemptContainer(cIdArgCaptor.capture());
      assertEquals(t1Cid, cIdArgCaptor.getValue());

    } finally {
      tsWrapper.shutdown();
    }

  }

  @Test(timeout = 10000)
  public void testForcedLocalityMultiplePreemptionsSameHost1() throws IOException,
      InterruptedException {
    Priority priority1 = Priority.newInstance(1);
    Priority priority2 = Priority.newInstance(2);

    String [] hosts = new String[] {HOST1, HOST2};

    String [] hostsH1 = new String[] {HOST1};

    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1, -1l);

    // Fill up host1 with p2 tasks.
    // Leave host2 empty
    // Try running p1 task on host1 - should preempt
    // Await preemption request.
    // Try running another p1 task on host1 - should preempt
    // Await preemption request.

    try {

      Object task1 = "task1";
      Object clientCookie1 = "cookie1";
      Object task2 = "task2";
      Object clientCookie2 = "cookie2";
      Object task3 = "task3";
      Object clientCookie3 = "cookie3";
      Object task4 = "task4";
      Object clientCookie4 = "cookie4";

      tsWrapper.controlScheduler(true);
      tsWrapper.allocateTask(task1, hostsH1, priority2, clientCookie1);
      tsWrapper.allocateTask(task2, hostsH1, priority2, clientCookie2);

      tsWrapper.awaitLocalTaskAllocations(2);
      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
      ArgumentCaptor<Container> cArgCaptor = ArgumentCaptor.forClass(Container.class);
      verify(tsWrapper.mockAppCallback, times(2))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), cArgCaptor.capture());
      assertEquals(2, argumentCaptor.getAllValues().size());
      assertEquals(task1, argumentCaptor.getAllValues().get(0));
      assertEquals(task2, argumentCaptor.getAllValues().get(1));
      assertEquals(2, cArgCaptor.getAllValues().size());
      ContainerId t1CId = cArgCaptor.getAllValues().get(0).getId();

      reset(tsWrapper.mockAppCallback);
      // At this point. 2 tasks running - both at priority 2.
      // Try running a priority 1 task
      tsWrapper.allocateTask(task3, hostsH1, priority1, clientCookie3);
      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numPreemptedTasks == 1) {
          break;
        }
      }
      ArgumentCaptor<ContainerId> cIdArgCaptor = ArgumentCaptor.forClass(ContainerId.class);
      verify(tsWrapper.mockAppCallback).preemptContainer(cIdArgCaptor.capture());

      // Determin which task has been preempted. Normally task2 would be preempted based on it starting
      // later. However - both may have the same start time, so either could be picked.
      Object deallocatedTask1; // De-allocated now
      Object deallocatedTask2; // Will be de-allocated later.
      if (cIdArgCaptor.getValue().equals(t1CId)) {
        deallocatedTask1 = task1;
        deallocatedTask2 = task2;
      } else {
        deallocatedTask1 = task2;
        deallocatedTask2 = task1;
      }

      tsWrapper.deallocateTask(deallocatedTask1, false, TaskAttemptEndReason.INTERNAL_PREEMPTION);

      tsWrapper.awaitLocalTaskAllocations(3);

      verify(tsWrapper.mockAppCallback, times(1)).taskAllocated(eq(task3),
          eq(clientCookie3), any(Container.class));

      reset(tsWrapper.mockAppCallback);
      // At this point. one p=2 task and task3(p=1) running. Ask for another p1 task.
      tsWrapper.allocateTask(task4, hostsH1, priority1, clientCookie4);
      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numPreemptedTasks == 2) {
          break;
        }
      }
      verify(tsWrapper.mockAppCallback).preemptContainer(any(ContainerId.class));

      tsWrapper.deallocateTask(deallocatedTask2, false, TaskAttemptEndReason.INTERNAL_PREEMPTION);

      tsWrapper.awaitLocalTaskAllocations(4);

      verify(tsWrapper.mockAppCallback, times(1)).taskAllocated(eq(task4),
          eq(clientCookie4), any(Container.class));


    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testForcedLocalityMultiplePreemptionsSameHost2() throws IOException,
      InterruptedException {
    Priority priority1 = Priority.newInstance(1);
    Priority priority2 = Priority.newInstance(2);

    String [] hosts = new String[] {HOST1, HOST2};

    String [] hostsH1 = new String[] {HOST1};

    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1, -1l);

    // Fill up host1 with p2 tasks.
    // Leave host2 empty
    // Try running both p1 tasks on host1.
    // R: Single preemption triggered, followed by allocation, followed by another preemption.
    //

    try {

      Object task1 = "task1";
      Object clientCookie1 = "cookie1";
      Object task2 = "task2";
      Object clientCookie2 = "cookie2";
      Object task3 = "task3";
      Object clientCookie3 = "cookie3";
      Object task4 = "task4";
      Object clientCookie4 = "cookie4";

      tsWrapper.controlScheduler(true);
      tsWrapper.allocateTask(task1, hostsH1, priority2, clientCookie1);
      tsWrapper.allocateTask(task2, hostsH1, priority2, clientCookie2);

      tsWrapper.awaitLocalTaskAllocations(2);
      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
      ArgumentCaptor<Container> cArgCaptor = ArgumentCaptor.forClass(Container.class);
      verify(tsWrapper.mockAppCallback, times(2))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), cArgCaptor.capture());
      assertEquals(2, argumentCaptor.getAllValues().size());
      assertEquals(task1, argumentCaptor.getAllValues().get(0));
      assertEquals(task2, argumentCaptor.getAllValues().get(1));
      assertEquals(2, cArgCaptor.getAllValues().size());
      ContainerId t1CId = cArgCaptor.getAllValues().get(0).getId();

      reset(tsWrapper.mockAppCallback);
      // At this point. 2 tasks running - both at priority 2.
      // Try running a priority 1 task
      tsWrapper.allocateTask(task3, hostsH1, priority1, clientCookie3);
      tsWrapper.allocateTask(task4, hostsH1, priority1, clientCookie4);
      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        if (tsWrapper.ts.dagStats.numPreemptedTasks == 1) {
          break;
        }
      }
      ArgumentCaptor<ContainerId> cIdArgCaptor = ArgumentCaptor.forClass(ContainerId.class);
      verify(tsWrapper.mockAppCallback).preemptContainer(cIdArgCaptor.capture());

      // Determin which task has been preempted. Normally task2 would be preempted based on it starting
      // later. However - both may have the same start time, so either could be picked.
      Object deallocatedTask1; // De-allocated now
      Object deallocatedTask2; // Will be de-allocated later.
      if (cIdArgCaptor.getValue().equals(t1CId)) {
        deallocatedTask1 = task1;
        deallocatedTask2 = task2;
      } else {
        deallocatedTask1 = task2;
        deallocatedTask2 = task1;
      }

      tsWrapper.deallocateTask(deallocatedTask1, false, TaskAttemptEndReason.INTERNAL_PREEMPTION);

      tsWrapper.awaitLocalTaskAllocations(3);

      verify(tsWrapper.mockAppCallback, times(1)).taskAllocated(eq(task3),
          eq(clientCookie3), any(Container.class));

      // At this point. one p=2 task and task3(p=1) running. Ask for another p1 task.
      while (true) {
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun(1000l);
        if (tsWrapper.ts.dagStats.numPreemptedTasks == 2) {
          break;
        }
      }
      verify(tsWrapper.mockAppCallback, times(2)).preemptContainer(any(ContainerId.class));

      tsWrapper.deallocateTask(deallocatedTask2, false, TaskAttemptEndReason.INTERNAL_PREEMPTION);

      tsWrapper.awaitLocalTaskAllocations(4);

      verify(tsWrapper.mockAppCallback, times(1)).taskAllocated(eq(task4),
          eq(clientCookie4), any(Container.class));


    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testForcedLocalityNotInDelayedQueue() throws IOException, InterruptedException {
    String[] hosts = new String[]{HOST1, HOST2};

    String[] hostsH1 = new String[]{HOST1};

    TestTaskSchedulerServiceWrapper tsWrapper =
        new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1, -1l);
    testNotInQueue(tsWrapper, hostsH1);
  }

  @Test(timeout = 10000)
  public void testNoLocalityNotInDelayedQueue() throws IOException, InterruptedException {
    String[] hosts = new String[]{HOST1};

    String[] hostsH1 = new String[]{HOST1};

    TestTaskSchedulerServiceWrapper tsWrapper =
        new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1, 0l);
    testNotInQueue(tsWrapper, hostsH1);
  }

  private void testNotInQueue(TestTaskSchedulerServiceWrapper tsWrapper, String[] hosts) throws
      InterruptedException {
    Priority priority1 = Priority.newInstance(1);
    try {
      tsWrapper.controlScheduler(true);
      tsWrapper.allocateTask(hosts, priority1);
      tsWrapper.allocateTask(hosts, priority1);
      tsWrapper.allocateTask(hosts, priority1); // 1 more than capacity.

      tsWrapper.awaitLocalTaskAllocations(2);

      assertEquals(0, tsWrapper.ts.delayedTaskQueue.size());

    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testDelayedLocalityFallbackToNonLocal() throws IOException, InterruptedException {
    Priority priority1 = Priority.newInstance(1);
    String [] hosts = new String[] {HOST1, HOST2};

    String [] hostsH1 = new String[] {HOST1};

    TestTaskSchedulerServiceWrapper tsWrapper =
        new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1, 10000l, true);
    LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled
        delayedTaskSchedulerCallableControlled =
        (LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled) tsWrapper.ts.delayedTaskSchedulerCallable;
    ControlledClock clock = tsWrapper.getClock();
    clock.setTime(clock.getTime());

    // Fill up host1 with tasks. Leave host2 empty.
    try {
      tsWrapper.controlScheduler(true);
      Object task1 = tsWrapper.allocateTask(hostsH1, priority1);
      Object task2 = tsWrapper.allocateTask(hostsH1, priority1);
      Object task3 = tsWrapper.allocateTask(hostsH1, priority1); // 1 more than capacity.

      tsWrapper.awaitLocalTaskAllocations(2);

      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
      verify(tsWrapper.mockAppCallback, times(2))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), any(Container.class));
      assertEquals(2, argumentCaptor.getAllValues().size());
      assertEquals(task1, argumentCaptor.getAllValues().get(0));
      assertEquals(task2, argumentCaptor.getAllValues().get(1));

      reset(tsWrapper.mockAppCallback);

      // No capacity left on node1. The next task should be allocated to node2 after it times out.
      clock.setTime(clock.getTime() + 10000l); // Past the timeout.

      assertEquals(
          LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled.STATE_NOT_RUN,
          delayedTaskSchedulerCallableControlled.lastState);

      delayedTaskSchedulerCallableControlled.triggerGetNextTask();
      delayedTaskSchedulerCallableControlled.awaitGetNextTaskProcessing();

      // Verify that an attempt was made to schedule the task, but the decision was to skip scheduling
      assertEquals(
          LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled.STATE_RETURNED_TASK,
          delayedTaskSchedulerCallableControlled.lastState);
      assertTrue(delayedTaskSchedulerCallableControlled.shouldScheduleTaskTriggered &&
          delayedTaskSchedulerCallableControlled.lastShouldScheduleTaskResult);

      tsWrapper.awaitChangeInTotalAllocations(2);

      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      argumentCaptor = ArgumentCaptor.forClass(Object.class);
      ArgumentCaptor<Container> containerCaptor = ArgumentCaptor.forClass(Container.class);
      verify(tsWrapper.mockAppCallback, times(1))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), containerCaptor.capture());
      assertEquals(1, argumentCaptor.getAllValues().size());
      assertEquals(task3, argumentCaptor.getAllValues().get(0));
      Container assignedContainer = containerCaptor.getValue();
      assertEquals(HOST2, assignedContainer.getNodeId().getHost());


      assertEquals(2, tsWrapper.ts.dagStats.numLocalAllocations);
      assertEquals(1, tsWrapper.ts.dagStats.numNonLocalAllocations);
      assertEquals(1, tsWrapper.ts.dagStats.numDelayedAllocations);
      assertEquals(2, tsWrapper.ts.dagStats.numAllocationsPerHost.get(HOST1).get());
      assertEquals(1, tsWrapper.ts.dagStats.numAllocationsPerHost.get(HOST2).get());

    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testDelayedLocalityDelayedAllocation() throws InterruptedException, IOException {
    Priority priority1 = Priority.newInstance(1);
    String [] hosts = new String[] {HOST1, HOST2};

    String [] hostsH1 = new String[] {HOST1};

    TestTaskSchedulerServiceWrapper tsWrapper =
        new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1, 10000l, true);
    LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled
        delayedTaskSchedulerCallableControlled =
        (LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled) tsWrapper.ts.delayedTaskSchedulerCallable;
    ControlledClock clock = tsWrapper.getClock();
    clock.setTime(clock.getTime());

    // Fill up host1 with tasks. Leave host2 empty.
    try {
      tsWrapper.controlScheduler(true);
      Object task1 = tsWrapper.allocateTask(hostsH1, priority1);
      Object task2 = tsWrapper.allocateTask(hostsH1, priority1);
      Object task3 = tsWrapper.allocateTask(hostsH1, priority1); // 1 more than capacity.

      tsWrapper.awaitLocalTaskAllocations(2);

      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
      verify(tsWrapper.mockAppCallback, times(2))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), any(Container.class));
      assertEquals(2, argumentCaptor.getAllValues().size());
      assertEquals(task1, argumentCaptor.getAllValues().get(0));
      assertEquals(task2, argumentCaptor.getAllValues().get(1));

      reset(tsWrapper.mockAppCallback);

      // Move the clock forward 2000ms, and check the delayed queue
      clock.setTime(clock.getTime() + 2000l); // Past the timeout.

      assertEquals(
          LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled.STATE_NOT_RUN,
          delayedTaskSchedulerCallableControlled.lastState);

      delayedTaskSchedulerCallableControlled.triggerGetNextTask();
      delayedTaskSchedulerCallableControlled.awaitGetNextTaskProcessing();

      // Verify that an attempt was made to schedule the task, but the decision was to skip scheduling
      assertEquals(
          LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled.STATE_TIMEOUT_NOT_EXPIRED,
          delayedTaskSchedulerCallableControlled.lastState);
      assertFalse(delayedTaskSchedulerCallableControlled.shouldScheduleTaskTriggered);

      tsWrapper.deallocateTask(task1, true, null);

      // Node1 now has free capacity. task1 should be allocated to it.
      tsWrapper.awaitChangeInTotalAllocations(2);

      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      argumentCaptor = ArgumentCaptor.forClass(Object.class);
      ArgumentCaptor<Container> containerCaptor = ArgumentCaptor.forClass(Container.class);
      verify(tsWrapper.mockAppCallback, times(1))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), containerCaptor.capture());
      assertEquals(1, argumentCaptor.getAllValues().size());
      assertEquals(task3, argumentCaptor.getAllValues().get(0));
      Container assignedContainer = containerCaptor.getValue();
      assertEquals(HOST1, assignedContainer.getNodeId().getHost());


      assertEquals(3, tsWrapper.ts.dagStats.numLocalAllocations);
      assertEquals(0, tsWrapper.ts.dagStats.numNonLocalAllocations);
      assertEquals(1, tsWrapper.ts.dagStats.numDelayedAllocations);
      assertEquals(3, tsWrapper.ts.dagStats.numAllocationsPerHost.get(HOST1).get());

    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testDelayedQueeTaskSelectionAfterScheduled() throws IOException,
      InterruptedException {
    Priority priority1 = Priority.newInstance(1);
    String [] hosts = new String[] {HOST1, HOST2};

    String [] hostsH1 = new String[] {HOST1};

    TestTaskSchedulerServiceWrapper tsWrapper =
        new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1, 10000l, true);
    LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled
        delayedTaskSchedulerCallableControlled =
        (LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled) tsWrapper.ts.delayedTaskSchedulerCallable;
    ControlledClock clock = tsWrapper.getClock();
    clock.setTime(clock.getTime());

    // Fill up host1 with tasks. Leave host2 empty.
    try {
      tsWrapper.controlScheduler(true);
      Object task1 = tsWrapper.allocateTask(hostsH1, priority1);
      Object task2 = tsWrapper.allocateTask(hostsH1, priority1);
      Object task3 = tsWrapper.allocateTask(hostsH1, priority1); // 1 more than capacity.

      tsWrapper.awaitLocalTaskAllocations(2);

      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
      verify(tsWrapper.mockAppCallback, times(2))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), any(Container.class));
      assertEquals(2, argumentCaptor.getAllValues().size());
      assertEquals(task1, argumentCaptor.getAllValues().get(0));
      assertEquals(task2, argumentCaptor.getAllValues().get(1));

      // Simulate a 2s delay before finishing the task.
      clock.setTime(clock.getTime() + 2000);

      assertEquals(
          LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled.STATE_NOT_RUN,
          delayedTaskSchedulerCallableControlled.lastState);

      delayedTaskSchedulerCallableControlled.triggerGetNextTask();
      delayedTaskSchedulerCallableControlled.awaitGetNextTaskProcessing();
      assertEquals(
          LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled.STATE_TIMEOUT_NOT_EXPIRED,
          delayedTaskSchedulerCallableControlled.lastState);
      assertFalse(delayedTaskSchedulerCallableControlled.shouldScheduleTaskTriggered);

      reset(tsWrapper.mockAppCallback);

      // Now finish task1, which will make capacity for task3 to run. Nothing is coming out of the delayed queue yet.
      tsWrapper.deallocateTask(task1, true, null);
      tsWrapper.awaitLocalTaskAllocations(3);

      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      argumentCaptor = ArgumentCaptor.forClass(Object.class);
      ArgumentCaptor<Container> containerCaptor = ArgumentCaptor.forClass(Container.class);
      verify(tsWrapper.mockAppCallback, times(1))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), containerCaptor.capture());
      assertEquals(1, argumentCaptor.getAllValues().size());
      assertEquals(task3, argumentCaptor.getAllValues().get(0));
      Container assignedContainer = containerCaptor.getValue();
      assertEquals(HOST1, assignedContainer.getNodeId().getHost());

      reset(tsWrapper.mockAppCallback);

      // Move the clock forward and trigger a run.
      clock.setTime(clock.getTime() + 8000); // Set to start + 10000 which is the timeout
      delayedTaskSchedulerCallableControlled.triggerGetNextTask();
      delayedTaskSchedulerCallableControlled.awaitGetNextTaskProcessing();
      assertEquals(
          LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled.STATE_RETURNED_TASK,
          delayedTaskSchedulerCallableControlled.lastState);
      // Verify that an attempt was made to schedule the task, but the decision was to skip scheduling
      assertTrue(delayedTaskSchedulerCallableControlled.shouldScheduleTaskTriggered &&
          !delayedTaskSchedulerCallableControlled.lastShouldScheduleTaskResult);

      // Ensure there's no more invocations.
      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      verify(tsWrapper.mockAppCallback, never()).taskAllocated(any(Object.class), any(Object.class), any(Container.class));

    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test(timeout = 10000)
  public void testTaskInfoDelay() {

    LlapTaskSchedulerService.LocalityDelayConf localityDelayConf1 =
        new LlapTaskSchedulerService.LocalityDelayConf(3000);

    ControlledClock clock = new ControlledClock(new MonotonicClock());
    clock.setTime(clock.getTime());


    // With a timeout of 3000.
    LlapTaskSchedulerService.TaskInfo taskInfo =
        new LlapTaskSchedulerService.TaskInfo(localityDelayConf1, clock, new Object(), new Object(),
            mock(Priority.class), mock(Resource.class), null, null, clock.getTime());

    assertFalse(taskInfo.shouldForceLocality());

    assertEquals(3000, taskInfo.getDelay(TimeUnit.MILLISECONDS));
    assertTrue(taskInfo.shouldDelayForLocality(clock.getTime()));

    clock.setTime(clock.getTime() + 500);
    assertEquals(2500, taskInfo.getDelay(TimeUnit.MILLISECONDS));
    assertTrue(taskInfo.shouldDelayForLocality(clock.getTime()));

    clock.setTime(clock.getTime() + 2500);
    assertEquals(0, taskInfo.getDelay(TimeUnit.MILLISECONDS));
    assertFalse(taskInfo.shouldDelayForLocality(clock.getTime()));


    // No locality delay
    LlapTaskSchedulerService.LocalityDelayConf localityDelayConf2 =
        new LlapTaskSchedulerService.LocalityDelayConf(0);
    taskInfo =
        new LlapTaskSchedulerService.TaskInfo(localityDelayConf2, clock, new Object(), new Object(),
            mock(Priority.class), mock(Resource.class), null, null, clock.getTime());
    assertFalse(taskInfo.shouldDelayForLocality(clock.getTime()));
    assertFalse(taskInfo.shouldForceLocality());
    assertTrue(taskInfo.getDelay(TimeUnit.MILLISECONDS) < 0);

    // Force locality
    LlapTaskSchedulerService.LocalityDelayConf localityDelayConf3 =
        new LlapTaskSchedulerService.LocalityDelayConf(-1);
    taskInfo =
        new LlapTaskSchedulerService.TaskInfo(localityDelayConf3, clock, new Object(), new Object(),
            mock(Priority.class), mock(Resource.class), null, null, clock.getTime());
    assertTrue(taskInfo.shouldDelayForLocality(clock.getTime()));
    assertTrue(taskInfo.shouldForceLocality());
    assertFalse(taskInfo.getDelay(TimeUnit.MILLISECONDS) < 0);
  }

  @Test(timeout = 10000)
  public void testLocalityDelayTaskOrdering() throws InterruptedException, IOException {

    LlapTaskSchedulerService.LocalityDelayConf localityDelayConf =
        new LlapTaskSchedulerService.LocalityDelayConf(3000);

    ControlledClock clock = new ControlledClock(new MonotonicClock());
    clock.setTime(clock.getTime());

    DelayQueue<LlapTaskSchedulerService.TaskInfo> delayedQueue = new DelayQueue<>();

    LlapTaskSchedulerService.TaskInfo taskInfo1 =
        new LlapTaskSchedulerService.TaskInfo(localityDelayConf, clock, new Object(), new Object(),
            mock(Priority.class), mock(Resource.class), null, null, clock.getTime());

    clock.setTime(clock.getTime() + 1000);
    LlapTaskSchedulerService.TaskInfo taskInfo2 =
        new LlapTaskSchedulerService.TaskInfo(localityDelayConf, clock, new Object(), new Object(),
            mock(Priority.class), mock(Resource.class), null, null, clock.getTime());

    delayedQueue.add(taskInfo1);
    delayedQueue.add(taskInfo2);

    assertEquals(taskInfo1, delayedQueue.peek());
  }

  @Test (timeout = 15000)
  public void testDelayedLocalityNodeCommErrorImmediateAllocation() throws IOException, InterruptedException {
    Priority priority1 = Priority.newInstance(1);
    String [] hosts = new String[] {HOST1, HOST2};

    String [] hostsH1 = new String[] {HOST1};

    // Node disable timeout higher than locality delay.
    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(20000, hosts, 1, 1, 10000l);

    // Fill up host1 with tasks. Leave host2 empty.
    try {
      long startTime = tsWrapper.getClock().getTime();
      tsWrapper.controlScheduler(true);
      Object task1 = tsWrapper.allocateTask(hostsH1, priority1);
      Object task2 = tsWrapper.allocateTask(hostsH1, priority1);
      Object task3 = tsWrapper.allocateTask(hostsH1, priority1); // 1 more than capacity.

      tsWrapper.awaitLocalTaskAllocations(2);

      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
      verify(tsWrapper.mockAppCallback, times(2))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), any(Container.class));
      assertEquals(2, argumentCaptor.getAllValues().size());
      assertEquals(task1, argumentCaptor.getAllValues().get(0));
      assertEquals(task2, argumentCaptor.getAllValues().get(1));

      reset(tsWrapper.mockAppCallback);
      // Flush any pending scheduler runs which may be blocked. Wait 2 seconds for the run to complete.
      tsWrapper.signalSchedulerRun();
      tsWrapper.awaitSchedulerRun(2000l);

      // Mark a task as failed due to a comm failure.
      tsWrapper.deallocateTask(task1, false, TaskAttemptEndReason.COMMUNICATION_ERROR);

      // Node1 marked as failed, node2 has capacity.
      // Timeout for nodes is larger than delay - immediate allocation
      tsWrapper.awaitChangeInTotalAllocations(2);

      long thirdAllocateTime = tsWrapper.getClock().getTime();
      long diff = thirdAllocateTime - startTime;
      // diffAfterSleep < total sleepTime
      assertTrue("Task not allocated in expected time window: duration=" + diff, diff < 10000l);

      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      argumentCaptor = ArgumentCaptor.forClass(Object.class);
      ArgumentCaptor<Container> containerCaptor = ArgumentCaptor.forClass(Container.class);
      verify(tsWrapper.mockAppCallback, times(1))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), containerCaptor.capture());
      assertEquals(1, argumentCaptor.getAllValues().size());
      assertEquals(task3, argumentCaptor.getAllValues().get(0));
      Container assignedContainer = containerCaptor.getValue();
      assertEquals(HOST2, assignedContainer.getNodeId().getHost());


      assertEquals(2, tsWrapper.ts.dagStats.numLocalAllocations);
      assertEquals(1, tsWrapper.ts.dagStats.numNonLocalAllocations);
      assertEquals(1, tsWrapper.ts.dagStats.numDelayedAllocations);
      assertEquals(2, tsWrapper.ts.dagStats.numAllocationsPerHost.get(HOST1).get());
      assertEquals(1, tsWrapper.ts.dagStats.numAllocationsPerHost.get(HOST2).get());

    } finally {
      tsWrapper.shutdown();
    }
  }

  @Test (timeout = 15000)
  public void testDelayedLocalityNodeCommErrorDelayedAllocation() throws IOException, InterruptedException {
    Priority priority1 = Priority.newInstance(1);
    String [] hosts = new String[] {HOST1, HOST2};

    String [] hostsH1 = new String[] {HOST1};

    TestTaskSchedulerServiceWrapper tsWrapper =
        new TestTaskSchedulerServiceWrapper(5000, hosts, 1, 1, 10000l, true);
    LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled
        delayedTaskSchedulerCallableControlled =
        (LlapTaskSchedulerServiceForTestControlled.DelayedTaskSchedulerCallableControlled) tsWrapper.ts.delayedTaskSchedulerCallable;
    ControlledClock clock = tsWrapper.getClock();
    clock.setTime(clock.getTime());

    // Fill up host1 with tasks. Leave host2 empty.
    try {
      tsWrapper.controlScheduler(true);
      Object task1 = tsWrapper.allocateTask(hostsH1, priority1);
      Object task2 = tsWrapper.allocateTask(hostsH1, priority1);
      Object task3 = tsWrapper.allocateTask(hostsH1, priority1); // 1 more than capacity.

      tsWrapper.awaitLocalTaskAllocations(2);

      verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
      ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
      verify(tsWrapper.mockAppCallback, times(2))
          .taskAllocated(argumentCaptor.capture(), any(Object.class), any(Container.class));
      assertEquals(2, argumentCaptor.getAllValues().size());
      assertEquals(task1, argumentCaptor.getAllValues().get(0));
      assertEquals(task2, argumentCaptor.getAllValues().get(1));

      reset(tsWrapper.mockAppCallback);

      // Mark a task as failed due to a comm failure.
      tsWrapper.deallocateTask(task1, false, TaskAttemptEndReason.COMMUNICATION_ERROR);

      // Node1 has free capacity but is disabled. Node 2 has capcaity. Delay > re-enable tiemout
      tsWrapper.ensureNoChangeInTotalAllocations(2, 2000l);
    } finally {
      tsWrapper.shutdown();
    }
  }

  private static class TestTaskSchedulerServiceWrapper {
    static final Resource resource = Resource.newInstance(1024, 1);
    Configuration conf;
    TaskSchedulerContext mockAppCallback = mock(TaskSchedulerContext.class);
    ServiceInstanceSet mockServiceInstanceSet = mock(ServiceInstanceSet.class);
    ControlledClock clock = new ControlledClock(new MonotonicClock());
    ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1000, 1), 1);
    LlapTaskSchedulerServiceForTest ts;

    TestTaskSchedulerServiceWrapper() throws IOException, InterruptedException {
      this(2000l);
    }

    TestTaskSchedulerServiceWrapper(long disableTimeoutMillis) throws IOException,
        InterruptedException {
      this(disableTimeoutMillis, new String[]{HOST1, HOST2, HOST3}, 4,
          ConfVars.LLAP_DAEMON_TASK_SCHEDULER_WAIT_QUEUE_SIZE.defaultIntVal);
    }

    TestTaskSchedulerServiceWrapper(long disableTimeoutMillis, String[] hosts, int numExecutors, int waitQueueSize) throws
        IOException, InterruptedException {
      this(disableTimeoutMillis, hosts, numExecutors, waitQueueSize, 0l);
    }

    TestTaskSchedulerServiceWrapper(long nodeDisableTimeoutMillis, String[] hosts, int numExecutors,
                                    int waitQueueSize, long localityDelayMs) throws
        IOException, InterruptedException {
      this(nodeDisableTimeoutMillis, hosts, numExecutors, waitQueueSize, localityDelayMs, false);
    }

    TestTaskSchedulerServiceWrapper(long nodeDisableTimeoutMillis, String[] hosts, int numExecutors,
                                    int waitQueueSize, long localityDelayMs, boolean controlledDelayedTaskQueue) throws
        IOException, InterruptedException {
      this(nodeDisableTimeoutMillis, hosts, numExecutors, waitQueueSize, localityDelayMs, controlledDelayedTaskQueue,
        hosts, false);
    }

    TestTaskSchedulerServiceWrapper(long nodeDisableTimeoutMillis, String[] hosts, int numExecutors,
      int waitQueueSize, long localityDelayMs, boolean controlledDelayedTaskQueue, String[] liveHosts,
      boolean useMockRegistry) throws
      IOException, InterruptedException {
      conf = new Configuration();
      conf.setStrings(ConfVars.LLAP_DAEMON_SERVICE_HOSTS.varname, hosts);
      conf.setInt(ConfVars.LLAP_DAEMON_NUM_EXECUTORS.varname, numExecutors);
      conf.setInt(ConfVars.LLAP_DAEMON_TASK_SCHEDULER_WAIT_QUEUE_SIZE.varname, waitQueueSize);
      conf.set(ConfVars.LLAP_TASK_SCHEDULER_NODE_REENABLE_MIN_TIMEOUT_MS.varname,
        nodeDisableTimeoutMillis + "ms");
      conf.setBoolean(LlapFixedRegistryImpl.FIXED_REGISTRY_RESOLVE_HOST_NAMES, false);
      conf.setLong(ConfVars.LLAP_TASK_SCHEDULER_LOCALITY_DELAY.varname, localityDelayMs);

      doReturn(appAttemptId).when(mockAppCallback).getApplicationAttemptId();
      doReturn(11111l).when(mockAppCallback).getCustomClusterIdentifier();
      UserPayload userPayload = TezUtils.createUserPayloadFromConf(conf);
      doReturn(userPayload).when(mockAppCallback).getInitialUserPayload();

      if (useMockRegistry) {
        List<ServiceInstance> liveInstances = new ArrayList<>();
        for (String host : liveHosts) {
          if (host == null) {
            ServiceInstance mockInactive = mock(InactiveServiceInstance.class);
            doReturn(host).when(mockInactive).getHost();
            doReturn("inactive-host-" + host).when(mockInactive).getWorkerIdentity();
            doReturn(ImmutableSet.builder().add(mockInactive).build()).when(mockServiceInstanceSet).getByHost(host);
            liveInstances.add(mockInactive);
          } else {
            ServiceInstance mockActive = mock(ServiceInstance.class);
            doReturn(host).when(mockActive).getHost();
            doReturn("host-" + host).when(mockActive).getWorkerIdentity();
            doReturn(ImmutableSet.builder().add(mockActive).build()).when(mockServiceInstanceSet).getByHost(host);
            liveInstances.add(mockActive);
          }
        }
        doReturn(liveInstances).when(mockServiceInstanceSet).getAllInstancesOrdered(true);

        List<ServiceInstance> allInstances = new ArrayList<>();
        for (String host : hosts) {
          ServiceInstance mockActive = mock(ServiceInstance.class);
          doReturn(host).when(mockActive).getHost();
          doReturn(Resource.newInstance(100, 1)).when(mockActive).getResource();
          doReturn("host-" + host).when(mockActive).getWorkerIdentity();
          allInstances.add(mockActive);
        }
        doReturn(allInstances).when(mockServiceInstanceSet).getAll();
      }
      if (controlledDelayedTaskQueue) {
        ts = new LlapTaskSchedulerServiceForTestControlled(mockAppCallback, clock);
      } else {
        ts = new LlapTaskSchedulerServiceForTest(mockAppCallback, clock);
      }
      controlScheduler(true);
      ts.initialize();
      ts.start();
      if (useMockRegistry) {
        ts.setServiceInstanceSet(mockServiceInstanceSet);
      }
      // One scheduler pass from the nodes that are added at startup
      signalSchedulerRun();
      controlScheduler(false);
      awaitSchedulerRun();
    }

    ControlledClock getClock() {
      return clock;
    }

    void controlScheduler(boolean val) {
      ts.forTestsetControlScheduling(val);
    }

    void signalSchedulerRun() throws InterruptedException {
      ts.forTestSignalSchedulingRun();
    }

    void awaitSchedulerRun() throws InterruptedException {
      ts.forTestAwaitSchedulingRun(-1);
    }

    /**
     *
     * @param timeoutMs
     * @return false if the time elapsed
     * @throws InterruptedException
     */
    boolean awaitSchedulerRun(long timeoutMs) throws InterruptedException {
      return ts.forTestAwaitSchedulingRun(timeoutMs);
    }

    void resetAppCallback() {
      reset(mockAppCallback);
    }

    void shutdown() {
      ts.shutdown();
    }

    void allocateTask(Object task, String[] hosts, Priority priority, Object clientCookie) {
      ts.allocateTask(task, resource, hosts, null, priority, null, clientCookie);
    }



    void deallocateTask(Object task, boolean succeeded, TaskAttemptEndReason endReason) {
      ts.deallocateTask(task, succeeded, endReason, null);
    }

    void rejectExecution(Object task) {
      ts.deallocateTask(task, false, TaskAttemptEndReason.EXECUTOR_BUSY, null);
    }


    // More complex methods which may wrap multiple operations
    Object allocateTask(String[] hosts, Priority priority) {
      Object task = new Object();
      Object clientCookie = new Object();
      allocateTask(task, hosts, priority, clientCookie);
      return task;
    }

    public void awaitTotalTaskAllocations(int numTasks) throws InterruptedException {
      while (true) {
        signalSchedulerRun();
        awaitSchedulerRun();
        if (ts.dagStats.numTotalAllocations == numTasks) {
          break;
        }
      }
    }

    public void awaitLocalTaskAllocations(int numTasks) throws InterruptedException {
      while (true) {
        signalSchedulerRun();
        awaitSchedulerRun();
        if (ts.dagStats.numLocalAllocations == numTasks) {
          break;
        }
      }
    }

    public void awaitChangeInTotalAllocations(int previousAllocations) throws InterruptedException {
      while (true) {
        signalSchedulerRun();
        awaitSchedulerRun();
        if (ts.dagStats.numTotalAllocations > previousAllocations) {
          break;
        }
        Thread.sleep(200l);
      }
    }

    public void ensureNoChangeInTotalAllocations(int previousAllocations, long timeout) throws
        InterruptedException {
      long startTime = Time.monotonicNow();
      long timeLeft = timeout;
      while (timeLeft > 0) {
        signalSchedulerRun();
        awaitSchedulerRun(Math.min(200, timeLeft));
        if (ts.dagStats.numTotalAllocations != previousAllocations) {
          throw new IllegalStateException("NumTotalAllocations expected to stay at " + previousAllocations + ". Actual=" + ts.dagStats.numTotalAllocations);
        }
        timeLeft = (startTime + timeout) - Time.monotonicNow();
      }
    }
  }

  private static class LlapTaskSchedulerServiceForTest extends LlapTaskSchedulerService {

    private AtomicBoolean controlScheduling = new AtomicBoolean(false);
    private final Lock testLock = new ReentrantLock();
    private final Condition schedulingCompleteCondition = testLock.newCondition();
    private boolean schedulingComplete = false;
    private final Condition triggerSchedulingCondition = testLock.newCondition();
    private boolean schedulingTriggered = false;
    private final AtomicInteger numSchedulerRuns = new AtomicInteger(0);


    public LlapTaskSchedulerServiceForTest(
        TaskSchedulerContext appClient, Clock clock) {
      super(appClient, clock, false);
    }

    @Override
    protected void schedulePendingTasks() {
      LOG.info("Attempted schedulPendingTasks");
      testLock.lock();
      try {
        if (controlScheduling.get()) {
          while (!schedulingTriggered) {
            try {
              triggerSchedulingCondition.await();
            } catch (InterruptedException e) {
              throw new RuntimeException(e);
            }
          }
        }
        numSchedulerRuns.incrementAndGet();
        super.schedulePendingTasks();
        schedulingTriggered = false;
        schedulingComplete = true;
        schedulingCompleteCondition.signal();
      } finally {
        testLock.unlock();
      }
    }

    // Enable or disable test scheduling control.
    void forTestsetControlScheduling(boolean control) {
      this.controlScheduling.set(control);
    }

    void forTestSignalSchedulingRun() throws InterruptedException {
      testLock.lock();
      try {
        schedulingTriggered = true;
        triggerSchedulingCondition.signal();
      } finally {
        testLock.unlock();
      }
    }

    boolean forTestAwaitSchedulingRun(long timeout) throws InterruptedException {
      testLock.lock();
      try {
        boolean success = true;
        while (!schedulingComplete) {
          if (timeout == -1) {
            schedulingCompleteCondition.await();
          } else {
            success = schedulingCompleteCondition.await(timeout, TimeUnit.MILLISECONDS);
            break;
          }
        }
        schedulingComplete = false;
        return success;
      } finally {
        testLock.unlock();
      }
    }

  }

  private static class LlapTaskSchedulerServiceForTestControlled extends LlapTaskSchedulerServiceForTest {

    private DelayedTaskSchedulerCallableControlled controlledTSCallable;

    public LlapTaskSchedulerServiceForTestControlled(
        TaskSchedulerContext appClient, Clock clock) {
      super(appClient, clock);
    }

    @Override
    LlapTaskSchedulerService.DelayedTaskSchedulerCallable createDelayedTaskSchedulerCallable() {
      controlledTSCallable = new DelayedTaskSchedulerCallableControlled();
      return controlledTSCallable;
    }

    class DelayedTaskSchedulerCallableControlled extends DelayedTaskSchedulerCallable {
      private final ReentrantLock lock = new ReentrantLock();
      private final Condition triggerRunCondition = lock.newCondition();
      private boolean shouldRun = false;
      private final Condition runCompleteCondition = lock.newCondition();
      private boolean runComplete = false;

      static final int STATE_NOT_RUN = 0;
      static final int STATE_NULL_FOUND = 1;
      static final int STATE_TIMEOUT_NOT_EXPIRED = 2;
      static final int STATE_RETURNED_TASK = 3;

      volatile int lastState = STATE_NOT_RUN;

      volatile boolean lastShouldScheduleTaskResult = false;
      volatile boolean shouldScheduleTaskTriggered = false;

      @Override
      public void processEvictedTask(TaskInfo taskInfo) {
        super.processEvictedTask(taskInfo);
        signalRunComplete();
      }

      @Override
      public TaskInfo getNextTask() throws InterruptedException {

        while (true) {
          lock.lock();
          try {
            while (!shouldRun) {
              triggerRunCondition.await();
            }
            // Preven subsequent runs until a new trigger is set.
            shouldRun = false;
          } finally {
            lock.unlock();
          }
          TaskInfo taskInfo = delayedTaskQueue.peek();
          if (taskInfo == null) {
            LOG.info("Triggered getTask but the queue is empty");
            lastState = STATE_NULL_FOUND;
            signalRunComplete();
            continue;
          }
          if (taskInfo.shouldDelayForLocality(
              LlapTaskSchedulerServiceForTestControlled.this.clock.getTime())) {
            LOG.info("Triggered getTask but the first element is not ready to execute");
            lastState = STATE_TIMEOUT_NOT_EXPIRED;
            signalRunComplete();
            continue;
          } else {
            delayedTaskQueue.poll(); // Remove the previously peeked element.
            lastState = STATE_RETURNED_TASK;
            return taskInfo;
          }
        }
      }

      @Override
      public boolean shouldScheduleTask(TaskInfo taskInfo) {
        shouldScheduleTaskTriggered = true;
        lastShouldScheduleTaskResult = super.shouldScheduleTask(taskInfo);
        return lastShouldScheduleTaskResult;
      }

      void resetShouldScheduleInformation() {
        shouldScheduleTaskTriggered = false;
        lastShouldScheduleTaskResult = false;
      }

      private void signalRunComplete() {
        lock.lock();
        try {
          runComplete = true;
          runCompleteCondition.signal();
        } finally {
          lock.unlock();
        }
      }

      void triggerGetNextTask() {
        lock.lock();
        try {
          shouldRun = true;
          triggerRunCondition.signal();
        } finally {
          lock.unlock();
        }
      }

      void awaitGetNextTaskProcessing() throws InterruptedException {
        lock.lock();
        try {
          while (!runComplete) {
            runCompleteCondition.await();
          }
          runComplete = false;
        } finally {
          lock.unlock();
        }
      }
    }
  }
}
