Examples with TaskRunnerWorkItem org.apache.druid.indexing.overlord.TaskRunnerWorkItem used on opensource projects

Example 1 with TaskRunnerWorkItem

use of org.apache.druid.indexing.overlord.TaskRunnerWorkItem in project druid by druid-io.

the class HttpRemoteTaskRunnerTest method testTaskRunnerRestart.

/*
  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.getListeners().get(0).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());
}

Example 2 with TaskRunnerWorkItem

use of org.apache.druid.indexing.overlord.TaskRunnerWorkItem in project druid by druid-io.

the class OverlordResourceTest method testGetTasksRequiresDatasourceRead.

@Test
public void testGetTasksRequiresDatasourceRead() {
    // Setup mocks for a user who has read access to "wikipedia"
    // and no access to "buzzfeed"
    expectAuthorizationTokenCheck(Users.WIKI_READER);
    // Setup mocks to return completed, active, known, pending and running tasks
    EasyMock.expect(taskStorageQueryAdapter.getCompletedTaskInfoByCreatedTimeDuration(null, null, null)).andStubReturn(ImmutableList.of(createTaskInfo("id_5", Datasources.WIKIPEDIA), createTaskInfo("id_6", Datasources.BUZZFEED)));
    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo(null)).andStubReturn(ImmutableList.of(createTaskInfo("id_1", Datasources.WIKIPEDIA), createTaskInfo("id_2", Datasources.BUZZFEED)));
    EasyMock.<Collection<? extends TaskRunnerWorkItem>>expect(taskRunner.getKnownTasks()).andReturn(ImmutableList.of(new MockTaskRunnerWorkItem("id_1", null), new MockTaskRunnerWorkItem("id_4", null))).atLeastOnce();
    EasyMock.<Collection<? extends TaskRunnerWorkItem>>expect(taskRunner.getPendingTasks()).andReturn(ImmutableList.of(new MockTaskRunnerWorkItem("id_4", null)));
    EasyMock.<Collection<? extends TaskRunnerWorkItem>>expect(taskRunner.getRunningTasks()).andReturn(ImmutableList.of(new MockTaskRunnerWorkItem("id_1", null)));
    // Replay all mocks
    EasyMock.replay(taskRunner, taskMaster, taskStorageQueryAdapter, indexerMetadataStorageAdapter, req, workerTaskRunnerQueryAdapter);
    // Verify that only the tasks of read access datasource are returned
    List<TaskStatusPlus> responseObjects = (List<TaskStatusPlus>) overlordResource.getTasks(null, null, null, null, null, req).getEntity();
    Assert.assertEquals(2, responseObjects.size());
    for (TaskStatusPlus taskStatus : responseObjects) {
        Assert.assertEquals(Datasources.WIKIPEDIA, taskStatus.getDataSource());
    }
}

Example 3 with TaskRunnerWorkItem

use of org.apache.druid.indexing.overlord.TaskRunnerWorkItem in project druid by druid-io.

the class OverlordResource method filterActiveTasks.

private List<AnyTask> filterActiveTasks(RunnerTaskState state, List<AnyTask> allTasks) {
    // divide active tasks into 3 lists : running, pending, waiting
    Optional<TaskRunner> taskRunnerOpt = taskMaster.getTaskRunner();
    if (!taskRunnerOpt.isPresent()) {
        throw new WebApplicationException(Response.serverError().entity("No task runner found").build());
    }
    TaskRunner runner = taskRunnerOpt.get();
    // counted twice but never skipped
    if (RunnerTaskState.WAITING.equals(state)) {
        Collection<? extends TaskRunnerWorkItem> runnersKnownTasks = runner.getKnownTasks();
        Set<String> runnerKnownTaskIds = runnersKnownTasks.stream().map(TaskRunnerWorkItem::getTaskId).collect(Collectors.toSet());
        final List<AnyTask> waitingTasks = new ArrayList<>();
        for (TaskRunnerWorkItem task : allTasks) {
            if (!runnerKnownTaskIds.contains(task.getTaskId())) {
                waitingTasks.add(((AnyTask) task).withTaskState(TaskState.RUNNING, RunnerTaskState.WAITING, task.getCreatedTime(), task.getQueueInsertionTime(), task.getLocation()));
            }
        }
        return waitingTasks;
    }
    if (RunnerTaskState.PENDING.equals(state)) {
        Collection<? extends TaskRunnerWorkItem> knownPendingTasks = runner.getPendingTasks();
        Set<String> pendingTaskIds = knownPendingTasks.stream().map(TaskRunnerWorkItem::getTaskId).collect(Collectors.toSet());
        Map<String, TaskRunnerWorkItem> workItemIdMap = knownPendingTasks.stream().collect(Collectors.toMap(TaskRunnerWorkItem::getTaskId, java.util.function.Function.identity(), (previousWorkItem, newWorkItem) -> newWorkItem));
        final List<AnyTask> pendingTasks = new ArrayList<>();
        for (TaskRunnerWorkItem task : allTasks) {
            if (pendingTaskIds.contains(task.getTaskId())) {
                pendingTasks.add(((AnyTask) task).withTaskState(TaskState.RUNNING, RunnerTaskState.PENDING, workItemIdMap.get(task.getTaskId()).getCreatedTime(), workItemIdMap.get(task.getTaskId()).getQueueInsertionTime(), workItemIdMap.get(task.getTaskId()).getLocation()));
            }
        }
        return pendingTasks;
    }
    if (RunnerTaskState.RUNNING.equals(state)) {
        Collection<? extends TaskRunnerWorkItem> knownRunningTasks = runner.getRunningTasks();
        Set<String> runningTaskIds = knownRunningTasks.stream().map(TaskRunnerWorkItem::getTaskId).collect(Collectors.toSet());
        Map<String, TaskRunnerWorkItem> workItemIdMap = knownRunningTasks.stream().collect(Collectors.toMap(TaskRunnerWorkItem::getTaskId, java.util.function.Function.identity(), (previousWorkItem, newWorkItem) -> newWorkItem));
        final List<AnyTask> runningTasks = new ArrayList<>();
        for (TaskRunnerWorkItem task : allTasks) {
            if (runningTaskIds.contains(task.getTaskId())) {
                runningTasks.add(((AnyTask) task).withTaskState(TaskState.RUNNING, RunnerTaskState.RUNNING, workItemIdMap.get(task.getTaskId()).getCreatedTime(), workItemIdMap.get(task.getTaskId()).getQueueInsertionTime(), workItemIdMap.get(task.getTaskId()).getLocation()));
            }
        }
        return runningTasks;
    }
    return allTasks;
}

Example 4 with TaskRunnerWorkItem

use of org.apache.druid.indexing.overlord.TaskRunnerWorkItem in project druid by druid-io.

the class OverlordResource method getTaskStatus.

@GET
@Path("/task/{taskid}/status")
@Produces(MediaType.APPLICATION_JSON)
@ResourceFilters(TaskResourceFilter.class)
public Response getTaskStatus(@PathParam("taskid") String taskid) {
    final TaskInfo<Task, TaskStatus> taskInfo = taskStorageQueryAdapter.getTaskInfo(taskid);
    TaskStatusResponse response = null;
    if (taskInfo != null) {
        if (taskMaster.getTaskRunner().isPresent()) {
            final TaskRunner taskRunner = taskMaster.getTaskRunner().get();
            final TaskRunnerWorkItem workItem = taskRunner.getKnownTasks().stream().filter(item -> item.getTaskId().equals(taskid)).findAny().orElse(null);
            if (workItem != null) {
                response = new TaskStatusResponse(workItem.getTaskId(), new TaskStatusPlus(taskInfo.getId(), taskInfo.getTask() == null ? null : taskInfo.getTask().getGroupId(), taskInfo.getTask() == null ? null : taskInfo.getTask().getType(), taskInfo.getCreatedTime(), // TaskStorage API doesn't yet allow it.
                DateTimes.EPOCH, taskInfo.getStatus().getStatusCode(), taskRunner.getRunnerTaskState(workItem.getTaskId()), taskInfo.getStatus().getDuration(), workItem.getLocation(), taskInfo.getDataSource(), taskInfo.getStatus().getErrorMsg()));
            }
        }
        if (response == null) {
            response = new TaskStatusResponse(taskid, new TaskStatusPlus(taskInfo.getId(), taskInfo.getTask() == null ? null : taskInfo.getTask().getGroupId(), taskInfo.getTask() == null ? null : taskInfo.getTask().getType(), taskInfo.getCreatedTime(), // TaskStorage API doesn't yet allow it.
            DateTimes.EPOCH, taskInfo.getStatus().getStatusCode(), RunnerTaskState.WAITING, taskInfo.getStatus().getDuration(), taskInfo.getStatus().getLocation() == null ? TaskLocation.unknown() : taskInfo.getStatus().getLocation(), taskInfo.getDataSource(), taskInfo.getStatus().getErrorMsg()));
        }
    } else {
        response = new TaskStatusResponse(taskid, null);
    }
    final Response.Status status = response.getStatus() == null ? Response.Status.NOT_FOUND : Response.Status.OK;
    return Response.status(status).entity(response).build();
}