contrib/src/main/java/org/apache/apex/malhar/contrib/kinesis/AbstractKinesisInputOperator.java

/**
 * 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 org.apache.apex.malhar.contrib.kinesis;

import java.io.IOException;
import java.lang.reflect.Array;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.validation.Valid;
import javax.validation.constraints.Min;
import javax.validation.constraints.NotNull;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.apex.malhar.lib.util.KryoCloneUtils;
import org.apache.apex.malhar.lib.wal.WindowDataManager;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.MutablePair;

import com.amazonaws.services.kinesis.model.Record;
import com.amazonaws.services.kinesis.model.Shard;
import com.amazonaws.services.kinesis.model.ShardIteratorType;
import com.google.common.collect.Sets;

import com.datatorrent.api.Context.OperatorContext;
import com.datatorrent.api.DefaultOutputPort;
import com.datatorrent.api.DefaultPartition;
import com.datatorrent.api.InputOperator;
import com.datatorrent.api.Operator;
import com.datatorrent.api.Operator.ActivationListener;
import com.datatorrent.api.Partitioner;
import com.datatorrent.api.Stats;
import com.datatorrent.api.StatsListener;
import com.datatorrent.api.annotation.Stateless;
import com.datatorrent.common.util.Pair;

/**
 * Base implementation of Kinesis Input Operator. Fetches records from kinesis and emits them as tuples.<br/>
 * <p>
 * <b>Partition Strategy:</b>
 * <p><b>1. ONE_TO_ONE partition</b> Each operator partition will consume from only one Kinesis shard </p>
 * <p><b>2. ONE_TO_MANY partition</b> Each operator partition will consume from more than one kinesis
 *    shard. Dynamic partition is enable by setting the {@link #shardsPerPartition} value > 1</p>
 * <p/>
 * Configurations:<br/>
 * {@link #accessKey} : AWS Credentials AccessKeyId <br/>
 * {@link #secretKey} : AWS Credentials SecretAccessKey <br/>
 * streamName : Name of the stream from where the records to be accessed
 *
 * @param <T>
 * @since 2.0.0
 */
@SuppressWarnings("rawtypes")
public abstract class AbstractKinesisInputOperator<T> implements InputOperator, ActivationListener<OperatorContext>, Partitioner<AbstractKinesisInputOperator>, StatsListener, Operator.CheckpointNotificationListener
{
  private static final Logger logger = LoggerFactory.getLogger(AbstractKinesisInputOperator.class);

  @Min(1)
  private int maxTuplesPerWindow = Integer.MAX_VALUE;
  private int emitCount = 0;
  @NotNull
  private String accessKey;
  @NotNull
  private String secretKey;

  private String endPoint;

  protected WindowDataManager windowDataManager;
  protected transient long currentWindowId;
  protected transient int operatorId;
  protected final transient Map<String, MutablePair<String, Integer>> currentWindowRecoveryState;
  @Valid
  protected KinesisConsumer consumer = new KinesisConsumer();

  // By default the partition policy is 1:1
  public PartitionStrategy strategy = PartitionStrategy.ONE_TO_ONE;

  private transient OperatorContext context = null;

  // Store the current partition info
  private transient Set<PartitionInfo> currentPartitionInfo = new HashSet<PartitionInfo>();

  protected transient Map<String, String> shardPosition = new HashMap<String, String>();
  private ShardManager shardManager = null;

  // Minimal interval between 2 (re)partition actions
  private long repartitionInterval = 30000L;

  // Minimal interval between checking collected stats and decide whether it needs to repartition or not.
  private long repartitionCheckInterval = 5000L;

  private transient long lastCheckTime = 0L;

  private transient long lastRepartitionTime = 0L;

  //No of shards per partition in dynamic MANY_TO_ONE strategy
  // If the value is more than 1, then it enables the dynamic partitioning
  @Min(1)
  private Integer shardsPerPartition = 1;

