contrib/src/main/java/org/apache/apex/malhar/contrib/kinesis/AbstractKinesisOutputOperator.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.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;

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

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

import com.amazonaws.AmazonClientException;

import com.amazonaws.services.kinesis.AmazonKinesisClient;
import com.amazonaws.services.kinesis.model.PutRecordRequest;
import com.amazonaws.services.kinesis.model.PutRecordsRequest;
import com.amazonaws.services.kinesis.model.PutRecordsRequestEntry;
import com.datatorrent.api.Context.OperatorContext;
import com.datatorrent.api.DefaultInputPort;
import com.datatorrent.api.Operator;
import com.datatorrent.common.util.Pair;

/**
 * Base implementation of Kinesis Output Operator. Convert tuples to records and emits to Kinesis.<br/>
 *
 * 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
 */
public abstract class AbstractKinesisOutputOperator<V, T> implements Operator, Operator.CheckpointNotificationListener
{
  private static final Logger logger = LoggerFactory.getLogger( AbstractKinesisOutputOperator.class );
  protected String streamName;
  @NotNull
  private String accessKey;
  @NotNull
  private String secretKey;
  private String endPoint;
  protected static transient AmazonKinesisClient client = null;
  protected int sendCount;
  protected boolean isBatchProcessing = true;

  /**
   * convert the value to record. the value is value of KeyValue pair.
   * @see tupleToKeyValue()
   * @param value
   * @return
   */
  protected abstract byte[] getRecord(V value);

  /**
   * convert tuple to pair of key and value. the key will be used as PartitionKey, and the value used as Data
   * @param tuple
   * @return
   */
  protected abstract Pair<String, V> tupleToKeyValue(T tuple);

  List<PutRecordsRequestEntry> putRecordsRequestEntryList = new ArrayList<PutRecordsRequestEntry>();
  // Max size of each record: 50KB, Max size of putRecords: 4.5MB
  // So, default capacity would be 4.5MB/50KB = 92
  @Min(2)
  @Max(500)
  protected int batchSize = 92;
  /**
   * Implement Operator Interface.
   */
  @Override
  public void beginWindow(long windowId)
  {
  }

  @Override
  public void endWindow()
  {
  }

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

  /**
   * Implement Operator Interface.
   */
  @Override
  public void beforeCheckpoint(long windowId)
  {
    if (isBatchProcessing && putRecordsRequestEntryList.size() != 0) {
      try {
        flushRecords();
      } catch (AmazonClientException e) {
        throw new RuntimeException(e);
      }
    }
  }

  @Override
  public void checkpointed(long windowId)
  {
  }

  @Override
  public void committed(long windowId)
  {
  }

  /**
   * Implement Component Interface.
   *
   * @param context
   */
  @Override
  public void setup(OperatorContext context)
  {
    try {
      KinesisUtil.getInstance().createKinesisClient(accessKey, secretKey, endPoint);
    } catch (Exception e) {
      throw new RuntimeException("Unable to load Credentials", e);
    }
    this.setClient(KinesisUtil.getInstance().getClient());
    if (isBatchProcessing) {
      putRecordsRequestEntryList.clear();
    }
  }

  /**
   * This input port receives tuples that will be written out to Kinesis.
   */
  public final transient DefaultInputPort<T> inputPort = new DefaultInputPort<T>()
  {
    @Override
    public void process(T tuple)
    {
      processTuple( tuple );
    }

  };

  public void processTuple(T tuple)
  {
    // Send out single data
    try {
      if (isBatchProcessing) {
        if (putRecordsRequestEntryList.size() == batchSize) {
          flushRecords();
          logger.debug( "flushed {} records.", batchSize );
        }
        addRecord(tuple);
      } else {
        Pair<String, V> keyValue = tupleToKeyValue(tuple);
        PutRecordRequest requestRecord = new PutRecordRequest();
        requestRecord.setStreamName(streamName);
        requestRecord.setPartitionKey(keyValue.first);
        requestRecord.setData(ByteBuffer.wrap(getRecord(keyValue.second)));
        client.putRecord(requestRecord);
      }
      sendCount++;
    } catch (AmazonClientException e) {
      throw new RuntimeException(e);
    }
  }


  private void addRecord(T tuple)
  {
    try {
      Pair<String, V> keyValue = tupleToKeyValue(tuple);
      PutRecordsRequestEntry putRecordsEntry = new PutRecordsRequestEntry();
      putRecordsEntry.setData(ByteBuffer.wrap(getRecord(keyValue.second)));
      putRecordsEntry.setPartitionKey(keyValue.first);
      putRecordsRequestEntryList.add(putRecordsEntry);
    } catch (AmazonClientException e) {
      throw new RuntimeException(e);
    }
  }

  private void flushRecords()
  {
    try {
      PutRecordsRequest putRecordsRequest = new PutRecordsRequest();
      putRecordsRequest.setStreamName(streamName);
      putRecordsRequest.setRecords(putRecordsRequestEntryList);
      client.putRecords(putRecordsRequest);
      putRecordsRequestEntryList.clear();
      logger.debug( "Records flushed." );
    } catch (AmazonClientException e) {
      logger.warn( "PutRecordsRequest exception.", e );
      throw new RuntimeException(e);
    }
  }

  public void setClient(AmazonKinesisClient _client)
  {
    client = _client;
  }

  public String getStreamName()
  {
    return streamName;
  }

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

  public int getBatchSize()
  {
    return batchSize;
  }

  public void setBatchSize(int batchSize)
  {
    this.batchSize = batchSize;
  }

  public boolean isBatchProcessing()
  {
    return isBatchProcessing;
  }

  public void setBatchProcessing(boolean isBatchProcessing)
  {
    this.isBatchProcessing = isBatchProcessing;
  }

  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;
  }
}