/**
* 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 backtype.storm.Config;
import backtype.storm.generated.*;
import backtype.storm.utils.ThriftTopologyUtils;
import com.alibaba.jstorm.client.ConfigExtension;
import com.alibaba.jstorm.cluster.Common;
import com.alibaba.jstorm.daemon.supervisor.SupervisorInfo;
import com.alibaba.jstorm.schedule.TopologyAssignContext;
import com.alibaba.jstorm.utils.FailedAssignTopologyException;
import com.alibaba.jstorm.utils.JStormUtils;
import org.apache.commons.lang.builder.ToStringBuilder;
import org.apache.commons.lang.builder.ToStringStyle;
import java.util.*;
import java.util.Map.Entry;
public class DefaultTopologyAssignContext extends TopologyAssignContext {
private final StormTopology sysTopology;
private final Map<String, String> sidToHostname;
private final Map<String, List<String>> hostToSid;
private final Set<ResourceWorkerSlot> oldWorkers;
private final Map<String, List<Integer>> componentTasks;
private final Set<ResourceWorkerSlot> unstoppedWorkers = new HashSet<>();
private final int totalWorkerNum;
private final int unstoppedWorkerNum;
private final int reserveWorkerNum;
private int computeWorkerNum() {
Integer settingNum = JStormUtils.parseInt(stormConf.get(Config.TOPOLOGY_WORKERS));
int ret, hintSum = 0, tmCount = 0;
Map<String, Object> components = ThriftTopologyUtils.getComponents(sysTopology);
for (Entry<String, Object> entry : components.entrySet()) {
String componentName = entry.getKey();
Object component = entry.getValue();
ComponentCommon common = null;
if (component instanceof Bolt) {
common = ((Bolt) component).get_common();
}
if (component instanceof SpoutSpec) {
common = ((SpoutSpec) component).get_common();
}
if (component instanceof StateSpoutSpec) {
common = ((StateSpoutSpec) component).get_common();
}
int hint = common.get_parallelism_hint();
if (componentName.equals(Common.TOPOLOGY_MASTER_COMPONENT_ID)) {
tmCount += hint;
continue;
}
hintSum += hint;
}
if (settingNum == null) {
ret = hintSum;
} else {
ret = Math.min(settingNum, hintSum);
}
Boolean isTmSingleWorker = ConfigExtension.getTopologyMasterSingleWorker(stormConf);
if (isTmSingleWorker != null) {
if (isTmSingleWorker) {
// Assign a single worker for topology master
ret += tmCount;
setAssignSingleWorkerForTM(true);
ConfigExtension.setTopologyMasterSingleWorker(stormConf, true);
}
} else {
// If not configured, judge this config by worker number
if (ret >= 10) {
ret += tmCount;
setAssignSingleWorkerForTM(true);
ConfigExtension.setTopologyMasterSingleWorker(stormConf, true);
}
}
return ret;
}
public int computeUnstoppedAssignments() {
for (Integer task : unstoppedTaskIds) {
// if unstoppedTasksIds isn't empty, it should be REASSIGN/Monitor
ResourceWorkerSlot worker = oldAssignment.getWorkerByTaskId(task);
unstoppedWorkers.add(worker);
}
return unstoppedWorkers.size();
}
private void refineDeadTasks() {
Set<Integer> rawDeadTasks = getDeadTaskIds();
Set<Integer> refineDeadTasks = new HashSet<>();
refineDeadTasks.addAll(rawDeadTasks);
Set<Integer> unstoppedTasks = getUnstoppedTaskIds();
// if type is ASSIGN_NEW, rawDeadTasks is empty
// then the oldWorkerTasks should be existingAssignment
for (Integer task : rawDeadTasks) {
if (unstoppedTasks.contains(task)) {
continue;
}
for (ResourceWorkerSlot worker : oldWorkers) {
if (worker.getTasks().contains(task)) {
refineDeadTasks.addAll(worker.getTasks());
}
}
}
setDeadTaskIds(refineDeadTasks);
}
/**
* Do we need just handle the case when type is ASSIGN_TYPE_NEW?
*/
private Map<String, String> generateSidToHost() {
Map<String, String> sidToHostname = new HashMap<>();
if (oldAssignment != null) {
sidToHostname.putAll(oldAssignment.getNodeHost());
}
for (Entry<String, SupervisorInfo> entry : cluster.entrySet()) {
String supervisorId = entry.getKey();
SupervisorInfo supervisorInfo = entry.getValue();
sidToHostname.put(supervisorId, supervisorInfo.getHostName());
}
return sidToHostname;
}
public DefaultTopologyAssignContext(TopologyAssignContext context) {
super(context);
try {
sysTopology = Common.system_topology(stormConf, rawTopology);
} catch (Exception e) {
throw new FailedAssignTopologyException("Failed to generate system topology");
}
sidToHostname = generateSidToHost();
hostToSid = JStormUtils.reverse_map(sidToHostname);
if (oldAssignment != null && oldAssignment.getWorkers() != null) {
oldWorkers = oldAssignment.getWorkers();
} else {
oldWorkers = new HashSet<>();
}
refineDeadTasks();
componentTasks = JStormUtils.reverse_map(context.getTaskToComponent());
for (Entry<String, List<Integer>> entry : componentTasks.entrySet()) {
List<Integer> componentTaskList = entry.getValue();
Collections.sort(componentTaskList);
}
totalWorkerNum = computeWorkerNum();
unstoppedWorkerNum = computeUnstoppedAssignments();
reserveWorkerNum = ConfigExtension.getReserveWorkers(stormConf);
}
public StormTopology getSysTopology() {
return sysTopology;
}
public Map<String, String> getSidToHostname() {
return sidToHostname;
}
public Map<String, List<String>> getHostToSid() {
return hostToSid;
}
public Map<String, List<Integer>> getComponentTasks() {
return componentTasks;
}
public int getTotalWorkerNum() {
return totalWorkerNum;
}
public int getUnstoppedWorkerNum() {
return unstoppedWorkerNum;
}
public Set<ResourceWorkerSlot> getOldWorkers() {
return oldWorkers;
}
public Set<ResourceWorkerSlot> getUnstoppedWorkers() {
return unstoppedWorkers;
}
@Override
public String toString() {
return (String) stormConf.get(Config.TOPOLOGY_NAME);
}
public String toDetailString() {
return ToStringBuilder.reflectionToString(this, ToStringStyle.SHORT_PREFIX_STYLE);
}
public int getReserveWorkerNum() {
return reserveWorkerNum;
}
}