Examples with TaskLocation org.apache.druid.indexer.TaskLocation used on opensource projects

Example 36 with TaskLocation

use of org.apache.druid.indexer.TaskLocation in project druid by druid-io.

the class TaskQueue method notifyStatus.

/**
 * Notify this queue that some task has an updated status. If this update is valid, the status will be persisted in
 * the task storage facility. If the status is a completed status, the task will be unlocked and no further
 * updates will be accepted.
 *
 * @param task       task to update
 * @param taskStatus new task status
 *
 * @throws NullPointerException     if task or status is null
 * @throws IllegalArgumentException if the task ID does not match the status ID
 * @throws IllegalStateException    if this queue is currently shut down
 */
private void notifyStatus(final Task task, final TaskStatus taskStatus, String reasonFormat, Object... args) {
    giant.lock();
    TaskLocation taskLocation = TaskLocation.unknown();
    try {
        Preconditions.checkNotNull(task, "task");
        Preconditions.checkNotNull(taskStatus, "status");
        Preconditions.checkState(active, "Queue is not active!");
        Preconditions.checkArgument(task.getId().equals(taskStatus.getId()), "Mismatching task ids[%s/%s]", task.getId(), taskStatus.getId());
        // Inform taskRunner that this task can be shut down
        try {
            taskLocation = taskRunner.getTaskLocation(task.getId());
            taskRunner.shutdown(task.getId(), reasonFormat, args);
        } catch (Exception e) {
            log.warn(e, "TaskRunner failed to cleanup task after completion: %s", task.getId());
        }
        // Remove from running tasks
        int removed = 0;
        for (int i = tasks.size() - 1; i >= 0; i--) {
            if (tasks.get(i).getId().equals(task.getId())) {
                removed++;
                removeTaskInternal(tasks.get(i));
                break;
            }
        }
        if (removed == 0) {
            log.warn("Unknown task completed: %s", task.getId());
        } else if (removed > 1) {
            log.makeAlert("Removed multiple copies of task").addData("count", removed).addData("task", task.getId()).emit();
        }
        // Remove from futures list
        taskFutures.remove(task.getId());
        if (removed > 0) {
            // If we thought this task should be running, save status to DB
            try {
                final Optional<TaskStatus> previousStatus = taskStorage.getStatus(task.getId());
                if (!previousStatus.isPresent() || !previousStatus.get().isRunnable()) {
                    log.makeAlert("Ignoring notification for already-complete task").addData("task", task.getId()).emit();
                } else {
                    taskStorage.setStatus(taskStatus.withLocation(taskLocation));
                    log.info("Task done: %s", task);
                    managementMayBeNecessary.signalAll();
                }
            } catch (Exception e) {
                log.makeAlert(e, "Failed to persist status for task").addData("task", task.getId()).addData("statusCode", taskStatus.getStatusCode()).emit();
            }
        }
    } finally {
        giant.unlock();
    }
}

Example 37 with TaskLocation

use of org.apache.druid.indexer.TaskLocation in project druid by druid-io.

the class ThreadingTaskRunner method run.

