Search in sources :

Example 1 with PutRecordsRequest

use of com.amazonaws.services.kinesis.model.PutRecordsRequest in project flink by apache.

the class ManualExactlyOnceWithStreamReshardingTest method main.

public static void main(String[] args) throws Exception {
    final ParameterTool pt = ParameterTool.fromArgs(args);
    LOG.info("Starting exactly once with stream resharding test");
    final String streamName = "flink-test-" + UUID.randomUUID().toString();
    final String accessKey = pt.getRequired("accessKey");
    final String secretKey = pt.getRequired("secretKey");
    final String region = pt.getRequired("region");
    final Properties configProps = new Properties();
    configProps.setProperty(ConsumerConfigConstants.AWS_ACCESS_KEY_ID, accessKey);
    configProps.setProperty(ConsumerConfigConstants.AWS_SECRET_ACCESS_KEY, secretKey);
    configProps.setProperty(ConsumerConfigConstants.AWS_REGION, region);
    configProps.setProperty(ConsumerConfigConstants.SHARD_DISCOVERY_INTERVAL_MILLIS, "0");
    final AmazonKinesisClient client = AWSUtil.createKinesisClient(configProps);
    // the stream is first created with 1 shard
    client.createStream(streamName, 1);
    // wait until stream has been created
    DescribeStreamResult status = client.describeStream(streamName);
    LOG.info("status {}", status);
    while (!status.getStreamDescription().getStreamStatus().equals("ACTIVE")) {
        status = client.describeStream(streamName);
        LOG.info("Status of stream {}", status);
        Thread.sleep(1000);
    }
    final Configuration flinkConfig = new Configuration();
    flinkConfig.setInteger(ConfigConstants.LOCAL_NUMBER_TASK_MANAGER, 1);
    flinkConfig.setInteger(ConfigConstants.TASK_MANAGER_NUM_TASK_SLOTS, 8);
    flinkConfig.setInteger(ConfigConstants.TASK_MANAGER_MEMORY_SIZE_KEY, 16);
    flinkConfig.setString(ConfigConstants.RESTART_STRATEGY_FIXED_DELAY_DELAY, "0 s");
    LocalFlinkMiniCluster flink = new LocalFlinkMiniCluster(flinkConfig, false);
    flink.start();
    final int flinkPort = flink.getLeaderRPCPort();
    try {
        // we have to use a manual generator here instead of the FlinkKinesisProducer
        // because the FlinkKinesisProducer currently has a problem where records will be resent to a shard
        // when resharding happens; this affects the consumer exactly-once validation test and will never pass
        final AtomicReference<Throwable> producerError = new AtomicReference<>();
        Runnable manualGenerate = new Runnable() {

            @Override
            public void run() {
                AmazonKinesisClient client = AWSUtil.createKinesisClient(configProps);
                int count = 0;
                final int batchSize = 30;
                while (true) {
                    try {
                        Thread.sleep(10);
                        Set<PutRecordsRequestEntry> batch = new HashSet<>();
                        for (int i = count; i < count + batchSize; i++) {
                            if (i >= TOTAL_EVENT_COUNT) {
                                break;
                            }
                            batch.add(new PutRecordsRequestEntry().withData(ByteBuffer.wrap(((i) + "-" + RandomStringUtils.randomAlphabetic(12)).getBytes(ConfigConstants.DEFAULT_CHARSET))).withPartitionKey(UUID.randomUUID().toString()));
                        }
                        count += batchSize;
                        PutRecordsResult result = client.putRecords(new PutRecordsRequest().withStreamName(streamName).withRecords(batch));
                        // and let this test fail
                        if (result.getFailedRecordCount() > 0) {
                            producerError.set(new RuntimeException("The producer has failed records in one of the put batch attempts."));
                            break;
                        }
                        if (count >= TOTAL_EVENT_COUNT) {
                            break;
                        }
                    } catch (Exception e) {
                        producerError.set(e);
                    }
                }
            }
        };
        Thread producerThread = new Thread(manualGenerate);
        producerThread.start();
        final AtomicReference<Throwable> consumerError = new AtomicReference<>();
        Thread consumerThread = ExactlyOnceValidatingConsumerThread.create(TOTAL_EVENT_COUNT, 10000, 2, 500, 500, accessKey, secretKey, region, streamName, consumerError, flinkPort, flinkConfig);
        consumerThread.start();
        // reshard the Kinesis stream while the producer / and consumers are running
        Runnable splitShard = new Runnable() {

            @Override
            public void run() {
                try {
                    // first, split shard in the middle of the hash range
                    Thread.sleep(5000);
                    LOG.info("Splitting shard ...");
                    client.splitShard(streamName, KinesisShardIdGenerator.generateFromShardOrder(0), "170141183460469231731687303715884105727");
                    // wait until the split shard operation finishes updating ...
                    DescribeStreamResult status;
                    Random rand = new Random();
                    do {
                        status = null;
                        while (status == null) {
                            // retry until we get status
                            try {
                                status = client.describeStream(streamName);
                            } catch (LimitExceededException lee) {
                                LOG.warn("LimitExceededException while describing stream ... retrying ...");
                                Thread.sleep(rand.nextInt(1200));
                            }
                        }
                    } while (!status.getStreamDescription().getStreamStatus().equals("ACTIVE"));
                    // then merge again
                    Thread.sleep(7000);
                    LOG.info("Merging shards ...");
                    client.mergeShards(streamName, KinesisShardIdGenerator.generateFromShardOrder(1), KinesisShardIdGenerator.generateFromShardOrder(2));
                } catch (InterruptedException iex) {
                //
                }
            }
        };
        Thread splitShardThread = new Thread(splitShard);
        splitShardThread.start();
        boolean deadlinePassed = false;
        // wait at most for five minutes
        long deadline = System.currentTimeMillis() + (1000 * 5 * 60);
        // wait until both producer and consumer finishes, or an unexpected error is thrown
        while ((consumerThread.isAlive() || producerThread.isAlive()) && (producerError.get() == null && consumerError.get() == null)) {
            Thread.sleep(1000);
            if (System.currentTimeMillis() >= deadline) {
                LOG.warn("Deadline passed");
                deadlinePassed = true;
                // enough waiting
                break;
            }
        }
        if (producerThread.isAlive()) {
            producerThread.interrupt();
        }
        if (consumerThread.isAlive()) {
            consumerThread.interrupt();
        }
        if (producerError.get() != null) {
            LOG.info("+++ TEST failed! +++");
            throw new RuntimeException("Producer failed", producerError.get());
        }
        if (consumerError.get() != null) {
            LOG.info("+++ TEST failed! +++");
            throw new RuntimeException("Consumer failed", consumerError.get());
        }
        if (!deadlinePassed) {
            LOG.info("+++ TEST passed! +++");
        } else {
            LOG.info("+++ TEST failed! +++");
        }
    } finally {
        client.deleteStream(streamName);
        client.shutdown();
        // stopping flink
        flink.stop();
    }
}
Also used : ParameterTool(org.apache.flink.api.java.utils.ParameterTool) PutRecordsRequestEntry(com.amazonaws.services.kinesis.model.PutRecordsRequestEntry) Configuration(org.apache.flink.configuration.Configuration) Properties(java.util.Properties) LocalFlinkMiniCluster(org.apache.flink.runtime.minicluster.LocalFlinkMiniCluster) PutRecordsResult(com.amazonaws.services.kinesis.model.PutRecordsResult) Random(java.util.Random) PutRecordsRequest(com.amazonaws.services.kinesis.model.PutRecordsRequest) HashSet(java.util.HashSet) AmazonKinesisClient(com.amazonaws.services.kinesis.AmazonKinesisClient) AtomicReference(java.util.concurrent.atomic.AtomicReference) LimitExceededException(com.amazonaws.services.kinesis.model.LimitExceededException) ExactlyOnceValidatingConsumerThread(org.apache.flink.streaming.connectors.kinesis.testutils.ExactlyOnceValidatingConsumerThread) LimitExceededException(com.amazonaws.services.kinesis.model.LimitExceededException) DescribeStreamResult(com.amazonaws.services.kinesis.model.DescribeStreamResult)

