/**
* 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 backtype.storm.coordination;
import static backtype.storm.utils.Utils.get;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import backtype.storm.Constants;
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.tuple.Fields;
import backtype.storm.tuple.Tuple;
import backtype.storm.tuple.Values;
import backtype.storm.utils.TimeCacheMap;
import backtype.storm.utils.Utils;
/**
* Coordination requires the request ids to be globally unique for a while
* so that it doesn't get confused in the case of retries.
*/
public class CoordinatedBolt implements IRichBolt {
public static Logger LOG = LoggerFactory.getLogger(CoordinatedBolt.class);
public interface FinishedCallback {
void finishedId(Object id);
}
public interface TimeoutCallback {
void timeoutId(Object id);
}
public static class SourceArgs implements Serializable {
public boolean singleCount;
protected SourceArgs(boolean singleCount) {
this.singleCount = singleCount;
}
public static SourceArgs single() {
return new SourceArgs(true);
}
public static SourceArgs all() {
return new SourceArgs(false);
}
@Override
public String toString() {
return "<Single: " + singleCount + ">";
}
}
public class CoordinatedOutputCollector implements IOutputCollector {
OutputCollector _delegate;
public CoordinatedOutputCollector(OutputCollector delegate) {
_delegate = delegate;
}
public List<Integer> emit(String stream, Collection<Tuple> anchors, List<Object> tuple) {
return _delegate.emit(stream, anchors, tuple, new CollectorCb(tuple.get(0)));
}
public void emitDirect(int task, String stream, Collection<Tuple> anchors, List<Object> tuple) {
updateTaskCounts(tuple.get(0), Arrays.asList(task));
_delegate.emitDirect(task, stream, anchors, tuple);
// due to here do updateTaskCounts, so we do flush operation
_delegate.flush();
}
public void ack(Tuple tuple) {
Object id = tuple.getValue(0);
synchronized (_tracked) {
TrackingInfo track = _tracked.get(id);
if (track != null)
track.receivedTuples++;
}
boolean failed = checkFinishId(tuple, TupleType.REGULAR);
if (failed) {
_delegate.fail(tuple);
} else {
_delegate.ack(tuple);
}
}
public void fail(Tuple tuple) {
Object id = tuple.getValue(0);
synchronized (_tracked) {
TrackingInfo track = _tracked.get(id);
if (track != null)
track.failed = true;
}
checkFinishId(tuple, TupleType.REGULAR);
_delegate.fail(tuple);
}
public void reportError(Throwable error) {
_delegate.reportError(error);
}
private void updateTaskCounts(Object id, List<Integer> tasks) {
synchronized (_tracked) {
TrackingInfo track = _tracked.get(id);
if (track != null) {
Map<Integer, Integer> taskEmittedTuples = track.taskEmittedTuples;
for (Integer task : tasks) {
int newCount = get(taskEmittedTuples, task, 0) + 1;
taskEmittedTuples.put(task, newCount);
}
}
}
}
class CollectorCb implements ICollectorCallback {
Object id;
public CollectorCb(Object id) {
this.id = id;
}
@Override
public void execute(String stream, List<Integer> outTasks, List values) {
updateTaskCounts(id, outTasks);
}
}
}
private Map<String, SourceArgs> _sourceArgs;
private IdStreamSpec _idStreamSpec;
private IRichBolt _delegate;
private Integer _numSourceReports;
private List<Integer> _countOutTasks = new ArrayList<>();
private OutputCollector _collector;
private TimeCacheMap<Object, TrackingInfo> _tracked;
public static class TrackingInfo {
int reportCount = 0;
int expectedTupleCount = 0;
int receivedTuples = 0;
boolean failed = false;
Map<Integer, Integer> taskEmittedTuples = new HashMap<>();
boolean receivedId = false;
boolean finished = false;
List<Tuple> ackTuples = new ArrayList<>();
@Override
public String toString() {
return "reportCount: " + reportCount + "\n" + "expectedTupleCount: " + expectedTupleCount + "\n" +
"receivedTuples: " + receivedTuples + "\n"
+ "failed: " + failed + "\n" + taskEmittedTuples.toString();
}
}
public static class IdStreamSpec implements Serializable {
GlobalStreamId _id;
public GlobalStreamId getGlobalStreamId() {
return _id;
}
public static IdStreamSpec makeDetectSpec(String component, String stream) {
return new IdStreamSpec(component, stream);
}
protected IdStreamSpec(String component, String stream) {
_id = new GlobalStreamId(component, stream);
}
}
public CoordinatedBolt(IRichBolt delegate) {
this(delegate, null, null);
}
public CoordinatedBolt(IRichBolt delegate, String sourceComponent, SourceArgs sourceArgs, IdStreamSpec idStreamSpec) {
this(delegate, singleSourceArgs(sourceComponent, sourceArgs), idStreamSpec);
}
public CoordinatedBolt(IRichBolt delegate, Map<String, SourceArgs> sourceArgs, IdStreamSpec idStreamSpec) {
_sourceArgs = sourceArgs;
if (_sourceArgs == null)
_sourceArgs = new HashMap<>();
_delegate = delegate;
_idStreamSpec = idStreamSpec;
}
public void prepare(Map config, TopologyContext context, OutputCollector collector) {
TimeCacheMap.ExpiredCallback<Object, TrackingInfo> callback = null;
if (_delegate instanceof TimeoutCallback) {
callback = new TimeoutItems();
}
_tracked = new TimeCacheMap<>(context.maxTopologyMessageTimeout(), callback);
_collector = collector;
_delegate.prepare(config, context, new OutputCollector(new CoordinatedOutputCollector(collector)));
for (String component : Utils.get(context.getThisTargets(), Constants.COORDINATED_STREAM_ID, new HashMap<String, Grouping>()).keySet()) {
for (Integer task : context.getComponentTasks(component)) {
_countOutTasks.add(task);
}
}
if (!_sourceArgs.isEmpty()) {
_numSourceReports = 0;
for (Entry<String, SourceArgs> entry : _sourceArgs.entrySet()) {
if (entry.getValue().singleCount) {
_numSourceReports += 1;
} else {
_numSourceReports += context.getComponentTasks(entry.getKey()).size();
}
}
}
}
private boolean checkFinishId(Tuple tup, TupleType type) {
_collector.flush();
Object id = tup.getValue(0);
boolean failed = false;
synchronized (_tracked) {
TrackingInfo track = _tracked.get(id);
try {
if (track != null) {
boolean delayed = false;
if (_idStreamSpec == null && type == TupleType.COORD || _idStreamSpec != null && type == TupleType.ID) {
track.ackTuples.add(tup);
delayed = true;
}
if (track.failed) {
failed = true;
for (Tuple t : track.ackTuples) {
_collector.fail(t);
}
_tracked.remove(id);
} else if (track.receivedId
&& (_sourceArgs.isEmpty() || track.reportCount == _numSourceReports && track.expectedTupleCount == track.receivedTuples)) {
if (_delegate instanceof FinishedCallback) {
((FinishedCallback) _delegate).finishedId(id);
}
if (!(_sourceArgs.isEmpty() || type != TupleType.REGULAR)) {
throw new IllegalStateException("Coordination condition met with a non-coordinating tuple. Should be impossible");
}
Iterator<Integer> outTasks = _countOutTasks.iterator();
_collector.flush();
while (outTasks.hasNext()) {
int task = outTasks.next();
int numTuples = get(track.taskEmittedTuples, task, 0);
_collector.emitDirect(task, Constants.COORDINATED_STREAM_ID, tup, new Values(id, numTuples));
}
for (Tuple t : track.ackTuples) {
_collector.ack(t);
}
track.finished = true;
_tracked.remove(id);
}
if (!delayed && type != TupleType.REGULAR) {
if (track.failed) {
_collector.fail(tup);
} else {
_collector.ack(tup);
}
}
} else {
if (type != TupleType.REGULAR)
_collector.fail(tup);
}
} catch (FailedException e) {
LOG.error("Failed to finish batch", e);
for (Tuple t : track.ackTuples) {
_collector.fail(t);
}
_tracked.remove(id);
failed = true;
} finally {
_collector.flush();
}
}
return failed;
}
public void execute(Tuple tuple) {
Object id = tuple.getValue(0);
TrackingInfo track;
TupleType type = getTupleType(tuple);
synchronized (_tracked) {
track = _tracked.get(id);
if (track == null) {
track = new TrackingInfo();
if (_idStreamSpec == null)
track.receivedId = true;
_tracked.put(id, track);
}
}
if (type == TupleType.ID) {
synchronized (_tracked) {
track.receivedId = true;
}
checkFinishId(tuple, type);
} else if (type == TupleType.COORD) {
int count = (Integer) tuple.getValue(1);
synchronized (_tracked) {
track.reportCount++;
track.expectedTupleCount += count;
}
checkFinishId(tuple, type);
} else {
synchronized (_tracked) {
_delegate.execute(tuple);
}
}
}
public void cleanup() {
_delegate.cleanup();
_tracked.cleanup();
}
public void declareOutputFields(OutputFieldsDeclarer declarer) {
_delegate.declareOutputFields(declarer);
declarer.declareStream(Constants.COORDINATED_STREAM_ID, true, new Fields("id", "count"));
}
@Override
public Map<String, Object> getComponentConfiguration() {
return _delegate.getComponentConfiguration();
}
private static Map<String, SourceArgs> singleSourceArgs(String sourceComponent, SourceArgs sourceArgs) {
Map<String, SourceArgs> ret = new HashMap<>();
ret.put(sourceComponent, sourceArgs);
return ret;
}
private class TimeoutItems implements TimeCacheMap.ExpiredCallback<Object, TrackingInfo> {
@Override
public void expire(Object id, TrackingInfo val) {
synchronized (_tracked) {
// the combination of the lock and the finished flag ensure that
// an id is never timed out if it has been finished
val.failed = true;
if (!val.finished) {
((TimeoutCallback) _delegate).timeoutId(id);
}
}
}
}
private TupleType getTupleType(Tuple tuple) {
if (_idStreamSpec != null && tuple.getSourceGlobalStreamid().equals(_idStreamSpec._id)) {
return TupleType.ID;
} else if (!_sourceArgs.isEmpty() && tuple.getSourceStreamId().equals(Constants.COORDINATED_STREAM_ID)) {
return TupleType.COORD;
} else {
return TupleType.REGULAR;
}
}
enum TupleType {
REGULAR, ID, COORD
}
}