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.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  
  import com.alibaba.jstorm.client.ConfigExtension;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import com.alibaba.jstorm.daemon.supervisor.SupervisorInfo;
  import com.alibaba.jstorm.schedule.Assignment;
  import com.alibaba.jstorm.schedule.IToplogyScheduler;
  import com.alibaba.jstorm.schedule.TopologyAssignContext;
  import com.alibaba.jstorm.utils.FailedAssignTopologyException;
  
  public class DefaultTopologyScheduler implements IToplogyScheduler {
      private static final Logger LOG = LoggerFactory.getLogger(DefaultTopologyScheduler.class);
  
      @SuppressWarnings("unused")
      private Map nimbusConf;
  
      @Override
      public void prepare(Map conf) {
          nimbusConf = conf;
      }
  
      /**
       * TODO: problematic: some dead slots have been freed
       *
       * @param context topology assign context
       */
      protected void freeUsed(TopologyAssignContext context) {
          Set<Integer> canFree = new HashSet<>();
          canFree.addAll(context.getAllTaskIds());
          canFree.removeAll(context.getUnstoppedTaskIds());
  
          Map<String, SupervisorInfo> cluster = context.getCluster();
          Assignment oldAssigns = context.getOldAssignment();
          for (Integer task : canFree) {
              ResourceWorkerSlot worker = oldAssigns.getWorkerByTaskId(task);
              if (worker == null) {
                  LOG.warn("No ResourceWorkerSlot of task " + task + " is found when freeing resource");
                  continue;
              }
  
              SupervisorInfo supervisorInfo = cluster.get(worker.getNodeId());
              if (supervisorInfo == null) {
                  continue;
              }
              supervisorInfo.getAvailableWorkerPorts().add(worker.getPort());
          }
      }
  
      private Set<Integer> getNeedAssignTasks(DefaultTopologyAssignContext context) {
          Set<Integer> needAssign = new HashSet<>();
  
          int assignType = context.getAssignType();
          if (assignType == TopologyAssignContext.ASSIGN_TYPE_NEW) {
              needAssign.addAll(context.getAllTaskIds());
          } else if (assignType == TopologyAssignContext.ASSIGN_TYPE_REBALANCE) {
              needAssign.addAll(context.getAllTaskIds());
              needAssign.removeAll(context.getUnstoppedTaskIds());
          } else { // ASSIGN_TYPE_MONITOR
              Set<Integer> deadTasks = context.getDeadTaskIds();
              needAssign.addAll(deadTasks);
          }
  
          return needAssign;
      }
  
      /**
       * Get the task Map which the task is alive and will be kept only when type is ASSIGN_TYPE_MONITOR, it is valid
       *
       * @param defaultContext default topology context
       * @param needAssigns    a set of tasks to be assigned
       * @return a set of assigned slots
       */
      public Set<ResourceWorkerSlot> getKeepAssign(DefaultTopologyAssignContext defaultContext, Set<Integer> needAssigns) {
          Set<Integer> keepAssignIds = new HashSet<>();
          keepAssignIds.addAll(defaultContext.getAllTaskIds());
          keepAssignIds.removeAll(defaultContext.getUnstoppedTaskIds());
          keepAssignIds.removeAll(needAssigns);
          Set<ResourceWorkerSlot> keeps = new HashSet<>();
          if (keepAssignIds.isEmpty()) {
              return keeps;
          }
  
          Assignment oldAssignment = defaultContext.getOldAssignment();
          if (oldAssignment == null) {
              return keeps;
          }
          keeps.addAll(defaultContext.getOldWorkers());
          for (ResourceWorkerSlot worker : defaultContext.getOldWorkers()) {
              for (Integer task : worker.getTasks()) {
                  if (!keepAssignIds.contains(task)) {
                      keeps.remove(worker);
                      break;
                  }
              }
          }
          return keeps;
      }
  
      @Override
      public Set<ResourceWorkerSlot> assignTasks(TopologyAssignContext context) throws FailedAssignTopologyException {
          int assignType = context.getAssignType();
          if (!TopologyAssignContext.isAssignTypeValid(assignType)) {
              throw new FailedAssignTopologyException("Invalid assign type " + assignType);
          }
  
          DefaultTopologyAssignContext defaultContext = new DefaultTopologyAssignContext(context);
          if (assignType == TopologyAssignContext.ASSIGN_TYPE_REBALANCE) {
              // Mark all current assigned worker as available. Current assignment will be restored in task scheduler.
              freeUsed(defaultContext);
          }
          LOG.info("Dead tasks:" + defaultContext.getDeadTaskIds());
          LOG.info("Unstopped tasks:" + defaultContext.getUnstoppedTaskIds());
  
          Set<Integer> needAssignTasks = getNeedAssignTasks(defaultContext);
  
          Set<ResourceWorkerSlot> keepAssigns = getKeepAssign(defaultContext, needAssignTasks);
  
          // todo: use tree map to sort tasks
          Set<ResourceWorkerSlot> ret = new HashSet<>();
          ret.addAll(keepAssigns);
          ret.addAll(defaultContext.getUnstoppedWorkers());
  
          int allocWorkerNum = defaultContext.getTotalWorkerNum() - defaultContext.getUnstoppedWorkerNum() - keepAssigns.size();
          LOG.info("allocWorkerNum=" + allocWorkerNum + ", totalWorkerNum=" +
                  defaultContext.getTotalWorkerNum() + ", keepWorkerNum=" + keepAssigns.size());
  
          if (allocWorkerNum <= 0) {
              LOG.warn("Don't need assign workers, all workers are fine " + defaultContext.toDetailString());
              throw new FailedAssignTopologyException("Don't need assign worker, all workers are fine ");
          }
  
          List<ResourceWorkerSlot> availableWorkers = WorkerScheduler.getInstance().getAvailableWorkers(
                  defaultContext, needAssignTasks, allocWorkerNum);
          TaskScheduler taskScheduler = new TaskScheduler(defaultContext, needAssignTasks, availableWorkers);
          Set<ResourceWorkerSlot> assignment = new HashSet<>(taskScheduler.assign());
  
          //setting worker's memory for TM
          int topologyMasterId = defaultContext.getTopologyMasterTaskId();
          Long tmWorkerMem = ConfigExtension.getMemSizePerTopologyMasterWorker(defaultContext.getStormConf());
          if (tmWorkerMem != null) {
              for (ResourceWorkerSlot resourceWorkerSlot : assignment) {
                  if (resourceWorkerSlot.getTasks().contains(topologyMasterId)) {
                      resourceWorkerSlot.setMemSize(tmWorkerMem);
                  }
              }
          }
          ret.addAll(assignment);
  
          LOG.info("Keep Alive slots:" + keepAssigns);
          LOG.info("Unstopped slots:" + defaultContext.getUnstoppedWorkers());
          LOG.info("New assign slots:" + assignment);
  
          return ret;
      }
  
  }
  