Example 2 with PutRecordsRequest

use of com.amazonaws.services.kinesis.model.PutRecordsRequest in project beam by apache.

the class KinesisUploader method uploadAll.

public static void uploadAll(List<String> data, KinesisTestOptions options) {
    AmazonKinesisClient client = new AmazonKinesisClient(new StaticCredentialsProvider(new BasicAWSCredentials(options.getAwsAccessKey(), options.getAwsSecretKey()))).withRegion(Regions.fromName(options.getAwsKinesisRegion()));
    List<List<String>> partitions = Lists.partition(data, MAX_NUMBER_OF_RECORDS_IN_BATCH);
    for (List<String> partition : partitions) {
        List<PutRecordsRequestEntry> allRecords = newArrayList();
        for (String row : partition) {
            allRecords.add(new PutRecordsRequestEntry().withData(ByteBuffer.wrap(row.getBytes(Charsets.UTF_8))).withPartitionKey(Integer.toString(row.hashCode())));
        }
        PutRecordsResult result;
        do {
            result = client.putRecords(new PutRecordsRequest().withStreamName(options.getAwsKinesisStream()).withRecords(allRecords));
            List<PutRecordsRequestEntry> failedRecords = newArrayList();
            int i = 0;
            for (PutRecordsResultEntry row : result.getRecords()) {
                if (row.getErrorCode() != null) {
                    failedRecords.add(allRecords.get(i));
                }
                ++i;
            }
            allRecords = failedRecords;
        } while (result.getFailedRecordCount() > 0);
    }
}
Also used : AmazonKinesisClient(com.amazonaws.services.kinesis.AmazonKinesisClient) PutRecordsRequestEntry(com.amazonaws.services.kinesis.model.PutRecordsRequestEntry) StaticCredentialsProvider(com.amazonaws.internal.StaticCredentialsProvider) BasicAWSCredentials(com.amazonaws.auth.BasicAWSCredentials) PutRecordsResult(com.amazonaws.services.kinesis.model.PutRecordsResult) List(java.util.List) Lists.newArrayList(com.google.common.collect.Lists.newArrayList) PutRecordsResultEntry(com.amazonaws.services.kinesis.model.PutRecordsResultEntry) PutRecordsRequest(com.amazonaws.services.kinesis.model.PutRecordsRequest)