@Override
public ListenableFuture<TaskStatus> run(Task task) {
    synchronized (tasks) {
        tasks.computeIfAbsent(task.getId(), k -> new ThreadingTaskRunnerWorkItem(task, taskExecutor.submit(new Callable<TaskStatus>() {

            @Override
            public TaskStatus call() {
                final String attemptUUID = UUID.randomUUID().toString();
                final File taskDir = taskConfig.getTaskDir(task.getId());
                final File attemptDir = new File(taskDir, attemptUUID);
                final TaskLocation taskLocation = TaskLocation.create(node.getHost(), node.getPlaintextPort(), node.getTlsPort());
                final ThreadingTaskRunnerWorkItem taskWorkItem;
                try {
                    FileUtils.mkdirp(attemptDir);
                    final File taskFile = new File(taskDir, "task.json");
                    final File reportsFile = new File(attemptDir, "report.json");
                    taskReportFileWriter.add(task.getId(), reportsFile);
                    // time to adjust process holders
                    synchronized (tasks) {
                        taskWorkItem = tasks.get(task.getId());
                        if (taskWorkItem == null) {
                            LOGGER.makeAlert("TaskInfo disappeared").addData("task", task.getId()).emit();
                            throw new ISE("TaskInfo disappeared for task[%s]!", task.getId());
                        }
                        if (taskWorkItem.shutdown) {
                            throw new IllegalStateException("Task has been shut down!");
                        }
                    }
                    if (!taskFile.exists()) {
                        jsonMapper.writeValue(taskFile, task);
                    }
                    // This will block for a while. So we append the thread information with more details
                    final String priorThreadName = Thread.currentThread().getName();
                    Thread.currentThread().setName(StringUtils.format("[%s]-%s", task.getId(), priorThreadName));
                    TaskStatus taskStatus;
                    final TaskToolbox toolbox = toolboxFactory.build(task);
                    TaskRunnerUtils.notifyLocationChanged(listeners, task.getId(), taskLocation);
                    TaskRunnerUtils.notifyStatusChanged(listeners, task.getId(), TaskStatus.running(task.getId()));
                    taskWorkItem.setState(RunnerTaskState.RUNNING);
                    try {
                        taskStatus = task.run(toolbox);
                    } catch (Throwable t) {
                        LOGGER.error(t, "Exception caught while running the task.");
                        taskStatus = TaskStatus.failure(task.getId(), "Failed with an exception. See indexer logs for more details.");
                    } finally {
                        taskWorkItem.setState(RunnerTaskState.NONE);
                        Thread.currentThread().setName(priorThreadName);
                        if (reportsFile.exists()) {
                            taskLogPusher.pushTaskReports(task.getId(), reportsFile);
                        }
                    }
                    TaskRunnerUtils.notifyStatusChanged(listeners, task.getId(), taskStatus);
                    return taskStatus;
                } catch (Throwable t) {
                    LOGGER.error(t, "Exception caught during execution");
                    throw new RuntimeException(t);
                } finally {
                    try {
                        taskReportFileWriter.delete(task.getId());
                        appenderatorsManager.removeAppenderatorsForTask(task.getId(), task.getDataSource());
                        synchronized (tasks) {
                            tasks.remove(task.getId());
                            if (!stopping) {
                                saveRunningTasks();
                            }
                        }
                        try {
                            if (!stopping && taskDir.exists()) {
                                FileUtils.deleteDirectory(taskDir);
                                LOGGER.info("Removed task directory: %s", taskDir);
                            }
                        } catch (Exception e) {
                            LOGGER.makeAlert(e, "Failed to delete task directory").addData("taskDir", taskDir.toString()).addData("task", task.getId()).emit();
                        }
                    } catch (Exception e) {
                        LOGGER.error(e, "Suppressing exception caught while cleaning up task");
                    }
                }
            }
        })));
        saveRunningTasks();
        return tasks.get(task.getId()).getResult();
    }
}

Example 38 with TaskLocation

use of org.apache.druid.indexer.TaskLocation in project druid by druid-io.

the class KinesisSupervisorTest method testResetRunningTasks.

