content: /**
   * 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 com.alibaba.jstorm.transactional.spout;
  
  import backtype.storm.spout.SpoutOutputCollectorCb;
  import backtype.storm.task.ICollectorCallback;
  import backtype.storm.tuple.Values;
  
  import com.alibaba.jstorm.transactional.BatchSnapshot;
  import com.alibaba.jstorm.transactional.TransactionCommon;
  import com.alibaba.jstorm.transactional.utils.CountValue;
  import com.alibaba.jstorm.utils.JStormUtils;
  
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Set;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  public class TransactionSpoutOutputCollector extends SpoutOutputCollectorCb {
      public static Logger LOG = LoggerFactory.getLogger(TransactionSpoutOutputCollector.class);
  
      SpoutOutputCollectorCb delegate;
      private TransactionSpout spout;
  
      private long currBatchId;
      private Map<Integer, CountValue> msgCount;
  
      private class CollectorCallback implements ICollectorCallback {
          @Override
          public void execute(String stream, List<Integer> outTasks, List values) {
              for (Integer task : outTasks) {
                  CountValue count = msgCount.get(task);
                  count.count++;
              }
          }
      }
  
      public TransactionSpoutOutputCollector(SpoutOutputCollectorCb delegate, TransactionSpout spout) {
          this.delegate = delegate;
          this.spout = spout;
      }
  
      public void init(long batchId, Set<Integer> targetTasks) {
          setCurrBatchId(batchId);
          initMsgCount(targetTasks);
          LOG.info("Init output: batchId={}", currBatchId);
      }
  
      public void initMsgCount(Set<Integer> targetTasks) {
          this.msgCount = new HashMap<>();
          for (Integer task : targetTasks) {
              this.msgCount.put(task, new CountValue());
          }
      }
  
      public void setCurrBatchId(long batchId) {
          this.currBatchId = batchId;
          delegate.setBatchId(batchId);
      }
  
      public long getCurrBatchId() {
          return currBatchId;
      }
  
      public void waitActive() {
          while (!spout.isActive()) {
              JStormUtils.sleepMs(1);
          }
      }
  
      @Override
      public List<Integer> emit(String streamId, List<Object> tuple, Object messageId) {
          return emit(streamId, tuple, null, new CollectorCallback());
      }
  
      @Override
      public void emitDirect(int taskId, String streamId, List<Object> tuple, Object messageId) {
          emitDirect(taskId, streamId, tuple, null, new CollectorCallback());
      }
  
      @Override
      public List<Integer> emit(String streamId, List<Object> tuple, Object messageId, ICollectorCallback callback) {
          delegate.emit(streamId, tuple, null, (callback != null) ? callback : new CollectorCallback());
          return null;
      }
  
      @Override
      public void emitDirect(int taskId, String streamId, List<Object> tuple, Object messageId, ICollectorCallback callback) {
          delegate.emitDirect(taskId, streamId, tuple, null, (callback != null) ? callback : new CollectorCallback());
      }
  
      public List<Integer> emitByDelegate(String streamId, List<Object> tuple, Object messageId) {
          return emitByDelegate(streamId, tuple, null, null);
      }
  
      public List<Integer> emitByDelegate(String streamId, List<Object> tuple, Object messageId, ICollectorCallback callback) {
          return delegate.emit(streamId, tuple, null, callback);
      }
  
      public void emitDirectByDelegate(int taskId, String streamId, List<Object> tuple, Object messageId) {
          emitDirectByDelegate(taskId, streamId, tuple, null, null);
      }
  
      public void emitDirectByDelegate(int taskId, String streamId, List<Object> tuple, Object messageId, ICollectorCallback callback) {
          delegate.emitDirect(taskId, streamId, tuple, null, callback);
      }
  
      @Override
      public void emitDirectCtrl(int taskId, String streamId, List<Object> tuple, Object messageId) {
          delegate.emitDirectCtrl(taskId, streamId, tuple, null);
      }
  
      @Override
      public List<Integer> emitCtrl(String streamId, List<Object> tuple, Object messageId) {
          return delegate.emitCtrl(streamId, tuple, null);
      }
  
      @Override
      public void reportError(Throwable error) {
          delegate.reportError(error);
      }
  
      @Override
      public void emitBarrier() {
          spout.commit();
      }
  
      public int flushBarrier() {
          delegate.flush();
  
          int ret = 0;
          // Emit and flush barrier message to all downstream tasks
          for (Entry<Integer, CountValue> entry : msgCount.entrySet()) {
              int taskId = entry.getKey();
              int count = entry.getValue().getValueAndSet(0);
              ret += count;
  
              BatchSnapshot barrierSnapshot = new BatchSnapshot(currBatchId, count);
              emitDirectByDelegate(taskId, TransactionCommon.BARRIER_STREAM_ID, new Values(barrierSnapshot), null, null);
          }
          delegate.flush();
  
          return ret;
      }
  
      public void moveToNextBatch() {
          currBatchId++;
          delegate.setBatchId(currBatchId);
      }
  
      public void flushInitBarrier() {
          // flush pending message in outputCollector
          delegate.flush();
          delegate.setBatchId(TransactionCommon.INIT_BATCH_ID);
          BatchSnapshot barrierSnapshot = new BatchSnapshot(TransactionCommon.INIT_BATCH_ID, 0);
          for (Entry<Integer, CountValue> entry : msgCount.entrySet()) {
              entry.getValue().setValue(0);
              emitDirectByDelegate(entry.getKey(), TransactionCommon.BARRIER_STREAM_ID, new Values(barrierSnapshot), null, null);
          }
          delegate.flush();
          delegate.setBatchId(currBatchId);
      }
  
  }