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

Example 1 with WorkerBehaviorConfig

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

the class OverlordResource method getTotalWorkerCapacity.

/**
 * Gets the total worker capacity of varies states of the cluster.
 */
@GET
@Path("/totalWorkerCapacity")
@Produces(MediaType.APPLICATION_JSON)
@ResourceFilters(ConfigResourceFilter.class)
public Response getTotalWorkerCapacity() {
    // Calculate current cluster capacity
    int currentCapacity;
    Optional<TaskRunner> taskRunnerOptional = taskMaster.getTaskRunner();
    if (!taskRunnerOptional.isPresent()) {
        // Cannot serve call as not leader
        return Response.status(Response.Status.SERVICE_UNAVAILABLE).build();
    }
    TaskRunner taskRunner = taskRunnerOptional.get();
    Collection<ImmutableWorkerInfo> workers;
    if (taskRunner instanceof WorkerTaskRunner) {
        workers = ((WorkerTaskRunner) taskRunner).getWorkers();
        currentCapacity = workers.stream().mapToInt(workerInfo -> workerInfo.getWorker().getCapacity()).sum();
    } else {
        log.debug("Cannot calculate capacity as task runner [%s] of type [%s] does not support listing workers", taskRunner, taskRunner.getClass().getName());
        workers = ImmutableList.of();
        currentCapacity = -1;
    }
    // Calculate maximum capacity with auto scale
    int maximumCapacity;
    if (workerConfigRef == null) {
        workerConfigRef = configManager.watch(WorkerBehaviorConfig.CONFIG_KEY, WorkerBehaviorConfig.class);
    }
    WorkerBehaviorConfig workerBehaviorConfig = workerConfigRef.get();
    if (workerBehaviorConfig == null) {
        // Auto scale not setup
        log.debug("Cannot calculate maximum worker capacity as worker behavior config is not configured");
        maximumCapacity = -1;
    } else if (workerBehaviorConfig instanceof DefaultWorkerBehaviorConfig) {
        DefaultWorkerBehaviorConfig defaultWorkerBehaviorConfig = (DefaultWorkerBehaviorConfig) workerBehaviorConfig;
        if (defaultWorkerBehaviorConfig.getAutoScaler() == null) {
            // Auto scale not setup
            log.debug("Cannot calculate maximum worker capacity as auto scaler not configured");
            maximumCapacity = -1;
        } else {
            int maxWorker = defaultWorkerBehaviorConfig.getAutoScaler().getMaxNumWorkers();
            int expectedWorkerCapacity = provisioningStrategy.getExpectedWorkerCapacity(workers);
            maximumCapacity = expectedWorkerCapacity == -1 ? -1 : maxWorker * expectedWorkerCapacity;
        }
    } else {
        // Auto scale is not using DefaultWorkerBehaviorConfig
        log.debug("Cannot calculate maximum worker capacity as WorkerBehaviorConfig [%s] of type [%s] does not support getting max capacity", workerBehaviorConfig, workerBehaviorConfig.getClass().getSimpleName());
        maximumCapacity = -1;
    }
    return Response.ok(new TotalWorkerCapacityResponse(currentCapacity, maximumCapacity)).build();
}

Example 2 with WorkerBehaviorConfig

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

the class HttpRemoteTaskRunner method findWorkerToRunTask.

private ImmutableWorkerInfo findWorkerToRunTask(Task task) {
    WorkerBehaviorConfig workerConfig = workerConfigRef.get();
    WorkerSelectStrategy strategy;
    if (workerConfig == null || workerConfig.getSelectStrategy() == null) {
        strategy = WorkerBehaviorConfig.DEFAULT_STRATEGY;
        log.debug("No worker selection strategy set. Using default of [%s]", strategy.getClass().getSimpleName());
    } else {
        strategy = workerConfig.getSelectStrategy();
    }
    return strategy.findWorkerForTask(config, ImmutableMap.copyOf(getWorkersEligibleToRunTasks()), task);
}

