Examples with TaskActionClientFactory org.apache.druid.indexing.common.actions.TaskActionClientFactory used on opensource projects

Example 1 with TaskActionClientFactory

use of org.apache.druid.indexing.common.actions.TaskActionClientFactory in project druid by druid-io.

the class AppenderatorDriverRealtimeIndexTaskTest method makeToolboxFactory.

private void makeToolboxFactory(final File directory) {
    taskStorage = new HeapMemoryTaskStorage(new TaskStorageConfig(null));
    publishedSegments = new CopyOnWriteArrayList<>();
    ObjectMapper mapper = new DefaultObjectMapper();
    mapper.registerSubtypes(LinearShardSpec.class);
    mapper.registerSubtypes(NumberedShardSpec.class);
    IndexerSQLMetadataStorageCoordinator mdc = new IndexerSQLMetadataStorageCoordinator(mapper, derbyConnectorRule.metadataTablesConfigSupplier().get(), derbyConnectorRule.getConnector()) {

        @Override
        public Set<DataSegment> announceHistoricalSegments(Set<DataSegment> segments) throws IOException {
            Set<DataSegment> result = super.announceHistoricalSegments(segments);
            Assert.assertFalse("Segment latch not initialized, did you forget to call expectPublishSegments?", segmentLatch == null);
            publishedSegments.addAll(result);
            segments.forEach(s -> segmentLatch.countDown());
            return result;
        }

        @Override
        public SegmentPublishResult announceHistoricalSegments(Set<DataSegment> segments, Set<DataSegment> segmentsToDrop, DataSourceMetadata startMetadata, DataSourceMetadata endMetadata) throws IOException {
            SegmentPublishResult result = super.announceHistoricalSegments(segments, segmentsToDrop, startMetadata, endMetadata);
            Assert.assertFalse("Segment latch not initialized, did you forget to call expectPublishSegments?", segmentLatch == null);
            publishedSegments.addAll(result.getSegments());
            result.getSegments().forEach(s -> segmentLatch.countDown());
            return result;
        }
    };
    taskLockbox = new TaskLockbox(taskStorage, mdc);
    final TaskConfig taskConfig = new TaskConfig(directory.getPath(), null, null, 50000, null, true, null, null, null, false, false, TaskConfig.BATCH_PROCESSING_MODE_DEFAULT.name());
    final TaskActionToolbox taskActionToolbox = new TaskActionToolbox(taskLockbox, taskStorage, mdc, EMITTER, EasyMock.createMock(SupervisorManager.class));
    final TaskActionClientFactory taskActionClientFactory = new LocalTaskActionClientFactory(taskStorage, taskActionToolbox, new TaskAuditLogConfig(false));
    final QueryRunnerFactoryConglomerate conglomerate = new DefaultQueryRunnerFactoryConglomerate(ImmutableMap.of(TimeseriesQuery.class, new TimeseriesQueryRunnerFactory(new TimeseriesQueryQueryToolChest(), new TimeseriesQueryEngine(), (query, future) -> {
    // do nothing
    })));
    handOffCallbacks = new ConcurrentHashMap<>();
    final SegmentHandoffNotifierFactory handoffNotifierFactory = dataSource -> new SegmentHandoffNotifier() {

        @Override
        public boolean registerSegmentHandoffCallback(SegmentDescriptor descriptor, Executor exec, Runnable handOffRunnable) {
            handOffCallbacks.put(descriptor, new Pair<>(exec, handOffRunnable));
            handoffLatch.countDown();
            return true;
        }

        @Override
        public void start() {
        // Noop
        }

        @Override
        public void close() {
        // Noop
        }
    };
    final TestUtils testUtils = new TestUtils();
    taskToolboxFactory = new TaskToolboxFactory(taskConfig, new DruidNode("druid/middlemanager", "localhost", false, 8091, null, true, false), taskActionClientFactory, EMITTER, new TestDataSegmentPusher(), new TestDataSegmentKiller(), // DataSegmentMover
    null, // DataSegmentArchiver
    null, new TestDataSegmentAnnouncer(), EasyMock.createNiceMock(DataSegmentServerAnnouncer.class), handoffNotifierFactory, () -> conglomerate, // queryExecutorService
    DirectQueryProcessingPool.INSTANCE, NoopJoinableFactory.INSTANCE, () -> EasyMock.createMock(MonitorScheduler.class), new SegmentCacheManagerFactory(testUtils.getTestObjectMapper()), testUtils.getTestObjectMapper(), testUtils.getTestIndexIO(), MapCache.create(1024), new CacheConfig(), new CachePopulatorStats(), testUtils.getTestIndexMergerV9(), EasyMock.createNiceMock(DruidNodeAnnouncer.class), EasyMock.createNiceMock(DruidNode.class), new LookupNodeService("tier"), new DataNodeService("tier", 1000, ServerType.INDEXER_EXECUTOR, 0), new SingleFileTaskReportFileWriter(reportsFile), null, AuthTestUtils.TEST_AUTHORIZER_MAPPER, new NoopChatHandlerProvider(), testUtils.getRowIngestionMetersFactory(), new TestAppenderatorsManager(), new NoopIndexingServiceClient(), null, null, null);
}

Example 2 with TaskActionClientFactory

use of org.apache.druid.indexing.common.actions.TaskActionClientFactory in project druid by druid-io.

the class TaskQueueTest method testShutdownReleasesTaskLock.

