/*
 * 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.druid.indexing.overlord.hrtr;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Optional;
import com.google.common.base.Predicates;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import org.apache.curator.framework.CuratorFramework;
import org.apache.druid.common.guava.DSuppliers;
import org.apache.druid.discovery.DiscoveryDruidNode;
import org.apache.druid.discovery.DruidNodeDiscovery;
import org.apache.druid.discovery.DruidNodeDiscoveryProvider;
import org.apache.druid.discovery.NodeRole;
import org.apache.druid.discovery.WorkerNodeService;
import org.apache.druid.indexer.TaskLocation;
import org.apache.druid.indexer.TaskState;
import org.apache.druid.indexer.TaskStatus;
import org.apache.druid.indexing.common.task.NoopTask;
import org.apache.druid.indexing.common.task.Task;
import org.apache.druid.indexing.overlord.TaskRunnerListener;
import org.apache.druid.indexing.overlord.TaskRunnerWorkItem;
import org.apache.druid.indexing.overlord.TaskStorage;
import org.apache.druid.indexing.overlord.autoscaling.NoopProvisioningStrategy;
import org.apache.druid.indexing.overlord.config.HttpRemoteTaskRunnerConfig;
import org.apache.druid.indexing.overlord.setup.DefaultWorkerBehaviorConfig;
import org.apache.druid.indexing.worker.TaskAnnouncement;
import org.apache.druid.indexing.worker.Worker;
import org.apache.druid.indexing.worker.config.WorkerConfig;
import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.concurrent.Execs;
import org.apache.druid.java.util.emitter.EmittingLogger;
import org.apache.druid.java.util.http.client.HttpClient;
import org.apache.druid.segment.TestHelper;
import org.apache.druid.server.DruidNode;
import org.apache.druid.server.initialization.IndexerZkConfig;
import org.apache.druid.server.initialization.ZkPathsConfig;
import org.apache.druid.server.metrics.NoopServiceEmitter;
import org.easymock.EasyMock;
import org.eclipse.jetty.util.ConcurrentHashSet;
import org.joda.time.Period;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

/**
 *
 */
public class HttpRemoteTaskRunnerTest
{
  @Before
  public void setup()
  {
    EmittingLogger.registerEmitter(new NoopServiceEmitter());
  }

  /*
  Simulates startup of Overlord and Workers being discovered with no previously known tasks. Fresh tasks are given
  and expected to be completed.
   */
  @Test(timeout = 60_000L)
  public void testFreshStart() throws Exception
  {
    TestDruidNodeDiscovery druidNodeDiscovery = new TestDruidNodeDiscovery();
    DruidNodeDiscoveryProvider druidNodeDiscoveryProvider = EasyMock.createMock(DruidNodeDiscoveryProvider.class);
    EasyMock.expect(druidNodeDiscoveryProvider.getForService(WorkerNodeService.DISCOVERY_SERVICE_KEY))
            .andReturn(druidNodeDiscovery);
    EasyMock.replay(druidNodeDiscoveryProvider);

    HttpRemoteTaskRunner taskRunner = new HttpRemoteTaskRunner(
        TestHelper.makeJsonMapper(),
        new HttpRemoteTaskRunnerConfig()
        {
          @Override
          public int getPendingTasksRunnerNumThreads()
          {
            return 3;
          }
        },
        EasyMock.createNiceMock(HttpClient.class),
        DSuppliers.of(new AtomicReference<>(DefaultWorkerBehaviorConfig.defaultConfig())),
        new NoopProvisioningStrategy<>(),
        druidNodeDiscoveryProvider,
        EasyMock.createNiceMock(TaskStorage.class),
        EasyMock.createNiceMock(CuratorFramework.class),
        new IndexerZkConfig(new ZkPathsConfig(), null, null, null, null)
    )
    {
      @Override
      protected WorkerHolder createWorkerHolder(
          ObjectMapper smileMapper,
          HttpClient httpClient,
          HttpRemoteTaskRunnerConfig config,
          ScheduledExecutorService workersSyncExec,
          WorkerHolder.Listener listener,
          Worker worker,
          List<TaskAnnouncement> knownAnnouncements
      )
      {
        return HttpRemoteTaskRunnerTest.createWorkerHolder(
            smileMapper,
            httpClient,
            config,
            workersSyncExec,
            listener,
            worker,
            ImmutableList.of(),
            ImmutableList.of(),
            ImmutableMap.of(),
            new AtomicInteger(),
            ImmutableSet.of()
        );
      }
    };

    taskRunner.start();

    DiscoveryDruidNode druidNode1 = new DiscoveryDruidNode(
        new DruidNode("service", "host1", false, 8080, null, true, false),
        NodeRole.MIDDLE_MANAGER,
        ImmutableMap.of(
            WorkerNodeService.DISCOVERY_SERVICE_KEY, new WorkerNodeService("ip1", 2, "0", WorkerConfig.DEFAULT_CATEGORY)
        )
    );

    DiscoveryDruidNode druidNode2 = new DiscoveryDruidNode(
        new DruidNode("service", "host2", false, 8080, null, true, false),
        NodeRole.MIDDLE_MANAGER,
        ImmutableMap.of(
            WorkerNodeService.DISCOVERY_SERVICE_KEY, new WorkerNodeService("ip2", 2, "0", WorkerConfig.DEFAULT_CATEGORY)
        )
    );

    druidNodeDiscovery.listener.nodesAdded(ImmutableList.of(druidNode1, druidNode2));

    int numTasks = 8;
    List<Future<TaskStatus>> futures = new ArrayList<>();
    for (int i = 0; i < numTasks; i++) {
      futures.add(taskRunner.run(NoopTask.create("task-id-" + i, 0)));
    }

    for (Future<TaskStatus> future : futures) {
      Assert.assertTrue(future.get().isSuccess());
    }

    Assert.assertEquals(numTasks, taskRunner.getKnownTasks().size());
    Assert.assertEquals(numTasks, taskRunner.getCompletedTasks().size());
  }