@Test
public void testResetRunningTasks() throws Exception {
    final TaskLocation location1 = new TaskLocation("testHost", 1234, -1);
    final TaskLocation location2 = new TaskLocation("testHost2", 145, -1);
    final DateTime startTime = DateTimes.nowUtc();
    supervisor = getTestableSupervisor(2, 1, true, "PT1H", null, null);
    supervisorRecordSupplier.assign(EasyMock.anyObject());
    EasyMock.expectLastCall().anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID1, SHARD_ID0)).anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(SHARD1_PARTITION, SHARD0_PARTITION)).anyTimes();
    supervisorRecordSupplier.seekToLatest(EasyMock.anyObject());
    EasyMock.expectLastCall().anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getEarliestSequenceNumber(EasyMock.anyObject())).andReturn("0").anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getLatestSequenceNumber(SHARD1_PARTITION)).andReturn("12").anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getLatestSequenceNumber(SHARD0_PARTITION)).andReturn("1").anyTimes();
    supervisorRecordSupplier.seek(EasyMock.anyObject(), EasyMock.anyString());
    EasyMock.expectLastCall().anyTimes();
    Task id1 = createKinesisIndexTask("id1", DATASOURCE, 0, new SeekableStreamStartSequenceNumbers<>("stream", ImmutableMap.of(SHARD_ID1, "0", SHARD_ID0, "0"), ImmutableSet.of()), new SeekableStreamEndSequenceNumbers<>("stream", ImmutableMap.of(SHARD_ID1, KinesisSequenceNumber.NO_END_SEQUENCE_NUMBER, SHARD_ID0, KinesisSequenceNumber.NO_END_SEQUENCE_NUMBER)), null, null);
    Task id2 = createKinesisIndexTask("id2", DATASOURCE, 0, new SeekableStreamStartSequenceNumbers<>("stream", ImmutableMap.of(SHARD_ID1, "3", SHARD_ID0, "1"), ImmutableSet.of(SHARD_ID0, SHARD_ID1)), new SeekableStreamEndSequenceNumbers<>("stream", ImmutableMap.of(SHARD_ID1, KinesisSequenceNumber.NO_END_SEQUENCE_NUMBER, SHARD_ID0, KinesisSequenceNumber.NO_END_SEQUENCE_NUMBER)), null, null);
    Task id3 = createKinesisIndexTask("id3", DATASOURCE, 0, new SeekableStreamStartSequenceNumbers<>("stream", ImmutableMap.of(SHARD_ID1, "3", SHARD_ID0, "1"), ImmutableSet.of(SHARD_ID0, SHARD_ID1)), new SeekableStreamEndSequenceNumbers<>("stream", ImmutableMap.of(SHARD_ID1, KinesisSequenceNumber.NO_END_SEQUENCE_NUMBER, SHARD_ID0, KinesisSequenceNumber.NO_END_SEQUENCE_NUMBER)), null, null);
    Collection workItems = new ArrayList<>();
    workItems.add(new TestTaskRunnerWorkItem(id1, null, location1));
    workItems.add(new TestTaskRunnerWorkItem(id2, null, location2));
    EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes();
    EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes();
    EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes();
    EasyMock.expect(taskStorage.getActiveTasksByDatasource(DATASOURCE)).andReturn(ImmutableList.of(id1, id2, id3)).anyTimes();
    EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes();
    EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes();
    EasyMock.expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes();
    EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes();
    EasyMock.expect(taskStorage.getTask("id2")).andReturn(Optional.of(id2)).anyTimes();
    EasyMock.expect(taskStorage.getTask("id3")).andReturn(Optional.of(id3)).anyTimes();
    EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn(new KinesisDataSourceMetadata(null)).anyTimes();
    EasyMock.expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(SeekableStreamIndexTaskRunner.Status.PUBLISHING));
    EasyMock.expect(taskClient.getStatusAsync("id2")).andReturn(Futures.immediateFuture(SeekableStreamIndexTaskRunner.Status.READING));
    EasyMock.expect(taskClient.getStatusAsync("id3")).andReturn(Futures.immediateFuture(SeekableStreamIndexTaskRunner.Status.READING));
    EasyMock.expect(taskClient.getStartTimeAsync("id2")).andReturn(Futures.immediateFuture(startTime));
    EasyMock.expect(taskClient.getStartTimeAsync("id3")).andReturn(Futures.immediateFuture(startTime));
    EasyMock.expect(taskClient.getEndOffsets("id1")).andReturn(ImmutableMap.of(SHARD_ID1, "3", SHARD_ID0, "1"));
    TreeMap<Integer, Map<String, String>> checkpoints = new TreeMap<>();
    checkpoints.put(0, ImmutableMap.of(SHARD_ID1, "3", SHARD_ID0, "1"));
    EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())).andReturn(Futures.immediateFuture(checkpoints)).times(1);
    EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), EasyMock.anyBoolean())).andReturn(Futures.immediateFuture(checkpoints)).times(1);
    taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class));
    replayAll();
    supervisor.start();
    supervisor.runInternal();
    verifyAll();
    EasyMock.reset(taskQueue, indexerMetadataStorageCoordinator);
    EasyMock.expect(indexerMetadataStorageCoordinator.deleteDataSourceMetadata(DATASOURCE)).andReturn(true);
    taskQueue.shutdown("id2", "DataSourceMetadata is not found while reset");
    taskQueue.shutdown("id3", "DataSourceMetadata is not found while reset");
    EasyMock.replay(taskQueue, indexerMetadataStorageCoordinator);
    supervisor.resetInternal(null);
    verifyAll();
}