  @Min(1)
  private int initialPartitionCount = 1;

  private transient List<String> newWaitingPartition = new LinkedList<String>();

  /**
   * This output port emits tuples extracted from Kinesis data records.
   */
  public final transient DefaultOutputPort<T> outputPort = new DefaultOutputPort<T>();

  public AbstractKinesisInputOperator()
  {
    /*
     * Application may override the windowDataManger behaviour but default
     * would be NoopWindowDataManager.
     */
    windowDataManager = new WindowDataManager.NoopWindowDataManager();
    currentWindowRecoveryState = new HashMap<String, MutablePair<String, Integer>>();
  }

  /**
   * Derived class has to implement this method, so that it knows what type of message it is going to send to Malhar.
   * It converts a ByteBuffer message into a Tuple. A Tuple can be of any type (derived from Java Object) that
   * operator user intends to.
   *
   * @param rc Record to convert into tuple
   */
  public abstract T getTuple(Record rc);

  /**
   * Any concrete class derived from AbstractKinesisInputOperator may implement this method to emit tuples to an output port.
   */
  public void emitTuple(Pair<String, Record> data)
  {
    outputPort.emit(getTuple(data.getSecond()));
  }

  @Override
  public void partitioned(Map<Integer, Partition<AbstractKinesisInputOperator>> partitions)
  {
    // update the last repartition time
    lastRepartitionTime = System.currentTimeMillis();
  }

