jstorm-core/src/main/java/com/alibaba/jstorm/callback/impl/DoRebalanceTransitionCallback.java - storm - Git at Google

 /**
  * 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 receive one rebalance command, it will do as following: 1. set topology status as rebalancing 2. delay 2 * timeout seconds 3. do this callback
  *
  * @author Xin.Li/Longda
  *
  */
 public class DoRebalanceTransitionCallback extends BaseCallback {

     private static 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<Integer>();
     }

     @Override
     public <T> Object execute(T... args) {
         boolean isSetTaskInfo = false;
         try {
             Boolean reassign = (Boolean) args[1];
             Map<Object, Object> conf = (Map<Object, Object>) args[2]; // args[0]:
                                                                       // delay,
                                                                       // args[1]:
                                                                       // reassign_flag,
                                                                       // args[2]:
                                                                       // conf
             if (conf != null) {
                 boolean isConfUpdate = false;
                 Map stormConf = data.getConf();

                 // Update topology code
                 Map topoConf = StormConfig.read_nimbus_topology_conf(stormConf, topologyid);
                 StormTopology rawOldTopology = StormConfig.read_nimbus_topology_code(stormConf, topologyid);
                 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(stormConf, topologyid, Utils.serialize(rawNewTopology));

                 // 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(stormConf, topologyid, topoConf);
                 }
             }

             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<Integer>(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<Integer>(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<Integer, TaskInfo>();

                 for (int i = 1; i <= delta; i++) {
                     cnt++;
                     TaskInfo taskInfo = new TaskInfo((String) 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<Integer>(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<Integer, TaskInfo>();

                 for (int i = 1; i <= delta; i++) {
                     cnt++;
                     TaskInfo taskInfo = new TaskInfo((String) 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<Integer>(taskId);
         clusterState.remove_task(topologyid, taskIds);
         Map<Integer, TkHbCacheTime> TkHbs = data.getTaskHeartbeatsCache(topologyid, false);
         if (TkHbs != null) {
             TkHbs.remove(taskId);
         }
     }
 }
	/**
	* 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 receive one rebalance command, it will do as following: 1. set topology status as rebalancing 2. delay 2 * timeout seconds 3. do this callback
	*
	* @author Xin.Li/Longda
	*
	*/
	public class DoRebalanceTransitionCallback extends BaseCallback {

	private static 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<Integer>();
	}

	@Override
	public <T> Object execute(T... args) {
	boolean isSetTaskInfo = false;
	try {
	Boolean reassign = (Boolean) args[1];
	Map<Object, Object> conf = (Map<Object, Object>) args[2]; // args[0]:
	// delay,
	// args[1]:
	// reassign_flag,
	// args[2]:
	// conf
	if (conf != null) {
	boolean isConfUpdate = false;
	Map stormConf = data.getConf();

	// Update topology code
	Map topoConf = StormConfig.read_nimbus_topology_conf(stormConf, topologyid);
	StormTopology rawOldTopology = StormConfig.read_nimbus_topology_code(stormConf, topologyid);
	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(stormConf, topologyid, Utils.serialize(rawNewTopology));

	// 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(stormConf, topologyid, topoConf);
	}
	}

	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<Integer>(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<Integer>(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<Integer, TaskInfo>();

	for (int i = 1; i <= delta; i++) {
	cnt++;
	TaskInfo taskInfo = new TaskInfo((String) 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<Integer>(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<Integer, TaskInfo>();

	for (int i = 1; i <= delta; i++) {
	cnt++;
	TaskInfo taskInfo = new TaskInfo((String) 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<Integer>(taskId);
	clusterState.remove_task(topologyid, taskIds);
	Map<Integer, TkHbCacheTime> TkHbs = data.getTaskHeartbeatsCache(topologyid, false);
	if (TkHbs != null) {
	TkHbs.remove(taskId);
	}
	}
	}