Example 39 with TaskLocation

use of org.apache.druid.indexer.TaskLocation in project druid by druid-io.

the class KinesisSupervisorTest method testDiscoverExistingPublishingTask.

@Test
public void testDiscoverExistingPublishingTask() throws Exception {
    final TaskLocation location = new TaskLocation("testHost", 1234, -1);
    final Map<String, Long> timeLag = ImmutableMap.of(SHARD_ID1, 0L, SHARD_ID0, 20000000L);
    supervisor = getTestableSupervisor(1, 1, true, "PT1H", null, null);
    supervisorRecordSupplier.assign(EasyMock.anyObject());
    EasyMock.expectLastCall().anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID1, SHARD_ID0)).anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(SHARD1_PARTITION, SHARD0_PARTITION)).anyTimes();
    supervisorRecordSupplier.seekToLatest(EasyMock.anyObject());
    EasyMock.expectLastCall().anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getEarliestSequenceNumber(EasyMock.anyObject())).andReturn("0").anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getLatestSequenceNumber(SHARD1_PARTITION)).andReturn("12").anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getLatestSequenceNumber(SHARD0_PARTITION)).andReturn("1").anyTimes();
    supervisorRecordSupplier.seek(EasyMock.anyObject(), EasyMock.anyString());
    EasyMock.expectLastCall().anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getPartitionsTimeLag(EasyMock.anyString(), EasyMock.anyObject())).andReturn(timeLag).atLeastOnce();
    Task task = createKinesisIndexTask("id1", DATASOURCE, 0, new SeekableStreamStartSequenceNumbers<>("stream", ImmutableMap.of(SHARD_ID1, "0", SHARD_ID0, "0"), ImmutableSet.of()), new SeekableStreamEndSequenceNumbers<>("stream", ImmutableMap.of(SHARD_ID1, KinesisSequenceNumber.NO_END_SEQUENCE_NUMBER, SHARD_ID0, KinesisSequenceNumber.NO_END_SEQUENCE_NUMBER)), null, null);
    Collection workItems = new ArrayList<>();
    workItems.add(new TestTaskRunnerWorkItem(task, null, location));
    Capture<KinesisIndexTask> captured = Capture.newInstance();
    EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes();
    EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes();
    EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes();
    EasyMock.expect(taskStorage.getActiveTasksByDatasource(DATASOURCE)).andReturn(ImmutableList.of(task)).anyTimes();
    EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes();
    EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(task)).anyTimes();
    EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn(new KinesisDataSourceMetadata(null)).anyTimes();
    EasyMock.expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(SeekableStreamIndexTaskRunner.Status.PUBLISHING));
    EasyMock.expect(taskClient.getCurrentOffsetsAsync("id1", false)).andReturn(Futures.immediateFuture(ImmutableMap.of(SHARD_ID1, "2", SHARD_ID0, "1")));
    EasyMock.expect(taskClient.getEndOffsets("id1")).andReturn(ImmutableMap.of(SHARD_ID1, "2", SHARD_ID0, "1"));
    EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true);
    TreeMap<Integer, Map<String, String>> checkpoints = new TreeMap<>();
    checkpoints.put(0, ImmutableMap.of(SHARD_ID1, "0", SHARD_ID0, "0"));
    EasyMock.expect(taskClient.getCheckpoints(EasyMock.anyString(), EasyMock.anyBoolean())).andReturn(checkpoints).anyTimes();
    taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class));
    replayAll();
    supervisor.start();
    supervisor.runInternal();
    supervisor.updateCurrentAndLatestOffsets();
    SupervisorReport<KinesisSupervisorReportPayload> report = supervisor.getStatus();
    verifyAll();
    Assert.assertEquals(DATASOURCE, report.getId());
    KinesisSupervisorReportPayload payload = report.getPayload();
    Assert.assertEquals(DATASOURCE, payload.getDataSource());
    Assert.assertEquals(3600L, payload.getDurationSeconds());
    Assert.assertEquals(2, payload.getPartitions());
    Assert.assertEquals(1, payload.getReplicas());
    Assert.assertEquals(STREAM, payload.getStream());
    Assert.assertEquals(0, payload.getActiveTasks().size());
    Assert.assertEquals(1, payload.getPublishingTasks().size());
    Assert.assertEquals(SupervisorStateManager.BasicState.RUNNING, payload.getDetailedState());
    Assert.assertEquals(0, payload.getRecentErrors().size());
    TaskReportData publishingReport = payload.getPublishingTasks().get(0);
    Assert.assertEquals("id1", publishingReport.getId());
    Assert.assertEquals(ImmutableMap.of(SHARD_ID1, "0", SHARD_ID0, "0"), publishingReport.getStartingOffsets());
    Assert.assertEquals(ImmutableMap.of(SHARD_ID1, "2", SHARD_ID0, "1"), publishingReport.getCurrentOffsets());
    KinesisIndexTask capturedTask = captured.getValue();
    Assert.assertEquals(dataSchema, capturedTask.getDataSchema());
    Assert.assertEquals(tuningConfig.convertToTaskTuningConfig(), capturedTask.getTuningConfig());
    KinesisIndexTaskIOConfig capturedTaskConfig = capturedTask.getIOConfig();
    Assert.assertEquals("awsEndpoint", capturedTaskConfig.getEndpoint());
    Assert.assertEquals("sequenceName-0", capturedTaskConfig.getBaseSequenceName());
    Assert.assertTrue("isUseTransaction", capturedTaskConfig.isUseTransaction());
    // check that the new task was created with starting sequences matching where the publishing task finished
    Assert.assertEquals(STREAM, capturedTaskConfig.getStartSequenceNumbers().getStream());
    Assert.assertEquals("2", capturedTaskConfig.getStartSequenceNumbers().getPartitionSequenceNumberMap().get(SHARD_ID1));
    Assert.assertEquals("1", capturedTaskConfig.getStartSequenceNumbers().getPartitionSequenceNumberMap().get(SHARD_ID0));
    Assert.assertEquals(STREAM, capturedTaskConfig.getEndSequenceNumbers().getStream());
    Assert.assertEquals(KinesisSequenceNumber.NO_END_SEQUENCE_NUMBER, capturedTaskConfig.getEndSequenceNumbers().getPartitionSequenceNumberMap().get(SHARD_ID1));
    Assert.assertEquals(KinesisSequenceNumber.NO_END_SEQUENCE_NUMBER, capturedTaskConfig.getEndSequenceNumbers().getPartitionSequenceNumberMap().get(SHARD_ID0));
}

