package backtype.storm.task;
import backtype.storm.Config;
import backtype.storm.generated.ComponentCommon;
import backtype.storm.generated.GlobalStreamId;
import backtype.storm.generated.Grouping;
import backtype.storm.generated.StormTopology;
import backtype.storm.generated.StreamInfo;
import backtype.storm.state.ISubscribedState;
import backtype.storm.tuple.Fields;
import backtype.storm.utils.ThriftTopologyUtils;
import backtype.storm.utils.Utils;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.commons.lang.NotImplementedException;
import org.json.simple.JSONValue;
/**
* A TopologyContext is given to bolts and spouts in their "prepare" and "open"
* methods, respectively. This object provides information about the component's
* place within the topology, such as task ids, inputs and outputs, etc.
*
* <p>The TopologyContext is also used to declare ISubscribedState objects to
* synchronize state with StateSpouts this object is subscribed to.</p>
*/
public class TopologyContext {
private StormTopology _topology;
private Map<Integer, String> _taskToComponent;
private Integer _taskId;
private Map<String, List<Integer>> _componentToTasks;
private String _codeDir;
private String _pidDir;
private String _stormId;
private Object _taskData = null;
public TopologyContext(StormTopology topology, Map<Integer, String> taskToComponent, String stormId, String codeDir, String pidDir, Integer taskId) {
_topology = topology;
_taskToComponent = taskToComponent;
_stormId = stormId;
_taskId = taskId;
_componentToTasks = new HashMap<String, List<Integer>>();
_pidDir = pidDir;
_codeDir = codeDir;
for(Integer task: taskToComponent.keySet()) {
String component = taskToComponent.get(task);
List<Integer> curr = _componentToTasks.get(component);
if(curr==null) curr = new ArrayList<Integer>();
curr.add(task);
_componentToTasks.put(component, curr);
}
for(String component: _componentToTasks.keySet()) {
List<Integer> tasks = _componentToTasks.get(component);
Collections.sort(tasks);
}
}
/**
* All state from all subscribed state spouts streams will be synced with
* the provided object.
*
* <p>It is recommended that your ISubscribedState object is kept as an instance
* variable of this object. The recommended usage of this method is as follows:</p>
*
* <p>
* _myState = context.setAllSubscribedState(new MyState());
* </p>
* @param obj Provided ISubscribedState implementation
* @return Returns the ISubscribedState object provided
*/
public <T extends ISubscribedState> T setAllSubscribedState(T obj) {
//check that only subscribed to one component/stream for statespout
//setsubscribedstate appropriately
throw new NotImplementedException();
}
/**
* Synchronizes the default stream from the specified state spout component
* id with the provided ISubscribedState object.
*
* <p>The recommended usage of this method is as follows:</p>
* <p>
* _myState = context.setSubscribedState(componentId, new MyState());
* </p>
*
* @param componentId the id of the StateSpout component to subscribe to
* @param obj Provided ISubscribedState implementation
* @return Returns the ISubscribedState object provided
*/
public <T extends ISubscribedState> T setSubscribedState(String componentId, T obj) {
return setSubscribedState(componentId, Utils.DEFAULT_STREAM_ID, obj);
}
/**
* Synchronizes the specified stream from the specified state spout component
* id with the provided ISubscribedState object.
*
* <p>The recommended usage of this method is as follows:</p>
* <p>
* _myState = context.setSubscribedState(componentId, streamId, new MyState());
* </p>
*
* @param componentId the id of the StateSpout component to subscribe to
* @param streamId the stream to subscribe to
* @param obj Provided ISubscribedState implementation
* @return Returns the ISubscribedState object provided
*/
public <T extends ISubscribedState> T setSubscribedState(String componentId, String streamId, T obj) {
throw new NotImplementedException();
}
/**
* Gets the unique id assigned to this topology. The id is the storm name with a
* unique nonce appended to it.
* @return the storm id
*/
public String getStormId() {
return _stormId;
}
/**
* Gets the task id of this task.
*
* @return the task id
*/
public int getThisTaskId() {
return _taskId;
}
/**
* Gets the Thrift object representing the topology.
*
* @return the Thrift definition representing the topology
*/
public StormTopology getRawTopology() {
return _topology;
}
/**
* Gets the component id for the specified task id. The component id maps
* to a component id specified for a Spout or Bolt in the topology definition.
*
* @param taskId the task id
* @return the component id for the input task id
*/
public String getComponentId(int taskId) {
return _taskToComponent.get(taskId);
}
/**
* Gets the component id for this task. The component id maps
* to a component id specified for a Spout or Bolt in the topology definition.
* @return
*/
public String getThisComponentId() {
return getComponentId(_taskId);
}
/**
* Gets the declared output fields for the specified stream id for the component
* this task is a part of.
*/
public Fields getThisOutputFields(String streamId) {
return getComponentOutputFields(getThisComponentId(), streamId);
}
/**
* Gets the set of streams declared for the component of this task.
*/
public Set<String> getThisStreams() {
return getComponentStreams(getThisComponentId());
}
/**
* Gets the set of streams declared for the specified component.
*/
public Set<String> getComponentStreams(String componentId) {
return getComponentCommon(componentId).get_streams().keySet();
}
/**
* Gets the task ids allocated for the given component id. The task ids are
* always returned in ascending order.
*/
public List<Integer> getComponentTasks(String componentId) {
List<Integer> ret = _componentToTasks.get(componentId);
if(ret==null) return new ArrayList<Integer>();
else return new ArrayList<Integer>(ret);
}
/**
* Gets the index of this task id in getComponentTasks(getThisComponentId()).
* An example use case for this method is determining which task
* accesses which resource in a distributed resource to ensure an even distribution.
*/
public int getThisTaskIndex() {
List<Integer> tasks = new ArrayList<Integer>(getComponentTasks(getThisComponentId()));
Collections.sort(tasks);
for(int i=0; i<tasks.size(); i++) {
if(tasks.get(i) == getThisTaskId()) {
return i;
}
}
throw new RuntimeException("Fatal: could not find this task id in this component");
}
/**
* Gets the declared output fields for the specified component/stream.
*/
public Fields getComponentOutputFields(String componentId, String streamId) {
StreamInfo streamInfo = getComponentCommon(componentId).get_streams().get(streamId);
if(streamInfo==null) {
throw new IllegalArgumentException("No output fields defined for component:stream " + componentId + ":" + streamId);
}
return new Fields(streamInfo.get_output_fields());
}
/**
* Gets the declared output fields for the specified global stream id.
*/
public Fields getComponentOutputFields(GlobalStreamId id) {
return getComponentOutputFields(id.get_componentId(), id.get_streamId());
}
/**
* Gets the declared inputs to this component.
*
* @return A map from subscribed component/stream to the grouping subscribed with.
*/
public Map<GlobalStreamId, Grouping> getThisSources() {
return getSources(getThisComponentId());
}
/**
* Gets the declared inputs to the specified component.
*
* @return A map from subscribed component/stream to the grouping subscribed with.
*/
public Map<GlobalStreamId, Grouping> getSources(String componentId) {
return getComponentCommon(componentId).get_inputs();
}
/**
* Gets information about who is consuming the outputs of this component, and how.
*
* @return Map from stream id to component id to the Grouping used.
*/
public Map<String, Map<String, Grouping>> getThisTargets() {
return getTargets(getThisComponentId());
}
/**
* Gets information about who is consuming the outputs of the specified component,
* and how.
*
* @return Map from stream id to component id to the Grouping used.
*/
public Map<String, Map<String, Grouping>> getTargets(String componentId) {
Map<String, Map<String, Grouping>> ret = new HashMap<String, Map<String, Grouping>>();
for(String otherComponentId: getComponentIds()) {
Map<GlobalStreamId, Grouping> inputs = getComponentCommon(otherComponentId).get_inputs();
for(GlobalStreamId id: inputs.keySet()) {
if(id.get_componentId().equals(componentId)) {
Map<String, Grouping> curr = ret.get(id.get_streamId());
if(curr==null) curr = new HashMap<String, Grouping>();
curr.put(otherComponentId, inputs.get(id));
ret.put(id.get_streamId(), curr);
}
}
}
return ret;
}
public String toJSONString() {
Map obj = new HashMap();
obj.put("taskid", _taskId);
obj.put("task->component", _taskToComponent);
// TODO: jsonify StormTopology
// at the minimum should send source info
return JSONValue.toJSONString(obj);
}
/**
* Gets the location of the external resources for this worker on the
* local filesystem. These external resources typically include bolts implemented
* in other languages, such as Ruby or Python.
*/
public String getCodeDir() {
return _codeDir;
}
/**
* If this task spawns any subprocesses, those subprocesses must immediately
* write their PID to this directory on the local filesystem to ensure that
* Storm properly destroys that process when the worker is shutdown.
*/
public String getPIDDir() {
return _pidDir;
}
/**
* Gets a map from task id to component id.
*/
public Map<Integer, String> getTaskToComponent() {
return _taskToComponent;
}
/**
* Gets a list of all component ids in this topology
*/
public Set<String> getComponentIds() {
return ThriftTopologyUtils.getComponentIds(getRawTopology());
}
public ComponentCommon getComponentCommon(String componentId) {
return ThriftTopologyUtils.getComponentCommon(getRawTopology(), componentId);
}
public void setTaskData(Object data) {
_taskData = data;
}
public Object getTaskData() {
return _taskData;
}
public int maxTopologyMessageTimeout(Map<String, Object> topologyConfig) {
Integer max = Utils.getInt(topologyConfig.get(Config.TOPOLOGY_MESSAGE_TIMEOUT_SECS));
for(String spout: getRawTopology().get_spouts().keySet()) {
ComponentCommon common = getComponentCommon(spout);
String jsonConf = common.get_json_conf();
if(jsonConf!=null) {
Map conf = (Map) JSONValue.parse(jsonConf);
Object comp = conf.get(Config.TOPOLOGY_MESSAGE_TIMEOUT_SECS);
if(comp!=null) {
max = Math.max(Utils.getInt(comp), max);
}
}
}
return max;
}
}