  /*
  Simulates one task not getting acknowledged to be running after assigning it to a worker. But, other tasks are
  successfully assigned to other worker and get completed.
   */
  @Test(timeout = 60_000L)
  public void testOneStuckTaskAssignmentDoesntBlockOthers() throws Exception
  {
    TestDruidNodeDiscovery druidNodeDiscovery = new TestDruidNodeDiscovery();
    DruidNodeDiscoveryProvider druidNodeDiscoveryProvider = EasyMock.createMock(DruidNodeDiscoveryProvider.class);
    EasyMock.expect(druidNodeDiscoveryProvider.getForService(WorkerNodeService.DISCOVERY_SERVICE_KEY))
            .andReturn(druidNodeDiscovery);
    EasyMock.replay(druidNodeDiscoveryProvider);

    Task task1 = NoopTask.create("task-id-1", 0);
    Task task2 = NoopTask.create("task-id-2", 0);
    Task task3 = NoopTask.create("task-id-3", 0);

    HttpRemoteTaskRunner taskRunner = new HttpRemoteTaskRunner(
        TestHelper.makeJsonMapper(),
        new HttpRemoteTaskRunnerConfig()
        {
          @Override
          public int getPendingTasksRunnerNumThreads()
          {
            return 3;
          }
        },
        EasyMock.createNiceMock(HttpClient.class),
        DSuppliers.of(new AtomicReference<>(DefaultWorkerBehaviorConfig.defaultConfig())),
        new NoopProvisioningStrategy<>(),
        druidNodeDiscoveryProvider,
        EasyMock.createNiceMock(TaskStorage.class),
        EasyMock.createNiceMock(CuratorFramework.class),
        new IndexerZkConfig(new ZkPathsConfig(), null, null, null, null)
    )
    {
      @Override
      protected WorkerHolder createWorkerHolder(
          ObjectMapper smileMapper,
          HttpClient httpClient,
          HttpRemoteTaskRunnerConfig config,
          ScheduledExecutorService workersSyncExec,
          WorkerHolder.Listener listener,
          Worker worker,
          List<TaskAnnouncement> knownAnnouncements
      )
      {
        return HttpRemoteTaskRunnerTest.createWorkerHolder(
            smileMapper,
            httpClient,
            config,
            workersSyncExec,
            listener,
            worker,
            ImmutableList.of(),
            ImmutableList.of(),
            ImmutableMap.of(task1, ImmutableList.of()), //no announcements would be received for task1
            new AtomicInteger(),
            ImmutableSet.of()
        );
      }
    };

    taskRunner.start();

    DiscoveryDruidNode druidNode1 = new DiscoveryDruidNode(
        new DruidNode("service", "host1", false, 8080, null, true, false),
        NodeRole.MIDDLE_MANAGER,
        ImmutableMap.of(
            WorkerNodeService.DISCOVERY_SERVICE_KEY, new WorkerNodeService("ip1", 2, "0", WorkerConfig.DEFAULT_CATEGORY)
        )
    );

    DiscoveryDruidNode druidNode2 = new DiscoveryDruidNode(
        new DruidNode("service", "host2", false, 8080, null, true, false),
        NodeRole.MIDDLE_MANAGER,
        ImmutableMap.of(
            WorkerNodeService.DISCOVERY_SERVICE_KEY, new WorkerNodeService("ip2", 2, "0", WorkerConfig.DEFAULT_CATEGORY)
        )
    );

    druidNodeDiscovery.listener.nodesAdded(ImmutableList.of(druidNode1, druidNode2));

    taskRunner.run(task1);
    Future<TaskStatus> future2 = taskRunner.run(task2);
    Future<TaskStatus> future3 = taskRunner.run(task3);

    Assert.assertTrue(future2.get().isSuccess());
    Assert.assertTrue(future3.get().isSuccess());

    Assert.assertEquals(task1.getId(), Iterables.getOnlyElement(taskRunner.getPendingTasks()).getTaskId());
  }

  /*
  Simulates restart of the Overlord where taskRunner, on start, discovers workers with prexisting tasks.
   */
  @Test(timeout = 60_000L)
  public void testTaskRunnerRestart() throws Exception
  {
    TestDruidNodeDiscovery druidNodeDiscovery = new TestDruidNodeDiscovery();
    DruidNodeDiscoveryProvider druidNodeDiscoveryProvider = EasyMock.createMock(DruidNodeDiscoveryProvider.class);
    EasyMock.expect(druidNodeDiscoveryProvider.getForService(WorkerNodeService.DISCOVERY_SERVICE_KEY))
            .andReturn(druidNodeDiscovery);
    EasyMock.replay(druidNodeDiscoveryProvider);

    ConcurrentMap<String, CustomFunction> workerHolders = new ConcurrentHashMap<>();

    Task task1 = NoopTask.create("task-id-1", 0);
    Task task2 = NoopTask.create("task-id-2", 0);
    Task task3 = NoopTask.create("task-id-3", 0);
    Task task4 = NoopTask.create("task-id-4", 0);
    Task task5 = NoopTask.create("task-id-5", 0);

    TaskStorage taskStorageMock = EasyMock.createStrictMock(TaskStorage.class);
    EasyMock.expect(taskStorageMock.getStatus(task1.getId())).andReturn(Optional.absent());
    EasyMock.expect(taskStorageMock.getStatus(task2.getId())).andReturn(Optional.absent()).times(2);
    EasyMock.expect(taskStorageMock.getStatus(task3.getId())).andReturn(Optional.of(TaskStatus.running(task3.getId())));
    EasyMock.expect(taskStorageMock.getStatus(task4.getId())).andReturn(Optional.of(TaskStatus.running(task4.getId())));
    EasyMock.expect(taskStorageMock.getStatus(task5.getId())).andReturn(Optional.of(TaskStatus.success(task5.getId())));
    EasyMock.replay(taskStorageMock);

    HttpRemoteTaskRunner taskRunner = new HttpRemoteTaskRunner(
        TestHelper.makeJsonMapper(),
        new HttpRemoteTaskRunnerConfig()
        {
          @Override
          public int getPendingTasksRunnerNumThreads()
          {
            return 3;
          }
        },
        EasyMock.createNiceMock(HttpClient.class),
        DSuppliers.of(new AtomicReference<>(DefaultWorkerBehaviorConfig.defaultConfig())),
        new NoopProvisioningStrategy<>(),
        druidNodeDiscoveryProvider,
        taskStorageMock,
        EasyMock.createNiceMock(CuratorFramework.class),
        new IndexerZkConfig(new ZkPathsConfig(), null, null, null, null)
    )
    {
      @Override
      protected WorkerHolder createWorkerHolder(
          ObjectMapper smileMapper,
          HttpClient httpClient,
          HttpRemoteTaskRunnerConfig config,
          ScheduledExecutorService workersSyncExec,
          WorkerHolder.Listener listener,
          Worker worker,
          List<TaskAnnouncement> knownAnnouncements
      )
      {
        if (workerHolders.containsKey(worker.getHost())) {
          return workerHolders.get(worker.getHost()).apply(
              smileMapper,
              httpClient,
              config,
              workersSyncExec,
              listener,
              worker,
              knownAnnouncements
          );
        } else {
          throw new ISE("No WorkerHolder for [%s].", worker.getHost());
        }
      }
    };

    taskRunner.start();

    DiscoveryDruidNode druidNode = new DiscoveryDruidNode(
        new DruidNode("service", "host", false, 1234, null, true, false),
        NodeRole.MIDDLE_MANAGER,
        ImmutableMap.of(
            WorkerNodeService.DISCOVERY_SERVICE_KEY, new WorkerNodeService("ip1", 2, "0", WorkerConfig.DEFAULT_CATEGORY)
        )
    );

    AtomicInteger ticks = new AtomicInteger();
    Set<String> taskShutdowns = new HashSet<>();

    workerHolders.put(
        "host:1234",
        (mapper, httpClient, config, exec, listener, worker, knownAnnouncements) -> createWorkerHolder(
            mapper,
            httpClient,
            config,
            exec,
            listener,
            worker,
            knownAnnouncements,
            ImmutableList.of(
                TaskAnnouncement.create(
                    task1,
                    TaskStatus.success(task1.getId()),
                    TaskLocation.create("host", 1234, 1235)
                ),
                TaskAnnouncement.create(
                    task2,
                    TaskStatus.running(task2.getId()),
                    TaskLocation.create("host", 1234, 1235)
                ),
                TaskAnnouncement.create(
                    task2,
                    TaskStatus.success(task2.getId()),
                    TaskLocation.create("host", 1234, 1235)
                ),
                TaskAnnouncement.create(
                    task3,
                    TaskStatus.success(task3.getId()),
                    TaskLocation.create("host", 1234, 1235)
                ),
                TaskAnnouncement.create(
                    task4,
                    TaskStatus.running(task4.getId()),
                    TaskLocation.create("host", 1234, 1235)
                ),
                TaskAnnouncement.create(
                    task5,
                    TaskStatus.running(task5.getId()),
                    TaskLocation.create("host", 1234, 1235)
                )
            ),
            ImmutableMap.of(),
            ticks,
            taskShutdowns
        )
    );

    druidNodeDiscovery.listener.nodesAdded(
        ImmutableList.of(
            druidNode
        )
    );

    while (ticks.get() < 1) {
      Thread.sleep(100);
    }

    EasyMock.verify(taskStorageMock);

    Assert.assertEquals(ImmutableSet.of(task2.getId(), task5.getId()), taskShutdowns);
    Assert.assertTrue(taskRunner.getPendingTasks().isEmpty());

    TaskRunnerWorkItem item = Iterables.getOnlyElement(taskRunner.getRunningTasks());
    Assert.assertEquals(task4.getId(), item.getTaskId());

    Assert.assertTrue(taskRunner.run(task3).get().isSuccess());

    Assert.assertEquals(2, taskRunner.getKnownTasks().size());

  }