Example 40 with TaskLocation

use of org.apache.druid.indexer.TaskLocation in project druid by druid-io.

the class KinesisSupervisorTest method testBeginPublishAndQueueNextTasks.

@Test
public void testBeginPublishAndQueueNextTasks() throws Exception {
    final TaskLocation location = new TaskLocation("testHost", 1234, -1);
    supervisor = getTestableSupervisor(2, 2, true, "PT1M", null, null);
    supervisorRecordSupplier.assign(EasyMock.anyObject());
    EasyMock.expectLastCall().anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID1, SHARD_ID0)).anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(SHARD1_PARTITION, SHARD0_PARTITION)).anyTimes();
    supervisorRecordSupplier.seekToLatest(EasyMock.anyObject());
    EasyMock.expectLastCall().anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getEarliestSequenceNumber(EasyMock.anyObject())).andReturn("0").anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getLatestSequenceNumber(SHARD1_PARTITION)).andReturn("12").anyTimes();
    EasyMock.expect(supervisorRecordSupplier.getLatestSequenceNumber(SHARD0_PARTITION)).andReturn("1").anyTimes();
    supervisorRecordSupplier.seek(EasyMock.anyObject(), EasyMock.anyString());
    EasyMock.expectLastCall().anyTimes();
    final Capture<Task> firstTasks = Capture.newInstance(CaptureType.ALL);
    EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes();
    EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes();
    EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.emptyList()).anyTimes();
    EasyMock.expect(taskStorage.getActiveTasksByDatasource(DATASOURCE)).andReturn(ImmutableList.of()).anyTimes();
    EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn(new KinesisDataSourceMetadata(null)).anyTimes();
    EasyMock.expect(taskQueue.add(EasyMock.capture(firstTasks))).andReturn(true).times(4);
    taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class));
    replayAll();
    supervisor.start();
    supervisor.runInternal();
    verifyAll();
    final List<Task> tasks = firstTasks.getValues();
    Collection workItems = new ArrayList<>();
    for (Task task : tasks) {
        workItems.add(new TestTaskRunnerWorkItem(task, null, location));
    }
    EasyMock.reset(taskStorage, taskRunner, taskClient, taskQueue);
    final Capture<Task> secondTasks = Capture.newInstance(CaptureType.ALL);
    EasyMock.expect(taskStorage.getActiveTasksByDatasource(DATASOURCE)).andReturn(tasks).anyTimes();
    for (Task task : tasks) {
        EasyMock.expect(taskStorage.getStatus(task.getId())).andReturn(Optional.of(TaskStatus.running(task.getId()))).anyTimes();
        EasyMock.expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes();
    }
    EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes();
    EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())).andReturn(Futures.immediateFuture(SeekableStreamIndexTaskRunner.Status.READING)).anyTimes();
    EasyMock.expect(taskClient.getStartTimeAsync(EasyMock.contains("sequenceName-0"))).andReturn(Futures.immediateFuture(DateTimes.nowUtc().minusMinutes(2))).andReturn(Futures.immediateFuture(DateTimes.nowUtc()));
    EasyMock.expect(taskClient.getStartTimeAsync(EasyMock.contains("sequenceName-1"))).andReturn(Futures.immediateFuture(DateTimes.nowUtc())).times(2);
    EasyMock.expect(taskClient.pauseAsync(EasyMock.contains("sequenceName-0"))).andReturn(Futures.immediateFuture(ImmutableMap.of(SHARD_ID1, "1"))).andReturn(Futures.immediateFuture(ImmutableMap.of(SHARD_ID1, "3")));
    EasyMock.expect(taskClient.setEndOffsetsAsync(EasyMock.contains("sequenceName-0"), EasyMock.eq(ImmutableMap.of(SHARD_ID1, "3")), EasyMock.eq(true))).andReturn(Futures.immediateFuture(true)).times(2);
    EasyMock.expect(taskQueue.add(EasyMock.capture(secondTasks))).andReturn(true).times(2);
    TreeMap<Integer, Map<String, String>> checkpoints1 = new TreeMap<>();
    checkpoints1.put(0, ImmutableMap.of(SHARD_ID1, "0"));
    TreeMap<Integer, Map<String, String>> checkpoints2 = new TreeMap<>();
    checkpoints2.put(0, ImmutableMap.of(SHARD_ID0, "0"));
    EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())).andReturn(Futures.immediateFuture(checkpoints1)).times(2);
    EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), EasyMock.anyBoolean())).andReturn(Futures.immediateFuture(checkpoints2)).times(2);
    EasyMock.replay(taskStorage, taskRunner, taskClient, taskQueue);
    supervisor.runInternal();
    verifyAll();
    for (Task task : secondTasks.getValues()) {
        KinesisIndexTask kinesisIndexTask = (KinesisIndexTask) task;
        Assert.assertEquals(dataSchema, kinesisIndexTask.getDataSchema());
        Assert.assertEquals(tuningConfig.convertToTaskTuningConfig(), kinesisIndexTask.getTuningConfig());
        KinesisIndexTaskIOConfig taskConfig = kinesisIndexTask.getIOConfig();
        Assert.assertEquals("sequenceName-0", taskConfig.getBaseSequenceName());
        Assert.assertTrue("isUseTransaction", taskConfig.isUseTransaction());
        Assert.assertEquals(STREAM, taskConfig.getStartSequenceNumbers().getStream());
        Assert.assertEquals("3", taskConfig.getStartSequenceNumbers().getPartitionSequenceNumberMap().get(SHARD_ID1));
        // start sequenceNumbers should be exclusive for the second batch of tasks
        Assert.assertEquals(ImmutableSet.of(SHARD_ID1), ((KinesisIndexTask) task).getIOConfig().getStartSequenceNumbers().getExclusivePartitions());
    }
}