  @Override
  public Collection<Partition<AbstractKinesisInputOperator>> definePartitions(Collection<Partition<AbstractKinesisInputOperator>> partitions, PartitioningContext context)
  {
    boolean isInitialParitition = partitions.iterator().next().getStats() == null;
    // Set the credentials to get the list of shards
    if (isInitialParitition) {
      try {
        KinesisUtil.getInstance().createKinesisClient(accessKey, secretKey, endPoint);
      } catch (Exception e) {
        throw new RuntimeException("[definePartitions]: Unable to load credentials. ", e);
      }
    }
    List<Shard> shards = KinesisUtil.getInstance().getShardList(getStreamName());

    // Operator partitions
    List<Partition<AbstractKinesisInputOperator>> newPartitions = null;
    Set<Integer> deletedOperators =  Sets.newHashSet();

    // initialize the shard positions
    Map<String, String> initShardPos = null;
    if (isInitialParitition && shardManager != null) {
      initShardPos = shardManager.loadInitialShardPositions();
    }

    switch (strategy) {
    // For the 1 to 1 mapping The framework will create number of operator partitions based on kinesis shards
    // Each operator partition will consume from only one kinesis shard
      case ONE_TO_ONE:
        if (isInitialParitition) {
          lastRepartitionTime = System.currentTimeMillis();
          logger.info("[ONE_TO_ONE]: Initializing partition(s)");
          // initialize the number of operator partitions according to number of shards
          newPartitions = new ArrayList<Partition<AbstractKinesisInputOperator>>(shards.size());
          for (int i = 0; i < shards.size(); i++) {
            logger.info("[ONE_TO_ONE]: Create operator partition for kinesis partition: " + shards.get(i).getShardId() + ", StreamName: " + this.getConsumer().streamName);
            newPartitions.add(createPartition(Sets.newHashSet(shards.get(i).getShardId()), initShardPos));
          }
        } else if (newWaitingPartition.size() != 0) {
          // Remove the partitions for the closed shards
          removePartitionsForClosedShards(partitions, deletedOperators);
          // add partition for new kinesis shard
          for (String pid : newWaitingPartition) {
            logger.info("[ONE_TO_ONE]: Add operator partition for kinesis partition " + pid);
            partitions.add(createPartition(Sets.newHashSet(pid), null));
          }
          newWaitingPartition.clear();
          List<WindowDataManager> managers = windowDataManager.partition(partitions.size(), deletedOperators);
          int i = 0;
          for (Partition<AbstractKinesisInputOperator> partition : partitions) {
            partition.getPartitionedInstance().setWindowDataManager(managers.get(i));
            i++;
          }
          return partitions;
        }
        break;
    // For the N to 1 mapping The initial partition number is defined by stream application
    // Afterwards, the framework will dynamically adjust the partition
      case MANY_TO_ONE:
      /* This case was handled into two ways.
         1. Dynamic Partition: Number of DT partitions is depends on the number of open shards.
         2. Static Partition: Number of DT partitions is fixed, whether the number of shards are increased/decreased.
      */
        int size = initialPartitionCount;
        if (newWaitingPartition.size() != 0) {
          // Get the list of open shards
          shards = getOpenShards(partitions);
          if (shardsPerPartition > 1) {
            size = (int)Math.ceil(shards.size() / (shardsPerPartition * 1.0));
          }
          initShardPos = shardManager.loadInitialShardPositions();
        }
        @SuppressWarnings("unchecked")
        Set<String>[] pIds = (Set<String>[])Array.newInstance((new HashSet<String>()).getClass(), size);

        newPartitions = new ArrayList<Partition<AbstractKinesisInputOperator>>(size);
        for (int i = 0; i < shards.size(); i++) {
          Shard pm = shards.get(i);
          if (pIds[i % size] == null) {
            pIds[i % size] = new HashSet<String>();
          }
          pIds[i % size].add(pm.getShardId());
        }
        if (isInitialParitition) {
          lastRepartitionTime = System.currentTimeMillis();
          logger.info("[MANY_TO_ONE]: Initializing partition(s)");
        } else {
          logger.info("[MANY_TO_ONE]: Add operator partition for kinesis partition(s): " + StringUtils.join(newWaitingPartition, ", ") + ", StreamName: " + this.getConsumer().streamName);
          newWaitingPartition.clear();
        }
        // Add the existing partition Ids to the deleted operators
        for (Partition<AbstractKinesisInputOperator> op : partitions) {
          deletedOperators.add(op.getPartitionedInstance().operatorId);
        }

        for (int i = 0; i < pIds.length; i++) {
          logger.info("[MANY_TO_ONE]: Create operator partition for kinesis partition(s): " + StringUtils.join(pIds[i], ", ") + ", StreamName: " + this.getConsumer().streamName);

          if (pIds[i] != null) {
            newPartitions.add(createPartition(pIds[i], initShardPos));
          }
        }
        break;
      default:
        break;
    }
    int i = 0;
    List<WindowDataManager> managers = windowDataManager.partition(partitions.size(), deletedOperators);
    for (Partition<AbstractKinesisInputOperator> partition : partitions) {
      partition.getPartitionedInstance().setWindowDataManager(managers.get(i++));
    }
    return newPartitions;
  }

  @Override
  public Response processStats(BatchedOperatorStats stats)
  {
    Response resp = new Response();
    List<KinesisConsumer.KinesisShardStats> kstats = extractkinesisStats(stats);
    resp.repartitionRequired = isPartitionRequired(kstats);
    return resp;
  }

  private void updateShardPositions(List<KinesisConsumer.KinesisShardStats> kstats)
  {
    //In every partition check interval, call shardmanager to update the positions
    if (shardManager != null) {
      shardManager.updatePositions(KinesisConsumer.KinesisShardStatsUtil.getShardStatsForPartitions(kstats));
    }
  }

  private List<KinesisConsumer.KinesisShardStats> extractkinesisStats(BatchedOperatorStats stats)
  {
    //preprocess the stats
    List<KinesisConsumer.KinesisShardStats> kmsList = new LinkedList<KinesisConsumer.KinesisShardStats>();
    for (Stats.OperatorStats os : stats.getLastWindowedStats()) {
      if (os != null && os.counters instanceof KinesisConsumer.KinesisShardStats) {
        kmsList.add((KinesisConsumer.KinesisShardStats)os.counters);
      }
    }
    return kmsList;
  }

