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

Example 1 with SurrogateAction

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

the class PartialSegmentMergeTask method runTask.

@Override
public TaskStatus runTask(TaskToolbox toolbox) throws Exception {
    // Group partitionLocations by interval and partitionId
    final Map<Interval, Int2ObjectMap<List<PartitionLocation>>> intervalToBuckets = new HashMap<>();
    for (PartitionLocation location : ioConfig.getPartitionLocations()) {
        intervalToBuckets.computeIfAbsent(location.getInterval(), k -> new Int2ObjectOpenHashMap<>()).computeIfAbsent(location.getBucketId(), k -> new ArrayList<>()).add(location);
    }
    final List<TaskLock> locks = toolbox.getTaskActionClient().submit(new SurrogateAction<>(supervisorTaskId, new LockListAction()));
    final Map<Interval, String> intervalToVersion = Maps.newHashMapWithExpectedSize(locks.size());
    locks.forEach(lock -> {
        if (lock.isRevoked()) {
            throw new ISE("Lock[%s] is revoked", lock);
        }
        final String mustBeNull = intervalToVersion.put(lock.getInterval(), lock.getVersion());
        if (mustBeNull != null) {
            throw new ISE("Unexpected state: Two versions([%s], [%s]) for the same interval[%s]", lock.getVersion(), mustBeNull, lock.getInterval());
        }
    });
    final Stopwatch fetchStopwatch = Stopwatch.createStarted();
    final Map<Interval, Int2ObjectMap<List<File>>> intervalToUnzippedFiles = fetchSegmentFiles(toolbox, intervalToBuckets);
    final long fetchTime = fetchStopwatch.elapsed(TimeUnit.SECONDS);
    fetchStopwatch.stop();
    LOG.info("Fetch took [%s] seconds", fetchTime);
    final ParallelIndexSupervisorTaskClient taskClient = toolbox.getSupervisorTaskClientFactory().build(new ClientBasedTaskInfoProvider(toolbox.getIndexingServiceClient()), getId(), // always use a single http thread
    1, getTuningConfig().getChatHandlerTimeout(), getTuningConfig().getChatHandlerNumRetries());
    final File persistDir = toolbox.getPersistDir();
    org.apache.commons.io.FileUtils.deleteQuietly(persistDir);
    FileUtils.mkdirp(persistDir);
    final Set<DataSegment> pushedSegments = mergeAndPushSegments(toolbox, getDataSchema(), getTuningConfig(), persistDir, intervalToVersion, intervalToUnzippedFiles);
    taskClient.report(supervisorTaskId, new PushedSegmentsReport(getId(), Collections.emptySet(), pushedSegments, ImmutableMap.of()));
    return TaskStatus.success(getId());
}

Example 2 with SurrogateAction

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

the class CountingLocalTaskActionClientForTest method submit.

@Override
public <RetType> RetType submit(TaskAction<RetType> taskAction) {
    final RetType result = delegate.submit(taskAction);
    final TaskAction actionKey;
    if (taskAction instanceof SurrogateAction) {
        actionKey = ((SurrogateAction) taskAction).getTaskAction();
    } else {
        actionKey = taskAction;
    }
    actionCountMap.computeIfAbsent(actionKey.getClass(), k -> new AtomicInteger()).incrementAndGet();
    return result;
}