  @Test(timeout = 60_000L)
  public void testWorkerDisapperAndReappearBeforeItsCleanup() throws Exception
  {
    TestDruidNodeDiscovery druidNodeDiscovery = new TestDruidNodeDiscovery();
    DruidNodeDiscoveryProvider druidNodeDiscoveryProvider = EasyMock.createMock(DruidNodeDiscoveryProvider.class);
    EasyMock.expect(druidNodeDiscoveryProvider.getForService(WorkerNodeService.DISCOVERY_SERVICE_KEY))
            .andReturn(druidNodeDiscovery);
    EasyMock.replay(druidNodeDiscoveryProvider);

    ConcurrentMap<String, CustomFunction> workerHolders = new ConcurrentHashMap<>();

    HttpRemoteTaskRunner taskRunner = new HttpRemoteTaskRunner(
        TestHelper.makeJsonMapper(),
        new HttpRemoteTaskRunnerConfig()
        {
          @Override
          public int getPendingTasksRunnerNumThreads()
          {
            return 3;
          }
        },
        EasyMock.createNiceMock(HttpClient.class),
        DSuppliers.of(new AtomicReference<>(DefaultWorkerBehaviorConfig.defaultConfig())),
        new NoopProvisioningStrategy<>(),
        druidNodeDiscoveryProvider,
        EasyMock.createNiceMock(TaskStorage.class),
        EasyMock.createNiceMock(CuratorFramework.class),
        new IndexerZkConfig(new ZkPathsConfig(), null, null, null, null)
    )
    {
      @Override
      protected WorkerHolder createWorkerHolder(
          ObjectMapper smileMapper,
          HttpClient httpClient,
          HttpRemoteTaskRunnerConfig config,
          ScheduledExecutorService workersSyncExec,
          WorkerHolder.Listener listener,
          Worker worker,
          List<TaskAnnouncement> knownAnnouncements
      )
      {
        if (workerHolders.containsKey(worker.getHost())) {
          return workerHolders.get(worker.getHost()).apply(
              smileMapper,
              httpClient,
              config,
              workersSyncExec,
              listener,
              worker,
              knownAnnouncements
          );
        } else {
          throw new ISE("No WorkerHolder for [%s].", worker.getHost());
        }
      }
    };

    taskRunner.start();

    Task task1 = NoopTask.create("task-id-1", 0);
    Task task2 = NoopTask.create("task-id-2", 0);

    DiscoveryDruidNode druidNode = new DiscoveryDruidNode(
        new DruidNode("service", "host", false, 1234, null, true, false),
        NodeRole.MIDDLE_MANAGER,
        ImmutableMap.of(
            WorkerNodeService.DISCOVERY_SERVICE_KEY, new WorkerNodeService("ip1", 2, "0", WorkerConfig.DEFAULT_CATEGORY)
        )
    );

    workerHolders.put(
        "host:1234",
        (mapper, httpClient, config, exec, listener, worker, knownAnnouncements) -> createWorkerHolder(
            mapper,
            httpClient,
            config,
            exec,
            listener,
            worker,
            knownAnnouncements,
            ImmutableList.of(),
            ImmutableMap.of(
                task1, ImmutableList.of(
                    TaskAnnouncement.create(
                        task1,
                        TaskStatus.running(task1.getId()),
                        TaskLocation.unknown()
                    ),
                    TaskAnnouncement.create(
                        task1,
                        TaskStatus.running(task1.getId()),
                        TaskLocation.create("host", 1234, 1235)
                    ),
                    TaskAnnouncement.create(
                        task1,
                        TaskStatus.success(task1.getId()),
                        TaskLocation.create("host", 1234, 1235)
                    )
                ),
                task2, ImmutableList.of(
                    TaskAnnouncement.create(
                        task2,
                        TaskStatus.running(task2.getId()),
                        TaskLocation.unknown()
                    ),
                    TaskAnnouncement.create(
                        task2,
                        TaskStatus.running(task2.getId()),
                        TaskLocation.create("host", 1234, 1235)
                    )
                )
            ),
            new AtomicInteger(),
            ImmutableSet.of()
        )
    );

    druidNodeDiscovery.listener.nodesAdded(
        ImmutableList.of(
            druidNode
        )
    );

    Future<TaskStatus> future1 = taskRunner.run(task1);
    Future<TaskStatus> future2 = taskRunner.run(task2);

    while (taskRunner.getPendingTasks().size() > 0) {
      Thread.sleep(100);
    }

    druidNodeDiscovery.listener.nodesRemoved(
        ImmutableList.of(
            druidNode
        )
    );

    workerHolders.put(
        "host:1234",
        (mapper, httpClient, config, exec, listener, worker, knownAnnouncements) -> createWorkerHolder(
            mapper,
            httpClient,
            config,
            exec,
            listener,
            worker,
            knownAnnouncements,
            ImmutableList.of(
                TaskAnnouncement.create(
                    task2,
                    TaskStatus.running(task2.getId()),
                    TaskLocation.create("host", 1234, 1235)
                ),
                TaskAnnouncement.create(
                    task2,
                    TaskStatus.success(task2.getId()),
                    TaskLocation.create("host", 1234, 1235)
                )

            ),
            ImmutableMap.of(),
            new AtomicInteger(),
            ImmutableSet.of()
        )
    );

    druidNodeDiscovery.listener.nodesAdded(
        ImmutableList.of(
            druidNode
        )
    );

    Assert.assertTrue(future1.get().isSuccess());
    Assert.assertTrue(future2.get().isSuccess());
  }