  private boolean isPartitionRequired(List<KinesisConsumer.KinesisShardStats> kstats)
  {

    long t = System.currentTimeMillis();

    if (t - lastCheckTime < repartitionCheckInterval) {
      // return false if it's within repartitionCheckInterval since last time it check the stats
      return false;
    }
    logger.debug("Use ShardManager to update the Shard Positions");
    updateShardPositions(kstats);
    if (repartitionInterval < 0) {
      // if repartition is disabled
      return false;
    }

    if (t - lastRepartitionTime < repartitionInterval) {
      // return false if it's still within repartitionInterval since last (re)partition
      return false;
    }

    try {
      // monitor if shards are repartitioned
      Set<String> existingIds = new HashSet<String>();
      for (PartitionInfo pio : currentPartitionInfo) {
        existingIds.addAll(pio.kpids);
      }
      List<Shard> shards = KinesisUtil.getInstance().getShardList(getStreamName());
      for (Shard shard :shards) {
        if (!existingIds.contains(shard.getShardId())) {
          newWaitingPartition.add(shard.getShardId());
        }
      }

      if (newWaitingPartition.size() != 0) {
        // found new kinesis partition
        lastRepartitionTime = t;
        return true;
      }
      return false;
    } finally {
      // update last  check time
      lastCheckTime = System.currentTimeMillis();
    }
  }

  // If all the shards in the partition are closed, then remove that partition
  private void removePartitionsForClosedShards(Collection<Partition<AbstractKinesisInputOperator>> partitions, Set<Integer> deletedOperators)
  {
    List<Partition<AbstractKinesisInputOperator>> closedPartitions = new ArrayList<Partition<AbstractKinesisInputOperator>>();
    for (Partition<AbstractKinesisInputOperator> op : partitions) {
      if (op.getPartitionedInstance().getConsumer().getClosedShards().size() ==
          op.getPartitionedInstance().getConsumer().getNumOfShards()) {
        closedPartitions.add(op);
        deletedOperators.add(op.getPartitionedInstance().operatorId);
      }
    }
    if (closedPartitions.size() != 0) {
      for (Partition<AbstractKinesisInputOperator> op : closedPartitions) {
        partitions.remove(op);
      }
    }
  }

  // Get the list of open shards
  private List<Shard> getOpenShards(Collection<Partition<AbstractKinesisInputOperator>> partitions)
  {
    List<Shard> closedShards = new ArrayList<Shard>();
    for (Partition<AbstractKinesisInputOperator> op : partitions) {
      closedShards.addAll(op.getPartitionedInstance().getConsumer().getClosedShards());
    }
    List<Shard> shards = KinesisUtil.getInstance().getShardList(getStreamName());
    List<Shard> openShards = new ArrayList<Shard>();
    for (Shard shard :shards) {
      if (!closedShards.contains(shard)) {
        openShards.add(shard);
      }
    }
    return openShards;
  }

  // Create a new partition with the shardIds and initial shard positions
  private Partition<AbstractKinesisInputOperator> createPartition(Set<String> shardIds, Map<String, String> initShardPos)
  {
    Partition<AbstractKinesisInputOperator> p = new DefaultPartition<AbstractKinesisInputOperator>(KryoCloneUtils.cloneObject(this));
    p.getPartitionedInstance().getConsumer().setShardIds(shardIds);
    p.getPartitionedInstance().getConsumer().resetShardPositions(initShardPos);

    PartitionInfo pif = new PartitionInfo();
    pif.kpids = shardIds;
    currentPartitionInfo.add(pif);
    return p;
  }

  /**
   * Implement Component Interface.
   *
   * @param context
   */
  @Override
  public void setup(OperatorContext context)
  {
    this.context = context;
    try {
      KinesisUtil.getInstance().createKinesisClient(accessKey, secretKey, endPoint);
    } catch (Exception e) {
      throw new RuntimeException(e);
    }
    consumer.create();
    operatorId = context.getId();
    windowDataManager.setup(context);
    shardPosition.clear();
  }