@Test
public void testShutdownReleasesTaskLock() throws Exception {
    final TaskActionClientFactory actionClientFactory = createActionClientFactory();
    final TaskQueue taskQueue = new TaskQueue(new TaskLockConfig(), new TaskQueueConfig(null, null, null, null), new DefaultTaskConfig(), getTaskStorage(), new SimpleTaskRunner(actionClientFactory), actionClientFactory, getLockbox(), new NoopServiceEmitter());
    taskQueue.setActive(true);
    // Create a Task and add it to the TaskQueue
    final TestTask task = new TestTask("t1", Intervals.of("2021-01/P1M"));
    taskQueue.add(task);
    // Acquire a lock for the Task
    getLockbox().lock(task, new TimeChunkLockRequest(TaskLockType.EXCLUSIVE, task, task.interval, null));
    final List<TaskLock> locksForTask = getLockbox().findLocksForTask(task);
    Assert.assertEquals(1, locksForTask.size());
    Assert.assertEquals(task.interval, locksForTask.get(0).getInterval());
    // Verify that locks are removed on calling shutdown
    taskQueue.shutdown(task.getId(), "Shutdown Task test");
    Assert.assertTrue(getLockbox().findLocksForTask(task).isEmpty());
    Optional<TaskStatus> statusOptional = getTaskStorage().getStatus(task.getId());
    Assert.assertTrue(statusOptional.isPresent());
    Assert.assertEquals(TaskState.FAILED, statusOptional.get().getStatusCode());
    Assert.assertNotNull(statusOptional.get().getErrorMsg());
    Assert.assertEquals("Shutdown Task test", statusOptional.get().getErrorMsg());
}

Example 3 with TaskActionClientFactory

use of org.apache.druid.indexing.common.actions.TaskActionClientFactory in project druid by druid-io.

the class TaskQueueTest method testUserProvidedContextOverrideLockConfig.

@Test
public void testUserProvidedContextOverrideLockConfig() throws EntryExistsException {
    final TaskActionClientFactory actionClientFactory = createActionClientFactory();
    final TaskQueue taskQueue = new TaskQueue(new TaskLockConfig(), new TaskQueueConfig(null, null, null, null), new DefaultTaskConfig(), getTaskStorage(), new SimpleTaskRunner(actionClientFactory), actionClientFactory, getLockbox(), new NoopServiceEmitter());
    taskQueue.setActive(true);
    final Task task = new TestTask("t1", Intervals.of("2021-01-01/P1D"), ImmutableMap.of(Tasks.FORCE_TIME_CHUNK_LOCK_KEY, false));
    taskQueue.add(task);
    final List<Task> tasks = taskQueue.getTasks();
    Assert.assertEquals(1, tasks.size());
    final Task queuedTask = tasks.get(0);
    Assert.assertFalse(queuedTask.getContextValue(Tasks.FORCE_TIME_CHUNK_LOCK_KEY));
}

Example 4 with TaskActionClientFactory

use of org.apache.druid.indexing.common.actions.TaskActionClientFactory in project druid by druid-io.

the class TaskQueueTest method testDefaultTaskContextOverrideDefaultLineageBasedSegmentAllocation.

@Test
public void testDefaultTaskContextOverrideDefaultLineageBasedSegmentAllocation() throws EntryExistsException {
    final TaskActionClientFactory actionClientFactory = createActionClientFactory();
    final TaskQueue taskQueue = new TaskQueue(new TaskLockConfig(), new TaskQueueConfig(null, null, null, null), new DefaultTaskConfig() {

        @Override
        public Map<String, Object> getContext() {
            return ImmutableMap.of(SinglePhaseParallelIndexTaskRunner.CTX_USE_LINEAGE_BASED_SEGMENT_ALLOCATION_KEY, false);
        }
    }, getTaskStorage(), new SimpleTaskRunner(actionClientFactory), actionClientFactory, getLockbox(), new NoopServiceEmitter());
    taskQueue.setActive(true);
    final Task task = new TestTask("t1", Intervals.of("2021-01-01/P1D"));
    taskQueue.add(task);
    final List<Task> tasks = taskQueue.getTasks();
    Assert.assertEquals(1, tasks.size());
    final Task queuedTask = tasks.get(0);
    Assert.assertFalse(queuedTask.getContextValue(SinglePhaseParallelIndexTaskRunner.CTX_USE_LINEAGE_BASED_SEGMENT_ALLOCATION_KEY));
}

Example 5 with TaskActionClientFactory

use of org.apache.druid.indexing.common.actions.TaskActionClientFactory in project druid by druid-io.

the class TaskQueueTest method testTaskStatusWhenExceptionIsThrownInIsReady.

@Test
public void testTaskStatusWhenExceptionIsThrownInIsReady() throws EntryExistsException {
    final TaskActionClientFactory actionClientFactory = createActionClientFactory();
    final TaskQueue taskQueue = new TaskQueue(new TaskLockConfig(), new TaskQueueConfig(null, null, null, null), new DefaultTaskConfig(), getTaskStorage(), new SimpleTaskRunner(actionClientFactory), actionClientFactory, getLockbox(), new NoopServiceEmitter());
    taskQueue.setActive(true);
    final Task task = new TestTask("t1", Intervals.of("2021-01-01/P1D")) {

        @Override
        public boolean isReady(TaskActionClient taskActionClient) {
            throw new RuntimeException("isReady failure test");
        }
    };
    taskQueue.add(task);
    taskQueue.manageInternal();
    Optional<TaskStatus> statusOptional = getTaskStorage().getStatus(task.getId());
    Assert.assertTrue(statusOptional.isPresent());
    Assert.assertEquals(TaskState.FAILED, statusOptional.get().getStatusCode());
    Assert.assertNotNull(statusOptional.get().getErrorMsg());
    Assert.assertTrue(StringUtils.format("Actual message is: %s", statusOptional.get().getErrorMsg()), statusOptional.get().getErrorMsg().startsWith("Failed while waiting for the task to be ready to run"));
}