Examples with WorkerSelectStrategy io.druid.indexing.overlord.setup.WorkerSelectStrategy used on opensource projects

Example 1 with WorkerSelectStrategy

use of io.druid.indexing.overlord.setup.WorkerSelectStrategy in project druid by druid-io.

the class RemoteTaskRunner method tryAssignTask.

/**
   * Ensures no workers are already running a task before assigning the task to a worker.
   * It is possible that a worker is running a task that the RTR has no knowledge of. This occurs when the RTR
   * needs to bootstrap after a restart.
   *
   * @param taskRunnerWorkItem - the task to assign
   *
   * @return true iff the task is now assigned
   */
private boolean tryAssignTask(final Task task, final RemoteTaskRunnerWorkItem taskRunnerWorkItem) throws Exception {
    Preconditions.checkNotNull(task, "task");
    Preconditions.checkNotNull(taskRunnerWorkItem, "taskRunnerWorkItem");
    Preconditions.checkArgument(task.getId().equals(taskRunnerWorkItem.getTaskId()), "task id != workItem id");
    if (runningTasks.containsKey(task.getId()) || findWorkerRunningTask(task.getId()) != null) {
        log.info("Task[%s] already running.", task.getId());
        return true;
    } else {
        // Nothing running this task, announce it in ZK for a worker to run it
        WorkerBehaviorConfig workerConfig = workerConfigRef.get();
        WorkerSelectStrategy strategy;
        if (workerConfig == null || workerConfig.getSelectStrategy() == null) {
            log.warn("No worker selections strategy set. Using default.");
            strategy = WorkerBehaviorConfig.DEFAULT_STRATEGY;
        } else {
            strategy = workerConfig.getSelectStrategy();
        }
        ZkWorker assignedWorker = null;
        Optional<ImmutableWorkerInfo> immutableZkWorker = null;
        try {
            synchronized (workersWithUnacknowledgedTask) {
                immutableZkWorker = strategy.findWorkerForTask(config, ImmutableMap.copyOf(Maps.transformEntries(Maps.filterEntries(zkWorkers, new Predicate<Map.Entry<String, ZkWorker>>() {

                    @Override
                    public boolean apply(Map.Entry<String, ZkWorker> input) {
                        return !lazyWorkers.containsKey(input.getKey()) && !workersWithUnacknowledgedTask.containsKey(input.getKey()) && !blackListedWorkers.contains(input.getValue());
                    }
                }), new Maps.EntryTransformer<String, ZkWorker, ImmutableWorkerInfo>() {

                    @Override
                    public ImmutableWorkerInfo transformEntry(String key, ZkWorker value) {
                        return value.toImmutable();
                    }
                })), task);
                if (immutableZkWorker.isPresent() && workersWithUnacknowledgedTask.putIfAbsent(immutableZkWorker.get().getWorker().getHost(), task.getId()) == null) {
                    assignedWorker = zkWorkers.get(immutableZkWorker.get().getWorker().getHost());
                }
            }
            if (assignedWorker != null) {
                return announceTask(task, assignedWorker, taskRunnerWorkItem);
            } else {
                log.debug("Unsuccessful task-assign attempt for task [%s] on workers [%s]. Workers to ack tasks are [%s].", task.getId(), zkWorkers.values(), workersWithUnacknowledgedTask);
            }
            return false;
        } finally {
            if (assignedWorker != null) {
                workersWithUnacknowledgedTask.remove(assignedWorker.getWorker().getHost());
                //if this attempt won the race to run the task then other task might be able to use this worker now after task ack.
                runPendingTasks();
            }
        }
    }
}

Example 2 with WorkerSelectStrategy

use of io.druid.indexing.overlord.setup.WorkerSelectStrategy in project druid by druid-io.

the class PendingTaskBasedWorkerResourceManagementStrategy method getWorkersNeededToAssignTasks.

int getWorkersNeededToAssignTasks(final WorkerTaskRunnerConfig workerTaskRunnerConfig, final WorkerBehaviorConfig workerConfig, final Collection<Task> pendingTasks, final Collection<ImmutableWorkerInfo> workers) {
    final Collection<ImmutableWorkerInfo> validWorkers = Collections2.filter(workers, ResourceManagementUtil.createValidWorkerPredicate(config));
    Map<String, ImmutableWorkerInfo> workersMap = Maps.newHashMap();
    for (ImmutableWorkerInfo worker : validWorkers) {
        workersMap.put(worker.getWorker().getHost(), worker);
    }
    WorkerSelectStrategy workerSelectStrategy = workerConfig.getSelectStrategy();
    int need = 0;
    int capacity = getExpectedWorkerCapacity(workers);
    // the number of additional workers needed to assign all the pending tasks is noted
    for (Task task : pendingTasks) {
        Optional<ImmutableWorkerInfo> selectedWorker = workerSelectStrategy.findWorkerForTask(workerTaskRunnerConfig, ImmutableMap.copyOf(workersMap), task);
        final ImmutableWorkerInfo workerRunningTask;
        if (selectedWorker.isPresent()) {
            workerRunningTask = selectedWorker.get();
        } else {
            // None of the existing worker can run this task, we need to provision one worker for it.
            // create a dummy worker and try to simulate assigning task to it.
            workerRunningTask = createDummyWorker("dummy" + need, capacity, workerTaskRunnerConfig.getMinWorkerVersion());
            need++;
        }
        // Update map with worker running task
        workersMap.put(workerRunningTask.getWorker().getHost(), workerWithTask(workerRunningTask, task));
    }
    return need;
}