  /**
   * Implement Component Interface.
   */
  @Override
  public void teardown()
  {
    windowDataManager.teardown();
    consumer.teardown();
  }

  /**
   * Implement Operator Interface.
   */
  @Override
  public void beginWindow(long windowId)
  {
    emitCount = 0;
    currentWindowId = windowId;
    if (windowId <= windowDataManager.getLargestCompletedWindow()) {
      replay(windowId);
    }
  }

  protected void replay(long windowId)
  {
    try {
      @SuppressWarnings("unchecked")
      Map<String, MutablePair<String, Integer>> recoveredData =
          (Map<String, MutablePair<String, Integer>>)windowDataManager.retrieve(windowId);
      if (recoveredData == null) {
        return;
      }
      for (Map.Entry<String, MutablePair<String, Integer>> rc: recoveredData.entrySet()) {
        logger.debug("Replaying the windowId: {}", windowId);
        logger.debug("ShardId: " + rc.getKey() + " , Start Sequence Id: " + rc.getValue().getLeft() + " , No Of Records: " + rc.getValue().getRight());
        try {
          List<Record> records = KinesisUtil.getInstance().getRecords(consumer.streamName, rc.getValue().getRight(),
              rc.getKey(), ShardIteratorType.AT_SEQUENCE_NUMBER, rc.getValue().getLeft());
          for (Record record : records) {
            emitTuple(new Pair<String, Record>(rc.getKey(), record));
            shardPosition.put(rc.getKey(), record.getSequenceNumber());
          }
        } catch (Exception e) {
          throw new RuntimeException(e);
        }
      }

      /*
       * Set the shard positions and start the consumer if last recovery windowid
       * match with current completed windowid.
       */
      if (windowId == windowDataManager.getLargestCompletedWindow()) {
        // Set the shard positions to the consumer
        Map<String, String> statsData = new HashMap<String, String>(getConsumer().getShardPosition());
        statsData.putAll(shardPosition);
        getConsumer().resetShardPositions(statsData);
        consumer.start();
      }
    } catch (IOException e) {
      throw new RuntimeException("replay", e);
    }
  }

  /**
   * Implement Operator Interface.
   */
  @Override
  public void endWindow()
  {
    if (currentWindowId > windowDataManager.getLargestCompletedWindow()) {
      context.setCounters(getConsumer().getConsumerStats(shardPosition));
      try {
        windowDataManager.save(currentWindowRecoveryState, currentWindowId);
      } catch (IOException e) {
        throw new RuntimeException("saving recovery", e);
      }
    }
    currentWindowRecoveryState.clear();

  }

  /**
   * Implement ActivationListener Interface.
   */
  @Override
  public void activate(OperatorContext ctx)
  {
    // If it is a replay state, don't start the consumer
    if (context.getValue(OperatorContext.ACTIVATION_WINDOW_ID) != Stateless.WINDOW_ID &&
        context.getValue(OperatorContext.ACTIVATION_WINDOW_ID) < windowDataManager.getLargestCompletedWindow()) {
      return;
    }
    consumer.start();
  }

  @Override
  public void committed(long windowId)
  {
    try {
      windowDataManager.committed(windowId);
    } catch (IOException e) {
      throw new RuntimeException("deleting state", e);
    }
  }

  @Override
  public void checkpointed(long windowId)
  {
  }

  @Override
  public void beforeCheckpoint(long windowId)
  {
  }

  /**
   * Implement ActivationListener Interface.
   */
  @Override
  public void deactivate()
  {
    consumer.stop();
  }

