content: package backtype.storm.coordination;
  
  import backtype.storm.topology.FailedException;
  import java.util.Map.Entry;
  import backtype.storm.tuple.Values;
  import backtype.storm.generated.GlobalStreamId;
  import backtype.storm.Config;
  import java.util.Collection;
  import backtype.storm.Constants;
  import backtype.storm.generated.Grouping;
  import backtype.storm.task.IOutputCollector;
  import backtype.storm.task.OutputCollector;
  import backtype.storm.task.TopologyContext;
  import backtype.storm.topology.IRichBolt;
  import backtype.storm.topology.OutputFieldsDeclarer;
  import backtype.storm.tuple.Fields;
  import backtype.storm.tuple.Tuple;
  import backtype.storm.utils.TimeCacheMap;
  import backtype.storm.utils.Utils;
  import java.io.Serializable;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.HashMap;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import org.apache.log4j.Logger;
  import static backtype.storm.utils.Utils.get;
  
  /**
   * Coordination requires the request ids to be globally unique for awhile. This is so it doesn't get confused
   * in the case of retries.
   */
  public class CoordinatedBolt implements IRichBolt {
      public static Logger LOG = Logger.getLogger(CoordinatedBolt.class);
  
      public static interface FinishedCallback {
          void finishedId(Object id);
      }
  
      public static 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 {
          IOutputCollector _delegate;
  
          public CoordinatedOutputCollector(IOutputCollector delegate) {
              _delegate = delegate;
          }
  
          public List<Integer> emit(String stream, Collection<Tuple> anchors, List<Object> tuple) {
              List<Integer> tasks = _delegate.emit(stream, anchors, tuple);
              updateTaskCounts(tuple.get(0), tasks);
              return tasks;
          }
  
          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);
          }
  
          public void ack(Tuple tuple) {
              Object id = tuple.getValue(0);
              synchronized(_tracked) {
                  _tracked.get(id).receivedTuples++;
              }
              boolean failed = checkFinishId(tuple);
              if(failed) {
                  _delegate.fail(tuple);
              } else {
                  _delegate.ack(tuple);
              }
          }
  
          public void fail(Tuple tuple) {
              Object id = tuple.getValue(0);
              synchronized(_tracked) {
                  _tracked.get(id).failed = true;
              }
              _delegate.fail(tuple);
          }
          
          public void reportError(Throwable error) {
              _delegate.reportError(error);
          }
  
  
          private void updateTaskCounts(Object id, List<Integer> tasks) {
              synchronized(_tracked) {
          	TrackingInfo tinfo=_tracked.get(id);
          	if(tinfo!=null)
          	{
                      Map<Integer, Integer> taskEmittedTuples = tinfo.taskEmittedTuples;
                      for(Integer task: tasks) {
                          int newCount = get(taskEmittedTuples, task, 0) + 1;
                          taskEmittedTuples.put(task, newCount);
                      }
          	}
              }
          }
      }
  
      private Map<String, SourceArgs> _sourceArgs;
      private IdStreamSpec _idStreamSpec;
      private IRichBolt _delegate;
      private Integer _numSourceReports;
      private List<Integer> _countOutTasks = new ArrayList<Integer>();;
      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<Integer, Integer>();
          boolean receivedId = false;
          boolean finished = false;
          
          @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<String, SourceArgs>();
          _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<Object, TrackingInfo>(context.maxTopologyMessageTimeout(config), 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) {
          boolean ret = false;
          Object id = tup.getValue(0);
          synchronized(_tracked) {
              TrackingInfo track = _tracked.get(id);
              try {
                  // if it timed out, then obviously it failed (hence the null check)
                  if(track==null || track.failed) ret = true;
                  if(track!=null
                          && !track.failed
                          && track.receivedId 
                          && (_sourceArgs.isEmpty()
                              ||
                             track.reportCount==_numSourceReports &&
                             track.expectedTupleCount == track.receivedTuples)) {
                      if(_delegate instanceof FinishedCallback) {
                          ((FinishedCallback)_delegate).finishedId(id);
                      }
                      if(!(_sourceArgs.isEmpty() ||
                           tup.getSourceStreamId().equals(Constants.COORDINATED_STREAM_ID) ||
                           (_idStreamSpec!=null && tup.getSourceGlobalStreamid().equals(_idStreamSpec._id))
                           )) {
                          throw new IllegalStateException("Coordination condition met on a non-coordinating tuple. Should be impossible");
                      }
                      Iterator<Integer> outTasks = _countOutTasks.iterator();
                      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));
                      }
                      track.finished = true;
                      _tracked.remove(id);
                  }
              } catch(FailedException e) {
                  LOG.error("Failed to finish batch", e);
                  track.failed = true;
                  ret = true;
              }
          }
          return ret;
      }
  
      public void execute(Tuple tuple) {
          Object id = tuple.getValue(0);
          TrackingInfo track;
          synchronized(_tracked) {
              track = _tracked.get(id);
              if(track==null) {
                  track = new TrackingInfo();
                  if(_idStreamSpec==null) track.receivedId = true;
                  _tracked.put(id, track);
              }
          }
          
          if(_idStreamSpec!=null
                  && tuple.getSourceGlobalStreamid().equals(_idStreamSpec._id)) {
              synchronized(_tracked) {
                  track.receivedId = true;
              }
              boolean failed = checkFinishId(tuple);
              if(failed) {
                  _collector.fail(tuple);
              } else {
                  _collector.ack(tuple);
              }
              
          } else if(!_sourceArgs.isEmpty()
                  && tuple.getSourceStreamId().equals(Constants.COORDINATED_STREAM_ID)) {
              int count = (Integer) tuple.getValue(1);
              synchronized(_tracked) {
                  track.reportCount++;
                  track.expectedTupleCount+=count;
              }
              boolean failed = checkFinishId(tuple);
              if(failed) {
                  _collector.fail(tuple);
              } else {
                  _collector.ack(tuple);
              }
          } else {            
              synchronized(_tracked) {
                  _delegate.execute(tuple);
              }
          }
      }
  
      public void cleanup() {
          _delegate.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<String, SourceArgs>();
          ret.put(sourceComponent, sourceArgs);
          return ret;
      }
      
      private class TimeoutItems implements TimeCacheMap.ExpiredCallback<Object, TrackingInfo> {
          @Override
          public void expire(Object id, TrackingInfo val) {
              synchronized(_tracked) {
                  // make sure we don't time out something that has been finished. the combination of 
                  // the flag and the lock ensure this
                  if(!val.finished) {
                      ((TimeoutCallback) _delegate).timeoutId(id);
                  }
              }
          }
      } 
  }
  