| package org.apache.helix.task; |
| |
| import org.apache.helix.controller.stages.CurrentStateOutput; |
| import org.apache.helix.model.IdealState; |
| import org.apache.helix.model.ResourceAssignment; |
| |
| import java.util.Collection; |
| import java.util.Map; |
| import java.util.Set; |
| import java.util.SortedSet; |
| import org.apache.helix.task.assigner.AssignableInstance; |
| |
| |
| public abstract class TaskAssignmentCalculator { |
| /** |
| * Get all the partitions/tasks that belong to this job. |
| * |
| * @param jobCfg the task configuration |
| * @param jobCtx the task context |
| * @param workflowCfg the workflow configuration |
| * @param workflowCtx the workflow context |
| * @param idealStateMap the map of resource name map to ideal state |
| * @return set of partition numbers |
| */ |
| public abstract Set<Integer> getAllTaskPartitions(JobConfig jobCfg, JobContext jobCtx, |
| WorkflowConfig workflowCfg, WorkflowContext workflowCtx, |
| Map<String, IdealState> idealStateMap); |
| |
| /** |
| * Compute an assignment of tasks to instances |
| * |
| * @param currStateOutput the current state of the instances |
| * @param prevAssignment the previous task partition assignment |
| * @param instances the instances |
| * @param jobCfg the task configuration |
| * @param jobContext the task context |
| * @param workflowCfg the workflow configuration |
| * @param workflowCtx the workflow context |
| * @param partitionSet the partitions to assign |
| * @param idealStateMap the map of resource name map to ideal state |
| * @return map of instances to set of partition numbers |
| */ |
| public abstract Map<String, SortedSet<Integer>> getTaskAssignment( |
| CurrentStateOutput currStateOutput, ResourceAssignment prevAssignment, |
| Collection<String> instances, JobConfig jobCfg, JobContext jobContext, |
| WorkflowConfig workflowCfg, WorkflowContext workflowCtx, Set<Integer> partitionSet, |
| Map<String, IdealState> idealStateMap); |
| |
| /** |
| * Returns the correct type for this job. Note that if the parent workflow has a type, then all of |
| * its jobs will inherit the type from the workflow. |
| * @param workflowConfig |
| * @param jobConfig |
| * @return |
| */ |
| public static String getQuotaType(WorkflowConfig workflowConfig, JobConfig jobConfig) { |
| String workflowType = workflowConfig.getWorkflowType(); |
| if (workflowType == null || workflowType.equals("")) { |
| // Workflow type is null, so we go by the job type |
| String jobType = jobConfig.getJobType(); |
| if (jobType == null || jobType.equals("")) { |
| // Job type is null, so we use DEFAULT |
| return AssignableInstance.DEFAULT_QUOTA_TYPE; |
| } |
| return jobType; |
| } |
| return workflowType; |
| } |
| } |