  /**
   * Implement InputOperator Interface.
   */
  @Override
  public void emitTuples()
  {
    if (currentWindowId <= windowDataManager.getLargestCompletedWindow()) {
      return;
    }
    int count = consumer.getQueueSize();
    if (maxTuplesPerWindow > 0) {
      count = Math.min(count, maxTuplesPerWindow - emitCount);
    }
    for (int i = 0; i < count; i++) {
      Pair<String, Record> data = consumer.pollRecord();
      String shardId = data.getFirst();
      String recordId = data.getSecond().getSequenceNumber();
      emitTuple(data);
      MutablePair<String, Integer> shardOffsetAndCount = currentWindowRecoveryState.get(shardId);
      if (shardOffsetAndCount == null) {
        currentWindowRecoveryState.put(shardId, new MutablePair<String, Integer>(recordId, 1));
      } else {
        shardOffsetAndCount.setRight(shardOffsetAndCount.right + 1);
      }
      shardPosition.put(shardId, recordId);
    }
    emitCount += count;
  }

  public static enum PartitionStrategy
  {
    /**
     * Each operator partition connect to only one kinesis Shard
     */
    ONE_TO_ONE,
    /**
     * Each operator consumes from several Shards.
     */
    MANY_TO_ONE
  }

  static class PartitionInfo
  {
    Set<String> kpids;
  }

  public void setConsumer(KinesisConsumer consumer)
  {
    this.consumer = consumer;
  }

  public KinesisConsumer getConsumer()
  {
    return consumer;
  }

  public String getStreamName()
  {
    return this.consumer.getStreamName();
  }

  public void setStreamName(String streamName)
  {
    this.consumer.setStreamName(streamName);
  }

  public int getMaxTuplesPerWindow()
  {
    return maxTuplesPerWindow;
  }

  public void setMaxTuplesPerWindow(int maxTuplesPerWindow)
  {
    this.maxTuplesPerWindow = maxTuplesPerWindow;
  }

  public PartitionStrategy getStrategy()
  {
    return strategy;
  }

  public void setStrategy(String policy)
  {
    this.strategy = PartitionStrategy.valueOf(policy.toUpperCase());
  }

  public OperatorContext getContext()
  {
    return context;
  }

  public void setContext(OperatorContext context)
  {
    this.context = context;
  }

  public ShardManager getShardManager()
  {
    return shardManager;
  }

  public void setShardManager(ShardManager shardManager)
  {
    this.shardManager = shardManager;
  }

  public long getRepartitionInterval()
  {
    return repartitionInterval;
  }

  public void setRepartitionInterval(long repartitionInterval)
  {
    this.repartitionInterval = repartitionInterval;
  }

  public long getRepartitionCheckInterval()
  {
    return repartitionCheckInterval;
  }

  public void setRepartitionCheckInterval(long repartitionCheckInterval)
  {
    this.repartitionCheckInterval = repartitionCheckInterval;
  }

  public Integer getShardsPerPartition()
  {
    return shardsPerPartition;
  }

  public void setShardsPerPartition(Integer shardsPerPartition)
  {
    this.shardsPerPartition = shardsPerPartition;
  }

  public int getInitialPartitionCount()
  {
    return initialPartitionCount;
  }

  public void setInitialPartitionCount(int initialPartitionCount)
  {
    this.initialPartitionCount = initialPartitionCount;
  }

  public void setInitialOffset(String initialOffset)
  {
    this.consumer.initialOffset = initialOffset;
  }

  public String getAccessKey()
  {
    return accessKey;
  }

  public void setAccessKey(String accessKey)
  {
    this.accessKey = accessKey;
  }

  public String getSecretKey()
  {
    return secretKey;
  }

  public void setSecretKey(String secretKey)
  {
    this.secretKey = secretKey;
  }

  public String getEndPoint()
  {
    return endPoint;
  }

  public void setEndPoint(String endPoint)
  {
    this.endPoint = endPoint;
  }

  public WindowDataManager getWindowDataManager()
  {
    return windowDataManager;
  }

  public void setWindowDataManager(WindowDataManager windowDataManager)
  {
    this.windowDataManager = windowDataManager;
  }
}