KinesisUtils.java

package com.sumologic.kinesis.utils;

import java.util.List;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import com.amazonaws.AmazonServiceException;
import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.regions.RegionUtils;
import com.amazonaws.services.kinesis.AmazonKinesisClient;
import com.amazonaws.services.kinesis.connectors.KinesisConnectorConfiguration;
import com.amazonaws.services.kinesis.model.CreateStreamRequest;
import com.amazonaws.services.kinesis.model.DeleteStreamRequest;
import com.amazonaws.services.kinesis.model.DescribeStreamRequest;
import com.amazonaws.services.kinesis.model.ListStreamsRequest;
import com.amazonaws.services.kinesis.model.ListStreamsResult;
import com.amazonaws.services.kinesis.model.ResourceNotFoundException;

/**
 * Utilities to create and delete Amazon Kinesis streams.
 */
public class KinesisUtils {

    private static Log LOG = LogFactory.getLog(KinesisUtils.class);

    /**
     * Creates the Amazon Kinesis stream specified by config.KINESIS_INPUT_STREAM
     * 
     * @param config
     *        The configuration with the specified input stream name and {@link AWSCredentialsProvider}
     * @param shardCount
     *        The shard count to create the stream with
     */
    public static void createInputStream(KinesisConnectorConfiguration config) {
        AmazonKinesisClient kinesisClient = new AmazonKinesisClient(config.AWS_CREDENTIALS_PROVIDER);
        kinesisClient.setRegion(RegionUtils.getRegion(config.REGION_NAME));
        if (config.KINESIS_ENDPOINT != null) {
            kinesisClient.setEndpoint(config.KINESIS_ENDPOINT);
        }
        createAndWaitForStreamToBecomeAvailable(kinesisClient,
                config.KINESIS_INPUT_STREAM,
                config.KINESIS_INPUT_STREAM_SHARD_COUNT);
    }

    /**
     * Creates the Amazon Kinesis stream specified by config.KINESIS_OUTPUT_STREAM.
     * 
     * @param config
     *        The configuration with the specified output stream name and {@link AWSCredentialsProvider}
     * @param shardCount
     *        The shard count to create the stream with
     */
    public static void createOutputStream(KinesisConnectorConfiguration config) {
        AmazonKinesisClient kinesisClient = new AmazonKinesisClient(config.AWS_CREDENTIALS_PROVIDER);
        kinesisClient.setRegion(RegionUtils.getRegion(config.REGION_NAME));
        if (config.KINESIS_ENDPOINT != null) {
            kinesisClient.setEndpoint(config.KINESIS_ENDPOINT);
        }
        createAndWaitForStreamToBecomeAvailable(kinesisClient,
                config.KINESIS_OUTPUT_STREAM,
                config.KINESIS_OUTPUT_STREAM_SHARD_COUNT);
    }

    /**
     * Creates an Amazon Kinesis stream if it does not exist and waits for it to become available
     * 
     * @param kinesisClient
     *        The {@link AmazonKinesisClient} with Amazon Kinesis read and write privileges
     * @param streamName
     *        The Amazon Kinesis stream name to create
     * @param shardCount
     *        The shard count to create the stream with
     * @throws IllegalStateException
     *         Invalid Amazon Kinesis stream state
     * @throws IllegalStateException
     *         Stream does not go active before the timeout
     */
    public static void createAndWaitForStreamToBecomeAvailable(AmazonKinesisClient kinesisClient,
            String streamName,
            int shardCount) {
        if (streamExists(kinesisClient, streamName)) {
            String state = streamState(kinesisClient, streamName);
            switch (state) {
                case "DELETING":
                    long startTime = System.currentTimeMillis();
                    long endTime = startTime + 1000 * 120;
                    while (System.currentTimeMillis() < endTime && streamExists(kinesisClient, streamName)) {
                        try {
                            LOG.info("...Deleting Stream " + streamName + "...");
                            Thread.sleep(1000 * 10);
                        } catch (InterruptedException e) {
                        }
                    }
                    if (streamExists(kinesisClient, streamName)) {
                        LOG.error("KinesisUtils timed out waiting for stream " + streamName + " to delete");
                        throw new IllegalStateException("KinesisUtils timed out waiting for stream " + streamName
                                + " to delete");
                    }
                case "ACTIVE":
                    LOG.info("Stream " + streamName + " is ACTIVE");
                    return;
                case "CREATING":
                    break;
                case "UPDATING":
                    LOG.info("Stream " + streamName + " is UPDATING");
                    return;
                default:
                    throw new IllegalStateException("Illegal stream state: " + state);
            }
        } else {
            CreateStreamRequest createStreamRequest = new CreateStreamRequest();
            createStreamRequest.setStreamName(streamName);
            createStreamRequest.setShardCount(shardCount);
            kinesisClient.createStream(createStreamRequest);
            LOG.info("Stream " + streamName + " created");
        }
        long startTime = System.currentTimeMillis();
        long endTime = startTime + (10 * 60 * 1000);
        while (System.currentTimeMillis() < endTime) {
            try {
                Thread.sleep(1000 * 10);
            } catch (Exception e) {
            }
            try {
                String streamStatus = streamState(kinesisClient, streamName);
                if (streamStatus.equals("ACTIVE")) {
                    LOG.info("Stream " + streamName + " is ACTIVE");
                    return;
                }
            } catch (ResourceNotFoundException e) {
                throw new IllegalStateException("Stream " + streamName + " never went active");
            }
        }
    }

