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();
}
}
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);
}
}
Aggregations