  @Test(timeout = 60_000L)
  public void testWorkerDisapperAndReappearAfterItsCleanup() throws Exception
  {
    TestDruidNodeDiscovery druidNodeDiscovery = new TestDruidNodeDiscovery();
    DruidNodeDiscoveryProvider druidNodeDiscoveryProvider = EasyMock.createMock(DruidNodeDiscoveryProvider.class);
    EasyMock.expect(druidNodeDiscoveryProvider.getForService(WorkerNodeService.DISCOVERY_SERVICE_KEY))
            .andReturn(druidNodeDiscovery);
    EasyMock.replay(druidNodeDiscoveryProvider);

    ConcurrentMap<String, CustomFunction> workerHolders = new ConcurrentHashMap<>();

    HttpRemoteTaskRunner taskRunner = new HttpRemoteTaskRunner(
        TestHelper.makeJsonMapper(),
        new HttpRemoteTaskRunnerConfig()
        {
          @Override
          public Period getTaskCleanupTimeout()
          {
            return Period.millis(1);
          }
        },
        EasyMock.createNiceMock(HttpClient.class),
        DSuppliers.of(new AtomicReference<>(DefaultWorkerBehaviorConfig.defaultConfig())),
        new NoopProvisioningStrategy<>(),
        druidNodeDiscoveryProvider,
        EasyMock.createNiceMock(TaskStorage.class),
        EasyMock.createNiceMock(CuratorFramework.class),
        new IndexerZkConfig(new ZkPathsConfig(), null, null, null, null)
    )
    {
      @Override
      protected WorkerHolder createWorkerHolder(
          ObjectMapper smileMapper,
          HttpClient httpClient,
          HttpRemoteTaskRunnerConfig config,
          ScheduledExecutorService workersSyncExec,
          WorkerHolder.Listener listener,
          Worker worker,
          List<TaskAnnouncement> knownAnnouncements
      )
      {
        if (workerHolders.containsKey(worker.getHost())) {
          return workerHolders.get(worker.getHost()).apply(
              smileMapper,
              httpClient,
              config,
              workersSyncExec,
              listener,
              worker,
              knownAnnouncements
          );
        } else {
          throw new ISE("No WorkerHolder for [%s].", worker.getHost());
        }
      }
    };

    taskRunner.start();

    Task task1 = NoopTask.create("task-id-1", 0);
    Task task2 = NoopTask.create("task-id-2", 0);

    DiscoveryDruidNode druidNode = new DiscoveryDruidNode(
        new DruidNode("service", "host", false, 1234, null, true, false),
        NodeRole.MIDDLE_MANAGER,
        ImmutableMap.of(
            WorkerNodeService.DISCOVERY_SERVICE_KEY,
            new WorkerNodeService("ip1", 2, "0", WorkerConfig.DEFAULT_CATEGORY)
        )
    );

    workerHolders.put(
        "host:1234",
        (mapper, httpClient, config, exec, listener, worker, knownAnnouncements) -> createWorkerHolder(
            mapper,
            httpClient,
            config,
            exec,
            listener,
            worker,
            knownAnnouncements,
            ImmutableList.of(),
            ImmutableMap.of(
                task1, ImmutableList.of(
                    TaskAnnouncement.create(
                        task1,
                        TaskStatus.running(task1.getId()),
                        TaskLocation.unknown()
                    ),
                    TaskAnnouncement.create(
                        task1,
                        TaskStatus.running(task1.getId()),
                        TaskLocation.create("host", 1234, 1235)
                    )
                ),
                task2, ImmutableList.of(
                    TaskAnnouncement.create(
                        task2,
                        TaskStatus.running(task2.getId()),
                        TaskLocation.unknown()
                    ),
                    TaskAnnouncement.create(
                        task2,
                        TaskStatus.running(task2.getId()),
                        TaskLocation.create("host", 1234, 1235)
                    )
                )
            ),
            new AtomicInteger(),
            ImmutableSet.of()
        )
    );

    druidNodeDiscovery.listener.nodesAdded(
        ImmutableList.of(
            druidNode
        )
    );

    Future<TaskStatus> future1 = taskRunner.run(task1);
    Future<TaskStatus> future2 = taskRunner.run(task2);

    while (taskRunner.getPendingTasks().size() > 0) {
      Thread.sleep(100);
    }

    druidNodeDiscovery.listener.nodesRemoved(
        ImmutableList.of(
            druidNode
        )
    );

    Assert.assertTrue(future1.get().isFailure());
    Assert.assertTrue(future2.get().isFailure());
    Assert.assertNotNull(future1.get().getErrorMsg());
    Assert.assertNotNull(future2.get().getErrorMsg());
    Assert.assertTrue(
        future1.get().getErrorMsg().startsWith(
            "The worker that this task was assigned disappeared and did not report cleanup within timeout"
        )
    );
    Assert.assertTrue(
        future2.get().getErrorMsg().startsWith(
            "The worker that this task was assigned disappeared and did not report cleanup within timeout"
        )
    );

    AtomicInteger ticks = new AtomicInteger();
    Set<String> actualShutdowns = new ConcurrentHashSet<>();

    workerHolders.put(
        "host:1234",
        (mapper, httpClient, config, exec, listener, worker, knownAnnouncements) -> createWorkerHolder(
            mapper,
            httpClient,
            config,
            exec,
            listener,
            worker,
            knownAnnouncements,
            ImmutableList.of(
                TaskAnnouncement.create(
                    task1,
                    TaskStatus.success(task1.getId()),
                    TaskLocation.create("host", 1234, 1235)
                ),
                TaskAnnouncement.create(
                    task2,
                    TaskStatus.running(task2.getId()),
                    TaskLocation.create("host", 1234, 1235)
                )
            ),
            ImmutableMap.of(),
            ticks,
            actualShutdowns
        )
    );

    druidNodeDiscovery.listener.nodesAdded(
        ImmutableList.of(
            druidNode
        )
    );

    while (ticks.get() < 1) {
      Thread.sleep(100);
    }

    Assert.assertEquals(ImmutableSet.of(task2.getId()), actualShutdowns);
    Assert.assertTrue(taskRunner.run(task1).get().isFailure());
    Assert.assertTrue(taskRunner.run(task2).get().isFailure());
  }