    /**
     * Helper method to determine if an Amazon Kinesis stream exists.
     * 
     * @param kinesisClient
     *        The {@link AmazonKinesisClient} with Amazon Kinesis read privileges
     * @param streamName
     *        The Amazon Kinesis stream to check for
     * @return true if the Amazon Kinesis stream exists, otherwise return false
     */
    private static boolean streamExists(AmazonKinesisClient kinesisClient, String streamName) {
        DescribeStreamRequest describeStreamRequest = new DescribeStreamRequest();
        describeStreamRequest.setStreamName(streamName);
        try {
            kinesisClient.describeStream(describeStreamRequest);
            return true;
        } catch (ResourceNotFoundException e) {
            return false;
        }
    }

    /**
     * Return the state of a Amazon Kinesis stream.
     * 
     * @param kinesisClient
     *        The {@link AmazonKinesisClient} with Amazon Kinesis read privileges
     * @param streamName
     *        The Amazon Kinesis stream to get the state of
     * @return String representation of the Stream state
     */
    private static String streamState(AmazonKinesisClient kinesisClient, String streamName) {
        DescribeStreamRequest describeStreamRequest = new DescribeStreamRequest();
        describeStreamRequest.setStreamName(streamName);
        try {
            return kinesisClient.describeStream(describeStreamRequest).getStreamDescription().getStreamStatus();
        } catch (AmazonServiceException e) {
            return null;
        }
    }

    /**
     * Gets a list of all Amazon Kinesis streams
     * 
     * @param kinesisClient
     *        The {@link AmazonKinesisClient} with Amazon Kinesis read privileges
     * @return list of Amazon Kinesis streams
     */
    public static List<String> listAllStreams(AmazonKinesisClient kinesisClient) {

        ListStreamsRequest listStreamsRequest = new ListStreamsRequest();
        listStreamsRequest.setLimit(10);
        ListStreamsResult listStreamsResult = kinesisClient.listStreams(listStreamsRequest);
        List<String> streamNames = listStreamsResult.getStreamNames();
        while (listStreamsResult.isHasMoreStreams()) {
            if (streamNames.size() > 0) {
                listStreamsRequest.setExclusiveStartStreamName(streamNames.get(streamNames.size() - 1));
            }

            listStreamsResult = kinesisClient.listStreams(listStreamsRequest);
            streamNames.addAll(listStreamsResult.getStreamNames());
        }
        return streamNames;
    }

    /**
     * Deletes the input stream specified by config.KINESIS_INPUT_STREAM
     * 
     * @param config
     *        The configuration containing the stream name and {@link AWSCredentialsProvider}
     */
    public static void deleteInputStream(KinesisConnectorConfiguration config) {
        AmazonKinesisClient kinesisClient = new AmazonKinesisClient(config.AWS_CREDENTIALS_PROVIDER);
        kinesisClient.setRegion(RegionUtils.getRegion(config.REGION_NAME));
        if (config.KINESIS_ENDPOINT != null) {
            kinesisClient.setEndpoint(config.KINESIS_ENDPOINT);
        }
        deleteStream(kinesisClient, config.KINESIS_INPUT_STREAM);
    }

    /**
     * Deletes the output stream specified by config.KINESIS_OUTPUT_STREAM
     * 
     * @param config
     *        The configuration containing the stream name and {@link AWSCredentialsProvider}
     */
    public static void deleteOutputStream(KinesisConnectorConfiguration config) {
        AmazonKinesisClient kinesisClient = new AmazonKinesisClient(config.AWS_CREDENTIALS_PROVIDER);
        kinesisClient.setRegion(RegionUtils.getRegion(config.REGION_NAME));
        if (config.KINESIS_ENDPOINT != null) {
            kinesisClient.setEndpoint(config.KINESIS_ENDPOINT);
        }
        deleteStream(kinesisClient, config.KINESIS_OUTPUT_STREAM);
    }

    /**
     * Deletes an Amazon Kinesis stream if it exists.
     * 
     * @param kinesisClient
     *        The {@link AmazonKinesisClient} with Amazon Kinesis read and write privileges
     * @param streamName
     *        The Amazon Kinesis stream to delete
     */
    public static void deleteStream(AmazonKinesisClient kinesisClient, String streamName) {
        if (streamExists(kinesisClient, streamName)) {
            DeleteStreamRequest deleteStreamRequest = new DeleteStreamRequest();
            deleteStreamRequest.setStreamName(streamName);
            kinesisClient.deleteStream(deleteStreamRequest);
            LOG.info("Deleting stream " + streamName);
        } else {
            LOG.warn("Stream " + streamName + " does not exist");
        }
    }

}