content: /**
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   * http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package com.alibaba.jstorm.schedule.default_assign;
  
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Set;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import com.alibaba.jstorm.client.ConfigExtension;
  import com.alibaba.jstorm.cluster.Common;
  import com.alibaba.jstorm.schedule.TopologyAssignContext;
  import com.alibaba.jstorm.schedule.default_assign.Selector.ComponentNumSelector;
  import com.alibaba.jstorm.schedule.default_assign.Selector.InputComponentNumSelector;
  import com.alibaba.jstorm.schedule.default_assign.Selector.Selector;
  import com.alibaba.jstorm.schedule.default_assign.Selector.TotalTaskNumSelector;
  import com.alibaba.jstorm.utils.FailedAssignTopologyException;
  
  public class TaskScheduler {
      public static Logger LOG = LoggerFactory.getLogger(TaskScheduler.class);
  
      private final TaskAssignContext taskContext;
      private List<ResourceWorkerSlot> assignments = new ArrayList<>();
      private int workerNum;
  
      /**
       * For balance purpose, default scheduler is trying to assign the same number of tasks to a worker.
       * e.g. There are 4 tasks and 3 available workers.
       * Each worker will be assigned one task first. And then one worker is chosen for the last one.
       */
      private int avgTaskNum;
      private int leftTaskNum;
  
      private Set<Integer> tasks;
  
      private DefaultTopologyAssignContext context;
      private Selector componentSelector;
      private Selector inputComponentSelector;
      private Selector totalTaskNumSelector;
  
      public TaskScheduler(DefaultTopologyAssignContext context, Set<Integer> tasks, List<ResourceWorkerSlot> workers) {
          this.tasks = tasks;
          LOG.info("Tasks " + tasks + " is going to be assigned in workers " + workers);
          this.context = context;
          this.taskContext = new TaskAssignContext(this.buildSupervisorToWorker(workers),
                  Common.buildSpoutOutoputAndBoltInputMap(context), context.getTaskToComponent());
          this.componentSelector = new ComponentNumSelector(taskContext);
          this.inputComponentSelector = new InputComponentNumSelector(taskContext);
          this.totalTaskNumSelector = new TotalTaskNumSelector(taskContext);
          if (tasks.size() == 0)
              return;
          if (context.getAssignType() != TopologyAssignContext.ASSIGN_TYPE_REBALANCE || context.isReassign()) {
              // warning ! it doesn't consider HA TM now!!
              if (context.getAssignSingleWorkerForTM() && tasks.contains(context.getTopologyMasterTaskId())) {
                  assignForTopologyMaster();
              }
          }
  
          int taskNum = tasks.size();
          Map<ResourceWorkerSlot, Integer> workerSlotIntegerMap = taskContext.getWorkerToTaskNum();
          Set<ResourceWorkerSlot> preAssignWorkers = new HashSet<>();
          for (Entry<ResourceWorkerSlot, Integer> worker : workerSlotIntegerMap.entrySet()) {
              if (worker.getValue() > 0) {
                  taskNum += worker.getValue();
                  preAssignWorkers.add(worker.getKey());
              }
          }
          setTaskNum(taskNum, workerNum);
  
          // Check the worker assignment status of pre-assigned workers, e.g user defined or old assignment workers.
          // Remove the workers which have been assigned with enough workers.
          for (ResourceWorkerSlot worker : preAssignWorkers) {
              if (taskContext.getWorkerToTaskNum().keySet().contains(worker)) {
  
                  Set<ResourceWorkerSlot> doneWorkers = removeWorkerFromSrcPool(taskContext.getWorkerToTaskNum().get(worker), worker);
                  if (doneWorkers != null) {
                      for (ResourceWorkerSlot doneWorker : doneWorkers) {
                          taskNum -= doneWorker.getTasks().size();
                          workerNum--;
                      }
                  }
              }
          }
          setTaskNum(taskNum, workerNum);
  
          // For Scale-out case, the old assignment should be kept.
          if (context.getAssignType() == TopologyAssignContext.ASSIGN_TYPE_REBALANCE && !context.isReassign()) {
              keepAssignment(taskNum, context.getOldAssignment().getWorkers());
          }
      }
  
      private void keepAssignment(int taskNum, Set<ResourceWorkerSlot> keepAssignments) {
          Set<Integer> keepTasks = new HashSet<>();
          ResourceWorkerSlot tmWorker = null;
          for (ResourceWorkerSlot worker : keepAssignments) {
              if (worker.getTasks().contains(context.getTopologyMasterTaskId()))
                  tmWorker = worker;
              for (Integer taskId : worker.getTasks()) {
                  if (tasks.contains(taskId)) {
                      ResourceWorkerSlot contextWorker = taskContext.getWorker(worker);
                      if (contextWorker != null) {
                          if (tmWorker != null && tmWorker.getTasks().contains(taskId) && context.getAssignSingleWorkerForTM()) {
                              if (context.getTopologyMasterTaskId() == taskId) {
                                  updateAssignedTasksOfWorker(taskId, contextWorker);
                                  taskContext.getWorkerToTaskNum().remove(contextWorker);
                                  contextWorker.getTasks().clear();
                                  contextWorker.getTasks().add(taskId);
                                  assignments.add(contextWorker);
                                  tasks.remove(taskId);
                                  taskNum--;
                                  workerNum--;
                                  LOG.info("assign for TopologyMaster: " + contextWorker);
                              }
                          } else {
                              String componentName = context.getTaskToComponent().get(taskId);
                              updateAssignedTasksOfWorker(taskId, contextWorker);
                              updateComponentsNumOfWorker(componentName, contextWorker);
                              keepTasks.add(taskId);
                          }
                      }
                  }
              }
          }
          if (tmWorker != null) {
              setTaskNum(taskNum, workerNum);
              keepAssignments.remove(tmWorker);
          }
  
          // Try to find the workers which have been assigned too much tasks
          // If found, remove the workers from worker resource pool and update
          // the avgNum and leftNum
          int doneAssignedTaskNum = 0;
          while (true) {
              boolean found = false;
              Set<ResourceWorkerSlot> doneAssignedWorkers = new HashSet<>();
              for (ResourceWorkerSlot worker : keepAssignments) {
                  ResourceWorkerSlot contextWorker = taskContext.getWorker(worker);
                  if (contextWorker != null && isTaskFullForWorker(contextWorker)) {
                      found = true;
                      workerNum--;
                      taskContext.getWorkerToTaskNum().remove(contextWorker);
                      assignments.add(contextWorker);
  
                      doneAssignedWorkers.add(worker);
                      doneAssignedTaskNum += contextWorker.getTasks().size();
                  }
              }
  
              if (found) {
                  taskNum -= doneAssignedTaskNum;
                  setTaskNum(taskNum, workerNum);
                  keepAssignments.removeAll(doneAssignedWorkers);
              } else {
                  break;
              }
          }
          tasks.removeAll(keepTasks);
          LOG.info("keep following assignment, " + assignments);
      }
  
      private boolean isTaskFullForWorker(ResourceWorkerSlot worker) {
          boolean ret = false;
          Set<Integer> tasks = worker.getTasks();
  
          if (tasks != null) {
              if ((leftTaskNum <= 0 && tasks.size() >= avgTaskNum) ||
                      (leftTaskNum > 0 && tasks.size() >= (avgTaskNum + 1))) {
                  ret = true;
              }
          }
          return ret;
      }
  
      private Set<ResourceWorkerSlot> getRestAssignedWorkers() {
          Set<ResourceWorkerSlot> ret = new HashSet<>();
          for (ResourceWorkerSlot worker : taskContext.getWorkerToTaskNum().keySet()) {
              if (worker.getTasks() != null && worker.getTasks().size() > 0) {
                  ret.add(worker);
              }
          }
          return ret;
      }
  
      public List<ResourceWorkerSlot> assign() {
          if (tasks.size() == 0) {
              assignments.addAll(getRestAssignedWorkers());
              return assignments;
          }
  
          // Firstly, assign workers to the components which are configured "task.on.differ.node"
          Set<Integer> assignedTasks = assignForDifferNodeTask();
  
          // Assign for the tasks except system tasks
          tasks.removeAll(assignedTasks);
          Map<Integer, String> systemTasks = new HashMap<>();
          for (Integer task : tasks) {
              String name = context.getTaskToComponent().get(task);
              if (Common.isSystemComponent(name)) {
                  systemTasks.put(task, name);
                  continue;
              }
              assignForTask(name, task);
          }
          
          /*
           * At last, make the assignment for system component, e.g. acker, topology master...
           */
          for (Entry<Integer, String> entry : systemTasks.entrySet()) {
              assignForTask(entry.getValue(), entry.getKey());
          }
  
          assignments.addAll(getRestAssignedWorkers());
          return assignments;
      }
  
      private void assignForTopologyMaster() {
          int taskId = context.getTopologyMasterTaskId();
  
          // Try to find a worker which is in a supervisor with most workers,
          // to avoid the balance problem when the assignment for other workers.
          ResourceWorkerSlot workerAssigned = null;
          int workerNumOfSuperv = 0;
          for (ResourceWorkerSlot workerSlot : taskContext.getWorkerToTaskNum().keySet()) {
              List<ResourceWorkerSlot> workers = taskContext.getSupervisorToWorker().get(workerSlot.getNodeId());
              if (workers != null && workers.size() > workerNumOfSuperv) {
                  for (ResourceWorkerSlot worker : workers) {
                      Set<Integer> tasks = worker.getTasks();
                      if (tasks == null || tasks.size() == 0) {
                          workerAssigned = worker;
                          workerNumOfSuperv = workers.size();
                          break;
                      }
                  }
              }
          }
  
          if (workerAssigned == null) {
              throw new FailedAssignTopologyException("there's no enough workers for the assignment of topology master");
          }
          updateAssignedTasksOfWorker(taskId, workerAssigned);
          taskContext.getWorkerToTaskNum().remove(workerAssigned);
          assignments.add(workerAssigned);
          tasks.remove(taskId);
          workerNum--;
          LOG.info("assignForTopologyMaster, assignments=" + assignments);
      }
  
      private void assignForTask(String name, Integer task) {
          ResourceWorkerSlot worker = chooseWorker(name, new ArrayList<>(taskContext.getWorkerToTaskNum().keySet()));
          pushTaskToWorker(task, name, worker);
      }
  
      private Set<Integer> assignForDifferNodeTask() {
          Set<Integer> ret = new HashSet<>();
          for (Integer task : tasks) {
              Map conf = Common.getComponentMap(context, task);
              if (ConfigExtension.isTaskOnDifferentNode(conf))
                  ret.add(task);
          }
          for (Integer task : ret) {
              String name = context.getTaskToComponent().get(task);
              ResourceWorkerSlot worker = chooseWorker(name, getDifferNodeTaskWokers(name));
              LOG.info("Due to task.on.differ.node, push task-{} to {}:{}", task, worker.getHostname(), worker.getPort());
              pushTaskToWorker(task, name, worker);
          }
          return ret;
      }
  
      private Map<String, List<ResourceWorkerSlot>> buildSupervisorToWorker(List<ResourceWorkerSlot> workers) {
          Map<String, List<ResourceWorkerSlot>> supervisorToWorker = new HashMap<>();
          for (ResourceWorkerSlot worker : workers) {
              List<ResourceWorkerSlot> supervisor = supervisorToWorker.get(worker.getNodeId());
              if (supervisor == null) {
                  supervisor = new ArrayList<>();
                  supervisorToWorker.put(worker.getNodeId(), supervisor);
              }
              supervisor.add(worker);
          }
          this.workerNum = workers.size();
          return supervisorToWorker;
      }
  
      private ResourceWorkerSlot chooseWorker(String name, List<ResourceWorkerSlot> workers) {
          List<ResourceWorkerSlot> result = componentSelector.select(workers, name);
          result = totalTaskNumSelector.select(result, name);
          if (Common.isSystemComponent(name))
              return result.iterator().next();
          result = inputComponentSelector.select(result, name);
          return result.iterator().next();
      }
  
      private void pushTaskToWorker(Integer task, String name, ResourceWorkerSlot worker) {
          LOG.debug("Push task-" + task + " to worker-" + worker.getPort());
          int taskNum = updateAssignedTasksOfWorker(task, worker);
          removeWorkerFromSrcPool(taskNum, worker);
          updateComponentsNumOfWorker(name, worker);
      }
  
      private int updateAssignedTasksOfWorker(Integer task, ResourceWorkerSlot worker) {
          int ret;
          Set<Integer> tasks = worker.getTasks();
          if (tasks == null) {
              tasks = new HashSet<>();
              worker.setTasks(tasks);
          }
          tasks.add(task);
  
          ret = taskContext.getWorkerToTaskNum().get(worker);
          taskContext.getWorkerToTaskNum().put(worker, ++ret);
          return ret;
      }
  
      /*
       * Remove the worker from source worker pool, if the worker is assigned with enough tasks
       */
      private Set<ResourceWorkerSlot> removeWorkerFromSrcPool(int taskNum, ResourceWorkerSlot worker) {
          Set<ResourceWorkerSlot> ret = new HashSet<>();
  
          if (leftTaskNum <= 0) {
              if (taskNum >= avgTaskNum) {
                  taskContext.getWorkerToTaskNum().remove(worker);
                  assignments.add(worker);
                  ret.add(worker);
              }
          } else {
              if (taskNum > avgTaskNum) {
                  taskContext.getWorkerToTaskNum().remove(worker);
                  leftTaskNum = leftTaskNum - (taskNum - avgTaskNum);
                  assignments.add(worker);
                  ret.add(worker);
              }
              if (leftTaskNum <= 0) {
                  List<ResourceWorkerSlot> needDelete = new ArrayList<>();
                  for (Entry<ResourceWorkerSlot, Integer> entry : taskContext.getWorkerToTaskNum().entrySet()) {
                      if (avgTaskNum != 0 && entry.getValue() == avgTaskNum)
                          needDelete.add(entry.getKey());
                  }
                  for (ResourceWorkerSlot workerToDelete : needDelete) {
                      taskContext.getWorkerToTaskNum().remove(workerToDelete);
                      assignments.add(workerToDelete);
                      ret.add(workerToDelete);
                  }
              }
          }
  
          return ret;
      }
  
      private void updateComponentsNumOfWorker(String name, ResourceWorkerSlot worker) {
          Map<String, Integer> components = taskContext.getWorkerToComponentNum().get(worker);
          if (components == null) {
              components = new HashMap<>();
              taskContext.getWorkerToComponentNum().put(worker, components);
          }
          Integer componentNum = components.get(name);
          if (componentNum == null) {
              componentNum = 0;
          }
          components.put(name, ++componentNum);
      }
  
      private void setTaskNum(int taskNum, int workerNum) {
          if (taskNum >= 0 && workerNum > 0) {
              this.avgTaskNum = taskNum / workerNum;
              this.leftTaskNum = taskNum % workerNum;
              LOG.debug("avgTaskNum=" + avgTaskNum + ", leftTaskNum=" + leftTaskNum);
          } else {
              LOG.warn("Illegal parameters, taskNum=" + taskNum + ", workerNum=" + workerNum);
          }
      }
  
      private List<ResourceWorkerSlot> getDifferNodeTaskWokers(String name) {
          List<ResourceWorkerSlot> workers = new ArrayList<>();
          workers.addAll(taskContext.getWorkerToTaskNum().keySet());
  
          for (Entry<String, List<ResourceWorkerSlot>> entry : taskContext.getSupervisorToWorker().entrySet()) {
              if (taskContext.getComponentNumOnSupervisor(entry.getKey(), name) != 0)
                  workers.removeAll(entry.getValue());
          }
          if (workers.size() == 0) {
              throw new FailedAssignTopologyException("there's no enough supervisor for making component: " +
                      name + " 's tasks on different node");
          }
          return workers;
      }
  }
  