  @Test(timeout = 60_000L)
  public void testMarkWorkersLazy() throws Exception
  {
    TestDruidNodeDiscovery druidNodeDiscovery = new TestDruidNodeDiscovery();
    DruidNodeDiscoveryProvider druidNodeDiscoveryProvider = EasyMock.createMock(DruidNodeDiscoveryProvider.class);
    EasyMock.expect(druidNodeDiscoveryProvider.getForService(WorkerNodeService.DISCOVERY_SERVICE_KEY))
            .andReturn(druidNodeDiscovery);
    EasyMock.replay(druidNodeDiscoveryProvider);

    Task task1 = NoopTask.create("task-id-1", 0);
    Task task2 = NoopTask.create("task-id-2", 0);

    ConcurrentMap<String, CustomFunction> workerHolders = new ConcurrentHashMap<>();

    HttpRemoteTaskRunner taskRunner = new HttpRemoteTaskRunner(
        TestHelper.makeJsonMapper(),
        new HttpRemoteTaskRunnerConfig()
        {
          @Override
          public int getPendingTasksRunnerNumThreads()
          {
            return 3;
          }
        },
        EasyMock.createNiceMock(HttpClient.class),
        DSuppliers.of(new AtomicReference<>(DefaultWorkerBehaviorConfig.defaultConfig())),
        new NoopProvisioningStrategy<>(),
        druidNodeDiscoveryProvider,
        EasyMock.createNiceMock(TaskStorage.class),
        EasyMock.createNiceMock(CuratorFramework.class),
        new IndexerZkConfig(new ZkPathsConfig(), null, null, null, null)
    )
    {
      @Override
      protected WorkerHolder createWorkerHolder(
          ObjectMapper smileMapper,
          HttpClient httpClient,
          HttpRemoteTaskRunnerConfig config,
          ScheduledExecutorService workersSyncExec,
          WorkerHolder.Listener listener,
          Worker worker,
          List<TaskAnnouncement> knownAnnouncements
      )
      {
        if (workerHolders.containsKey(worker.getHost())) {
          return workerHolders.get(worker.getHost()).apply(
              smileMapper,
              httpClient,
              config,
              workersSyncExec,
              listener,
              worker,
              knownAnnouncements
          );
        } else {
          throw new ISE("No WorkerHolder for [%s].", worker.getHost());
        }
      }
    };

    taskRunner.start();

    Assert.assertEquals(0, taskRunner.getTotalTaskSlotCount());
    Assert.assertEquals(0, taskRunner.getIdleTaskSlotCount());
    Assert.assertEquals(0, taskRunner.getUsedTaskSlotCount());

    AtomicInteger ticks = new AtomicInteger();

    DiscoveryDruidNode druidNode1 = new DiscoveryDruidNode(
        new DruidNode("service", "host1", false, 8080, null, true, false),
        NodeRole.MIDDLE_MANAGER,
        ImmutableMap.of(
            WorkerNodeService.DISCOVERY_SERVICE_KEY, new WorkerNodeService("ip1", 1, "0", WorkerConfig.DEFAULT_CATEGORY)
        )
    );

    workerHolders.put(
        "host1:8080",
        (mapper, httpClient, config, exec, listener, worker, knownAnnouncements) -> createWorkerHolder(
            mapper,
            httpClient,
            config,
            exec,
            listener,
            worker,
            knownAnnouncements,
            ImmutableList.of(),
            ImmutableMap.of(
                task1, ImmutableList.of(
                    TaskAnnouncement.create(
                        task1,
                        TaskStatus.running(task1.getId()),
                        TaskLocation.unknown()
                    ),
                    TaskAnnouncement.create(
                        task1,
                        TaskStatus.running(task1.getId()),
                        TaskLocation.create("host1", 8080, -1)
                    )
                )
            ),
            ticks,
            ImmutableSet.of()
        )
    );

    druidNodeDiscovery.listener.nodesAdded(ImmutableList.of(druidNode1));

    Assert.assertEquals(1, taskRunner.getTotalTaskSlotCount());
    Assert.assertEquals(1, taskRunner.getIdleTaskSlotCount());
    Assert.assertEquals(0, taskRunner.getUsedTaskSlotCount());

    taskRunner.run(task1);

    while (ticks.get() < 1) {
      Thread.sleep(100);
    }

    Assert.assertEquals(1, taskRunner.getTotalTaskSlotCount());
    Assert.assertEquals(0, taskRunner.getIdleTaskSlotCount());
    Assert.assertEquals(1, taskRunner.getUsedTaskSlotCount());

    DiscoveryDruidNode druidNode2 = new DiscoveryDruidNode(
        new DruidNode("service", "host2", false, 8080, null, true, false),
        NodeRole.MIDDLE_MANAGER,
        ImmutableMap.of(
            WorkerNodeService.DISCOVERY_SERVICE_KEY,
            new WorkerNodeService("ip2", 1, "0", WorkerConfig.DEFAULT_CATEGORY)
        )
    );

    workerHolders.put(
        "host2:8080",
        (mapper, httpClient, config, exec, listener, worker, knownAnnouncements) -> createWorkerHolder(
            mapper,
            httpClient,
            config,
            exec,
            listener,
            worker,
            knownAnnouncements,
            ImmutableList.of(),
            ImmutableMap.of(task2, ImmutableList.of()),
            ticks,
            ImmutableSet.of()
        )
    );

    druidNodeDiscovery.listener.nodesAdded(ImmutableList.of(druidNode2));

    Assert.assertEquals(2, taskRunner.getTotalTaskSlotCount());
    Assert.assertEquals(1, taskRunner.getIdleTaskSlotCount());
    Assert.assertEquals(1, taskRunner.getUsedTaskSlotCount());

    taskRunner.run(task2);

    while (ticks.get() < 2) {
      Thread.sleep(100);
    }

    Assert.assertEquals(2, taskRunner.getTotalTaskSlotCount());
    Assert.assertEquals(0, taskRunner.getIdleTaskSlotCount());
    Assert.assertEquals(1, taskRunner.getUsedTaskSlotCount());

    DiscoveryDruidNode druidNode3 = new DiscoveryDruidNode(
        new DruidNode("service", "host3", false, 8080, null, true, false),
        NodeRole.MIDDLE_MANAGER,
        ImmutableMap.of(
            WorkerNodeService.DISCOVERY_SERVICE_KEY,
            new WorkerNodeService("ip2", 1, "0", WorkerConfig.DEFAULT_CATEGORY)
        )
    );

    workerHolders.put(
        "host3:8080",
        (mapper, httpClient, config, exec, listener, worker, knownAnnouncements) -> createWorkerHolder(
            mapper,
            httpClient,
            config,
            exec,
            listener,
            worker,
            knownAnnouncements,
            ImmutableList.of(),
            ImmutableMap.of(),
            new AtomicInteger(),
            ImmutableSet.of()
        )
    );

    druidNodeDiscovery.listener.nodesAdded(ImmutableList.of(druidNode3));

    Assert.assertEquals(3, taskRunner.getTotalTaskSlotCount());
    Assert.assertEquals(1, taskRunner.getIdleTaskSlotCount());
    Assert.assertEquals(1, taskRunner.getUsedTaskSlotCount());
    Assert.assertEquals(0, taskRunner.getLazyTaskSlotCount());

    Assert.assertEquals(task1.getId(), Iterables.getOnlyElement(taskRunner.getRunningTasks()).getTaskId());
    Assert.assertEquals(task2.getId(), Iterables.getOnlyElement(taskRunner.getPendingTasks()).getTaskId());

    Assert.assertEquals(
        "host3:8080",
        Iterables.getOnlyElement(taskRunner.markWorkersLazy(Predicates.alwaysTrue(), Integer.MAX_VALUE))
                 .getHost()
    );

    Assert.assertEquals(3, taskRunner.getTotalTaskSlotCount());
    Assert.assertEquals(0, taskRunner.getIdleTaskSlotCount());
    Assert.assertEquals(1, taskRunner.getUsedTaskSlotCount());
    Assert.assertEquals(1, taskRunner.getLazyTaskSlotCount());
  }