Aggregations

AmazonKinesisClient (com.amazonaws.services.kinesis.AmazonKinesisClient)2 PutRecordsRequest (com.amazonaws.services.kinesis.model.PutRecordsRequest)2 PutRecordsRequestEntry (com.amazonaws.services.kinesis.model.PutRecordsRequestEntry)2 PutRecordsResult (com.amazonaws.services.kinesis.model.PutRecordsResult)2 BasicAWSCredentials (com.amazonaws.auth.BasicAWSCredentials)1 StaticCredentialsProvider (com.amazonaws.internal.StaticCredentialsProvider)1 DescribeStreamResult (com.amazonaws.services.kinesis.model.DescribeStreamResult)1 LimitExceededException (com.amazonaws.services.kinesis.model.LimitExceededException)1 PutRecordsResultEntry (com.amazonaws.services.kinesis.model.PutRecordsResultEntry)1 Lists.newArrayList (com.google.common.collect.Lists.newArrayList)1 HashSet (java.util.HashSet)1 List (java.util.List)1 Properties (java.util.Properties)1 Random (java.util.Random)1 AtomicReference (java.util.concurrent.atomic.AtomicReference)1 ParameterTool (org.apache.flink.api.java.utils.ParameterTool)1 Configuration (org.apache.flink.configuration.Configuration)1 LocalFlinkMiniCluster (org.apache.flink.runtime.minicluster.LocalFlinkMiniCluster)1 ExactlyOnceValidatingConsumerThread (org.apache.flink.streaming.connectors.kinesis.testutils.ExactlyOnceValidatingConsumerThread)1