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.callback.impl;
  
  import backtype.storm.Config;
  import backtype.storm.generated.Bolt;
  import backtype.storm.generated.SpoutSpec;
  import backtype.storm.generated.StormTopology;
  import backtype.storm.utils.Utils;
  import com.alibaba.jstorm.callback.BaseCallback;
  import com.alibaba.jstorm.cluster.Common;
  import com.alibaba.jstorm.cluster.StormClusterState;
  import com.alibaba.jstorm.cluster.StormConfig;
  import com.alibaba.jstorm.cluster.StormStatus;
  import com.alibaba.jstorm.daemon.nimbus.*;
  import com.alibaba.jstorm.task.TaskInfo;
  import com.alibaba.jstorm.task.TkHbCacheTime;
  import com.alibaba.jstorm.utils.JStormUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.*;
  import java.util.Map.Entry;
  
  /**
   * Do real rebalance action.
   *
   * After nimbus receives a rebalance command, it will do the following:
   * 1. set topology status to rebalancing
   * 2. delay 2 * timeout seconds
   * 3. execute this callback
   *
   * @author Xin.Li/Longda
   */
  public class DoRebalanceTransitionCallback extends BaseCallback {
      private static final Logger LOG = LoggerFactory.getLogger(DoRebalanceTransitionCallback.class);
  
      private NimbusData data;
      private String topologyId;
      private StormStatus oldStatus;
      private Set<Integer> newTasks;
  
      public DoRebalanceTransitionCallback(NimbusData data, String topologyId, StormStatus status) {
          this.data = data;
          this.topologyId = topologyId;
          this.oldStatus = status;
          this.newTasks = new HashSet<>();
      }
  
      @Override
      public <T> Object execute(T... args) {
          boolean isSetTaskInfo = false;
          try {
              Boolean reassign = (Boolean) args[1];
              // args[0]: delay, args[1]: reassign_flag, args[2]: conf
              Map<Object, Object> conf = (Map<Object, Object>) args[2];
              if (conf != null) {
                  boolean isConfUpdate = false;
  
                  // Update topology code
                  Map topoConf = StormConfig.read_nimbus_topology_conf(topologyId, data.getBlobStore());
                  StormTopology rawOldTopology = StormConfig.read_nimbus_topology_code(topologyId, data.getBlobStore());
                  StormTopology rawNewTopology = NimbusUtils.normalizeTopology(conf, rawOldTopology, true);
                  StormTopology sysOldTopology = rawOldTopology.deepCopy();
                  StormTopology sysNewTopology = rawNewTopology.deepCopy();
                  if (conf.get(Config.TOPOLOGY_ACKER_EXECUTORS) != null) {
                      Common.add_acker(topoConf, sysOldTopology);
                      Common.add_acker(conf, sysNewTopology);
                      int ackerNum = JStormUtils.parseInt(conf.get(Config.TOPOLOGY_ACKER_EXECUTORS));
                      int oldAckerNum = JStormUtils.parseInt(topoConf.get(Config.TOPOLOGY_ACKER_EXECUTORS));
                      LOG.info("Update acker from oldAckerNum=" + oldAckerNum + " to ackerNum=" + ackerNum);
                      topoConf.put(Config.TOPOLOGY_ACKER_EXECUTORS, ackerNum);
                      isConfUpdate = true;
                  }
  
                  // If scale-out, setup task info for new added tasks
                  setTaskInfo(sysOldTopology, sysNewTopology);
                  isSetTaskInfo = true;
  
                  // If everything is OK, write topology code into disk
                  StormConfig.write_nimbus_topology_code(topologyId, Utils.serialize(rawNewTopology), data);
  
                  // Update topology conf if worker num has been updated
                  Set<Object> keys = conf.keySet();
                  Integer workerNum = JStormUtils.parseInt(conf.get(Config.TOPOLOGY_WORKERS));
                  if (workerNum != null) {
                      Integer oldWorkerNum = JStormUtils.parseInt(topoConf.get(Config.TOPOLOGY_WORKERS));
                      topoConf.put(Config.TOPOLOGY_WORKERS, workerNum);
                      isConfUpdate = true;
  
                      LOG.info("Update worker num from " + oldWorkerNum + " to " + workerNum);
                  }
  
                  if (keys.contains(Config.ISOLATION_SCHEDULER_MACHINES)) {
                      topoConf.put(Config.ISOLATION_SCHEDULER_MACHINES, conf.get(Config.ISOLATION_SCHEDULER_MACHINES));
                  }
  
                  if (isConfUpdate) {
                      StormConfig.write_nimbus_topology_conf(topologyId, topoConf, data);
                  }
              }
  
              TopologyAssignEvent event = new TopologyAssignEvent();
  
              event.setTopologyId(topologyId);
              event.setScratch(true);
              event.setOldStatus(oldStatus);
              event.setReassign(reassign);
              if (conf != null)
                  event.setScaleTopology(true);
              TopologyAssign.push(event);
              event.waitFinish();
          } catch (Exception e) {
              LOG.error("do-rebalance error!", e);
              // Rollback the changes on ZK
              if (isSetTaskInfo) {
                  try {
                      StormClusterState clusterState = data.getStormClusterState();
                      clusterState.remove_task(topologyId, newTasks);
                  } catch (Exception e1) {
                      LOG.error("Failed to rollback the changes on ZK for task-" + newTasks, e);
                  }
              }
          }
  
          DelayStatusTransitionCallback delayCallback =
                  new DelayStatusTransitionCallback(data, topologyId, oldStatus, StatusType.rebalancing, StatusType.done_rebalance);
          return delayCallback.execute();
      }
  
      private void setTaskInfo(StormTopology oldTopology, StormTopology newTopology) throws Exception {
          StormClusterState clusterState = data.getStormClusterState();
          // Retrieve the max task ID
          TreeSet<Integer> taskIds = new TreeSet<>(clusterState.task_ids(topologyId));
          int cnt = taskIds.descendingIterator().next();
  
          cnt = setBoltInfo(oldTopology, newTopology, cnt, clusterState);
          cnt = setSpoutInfo(oldTopology, newTopology, cnt, clusterState);
      }
  
      private int setBoltInfo(StormTopology oldTopology, StormTopology newTopology,
                              int cnt, StormClusterState clusterState) throws Exception {
          Map<String, Bolt> oldBolts = oldTopology.get_bolts();
          Map<String, Bolt> bolts = newTopology.get_bolts();
          for (Entry<String, Bolt> entry : oldBolts.entrySet()) {
              String boltName = entry.getKey();
              Bolt oldBolt = entry.getValue();
              Bolt bolt = bolts.get(boltName);
              if (oldBolt.get_common().get_parallelism_hint() > bolt.get_common().get_parallelism_hint()) {
                  int removedTaskNum = oldBolt.get_common().get_parallelism_hint() - bolt.get_common().get_parallelism_hint();
                  TreeSet<Integer> taskIds = new TreeSet<>(clusterState.task_ids_by_componentId(topologyId, boltName));
                  Iterator<Integer> descendIterator = taskIds.descendingIterator();
                  while (--removedTaskNum >= 0) {
                      int taskId = descendIterator.next();
                      removeTask(topologyId, taskId, clusterState);
                      LOG.info("Remove bolt task, taskId=" + taskId + " for " + boltName);
                  }
              } else if (oldBolt.get_common().get_parallelism_hint() == bolt.get_common().get_parallelism_hint()) {
                  continue;
              } else {
                  int delta = bolt.get_common().get_parallelism_hint() - oldBolt.get_common().get_parallelism_hint();
                  Map<Integer, TaskInfo> taskInfoMap = new HashMap<>();
  
                  for (int i = 1; i <= delta; i++) {
                      cnt++;
                      TaskInfo taskInfo = new TaskInfo(entry.getKey(), "bolt");
                      taskInfoMap.put(cnt, taskInfo);
                      newTasks.add(cnt);
                      LOG.info("Setup new bolt task, taskId=" + cnt + " for " + boltName);
                  }
                  clusterState.add_task(topologyId, taskInfoMap);
              }
          }
  
          return cnt;
      }
  
      private int setSpoutInfo(StormTopology oldTopology, StormTopology newTopology,
                               int cnt, StormClusterState clusterState) throws Exception {
          Map<String, SpoutSpec> oldSpouts = oldTopology.get_spouts();
          Map<String, SpoutSpec> spouts = newTopology.get_spouts();
          for (Entry<String, SpoutSpec> entry : oldSpouts.entrySet()) {
              String spoutName = entry.getKey();
              SpoutSpec oldSpout = entry.getValue();
              SpoutSpec spout = spouts.get(spoutName);
              if (oldSpout.get_common().get_parallelism_hint() > spout.get_common().get_parallelism_hint()) {
                  int removedTaskNum = oldSpout.get_common().get_parallelism_hint() - spout.get_common().get_parallelism_hint();
                  TreeSet<Integer> taskIds = new TreeSet<>(clusterState.task_ids_by_componentId(topologyId, spoutName));
                  Iterator<Integer> descendIterator = taskIds.descendingIterator();
                  while (--removedTaskNum >= 0) {
                      int taskId = descendIterator.next();
                      removeTask(topologyId, taskId, clusterState);
                      LOG.info("Remove spout task, taskId=" + taskId + " for " + spoutName);
                  }
              } else if (oldSpout.get_common().get_parallelism_hint() == spout.get_common().get_parallelism_hint()) {
                  continue;
              } else {
                  int delta = spout.get_common().get_parallelism_hint() - oldSpout.get_common().get_parallelism_hint();
                  Map<Integer, TaskInfo> taskInfoMap = new HashMap<>();
  
                  for (int i = 1; i <= delta; i++) {
                      cnt++;
                      TaskInfo taskInfo = new TaskInfo(entry.getKey(), "spout");
                      taskInfoMap.put(cnt, taskInfo);
                      newTasks.add(cnt);
                      LOG.info("Setup new spout task, taskId=" + cnt + " for " + spoutName);
                  }
                  clusterState.add_task(topologyId, taskInfoMap);
              }
          }
  
          return cnt;
      }
  
      private void removeTask(String topologyId, int taskId, StormClusterState clusterState) throws Exception {
          Set<Integer> taskIds = new HashSet<>(taskId);
          clusterState.remove_task(this.topologyId, taskIds);
          Map<Integer, TkHbCacheTime> tkHbs = data.getTaskHeartbeatsCache(this.topologyId, false);
          if (tkHbs != null) {
              tkHbs.remove(taskId);
          }
      }
  }
  