  /*
   * Task goes PENDING -> RUNNING -> SUCCESS and few more useless notifications in between.
   */
  @Test
  public void testTaskAddedOrUpdated1() throws Exception
  {
    Task task = NoopTask.create("task");
    List<Object> listenerNotificationsAccumulator = new ArrayList<>();
    HttpRemoteTaskRunner taskRunner = createTaskRunnerForTestTaskAddedOrUpdated(
        EasyMock.createStrictMock(TaskStorage.class),
        listenerNotificationsAccumulator
    );

    WorkerHolder workerHolder = EasyMock.createMock(WorkerHolder.class);
    EasyMock.expect(workerHolder.getWorker()).andReturn(new Worker("http", "worker", "127.0.0.1", 1, "v1", WorkerConfig.DEFAULT_CATEGORY)).anyTimes();
    workerHolder.setLastCompletedTaskTime(EasyMock.anyObject());
    workerHolder.resetContinuouslyFailedTasksCount();
    EasyMock.expect(workerHolder.getContinuouslyFailedTasksCount()).andReturn(0);
    EasyMock.replay(workerHolder);

    Future<TaskStatus> future = taskRunner.run(task);
    Assert.assertEquals(task.getId(), Iterables.getOnlyElement(taskRunner.getPendingTasks()).getTaskId());

    // RUNNING notification from worker
    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task,
        TaskStatus.running(task.getId()),
        TaskLocation.create("worker", 1000, 1001)
    ), workerHolder);
    Assert.assertEquals(task.getId(), Iterables.getOnlyElement(taskRunner.getRunningTasks()).getTaskId());

    // Another RUNNING notification from worker, notifying change in location
    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task,
        TaskStatus.running(task.getId()),
        TaskLocation.create("worker", 1, 2)
    ), workerHolder);
    Assert.assertEquals(task.getId(), Iterables.getOnlyElement(taskRunner.getRunningTasks()).getTaskId());

    // Redundant RUNNING notification from worker, ignored
    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task,
        TaskStatus.running(task.getId()),
        TaskLocation.create("worker", 1, 2)
    ), workerHolder);
    Assert.assertEquals(task.getId(), Iterables.getOnlyElement(taskRunner.getRunningTasks()).getTaskId());

    // Another "rogue-worker" reports running it, and gets asked to shutdown the task
    WorkerHolder rogueWorkerHolder = EasyMock.createMock(WorkerHolder.class);
    EasyMock.expect(rogueWorkerHolder.getWorker())
            .andReturn(new Worker("http", "rogue-worker", "127.0.0.1", 5, "v1", WorkerConfig.DEFAULT_CATEGORY))
            .anyTimes();
    rogueWorkerHolder.shutdownTask(task.getId());
    EasyMock.replay(rogueWorkerHolder);
    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task,
        TaskStatus.running(task.getId()),
        TaskLocation.create("rogue-worker", 1, 2)
    ), rogueWorkerHolder);
    Assert.assertEquals(task.getId(), Iterables.getOnlyElement(taskRunner.getRunningTasks()).getTaskId());
    EasyMock.verify(rogueWorkerHolder);

    // "rogue-worker" reports FAILURE for the task, ignored
    rogueWorkerHolder = EasyMock.createMock(WorkerHolder.class);
    EasyMock.expect(rogueWorkerHolder.getWorker())
            .andReturn(new Worker("http", "rogue-worker", "127.0.0.1", 5, "v1", WorkerConfig.DEFAULT_CATEGORY))
            .anyTimes();
    EasyMock.replay(rogueWorkerHolder);
    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task,
        TaskStatus.failure(task.getId()),
        TaskLocation.create("rogue-worker", 1, 2)
    ), rogueWorkerHolder);
    Assert.assertEquals(task.getId(), Iterables.getOnlyElement(taskRunner.getRunningTasks()).getTaskId());
    EasyMock.verify(rogueWorkerHolder);

    // workers sends SUCCESS notification, task is marked SUCCESS now.
    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task,
        TaskStatus.success(task.getId()),
        TaskLocation.create("worker", 1, 2)
    ), workerHolder);
    Assert.assertEquals(task.getId(), Iterables.getOnlyElement(taskRunner.getCompletedTasks()).getTaskId());
    Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode());

    // "rogue-worker" reports running it, and gets asked to shutdown the task
    rogueWorkerHolder = EasyMock.createMock(WorkerHolder.class);
    EasyMock.expect(rogueWorkerHolder.getWorker())
            .andReturn(new Worker("http", "rogue-worker", "127.0.0.1", 5, "v1", WorkerConfig.DEFAULT_CATEGORY))
            .anyTimes();
    rogueWorkerHolder.shutdownTask(task.getId());
    EasyMock.replay(rogueWorkerHolder);
    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task,
        TaskStatus.running(task.getId()),
        TaskLocation.create("rogue-worker", 1, 2)
    ), rogueWorkerHolder);
    Assert.assertEquals(task.getId(), Iterables.getOnlyElement(taskRunner.getCompletedTasks()).getTaskId());
    EasyMock.verify(rogueWorkerHolder);

    // "rogue-worker" reports FAILURE for the tasks, ignored
    rogueWorkerHolder = EasyMock.createMock(WorkerHolder.class);
    EasyMock.expect(rogueWorkerHolder.getWorker())
            .andReturn(new Worker("http", "rogue-worker", "127.0.0.1", 5, "v1", WorkerConfig.DEFAULT_CATEGORY))
            .anyTimes();
    EasyMock.replay(rogueWorkerHolder);
    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task,
        TaskStatus.failure(task.getId()),
        TaskLocation.create("rogue-worker", 1, 2)
    ), rogueWorkerHolder);
    Assert.assertEquals(task.getId(), Iterables.getOnlyElement(taskRunner.getCompletedTasks()).getTaskId());
    EasyMock.verify(rogueWorkerHolder);

    Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode());

    EasyMock.verify(workerHolder);

    Assert.assertEquals(
        listenerNotificationsAccumulator,
        ImmutableList.of(
            ImmutableList.of(task.getId(), TaskLocation.create("worker", 1000, 1001)),
            ImmutableList.of(task.getId(), TaskLocation.create("worker", 1, 2)),
            ImmutableList.of(task.getId(), TaskStatus.success(task.getId()))
        )
    );
  }

  /*
   * Task goes from PENDING -> SUCCESS . Happens when TaskRunner is given task but a worker reported it being already
   * completed with SUCCESS.
   */
  @Test
  public void testTaskAddedOrUpdated2() throws Exception
  {
    Task task = NoopTask.create("task");
    List<Object> listenerNotificationsAccumulator = new ArrayList<>();
    HttpRemoteTaskRunner taskRunner = createTaskRunnerForTestTaskAddedOrUpdated(
        EasyMock.createStrictMock(TaskStorage.class),
        listenerNotificationsAccumulator
    );

    Worker worker = new Worker("http", "localhost", "127.0.0.1", 1, "v1", WorkerConfig.DEFAULT_CATEGORY);

    WorkerHolder workerHolder = EasyMock.createMock(WorkerHolder.class);
    EasyMock.expect(workerHolder.getWorker()).andReturn(worker).anyTimes();
    workerHolder.setLastCompletedTaskTime(EasyMock.anyObject());
    workerHolder.resetContinuouslyFailedTasksCount();
    EasyMock.expect(workerHolder.getContinuouslyFailedTasksCount()).andReturn(0);
    EasyMock.replay(workerHolder);

    Future<TaskStatus> future = taskRunner.run(task);
    Assert.assertEquals(task.getId(), Iterables.getOnlyElement(taskRunner.getPendingTasks()).getTaskId());

    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task,
        TaskStatus.success(task.getId()),
        TaskLocation.create("worker", 1, 2)
    ), workerHolder);
    Assert.assertEquals(task.getId(), Iterables.getOnlyElement(taskRunner.getCompletedTasks()).getTaskId());

    Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode());

    EasyMock.verify(workerHolder);

    Assert.assertEquals(
        listenerNotificationsAccumulator,
        ImmutableList.of(
            ImmutableList.of(task.getId(), TaskLocation.create("worker", 1, 2)),
            ImmutableList.of(task.getId(), TaskStatus.success(task.getId()))
        )
    );
  }

  /*
   * Notifications received for tasks not known to TaskRunner maybe known to TaskStorage.
   * This could happen when TaskRunner starts and workers reports running/completed tasks on them.
   */
  @Test
  public void testTaskAddedOrUpdated3()
  {
    Task task1 = NoopTask.create("task1");
    Task task2 = NoopTask.create("task2");
    Task task3 = NoopTask.create("task3");
    Task task4 = NoopTask.create("task4");
    Task task5 = NoopTask.create("task5");
    Task task6 = NoopTask.create("task6");

    TaskStorage taskStorage = EasyMock.createMock(TaskStorage.class);
    EasyMock.expect(taskStorage.getStatus(task1.getId())).andReturn(Optional.of(TaskStatus.running(task1.getId())));
    EasyMock.expect(taskStorage.getStatus(task2.getId())).andReturn(Optional.of(TaskStatus.running(task2.getId())));
    EasyMock.expect(taskStorage.getStatus(task3.getId())).andReturn(Optional.of(TaskStatus.success(task3.getId())));
    EasyMock.expect(taskStorage.getStatus(task4.getId())).andReturn(Optional.of(TaskStatus.success(task4.getId())));
    EasyMock.expect(taskStorage.getStatus(task5.getId())).andReturn(Optional.absent());
    EasyMock.expect(taskStorage.getStatus(task6.getId())).andReturn(Optional.absent());
    EasyMock.replay(taskStorage);

    List<Object> listenerNotificationsAccumulator = new ArrayList<>();
    HttpRemoteTaskRunner taskRunner =
        createTaskRunnerForTestTaskAddedOrUpdated(taskStorage, listenerNotificationsAccumulator);

    Worker worker = new Worker("http", "localhost", "127.0.0.1", 1, "v1", WorkerConfig.DEFAULT_CATEGORY);

    WorkerHolder workerHolder = EasyMock.createMock(WorkerHolder.class);
    EasyMock.expect(workerHolder.getWorker()).andReturn(worker).anyTimes();
    workerHolder.setLastCompletedTaskTime(EasyMock.anyObject());
    workerHolder.resetContinuouslyFailedTasksCount();
    EasyMock.expect(workerHolder.getContinuouslyFailedTasksCount()).andReturn(0);
    workerHolder.shutdownTask(task3.getId());
    workerHolder.shutdownTask(task5.getId());
    EasyMock.replay(workerHolder);

    Assert.assertEquals(0, taskRunner.getKnownTasks().size());

    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task1,
        TaskStatus.running(task1.getId()),
        TaskLocation.create("worker", 1, 2)
    ), workerHolder);

    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task2,
        TaskStatus.success(task2.getId()),
        TaskLocation.create("worker", 3, 4)
    ), workerHolder);

    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task3,
        TaskStatus.running(task3.getId()),
        TaskLocation.create("worker", 5, 6)
    ), workerHolder);

    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task4,
        TaskStatus.success(task4.getId()),
        TaskLocation.create("worker", 7, 8)
    ), workerHolder);

    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task5,
        TaskStatus.running(task5.getId()),
        TaskLocation.create("worker", 9, 10)
    ), workerHolder);

    taskRunner.taskAddedOrUpdated(TaskAnnouncement.create(
        task6,
        TaskStatus.success(task6.getId()),
        TaskLocation.create("worker", 11, 12)
    ), workerHolder);

    EasyMock.verify(workerHolder, taskStorage);

    Assert.assertEquals(
        listenerNotificationsAccumulator,
        ImmutableList.of(
            ImmutableList.of(task1.getId(), TaskLocation.create("worker", 1, 2)),
            ImmutableList.of(task2.getId(), TaskLocation.create("worker", 3, 4)),
            ImmutableList.of(task2.getId(), TaskStatus.success(task2.getId()))
        )
    );
  }

  @Test
  public void testTimeoutInAssigningTasks() throws Exception
  {
    TestDruidNodeDiscovery druidNodeDiscovery = new TestDruidNodeDiscovery();
    DruidNodeDiscoveryProvider druidNodeDiscoveryProvider = EasyMock.createMock(DruidNodeDiscoveryProvider.class);
    EasyMock.expect(druidNodeDiscoveryProvider.getForService(WorkerNodeService.DISCOVERY_SERVICE_KEY))
            .andReturn(druidNodeDiscovery);
    EasyMock.replay(druidNodeDiscoveryProvider);

    HttpRemoteTaskRunner taskRunner = new HttpRemoteTaskRunner(
        TestHelper.makeJsonMapper(),
        new HttpRemoteTaskRunnerConfig()
        {
          @Override
          public int getPendingTasksRunnerNumThreads()
          {
            return 1;
          }

          @Override
          public Period getTaskAssignmentTimeout()
          {
            return new Period("PT1S");
          }
        },
        EasyMock.createNiceMock(HttpClient.class),
        DSuppliers.of(new AtomicReference<>(DefaultWorkerBehaviorConfig.defaultConfig())),
        new NoopProvisioningStrategy<>(),
        druidNodeDiscoveryProvider,
        EasyMock.createNiceMock(TaskStorage.class),
        EasyMock.createNiceMock(CuratorFramework.class),
        new IndexerZkConfig(new ZkPathsConfig(), null, null, null, null)
    )
    {
      @Override
      protected WorkerHolder createWorkerHolder(
          ObjectMapper smileMapper,
          HttpClient httpClient,
          HttpRemoteTaskRunnerConfig config,
          ScheduledExecutorService workersSyncExec,
          WorkerHolder.Listener listener,
          Worker worker,
          List<TaskAnnouncement> knownAnnouncements
      )
      {
        return new WorkerHolder(
            smileMapper,
            httpClient,
            config,
            workersSyncExec,
            listener,
            worker,
            ImmutableList.of()
        )
        {
          @Override
          public void start()
          {
            disabled.set(false);
          }

          @Override
          public void stop()
          {
          }

          @Override
          public boolean isInitialized()
          {
            return true;
          }

          @Override
          public void waitForInitialization()
          {
          }

          @Override
          public boolean assignTask(Task task)
          {
            // Always returns true
            return true;
          }

          @Override
          public void shutdownTask(String taskId)
          {
          }
        };
      }
    };

    taskRunner.start();

    DiscoveryDruidNode druidNode1 = new DiscoveryDruidNode(
        new DruidNode("service", "host1", false, 8080, null, true, false),
        NodeRole.MIDDLE_MANAGER,
        ImmutableMap.of(
            WorkerNodeService.DISCOVERY_SERVICE_KEY, new WorkerNodeService("ip1", 2, "0", WorkerConfig.DEFAULT_CATEGORY)
        )
    );

    druidNodeDiscovery.listener.nodesAdded(ImmutableList.of(druidNode1));

    Future<TaskStatus> future = taskRunner.run(NoopTask.create("task-id", 0));
    Assert.assertTrue(future.get().isFailure());
    Assert.assertNotNull(future.get().getErrorMsg());
    Assert.assertTrue(
        future.get().getErrorMsg().startsWith("The worker that this task is assigned did not start it in timeout")
    );
  }

  @Test
  public void testExceptionThrownInAssigningTasks() throws Exception
  {
    TestDruidNodeDiscovery druidNodeDiscovery = new TestDruidNodeDiscovery();
    DruidNodeDiscoveryProvider druidNodeDiscoveryProvider = EasyMock.createMock(DruidNodeDiscoveryProvider.class);
    EasyMock.expect(druidNodeDiscoveryProvider.getForService(WorkerNodeService.DISCOVERY_SERVICE_KEY))
            .andReturn(druidNodeDiscovery);
    EasyMock.replay(druidNodeDiscoveryProvider);

    HttpRemoteTaskRunner taskRunner = new HttpRemoteTaskRunner(
        TestHelper.makeJsonMapper(),
        new HttpRemoteTaskRunnerConfig()
        {
          @Override
          public int getPendingTasksRunnerNumThreads()
          {
            return 1;
          }

          @Override
          public Period getTaskAssignmentTimeout()
          {
            return new Period("PT1S");
          }
        },
        EasyMock.createNiceMock(HttpClient.class),
        DSuppliers.of(new AtomicReference<>(DefaultWorkerBehaviorConfig.defaultConfig())),
        new NoopProvisioningStrategy<>(),
        druidNodeDiscoveryProvider,
        EasyMock.createNiceMock(TaskStorage.class),
        EasyMock.createNiceMock(CuratorFramework.class),
        new IndexerZkConfig(new ZkPathsConfig(), null, null, null, null)
    )
    {
      @Override
      protected WorkerHolder createWorkerHolder(
          ObjectMapper smileMapper,
          HttpClient httpClient,
          HttpRemoteTaskRunnerConfig config,
          ScheduledExecutorService workersSyncExec,
          WorkerHolder.Listener listener,
          Worker worker,
          List<TaskAnnouncement> knownAnnouncements
      )
      {
        return new WorkerHolder(
            smileMapper,
            httpClient,
            config,
            workersSyncExec,
            listener,
            worker,
            ImmutableList.of()
        )
        {
          @Override
          public void start()
          {
            disabled.set(false);
          }

          @Override
          public void stop()
          {
          }

          @Override
          public boolean isInitialized()
          {
            return true;
          }

          @Override
          public void waitForInitialization()
          {
          }

          @Override
          public boolean assignTask(Task task)
          {
            throw new RuntimeException("Assign failure test");
          }

          @Override
          public void shutdownTask(String taskId)
          {
          }
        };
      }
    };

    taskRunner.start();

    DiscoveryDruidNode druidNode1 = new DiscoveryDruidNode(
        new DruidNode("service", "host1", false, 8080, null, true, false),
        NodeRole.MIDDLE_MANAGER,
        ImmutableMap.of(
            WorkerNodeService.DISCOVERY_SERVICE_KEY, new WorkerNodeService("ip1", 2, "0", WorkerConfig.DEFAULT_CATEGORY)
        )
    );

    druidNodeDiscovery.listener.nodesAdded(ImmutableList.of(druidNode1));

    Future<TaskStatus> future = taskRunner.run(NoopTask.create("task-id", 0));
    Assert.assertTrue(future.get().isFailure());
    Assert.assertNotNull(future.get().getErrorMsg());
    Assert.assertTrue(
        StringUtils.format("Actual message is: %s", future.get().getErrorMsg()),
        future.get().getErrorMsg().startsWith("Failed to assign this task")
    );
  }

  private HttpRemoteTaskRunner createTaskRunnerForTestTaskAddedOrUpdated(
      TaskStorage taskStorage,
      List<Object> listenerNotificationsAccumulator
  )
  {
    TestDruidNodeDiscovery druidNodeDiscovery = new TestDruidNodeDiscovery();
    DruidNodeDiscoveryProvider druidNodeDiscoveryProvider = EasyMock.createMock(DruidNodeDiscoveryProvider.class);
    EasyMock.expect(druidNodeDiscoveryProvider.getForService(WorkerNodeService.DISCOVERY_SERVICE_KEY))
            .andReturn(druidNodeDiscovery);
    EasyMock.replay(druidNodeDiscoveryProvider);

    HttpRemoteTaskRunner taskRunner = new HttpRemoteTaskRunner(
        TestHelper.makeJsonMapper(),
        new HttpRemoteTaskRunnerConfig()
        {
          @Override
          public int getPendingTasksRunnerNumThreads()
          {
            return 3;
          }
        },
        EasyMock.createNiceMock(HttpClient.class),
        DSuppliers.of(new AtomicReference<>(DefaultWorkerBehaviorConfig.defaultConfig())),
        new NoopProvisioningStrategy<>(),
        druidNodeDiscoveryProvider,
        taskStorage,
        EasyMock.createNiceMock(CuratorFramework.class),
        new IndexerZkConfig(new ZkPathsConfig(), null, null, null, null)
    );

    taskRunner.start();

    if (listenerNotificationsAccumulator != null) {
      taskRunner.registerListener(
          new TaskRunnerListener()
          {
            @Override
            public String getListenerId()
            {
              return "test-listener";
            }

            @Override
            public void locationChanged(String taskId, TaskLocation newLocation)
            {
              listenerNotificationsAccumulator.add(ImmutableList.of(taskId, newLocation));
            }

            @Override
            public void statusChanged(String taskId, TaskStatus status)
            {
              listenerNotificationsAccumulator.add(ImmutableList.of(taskId, status));
            }
          },
          Execs.directExecutor()
      );
    }

    return taskRunner;
  }

  private static WorkerHolder createWorkerHolder(
      ObjectMapper smileMapper,
      HttpClient httpClient,
      HttpRemoteTaskRunnerConfig config,
      ScheduledExecutorService workersSyncExec,
      WorkerHolder.Listener listener,
      Worker worker,
      List<TaskAnnouncement> knownAnnouncements,

      // simulates task announcements received from worker on first sync call for the tasks that are already
      // running/completed on the worker.
      List<TaskAnnouncement> preExistingTaskAnnouncements,

      // defines behavior for what to do when a particular task is assigned
      Map<Task, List<TaskAnnouncement>> toBeAssignedTasks,

      // incremented on each runnable completion in workersSyncExec, useful for deterministically watching that some
      // work completed
      AtomicInteger ticks,

      // Updated each time a shutdown(taskId) call is received, useful for asserting that expected shutdowns indeed
      // happened.
      Set<String> actualShutdowns
  )
  {
    return new WorkerHolder(smileMapper, httpClient, config, workersSyncExec, listener, worker, knownAnnouncements)
    {
      private final String workerHost;
      private final int workerPort;

      {
        String hostAndPort = worker.getHost();
        int colonIndex = hostAndPort.indexOf(':');
        if (colonIndex == -1) {
          throw new IAE("Invalid host and port: [%s]", colonIndex);
        }
        workerHost = hostAndPort.substring(0, colonIndex);
        workerPort = Integer.parseInt(hostAndPort.substring(colonIndex + 1));
      }

      @Override
      public void start()
      {
        disabled.set(false);

        if (!preExistingTaskAnnouncements.isEmpty()) {
          workersSyncExec.execute(
              () -> {
                for (TaskAnnouncement announcement : preExistingTaskAnnouncements) {
                  tasksSnapshotRef.get().put(announcement.getTaskId(), announcement);
                  listener.taskAddedOrUpdated(announcement, this);
                }
                ticks.incrementAndGet();
              }
          );
        }
      }

      @Override
      public void stop()
      {

      }

      @Override
      public boolean isInitialized()
      {
        return true;
      }

      @Override
      public void waitForInitialization()
      {
      }

      @Override
      public boolean assignTask(Task task)
      {
        // artificial sleeps are introduced to simulate some latency.

        try {
          Thread.sleep(500);
        }
        catch (InterruptedException ex) {
          throw new RuntimeException(ex);
        }

        if (toImmutable().getCurrCapacityUsed() > worker.getCapacity()) {
          throw new ISE("Got assigned tasks more than capacity.");
        }

        final List<TaskAnnouncement> announcements;
        if (toBeAssignedTasks.containsKey(task)) {
          announcements = toBeAssignedTasks.get(task);
        } else {
          // no behavior specified for the task, so do default behavior of completing the task
          announcements = new ArrayList<>();
          announcements.add(
              TaskAnnouncement.create(
                  task,
                  TaskStatus.running(task.getId()),
                  TaskLocation.unknown()
              )
          );
          announcements.add(
              TaskAnnouncement.create(
                  task,
                  TaskStatus.running(task.getId()),
                  TaskLocation.create(workerHost, workerPort, -1)
              )
          );
          announcements.add(
              TaskAnnouncement.create(
                  task,
                  TaskStatus.success(task.getId()),
                  TaskLocation.create(workerHost, workerPort, -1)
              )
          );
        }

        workersSyncExec.execute(
            () -> {
              for (TaskAnnouncement announcement : announcements) {
                try {
                  Thread.sleep(100);
                }
                catch (InterruptedException ex) {
                  throw new RuntimeException(ex);
                }

                tasksSnapshotRef.get().put(announcement.getTaskId(), announcement);
                listener.taskAddedOrUpdated(announcement, this);
              }

              ticks.incrementAndGet();
            }
        );
        return true;
      }

      @Override
      public void shutdownTask(String taskId)
      {
        actualShutdowns.add(taskId);
      }
    };
  }

  private static class TestDruidNodeDiscovery implements DruidNodeDiscovery
  {
    Listener listener;

    @Override
    public Collection<DiscoveryDruidNode> getAllNodes()
    {
      throw new UnsupportedOperationException("Not Implemented.");
    }

    @Override
    public void registerListener(Listener listener)
    {
      listener.nodesAdded(ImmutableList.of());
      listener.nodeViewInitialized();
      this.listener = listener;
    }
  }

  private interface CustomFunction
  {
    WorkerHolder apply(
        ObjectMapper smileMapper,
        HttpClient httpClient,
        HttpRemoteTaskRunnerConfig config,
        ScheduledExecutorService workersSyncExec,
        WorkerHolder.Listener listener,
        Worker worker,
        List<TaskAnnouncement> knownAnnouncements
    );
  }
}