Example 3 with WorkerBehaviorConfig

use of org.apache.druid.indexing.overlord.setup.WorkerBehaviorConfig 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) {
            strategy = WorkerBehaviorConfig.DEFAULT_STRATEGY;
            log.debug("No worker selection strategy set. Using default of [%s]", strategy.getClass().getSimpleName());
        } else {
            strategy = workerConfig.getSelectStrategy();
        }
        ZkWorker assignedWorker = null;
        final ImmutableWorkerInfo immutableZkWorker;
        try {
            synchronized (workersWithUnacknowledgedTask) {
                immutableZkWorker = strategy.findWorkerForTask(config, ImmutableMap.copyOf(getWorkersEligibleToRunTasks()), task);
                if (immutableZkWorker != null && workersWithUnacknowledgedTask.putIfAbsent(immutableZkWorker.getWorker().getHost(), task.getId()) == null) {
                    assignedWorker = zkWorkers.get(immutableZkWorker.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 4 with WorkerBehaviorConfig

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

the class PendingTaskBasedWorkerProvisioningStrategy method getDefaultWorkerBehaviorConfig.

@VisibleForTesting
@Nullable
public static DefaultWorkerBehaviorConfig getDefaultWorkerBehaviorConfig(Supplier<WorkerBehaviorConfig> workerConfigRef, SimpleWorkerProvisioningConfig config, String action, EmittingLogger log) {
    final WorkerBehaviorConfig workerBehaviorConfig = workerConfigRef.get();
    if (workerBehaviorConfig == null) {
        log.error("No workerConfig available, cannot %s workers.", action);
        return null;
    }
    if (!(workerBehaviorConfig instanceof DefaultWorkerBehaviorConfig)) {
        log.error("Only DefaultWorkerBehaviorConfig is supported as WorkerBehaviorConfig, [%s] given, cannot %s workers", workerBehaviorConfig, action);
        return null;
    }
    final DefaultWorkerBehaviorConfig workerConfig = (DefaultWorkerBehaviorConfig) workerBehaviorConfig;
    if (workerConfig.getAutoScaler() == null) {
        log.error("No autoScaler available, cannot %s workers", action);
        return null;
    }
    if (config instanceof PendingTaskBasedWorkerProvisioningConfig && workerConfig.getAutoScaler().getMinNumWorkers() == 0 && ((PendingTaskBasedWorkerProvisioningConfig) config).getWorkerCapacityHint() <= 0) {
        log.error(ERROR_MESSAGE_MIN_WORKER_ZERO_HINT_UNSET, ((PendingTaskBasedWorkerProvisioningConfig) config).getWorkerCapacityHint());
        return null;
    }
    return workerConfig;
}

Example 5 with WorkerBehaviorConfig

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

the class OverlordResourceTest method testGetTotalWorkerCapacityWithUnknown.

@Test
public void testGetTotalWorkerCapacityWithUnknown() {
    WorkerBehaviorConfig workerBehaviorConfig = EasyMock.createMock(WorkerBehaviorConfig.class);
    AtomicReference<WorkerBehaviorConfig> workerBehaviorConfigAtomicReference = new AtomicReference<>(workerBehaviorConfig);
    EasyMock.expect(configManager.watch(WorkerBehaviorConfig.CONFIG_KEY, WorkerBehaviorConfig.class)).andReturn(workerBehaviorConfigAtomicReference);
    EasyMock.replay(taskRunner, taskMaster, taskStorageQueryAdapter, indexerMetadataStorageAdapter, req, workerTaskRunnerQueryAdapter, configManager);
    final Response response = overlordResource.getTotalWorkerCapacity();
    Assert.assertEquals(HttpResponseStatus.OK.getCode(), response.getStatus());
    Assert.assertEquals(-1, ((TotalWorkerCapacityResponse) response.getEntity()).getCurrentClusterCapacity());
    Assert.assertEquals(-1, ((TotalWorkerCapacityResponse) response.getEntity()).getMaximumCapacityWithAutoScale());
}