package backtype.storm.transactional.partitioned;
import backtype.storm.Config;
import backtype.storm.task.TopologyContext;
import backtype.storm.topology.OutputFieldsDeclarer;
import backtype.storm.transactional.ITransactionalSpout;
import backtype.storm.transactional.TransactionAttempt;
import backtype.storm.coordination.BatchOutputCollector;
import backtype.storm.transactional.state.RotatingTransactionalState;
import backtype.storm.transactional.state.TransactionalState;
import java.math.BigInteger;
import java.util.HashMap;
import java.util.Map;
public class PartitionedTransactionalSpoutExecutor implements ITransactionalSpout<Integer> {
IPartitionedTransactionalSpout _spout;
public PartitionedTransactionalSpoutExecutor(IPartitionedTransactionalSpout spout) {
_spout = spout;
}
public IPartitionedTransactionalSpout getPartitionedSpout() {
return _spout;
}
class Coordinator implements ITransactionalSpout.Coordinator<Integer> {
private IPartitionedTransactionalSpout.Coordinator _coordinator;
public Coordinator(Map conf, TopologyContext context) {
_coordinator = _spout.getCoordinator(conf, context);
}
@Override
public Integer initializeTransaction(BigInteger txid, Integer prevMetadata) {
return _coordinator.numPartitions();
}
@Override
public void close() {
_coordinator.close();
}
}
class Emitter implements ITransactionalSpout.Emitter<Integer> {
private IPartitionedTransactionalSpout.Emitter _emitter;
private TransactionalState _state;
private Map<Integer, RotatingTransactionalState> _partitionStates = new HashMap<Integer, RotatingTransactionalState>();
private int _index;
private int _numTasks;
public Emitter(Map conf, TopologyContext context) {
_emitter = _spout.getEmitter(conf, context);
_state = TransactionalState.newUserState(conf, (String) conf.get(Config.TOPOLOGY_TRANSACTIONAL_ID), getComponentConfiguration());
_index = context.getThisTaskIndex();
_numTasks = context.getComponentTasks(context.getThisComponentId()).size();
}
@Override
public void emitBatch(final TransactionAttempt tx, final Integer partitions,
final BatchOutputCollector collector) {
for(int i=_index; i < partitions; i+=_numTasks) {
if(!_partitionStates.containsKey(i)) {
_partitionStates.put(i, new RotatingTransactionalState(_state, "" + i));
}
RotatingTransactionalState state = _partitionStates.get(i);
final int partition = i;
Object meta = state.getStateOrCreate(tx.getTransactionId(),
new RotatingTransactionalState.StateInitializer() {
@Override
public Object init(BigInteger txid, Object lastState) {
return _emitter.emitPartitionBatchNew(tx, collector, partition, lastState);
}
});
// it's null if one of:
// a) a later transaction batch was emitted before this, so we should skip this batch
// b) if didn't exist and was created (in which case the StateInitializer was invoked and
// it was emitted
if(meta!=null) {
_emitter.emitPartitionBatch(tx, collector, partition, meta);
}
}
}
@Override
public void cleanupBefore(BigInteger txid) {
for(RotatingTransactionalState state: _partitionStates.values()) {
state.cleanupBefore(txid);
}
}
@Override
public void close() {
_state.close();
_emitter.close();
}
}
@Override
public ITransactionalSpout.Coordinator getCoordinator(Map conf, TopologyContext context) {
return new Coordinator(conf, context);
}
@Override
public ITransactionalSpout.Emitter getEmitter(Map conf, TopologyContext context) {
return new Emitter(conf, context);
}
@Override
public void declareOutputFields(OutputFieldsDeclarer declarer) {
_spout.declareOutputFields(declarer);
}
@Override
public Map<String, Object> getComponentConfiguration() {
return _spout.getComponentConfiguration();
}
}