/**
* 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 storm.trident.topology;
import backtype.storm.Config;
import backtype.storm.Constants;
import backtype.storm.coordination.BatchOutputCollector;
import backtype.storm.coordination.BatchOutputCollectorImpl;
import backtype.storm.generated.GlobalStreamId;
import backtype.storm.generated.Grouping;
import backtype.storm.task.ICollectorCallback;
import backtype.storm.task.IOutputCollector;
import backtype.storm.task.OutputCollector;
import backtype.storm.task.TopologyContext;
import backtype.storm.topology.FailedException;
import backtype.storm.topology.IRichBolt;
import backtype.storm.topology.OutputFieldsDeclarer;
import backtype.storm.topology.ReportedFailedException;
import backtype.storm.tuple.Fields;
import backtype.storm.tuple.Tuple;
import backtype.storm.tuple.Values;
import backtype.storm.utils.RotatingMap;
import backtype.storm.utils.TupleUtils;
import backtype.storm.utils.Utils;
import java.io.Serializable;
import java.util.*;
import org.apache.commons.lang.builder.ToStringBuilder;
import storm.trident.spout.IBatchID;
public class TridentBoltExecutor implements IRichBolt {
public static final String COORD_STREAM_PREFIX = "$coord-";
public static String COORD_STREAM(String batch) {
return COORD_STREAM_PREFIX + batch;
}
public static class CoordType implements Serializable {
public boolean singleCount;
protected CoordType(boolean singleCount) {
this.singleCount = singleCount;
}
public static CoordType single() {
return new CoordType(true);
}
public static CoordType all() {
return new CoordType(false);
}
@Override
public int hashCode() {
final int prime = 31;
int result = 1;
result = prime * result + (singleCount ? 1231 : 1237);
return result;
}
@Override
public boolean equals(Object obj) {
if (this == obj)
return true;
if (obj == null)
return false;
if (getClass() != obj.getClass())
return false;
CoordType other = (CoordType) obj;
if (singleCount != other.singleCount)
return false;
return true;
}
@Override
public String toString() {
return "<Single: " + singleCount + ">";
}
}
public static class CoordSpec implements Serializable {
public GlobalStreamId commitStream = null;
public Map<String, CoordType> coords = new HashMap<>();
public CoordSpec() {
}
}
public static class CoordCondition implements Serializable {
public GlobalStreamId commitStream;
public int expectedTaskReports;
Set<Integer> targetTasks;
@Override
public String toString() {
return ToStringBuilder.reflectionToString(this);
}
}
Map<GlobalStreamId, String> _batchGroupIds;
Map<String, CoordSpec> _coordSpecs;
Map<String, CoordCondition> _coordConditions;
ITridentBatchBolt _bolt;
long _messageTimeoutMs;
long _lastRotate;
RotatingMap<Object, TrackedBatch> _batches;
// map from batch group id to coord spec
public TridentBoltExecutor(ITridentBatchBolt bolt, Map<GlobalStreamId, String> batchGroupIds,
Map<String, CoordSpec> coordinationSpecs) {
_batchGroupIds = batchGroupIds;
_coordSpecs = coordinationSpecs;
_bolt = bolt;
}
public static class TrackedBatch {
int attemptId;
BatchInfo info;
CoordCondition condition;
int reportedTasks = 0;
int expectedTupleCount = 0;
int receivedTuples = 0;
Map<Integer, Integer> taskEmittedTuples = new HashMap<>();
boolean failed = false;
boolean receivedCommit;
Tuple delayedAck = null;
public TrackedBatch(BatchInfo info, CoordCondition condition, int attemptId) {
this.info = info;
this.condition = condition;
this.attemptId = attemptId;
receivedCommit = condition.commitStream == null;
}
@Override
public String toString() {
return ToStringBuilder.reflectionToString(this);
}
}
private static class CoordinatedOutputCollector implements IOutputCollector {
OutputCollector _delegate;
TrackedBatch _currBatch = null;
class CollectorCb implements ICollectorCallback {
public CollectorCb() {
}
@Override
public void execute(String stream, List<Integer> outTasks, List values) {
updateTaskCounts(outTasks);
}
}
public void setCurrBatch(TrackedBatch batch) {
_currBatch = batch;
}
public CoordinatedOutputCollector(OutputCollector delegate) {
_delegate = delegate;
}
public List<Integer> emit(String stream, Collection<Tuple> anchors, List<Object> tuple) {
List<Integer> tasks = _delegate.emit(stream, anchors, tuple, new CollectorCb());
return tasks;
}
public void emitDirect(int task, String stream, Collection<Tuple> anchors, List<Object> tuple) {
updateTaskCounts(Arrays.asList(task));
_delegate.emitDirect(task, stream, anchors, tuple);
_delegate.flush();
}
public void ack(Tuple tuple) {
throw new IllegalStateException("Method should never be called");
}
public void fail(Tuple tuple) {
throw new IllegalStateException("Method should never be called");
}
public void resetTimeout(Tuple tuple) {
throw new IllegalStateException("Method should never be called");
}
public void reportError(Throwable error) {
_delegate.reportError(error);
}
private void updateTaskCounts(List<Integer> tasks) {
synchronized (_currBatch) {
if (_currBatch != null) {
Map<Integer, Integer> taskEmittedTuples = _currBatch.taskEmittedTuples;
for (Integer task : tasks) {
int newCount = Utils.get(taskEmittedTuples, task, 0) + 1;
taskEmittedTuples.put(task, newCount);
}
}
}
}
}
OutputCollector _collector;
CoordinatedOutputCollector _coordCollector;
BatchOutputCollector _coordOutputCollector;
TopologyContext _context;
@Override
public void prepare(Map conf, TopologyContext context, OutputCollector collector) {
_messageTimeoutMs = context.maxTopologyMessageTimeout() * 1000L;
_lastRotate = System.currentTimeMillis();
_batches = new RotatingMap<>(2);
_context = context;
_collector = collector;
_coordCollector = new CoordinatedOutputCollector(new OutputCollector(collector));
_coordOutputCollector = new BatchOutputCollectorImpl(new OutputCollector(_coordCollector));
_coordConditions = (Map) context.getExecutorData("__coordConditions");
if (_coordConditions == null) {
_coordConditions = new HashMap<>();
for (String batchGroup : _coordSpecs.keySet()) {
CoordSpec spec = _coordSpecs.get(batchGroup);
CoordCondition cond = new CoordCondition();
cond.commitStream = spec.commitStream;
cond.expectedTaskReports = 0;
for (String comp : spec.coords.keySet()) {
CoordType ct = spec.coords.get(comp);
if (ct.equals(CoordType.single())) {
cond.expectedTaskReports += 1;
} else {
cond.expectedTaskReports += context.getComponentTasks(comp).size();
}
}
cond.targetTasks = new HashSet<>();
for (String component : Utils.get(context.getThisTargets(),
COORD_STREAM(batchGroup),
new HashMap<String, Grouping>()).keySet()) {
cond.targetTasks.addAll(context.getComponentTasks(component));
}
_coordConditions.put(batchGroup, cond);
}
context.setExecutorData("_coordConditions", _coordConditions);
}
_bolt.prepare(conf, context, _coordOutputCollector);
}
private void failBatch(TrackedBatch tracked, FailedException e) {
if (e != null && e instanceof ReportedFailedException) {
_collector.reportError(e);
}
tracked.failed = true;
if (tracked.delayedAck != null) {
_collector.fail(tracked.delayedAck);
tracked.delayedAck = null;
}
}
private void failBatch(TrackedBatch tracked) {
failBatch(tracked, null);
}
private boolean finishBatch(TrackedBatch tracked, Tuple finishTuple) {
boolean success = true;
try {
_bolt.finishBatch(tracked.info);
String stream = COORD_STREAM(tracked.info.batchGroup);
_collector.flush();
for (Integer task : tracked.condition.targetTasks) {
_collector.emitDirect(task, stream, finishTuple, new Values(tracked.info.batchId,
Utils.get(tracked.taskEmittedTuples, task, 0)));
}
if (tracked.delayedAck != null) {
_collector.ack(tracked.delayedAck);
tracked.delayedAck = null;
}
} catch (FailedException e) {
failBatch(tracked, e);
success = false;
}
_batches.remove(tracked.info.batchId.getId());
return success;
}
private void checkFinish(TrackedBatch tracked, Tuple tuple, TupleType type) {
if (tracked.failed) {
failBatch(tracked);
_collector.fail(tuple);
return;
}
CoordCondition cond = tracked.condition;
boolean delayed = tracked.delayedAck == null &&
(cond.commitStream != null && type == TupleType.COMMIT
|| cond.commitStream == null);
if (delayed) {
tracked.delayedAck = tuple;
}
boolean failed = false;
if (tracked.receivedCommit && tracked.reportedTasks == cond.expectedTaskReports) {
if (tracked.receivedTuples == tracked.expectedTupleCount) {
finishBatch(tracked, tuple);
} else {
//TODO: add logging that not all tuples were received
failBatch(tracked);
_collector.fail(tuple);
failed = true;
}
}
if (!delayed && !failed) {
_collector.ack(tuple);
}
}
@Override
public void execute(Tuple tuple) {
if (TupleUtils.isTick(tuple)) {
long now = System.currentTimeMillis();
if (now - _lastRotate > _messageTimeoutMs) {
_batches.rotate();
_lastRotate = now;
}
return;
}
String batchGroup = _batchGroupIds.get(tuple.getSourceGlobalStreamid());
if (batchGroup == null) {
// this is so we can do things like have simple DRPC that doesn't need to use batch processing
_coordCollector.setCurrBatch(null);
_bolt.execute(null, tuple);
_collector.ack(tuple);
return;
}
IBatchID id = (IBatchID) tuple.getValue(0);
//get transaction id
//if it already exists and attempt id is greater than the attempt there
TrackedBatch tracked = (TrackedBatch) _batches.get(id.getId());
// if(_batches.size() > 10 && _context.getThisTaskIndex() == 0) {
// System.out.println("Received in " + _context.getThisComponentId() + " " + _context.getThisTaskIndex()
// + " (" + _batches.size() + ")" +
// "\ntuple: " + tuple +
// "\nwith tracked " + tracked +
// "\nwith id " + id +
// "\nwith group " + batchGroup
// + "\n");
//
// }
//System.out.println("Num tracked: " + _batches.size() + " " + _context.getThisComponentId() + " " + _context.getThisTaskIndex());
// this code here ensures that only one attempt is ever tracked for a batch, so when
// failures happen you don't get an explosion in memory usage in the tasks
if (tracked != null) {
if (id.getAttemptId() > tracked.attemptId) {
_batches.remove(id.getId());
tracked = null;
} else if (id.getAttemptId() < tracked.attemptId) {
// no reason to try to execute a previous attempt than we've already seen
return;
}
}
if (tracked == null) {
tracked = new TrackedBatch(new BatchInfo(batchGroup, id, _bolt.initBatchState(batchGroup, id)),
_coordConditions.get(batchGroup), id.getAttemptId());
_batches.put(id.getId(), tracked);
}
_coordCollector.setCurrBatch(tracked);
//System.out.println("TRACKED: " + tracked + " " + tuple);
TupleType t = getTupleType(tuple, tracked);
if (t == TupleType.COMMIT) {
tracked.receivedCommit = true;
checkFinish(tracked, tuple, t);
} else if (t == TupleType.COORD) {
int count = tuple.getInteger(1);
tracked.reportedTasks++;
tracked.expectedTupleCount += count;
checkFinish(tracked, tuple, t);
} else {
tracked.receivedTuples++;
boolean success = true;
try {
_bolt.execute(tracked.info, tuple);
if (tracked.condition.expectedTaskReports == 0) {
success = finishBatch(tracked, tuple);
}
} catch (FailedException e) {
failBatch(tracked, e);
}
if (success) {
_collector.ack(tuple);
} else {
_collector.fail(tuple);
}
}
_coordCollector.setCurrBatch(null);
}
@Override
public void cleanup() {
_bolt.cleanup();
}
@Override
public void declareOutputFields(OutputFieldsDeclarer declarer) {
_bolt.declareOutputFields(declarer);
for (String batchGroup : _coordSpecs.keySet()) {
declarer.declareStream(COORD_STREAM(batchGroup), true, new Fields("id", "count"));
}
}
@Override
public Map<String, Object> getComponentConfiguration() {
Map<String, Object> ret = _bolt.getComponentConfiguration();
if (ret == null) ret = new HashMap<>();
ret.put(Config.TOPOLOGY_TICK_TUPLE_FREQ_SECS, 5);
// TODO: Need to be able to set the tick tuple time to the message timeout, ideally without parameterization
return ret;
}
private TupleType getTupleType(Tuple tuple, TrackedBatch batch) {
CoordCondition cond = batch.condition;
if (cond.commitStream != null
&& tuple.getSourceGlobalStreamid().equals(cond.commitStream)) {
return TupleType.COMMIT;
} else if (cond.expectedTaskReports > 0
&& tuple.getSourceStreamId().startsWith(COORD_STREAM_PREFIX)) {
return TupleType.COORD;
} else {
return TupleType.REGULAR;
}
}
enum TupleType {
REGULAR,
COMMIT,
COORD
}
}